|
) N. E6 G- |9 Y8 r; v6 r9 h
气象海洋领域,常常会涉及到大规模的数据处理,比如高时空分辨率的模式数据、雷达数据和卫星数据等。
2 C: _4 F9 ?7 h, X! w0 p8 w 1 x8 p& Z( b, U9 O: d
大部分情况,我们一般只会涉及小规模数据处理,对计算效率并不会太过追求。但是当数据量变大时,低效的数据处理所耗费的时间非常明显,因此高效的数据处理方式尤为重要。
6 ?3 P- T* B& i! z# Y9 R* [. D ~1 G4 m+ i
本篇以拟合一个高维数据的正太分布参数为例,介绍如何使用xarray+dask加速数据处理。 8 m6 D4 T( b0 G) E2 g6 G
$ h. H- o- g8 x7 p9 v 数据维度X[time, lev, lat, lon],需要对三维空间每一点,沿着时间维度做正太分布拟合(正太分布拟合只是作为例子,这里可以定义你需要的操作函数)。 . v! t( A# v% L
其中几个关键点解释一下: / ]0 t, l2 b5 d: q& ~
(1)首先定义拟合正太分布的函数
% o! t% E k& \$ O def norm_fit(data):+ D- T( d {; R. Q. a" j2 K
loc, scale = st.norm.fit(data)0 Z5 a8 v. G: A4 c7 A# E9 S+ H
return np.array([loc, scale]) p+ z+ `# G9 u6 K: a( A
这里需要注意的是,拟合的函数,输出参数,需要打包为一个数组。并且它的维度需要和后面aplly_ufun定义的输出维度一致。
3 T1 y" h& p. m; k' A' S5 u (2)xarray分块 2 w% b2 d$ T/ s' u3 Y+ q" l
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])
% {2 }4 l: O5 c9 P+ [0 n0 r, h" @ x = x.chunk({"lev":5, "lat":100,"lon":100}) . P, \' S8 k& o
xarray结合dask可以将一个大型数组分成一个个数据块(chunk),需要注意的是我们需要沿着时间维操作,拟合需要整个时间维度的数据,因此时间维time不能分块,只能对其他维度分块。
; `" f, y5 y7 N: @ (3)xarray.apply_ufunc函数 , _% Z7 [0 y/ q; z# w+ H( n: C; o% h
result = xr.apply_ufunc(norm_fit,7 q Q* G7 }3 D1 m% v; m& ~
x,
& W4 P) A/ h" p. H, y/ T input_core_dims=[["time"]],
1 h, J( h2 [4 _8 Y, ` output_core_dims=[["paras"]],
" i# r! ], B2 h2 t dask="parallelized",
& B( n0 E# v7 c) D) N output_dtypes=[np.float],
; L8 P; y! ?5 j* g/ W/ x4 i dask_gufunc_kwargs={output_sizes:{"paras":2}}
: P1 m% w1 E$ a7 G2 s )
# Y6 h6 _/ R) Q. T) n; ^) S apply_ufunc函数具体可以参考官网教程,这里只说使用时遇到的困难,即如何定义输出维度:输出维度是用output_core_dim定义的,将输出的拟合参数(期望和标准差)定义为paras维度,维度的大小为2,通过output_sizes参数设置。这样我们输入[time, lev, lat, lon]的数据,在每个空间点对时间维度拟合之后,输出的数据为[lev, lat, lon, paras]。(PS: 这里感谢深雨露大佬的指点)
+ W' Q) ]( D1 j0 O2 _* e/ q0 h$ s+ B) ~
以下是全部代码:
5 B2 X7 t( G9 i1 F$ I from scipy import stats as st
6 W" P0 @9 i* r* M ?0 _$ B) s
import xarray as xr
1 Y+ L+ N; t- t- h import dask+ k3 y3 q3 W2 c) D1 Q: B% ^" x8 u
import numpy as np
Z4 W* ^, S8 e8 E0 d from dask.diagnostics import ProgressBar" t* t% |& E% I. l
: q$ D8 q/ f; s5 x! _* T def norm_fit(data):
, s5 {" C: }" J) ]+ h loc, scale = st.norm.fit(data)
3 G: D7 v8 b F* ?4 \$ O' u return np.array([loc, scale])& y9 R+ h1 _" J6 L8 p/ X' z/ ^
; w& a$ X9 b+ l$ D( \6 h, \
rs = np.random.RandomState(0)) \/ P/ h% v, F3 p
x = xr.DataArray(rs.randn(500, 20, 1500, 1500), dims=["time", "lev", "lat","lon"])
: P; p; o8 f4 H6 Q x = x.chunk({"lev":5, "lat":100,"lon":100})) K @2 x9 C1 W3 E5 d! k
! u7 `- b) {5 A1 K/ `2 W+ L0 M/ d #使用apply_ufunc计算,并用dask的并行计算
% }5 ~3 i5 a2 c8 \2 o result = xr.apply_ufunc(norm_fit,
) ^8 r3 g0 Y# Q" B, C/ } x,$ l# j6 _, w9 n
input_core_dims=[["time"]],
: k) ~7 ~ J2 F output_core_dims=[["paras"]],2 R0 J% z* v# o. `6 t% x, S7 A! O! G
dask="parallelized",
- P M+ \6 j9 ] output_dtypes=[np.float],- o4 G0 f( b: W8 B
dask_gufunc_kwargs={output_sizes:{"paras":2}}- Q/ V0 ]; J" y- c7 n; ]
)& [4 r7 T, ?* f5 c. Q5 s# K2 q6 o
g& I8 ^* E% s
#compute进行真实的计算并显示进度
; R4 {4 W1 a/ ]0 Q7 L with ProgressBar(): f! f, Q' A& m' e* U
result = results.compute()
9 v( o" q7 {+ P0 g1 U/ I& c, s. H; ]; Y
#结果冲命名保存到nc文件- _7 p+ a+ Q2 P9 U0 C+ j6 s
result = result.rename("norm_paras")' O% f/ T1 l) b: y/ u
result = result.to_netcdf("norm_fit_paras.nc",compute=False)
/ Q* O& I" G( s4 v" c" w with ProgressBar():# R1 R8 _& j( P4 g! s& k/ i
result.compute()
3 _. z, Y3 P9 O# @0 f. j% p 转自:气海同途 # _; O: c5 G: o+ F3 d
关注【Ai尚研修科研技术平台】公众号,查看更多课程安排及免费数据资料
7 e% B E5 o. A! N; B5 K4 M & ]0 K: U) _1 I5 r
推荐:
5 A7 v e5 |1 B, s 1、Python语言在地球科学领域中的应用实践技术应用
6 A. o/ g4 c! X2 j- i 2、【夏令营】针对课题组人员AI培养计划:“开启AI科研之路”
5 f& ?# E) h; g% ~6 L( O l 3、Python在气象与海洋中的实践技术应用精品课程
7 w g9 N3 p1 G2 j) I6 ^# d 4、Python在WRF模型自动化运行及前后处理中的实践技术应用 9 a* {3 }3 D# ^: ]# D
5、全套Python机器学习核心技术与案例分析实践应用视频
9 _4 b/ ]$ e0 }: q+ p 6、全套区域高精度地学模拟-WRF气象建模、多案例应用与精美制图精品课程 5 \+ [6 g6 y9 B
7、WRF DA资料同化系统理论、运行与与变分、混合同化新方法技术应用视频教程 " M, D# c% q$ Y4 p y
# I( o5 u! n: m
* ?9 P6 F$ d Y9 i0 M6 A$ z8 X% y! J, l; U! ~0 J
1 o# @$ _, d- m |