|
# B8 f+ W- f% p' @6 g
气象海洋领域，常常会涉及到大规模的数据处理，比如高时空分辨率的模式数据、雷达数据和卫星数据等。 " S( G- @: v+ l3 X: J; q7 j: I
5 Z$ ?' W9 x0 U% O. q 大部分情况，我们一般只会涉及小规模数据处理，对计算效率并不会太过追求。但是当数据量变大时，低效的数据处理所耗费的时间非常明显，因此高效的数据处理方式尤为重要。 5 _7 _4 x9 J3 J4 d$ e6 S! Y9 C
3 x: D4 C" w" q+ ~" }9 ~
本篇以拟合一个高维数据的正太分布参数为例，介绍如何使用xarray+dask加速数据处理。 / `* J; \2 `& ^. J. A: c/ N
, h" U9 C6 l; J% ]+ e* M 数据维度X[time, lev, lat, lon]，需要对三维空间每一点，沿着时间维度做正太分布拟合(正太分布拟合只是作为例子，这里可以定义你需要的操作函数)。
- o7 ^) }9 w/ w* y. [ 其中几个关键点解释一下：
, J3 r# o6 Q2 q' ` （1）首先定义拟合正太分布的函数 : Y2 ~; Q* M+ o& V9 c3 p1 t/ @. b3 X
def norm_fit(data):2 p+ r2 A) `3 _5 `
loc, scale = st.norm.fit(data)- c7 B. j- w7 O2 i8 m3 S
return np.array([loc, scale]) " n2 G7 U R( D
这里需要注意的是，拟合的函数，输出参数，需要打包为一个数组。并且它的维度需要和后面aplly_ufun定义的输出维度一致。 - S0 F. G( l% Y6 k7 I: ?# T# v
（2）xarray分块 3 C" W* w) ]! `4 y* d1 C5 W" t8 S
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])
8 M; `, T- L' X4 x. b# n x = x.chunk({"lev":5, "lat":100,"lon":100}) " R- y! m' S: q4 \7 y j: r
xarray结合dask可以将一个大型数组分成一个个数据块（chunk），需要注意的是我们需要沿着时间维操作，拟合需要整个时间维度的数据，因此时间维time不能分块，只能对其他维度分块。
1 P1 k- c% n3 u& Q- I( ]: G （3）xarray.apply_ufunc函数
9 R% U; i9 {' R& ~( V# D8 _ result = xr.apply_ufunc(norm_fit,
( e6 Q, I+ U8 t( k; z0 h" z) c5 x x,
; v8 Q! u* k3 B( R7 z) `, f4 S: V input_core_dims=[["time"]],1 p" i- q) r' F
output_core_dims=[["paras"]],$ U; O/ D' d" \( ]- }: H- `
dask="parallelized",
7 B& p" H* t) o% x5 ~ output_dtypes=[np.float],2 \" e: s5 ]# R0 C# A, v
dask_gufunc_kwargs={output_sizes:{"paras":2}}- c8 A4 r! O7 m, J+ K7 Q h
)
1 F: K5 d7 Z, P8 q! ^ apply_ufunc函数具体可以参考官网教程，这里只说使用时遇到的困难，即如何定义输出维度：输出维度是用output_core_dim定义的，将输出的拟合参数（期望和标准差）定义为paras维度，维度的大小为2，通过output_sizes参数设置。这样我们输入[time, lev, lat, lon]的数据，在每个空间点对时间维度拟合之后，输出的数据为[lev, lat, lon, paras]。（PS: 这里感谢深雨露大佬的指点）
7 i; H0 E: V' f/ V- a& v4 a2 ~2 A! B+ _
以下是全部代码： 5 R- b. Z, e8 f( l
from scipy import stats as st
; n; ^% q( P# S# G/ J- M P3 B# M6 l
import xarray as xr! ^5 @ a! M9 ?, V0 T9 \
import dask9 |; G, N, b3 r0 B% e8 x
import numpy as np9 G& o V0 }' A( S2 j3 v7 I
from dask.diagnostics import ProgressBar! _; r% G8 u: D @7 }! @
% S/ B' o2 x4 O* U
def norm_fit(data):
* @0 o* ?5 \8 r3 V O q loc, scale = st.norm.fit(data)& `# v3 X/ [( i
return np.array([loc, scale])/ h# M& y0 }0 s" {6 W' T7 C
, J, j# }3 S8 d: \0 N# x. Q
rs = np.random.RandomState(0)- Y( F8 }9 c; z6 G+ d
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"]) o% n# o: C' z- C7 P' s
x = x.chunk({"lev":5, "lat":100,"lon":100}). i" a# U2 ]5 q& W0 Y5 ?: |
9 D6 r5 w3 I* R) ~1 a& X
#使用apply_ufunc计算，并用dask的并行计算1 K, y: y) ~+ j" X
result = xr.apply_ufunc(norm_fit,
( K v! L! |+ B2 r0 `" C$ Y4 T x,
% G# P6 {& M6 C+ T2 U input_core_dims=[["time"]],
4 r- b, A6 G8 {3 v2 h( F) ^ output_core_dims=[["paras"]],# F- [ J6 V: T
dask="parallelized",: z& Q* i, S- d) A& I; G
output_dtypes=[np.float],3 T( i( C" [; k* W/ `' }& j
dask_gufunc_kwargs={output_sizes:{"paras":2}}7 | N7 K' N! n) ]
)0 A9 X0 P {9 {1 ?: Q: I
, o, \$ k6 n2 b# D7 p% Y #compute进行真实的计算并显示进度
3 u6 f0 F" _( T1 Z with ProgressBar():
' y& x/ a) Q, U8 y result = results.compute()
9 S$ ?* Q1 w# L
; x% P4 \2 w- x #结果冲命名保存到nc文件9 [8 g' Z- C7 J" n [. |3 R# t
result = result.rename("norm_paras")9 r: @7 F/ y, y( U
result = result.to_netcdf("norm_fit_paras.nc",compute=False), h9 V, W* ^' Q# o* _2 a6 ?( g7 b
with ProgressBar():8 A+ d* I O8 S2 X4 r" ]0 B
result.compute()
Q: H( E+ k+ c% a 转自：气海同途 , W- ]2 L2 }. S* M
关注【Ai尚研修科研技术平台】公众号，查看更多课程安排及免费数据资料 % C- |/ t) X6 ~$ `
+ M: O5 `6 n! B# I8 J
推荐：
' T& i2 |1 t0 X1 U 1、Python语言在地球科学领域中的应用实践技术应用
. d, r% v4 o7 v' U$ N 2、【夏令营】针对课题组人员AI培养计划:“开启AI科研之路” ( N( A$ l6 O% j4 H u: h5 P
3、Python在气象与海洋中的实践技术应用精品课程
3 Y4 ~2 z. x3 U. H( ^2 `% b6 S 4、Python在WRF模型自动化运行及前后处理中的实践技术应用 ( W/ y6 a4 @7 k f
5、全套Python机器学习核心技术与案例分析实践应用视频 # y% ^9 f: W" ^) X
6、全套区域高精度地学模拟-WRF气象建模、多案例应用与精美制图精品课程
) y2 P4 ]& c9 M: M; b 7、WRF DA资料同化系统理论、运行与与变分、混合同化新方法技术应用视频教程
% X A9 L9 R* F: E* e; y* ~) k$ @# U! B& z/ p: c8 I7 L
& E! Y4 U' O* s- Z3 R5 U) m
4 @* w$ [* i: @* n$ X Q
( c( z) s/ @' N1 A9 s, r | 
