coordax.Field.tag

Field.tag(*names: str | Coordinate | ellipsis | None) Field[source]

Returns a Field with attached coordinates to the positional axes.

Parameters:

*names – Names or coordinates to assign to the positional axes. The total number of dimensions corresponding to these objects must match the number of positional axes in the field unless ... is used to indicate positions of untagged axes.

Returns:

A new Field with the specified dimensions tagged.

Examples

>>> import coordax as cx
>>> import jax.numpy as jnp
>>> field = cx.Field(jnp.ones((2, 3)))
>>> field.tag('x', 'y')
<Field dims=('x', 'y') shape=(2, 3) axes={} >

Tagging with Coordinate objects adds them to the field:

>>> x = cx.SizedAxis('x', 2)
>>> field.tag(x, 'y')
<Field dims=('x', 'y') shape=(2, 3) axes={'x': SizedAxis} >

None leaves a dimension untagged:

>>> field.tag('x', None)
<Field dims=('x', None) shape=(2, 3) axes={} >

Ellipsis ... can be indicate dimensions that should not be named:

>>> field.tag(..., 'y')
<Field dims=(None, 'y') shape=(2, 3) axes={} >

Tagging with the wrong number of arguments raises an error:

>>> field.tag('x')
Traceback (most recent call last):
...
ValueError: there must be exactly as many dimensions given to ``tag`` as
there are positional axes in the array, but got ('x',) for 2 positional
axes.

You can also tag with multi-dimensional coordinates corresponding to multiple array axes. If the multi-dimensional coordinate is a coordax.CartesianProduct, it is unpacked:

>>> x = cx.SizedAxis('x', 2)
>>> y = cx.SizedAxis('y', 3)
>>> xy = cx.compose(x, y)
>>> xy
CartesianProduct(axes={'x': SizedAxis, 'y': SizedAxis})
>>> field = cx.Field(jnp.zeros((2, 3)))
>>> field.tag(xy)
<Field dims=('x', 'y') shape=(2, 3) axes={'x': SizedAxis, 'y': SizedAxis}
>