|
0 J" A, ^2 e4 ]5 G, q8 d! d
气象海洋领域，常常会涉及到大规模的数据处理，比如高时空分辨率的模式数据、雷达数据和卫星数据等。 + Z+ \$ u* Z c0 \, [) e. E- T5 i
/ B% B+ f$ K, Y- {$ p: m
大部分情况，我们一般只会涉及小规模数据处理，对计算效率并不会太过追求。但是当数据量变大时，低效的数据处理所耗费的时间非常明显，因此高效的数据处理方式尤为重要。
/ y. N6 i3 ~# Q2 g, k6 z% @2 Z( |" t
$ l5 l1 m* D- M2 ?3 K4 x* E, } 本篇以拟合一个高维数据的正太分布参数为例，介绍如何使用xarray+dask加速数据处理。
, l3 B+ S6 E% w8 O; }4 Q) P4 A0 u- w: o- z( ^# ]0 y# g8 B
数据维度X[time, lev, lat, lon]，需要对三维空间每一点，沿着时间维度做正太分布拟合(正太分布拟合只是作为例子，这里可以定义你需要的操作函数)。
4 I3 y; T- _- r" q/ l3 Z+ y 其中几个关键点解释一下：
8 K& Z/ ?# w' b* P （1）首先定义拟合正太分布的函数
6 X+ L$ q& i/ K def norm_fit(data):- P/ b5 m8 Y& f5 ]! m: D
loc, scale = st.norm.fit(data)
3 S3 _, ~+ p0 G return np.array([loc, scale])
5 H7 X& ]0 f; i5 r% }1 t' S d 这里需要注意的是，拟合的函数，输出参数，需要打包为一个数组。并且它的维度需要和后面aplly_ufun定义的输出维度一致。
/ g/ ]7 L3 E1 l8 R C( ]9 B （2）xarray分块 , Y7 I! o/ c9 G q4 F
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"]), I0 {# W9 s5 ?8 F' d; u9 X
x = x.chunk({"lev":5, "lat":100,"lon":100}) ' M" w7 |# B! e, H) a- Y
xarray结合dask可以将一个大型数组分成一个个数据块（chunk），需要注意的是我们需要沿着时间维操作，拟合需要整个时间维度的数据，因此时间维time不能分块，只能对其他维度分块。 1 O( C- r& D& K* V4 ~$ h' \
（3）xarray.apply_ufunc函数 ) H1 p# k+ ?$ k' \
result = xr.apply_ufunc(norm_fit,; [) n/ q' g- u) _( ? Q. h7 K; r
x,
3 P. K, f" k9 k8 z( | input_core_dims=[["time"]],
/ q( t- L3 Q7 z) X* ` output_core_dims=[["paras"]],
X, L+ n! ~8 [$ G/ y dask="parallelized",: C0 `+ W; M2 v9 t; P3 _
output_dtypes=[np.float],
% ~9 U" x S# U9 j dask_gufunc_kwargs={output_sizes:{"paras":2}}
8 h1 M# i( t7 ]3 x ) 8 h e* Z$ v p$ q j9 Y7 @4 U6 g
apply_ufunc函数具体可以参考官网教程，这里只说使用时遇到的困难，即如何定义输出维度：输出维度是用output_core_dim定义的，将输出的拟合参数（期望和标准差）定义为paras维度，维度的大小为2，通过output_sizes参数设置。这样我们输入[time, lev, lat, lon]的数据，在每个空间点对时间维度拟合之后，输出的数据为[lev, lat, lon, paras]。（PS: 这里感谢深雨露大佬的指点） 2 Q- l* y6 }% [0 z
; a5 ]- t$ E! D: G @. A# G
以下是全部代码：
/ ^: B% J4 b. D& h# s) ] from scipy import stats as st
4 ~7 \- @) c- R! @4 V: y( W
- M/ U3 _3 J O; A" G/ H import xarray as xr
- U: D4 c. ^3 a# I8 ]- @/ | import dask
" e! q, `' v5 q% X* B import numpy as np6 ]. i$ m1 Z4 Y; ?4 K- @& p& ]' x
from dask.diagnostics import ProgressBar S# K6 `. H! y# S
; `' V- f( g/ m) w8 x$ `: t def norm_fit(data):
7 i: k' }: t: s/ M# k; ~8 H& F loc, scale = st.norm.fit(data); Q/ b; ]/ W9 ?! @2 e
return np.array([loc, scale])
8 `, A. E+ K0 P# C9 Z
' D# c6 D0 n( R9 K* c" p6 x rs = np.random.RandomState(0) T' C2 J/ O7 \4 L1 k2 y5 A
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])
4 z! F$ O9 l$ V x = x.chunk({"lev":5, "lat":100,"lon":100})
- O, ?3 Q4 A3 p* i. W9 `5 ]3 K& F( B
#使用apply_ufunc计算，并用dask的并行计算5 A- }( v* p- S8 w
result = xr.apply_ufunc(norm_fit,; B t7 c, D6 e( w7 M7 y) [: f* f6 \
x,8 P/ l3 g, S8 U* r
input_core_dims=[["time"]],
2 [: ~. I' R4 A! V" Z9 B1 M output_core_dims=[["paras"]],
$ i( \4 w$ k* [- _ dask="parallelized",, }3 m6 G+ h* j5 r
output_dtypes=[np.float],* x6 F! [ t5 O1 S' }& ]1 Z. Z( G. A
dask_gufunc_kwargs={output_sizes:{"paras":2}}! G( u! C8 h% m! i3 G: O$ e
)0 k1 U( ]1 H: h* \
' `. |4 |- Q/ d) z1 _8 |8 `
#compute进行真实的计算并显示进度
/ @) q4 m2 v/ ]' E with ProgressBar():
4 ^4 z- H& R5 R result = results.compute()
0 y8 e* L+ I! t* R: k% X/ a+ @4 D, J* f" J0 X0 ?- G* F5 E
#结果冲命名保存到nc文件 a- o x' p) H, S/ G# v1 k7 R
result = result.rename("norm_paras")
0 K& |) S# l2 ]/ f result = result.to_netcdf("norm_fit_paras.nc",compute=False)+ X/ w2 l& j; r( I
with ProgressBar():
, k! |+ r# l% T( c- n! b! Q0 @ result.compute()
/ i: D$ w. d& W! | N4 G9 @ 转自：气海同途 . S3 H$ b" Z; s, u& N2 H
关注【Ai尚研修科研技术平台】公众号，查看更多课程安排及免费数据资料
3 W2 k, s( W* e7 r9 s$ j1 n& i5 p ; L. s2 O( w) ^. m' Z
推荐： 1 S& W6 i- h0 m3 H8 b
1、Python语言在地球科学领域中的应用实践技术应用 4 d) t+ g* q- [- Y6 n( k% ?3 b4 V
2、【夏令营】针对课题组人员AI培养计划:“开启AI科研之路”
* j2 b3 O6 ]9 f2 @ 3、Python在气象与海洋中的实践技术应用精品课程
' B. W; u# {9 k9 g( {8 F9 T* L6 R 4、Python在WRF模型自动化运行及前后处理中的实践技术应用
7 Z! \3 b& a! `2 V$ ^ 5、全套Python机器学习核心技术与案例分析实践应用视频 % [- u3 N3 g8 X8 w
6、全套区域高精度地学模拟-WRF气象建模、多案例应用与精美制图精品课程
) C+ j/ j: S, w- G X1 V: V 7、WRF DA资料同化系统理论、运行与与变分、混合同化新方法技术应用视频教程 0 z' C( }2 F% n
9 A% P# Q- [# ^1 Y3 L
. _$ S3 h4 d0 c# {4 A2 c$ l3 ?4 a- N6 n
3 q' N. l2 `& ?3 n# A | 
