Computation¶
The labels associated with DataArray
and
Dataset
objects enables some powerful shortcuts for
computation, notably including aggregation and broadcasting by dimension
names.
Basic array math¶
Arithmetic operations with a single DataArray automatically vectorize (like numpy) over all array values:
In [1]: arr = xr.DataArray(np.random.RandomState(0).randn(2, 3), ...: [('x', ['a', 'b']), ('y', [10, 20, 30])]) ...: In [2]: 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 In [3]: 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
You can also use any of numpy’s or scipy’s many ufunc functions directly on a DataArray:
In [4]: 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
Use where()
to conditionally switch between values:
In [5]: 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
Use @ to perform matrix multiplication:
In [6]: arr @ arr
Out[6]:
<xarray.DataArray ()>
array(13.69438174)
Data arrays also implement many numpy.ndarray
methods:
In [7]: 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 In [8]: 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
Missing values¶
xarray objects borrow the isnull()
,
notnull()
, count()
,
dropna()
, fillna()
,
ffill()
, and bfill()
methods for working with missing data from pandas:
In [9]: x = xr.DataArray([0, 1, np.nan, np.nan, 2], dims=['x']) In [10]: x.isnull() Out[10]: <xarray.DataArray (x: 5)> array([False, False, True, True, False]) Dimensions without coordinates: x In [11]: x.notnull() Out[11]: <xarray.DataArray (x: 5)> array([ True, True, False, False, True]) Dimensions without coordinates: x In [12]: x.count() Out[12]: <xarray.DataArray ()> array(3) In [13]: x.dropna(dim='x') Out[13]: <xarray.DataArray (x: 3)> array([0., 1., 2.]) Dimensions without coordinates: x In [14]: x.fillna(-1) Out[14]: <xarray.DataArray (x: 5)> array([ 0., 1., -1., -1., 2.]) Dimensions without coordinates: x In [15]: x.ffill('x') Out[15]: <xarray.DataArray (x: 5)> array([0., 1., 1., 1., 2.]) Dimensions without coordinates: x In [16]: x.bfill('x') Out[16]: <xarray.DataArray (x: 5)> array([0., 1., 2., 2., 2.]) Dimensions without coordinates: x
Like pandas, xarray uses the float value np.nan
(not-a-number) to represent
missing values.
xarray objects also have an interpolate_na()
method
for filling missing values via 1D interpolation.
In [17]: x = xr.DataArray([0, 1, np.nan, np.nan, 2], dims=['x'],
....: coords={'xx': xr.Variable('x', [0, 1, 1.1, 1.9, 3])})
....:
In [18]: x.interpolate_na(dim='x', method='linear', use_coordinate='xx')
Out[18]:
<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
Note that xarray slightly diverges from the pandas interpolate
syntax by
providing the use_coordinate
keyword which facilitates a clear specification
of which values to use as the index in the interpolation.
Aggregation¶
Aggregation methods have been updated to take a dim argument instead of axis. This allows for very intuitive syntax for aggregation methods that are applied along particular dimension(s):
In [19]: arr.sum(dim='x') Out[19]: <xarray.DataArray (y: 3)> array([4.00494555e+00, 2.26771520e+00, 1.46010423e-03]) Coordinates: * y (y) int64 10 20 30 In [20]: arr.std(['x', 'y']) Out[20]: <xarray.DataArray ()> array(1.09038344) In [21]: arr.min() Out[21]: <xarray.DataArray ()> array(-0.97727788)
If you need to figure out the axis number for a dimension yourself (say,
for wrapping code designed to work with numpy arrays), you can use the
get_axis_num()
method:
In [22]: arr.get_axis_num('y')
Out[22]: 1
These operations automatically skip missing values, like in pandas:
In [23]: xr.DataArray([1, 2, np.nan, 3]).mean()
Out[23]:
<xarray.DataArray ()>
array(2.)
If desired, you can disable this behavior by invoking the aggregation method
with skipna=False
.
Rolling window operations¶
DataArray
objects include a rolling()
method. This
method supports rolling window aggregation:
In [24]: arr = xr.DataArray(np.arange(0, 7.5, 0.5).reshape(3, 5),
....: dims=('x', 'y'))
....:
In [25]: arr
Out[25]:
<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()
is applied along one dimension using the
name of the dimension as a key (e.g. y
) and the window size as the value
(e.g. 3
). We get back a Rolling
object:
In [26]: arr.rolling(y=3)
Out[26]: DataArrayRolling [window->3,center->False,dim->y]
Aggregation and summary methods can be applied directly to the Rolling
object:
In [27]: r = arr.rolling(y=3) In [28]: r.reduce(np.std) Out[28]: <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 In [29]: r.mean() Out[29]: <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
Aggregation results are assigned the coordinate at the end of each window by
default, but can be centered by passing center=True
when constructing the
Rolling
object:
In [30]: r = arr.rolling(y=3, center=True)
In [31]: r.mean()
Out[31]:
<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
As can be seen above, aggregations of windows which overlap the border of the
array produce nan
s. Setting min_periods
in the call to rolling
changes the minimum number of observations within the window required to have
a value when aggregating:
In [32]: r = arr.rolling(y=3, min_periods=2)
In [33]: r.mean()
Out[33]:
<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
In [34]: r = arr.rolling(y=3, center=True, min_periods=2)
In [35]: r.mean()
Out[35]:
<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
Tip
Note that rolling window aggregations are faster and use less memory when bottleneck is installed. This only applies to numpy-backed xarray objects.
We can also manually iterate through Rolling
objects:
for label, arr_window in r:
# arr_window is a view of x
While rolling
provides a simple moving average, DataArray
also supports
an exponential moving average with rolling_exp()
.
This is similiar to pandas’ ewm
method. numbagg is required.
arr.rolling_exp(y=3).mean()
The rolling_exp
method takes a window_type
kwarg, which can be 'alpha'
,
'com'
(for center-of-mass
), 'span'
, and 'halflife'
. The default is
span
.
Finally, the rolling object has a construct
method which returns a
view of the original DataArray
with the windowed dimension in
the last position.
You can use this for more advanced rolling operations such as strided rolling,
windowed rolling, convolution, short-time FFT etc.
# rolling with 2-point stride In [36]: rolling_da = r.construct('window_dim', stride=2) In [37]: rolling_da Out[37]: <xarray.DataArray (x: 3, y: 3, window_dim: 3)> array([[[nan, 0. , 0.5], [0.5, 1. , 1.5], [1.5, 2. , nan]], [[nan, 2.5, 3. ], [3. , 3.5, 4. ], [4. , 4.5, nan]], [[nan, 5. , 5.5], [5.5, 6. , 6.5], [6.5, 7. , nan]]]) Dimensions without coordinates: x, y, window_dim In [38]: rolling_da.mean('window_dim', skipna=False) Out[38]: <xarray.DataArray (x: 3, y: 3)> array([[nan, 1. , nan], [nan, 3.5, nan], [nan, 6. , nan]]) Dimensions without coordinates: x, y
Because the DataArray
given by r.construct('window_dim')
is a view
of the original array, it is memory efficient.
You can also use construct
to compute a weighted rolling sum:
In [39]: weight = xr.DataArray([0.25, 0.5, 0.25], dims=['window'])
In [40]: arr.rolling(y=3).construct('window').dot(weight)
Out[40]:
<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
Note
numpy’s Nan-aggregation functions such as nansum
copy the original array.
In xarray, we internally use these functions in our aggregation methods
(such as .sum()
) if skipna
argument is not specified or set to True.
This means rolling_da.mean('window_dim')
is memory inefficient.
To avoid this, use skipna=False
as the above example.
Coarsen large arrays¶
DataArray
and Dataset
objects include a
coarsen()
and coarsen()
methods. This supports the block aggregation along multiple dimensions,
In [41]: x = np.linspace(0, 10, 300)
In [42]: t = pd.date_range('15/12/1999', periods=364)
In [43]: da = xr.DataArray(np.sin(x) * np.cos(np.linspace(0, 1, 364)[:, np.newaxis]),
....: dims=['time', 'x'], coords={'time': t, 'x': x})
....:
In [44]: da
Out[44]:
<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
In order to take a block mean for every 7 days along time
dimension and
every 2 points along x
dimension,
In [45]: da.coarsen(time=7, x=2).mean()
Out[45]:
<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
coarsen()
raises an ValueError
if the data
length is not a multiple of the corresponding window size.
You can choose boundary='trim'
or boundary='pad'
options for trimming
the excess entries or padding nan
to insufficient entries,
In [46]: da.coarsen(time=30, x=2, boundary='trim').mean()
Out[46]:
<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
If you want to apply a specific function to coordinate, you can pass the
function or method name to coord_func
option,
In [47]: da.coarsen(time=7, x=2, coord_func={'time': 'min'}).mean()
Out[47]:
<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 objects have some handy methods for the computation with their
coordinates. differentiate()
computes derivatives by
central finite differences using their coordinates,
In [48]: a = xr.DataArray([0, 1, 2, 3], dims=['x'], coords=[[0.1, 0.11, 0.2, 0.3]]) In [49]: a Out[49]: <xarray.DataArray (x: 4)> array([0, 1, 2, 3]) Coordinates: * x (x) float64 0.1 0.11 0.2 0.3 In [50]: a.differentiate('x') Out[50]: <xarray.DataArray (x: 4)> array([100. , 91.11111111, 10.58479532, 10. ]) Coordinates: * x (x) float64 0.1 0.11 0.2 0.3
This method can be used also for multidimensional arrays,
In [51]: a = xr.DataArray(np.arange(8).reshape(4, 2), dims=['x', 'y'],
....: coords={'x': [0.1, 0.11, 0.2, 0.3]})
....:
In [52]: a.differentiate('x')
Out[52]:
<xarray.DataArray (x: 4, y: 2)>
array([[200. , 200. ],
[182.22222222, 182.22222222],
[ 21.16959064, 21.16959064],
[ 20. , 20. ]])
Coordinates:
* x (x) float64 0.1 0.11 0.2 0.3
Dimensions without coordinates: y
integrate()
computes integration based on
trapezoidal rule using their coordinates,
In [53]: a.integrate('x')
Out[53]:
<xarray.DataArray (y: 2)>
array([0.78, 0.98])
Dimensions without coordinates: y
Note
These methods are limited to simple cartesian geometry. Differentiation and integration along multidimensional coordinate are not supported.
Broadcasting by dimension name¶
DataArray
objects are automatically align themselves (“broadcasting” in
the numpy parlance) by dimension name instead of axis order. With xarray, you
do not need to transpose arrays or insert dimensions of length 1 to get array
operations to work, as commonly done in numpy with np.reshape()
or
np.newaxis
.
This is best illustrated by a few examples. Consider two one-dimensional arrays with different sizes aligned along different dimensions:
In [54]: a = xr.DataArray([1, 2], [('x', ['a', 'b'])])
In [55]: a
Out[55]:
<xarray.DataArray (x: 2)>
array([1, 2])
Coordinates:
* x (x) <U1 'a' 'b'
In [56]: b = xr.DataArray([-1, -2, -3], [('y', [10, 20, 30])])
In [57]: b
Out[57]:
<xarray.DataArray (y: 3)>
array([-1, -2, -3])
Coordinates:
* y (y) int64 10 20 30
With xarray, we can apply binary mathematical operations to these arrays, and their dimensions are expanded automatically:
In [58]: a * b
Out[58]:
<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
Moreover, dimensions are always reordered to the order in which they first appeared:
In [59]: c = xr.DataArray(np.arange(6).reshape(3, 2), [b['y'], a['x']]) In [60]: c Out[60]: <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' In [61]: a + c Out[61]: <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
This means, for example, that you always subtract an array from its transpose:
In [62]: c - c.T
Out[62]:
<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'
You can explicitly broadcast xarray data structures by using the
broadcast()
function:
In [63]: a2, b2 = xr.broadcast(a, b) In [64]: a2 Out[64]: <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 In [65]: b2 Out[65]: <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 enforces alignment between index Coordinates (that is,
coordinates with the same name as a dimension, marked by *
) on objects used
in binary operations.
Similarly to pandas, this alignment is automatic for arithmetic on binary operations. The default result of a binary operation is by the intersection (not the union) of coordinate labels:
In [66]: arr = xr.DataArray(np.arange(3), [('x', range(3))])
In [67]: arr + arr[:-1]
Out[67]:
<xarray.DataArray (x: 2)>
array([0, 2])
Coordinates:
* x (x) int64 0 1
If coordinate values for a dimension are missing on either argument, all matching dimensions must have the same size:
In [68]: 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
However, one can explicitly change this default automatic alignment type (“inner”)
via set_options()
in context manager:
In [69]: with xr.set_options(arithmetic_join="outer"):
....: arr + arr[:1]
....:
In [70]: arr + arr[:1]
Out[70]:
<xarray.DataArray (x: 1)>
array([0])
Coordinates:
* x (x) int64 0
Before loops or performance critical code, it’s a good idea to align arrays
explicitly (e.g., by putting them in the same Dataset or using
align()
) to avoid the overhead of repeated alignment with each
operation. See Align and reindex for more details.
Note
There is no automatic alignment between arguments when performing in-place
arithmetic operations such as +=
. You will need to use
manual alignment. This ensures in-place
arithmetic never needs to modify data types.
Coordinates¶
Although index coordinates are aligned, other coordinates are not, and if their values conflict, they will be dropped. This is necessary, for example, because indexing turns 1D coordinates into scalar coordinates:
In [71]: arr[0] Out[71]: <xarray.DataArray ()> array(0) Coordinates: x int64 0 In [72]: arr[1] Out[72]: <xarray.DataArray ()> array(1) Coordinates: x int64 1 # notice that the scalar coordinate 'x' is silently dropped In [73]: arr[1] - arr[0] Out[73]: <xarray.DataArray ()> array(1)
Still, xarray will persist other coordinates in arithmetic, as long as there are no conflicting values:
# only one argument has the 'x' coordinate In [74]: arr[0] + 1 Out[74]: <xarray.DataArray ()> array(1) Coordinates: x int64 0 # both arguments have the same 'x' coordinate In [75]: arr[0] - arr[0] Out[75]: <xarray.DataArray ()> array(0) Coordinates: x int64 0
Math with datasets¶
Datasets support arithmetic operations by automatically looping over all data variables:
In [76]: ds = xr.Dataset({'x_and_y': (('x', 'y'), np.random.randn(3, 5)),
....: 'x_only': ('x', np.random.randn(3))},
....: coords=arr.coords)
....:
In [77]: ds > 0
Out[77]:
<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 False False False True ... True True False False
x_only (x) bool True False True
Datasets support most of the same methods found on data arrays:
In [78]: ds.mean(dim='x') Out[78]: <xarray.Dataset> Dimensions: (y: 5) Dimensions without coordinates: y Data variables: x_and_y (y) float64 -0.06634 0.3027 -0.6106 -0.9014 -0.644 x_only float64 0.138 In [79]: abs(ds) Out[79]: <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.4691 0.2829 1.509 1.136 ... 0.7216 0.7068 1.04 x_only (x) float64 0.2719 0.425 0.567
Datasets also support NumPy ufuncs (requires NumPy v1.13 or newer), or
alternatively you can use apply()
to apply a function
to each variable in a dataset:
In [80]: np.sin(ds) Out[80]: <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.4521 -0.2791 -0.9981 ... 0.6606 -0.6494 -0.8622 x_only (x) float64 0.2685 -0.4123 0.5371 In [81]: ds.apply(np.sin) Out[81]: <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.4521 -0.2791 -0.9981 ... 0.6606 -0.6494 -0.8622 x_only (x) float64 0.2685 -0.4123 0.5371
Datasets also use looping over variables for broadcasting in binary
arithmetic. You can do arithmetic between any DataArray
and a dataset:
In [82]: ds + arr
Out[82]:
<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.4691 -0.2829 -1.509 -1.136 ... 2.722 1.293 0.9604
x_only (x) float64 0.2719 0.575 2.567
Arithmetic between two datasets matches data variables of the same name:
In [83]: ds2 = xr.Dataset({'x_and_y': 0, 'x_only': 100})
In [84]: ds - ds2
Out[84]:
<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.4691 -0.2829 -1.509 ... 0.7216 -0.7068 -1.04
x_only (x) float64 -99.73 -100.4 -99.43
Similarly to index based alignment, the result has the intersection of all matching data variables.
Wrapping custom computation¶
It doesn’t always make sense to do computation directly with xarray objects:
In the inner loop of performance limited code, using xarray can add considerable overhead compared to using NumPy or native Python types. This is particularly true when working with scalars or small arrays (less than ~1e6 elements). Keeping track of labels and ensuring their consistency adds overhead, and xarray’s core itself is not especially fast, because it’s written in Python rather than a compiled language like C. Also, xarray’s high level label-based APIs removes low-level control over how operations are implemented.
Even if speed doesn’t matter, it can be important to wrap existing code, or to support alternative interfaces that don’t use xarray objects.
For these reasons, it is often well-advised to write low-level routines that
work with NumPy arrays, and to wrap these routines to work with xarray objects.
However, adding support for labels on both Dataset
and
DataArray
can be a bit of a chore.
To make this easier, xarray supplies the apply_ufunc()
helper
function, designed for wrapping functions that support broadcasting and
vectorization on unlabeled arrays in the style of a NumPy
universal function (“ufunc” for short).
apply_ufunc
takes care of everything needed for an idiomatic xarray wrapper,
including alignment, broadcasting, looping over Dataset
variables (if
needed), and merging of coordinates. In fact, many internal xarray
functions/methods are written using apply_ufunc
.
Simple functions that act independently on each value should work without any additional arguments:
In [85]: squared_error = lambda x, y: (x - y) ** 2
In [86]: arr1 = xr.DataArray([0, 1, 2, 3], dims='x')
In [87]: xr.apply_ufunc(squared_error, arr1, 1)
Out[87]:
<xarray.DataArray (x: 4)>
array([1, 0, 1, 4])
Dimensions without coordinates: x
For using more complex operations that consider some array values collectively,
it’s important to understand the idea of “core dimensions” from NumPy’s
generalized ufuncs. Core dimensions are defined as dimensions
that should not be broadcast over. Usually, they correspond to the fundamental
dimensions over which an operation is defined, e.g., the summed axis in
np.sum
. A good clue that core dimensions are needed is the presence of an
axis
argument on the corresponding NumPy function.
With apply_ufunc
, core dimensions are recognized by name, and then moved to
the last dimension of any input arguments before applying the given function.
This means that for functions that accept an axis
argument, you usually need
to set axis=-1
. As an example, here is how we would wrap
numpy.linalg.norm()
to calculate the vector norm:
def vector_norm(x, dim, ord=None):
return xr.apply_ufunc(np.linalg.norm, x,
input_core_dims=[[dim]],
kwargs={'ord': ord, 'axis': -1})
In [88]: vector_norm(arr1, dim='x')
Out[88]:
<xarray.DataArray ()>
array(3.74165739)
Because apply_ufunc
follows a standard convention for ufuncs, it plays
nicely with tools for building vectorized functions, like
numpy.broadcast_arrays()
and numpy.vectorize()
. For high performance
needs, consider using Numba’s vectorize and guvectorize.
In addition to wrapping functions, apply_ufunc
can automatically parallelize
many functions when using dask by setting dask='parallelized'
. See
Automatic parallelization for details.
apply_ufunc()
also supports some advanced options for
controlling alignment of variables and the form of the result. See the
docstring for full details and more examples.