AxisArray module#

class ezmsg.util.messages.axisarray.ArrayWithNamedDims(data, dims)[source]#

Bases: object

Base class for arrays with named dimensions.

This class provides a foundation for arrays where each dimension has a name, enabling more intuitive array manipulation and access patterns.

Parameters:

  • data (npt.NDArray) – The underlying numpy array data

  • dims (list[str]) – List of dimension names, must match the array’s number of dimensions

Raises:

ValueError – If dims length doesn’t match data.ndim or contains duplicate names

class ezmsg.util.messages.axisarray.AxisArray(data, dims, axes=<factory>, attrs=<factory>, key='')[source]#

Bases: ArrayWithNamedDims

A lightweight message class comprising a numpy ndarray and its metadata.

AxisArray extends ArrayWithNamedDims to provide a complete data structure for scientific computing with named dimensions, axis coordinate systems, and metadata. It’s designed to be similar to xarray.DataArray but optimized for message passing in streaming applications.

Parameters:

  • data (npt.NDArray) – The underlying numpy array data (inherited from ArrayWithNamedDims)

  • dims (list[str]) – List of dimension names (inherited from ArrayWithNamedDims)

  • axes (dict[str, AxisBase]) – Dictionary mapping dimension names to their axis coordinate systems

  • attrs (dict[str, Any]) – Dictionary of arbitrary metadata attributes

  • key (str) – Optional key identifier for this array, typically used to specify source device (default is empty string)

class Axis(unit='', gain=1.0, offset=0.0)[source]#

Bases: LinearAxis

Deprecated alias for LinearAxis.

Deprecated since version 3.6.0: Use LinearAxis instead.

index(val)[source]#

Convert axis value(s) to index(es) using the inverse linear transformation.

Parameters:

  • v (float | npt.NDArray[np.float64]) – Axis value or array of values to convert to indices

  • fn (Callable) – Function to apply for rounding (default: np.rint)

Returns:

Corresponding index or array of indices

Return type:

int | npt.NDArray[numpy.int_]

class AxisInfo(axis, idx, size=None)[source]#

Bases: object

Container for axis information including the axis object, index, and size.

This class provides a convenient way to access both the axis coordinate system and metadata about where that axis appears in the array structure.

Parameters:

  • axis (AxisBase) – The axis coordinate system

  • idx (int) – The index of this axis in the AxisArray object’s dimension list

  • size (int | None) – The size of this dimension (stored as None for CoordinateAxis which determines size from data)

Raises:

ValueError – If size rules are violated for the axis type

property indices: ndarray[tuple[Any, ...], dtype[int64]][source]#

Get array of all valid indices for this axis.

Returns:

Array of indices from 0 to len(self)-1

Return type:

npt.NDArray[numpy.int_]

property values: ndarray[tuple[Any, ...], dtype[_ScalarT]][source]#

Get array of coordinate values for all indices of this axis.

Returns:

Array of coordinate values computed from axis.value(indices)

Return type:

npt.NDArray

class CoordinateAxis(data, dims, unit='')[source]#

Bases: AxisBase, ArrayWithNamedDims

An axis implementation that uses explicit coordinate values stored in an array.

This class allows for non-linear or irregularly spaced coordinate systems by storing the actual coordinate values in a data array.

Inherits from both AxisBase and ArrayWithNamedDims, combining axis functionality with named dimension support.

Parameters:

  • unit (str) – The unit of measurement for this axis (inherited from AxisBase)

  • data (npt.NDArray) – The underlying numpy array data (inherited from ArrayWithNamedDims)

  • dims (list[str]) – List of dimension names, must match the array’s number of dimensions (inherited from ArrayWithNamedDims)

value(x)[source]#

Get coordinate value(s) at the given index(es).

Parameters:

x (int | npt.NDArray[numpy.int_]) – Index or array of indices to lookup

Returns:

Coordinate value(s) at the specified index(es)

Return type:

Any | npt.NDArray

class LinearAxis(unit='', gain=1.0, offset=0.0)[source]#

Bases: AxisBase

An axis implementation for sparse axes with regular intervals between elements.

It is called “linear” because it provides a simple linear mapping between indices and element values: value = (index * gain) + offset.

