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.235948, -2.599843, -2.021262],
       [-0.759107, -1.132442, -3.977278]])
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.764052,  0.400157,  0.978738],
       [ 2.240893,  1.867558,  0.977278]])
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.981384,  0.389563,  0.829794],
       [ 0.783762,  0.956288, -0.828978]])
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

Data arrays also implement many numpy.ndarray methods:

In [6]: arr.round(2)
Out[6]: 
<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 [7]: arr.T
Out[7]: 
<xarray.DataArray (y: 3, x: 2)>
array([[ 1.764052,  2.240893],
       [ 0.400157,  1.867558],
       [ 0.978738, -0.977278]])
Coordinates:
  * x        (x) <U1 'a' 'b'
  * y        (y) int64 10 20 30

Missing values

xarray objects borrow the isnull(), notnull(), count(), dropna() and fillna() methods for working with missing data from pandas:

In [8]: x = xr.DataArray([0, 1, np.nan, np.nan, 2], dims=['x'])

In [9]: x.isnull()
Out[9]: 
<xarray.DataArray (x: 5)>
array([False, False,  True,  True, False], dtype=bool)
Dimensions without coordinates: x

In [10]: x.notnull()
Out[10]: 
<xarray.DataArray (x: 5)>
array([ True,  True, False, False,  True], dtype=bool)
Dimensions without coordinates: x

In [11]: x.count()
Out[11]: 
<xarray.DataArray ()>
array(3)

In [12]: x.dropna(dim='x')
Out[12]: 
<xarray.DataArray (x: 3)>
array([ 0.,  1.,  2.])
Dimensions without coordinates: x

In [13]: x.fillna(-1)
Out[13]: 
<xarray.DataArray (x: 5)>
array([ 0.,  1., -1., -1.,  2.])
Dimensions without coordinates: x

Like pandas, xarray uses the float value np.nan (not-a-number) to represent missing values.

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 [14]: arr.sum(dim='x')
Out[14]: 
<xarray.DataArray (y: 3)>
array([  4.004946e+00,   2.267715e+00,   1.460104e-03])
Coordinates:
  * y        (y) int64 10 20 30

In [15]: arr.std(['x', 'y'])
Out[15]: 
<xarray.DataArray ()>
array(1.0903834448772864)

In [16]: arr.min()
Out[16]: 
<xarray.DataArray ()>
array(-0.977277879876411)

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 [17]: arr.get_axis_num('y')
Out[17]: 1

These operations automatically skip missing values, like in pandas:

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

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 [19]: arr = xr.DataArray(np.arange(0, 7.5, 0.5).reshape(3, 5),
   ....:                    dims=('x', 'y'))
   ....: 

In [20]: arr
Out[20]: 
<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 [21]: arr.rolling(y=3)
Out[21]: DataArrayRolling [window->3,center->False,dim->y]

The label position and minimum number of periods in the rolling window are controlled by the center and min_periods arguments:

In [22]: arr.rolling(y=3, min_periods=2, center=True)
Out[22]: DataArrayRolling [window->3,min_periods->2,center->True,dim->y]

Aggregation and summary methods can be applied directly to the Rolling object:

In [23]: r = arr.rolling(y=3)

In [24]: r.mean()
Out[24]: 
<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

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

Note that rolling window aggregations are much faster (both asymptotically and because they avoid a loop in Python) when bottleneck is installed. Otherwise, we fall back to a slower, pure Python implementation.

Finally, we can manually iterate through Rolling objects:

In [26]: for label, arr_window in r:
   # arr_window is a view of x

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 [27]: a = xr.DataArray([1, 2], [('x', ['a', 'b'])])

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

In [29]: b = xr.DataArray([-1, -2, -3], [('y', [10, 20, 30])])

In [30]: b
Out[30]: 
<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 [31]: a * b
Out[31]: 
<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 [32]: c = xr.DataArray(np.arange(6).reshape(3, 2), [b['y'], a['x']])

In [33]: c
Out[33]: 
<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 [34]: a + c
Out[34]: 
<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 [35]: c - c.T
Out[35]: 
<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 xaray data structures by using the broadcast() function:

In [36]: a2, b2 = xr.broadcast(a, b)

In [37]: a2
Out[37]: 
<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 [38]: b2
Out[38]: 
<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 [39]: arr = xr.DataArray(np.arange(3), [('x', range(3))])

