xarray (教程)第五章 计算

Computation

与DataArray和Dataset对象相关联的标签为计算提供了一些强大的快捷方式,特别是包括按维度名称进行聚合和广播。

Basic array math

使用单个DataArray的算术运算会自动对所有数组值进行矢量化处理(如numpy ):

arr = xr.DataArray(
    np.random.RandomState(0).randn(2, 3), [("x", ["a", "b"]), ("y", [10, 20, 30])]
)

arr - 3
"""Out[2]: 
<xarray.DataArray (x: 2, y: 3)>
array([[-1.23594765, -2.59984279, -2.02126202],
       [-0.7591068 , -1.13244201, -3.97727788]])
Coordinates:
  * x        (x) <U1 'a' 'b'
  * y        (y) int64 10 20 30"""

abs(arr)
"""Out[3]: 
<xarray.DataArray (x: 2, y: 3)>
array([[1.76405235, 0.40015721, 0.97873798],
       [2.2408932 , 1.86755799, 0.97727788]])
Coordinates:
  * x        (x) <U1 'a' 'b'
  * y        (y) int64 10 20 30"""

您还可以直接在DataArray上使用numpy或scipy的许多ufunc函数中的任何一个:

np.sin(arr)
"""Out[4]: 
<xarray.DataArray (x: 2, y: 3)>
array([[ 0.9813841 ,  0.38956314,  0.82979374],
       [ 0.78376151,  0.95628847, -0.82897801]])
Coordinates:
  * x        (x) <U1 'a' 'b'
  * y        (y) int64 10 20 30"""

使用where()有条件地在值之间切换:

xr.where(arr > 0, "positive", "negative")
"""Out[5]: 
<xarray.DataArray (x: 2, y: 3)>
array([['positive', 'positive', 'positive'],
       ['positive', 'positive', 'negative']], dtype='<U8')
Coordinates:
  * x        (x) <U1 'a' 'b'
  * y        (y) int64 10 20 30"""

使用@执行矩阵乘法:

arr @ arr
"""Out[6]: 
<xarray.DataArray ()>
array(13.69438174)"""

数据数组还实现了许多numpy.ndarray方法:

arr.round(2)
"""Out[7]: 
<xarray.DataArray (x: 2, y: 3)>
array([[ 1.76,  0.4 ,  0.98],
       [ 2.24,  1.87, -0.98]])
Coordinates:
  * x        (x) <U1 'a' 'b'
  * y        (y) int64 10 20 30"""

arr.T
"""Out[8]: 
<xarray.DataArray (y: 3, x: 2)>
array([[ 1.76405235,  2.2408932 ],
       [ 0.40015721,  1.86755799],
       [ 0.97873798, -0.97727788]])
Coordinates:
  * x        (x) <U1 'a' 'b'
  * y        (y) int64 10 20 30"""

intarr = xr.DataArray([0, 1, 2, 3, 4, 5])

intarr << 2  # only supported for int types
"""Out[10]: 
<xarray.DataArray (dim_0: 6)>
array([ 0,  4,  8, 12, 16, 20])
Dimensions without coordinates: dim_0"""

intarr >> 1
"""Out[11]: 
<xarray.DataArray (dim_0: 6)>
array([0, 0, 1, 1, 2, 2])
Dimensions without coordinates: dim_0"""

Missing values

Xarray使用NumPy中的“NaN”(非数字)值表示缺少的值,该值是一种特殊的浮点值,表示未定义或不可表示的值。
有几种方法可以处理xarray中的缺失值:
Xarray对象借用isnull()、notnull()、count()、dropna()、fillna()、ffill()和bfill()方法来处理pandas中缺失的数据:

  • isnull()是xarray中的一种方法,可用于检查xarray对象中是否缺少值或空值。它返回一个新的xarray对象,该对象与原始对象具有相同的尺寸,但使用布尔值来指示缺少值的位置。
x = xr.DataArray([0, 1, np.nan, np.nan, 2], dims=["x"])

x.isnull()
"""Out[13]: 
<xarray.DataArray (x: 5)>
array([False, False,  True,  True, False])
Dimensions without coordinates: x"""
  • notnull()是xarray中的一种方法,可用于检查xarray对象中的非缺失或非空值。它返回一个新的xarray对象,该对象具有与原始对象相同的尺寸,但具有布尔值,指示存在非缺失值的位置。
x = xr.DataArray([0, 1, np.nan, np.nan, 2], dims=["x"])

x.notnull()
"""Out[15]: 
<xarray.DataArray (x: 5)>
array([ True,  True, False, False,  True])
Dimensions without coordinates: x"""
  • count()是xarray中的一种方法,可用于计算xarray对象的一个或多个维度上非缺失值的数量。它返回一个新的xarray对象,该对象与原始对象具有相同的维度,但每个元素都被指定维度上的非缺失值的计数所替换。
x = xr.DataArray([0, 1, np.nan, np.nan, 2], dims=["x"])