A typical example is a time axis (TimeAxis), with regular sampling rate.

Parameters:

  • gain (float) – Step size (scaling factor) for the linear axis

  • offset (float) – The offset (value) of the first sample

  • unit (str) – The unit of measurement for this axis (inherited from AxisBase)

classmethod create_time_axis(fs, offset=0.0)[source]#

Convenience method to construct a LinearAxis for time.

Parameters:

  • fs (float) – Sampling frequency in Hz

  • offset (float) – Time offset in seconds (default: 0.0)

Returns:

A LinearAxis configured for time with gain=1/fs and unit=’s’

Return type:

LinearAxis

index(v, fn=<ufunc 'rint'>)[source]#

Convert axis value(s) to index(es) using the inverse linear transformation.

Parameters:

  • v (float | npt.NDArray[np.float64]) – Axis value or array of values to convert to indices

  • fn (Callable) – Function to apply for rounding (default: np.rint)

Returns:

Corresponding index or array of indices

Return type:

int | npt.NDArray[numpy.int_]

value(x)[source]#

Convert index(es) to axis value(s) using the linear transformation.

Parameters:

x (int | npt.NDArray[numpy.int_]) – Index or array of indices to convert

Returns:

Corresponding axis value(s)

Return type:

float | npt.NDArray[np.float64]

classmethod TimeAxis(fs, offset=0.0)[source]#

Convenience method to construct a LinearAxis for time.

Parameters:

  • fs (float) – Sampling frequency in Hz

  • offset (float) – Time offset in seconds (default: 0.0)

Returns:

A LinearAxis configured for time with gain=1/fs and unit=’s’

Return type:

LinearAxis

as2d(dim)[source]#

Get a 2D view of the data with the specified dimension as the first axis.

Parameters:

dim (str | int) – Dimension name or index to move to first axis

Returns:

2D array view with shape (dim_size, remaining_elements)

Return type:

npt.NDArray

ax(dim)[source]#

Get AxisInfo for a specified dimension.

Parameters:

dim (str | int) – Dimension name or index

Returns:

AxisInfo containing axis, index, and size information

Return type:

AxisInfo

axis_idx(dim)[source]#

Get the axis index for a given dimension name or pass through if already an int.

Parameters:

dim (str | int) – Dimension name or index

Returns:

The axis index

Return type:

int

static concatenate(*aas, dim, axis=None, filter_key=None)[source]#

Concatenate multiple AxisArray objects along a specified dimension.

Parameters:

  • aas (T) – Variable number of AxisArray objects to concatenate

  • dim (str) – Dimension name along which to concatenate

  • axis (AxisBase | None) – Optional axis coordinate system for the concatenated dimension

  • filter_key (str | None) – Optional key filter - only concatenate arrays with matching key

Returns:

New AxisArray with concatenated data

Return type:

T

Raises:

ValueError – If objects have incompatible dimensions

get_axis(dim)[source]#

Get the axis coordinate system for a specified dimension.

Parameters:

dim (str | int) – Dimension name or index

Returns:

The axis coordinate system (defaults to LinearAxis if not specified)

Return type:

AxisBase

Raises:

ValueError – If dimension is not present in the object

get_axis_idx(dim)[source]#

Get the axis index for a given dimension name.

Parameters:

dim (str) – The dimension name

Returns:

The axis index

Return type:

int

Raises:

ValueError – If dimension is not present in the object

get_axis_name(dim)[source]#

Get the dimension name for a given axis index.

Parameters:

dim (int) – The axis index

Returns:

The dimension name

Return type:

str

isel(indexers=None, **indexers_kwargs)[source]#

Select data using integer-based indexing along specified dimensions.

This method allows for flexible indexing using integers, slices, or arrays of integers to select subsets of the data along named dimensions.

Parameters:

  • indexers (Any | None) – Dictionary of {dimension_name: indexer} pairs

  • indexers_kwargs (Any) – Alternative way to specify indexers as keyword arguments

Returns:

New AxisArray with selected data

Return type:

T

Raises:

ValueError – If indexer types are not supported (arrays, ints, or slices only)

iter_over_axis(axis)[source]#

