|
1 u) K: z. Y% O5 W& e& h 气象海洋领域，常常会涉及到大规模的数据处理，比如高时空分辨率的模式数据、雷达数据和卫星数据等。 / A- A7 ?* S+ }( l& o) z: v3 B5 x
b7 \/ y* ^3 t+ W9 Q 大部分情况，我们一般只会涉及小规模数据处理，对计算效率并不会太过追求。但是当数据量变大时，低效的数据处理所耗费的时间非常明显，因此高效的数据处理方式尤为重要。 $ Q" l, Z, |5 q9 W- c
$ P/ g+ j1 t4 X1 ^4 A1 c 本篇以拟合一个高维数据的正太分布参数为例，介绍如何使用xarray+dask加速数据处理。
6 z6 d" G3 |8 j% n. x. S8 ~5 }) A( J. w/ y! U8 E7 L# Y2 N
数据维度X[time, lev, lat, lon]，需要对三维空间每一点，沿着时间维度做正太分布拟合(正太分布拟合只是作为例子，这里可以定义你需要的操作函数)。 1 D' i% M7 w P
其中几个关键点解释一下：
& ^7 r9 o5 N6 M3 Q8 o （1）首先定义拟合正太分布的函数
' \% G7 c7 D5 D3 u- x- y6 }% z5 F def norm_fit(data):
, |- V& |( L( ~" _7 W5 x3 Y loc, scale = st.norm.fit(data)
% I& f# ~1 |1 D! ~9 D, } return np.array([loc, scale]) 2 j4 D4 K ]0 T
这里需要注意的是，拟合的函数，输出参数，需要打包为一个数组。并且它的维度需要和后面aplly_ufun定义的输出维度一致。 0 [8 U* t7 ^9 K, |! G) ~, n T' e
（2）xarray分块
+ `9 l. q9 A$ S+ o- j x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])7 W4 z% j: z( u' {1 ?8 g6 v6 N, [
x = x.chunk({"lev":5, "lat":100,"lon":100})
; r( L" k* {; o7 e j xarray结合dask可以将一个大型数组分成一个个数据块（chunk），需要注意的是我们需要沿着时间维操作，拟合需要整个时间维度的数据，因此时间维time不能分块，只能对其他维度分块。 & T7 w9 W6 h! F. _% f1 s+ ?3 i% I
（3）xarray.apply_ufunc函数
8 }% p5 j: d# p& }1 c. p result = xr.apply_ufunc(norm_fit,) ^/ ]% i9 y0 A+ G f
x,; r, k, u. M7 k( n6 i: d
input_core_dims=[["time"]],1 r2 W: b3 d! r* \3 e2 r& t5 k* J
output_core_dims=[["paras"]],4 M/ v. e' y7 p' {7 ]
dask="parallelized",
6 j3 i" O- W. n8 D' _* i9 K output_dtypes=[np.float],9 ~" B% M; c. _" D4 N
dask_gufunc_kwargs={output_sizes:{"paras":2}}
- x! G% G# n: [ ) ) i( v4 ^) L* W. _8 [4 p
apply_ufunc函数具体可以参考官网教程，这里只说使用时遇到的困难，即如何定义输出维度：输出维度是用output_core_dim定义的，将输出的拟合参数（期望和标准差）定义为paras维度，维度的大小为2，通过output_sizes参数设置。这样我们输入[time, lev, lat, lon]的数据，在每个空间点对时间维度拟合之后，输出的数据为[lev, lat, lon, paras]。（PS: 这里感谢深雨露大佬的指点）
$ Q* ]+ s I( G- u
9 X8 z2 j) ?: ^ 以下是全部代码： & Q: ~/ d5 n2 v! q6 N
from scipy import stats as st4 g# M5 {3 u6 p
8 e- M6 B6 Z* K" W! x3 r6 M& Y0 H
import xarray as xr
$ G6 p# b$ M: q, b( u0 h( X import dask9 \* ? J' \; H. c$ U' P
import numpy as np
0 C' \, n, D# G- _ from dask.diagnostics import ProgressBar+ H3 J s: Q: i1 z0 A
4 R- y4 d4 r5 Z; g' J$ M def norm_fit(data):' Q$ e) i4 T. i# _) v
loc, scale = st.norm.fit(data)
! V. A: }( k* I9 T {& t+ J. K return np.array([loc, scale])1 O3 b, y/ `, C0 o7 g# @) N
1 I- U% K+ F- z8 u rs = np.random.RandomState(0)
7 m& h: b: B6 t- A! G/ L, y# [ x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])
1 B! j: H3 l, R x = x.chunk({"lev":5, "lat":100,"lon":100})
5 u$ P+ X. J, I% L) u7 X2 @ A
- W$ A6 W5 v, a" I: e6 b: p #使用apply_ufunc计算，并用dask的并行计算% k. Q7 N% }& h! r0 ?" }
result = xr.apply_ufunc(norm_fit,
" k. U& [% g% L/ _4 }! ~ x," V2 ~- z( X- m
input_core_dims=[["time"]],
; @( }$ j: o) t+ `- h8 k. r output_core_dims=[["paras"]],
1 @1 g; l' |, R: I+ C dask="parallelized",
" I% N4 d0 D2 h _9 ~ output_dtypes=[np.float],
7 x! W, L/ Z8 n$ a dask_gufunc_kwargs={output_sizes:{"paras":2}}* B" [; P$ I( M) T
)
9 E/ Z/ e- n- K0 I
8 W- n- O1 ^" o #compute进行真实的计算并显示进度
4 x9 K8 o, t1 n+ I: j8 p with ProgressBar():/ z3 |3 _. y& M5 e" i# N
result = results.compute()% M6 ~- C6 l9 Z" b
0 j% ~1 I, q% c; p d4 m #结果冲命名保存到nc文件# B* j: v, ]4 u* y/ ^) X
result = result.rename("norm_paras")
' v2 s9 R+ F' Q4 V+ [# [; r" ~ result = result.to_netcdf("norm_fit_paras.nc",compute=False)& H# E7 c: @) O! U
with ProgressBar():
* c% {# G" F4 V6 N# w2 z2 B u result.compute() 2 o1 A1 f3 S$ Y; {
转自：气海同途
9 i3 i) K9 z" @& [/ r8 _9 T+ Q 关注【Ai尚研修科研技术平台】公众号，查看更多课程安排及免费数据资料 . }% F" n2 ^+ B! Y X' D' b$ Y
9 A+ G( ]* ^8 Z+ h# T4 X ~ 推荐： , d8 V4 ^: z |/ f* F
1、Python语言在地球科学领域中的应用实践技术应用
* B) d/ b' E0 R0 n( v 2、【夏令营】针对课题组人员AI培养计划:“开启AI科研之路”
+ ^: u/ p( F& z$ k5 u' v9 i$ P 3、Python在气象与海洋中的实践技术应用精品课程 $ Y2 N, t1 U! e# @; l
4、Python在WRF模型自动化运行及前后处理中的实践技术应用 ) b/ ~3 b3 G9 d- e( Q2 C
5、全套Python机器学习核心技术与案例分析实践应用视频
, ~( f: r |' R( N: Q 6、全套区域高精度地学模拟-WRF气象建模、多案例应用与精美制图精品课程
; W9 E3 N' c) I# O2 x& w 7、WRF DA资料同化系统理论、运行与与变分、混合同化新方法技术应用视频教程
# e# D, |- G ?+ O: Q0 ^; S7 B% w" q3 l/ {' S0 a
$ S, L& J6 n. ], a4 A5 Y9 }1 s# h6 W7 O7 y: O' C3 ^
4 q# z$ E# f) q) c4 x, X9 G7 I | 
