|
y+ r: L1 N+ O4 H t! K. M 气象海洋领域，常常会涉及到大规模的数据处理，比如高时空分辨率的模式数据、雷达数据和卫星数据等。
8 I8 C! |0 D/ a/ [
' w; Z& Y1 I# F% W 大部分情况，我们一般只会涉及小规模数据处理，对计算效率并不会太过追求。但是当数据量变大时，低效的数据处理所耗费的时间非常明显，因此高效的数据处理方式尤为重要。 8 M0 U" `, V# w U1 g8 B4 t
. N. {. N2 c' G' U5 o7 I
本篇以拟合一个高维数据的正太分布参数为例，介绍如何使用xarray+dask加速数据处理。
% \5 J7 S; F' Z+ E5 l0 F9 K) t9 A/ l1 ^$ C+ \& B
数据维度X[time, lev, lat, lon]，需要对三维空间每一点，沿着时间维度做正太分布拟合(正太分布拟合只是作为例子，这里可以定义你需要的操作函数)。 ' T. l1 n, p) G3 h, N
其中几个关键点解释一下：
7 H: `) z7 t/ X% L- R （1）首先定义拟合正太分布的函数 [4 j- l/ f ~ }
def norm_fit(data):9 O2 v' f9 h$ R6 T. g4 m6 C) J
loc, scale = st.norm.fit(data)$ }; b9 g/ x H7 ]
return np.array([loc, scale])
3 E+ i- _3 [, W, K* `% Q 这里需要注意的是，拟合的函数，输出参数，需要打包为一个数组。并且它的维度需要和后面aplly_ufun定义的输出维度一致。 * q! o3 w6 |2 }4 O5 w
（2）xarray分块
G. i# @; d3 G3 `' `. ] x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])
4 V' E- k# s- L& O x = x.chunk({"lev":5, "lat":100,"lon":100})
1 `+ Z* U6 Z9 o T" O$ T xarray结合dask可以将一个大型数组分成一个个数据块（chunk），需要注意的是我们需要沿着时间维操作，拟合需要整个时间维度的数据，因此时间维time不能分块，只能对其他维度分块。
4 b( v2 L0 S+ w （3）xarray.apply_ufunc函数
; g# M* T( Q7 ~ h result = xr.apply_ufunc(norm_fit,: n# v5 e8 |. ]1 l& a* h
x,
- p1 Y# ~; y' {8 x4 o S- g9 \, M input_core_dims=[["time"]],. G/ M3 ~ D) G' q, F& T" `" ?
output_core_dims=[["paras"]],! s( m f+ ~8 o
dask="parallelized",2 A7 g* }1 N$ e9 F5 X; m# u
output_dtypes=[np.float],0 `( x7 r1 P- X$ h
dask_gufunc_kwargs={output_sizes:{"paras":2}}
: f+ w0 r3 K U2 Q )
! T* w! I, p" E. U: Z, T) W apply_ufunc函数具体可以参考官网教程，这里只说使用时遇到的困难，即如何定义输出维度：输出维度是用output_core_dim定义的，将输出的拟合参数（期望和标准差）定义为paras维度，维度的大小为2，通过output_sizes参数设置。这样我们输入[time, lev, lat, lon]的数据，在每个空间点对时间维度拟合之后，输出的数据为[lev, lat, lon, paras]。（PS: 这里感谢深雨露大佬的指点）
1 ?9 n% w9 [' q" e8 i: X4 W2 ]# [& l4 ^* G# M8 X& \
以下是全部代码：
+ ~* Y% `5 `# u$ `" c0 T5 C" H from scipy import stats as st6 V( `2 Y; |1 s
4 e; U8 _2 s1 S; w- ^ import xarray as xr, G* R1 I8 d; _2 i2 w2 V( d
import dask
8 b; J1 {; i+ t' t' M import numpy as np9 V4 {: a" W+ r W, R
from dask.diagnostics import ProgressBar
; ?; ~/ f6 I0 I3 }
8 v- f+ `9 @8 j- N" j- K def norm_fit(data):
! B) O2 {5 D! y5 W, B! L$ F8 w loc, scale = st.norm.fit(data)
/ g) |+ b! h6 {% m; g- ]* Q return np.array([loc, scale]); F( f6 F' h% j1 M' E$ w2 [. O- d
! |" k1 s4 A: M# ~( }
rs = np.random.RandomState(0) n5 E/ H$ Z2 N# d1 o) r/ ]
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])) F; B: v u5 {8 {) S
x = x.chunk({"lev":5, "lat":100,"lon":100})
5 h' c h: e7 D' w! U6 M
5 s) m0 D8 w7 e8 z/ t* i #使用apply_ufunc计算，并用dask的并行计算
$ }( s6 p# Q' d" d4 ]1 |" x' n; L result = xr.apply_ufunc(norm_fit,
- M9 Z0 \/ W+ K, t8 J0 S( J4 C" Q x,
6 G0 q0 D! \; e+ [ input_core_dims=[["time"]],8 O: K, D2 ~" R0 J- K8 I; S
output_core_dims=[["paras"]],$ a4 @& w- [7 x+ a* }' ~, b
dask="parallelized",
" V( C$ a: y( l: C2 E output_dtypes=[np.float],% E8 C7 F5 ^& B9 t! _; W
dask_gufunc_kwargs={output_sizes:{"paras":2}}- J' l9 Z- x1 a- ^
)
0 T1 Y. h8 p$ \2 B) J) T
: G9 M/ U% A: z$ j #compute进行真实的计算并显示进度/ [+ T, k, l4 i3 t9 a
with ProgressBar():- H$ b2 @: ^3 a R) d
result = results.compute()
; Q- i2 D0 r6 ?! Q7 @' a0 W1 R" b! J2 ^. Y/ P
#结果冲命名保存到nc文件
0 z" E/ X) X& `2 r5 v7 j+ M7 J( [ result = result.rename("norm_paras")8 E& e2 M. }/ c% t
result = result.to_netcdf("norm_fit_paras.nc",compute=False)
: `; g& J3 W' C7 u! L with ProgressBar():& l# @/ W3 Y4 |) N4 H' o$ S' G* d
result.compute()
7 l& q, K& c1 d 转自：气海同途
' |. p1 A9 ~+ d 关注【Ai尚研修科研技术平台】公众号，查看更多课程安排及免费数据资料
! u* ]/ G8 M7 L8 E; u \, [
7 B. j. r& n* X% h! R% K 推荐： 8 U8 h$ T6 b+ i5 U* L9 Q% y8 d# y
1、Python语言在地球科学领域中的应用实践技术应用
9 K$ R$ _9 N3 n: t7 d) C 2、【夏令营】针对课题组人员AI培养计划:“开启AI科研之路” ' M0 T/ v" a. M% r7 h. p+ ]
3、Python在气象与海洋中的实践技术应用精品课程 5 V; ]5 L" r. j
4、Python在WRF模型自动化运行及前后处理中的实践技术应用
2 m8 _1 U+ P" [- a4 @ 5、全套Python机器学习核心技术与案例分析实践应用视频 8 ~, f K1 p- F) b$ O
6、全套区域高精度地学模拟-WRF气象建模、多案例应用与精美制图精品课程 % d' L. }5 ?( X) D6 o5 V
7、WRF DA资料同化系统理论、运行与与变分、混合同化新方法技术应用视频教程 ) \. ?9 G# t0 M" S. y+ g! i
) n, ~0 g+ f' y) R4 M0 M
$ F5 L; W, t; B( a: v. ~9 r
5 p" L, H' J4 `
% J7 Y: f/ T. l5 v' ~" z7 R% P1 f" x( g | 
