|
U0 \; ^5 o( w, R: F+ Z
气象海洋领域，常常会涉及到大规模的数据处理，比如高时空分辨率的模式数据、雷达数据和卫星数据等。
' N" F$ s8 E& b * O8 \8 [) s% j8 S
大部分情况，我们一般只会涉及小规模数据处理，对计算效率并不会太过追求。但是当数据量变大时，低效的数据处理所耗费的时间非常明显，因此高效的数据处理方式尤为重要。
0 L( y, z. Y6 v4 V5 [; _6 l) p' Q: Z# j! k, `( z
本篇以拟合一个高维数据的正太分布参数为例，介绍如何使用xarray+dask加速数据处理。 - C$ n+ g, m7 a% [% i3 j. t; v
$ R; E( \. h7 T" Z 数据维度X[time, lev, lat, lon]，需要对三维空间每一点，沿着时间维度做正太分布拟合(正太分布拟合只是作为例子，这里可以定义你需要的操作函数)。
* c- v/ t- J! V' C9 d 其中几个关键点解释一下：
: I. Z+ P" G1 R: A6 X% } （1）首先定义拟合正太分布的函数 $ |- I/ l% f# B/ _4 o, C
def norm_fit(data):# E- P0 v. e% U7 U3 [9 K) H) F g
loc, scale = st.norm.fit(data)
: G$ ?$ m4 ?$ ?9 ?7 n8 Q return np.array([loc, scale]) " K# a. L \5 x% B
这里需要注意的是，拟合的函数，输出参数，需要打包为一个数组。并且它的维度需要和后面aplly_ufun定义的输出维度一致。 & V# e, F+ \- L. B7 ]
（2）xarray分块 / e9 U' O( L" M* i# W& f
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])
5 S) c) O2 a+ U3 r2 r A: S x = x.chunk({"lev":5, "lat":100,"lon":100})
; N: U% s A) T3 F' u) j5 u5 _& j' [ xarray结合dask可以将一个大型数组分成一个个数据块（chunk），需要注意的是我们需要沿着时间维操作，拟合需要整个时间维度的数据，因此时间维time不能分块，只能对其他维度分块。
9 ?2 i1 K7 L4 E0 r3 }( ^ （3）xarray.apply_ufunc函数 , U+ M6 p# j; ~9 L |* ? K
result = xr.apply_ufunc(norm_fit,
6 Z' {+ ~- s1 I: H! n1 I& h5 [) Y x,3 I3 d& M# h0 x* x3 _, G- N
input_core_dims=[["time"]],; t' s- x, s3 ^/ A. P3 Y- J* p
output_core_dims=[["paras"]],$ [# [8 P4 w! q" r) p8 g( S/ f
dask="parallelized",9 Z, n% M$ R# m: b7 Z
output_dtypes=[np.float],
0 \" b% o5 h+ n dask_gufunc_kwargs={output_sizes:{"paras":2}}4 b9 k$ Y; P/ I8 Y- l7 c5 K
) ' ^8 X8 G" y) ^3 C. d0 U7 W
apply_ufunc函数具体可以参考官网教程，这里只说使用时遇到的困难，即如何定义输出维度：输出维度是用output_core_dim定义的，将输出的拟合参数（期望和标准差）定义为paras维度，维度的大小为2，通过output_sizes参数设置。这样我们输入[time, lev, lat, lon]的数据，在每个空间点对时间维度拟合之后，输出的数据为[lev, lat, lon, paras]。（PS: 这里感谢深雨露大佬的指点） 6 d! w/ g; p2 e& J7 x1 `+ @
) t/ J) H* e( Z( ?6 R+ W' p 以下是全部代码： / n Y$ P: Z! V8 n7 Y% x
from scipy import stats as st
3 i. y* L& b G/ L3 @: M/ N* j
8 M8 F7 C: |; f+ \, R" x/ g import xarray as xr" a( h2 \ G. Q- g* ^9 T! E
import dask; w" g9 R9 W7 b
import numpy as np8 K$ H7 E& b2 ^/ Z9 @' T$ c
from dask.diagnostics import ProgressBar
0 q1 ]! W2 k7 B4 N. N0 o0 [
; E$ ]8 y% c6 K$ J: j" @ def norm_fit(data):& q* t, H5 x# O) F! b$ V" _# k; O
loc, scale = st.norm.fit(data)1 j) G2 G v5 V y0 q# I4 m
return np.array([loc, scale])( W0 m8 ]6 }4 K, P* N) ?
3 E! ] C0 t3 D( n- b
rs = np.random.RandomState(0)
, N. Q- Z. O8 W: y x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])) h8 |9 A! U3 z; S, U; Y6 E( g
x = x.chunk({"lev":5, "lat":100,"lon":100})( |! t. r1 s5 j1 `% w
8 b2 K0 E+ `9 s7 X R) d3 e
#使用apply_ufunc计算，并用dask的并行计算- V9 m3 I& Z2 v% G
result = xr.apply_ufunc(norm_fit,
" g5 [1 o: @8 J8 L1 Y5 ? x,$ I# p3 v' P# A3 Y/ e
input_core_dims=[["time"]],. J! u1 F, @/ u! R' y
output_core_dims=[["paras"]],
4 a) Y& i, Y) \0 ?+ b A/ h dask="parallelized",
0 h( b1 j8 X, W7 w% z output_dtypes=[np.float],
+ o/ T3 a1 Q3 z dask_gufunc_kwargs={output_sizes:{"paras":2}}* i( i) n4 J$ _
)
' a+ r, L9 Z( W+ w2 K9 B0 Y ?( m* k' N5 R4 J
#compute进行真实的计算并显示进度
, k9 A7 d! t' H7 B! Q; g with ProgressBar():
* t- M/ L. W3 R9 l- Q* a result = results.compute()
" Y5 j x6 M! g* x9 t8 H/ j- F- I5 e% _
! T& A! u& f+ b, @ #结果冲命名保存到nc文件" o# M9 z! {+ l) l& J" ]2 _
result = result.rename("norm_paras")
! T2 c9 f; q. J& {1 Z! B$ r' ] result = result.to_netcdf("norm_fit_paras.nc",compute=False)3 I3 I" b! y8 n% s
with ProgressBar():
. y4 V! C5 f: y result.compute()
: |6 B2 B K* r% ` 转自：气海同途 8 \# ]( u. X. P) }
关注【Ai尚研修科研技术平台】公众号，查看更多课程安排及免费数据资料 # r' H! c5 B: H; `4 X' i
P5 N1 ]. w% R% R2 r4 M
推荐：
6 ]# z9 b2 M+ w4 @5 J. O+ t 1、Python语言在地球科学领域中的应用实践技术应用 $ z6 M( G- B) Q+ i: X& W/ ^
2、【夏令营】针对课题组人员AI培养计划:“开启AI科研之路” # i3 H, N* {$ ?. W" m
3、Python在气象与海洋中的实践技术应用精品课程 ! ]- _, k. e; \, g2 R2 d- Y; n
4、Python在WRF模型自动化运行及前后处理中的实践技术应用
* W) G: D7 P/ }5 Y X1 Z 5、全套Python机器学习核心技术与案例分析实践应用视频 4 b. B$ m8 B; n B) I) B
6、全套区域高精度地学模拟-WRF气象建模、多案例应用与精美制图精品课程
6 g5 D- Q; V8 Z5 W- X 7、WRF DA资料同化系统理论、运行与与变分、混合同化新方法技术应用视频教程
' |1 U3 t5 g( `) ?5 Z8 C- A4 S3 T, T; c% I2 @( D
1 n3 v' m' N5 D& q h5 ^: H1 N# _3 a1 H
7 V4 c) g- ]. z0 Q5 k
| 
