|
, x8 I+ B5 \ l9 W: g' V
气象海洋领域，常常会涉及到大规模的数据处理，比如高时空分辨率的模式数据、雷达数据和卫星数据等。 9 d+ \) G2 b8 ]' }4 X) N" k
& @# F5 W9 e% |2 Q6 @: j) s
大部分情况，我们一般只会涉及小规模数据处理，对计算效率并不会太过追求。但是当数据量变大时，低效的数据处理所耗费的时间非常明显，因此高效的数据处理方式尤为重要。
+ B$ k3 F3 U5 C$ {/ @+ C: ? S
6 R% S3 \$ R2 b8 a$ F7 P 本篇以拟合一个高维数据的正太分布参数为例，介绍如何使用xarray+dask加速数据处理。 $ E* W$ w. Z- o2 }- }3 P! n; C
. e ]: y% W ?) H
数据维度X[time, lev, lat, lon]，需要对三维空间每一点，沿着时间维度做正太分布拟合(正太分布拟合只是作为例子，这里可以定义你需要的操作函数)。
' Q# v, j* j' f9 ? 其中几个关键点解释一下：
4 b$ x; _8 k! |0 j2 c8 U （1）首先定义拟合正太分布的函数 5 z. R! @+ o5 g, x7 l
def norm_fit(data):8 O* F8 s( e: {( r. ^ o
loc, scale = st.norm.fit(data), j' F* z' P [& h( ]2 J
return np.array([loc, scale])
# r5 @& s1 K$ Q3 S 这里需要注意的是，拟合的函数，输出参数，需要打包为一个数组。并且它的维度需要和后面aplly_ufun定义的输出维度一致。
5 a, W6 o5 \* Z8 G! D （2）xarray分块 : v# P& w6 o' m$ I. O
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"]); L8 o4 T0 {0 J' T
x = x.chunk({"lev":5, "lat":100,"lon":100})
, O* x& r! x. h* D* ]# Q& S xarray结合dask可以将一个大型数组分成一个个数据块（chunk），需要注意的是我们需要沿着时间维操作，拟合需要整个时间维度的数据，因此时间维time不能分块，只能对其他维度分块。
( M% u: d# U. T, C. ]# }" ^3 T （3）xarray.apply_ufunc函数 4 Z6 k! `7 W/ a8 |! h) |& [
result = xr.apply_ufunc(norm_fit,
V" L v( }6 Y x,
/ w% A; d* V" g input_core_dims=[["time"]],: K% _- M& j5 j1 r8 @0 D
output_core_dims=[["paras"]],9 h& N/ Z' U* N3 W
dask="parallelized",3 h1 @: M; k2 V+ h" l" D
output_dtypes=[np.float],6 x4 n: q: M! U! ?# ]
dask_gufunc_kwargs={output_sizes:{"paras":2}}
7 b" M# e& |( M: ?* w )
+ p4 e7 k' b8 }9 z5 P1 b9 e# I( b apply_ufunc函数具体可以参考官网教程，这里只说使用时遇到的困难，即如何定义输出维度：输出维度是用output_core_dim定义的，将输出的拟合参数（期望和标准差）定义为paras维度，维度的大小为2，通过output_sizes参数设置。这样我们输入[time, lev, lat, lon]的数据，在每个空间点对时间维度拟合之后，输出的数据为[lev, lat, lon, paras]。（PS: 这里感谢深雨露大佬的指点）
% _/ A& |: Q5 U8 k6 k$ F! k
% \, K X4 Z$ P# |3 [ 以下是全部代码：
j3 P8 W" |# a G( i5 h5 w/ b from scipy import stats as st |# e' \, e! F! p1 u2 I& s
, M' z" e8 Q* Q5 l
import xarray as xr$ _1 T2 D- u2 S- w, A$ d6 I9 m9 B0 {
import dask$ {! N! b( Y1 H: ?7 }
import numpy as np/ d# F: X8 N0 n4 h9 q, C! }0 n8 a) q( P
from dask.diagnostics import ProgressBar
5 T* \, {4 Y. g
4 v2 W y& q8 m% ^ def norm_fit(data):3 k0 Q2 w5 b* `7 U* b
loc, scale = st.norm.fit(data)) D# i) ~5 r5 c% G7 J9 S3 ]; e
return np.array([loc, scale])* Z* f$ J2 U: t0 T
/ a. b) P( u3 T4 K( R2 z3 D
rs = np.random.RandomState(0)" c/ I# s9 D( d2 _$ @8 J+ z
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])
2 n- y R0 @$ K x = x.chunk({"lev":5, "lat":100,"lon":100})
1 E; d( s, t9 u9 G" o4 C0 G" \
#使用apply_ufunc计算，并用dask的并行计算
, x9 l8 }8 l1 W6 F0 W, @ result = xr.apply_ufunc(norm_fit,7 \% M& i' B# Y, C/ t2 c
x,) x5 D% k4 h- |' i* v
input_core_dims=[["time"]],
4 f" _6 ^6 y. I% e1 M output_core_dims=[["paras"]],6 r& v+ U- [1 v4 H0 H% W
dask="parallelized",/ h2 U% ~% a$ w
output_dtypes=[np.float],- O3 K, ^# v- \1 u0 ]8 t; e* Q* G9 H
dask_gufunc_kwargs={output_sizes:{"paras":2}}+ R, ~3 `! v3 f# A
): o& p0 x. c* ?' n& Z: m- P
7 E% s! F. q; d2 K3 M9 F! A #compute进行真实的计算并显示进度! @; r" W$ r& y* {' ^- Y# [5 W- z7 \
with ProgressBar():
2 o3 u8 Q& W( I result = results.compute()2 q( \5 I: @9 `8 C. M5 A8 q
3 M9 Z& @, P" | #结果冲命名保存到nc文件: H0 u! f" y. {2 m$ k
result = result.rename("norm_paras")0 Q U& T5 E% i9 k' K
result = result.to_netcdf("norm_fit_paras.nc",compute=False)
4 w- c/ _ g( \4 T" a/ `, H% W8 J with ProgressBar():. ^3 W9 @. ` r& ]' y7 M
result.compute() & I, O2 u+ l7 v1 ]# Z
转自：气海同途 " [# H0 v1 L- Z$ s2 R- |
关注【Ai尚研修科研技术平台】公众号，查看更多课程安排及免费数据资料 ( ` l# X) g& E8 i5 S, k: C! p
& L" I! X; c& }+ N. L7 g% V 推荐：
+ l) p) O0 C- r% b0 { 1、Python语言在地球科学领域中的应用实践技术应用
9 L: r1 `5 O7 H' ^* R: U 2、【夏令营】针对课题组人员AI培养计划:“开启AI科研之路” ; F- c' k( ]5 @
3、Python在气象与海洋中的实践技术应用精品课程 : a% k: G x9 {6 S2 w* f
4、Python在WRF模型自动化运行及前后处理中的实践技术应用 $ M5 s& o$ G/ P! b
5、全套Python机器学习核心技术与案例分析实践应用视频 * Y- U* e+ l4 y& k5 p& G- o" ^, ^
6、全套区域高精度地学模拟-WRF气象建模、多案例应用与精美制图精品课程
, S: k) B! G5 M$ m' u* Q2 {5 Q 7、WRF DA资料同化系统理论、运行与与变分、混合同化新方法技术应用视频教程 % K8 K/ ?/ Z2 Y4 `( d) Q, |* h
2 @6 z1 v" S0 ], I
5 D* P8 h S4 x# F% E7 x, I$ ?$ N& h9 Q0 j F
3 [& n/ h! @" [1 P: F | 