Iterate over slices along the specified axis.

Yields AxisArray objects for each slice along the given axis, with that dimension removed from the resulting arrays.

Parameters:

axis (str | int) – Dimension name or index to iterate over

Yields:

AxisArray objects for each slice along the axis

Return type:

Generator[T, None, None]

sel(indexers=None, **indexers_kwargs)[source]#

Select data using label-based indexing along specified dimensions.

This method allows selection using real-world coordinate values rather than integer indices. Currently supports only slice objects and LinearAxis.

Parameters:

  • indexers (Any | None) – Dictionary of {dimension_name: slice_indexer} pairs

  • indexers_kwargs (Any) – Alternative way to specify indexers as keyword arguments

Returns:

New AxisArray with selected data

Return type:

T

Raises:

ValueError – If indexer is not a slice or axis is not a LinearAxis

property shape: tuple[int, ...][source]#

Shape of data.

Returns:

Tuple representing the shape of the underlying data array

Return type:

tuple[int, …]

shape2d(dim)[source]#

Get the 2D shape when viewing data with specified dimension first.

Parameters:

dim (str | int) – Dimension name or index

Returns:

Tuple of (dim_size, remaining_elements)

Return type:

tuple[int, int]

to_xr_dataarray()[source]#

Convert this AxisArray to an xarray DataArray.

This method creates an xarray DataArray with equivalent data, coordinates, dimensions, and attributes. Useful for interoperability with the xarray ecosystem.

Returns:

xarray DataArray representation of this AxisArray

Return type:

xarray.DataArray

static transpose(aa, dims=None)[source]#

Transpose (reorder) the dimensions of an AxisArray.

Parameters:

  • aa (T) – The AxisArray to transpose

  • dims (Iterable[str] | Iterable[int] | None) – New dimension order (names or indices). If None, reverses all dimensions

Returns:

New AxisArray with transposed dimensions

Return type:

T

view2d(dim)[source]#

Context manager providing a 2D view of the data.

Yields a 2D array view with the specified dimension as the first axis. Changes to the yielded array may be reflected in the original data.

Parameters:

dim (str | int) – Dimension name or index to move to first axis

Yields:

2D array view with shape (dim_size, remaining_elements)

Return type:

Generator[npt.NDArray, None, None]

class ezmsg.util.messages.axisarray.AxisBase(unit='')[source]#

Bases: ABC

Abstract base class for axes types used by AxisArray.

class ezmsg.util.messages.axisarray.CoordinateAxis(data, dims, unit='')[source]#

Bases: AxisBase, ArrayWithNamedDims

An axis implementation that uses explicit coordinate values stored in an array.

This class allows for non-linear or irregularly spaced coordinate systems by storing the actual coordinate values in a data array.

Inherits from both AxisBase and ArrayWithNamedDims, combining axis functionality with named dimension support.

Parameters:

  • unit (str) – The unit of measurement for this axis (inherited from AxisBase)

  • data (npt.NDArray) – The underlying numpy array data (inherited from ArrayWithNamedDims)

  • dims (list[str]) – List of dimension names, must match the array’s number of dimensions (inherited from ArrayWithNamedDims)

value(x)[source]#

Get coordinate value(s) at the given index(es).

Parameters:

x (int | npt.NDArray[numpy.int_]) – Index or array of indices to lookup

Returns:

Coordinate value(s) at the specified index(es)

Return type:

Any | npt.NDArray

class ezmsg.util.messages.axisarray.LinearAxis(unit='', gain=1.0, offset=0.0)[source]#

Bases: AxisBase

An axis implementation for sparse axes with regular intervals between elements.

It is called “linear” because it provides a simple linear mapping between indices and element values: value = (index * gain) + offset.

A typical example is a time axis (TimeAxis), with regular sampling rate.

Parameters:

  • gain (float) – Step size (scaling factor) for the linear axis

  • offset (float) – The offset (value) of the first sample

  • unit (str) – The unit of measurement for this axis (inherited from AxisBase)

classmethod create_time_axis(fs, offset=0.0)[source]#

Convenience method to construct a LinearAxis for time.

