make_broadcastable#

movement.utils.broadcasting.make_broadcastable(is_classmethod=False, only_broadcastable_along=None, new_dimension_name=None)[source]#

Create a decorator that allows a function to be broadcast.

Parameters:

  • is_classmethod (bool) – Whether the target of the decoration is a class method which takes the self argument, or a standalone function that receives no implicit arguments.

  • only_broadcastable_along (str, optional) – Whether the decorated function should only support broadcasting along this dimension. The returned function will not take the broadcast_dimension argument, and will use the dimension provided here as the value for this argument.

  • new_dimension_name (str, optional) – Passed to apply_along_da_axis().

Returns:

Decorator function that can be applied with the @make_broadcastable(...) syntax. See Notes for a description of the action of the returned decorator.

Return type:

Decorator

Notes

The returned decorator (the “r_decorator”) extends a function that acts on a 1D sequence of values, allowing it to be broadcast along the axes of an input xarray.DataArray.

The r_decorator takes a single parameter, f. f should be a Callable that acts on 1D inputs, that is to be converted into a broadcast-able function fr, applying the action of f along an axis of an xarray.DataArray. f should return either scalar or 1D outputs.

If f is a class method, it should be callable as f(self, [x, y, ...], *args, **kwargs). Otherwise, f should be callable as f([x, y, ...], *args, **kwargs).

The function fr returned by the r_decorator is callable with the signature fr([self,] data, *args, broadcast_dimension = str, **kwargs), where the self argument is present only if f was a class method. fr applies f along the broadcast_dimension of data. The *args and **kwargs match those passed to f, and retain the same interpretations and effects on the result. If data provided to fr is not an xarray.DataArray, it will fall back on the behaviour of f (and ignore the broadcast_dimension argument).

See the docstring of make_broadcastable_inner in the source code for a more explicit explanation of the returned decorator.

See also

broadcastable_method

Convenience alias for is_classmethod = True.

space_broadcastable

Convenience alias for only_broadcastable_along = "space".

Examples

Make a standalone function broadcast along the "space" axis of an xarray.DataArray.

>>> @make_broadcastable(is_classmethod=False, only_broadcast_along="space")
... def my_function(xyz_data, *args, **kwargs)
...
... # Call via the usual arguments, replacing the xyz_data argument with
... # the xarray.DataArray to broadcast over
... my_function(data_array, *args, **kwargs)
```

Make a class method broadcast along any axis of an xarray.DataArray.

>>> from dataclasses import dataclass
>>>
>>> @dataclass
... class MyClass:
...     factor: float
...     offset: float
...
...     @make_broadcastable(is_classmethod=True)
...     def manipulate_values(self, xyz_values, *args, **kwargs):
...         return self.factor * sum(xyz_values) + self.offset
...
>>> m = MyClass(factor=5.9, offset=1.0)
>>> m.manipulate_values(
...     data_array, *args, broadcast_dimension="time", **kwargs
... )
```

Examples using make_broadcastable#

Broadcast functions across multi-dimensional data

Broadcast functions across multi-dimensional data