|
5 H4 H* {# o+ M4 u' f- K( n
气象海洋领域，常常会涉及到大规模的数据处理，比如高时空分辨率的模式数据、雷达数据和卫星数据等。 7 d0 N! c/ W' W! F# H2 `* N) R) I4 ^
; `2 u. D, }( n/ s- P 大部分情况，我们一般只会涉及小规模数据处理，对计算效率并不会太过追求。但是当数据量变大时，低效的数据处理所耗费的时间非常明显，因此高效的数据处理方式尤为重要。 2 T) ~9 ^; Y( d! R1 A
7 [. o% E4 ^/ C4 s: t/ @7 r' x 本篇以拟合一个高维数据的正太分布参数为例，介绍如何使用xarray+dask加速数据处理。
' y) a( p: @- W1 `& f; g8 ]+ U3 L8 ~1 W5 j
数据维度X[time, lev, lat, lon]，需要对三维空间每一点，沿着时间维度做正太分布拟合(正太分布拟合只是作为例子，这里可以定义你需要的操作函数)。
3 a, q( N, W8 ?" Y 其中几个关键点解释一下：
" U; o! m+ |) Z& k （1）首先定义拟合正太分布的函数
8 M0 u) y) c$ t& ~. Y& z. Q0 b def norm_fit(data):6 _- a& U" Z2 ?) }, }. ]5 `" v" t6 ]
loc, scale = st.norm.fit(data)
7 u9 U0 z+ \2 f* k7 n4 w* h2 D return np.array([loc, scale])
! |0 |& F. ^* t- ? 这里需要注意的是，拟合的函数，输出参数，需要打包为一个数组。并且它的维度需要和后面aplly_ufun定义的输出维度一致。 0 y% F* W9 E: D; [
（2）xarray分块 # A3 Y, ~: o. `
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])
/ s* b. w) F: \0 U+ h; A x = x.chunk({"lev":5, "lat":100,"lon":100}) 6 s }0 E" a* d1 F
xarray结合dask可以将一个大型数组分成一个个数据块（chunk），需要注意的是我们需要沿着时间维操作，拟合需要整个时间维度的数据，因此时间维time不能分块，只能对其他维度分块。
& M, ~" ?" A# s5 W% P( k5 y8 t! J （3）xarray.apply_ufunc函数
1 }( l2 L- g3 a; W0 V# M result = xr.apply_ufunc(norm_fit,
+ X7 U9 _ P3 B0 _: o5 [ x,
9 ^8 l H$ r: p input_core_dims=[["time"]],5 ?' T( m$ l! X1 X$ X
output_core_dims=[["paras"]],
! H8 }" r% A' O) Q4 U9 y dask="parallelized",
9 Z0 w, p7 P, `, h, j output_dtypes=[np.float],
/ W n/ o5 L5 ~! z dask_gufunc_kwargs={output_sizes:{"paras":2}}
, Q4 r3 p" C% {0 ` )
+ D9 \# E: C# {) x apply_ufunc函数具体可以参考官网教程，这里只说使用时遇到的困难，即如何定义输出维度：输出维度是用output_core_dim定义的，将输出的拟合参数（期望和标准差）定义为paras维度，维度的大小为2，通过output_sizes参数设置。这样我们输入[time, lev, lat, lon]的数据，在每个空间点对时间维度拟合之后，输出的数据为[lev, lat, lon, paras]。（PS: 这里感谢深雨露大佬的指点）
4 P& s+ b$ K1 e0 o, r0 @% i! t0 o
以下是全部代码：
! C4 I4 X; e) t+ s8 L5 ]& V! E; U s- | from scipy import stats as st
3 e: h4 {# L, x* i2 |+ k$ W
7 g, ]+ H- p$ l- F+ X( O3 C import xarray as xr7 c. G' W3 z; t, T; r
import dask4 B) ^2 W$ }8 d9 O+ A/ a
import numpy as np$ W! N% D+ `0 K: A
from dask.diagnostics import ProgressBar
% B* H' H1 ?/ S( t5 T* M4 J& Z, g/ Z# |
def norm_fit(data):
0 t; v; R" {; Y' [, I loc, scale = st.norm.fit(data)
8 Q6 L0 U- z: \# ]* u( \1 k. P: G: e return np.array([loc, scale])/ f# s7 E+ s' ^9 M1 O
9 `' r/ r0 f4 ?! y* _ rs = np.random.RandomState(0)% F$ h" \$ ~0 y5 ?9 ?: U( W+ Q
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])
% |2 V( s. E1 G7 r x = x.chunk({"lev":5, "lat":100,"lon":100})* Q' F$ t1 {* U* Q9 O
4 }4 v% c; S4 U% S' g, J% N, T #使用apply_ufunc计算，并用dask的并行计算
, ^( \" U) Z6 g result = xr.apply_ufunc(norm_fit,
" R9 Y% `9 s- d) k$ i# `7 X: b( z, u+ P x,2 S, p/ |0 K! o
input_core_dims=[["time"]]," M; }2 T. @. T6 e* I$ c$ G* |
output_core_dims=[["paras"]],
^- P6 Y8 N8 U8 a4 L dask="parallelized",- i% [" M0 R6 I# r, `1 `, g$ q
output_dtypes=[np.float], e+ |3 D9 V, X" d! o
dask_gufunc_kwargs={output_sizes:{"paras":2}}8 l6 O0 O. V N/ T3 h$ |0 T
); s; g$ n2 o, f0 Z; G
3 A; D" x# B* o1 s$ T #compute进行真实的计算并显示进度& z2 r. s W2 Y$ ^( Y
with ProgressBar():& \* R6 ]7 t/ `; A
result = results.compute()
: t2 i; P' i* w) B$ M9 h
* R; }5 n2 q. u: ?5 Q/ i #结果冲命名保存到nc文件
# u$ ]+ W; N9 L D4 ~' P result = result.rename("norm_paras")' R8 t; G! Z( O" H
result = result.to_netcdf("norm_fit_paras.nc",compute=False)' q8 H" v) z5 V& y# Y5 i
with ProgressBar():8 S4 ^3 |$ n. k, `. J- r0 D
result.compute() * a- Q( ^1 l& j/ e3 R {
转自：气海同途
. W) p) h! [8 |0 G- P' ~! ] 关注【Ai尚研修科研技术平台】公众号，查看更多课程安排及免费数据资料
6 K9 E' `3 s+ k6 a% ~ . n0 ]. ]! v8 H
推荐：
/ I1 \2 c* ^2 p# A6 t9 S 1、Python语言在地球科学领域中的应用实践技术应用 % r3 F8 a& J# J3 K' x( M2 ^
2、【夏令营】针对课题组人员AI培养计划:“开启AI科研之路”
- m. }! p- }8 m" j2 S5 r 3、Python在气象与海洋中的实践技术应用精品课程
+ A/ [& c! y& C: K" O% f; \" B0 B 4、Python在WRF模型自动化运行及前后处理中的实践技术应用 , X3 |3 ]8 D) Y" U$ Q$ o
5、全套Python机器学习核心技术与案例分析实践应用视频
/ |5 n, ?2 R+ r' e) t( [ 6、全套区域高精度地学模拟-WRF气象建模、多案例应用与精美制图精品课程
1 {' e( l" l9 q3 h. m" P, Q 7、WRF DA资料同化系统理论、运行与与变分、混合同化新方法技术应用视频教程
" ~% [5 ]- Y2 d9 D) c: d/ B, M" A- r& i0 @5 g3 U) z; Q' F; t
- [- n9 G. i6 W: }- y
4 |) d. w. r8 R4 u/ I& |' d3 K: B1 d, t$ O |) Y( c# ~# X3 B
| 
