|
9 x+ F7 }' u4 Q8 N
气象海洋领域，常常会涉及到大规模的数据处理，比如高时空分辨率的模式数据、雷达数据和卫星数据等。
4 H# E# V: U& k# ^+ E6 q9 T4 U / Q9 d; f1 F/ a( W5 h: \
大部分情况，我们一般只会涉及小规模数据处理，对计算效率并不会太过追求。但是当数据量变大时，低效的数据处理所耗费的时间非常明显，因此高效的数据处理方式尤为重要。
% u1 c# `: p `: ^' ]5 I! [- \+ W! ?) |* ^3 ]1 @! J
本篇以拟合一个高维数据的正太分布参数为例，介绍如何使用xarray+dask加速数据处理。 . ^) M/ t# J$ f2 b) _- H8 i+ i# Y' r! l
; t& I+ `& S' k ?
数据维度X[time, lev, lat, lon]，需要对三维空间每一点，沿着时间维度做正太分布拟合(正太分布拟合只是作为例子，这里可以定义你需要的操作函数)。 # V+ k+ h5 d# { E5 ?# H( |9 Q
其中几个关键点解释一下：
7 Y% c! U) t4 U' E3 ` （1）首先定义拟合正太分布的函数 2 p* Y: d2 W' B
def norm_fit(data):
1 h" }( \0 o5 g loc, scale = st.norm.fit(data)
! z% A/ q/ j0 g! J1 ~ \ return np.array([loc, scale]) / g0 C {* z# d I9 q
这里需要注意的是，拟合的函数，输出参数，需要打包为一个数组。并且它的维度需要和后面aplly_ufun定义的输出维度一致。
, b6 d& q, d+ o2 n1 {; I （2）xarray分块 ; R+ t$ |8 E2 P# u, f
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])- i6 q/ m5 q7 [% s/ U" k9 ~0 `' N
x = x.chunk({"lev":5, "lat":100,"lon":100})
9 N. l3 Q+ D, \& K& h$ U xarray结合dask可以将一个大型数组分成一个个数据块（chunk），需要注意的是我们需要沿着时间维操作，拟合需要整个时间维度的数据，因此时间维time不能分块，只能对其他维度分块。 : s$ E" e1 T6 m( @4 r% y
（3）xarray.apply_ufunc函数
+ R" u) K- f/ i2 X0 ]; v result = xr.apply_ufunc(norm_fit,
* ^$ ^) r) P. o. o5 Y* |8 w x,
% y, \) }. N# W! ?( K: T4 M input_core_dims=[["time"]],
. Q0 ?7 f9 @0 ]5 w5 y output_core_dims=[["paras"]],2 k) U. I4 M; j3 v. ~
dask="parallelized",
9 Q8 k- |( f* F! P7 W( b# | output_dtypes=[np.float],
& I3 }! S: c, k* T dask_gufunc_kwargs={output_sizes:{"paras":2}}
" n5 a6 E0 `7 V& X' O2 A ) h2 t: b& T0 e% x6 z
apply_ufunc函数具体可以参考官网教程，这里只说使用时遇到的困难，即如何定义输出维度：输出维度是用output_core_dim定义的，将输出的拟合参数（期望和标准差）定义为paras维度，维度的大小为2，通过output_sizes参数设置。这样我们输入[time, lev, lat, lon]的数据，在每个空间点对时间维度拟合之后，输出的数据为[lev, lat, lon, paras]。（PS: 这里感谢深雨露大佬的指点） 3 d! ?. a1 q- \+ h5 H
. t. H" C3 N6 D. l
以下是全部代码：
8 G0 }" f" m A1 o. c from scipy import stats as st
2 o8 S/ Q: S. L5 o: f; l! f! t
: U0 t; _3 v$ b; Z2 b, x# b import xarray as xr" T8 P4 W/ |/ Y
import dask
: v1 x1 t2 F# ]- E$ j! M import numpy as np5 h" _/ ]: N! ]$ i- E: D2 s9 I) r7 k
from dask.diagnostics import ProgressBar
2 W6 m( N' l \% I ?% c- e& A1 r: z3 H
def norm_fit(data):" ~% E5 e/ F' Z6 \
loc, scale = st.norm.fit(data)& A# f1 ~' A5 ]) y" G, ~+ O
return np.array([loc, scale])' |; P( d3 K) S9 d
! Y3 s+ M9 ]* Y; u rs = np.random.RandomState(0)
1 g ?; _; f1 w, T( r$ u3 P x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])
' E) s4 k. H5 o6 H- r/ a3 p x = x.chunk({"lev":5, "lat":100,"lon":100})9 [4 I- i" c1 ?
6 u6 X% j3 Q) D% H
#使用apply_ufunc计算，并用dask的并行计算: i2 h6 G% K5 V% _8 k% L+ F
result = xr.apply_ufunc(norm_fit,
# P2 x2 t; B0 _! i- g x,
+ P9 ^3 f( W" ~ input_core_dims=[["time"]],
9 {: f6 S8 c& n6 G+ G: Y output_core_dims=[["paras"]],
7 Z2 U. `, m* E2 a7 Q( F) {! I) H* E dask="parallelized",- b2 @( h6 l v/ U+ I
output_dtypes=[np.float],
+ H% o! l2 W% s/ I3 ~3 d+ F dask_gufunc_kwargs={output_sizes:{"paras":2}}
- z: E+ [! k. e7 {/ a )
& B* o4 C5 ~! d2 ^) T! d: u
2 p+ A3 a1 M4 C6 O" z #compute进行真实的计算并显示进度
6 A" S0 N' ^1 `$ ?/ }- ] with ProgressBar():9 W. A. X0 a; ^6 B$ U
result = results.compute()$ j$ G. z4 }4 o" K4 \1 [
0 A+ L& Y; B' ^/ h* ~+ c% r6 Z #结果冲命名保存到nc文件4 K/ {- x( a. \! n; V/ o5 k
result = result.rename("norm_paras")+ @, z! C$ k! G
result = result.to_netcdf("norm_fit_paras.nc",compute=False)
! P' y% n' B9 Q6 X& C$ f' [ with ProgressBar():$ c: |. V. R1 p! q0 X
result.compute() + |: X, u8 J! s z" k/ j7 K
转自：气海同途 ; c$ \% }. C; z( r
关注【Ai尚研修科研技术平台】公众号，查看更多课程安排及免费数据资料 ; Z: l& n! ?3 C4 f7 ?% ]; j- q6 W
( z5 ]- _5 c; Q2 V' }
推荐： + J3 z8 \* y T
1、Python语言在地球科学领域中的应用实践技术应用 . h. F- `. s5 z7 X
2、【夏令营】针对课题组人员AI培养计划:“开启AI科研之路”
) d9 C9 Y p- m% P7 Q' r+ l4 v 3、Python在气象与海洋中的实践技术应用精品课程
- A! j) K7 D, H. @1 j3 h4 m 4、Python在WRF模型自动化运行及前后处理中的实践技术应用 % y( d% |1 G1 Q7 _1 R. I
5、全套Python机器学习核心技术与案例分析实践应用视频
0 P# w" V( |0 C4 J& l7 ~ 6、全套区域高精度地学模拟-WRF气象建模、多案例应用与精美制图精品课程 1 d# T! Z a; @8 Z5 V
7、WRF DA资料同化系统理论、运行与与变分、混合同化新方法技术应用视频教程
0 v* L ^( |: S8 H& V6 ^8 s; `
- W7 {: G7 T5 a8 b: Q
; s" _2 M# ~9 ^: t' R' [/ V
4 A* {& s, n1 m3 e4 z+ B" |* K1 U5 p# N6 a ?, P1 [7 j5 _5 {8 r
| 