x.count()
"""Out[17]: 
<xarray.DataArray ()>
array(3)"""
  • dropna()是xarray中的一个方法,可用于从xarray对象中删除缺失值或空值。它返回一个新的xarray对象,其尺寸与原始对象相同,但缺少的值已被删除。
x = xr.DataArray([0, 1, np.nan, np.nan, 2], dims=["x"])

x.dropna(dim="x")
"""Out[19]: 
<xarray.DataArray (x: 3)>
array([0., 1., 2.])
Dimensions without coordinates: x"""
  • fillna()是xarray中的一个方法,可用于用指定的值或方法填充xarray对象中缺失的值或空值。它返回一个新的xarray对象,其尺寸与原始对象相同,但缺少的值已被填充。
x = xr.DataArray([0, 1, np.nan, np.nan, 2], dims=["x"])

x.fillna(-1)
"""Out[21]: 
<xarray.DataArray (x: 5)>
array([ 0.,  1., -1., -1.,  2.])
Dimensions without coordinates: x"""
  • ffill()是xarray中的一种方法,可用于沿一维或多维向前填充(或向前填充)xarray对象中的缺失值。它返回一个新的xarray对象,该对象与原始对象具有相同的维度,但缺失值由指定维度上最后一个非缺失值替换。
x = xr.DataArray([0, 1, np.nan, np.nan, 2], dims=["x"])

x.ffill("x")
Out[23]:
< xarray.DataArray(x: 5) >
array([0., 1., 1., 1., 2.])
Dimensions
without
coordinates: x
  • bfill()是xarray中的一种方法,可用于沿一维或多维向后填充xarray对象中的缺失值。它返回一个新的xarray对象,该对象与原始对象具有相同的维度,但缺失值由指定维度上的下一个非缺失值替换。
x = xr.DataArray([0, 1, np.nan, np.nan, 2], dims=["x"])

x.bfill("x")
"""Out[25]: 
<xarray.DataArray (x: 5)>
array([0., 1., 2., 2., 2.])
Dimensions without coordinates: x"""

和pandas一样,xarray使用浮点值np.nan(非数字)来表示缺失值.
Xarray对象还有一个interpolate_na()方法,用于通过1D插值来填充缺失值。它返回一个新的xarray对象,该对象与原始对象具有相同的尺寸,但丢失的值被插值。

x = xr.DataArray(
    [0, 1, np.nan, np.nan, 2],
    dims=["x"],
    coords={"xx": xr.Variable("x", [0, 1, 1.1, 1.9, 3])},
)

x.interpolate_na(dim="x", method="linear", use_coordinate="xx")
"""Out[27]: 
<xarray.DataArray (x: 5)>
array([0.  , 1.  , 1.05, 1.45, 2.  ])
Coordinates:
    xx       (x) float64 0.0 1.0 1.1 1.9 3.0
Dimensions without coordinates: x"""

Aggregation

聚合方法已更新为采用dim参数而不是axis。这为应用于特定维度的聚合方法提供了非常直观的语法:

arr.sum(dim="x")
"""Out[28]: 
<xarray.DataArray (y: 3)>
array([4.00494555e+00, 2.26771520e+00, 1.46010423e-03])
Coordinates:
  * y        (y) int64 10 20 30"""
# 标准差
arr.std(["x", "y"])
"""Out[29]: 
<xarray.DataArray ()>
array(1.09038344)"""

arr.min()
"""Out[30]: 
<xarray.DataArray ()>
array(-0.97727788)"""

如果您需要自己计算维度的轴数(例如,包装用于numpy数组的代码),您可以使用get_axis_num()方法:


arr.get_axis_num("y")
# Out[31]: 1

这些操作会自动跳过缺少的值,例如pandas中的:

xr.DataArray([1, 2, np.nan, 3]).mean()
"""Out[32]: 
<xarray.DataArray ()>
array(2.)"""

如果需要,可以通过调用skipna=False的聚合方法来禁用此行为。

Rolling window operations

DataArray对象包含一个rolling()方法。此方法支持滚动窗口聚合:

arr = xr.DataArray(np.arange(0, 7.5, 0.5).reshape(3, 5), dims=("x", "y"))

"""arr
Out[34]: 
<xarray.DataArray (x: 3, y: 5)>
array([[0. , 0.5, 1. , 1.5, 2. ],
       [2.5, 3. , 3.5, 4. , 4.5],
       [5. , 5.5, 6. , 6.5, 7. ]])
Dimensions without coordinates: x, y"""

rolling()沿一个维度应用,使用维度的名称作为键(例如y),窗口大小作为值(例如3)。我们得到一个滚动的物体:

arr.rolling(y=3)
# Out[35]: DataArrayRolling [y->3]

聚合和汇总方法可以直接应用于滚动对象:

r = arr.rolling(y=3)

r.reduce(np.std)
"""Out[37]: 
<xarray.DataArray (x: 3, y: 5)>
array([[       nan,        nan, 0.40824829, 0.40824829, 0.40824829],
       [       nan,        nan, 0.40824829, 0.40824829, 0.40824829],
       [       nan,        nan, 0.40824829, 0.40824829, 0.40824829]])
Dimensions without coordinates: x, y"""

