Fields

Key concepts:

  • Field is Coordax’s core data structure, an array with optional labeled axes.

  • Field methods tag() and untag() are used to attach/remove coordinates

  • Array dimensions labeled with None correspond to locally positional axes

  • Labeled dimensions facilitate reordering and broadcasting of underlying data

  • Coordinates are checked for consistency to catch alignment errors

Instantiating Field

To instantiate a Field object, either use the Field constructor directly (Field(data: Array, dims: tuple[str, ...], axes: dict[str, Coordinate])) or the field() helper function (cx.field(data: Array, *dims: str | Coordinate)).

Here’s a minimal example, using string axis names:

import coordax as cx
import numpy as np

field_data = np.random.RandomState(0).uniform(size=(3, 4))

# or equivalently: cx.Field(field_data, dims=('x', 'y'))
field = cx.field(field_data, 'x', 'y')

field

<Field dims=('x', 'y') shape=(3, 4) axes={} >

The field’s data and dimension names are stored in the data and dims attributes:

field.data

array([[0.5488135 , 0.71518937, 0.60276338, 0.54488318],
       [0.4236548 , 0.64589411, 0.43758721, 0.891773  ],
       [0.96366276, 0.38344152, 0.79172504, 0.52889492]])

field.dims

('x', 'y')

field.axes  # empty

{}

To specify an axis of a Field as unlabeled, use None rather than a string.

partially_labeled_field = cx.field(field_data, 'x', None)
partially_labeled_field

<Field dims=('x', None) shape=(3, 4) axes={} >

When some of the axes of a Field are not labeled, we say that the Field has positional axes. Positional axes play a crucial role in how computation is performed on Field objects.

The overall shape of a Field is partitioned into positional_shape and named_shape properties, indicating the sizes of axes associated with each.

print(f'{field.shape=}')
print(f'{field.positional_shape=}')
print(f'{field.named_shape=}')

print(f'{partially_labeled_field.shape=}')
print(f'{partially_labeled_field.positional_shape=}')
print(f'{partially_labeled_field.named_shape=}')

field.shape=(3, 4)
field.positional_shape=()
field.named_shape={'x': 3, 'y': 4}
partially_labeled_field.shape=(3, 4)
partially_labeled_field.positional_shape=(4,)
partially_labeled_field.named_shape={'x': 3}

You can use Coordinate objects as an alternative to strings to annotate axes in more detail (e.g., to add tick labels). Here’s a simple example, using the built-in LabeledAxis coordinate:

y_axis = cx.LabeledAxis('y', np.linspace(0, 1, 4))
field = cx.field(field_data, 'x', y_axis)
field

<Field dims=('x', 'y') shape=(3, 4) axes={'y': LabeledAxis} >

We’ll cover Coordinate objects in more detail later, but for now, note that they exist and are used by Coordax for alignment checks. They are listed in the axes attribute:

field.axes

{'y': coordax.LabeledAxis('y', ticks=array([0.        , 0.33333333, 0.66666667, 1.        ]))}

Updating coordinate labels using tag/untag

To instantiate {py.class}~cx.Field with a different set of coordinate labels from another Field, it is convenient to use tag/untag methods.

These restrictions are necessary to avoid ambiguity in which positional axis is being tagged or in which order should untagged axes appear in the positional_shape. Enforcing the above rules removes such ambiguity.

partially_labeled_field = field.untag('x')
labeled_field = partially_labeled_field.tag('x')  # same as field
cx.testing.assert_fields_equal(field, labeled_field)

When dealing with a pytee of Field objects, it may be desired to relabel all entries. This can be done using simple helpers coordax.tag() and coordax.untag().

tree = {
    'a': cx.field(np.ones((1, 1, 1)), 'x', 'y', 'z'),
    'tuple': (cx.field(np.array([np.pi]), 'x'), cx.field(np.array([np.e]), 'x')),
}
tree

{'a': <Field dims=('x', 'y', 'z') shape=(1, 1, 1) axes={} >,
 'tuple': (<Field dims=('x',) shape=(1,) axes={} >,
  <Field dims=('x',) shape=(1,) axes={} >)}

cx.untag(tree, 'x')

{'a': <Field dims=(None, 'y', 'z') shape=(1, 1, 1) axes={} >,
 'tuple': (<Field dims=(None,) shape=(1,) axes={} >,
  <Field dims=(None,) shape=(1,) axes={} >)}

cx.tag(cx.untag(tree, 'x'), 'batch')

{'a': <Field dims=('batch', 'y', 'z') shape=(1, 1, 1) axes={} >,
 'tuple': (<Field dims=('batch',) shape=(1,) axes={} >,
  <Field dims=('batch',) shape=(1,) axes={} >)}

Later we will see that a common pattern with coordax is to:

  1. untag dimension on which we want to operate

  2. perform the desired computation

  3. retag the result with approriate coordinates

Reordering, broadcasting and basic operations on Field

Reordering and broadcasting of Field objects with order_as() and broadcast_like() can be done using their coordinates or dimension names:

f_xy = cx.field(np.arange(2 * 3).reshape((2, 3)), 'x', 'y')
f_yx = f_xy.order_as('y', 'x')  # transpose to a new order.
print(f'{f_xy.dims=}, {f_yx.dims=}')

f_xy.dims=('x', 'y'), f_yx.dims=('y', 'x')

f = cx.field(np.ones((4, 7)), 'x', 'y')
arange = cx.field(np.arange(4), 'x')
b_arange = arange.broadcast_like(f)  # broadcast to another
b_arange

<Field dims=('x', 'y') shape=(4, 7) axes={} >

Ellipsis can also be used in tag() and order_as() to indicate “all other dimensions:”

f = cx.field(np.ones((1, 2, 4, 3)), 'x', 'y', 'z', 'q')
f_qy = f.order_as('q', 'y', ...)
f_xq = f.untag('y', 'z')
another_f_xq = cx.field(np.ones((1, 2, 4, 3)), 'x', ..., 'q')
assert another_f_xq.dims == f_xq.dims
f_qy, f_xq

(<Field dims=('q', 'y', 'x', 'z') shape=(3, 2, 1, 4) axes={} >,
 <Field dims=('x', None, None, 'q') shape=(1, 2, 4, 3) axes={} >)

Field also supports a small set of operations directly, including arithmetic with python numeric types (e.g., scaling):

cx.field(np.arange(5), 'x') * 2

<Field dims=('x',) shape=(5,) axes={} >