|
气象海洋领域，常常会涉及到大规模的数据处理，比如高时空分辨率的模式数据、雷达数据和卫星数据等。
- S8 t' X0 b* k
# K0 G2 P) ~! b 大部分情况，我们一般只会涉及小规模数据处理，对计算效率并不会太过追求。但是当数据量变大时，低效的数据处理所耗费的时间非常明显，因此高效的数据处理方式尤为重要。
6 `! d, o/ G: `, R8 Q7 S; s. W% s
本篇以拟合一个高维数据的正太分布参数为例，介绍如何使用xarray+dask加速数据处理。
6 H' @* n2 ^/ X4 j% _- v5 w
" I2 M; Q9 j% a6 K 数据维度X[time, lev, lat, lon]，需要对三维空间每一点，沿着时间维度做正太分布拟合(正太分布拟合只是作为例子，这里可以定义你需要的操作函数)。 6 H& R v$ N5 L& O" N' [0 f/ c
其中几个关键点解释一下：
: {* V7 Z# A- r) y ~" M9 y （1）首先定义拟合正太分布的函数 ! O8 c' Y" }; u$ a" @& a Y# `7 f
def norm_fit(data):
6 Z+ Q3 t9 N2 ]6 c+ W$ z loc, scale = st.norm.fit(data)
; ~: v- D2 {6 V% q W3 j# G/ t5 U return np.array([loc, scale]) , p4 Q8 H& x' h6 Z/ K& M6 a
这里需要注意的是，拟合的函数，输出参数，需要打包为一个数组。并且它的维度需要和后面aplly_ufun定义的输出维度一致。
% k9 R# l9 j" U! p# D （2）xarray分块
8 N( ~3 y/ }- b" t, H4 z( H x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])
6 ?0 m' j% B8 q% m x = x.chunk({"lev":5, "lat":100,"lon":100}) 9 O. D/ w+ ~# J4 v
xarray结合dask可以将一个大型数组分成一个个数据块（chunk），需要注意的是我们需要沿着时间维操作，拟合需要整个时间维度的数据，因此时间维time不能分块，只能对其他维度分块。 ' i7 A' A0 q) E
（3）xarray.apply_ufunc函数
) k- o- Z; z. S/ W, a8 g3 P& Y result = xr.apply_ufunc(norm_fit,
8 U' C' C8 |( J x, @( o1 |' B2 B
input_core_dims=[["time"]],
5 C& W% c, v5 H6 | output_core_dims=[["paras"]],
7 V. G' R8 s9 Y+ i& { dask="parallelized",* Z5 v1 G& D& `
output_dtypes=[np.float],, W# P, s* G; ^1 K0 q
dask_gufunc_kwargs={output_sizes:{"paras":2}}4 ?7 g- Q2 R4 Z! v
)
3 { J P; c7 {$ K- i apply_ufunc函数具体可以参考官网教程，这里只说使用时遇到的困难，即如何定义输出维度：输出维度是用output_core_dim定义的，将输出的拟合参数（期望和标准差）定义为paras维度，维度的大小为2，通过output_sizes参数设置。这样我们输入[time, lev, lat, lon]的数据，在每个空间点对时间维度拟合之后，输出的数据为[lev, lat, lon, paras]。（PS: 这里感谢深雨露大佬的指点） 9 p7 P: d g/ O ^0 u( v
/ b: F# |1 U. S4 `/ k
以下是全部代码： # m0 X1 y \9 O# R2 f% E
from scipy import stats as st
$ S+ t( K J6 _) u7 i, H/ Y/ ~5 {1 G; L0 D# T' b
import xarray as xr/ G i8 c7 {7 ?3 r; |. _8 z
import dask- _& Y8 c `% x/ [8 k; h
import numpy as np
8 R! g7 q5 p& O from dask.diagnostics import ProgressBar2 x1 g3 e) [) c$ |/ e
3 m0 d8 C8 A0 ~3 z/ `" \9 J
def norm_fit(data):" Z5 ^: D& l. C6 x2 d- [8 f+ W
loc, scale = st.norm.fit(data); H7 J6 @4 |; r, R- Y
return np.array([loc, scale])2 F1 {2 U0 k# D0 t5 H7 _8 G* S
; E( _. P5 m: c+ v4 t rs = np.random.RandomState(0)
4 a' s9 ~5 U5 U* A x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])# L% N( j8 J$ q% b c3 B# l% r8 y- c
x = x.chunk({"lev":5, "lat":100,"lon":100})2 V( P/ T% n$ m
, M x* t, a) G
#使用apply_ufunc计算，并用dask的并行计算
/ g: Z6 \/ p1 h result = xr.apply_ufunc(norm_fit,
; `: [# L: x9 |' u, n x,, b, B1 d' p/ D
input_core_dims=[["time"]],8 c s' f8 {2 V6 e: x6 E; |6 C
output_core_dims=[["paras"]],
& r: y; ? r+ X# [$ K4 H/ y# r dask="parallelized",8 I3 L, o0 s; z/ q C% \9 ?
output_dtypes=[np.float],
6 C, W! p$ X) l+ Q ^& ~; @8 o! h% @; T dask_gufunc_kwargs={output_sizes:{"paras":2}}1 Y% C6 f* P2 V9 l5 C# Y* M9 N
)
# x% I" Z, D) f. R: \$ E
% s; f6 H5 f" d$ @2 } #compute进行真实的计算并显示进度* H9 Q* l9 v0 R1 Y+ ?6 o
with ProgressBar():, z& G- Y1 r, u
result = results.compute()' c: Y# x( ^5 c2 R% V1 V
! J$ K" o, c6 ^8 s1 X$ m# `% g #结果冲命名保存到nc文件
/ q# p% _' W' I5 a% V) _ result = result.rename("norm_paras")
: G9 K* }* @; S9 ]5 a' O; ^ result = result.to_netcdf("norm_fit_paras.nc",compute=False)
3 B2 e1 @" v: [, t" g with ProgressBar():
8 L$ e- G, z- ~+ f3 \' t: x result.compute() # }$ }" |) |' l& G* ], r
转自：气海同途 # ?4 Q) Q3 Q- o' D9 u* z
关注【Ai尚研修科研技术平台】公众号，查看更多课程安排及免费数据资料 | 