r.mean()
"""Out[38]: 
<xarray.DataArray (x: 3, y: 5)>
array([[nan, nan, 0.5, 1. , 1.5],
       [nan, nan, 3. , 3.5, 4. ],
       [nan, nan, 5.5, 6. , 6.5]])
Dimensions without coordinates: x, y"""

默认情况下,聚合结果在每个窗口的末尾被分配坐标,但在构造滚动对象时,可以通过传递center=True来居中:

r = arr.rolling(y=3, center=True)

r.mean()
"""Out[40]: 
<xarray.DataArray (x: 3, y: 5)>
array([[nan, 0.5, 1. , 1.5, nan],
       [nan, 3. , 3.5, 4. , nan],
       [nan, 5.5, 6. , 6.5, nan]])
Dimensions without coordinates: x, y"""

如上所述,与阵列边界重叠的窗口集合产生了nan。在滚动调用中设置min_periods会更改聚集时窗口内需要有值的最小观察数:

r = arr.rolling(y=3, min_periods=2)

r.mean()
"""Out[42]: 
<xarray.DataArray (x: 3, y: 5)>
array([[ nan, 0.25, 0.5 , 1.  , 1.5 ],
       [ nan, 2.75, 3.  , 3.5 , 4.  ],
       [ nan, 5.25, 5.5 , 6.  , 6.5 ]])
Dimensions without coordinates: x, y"""

r = arr.rolling(y=3, center=True, min_periods=2)

r.mean()
"""Out[44]: 
<xarray.DataArray (x: 3, y: 5)>
array([[0.25, 0.5 , 1.  , 1.5 , 1.75],
       [2.75, 3.  , 3.5 , 4.  , 4.25],
       [5.25, 5.5 , 6.  , 6.5 , 6.75]])
Dimensions without coordinates: x, y"""

虽然rolling提供了简单的移动平均线,但DataArray还支持具有rolling_exp()的指数移动平均线。这类似于pandas的ewm方法。numbagg是必需的。
rolling_exp方法采用window_type kwarg,它可以是“alpha”、“com”(代表质心)、“span”和“halflife”。默认值为span。
最后,rolling对象有一个构造方法,该方法返回最后一个位置有窗口维度的原始DataArray视图。您可以将其用于更高级的滚动操作,如步进滚动、窗口滚动、卷积、短时FFT等。

# rolling with 2-point stride
rolling_da = r.construct(x="x_win", y="y_win", stride=2)

rolling_da
"""Out[48]: 
<xarray.DataArray (x: 2, y: 3, x_win: 2, y_win: 3)>
array([[[[nan, nan, nan],
         [nan, nan, 0. ]],

        [[nan, nan, nan],
         [0. , 0.5, 1. ]],

        [[nan, nan, nan],
         [1. , 1.5, 2. ]]],


       [[[nan, nan, 2.5],
         [nan, nan, 5. ]],

        [[2.5, 3. , 3.5],
         [5. , 5.5, 6. ]],

        [[3.5, 4. , 4.5],
         [6. , 6.5, 7. ]]]])
Dimensions without coordinates: x, y, x_win, y_win"""

rolling_da.mean(["x_win", "y_win"], skipna=False)
"""Out[49]: 
<xarray.DataArray (x: 2, y: 3)>
array([[ nan,  nan,  nan],
       [ nan, 4.25, 5.25]])
Dimensions without coordinates: x, y"""

因为r . construct(‘window _
dim‘)给出的DataArray是原始数组的视图,所以它的内存效率很高。您还可以使用construct计算加权滚动和:

weight = xr.DataArray([0.25, 0.5, 0.25], dims=["window"])

arr.rolling(y=3).construct(y="window").dot(weight)
"""Out[51]: 
<xarray.DataArray (x: 3, y: 5)>
array([[nan, nan, 0.5, 1. , 1.5],
       [nan, nan, 3. , 3.5, 4. ],
       [nan, nan, 5.5, 6. , 6.5]])
Dimensions without coordinates: x, y"""

Weighted array reductions

DataArray和Dataset对象包括DataArray.weighted()和Dataset.weighted()数组归约方法。它们目前支持加权和、均值、标准差、var和分位数。

coords = dict(month=("month", [1, 2, 3]))

prec = xr.DataArray([1.1, 1.0, 0.9], dims=("month",), coords=coords)

weights = xr.DataArray([31, 28, 31], dims=("month",), coords=coords)
# Create a weighted object:
weighted_prec = prec.weighted(weights)

weighted_prec
# Out[56]: DataArrayWeighted with weights along dimensions: month
# Calculate the weighted sum:
weighted_prec.sum()
"""Out[57]: 
<xarray.DataArray ()>
array(90.)"""
# Calculate the weighted mean:
weighted_prec.mean(dim="month")
"""Out[58]: 
<xarray.DataArray ()>
array(1.)"""
# Calculate the weighted quantile:
weighted_prec.quantile(q=0.5, dim="month")
"""Out[59]: 
<xarray.DataArray ()>
array(1.)
Coordinates:
    quantile  float64 0.5"""