Parameters:

  • fs (float) – Sampling frequency in Hz

  • offset (float) – Time offset in seconds (default: 0.0)

Returns:

A LinearAxis configured for time with gain=1/fs and unit=’s’

Return type:

LinearAxis

index(v, fn=<ufunc 'rint'>)[source]#

Convert axis value(s) to index(es) using the inverse linear transformation.

Parameters:

  • v (float | npt.NDArray[np.float64]) – Axis value or array of values to convert to indices

  • fn (Callable) – Function to apply for rounding (default: np.rint)

Returns:

Corresponding index or array of indices

Return type:

int | npt.NDArray[numpy.int_]

value(x)[source]#

Convert index(es) to axis value(s) using the linear transformation.

Parameters:

x (int | npt.NDArray[numpy.int_]) – Index or array of indices to convert

Returns:

Corresponding axis value(s)

Return type:

float | npt.NDArray[np.float64]

ezmsg.util.messages.axisarray.as2d(in_arr, axis=0, *, xp)[source]#

Reshape array to 2D with specified axis first.

Parameters:

  • in_arr (npt.NDArray) – Input array

  • axis (int) – Axis to move to first position (default: 0)

  • xp – Array namespace (numpy, cupy, etc.)

Returns:

2D reshaped array

Return type:

npt.NDArray

ezmsg.util.messages.axisarray.shape2d(arr, axis=0)[source]#

Calculate the 2D shape when viewing array with specified axis first.

Parameters:

  • arr (npt.NDArray) – Input array

  • axis (int) – Axis to move to first position (default: 0)

Returns:

Tuple of (axis_size, remaining_elements)

Return type:

tuple[int, int]

ezmsg.util.messages.axisarray.slice_along_axis(in_arr, sl, axis)[source]#
Slice the input array along a specified axis using the given slice object or integer index.

Integer arguments to sl will cause the sliced dimension to be dropped. Use slice(my_int, my_int+1, None) to keep the sliced dimension.

Parameters:

  • in_arr (npt.NDArray) – The input array to be sliced

  • sl (slice | int) – The slice object or integer index to use for slicing

  • axis (int) – The axis along which to slice the array

Returns:

The sliced array (view)

Return type:

npt.NDArray

Raises:

ValueError – If the axis value is invalid for the input array

ezmsg.util.messages.axisarray.sliding_win_oneaxis(in_arr, nwin, axis, step=1)[source]#

Generates a view of an array using a sliding window of specified length along a specified axis of the input array. This is a slightly optimized version of nps.sliding_window_view with a few important differences:

  • This only accepts a single nwin and a single axis, thus we can skip some checks.

  • The new win axis precedes immediately the original target axis, unlike sliding_window_view where the

    target axis is moved to the end of the output.

Combine this with slice_along_axis(…, sl=slice(None, None, step), axis=axis) to step the window

by more than 1 sample at a time.

Parameters:

  • in_arr (npt.NDArray) – The input array

  • nwin (int) – The size of the sliding window

  • axis (int) – The axis along which the sliding window will be applied

  • step (int) – The size of the step between windows. If > 1, the strided window will be sliced with slice_along_axis (default = 1)

Returns:

A view to the input array with the sliding window applied

Return type:

npt.NDArray

Note: There is a known edge case when nwin == shape[axis] + 1. While this should raise

an error because the window is larger than the input, the implementation ends up returning a 0-length window. We could check for this but this function is intended to have minimal latency so we have decided to skip the checks and deal with the support issues as they arise.

ezmsg.util.messages.axisarray.view2d(in_arr, axis=0)[source]#

Context manager providing 2D view of the array, no matter what input dimensionality is. Yields a view of underlying data when possible, changes to data in yielded array will always be reflected in self.data, either via the view or copying.

NOTE: In practice, I’m not sure this is very useful because it requires modifying the numpy array data in-place, which limits its application to zero-copy messages

NOTE: The context manager allows the use of with when calling view2d.

Parameters:

  • in_arr (npt.NDArray) – Input array to be viewed as 2D

  • axis (int) – Dimension index to move to first axis (Default = 0)

Yields:

2D array view with shape (dim_size, remaining_elements)

Return type:

Generator[npt.NDArray, None, None]