|
3 b1 M3 J/ f, D3 W1 d( e- O- v 气象海洋领域，常常会涉及到大规模的数据处理，比如高时空分辨率的模式数据、雷达数据和卫星数据等。 3 a- n5 I d( \
7 A$ y V% x6 J# ? 大部分情况，我们一般只会涉及小规模数据处理，对计算效率并不会太过追求。但是当数据量变大时，低效的数据处理所耗费的时间非常明显，因此高效的数据处理方式尤为重要。
5 y& e. U- G0 \7 u2 S+ ~: d. L1 n: W6 t, A% K/ e# r3 w1 |' a
本篇以拟合一个高维数据的正太分布参数为例，介绍如何使用xarray+dask加速数据处理。 9 ^. L# O0 p2 w# @/ o8 T, S( ^
* O( |% ]6 u z) k* t
数据维度X[time, lev, lat, lon]，需要对三维空间每一点，沿着时间维度做正太分布拟合(正太分布拟合只是作为例子，这里可以定义你需要的操作函数)。 1 J! J) T6 `( d, W6 _& {
其中几个关键点解释一下：
" {* `- ~! O: S$ h/ j% d- V （1）首先定义拟合正太分布的函数
& Y C! r; p. T% q5 N* f: D5 W$ | def norm_fit(data):, t, X! k* d7 P/ q
loc, scale = st.norm.fit(data)9 `1 y9 C; E+ I
return np.array([loc, scale])
/ I- w7 l/ e H, Q9 e 这里需要注意的是，拟合的函数，输出参数，需要打包为一个数组。并且它的维度需要和后面aplly_ufun定义的输出维度一致。
2 [5 c% p8 B' S/ o （2）xarray分块
: A. f2 G! q5 Y& v% l+ n) q, \ x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])2 P- F/ i9 K( [# R; k2 Z
x = x.chunk({"lev":5, "lat":100,"lon":100})
8 Z; g% m$ {( Y- h. E. i xarray结合dask可以将一个大型数组分成一个个数据块（chunk），需要注意的是我们需要沿着时间维操作，拟合需要整个时间维度的数据，因此时间维time不能分块，只能对其他维度分块。 + L' k+ Q; D+ J& C2 \* @
（3）xarray.apply_ufunc函数
, Q/ e2 U9 y4 |6 [ result = xr.apply_ufunc(norm_fit,
" L; v; ~) |4 S6 w x,2 V0 O) f( p# n% u: f" _2 _# K
input_core_dims=[["time"]],
' s& W) R p: ~& q; ~; O, U output_core_dims=[["paras"]],8 Y9 W0 W! O1 B0 E& v% r
dask="parallelized",, W9 e5 l7 ~! o0 L
output_dtypes=[np.float],, S/ @7 _- V8 `
dask_gufunc_kwargs={output_sizes:{"paras":2}}4 T. L) R5 i7 Y0 b& J, c2 C9 K
) # w+ }3 \9 [, f1 Y1 E! o/ o! [
apply_ufunc函数具体可以参考官网教程，这里只说使用时遇到的困难，即如何定义输出维度：输出维度是用output_core_dim定义的，将输出的拟合参数（期望和标准差）定义为paras维度，维度的大小为2，通过output_sizes参数设置。这样我们输入[time, lev, lat, lon]的数据，在每个空间点对时间维度拟合之后，输出的数据为[lev, lat, lon, paras]。（PS: 这里感谢深雨露大佬的指点） 9 O1 ~5 S5 d1 ~% @* [# l7 h9 {, S6 g
9 w9 J6 H. R# C) O
以下是全部代码： " L( p/ Q' ]8 z* C/ L
from scipy import stats as st
) S6 ~6 I8 X* g, z" l! E) _& ]
( F2 H' g# c: W$ B6 e/ H& l' q: w import xarray as xr
" ~5 p$ d, k7 s8 Q5 } import dask- ~3 y# }) ~: U3 n' S! U* g$ C2 G
import numpy as np6 M$ \9 V/ e) d! H- y
from dask.diagnostics import ProgressBar
+ c0 W" e; M$ L1 e' J3 q: o% [$ i6 p, T v$ q+ |8 j# L
def norm_fit(data):
3 n* l% R" ?- i$ U2 b% d- | loc, scale = st.norm.fit(data)+ L; e3 G1 r5 G1 h. N3 A
return np.array([loc, scale])0 z9 t' L1 c/ x3 Q3 N
) ?5 |! {2 e$ ~. P rs = np.random.RandomState(0)
$ _) t7 W G7 E x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])7 x( t% {2 Z1 V! J! @( f
x = x.chunk({"lev":5, "lat":100,"lon":100})
0 w. [9 I; C0 N* G+ H+ p
" F' h: I3 P4 I. ?0 S! v #使用apply_ufunc计算，并用dask的并行计算" a. T; V U1 [* m& X4 I
result = xr.apply_ufunc(norm_fit,
2 [+ j) @+ e( { x,
. i# p3 c h3 z3 I9 y: w% L8 d input_core_dims=[["time"]],3 M, B5 J2 E/ U Z) I
output_core_dims=[["paras"]],
: |" w( H8 h t dask="parallelized",
% r( Z9 |. j. G- B+ F output_dtypes=[np.float],
, T' s) d1 L5 F+ H2 W3 E+ j$ _& u2 G6 a2 n dask_gufunc_kwargs={output_sizes:{"paras":2}}
2 v' Z* c; @, n( o0 J )
( t b+ Z+ R8 G
. P2 }+ L- d- C8 u# ] #compute进行真实的计算并显示进度+ L; m3 M1 A A, _3 n! p; P; }
with ProgressBar():& h( T% J7 b1 {& q
result = results.compute()
! r5 h& N: {5 R R2 G' p' I5 ^
3 {/ U2 ^8 O7 F5 P #结果冲命名保存到nc文件+ B q6 a* f: N
result = result.rename("norm_paras")! B+ d, D/ e8 C }5 [3 f
result = result.to_netcdf("norm_fit_paras.nc",compute=False)
$ o4 M5 Q! A, o- k+ O with ProgressBar():
( t0 ^# A# }% ? }8 i8 e$ q; H result.compute() % b! B/ i) C( _0 C
转自：气海同途
- N% q3 N# e/ I) S0 f 关注【Ai尚研修科研技术平台】公众号，查看更多课程安排及免费数据资料
- t3 |6 l0 W$ j" J; l" I S$ |# H
) U7 d1 Q7 H5 g1 Z3 S! y 推荐：
4 m7 v8 |) i9 X- z( ?" t 1、Python语言在地球科学领域中的应用实践技术应用 ( j1 i7 N+ ~% Z! y$ n) c: c
2、【夏令营】针对课题组人员AI培养计划:“开启AI科研之路” 6 m; N1 [- `# Z: ?& y- u, u
3、Python在气象与海洋中的实践技术应用精品课程
- o. D: O3 u! W( v7 S! R' p 4、Python在WRF模型自动化运行及前后处理中的实践技术应用
9 j' Y% K( R& ^' D 5、全套Python机器学习核心技术与案例分析实践应用视频
+ H. R; O3 I# t5 k- Y5 G 6、全套区域高精度地学模拟-WRF气象建模、多案例应用与精美制图精品课程
+ O) Y2 o; u4 y. W) o 7、WRF DA资料同化系统理论、运行与与变分、混合同化新方法技术应用视频教程 L5 C8 a5 }2 C, b* T4 K
# h9 ?) E) v4 X8 a
9 ~. v. `8 m9 k
# i& ?! V5 u1 l0 B9 J
- s m2 _0 m' G8 q4 A. G/ G) x( j | 