# The weighted sum corresponds to:
weighted_sum = (prec * weights).sum()

weighted_sum
"""Out[61]: 
<xarray.DataArray ()>
array(90.)"""
# the weighted mean to:
weighted_mean = weighted_sum / weights.sum()

weighted_mean
"""Out[63]: 
<xarray.DataArray ()>
array(1.)"""
# the weighted variance to:
weighted_var = weighted_prec.sum_of_squares() / weights.sum()

weighted_var
"""Out[65]: 
<xarray.DataArray ()>
array(0.00688889)"""
# and the weighted standard deviation to:
weighted_std = np.sqrt(weighted_var)

weighted_std
"""Out[67]: 
<xarray.DataArray ()>
array(0.08299933)"""
# 但是,这些函数也会考虑数据中缺失的值:
data = xr.DataArray([np.NaN, 2, 4])

weights = xr.DataArray([8, 1, 1])

data.weighted(weights).mean()
"""Out[70]: 
<xarray.DataArray ()>
array(3.)"""
# 使用(数据*权重).sum()/weights . sum()将(不正确地)得出0.6。
# 如果权重相加为0,sum返回0:
data = xr.DataArray([1.0, 1.0])

weights = xr.DataArray([-1.0, 1.0])

data.weighted(weights).sum()
"""Out[73]: 
<xarray.DataArray ()>
array(0.)"""

Coarsen large arrays

DataArray和Dataset对象包含一个 coarsen()和 coarsen()方法。这支持沿着多个维度的块聚合,

x = np.linspace(0, 10, 300)

t = pd.date_range("1999-12-15", periods=364)

da = xr.DataArray(
    np.sin(x) * np.cos(np.linspace(0, 1, 364)[:, np.newaxis]),
    dims=["time", "x"],
    coords={"time": t, "x": x},
)

da
"""Out[78]: 
<xarray.DataArray (time: 364, x: 300)>
array([[ 0.        ,  0.03343858,  0.06683976, ..., -0.48672119,
        -0.51565952, -0.54402111],
       [ 0.        ,  0.03343845,  0.06683951, ..., -0.48671934,
        -0.51565756, -0.54401905],
       [ 0.        ,  0.03343807,  0.06683875, ..., -0.4867138 ,
        -0.51565169, -0.54401285],
       ...,
       [ 0.        ,  0.0182217 ,  0.03642301, ..., -0.26522911,
        -0.28099849, -0.29645358],
       [ 0.        ,  0.01814439,  0.03626848, ..., -0.26410385,
        -0.27980632, -0.29519584],
       [ 0.        ,  0.01806694,  0.03611368, ..., -0.26297658,
        -0.27861203, -0.29393586]])
Coordinates:
  * time     (time) datetime64[ns] 1999-12-15 1999-12-16 ... 2000-12-12
  * x        (x) float64 0.0 0.03344 0.06689 0.1003 ... 9.9 9.933 9.967 10.0"""

为了沿时间维度每7天和沿x维度每2个点取一个块平均值,

da.coarsen(time=7, x=2).mean()
"""Out[79]: 
<xarray.DataArray (time: 52, x: 150)>
array([[ 0.01671847,  0.08349886,  0.14990579, ..., -0.41198807,
        -0.47195655, -0.52981418],
       [ 0.01671269,  0.08347003,  0.14985403, ..., -0.41184582,
        -0.47179359, -0.52963124],
       [ 0.01670071,  0.08341016,  0.14974655, ..., -0.41155042,
        -0.47145519, -0.52925136],
       ...,
       [ 0.00968205,  0.04835611,  0.0868139 , ..., -0.23859177,
        -0.2733209 , -0.30682759],
       [ 0.00941742,  0.04703446,  0.08444113, ..., -0.23207067,
        -0.26585059, -0.29844148],
       [ 0.00914929,  0.04569531,  0.08203696, ..., -0.22546326,
        -0.25828142, -0.2899444 ]])
Coordinates:
  * time     (time) datetime64[ns] 1999-12-18 1999-12-25 ... 2000-12-09
  * x        (x) float64 0.01672 0.08361 0.1505 0.2174 ... 9.849 9.916 9.983"""

如果数据长度不是相应窗口大小的倍数,则粗化()会引发ValueError。您可以选择boundary =‘trim‘或boundary
=‘pad‘选项来修剪多余的条目或填充不足的条目,

da.coarsen(time=30, x=2, boundary="trim").mean()
"""Out[80]: 
<xarray.DataArray (time: 12, x: 150)>
array([[ 0.01670121,  0.08341265,  0.14975103, ..., -0.41156272,
        -0.47146929, -0.52926718],
       [ 0.0165891 ,  0.08285275,  0.14874584, ..., -0.40880017,
        -0.46830462, -0.52571455],
       [ 0.01636376,  0.08172729,  0.14672529, ..., -0.40324704,
        -0.46194319, -0.51857326],
       ...,
       [ 0.01183847,  0.05912615,  0.10614938, ..., -0.29173175,
        -0.33419587, -0.37516528],
       [ 0.01082401,  0.05405954,  0.09705329, ..., -0.26673283,
        -0.30555813, -0.34301681],
       [ 0.00973567,  0.04862391,  0.08729468, ..., -0.23991312,
        -0.27483458, -0.30852683]])
Coordinates:
  * time     (time) datetime64[ns] 1999-12-29T12:00:00 ... 2000-11-23T12:00:00
  * x        (x) float64 0.01672 0.08361 0.1505 0.2174 ... 9.849 9.916 9.983"""

