|
/ g0 x4 B5 a/ K. Y* O/ t 气象海洋领域，常常会涉及到大规模的数据处理，比如高时空分辨率的模式数据、雷达数据和卫星数据等。 4 U5 H7 v5 t. y
) V3 {5 z" Y5 I6 N. k 大部分情况，我们一般只会涉及小规模数据处理，对计算效率并不会太过追求。但是当数据量变大时，低效的数据处理所耗费的时间非常明显，因此高效的数据处理方式尤为重要。
* O. M- z$ V/ \1 U9 L, K1 c
; X$ `8 R/ c" X& x# T4 Y 本篇以拟合一个高维数据的正太分布参数为例，介绍如何使用xarray+dask加速数据处理。 " W" W, E( B! Y% z( G
, T, s; h U; }4 ~1 ?' y 数据维度X[time, lev, lat, lon]，需要对三维空间每一点，沿着时间维度做正太分布拟合(正太分布拟合只是作为例子，这里可以定义你需要的操作函数)。
0 f; A6 l8 P2 X- B 其中几个关键点解释一下：
- o3 \! e; A; X2 M( W. E （1）首先定义拟合正太分布的函数 * F9 W" N- \4 i" w0 ]2 Q
def norm_fit(data):
, q7 L5 q4 e/ j& z3 z loc, scale = st.norm.fit(data): T+ |! z Q. n& R
return np.array([loc, scale])
- j6 J" X+ Z7 }* c+ H$ {6 s. I/ N& j 这里需要注意的是，拟合的函数，输出参数，需要打包为一个数组。并且它的维度需要和后面aplly_ufun定义的输出维度一致。
+ X5 w9 a( d7 ~, L+ j( R4 ] （2）xarray分块
3 h' E" N; j) _. U$ f1 p0 S x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])
$ w, ~: @3 p$ y3 O2 c8 W$ h/ s- c x = x.chunk({"lev":5, "lat":100,"lon":100}) * X7 _! I# c, j% s( U) X
xarray结合dask可以将一个大型数组分成一个个数据块（chunk），需要注意的是我们需要沿着时间维操作，拟合需要整个时间维度的数据，因此时间维time不能分块，只能对其他维度分块。 0 d- L t$ b4 \; j' K) l8 g
（3）xarray.apply_ufunc函数 E, j1 S7 K9 R+ t
result = xr.apply_ufunc(norm_fit,3 v8 K0 R& k7 j3 W5 C; ?- p
x,
7 ]9 U$ z+ F3 L! e9 R: H4 E: f/ K input_core_dims=[["time"]],1 m4 O# k6 D3 W0 `
output_core_dims=[["paras"]],9 q: S2 p( s$ i' j* j) W
dask="parallelized",' K( Y' j: o# n+ z$ U
output_dtypes=[np.float],
$ b0 a- w1 z7 W0 p dask_gufunc_kwargs={output_sizes:{"paras":2}}
+ i* j6 D' C8 t2 K0 Q' u3 V- z- t )
e2 m+ T/ T2 |2 ~ apply_ufunc函数具体可以参考官网教程，这里只说使用时遇到的困难，即如何定义输出维度：输出维度是用output_core_dim定义的，将输出的拟合参数（期望和标准差）定义为paras维度，维度的大小为2，通过output_sizes参数设置。这样我们输入[time, lev, lat, lon]的数据，在每个空间点对时间维度拟合之后，输出的数据为[lev, lat, lon, paras]。（PS: 这里感谢深雨露大佬的指点）
4 S6 J9 ~2 l% M, w& Q3 p, H( k& V& \0 u* b
以下是全部代码： 0 y5 Z; ]0 P5 F& P
from scipy import stats as st
( a2 D/ z. b2 a
; ?5 y# C9 O! v* _( T9 ?6 C import xarray as xr
2 W: F8 k2 f& N3 Q4 ? H import dask9 V* G1 O8 m" v& |- P: P8 M
import numpy as np: t7 R0 M, E- D$ w
from dask.diagnostics import ProgressBar5 N5 W' R3 A3 V
$ r- T3 v4 e, e1 G- S6 C
def norm_fit(data):
! C! N3 v7 y9 v7 [ d% e- [% h loc, scale = st.norm.fit(data)
: @, ?1 z# V. x. a8 W return np.array([loc, scale])& K% j, q, Z5 t$ W
/ ?4 c- e; E3 ?6 {$ j0 n
rs = np.random.RandomState(0). f$ T4 S. D. u0 f" ]7 R H
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])! b% H% q% A. _7 |. C+ @ l8 |
x = x.chunk({"lev":5, "lat":100,"lon":100})
) m9 _. D9 Q+ d/ @7 M R* J% j
3 A2 R+ z9 W. O7 h, F" ~9 g8 u #使用apply_ufunc计算，并用dask的并行计算
9 b% _5 L; @# x& C8 ` result = xr.apply_ufunc(norm_fit,
9 n( j1 `! Q P& I, A2 w9 q; @& J x,
) _+ u4 j2 d4 ]' I input_core_dims=[["time"]],
: P0 `& Y5 v! ~, r+ F' E* [9 A9 T6 a output_core_dims=[["paras"]],5 o( E) j8 W1 P. E# b
dask="parallelized",9 i0 g- O( P' D& |2 I2 l0 h+ X1 q0 a
output_dtypes=[np.float],- l8 a9 o& k& X/ i; p/ a' \3 P7 N
dask_gufunc_kwargs={output_sizes:{"paras":2}}
' c; A+ C, ~& f )+ `4 ]8 i4 ?+ k
5 g: b& t4 \' z7 V5 {* p5 U #compute进行真实的计算并显示进度: J b2 r) e2 t; p! U
with ProgressBar():
5 t& ~, U# H3 ^! q' e/ f0 S result = results.compute()9 N3 K }* W& M3 R7 u
6 D) j* v2 n1 ` J7 L #结果冲命名保存到nc文件
; ]' c$ a3 T! L( Y G/ ^ result = result.rename("norm_paras"). P4 P. i& m/ o0 R- y5 J
result = result.to_netcdf("norm_fit_paras.nc",compute=False)0 D$ F, q& Z( L% N- R
with ProgressBar():
; r0 C- b. F. [: ]& h) d2 t result.compute()
( e5 r; E1 o( M" P d' N 转自：气海同途 ( P) _1 E& q; _, f
关注【Ai尚研修科研技术平台】公众号，查看更多课程安排及免费数据资料
( G% D6 k, ^4 x# D" t, W
) Y4 E s. F1 i J 推荐：
% t) T- e2 O8 C 1、Python语言在地球科学领域中的应用实践技术应用 9 ?; C/ {8 x9 ]3 D4 q
2、【夏令营】针对课题组人员AI培养计划:“开启AI科研之路”
% z* f5 |* k+ o. q; V3 h 3、Python在气象与海洋中的实践技术应用精品课程
# w$ `: {! ]! I3 l) m, D" s+ ?' H 4、Python在WRF模型自动化运行及前后处理中的实践技术应用
/ M, I! o6 v. g5 P2 V/ E 5、全套Python机器学习核心技术与案例分析实践应用视频
! C% [1 x9 o8 I( P' f" h 6、全套区域高精度地学模拟-WRF气象建模、多案例应用与精美制图精品课程
$ O# m1 l# O, v N/ M: k. w 7、WRF DA资料同化系统理论、运行与与变分、混合同化新方法技术应用视频教程 # h/ ?( `' k" o5 U
$ j2 }) ~$ C# ?
$ v0 g8 {+ z" W j2 H# b$ g% u: F+ H |
9 {, x# V8 _/ o4 n0 S8 ], @/ E5 \" _/ |, j
| 
