bitPhile

xr.apply_ufunc

2 min read

apply_ufunc

Xarray method that applies custom function on underlying array in xarray. The underlying array can be a numpy or a dask array.

Syntax

xr.apply_ufunc(custom_function, dataset/dataarray..., input_core_dims, output_core_dims, kwargs={}, keep_attrs=False, vectorize=False, dask="allowed|parallelized|forbidden")
  • custom_function is the function we want to apply on xarray.
  • followed by a list of DataArrays of Datasets.
  • input_core_dims and output_core_dims are effective when vectorize is True. input_core_dims provides dimensions that will be going to custom function. output_core_dims provides dimensions that custom function will return. Please refer.
  • vectorize enables numpy vectorize to be used when dask=parallelized.

Example

>>> data
<xarray.DataArray 'vis' (time: 2, baselineid: 1, frequency: 3, polarisation: 4)> Size: 192B
array([[[[0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
         [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
         [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j]]],
 
 
       [[[0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
         [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j],
         [0.+0.j, 0.+0.j, 0.+0.j, 0.+0.j]]]], dtype=complex64)
Coordinates:
  * time          (time) float64 16B 4.454e+09 4.454e+09
  * frequency     (frequency) float64 24B 1.281e+08 1.282e+08 1.282e+08
  * polarisation  (polarisation) <U2 32B 'XX' 'XY' 'YX' 'YY'
  * baselineid    (baselineid) int64 8B 0
Attributes:
    units:    Jy
 
>>> def adder(arr, value):
          print('------------------- START ---------------------')
          print(type(arr))
          print(arr.shape)
          print(arr)
          print('-------------------- END ------------------------')
          return arr + value

data is a numpy xarray.

>>> xr.apply_ufunc(adder, data, 1)
------------------- START ---------------------
<class 'numpy.ndarray'>
(2, 1, 3, 4)
[[[[0.+0.j 0.+0.j 0.+0.j 0.+0.j]
   [0.+0.j 0.+0.j 0.+0.j 0.+0.j]
   [0.+0.j 0.+0.j 0.+0.j 0.+0.j]]]
 
 
 [[[0.+0.j 0.+0.j 0.+0.j 0.+0.j]
   [0.+0.j 0.+0.j 0.+0.j 0.+0.j]
   [0.+0.j 0.+0.j 0.+0.j 0.+0.j]]]]
-------------------- END ------------------------
Out[51]: 
<xarray.DataArray 'vis' (time: 2, baselineid: 1, frequency: 3, polarisation: 4)> Size: 192B
array([[[[1.+0.j, 1.+0.j, 1.+0.j, 1.+0.j],
         [1.+0.j, 1.+0.j, 1.+0.j, 1.+0.j],
         [1.+0.j, 1.+0.j, 1.+0.j, 1.+0.j]]],
 
 
       [[[1.+0.j, 1.+0.j, 1.+0.j, 1.+0.j],
         [1.+0.j, 1.+0.j, 1.+0.j, 1.+0.j],
         [1.+0.j, 1.+0.j, 1.+0.j, 1.+0.j]]]], dtype=complex64)
Coordinates:
  * time          (time) float64 16B 4.454e+09 4.454e+09
  * frequency     (frequency) float64 24B 1.281e+08 1.282e+08 1.282e+08
  * polarisation  (polarisation) <U2 32B 'XX' 'XY' 'YX' 'YY'
  * baselineid    (baselineid) int64 8B 0

How it works?

  1. When xarray DataArray is given to apply_ufunc, it gets the underlying array, and passes to the custom function. Custom function returns the result array which again wrapped by apply_ufunc and returned to us.
  2. When xarray DataSet is given, it loops over the data variables(DataArray), passes each of the data array to the custom function, gets the result array and wraps in Data Array and then in Data set and returns.
  3. All the input DataSet/DataArray metadata gets populated on the result by apply_ufunc.
  4. By default, attributes are not populated on the result but can retained using keep_attrs=True.

References

  1. https://tutorial.xarray.dev/advanced/apply_ufunc/simple_numpy_apply_ufunc.html

Graph View

Backlinks