如果要将特定的函数应用于坐标,可以将函数或方法名传递给coord_func选项,

da.coarsen(time=7, x=2, coord_func={"time": "min"}).mean()
"""Out[81]: 
<xarray.DataArray (time: 52, x: 150)>
array([[ 0.01671847,  0.08349886,  0.14990579, ..., -0.41198807,
        -0.47195655, -0.52981418],
       [ 0.01671269,  0.08347003,  0.14985403, ..., -0.41184582,
        -0.47179359, -0.52963124],
       [ 0.01670071,  0.08341016,  0.14974655, ..., -0.41155042,
        -0.47145519, -0.52925136],
       ...,
       [ 0.00968205,  0.04835611,  0.0868139 , ..., -0.23859177,
        -0.2733209 , -0.30682759],
       [ 0.00941742,  0.04703446,  0.08444113, ..., -0.23207067,
        -0.26585059, -0.29844148],
       [ 0.00914929,  0.04569531,  0.08203696, ..., -0.22546326,
        -0.25828142, -0.2899444 ]])
Coordinates:
  * time     (time) datetime64[ns] 1999-12-15 1999-12-22 ... 2000-12-06
  * x        (x) float64 0.01672 0.08361 0.1505 0.2174 ... 9.849 9.916 9.983"""

Computation using Coordinates

Xarray对象有一些方便的方法来计算它们的坐标。differentiate()使用坐标通过中心有限差分计算导数,

a = xr.DataArray([0, 1, 2, 3], dims=["x"], coords=[[0.1, 0.11, 0.2, 0.3]])

a
"""Out[83]: 
<xarray.DataArray (x: 4)>
array([0, 1, 2, 3])
Coordinates:
  * x        (x) float64 0.1 0.11 0.2 0.3"""

a.differentiate("x")
"""Out[84]: 
<xarray.DataArray (x: 4)>
array([100.        ,  91.11111111,  10.58479532,  10.        ])
Coordinates:
  * x        (x) float64 0.1 0.11 0.2 0.3"""

integrate()使用坐标根据梯形规则计算积分

a.integrate("x")
"""Out[87]: 
<xarray.DataArray (y: 2)>
array([0.78, 0.98])
Dimensions without coordinates: y"""

Fitting polynomials

Xarray对象提供了一个使用最小二乘法执行线性或多项式回归的接口。polyfit()计算给定维度和给定顺序的最佳拟合系数,

x = xr.DataArray(np.arange(10), dims=["x"], name="x")

a = xr.DataArray(3 + 4 * x, dims=["x"], coords={"x": x})

out = a.polyfit(dim="x", deg=1, full=True)

out
"""Out[91]: 
<xarray.Dataset>
Dimensions:               (degree: 2)
Coordinates:
  * degree                (degree) int64 1 0
Data variables:
    x_matrix_rank         int64 2
    x_singular_values     (degree) float64 1.358 0.3963
    polyfit_coefficients  (degree) float64 4.0 3.0
    polyfit_residuals     float64 4.522e-28"""

该方法输出包含系数的数据集(如果full=True,则包含更多系数)。相反的操作是用polyval()完成的,

xr.polyval(coord=x, coeffs=out.polyfit_coefficients)
"""Out[92]: 
<xarray.DataArray (x: 10)>
array([ 3.,  7., 11., 15., 19., 23., 27., 31., 35., 39.])
Dimensions without coordinates: x"""

Fitting arbitrary functions

Xarray对象还提供了使用scipy.optimize.curve_fit()拟合更复杂函数的接口。curvefit()接受用户定义的函数,可以沿多个坐标拟合。
例如,我们可以拟合两个DataArray对象之间的关系,在每个空间坐标上保持唯一的拟合,但在时间维度上进行聚合:

def exponential(x, a, xc):
    return np.exp((x - xc) / a)


x = np.arange(-5, 5, 0.1)

t = np.arange(-5, 5, 0.1)

X, T = np.meshgrid(x, t)

Z1 = np.random.uniform(low=-5, high=5, size=X.shape)

Z2 = exponential(Z1, 3, X)

Z3 = exponential(Z1, 1, -X)

ds = xr.Dataset(
    data_vars=dict(
        var1=(["t", "x"], Z1), var2=(["t", "x"], Z2), var3=(["t", "x"], Z3)
    ),
    coords={"t": t, "x": x},
)

ds[["var2", "var3"]].curvefit(
    coords=ds.var1,
    func=exponential,
    reduce_dims="t",
    bounds={"a": (0.5, 5), "xc": (-5, 5)},
)