In [40]: arr + arr[:-1]
Out[40]: 
<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 [41]: In [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
   ....: 

However, one can explicitly change this default automatic alignment type (“inner”) via set_options() in context manager:

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

In [43]: arr + arr[:1]
Out[43]: 
<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 [44]: arr[0]
Out[44]: 
<xarray.DataArray ()>
array(0)
Coordinates:
    x        int64 0

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

# notice that the scalar coordinate 'x' is silently dropped
In [46]: arr[1] - arr[0]
Out[46]: 
<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 [47]: arr[0] + 1
Out[47]: 
<xarray.DataArray ()>
array(1)
Coordinates:
    x        int64 0

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

Math with datasets

Datasets support arithmetic operations by automatically looping over all data variables:

In [49]: ds = xr.Dataset({'x_and_y': (('x', 'y'), np.random.randn(3, 5)),
   ....:                  'x_only': ('x', np.random.randn(3))},
   ....:                  coords=arr.coords)
   ....: 

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

Datasets support most of the same methods found on data arrays:

In [51]: ds.mean(dim='x')
Out[51]: 
<xarray.Dataset>
Dimensions:  (y: 5)
Dimensions without coordinates: y
Data variables:
    x_only   float64 0.138
    x_and_y  (y) float64 -0.06634 0.3027 -0.6106 -0.9014 -0.644

In [52]: abs(ds)
Out[52]: 
<xarray.Dataset>
Dimensions:  (x: 3)
Coordinates:
  * x        (x) int64 0 1 2
Data variables:
    x_only   (x) float64 0.2719 0.425 0.567
    x_and_y  (x, y) float64 0.4691 0.2829 1.509 1.136 1.212 0.1732 0.1192 ...

Unfortunately, we currently do not support NumPy ufuncs for datasets [1]. apply() works around this limitation, by applying the given function to each variable in the dataset:

In [53]: ds.apply(np.sin)
Out[53]: 
<xarray.Dataset>
Dimensions:  (x: 3, y: 5)
Coordinates:
  * x        (x) int64 0 1 2
Dimensions without coordinates: y
Data variables:
    x_only   (x) float64 0.2685 -0.4123 0.5371
    x_and_y  (x, y) float64 0.4521 -0.2791 -0.9981 -0.9068 0.9364 -0.1723 ...

You can also use the wrapped functions in the xarray.ufuncs module:

In [54]: import xarray.ufuncs as xu

In [55]: xu.sin(ds)
Out[55]: 
<xarray.Dataset>
Dimensions:  (x: 3)
Coordinates:
  * x        (x) int64 0 1 2
Data variables:
    x_only   (x) float64 0.2685 -0.4123 0.5371
    x_and_y  (x, y) float64 0.4521 -0.2791 -0.9981 -0.9068 0.9364 -0.1723 ...

Datasets also use looping over variables for broadcasting in binary arithmetic. You can do arithmetic between any DataArray and a dataset:

In [56]: ds + arr
Out[56]: 
<xarray.Dataset>
Dimensions:  (x: 3, y: 5)
Coordinates:
  * x        (x) int64 0 1 2
Dimensions without coordinates: y
Data variables:
    x_only   (x) float64 0.2719 0.575 2.567
    x_and_y  (x, y) float64 0.4691 -0.2829 -1.509 -1.136 1.212 0.8268 1.119 ...

Arithmetic between two datasets matches data variables of the same name:

In [57]: ds2 = xr.Dataset({'x_and_y': 0, 'x_only': 100})

In [58]: ds - ds2
Out[58]: 
<xarray.Dataset>
Dimensions:  (x: 3, y: 5)
Coordinates:
  * x        (x) int64 0 1 2
Dimensions without coordinates: y
Data variables:
    x_only   (x) float64 -99.73 -100.4 -99.43
    x_and_y  (x, y) float64 0.4691 -0.2829 -1.509 -1.136 1.212 -0.1732 ...

Similarly to index based alignment, the result has the intersection of all matching data variables.

[1]This was previously due to a limitation of NumPy, but with NumPy 1.13 we should be able to support this by leveraging __array_ufunc__ (GH1617).

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 [59]: squared_error = lambda x, y: (x - y) ** 2

In [60]: arr1 = xr.DataArray([0, 1, 2, 3], dims='x')

In [61]: xr.apply_ufunc(squared_error, arr1, 1)
Out[61]: 
<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 [62]: vector_norm(arr1, dim='x')
Out[62]: 
<xarray.DataArray ()>
array(3.7416573867739413)

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.