|
( D1 J% D0 e8 U; _$ v7 ?% K" | 气象海洋领域，常常会涉及到大规模的数据处理，比如高时空分辨率的模式数据、雷达数据和卫星数据等。
5 a2 K6 a8 M! y: i6 G% p, C9 W 9 M. h% x- T0 {5 p" \ H
大部分情况，我们一般只会涉及小规模数据处理，对计算效率并不会太过追求。但是当数据量变大时，低效的数据处理所耗费的时间非常明显，因此高效的数据处理方式尤为重要。
# Q& L4 C! e* \+ K6 a- N' y+ k4 j4 k$ Z$ [- _' u
本篇以拟合一个高维数据的正太分布参数为例，介绍如何使用xarray+dask加速数据处理。
& R4 p" a4 O- b& U4 \' n- ]7 j* s) t' q D J8 w
数据维度X[time, lev, lat, lon]，需要对三维空间每一点，沿着时间维度做正太分布拟合(正太分布拟合只是作为例子，这里可以定义你需要的操作函数)。
! g9 D6 U" u) W. Z y+ p 其中几个关键点解释一下： 6 e" f/ ~- T8 q( c' j, n
（1）首先定义拟合正太分布的函数 / V" u) c0 A. V8 M; N
def norm_fit(data):+ s1 v5 B1 q- Y) X% O
loc, scale = st.norm.fit(data)
s& }3 J0 j1 Q, t- \ return np.array([loc, scale]) / u4 T+ ~3 ?# z# l, B
这里需要注意的是，拟合的函数，输出参数，需要打包为一个数组。并且它的维度需要和后面aplly_ufun定义的输出维度一致。
h$ M; k! P3 H （2）xarray分块
3 s7 h- c; y5 W3 ~ x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])
, L. t& F$ h' L x = x.chunk({"lev":5, "lat":100,"lon":100})
: M5 ^- P" T. B/ h xarray结合dask可以将一个大型数组分成一个个数据块（chunk），需要注意的是我们需要沿着时间维操作，拟合需要整个时间维度的数据，因此时间维time不能分块，只能对其他维度分块。
. E1 x4 A f. X （3）xarray.apply_ufunc函数 ( A, e3 |. ] I7 x6 o( c
result = xr.apply_ufunc(norm_fit,
. X4 @2 a* y9 B0 I x,
* W# {# d& j3 n8 b/ e$ t input_core_dims=[["time"]],1 f2 N# _$ Y- ~) _
output_core_dims=[["paras"]],
- U" G% P8 y1 h1 T7 h' E; x dask="parallelized",
3 A# i; D- S, `' E# x" ` output_dtypes=[np.float],( Y: O& @7 T/ Z2 i, ^: R* T5 V
dask_gufunc_kwargs={output_sizes:{"paras":2}}
v, k! f: P' G l: [$ [6 `! j" r ) ' }/ e7 g! \4 W' _
apply_ufunc函数具体可以参考官网教程，这里只说使用时遇到的困难，即如何定义输出维度：输出维度是用output_core_dim定义的，将输出的拟合参数（期望和标准差）定义为paras维度，维度的大小为2，通过output_sizes参数设置。这样我们输入[time, lev, lat, lon]的数据，在每个空间点对时间维度拟合之后，输出的数据为[lev, lat, lon, paras]。（PS: 这里感谢深雨露大佬的指点）
7 g. [7 |& {2 O
# Y5 d u- ]0 Z+ G7 t$ D 以下是全部代码：
: s" ^# D+ D4 q from scipy import stats as st
: u! x( j8 n0 {0 {2 i
2 S7 \" s& V: y5 a: o0 c f import xarray as xr
3 ^. s0 _ H8 Y2 u' _0 _' ? import dask3 D; M* f" C0 y; l+ K3 E0 J
import numpy as np
+ n. h( A1 m3 K from dask.diagnostics import ProgressBar* r4 v8 [$ q$ ] C* {
. ~; h+ a! ^% I# {; F# N def norm_fit(data):
) Z: J/ W4 x1 ^$ t, U loc, scale = st.norm.fit(data), s9 D' B5 d2 {
return np.array([loc, scale])
+ Q4 U6 e8 ^* U' Y
6 L. h1 `0 c2 Q4 V ~ rs = np.random.RandomState(0)
1 C5 Z8 g5 y1 W( C0 ~7 V x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])
& m# M9 a& E5 O" M' o: a V x = x.chunk({"lev":5, "lat":100,"lon":100}). r$ ?0 W0 t9 _3 s" F
& P l5 [2 V; l: r5 @ #使用apply_ufunc计算，并用dask的并行计算8 O" D( c2 [5 w, D# l
result = xr.apply_ufunc(norm_fit,
% W; W1 m7 p( Z$ `0 S: [ x,6 o7 [9 q0 ~4 V2 h% k8 [5 b
input_core_dims=[["time"]],4 T; u$ A5 @$ ~1 d6 D1 ]- m) a% Y
output_core_dims=[["paras"]],
: k. X; R) i8 Q) M0 H$ @1 K8 ^ dask="parallelized",: U) _' v: S" |3 H0 `
output_dtypes=[np.float],, X2 V+ ^. l& q6 C
dask_gufunc_kwargs={output_sizes:{"paras":2}}
7 T/ [7 b$ n; E! }/ C8 a )0 V0 ^& o$ c5 d
7 t! k7 w1 y0 ^6 N8 e; p #compute进行真实的计算并显示进度
/ I$ s/ y% G. G- O; z+ @7 P: w6 H with ProgressBar():9 ^; k( I' S6 e; }6 u
result = results.compute()
2 w3 `' F; c0 t3 [ @7 t; I* |8 ~ k) f7 v, [: Q
#结果冲命名保存到nc文件1 U+ j0 E4 T6 N* c# [) |
result = result.rename("norm_paras")6 Z% ?6 a: V1 f( W
result = result.to_netcdf("norm_fit_paras.nc",compute=False)
8 e* f0 C. Y; i5 m7 t2 H9 ^8 L& { with ProgressBar():( d6 m Y! v" n9 M9 R! w7 W
result.compute() , J) p- D8 i' {
转自：气海同途 & }3 P7 m: d! s/ B- o
关注【Ai尚研修科研技术平台】公众号，查看更多课程安排及免费数据资料
% w* S5 R }0 L) F & m% c* b+ Z' H
推荐： ! g% R8 L& R9 ?+ [* K$ l
1、Python语言在地球科学领域中的应用实践技术应用 9 @% F. a$ H- W E+ d" q
2、【夏令营】针对课题组人员AI培养计划:“开启AI科研之路” & U3 S. p" Z; C" T! e- U
3、Python在气象与海洋中的实践技术应用精品课程
. m( Y) I- ~/ ^8 W6 W8 o 4、Python在WRF模型自动化运行及前后处理中的实践技术应用
% X" F' Q. N0 g9 d; C5 n 5、全套Python机器学习核心技术与案例分析实践应用视频 . ~' U$ H ^. S$ c- Z
6、全套区域高精度地学模拟-WRF气象建模、多案例应用与精美制图精品课程
9 i# f. @" w& ?8 u, Z. N 7、WRF DA资料同化系统理论、运行与与变分、混合同化新方法技术应用视频教程
" P% J; P2 P: i$ H3 ~" r
$ G9 U+ X3 e: |# O3 n1 V4 r9 q$ u, C5 Q
0 q4 R9 K- T; y. }
( ?( Z4 m; c. p1 `( q
| 