"""Out[101]: 
<xarray.Dataset>
Dimensions:                     (x: 100, param: 2, cov_i: 2, cov_j: 2)
Coordinates:
  * x                           (x) float64 -5.0 -4.9 -4.8 -4.7 ... 4.7 4.8 4.9
  * param                       (param) <U2 'a' 'xc'
  * cov_i                       (cov_i) <U2 'a' 'xc'
  * cov_j                       (cov_j) <U2 'a' 'xc'
Data variables:
    var2_curvefit_coefficients  (x, param) float64 3.0 -5.0 3.0 ... 4.8 3.0 4.9
    var2_curvefit_covariance    (x, cov_i, cov_j) float64 9.286e-14 ... 1.104...
    var3_curvefit_coefficients  (x, param) float64 0.9999 5.0 1.0 ... 1.0 -4.9
    var3_curvefit_covariance    (x, cov_i, cov_j) float64 5.825e-11 ... 1.474..."""

我们还可以拟合多维函数,甚至使用包装函数同时拟合几个函数的和,例如包含两个高斯峰值的字段:

def gaussian_2d(coords, a, xc, yc, xalpha, yalpha):
    x, y = coords
    z = a * np.exp(
        -np.square(x - xc) / 2 / np.square(xalpha)
        - np.square(y - yc) / 2 / np.square(yalpha)
    )
    return z


def multi_peak(coords, *args):
    z = np.zeros(coords[0].shape)
    for i in range(len(args) // 5):
        z += gaussian_2d(coords, *args[i * 5: i * 5 + 5])
    return z


x = np.arange(-5, 5, 0.1)

y = np.arange(-5, 5, 0.1)

X, Y = np.meshgrid(x, y)

n_peaks = 2

names = ["a", "xc", "yc", "xalpha", "yalpha"]

names = [f"{name}{i}" for i in range(n_peaks) for name in names]

Z = gaussian_2d((X, Y), 3, 1, 1, 2, 1) + gaussian_2d((X, Y), 2, -1, -2, 1, 1)

Z += np.random.normal(scale=0.1, size=Z.shape)

da = xr.DataArray(Z, dims=["y", "x"], coords={"y": y, "x": x})

da.curvefit(
    coords=["x", "y"],
    func=multi_peak,
    param_names=names,
    kwargs={"maxfev": 10000},
)

"""Out[113]: 
<xarray.Dataset>
Dimensions:                (param: 10, cov_i: 10, cov_j: 10)
Coordinates:
  * param                  (param) <U7 'a0' 'xc0' 'yc0' ... 'xalpha1' 'yalpha1'
  * cov_i                  (cov_i) <U7 'a0' 'xc0' 'yc0' ... 'xalpha1' 'yalpha1'
  * cov_j                  (cov_j) <U7 'a0' 'xc0' 'yc0' ... 'xalpha1' 'yalpha1'
Data variables:
    curvefit_coefficients  (param) float64 3.0 1.004 1.003 ... 1.007 1.008
    curvefit_covariance    (cov_i, cov_j) float64 3.362e-05 ... 2.125e-05"""

Broadcasting by dimension name

DataArray对象通过维度名称而不是轴顺序自动对齐自身(在numpy术语中称为“广播”)。使用xarray,您不需要转置数组或插入长度为1的维度来进行数组操作,而在numpy中通常使用numpy.reshape()或numpy.newaxis来进行这些操作。
有几个例子可以很好地说明这一点。考虑两个沿不同维度排列的不同大小的一维数组:
使用xarray,我们可以对这些数组应用二进制数学运算,并且它们的维数会自动扩展:
此外,尺寸总是按照它们首次出现的顺序重新排序:

a = xr.DataArray([1, 2], [("x", ["a", "b"])])

a
"""Out[115]: 
<xarray.DataArray (x: 2)>
array([1, 2])
Coordinates:
  * x        (x) <U1 'a' 'b'"""

b = xr.DataArray([-1, -2, -3], [("y", [10, 20, 30])])

b
"""Out[117]: 
<xarray.DataArray (y: 3)>
array([-1, -2, -3])
Coordinates:
  * y        (y) int64 10 20 30"""
a * b
"""Out[118]: 
<xarray.DataArray (x: 2, y: 3)>
array([[-1, -2, -3],
       [-2, -4, -6]])
Coordinates:
  * x        (x) <U1 'a' 'b'
  * y        (y) int64 10 20 30"""
c = xr.DataArray(np.arange(6).reshape(3, 2), [b["y"], a["x"]])

c
"""Out[120]: 
<xarray.DataArray (y: 3, x: 2)>
array([[0, 1],
       [2, 3],
       [4, 5]])
Coordinates:
  * y        (y) int64 10 20 30
  * x        (x) <U1 'a' 'b'"""

a + c
"""Out[121]: 
<xarray.DataArray (x: 2, y: 3)>
array([[1, 3, 5],
       [3, 5, 7]])
Coordinates:
  * x        (x) <U1 'a' 'b'
  * y        (y) int64 10 20 30"""
c+a
"""
<xarray.DataArray (y: 3, x: 2)>
array([[1, 3],
       [3, 5],
       [5, 7]])
Coordinates:
  * y        (y) int32 10 20 30
  * x        (x) <U1 'a' 'b'
  """
#这意味着,例如,您总是从数组的转置中减去数组:
c - c.T
"""Out[122]: 
<xarray.DataArray (y: 3, x: 2)>
array([[0, 0],
       [0, 0],
       [0, 0]])
Coordinates:
  * y        (y) int64 10 20 30
  * x        (x) <U1 'a' 'b'"""
#您可以使用broadcast()函数显式广播xarray数据结构:
a2, b2 = xr.broadcast(a, b)

a2
"""Out[124]: 
<xarray.DataArray (x: 2, y: 3)>
array([[1, 1, 1],
       [2, 2, 2]])
Coordinates:
  * x        (x) <U1 'a' 'b'
  * y        (y) int64 10 20 30"""

b2
"""Out[125]: 
<xarray.DataArray (x: 2, y: 3)>
array([[-1, -2, -3],
       [-1, -2, -3]])
Coordinates:
  * y        (y) int64 10 20 30
  * x        (x) <U1 'a' 'b'"""

Automatic alignment

Xarray在二元运算中使用的对象上强制索引坐标(即与维度同名的坐标,用*标记)之间对齐。
与pandas类似,对于二元运算的算术运算,这种对齐是自动的。二元运算的默认结果是坐标标签的交集(不是并集):

arr = xr.DataArray(np.arange(3), [("x", range(3))])

arr + arr[:-1]
"""Out[127]: 
<xarray.DataArray (x: 2)>
array([0, 2])
Coordinates:
  * x        (x) int64 0 1"""

如果某个尺寸的坐标值在任一参数中缺失,则所有匹配的尺寸必须具有相同的大小:
arr + xr.DataArray([1, 2], dims=“x”)
ValueError: arguments without labels along dimension ‘x’ cannot be aligned because they have different dimension size(s) {2} than the size of the aligned dimension labels: 3
但是,用户可以通过上下文管理器中的set_options()显式更改此默认自动对齐类型(“内部”):

with xr.set_options(arithmetic_join="outer"):
    arr + arr[:1]


arr + arr[:1]
"""Out[130]: 
<xarray.DataArray (x: 1)>
array([0])
Coordinates:
  * x        (x) int64 0"""

在循环或性能关键代码之前,最好显式对齐数组(例如,将数组放在同一个数据集中或使用align()),以避免每次操作重复对齐的开销。有关更多详细信息,请参见对齐和重新索引。

Coordinates

尽管索引坐标是对齐的,但其他坐标不是对齐的,如果它们的值冲突,它们将被删除。这是必要的,例如,因为索引将1D坐标转换为标量坐标:

arr[0]
"""Out[131]: 
<xarray.DataArray ()>
array(0)
Coordinates:
    x        int64 0"""

arr[1]
"""Out[132]: 
<xarray.DataArray ()>
array(1)
Coordinates:
    x        int64 1"""

# notice that the scalar coordinate 'x' is silently dropped
arr[1] - arr[0]
"""Out[133]: 
<xarray.DataArray ()>
array(1)"""

不过,只要没有冲突值,xarray将在算术中保存其他坐标:

# only one argument has the 'x' coordinate
arr[0] + 1
"""Out[134]: 
<xarray.DataArray ()>
array(1)
Coordinates:
    x        int64 0"""

# both arguments have the same 'x' coordinate
arr[0] - arr[0]
"""Out[135]: 
<xarray.DataArray ()>
array(0)
Coordinates:
    x        int64 0"""

Math with datasets

数据集通过自动循环所有数据变量来支持算术运算:

ds = xr.Dataset(
    {
        "x_and_y": (("x", "y"), np.random.randn(3, 5)),
        "x_only": ("x", np.random.randn(3)),
    },
    coords=arr.coords,
)


ds > 0
"""Out[137]: 
<xarray.Dataset>
Dimensions:  (x: 3, y: 5)
Coordinates:
  * x        (x) int64 0 1 2
Dimensions without coordinates: y
Data variables:
    x_and_y  (x, y) bool True True False True False ... True False False False
    x_only   (x) bool False True False"""

数据集支持数据数组中的大多数相同方法

ds.mean(dim="x")
"""Out[138]: 
<xarray.Dataset>
Dimensions:  (y: 5)
Dimensions without coordinates: y
Data variables:
    x_and_y  (y) float64 -0.1779 0.449 -0.6525 0.2515 0.09179
    x_only   float64 -0.371
"""
abs(ds)
"""Out[139]: 
<xarray.Dataset>
Dimensions:  (x: 3, y: 5)
Coordinates:
  * x        (x) int64 0 1 2
Dimensions without coordinates: y
Data variables:
    x_and_y  (x, y) float64 0.4281 0.9399 0.884 0.3229 ... 0.7523 0.1212 0.3989
    x_only   (x) float64 0.5093 0.2509 0.8548"""

数据集还支持NumPy ufuncs(需要NumPy v1.13或更高版本),或者您可以使用map()将函数映射到数据集中的每个变量:

np.sin(ds)
"""Out[140]: 
<xarray.Dataset>
Dimensions:  (x: 3, y: 5)
Coordinates:
  * x        (x) int64 0 1 2
Dimensions without coordinates: y
Data variables:
    x_and_y  (x, y) float64 0.4152 0.8075 -0.7733 ... -0.6833 -0.1209 -0.3884
    x_only   (x) float64 -0.4875 0.2483 -0.7544"""

ds.map(np.sin)
"""Out[141]: 
<xarray.Dataset>
Dimensions:  (x: 3, y: 5)
Coordinates:
  * x        (x) int64 0 1 2
Dimensions without coordinates: y
Data variables:
    x_and_y  (x, y) float64 0.4152 0.8075 -0.7733 ... -0.6833 -0.1209 -0.3884
    x_only   (x) float64 -0.4875 0.2483 -0.7544"""

数据集还在二进制算术中使用变量循环进行广播。您可以在任何数据数组和数据集之间进行运算:

ds + arr
"""Out[142]: 
<xarray.Dataset>
Dimensions:  (x: 3, y: 5)
Coordinates:
  * x        (x) int64 0 1 2
Dimensions without coordinates: y
Data variables:
    x_and_y  (x, y) float64 0.4281 0.9399 -0.884 0.3229 ... 1.248 1.879 1.601
    x_only   (x) float64 -0.5093 1.251 1.145"""

两个数据集之间的运算匹配同名的数据变量:

ds2 = xr.Dataset({"x_and_y": 0, "x_only": 100})

ds - ds2
"""Out[144]: 
<xarray.Dataset>
Dimensions:  (x: 3, y: 5)
Coordinates:
  * x        (x) int64 0 1 2
Dimensions without coordinates: y
Data variables:
    x_and_y  (x, y) float64 0.4281 0.9399 -0.884 ... -0.7523 -0.1212 -0.3989
    x_only   (x) float64 -100.5 -99.75 -100.9"""

与基于索引的对齐类似,结果是所有匹配数据变量的交集。

Wrapping custom computation

直接用xarray对象进行计算并不总是有意义的:

  • 在性能有限的代码的内循环中,与使用NumPy或原生Python类型相比,使用xarray会增加相当大的开销。当处理标量或小数组(少于1e6个元素)时尤其如此。跟踪标签并确保它们的一致性会增加开销,xarray的内核本身也不是特别快,因为它是用Python而不是c等编译语言编写的。此外,xarray基于标签的高级API消除了对操作实现方式的低级控制。
  • 即使速度不重要,包装现有代码或支持不使用xarray对象的替代接口也很重要。
    由于这些原因,通常最好编写使用NumPy数组的低级例程,并包装这些例程以使用xarray对象。然而,在数据集和数据数组上添加标签支持可能有点麻烦。
    为了使这变得更容易,xarray提供了apply _ ufunc()helper函数,该函数旨在以NumPy通用函数(简称“ufunc”)的风格包装支持对无标签数组进行广播和矢量化的函数。apply_ufunc负责惯用xarray包装器所需的一切,包括对齐、广播、数据集变量循环(如果需要)和坐标合并。事实上,许多内部xarray函数/方法都是使用apply_ufunc编写的。
    独立作用于每个值的简单函数应该无需任何附加参数即可工作:
squared_error = lambda x, y: (x - y) ** 2

arr1 = xr.DataArray([0, 1, 2, 3], dims="x")

xr.apply_ufunc(squared_error, arr1, 1)
"""Out[147]: 
<xarray.DataArray (x: 4)>
array([1, 0, 1, 4])
Dimensions without coordinates: x"""

为了使用更复杂的运算来综合考虑一些数组值,理解NumPy的广义ufuncs中的“核心维度”的概念非常重要。核心维度是指不应该广播的维度。通常,它们对应于定义操作的基本维度,例如np.sum中的求和轴。需要核心维度的一个很好的线索是对应NumPy函数中存在轴参数。

def vector_norm(x, dim, ord=None):
    return xr.apply_ufunc(
        np.linalg.norm, x, input_core_dims=[[dim]], kwargs={"ord": ord, "axis": -1}
    )
vector_norm(arr1, dim="x")
"""Out[148]: 
<xarray.DataArray ()>
array(3.74165739)"""

因为apply_ufunc遵循ufunc的标准惯例,所以它可以很好地与构建矢量化函数的工具配合使用,如numpy.broadcast_arrays()和numpy.vectorize .对于高性能需求,请考虑使用Numba的矢量化和guvectorize。
除了包装函数之外,通过设置dask =‘parallelised‘,apply_ufunc还可以在使用dask时自动并行化许多函数。有关详细信息,请参见使用apply_ufunc和map_blocks进行自动并行化。
apply_ufunc()还支持一些用于控制变量对齐和结果形式的高级选项。有关完整的详细信息和更多示例,请参见docstring。

  • 13
    点赞
  • 14
    收藏
    觉得还不错? 一键收藏
  • 1
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值