AxisManager

The AxisManager is a container class for numpy arrays and similar structures, that tracks relationships between axes. The object can be made aware, for example, that the 2nd axis of an array called “data” is tied to the 1st axis of an array called “azimuth”; slices / sub-selection operations can be performed on a named on that axis and the entire object will be modified self-consistently.

The data model here is similar to what is provided by the xarray library. However, AxisManager makes it possible to store objects other than numpy ndarrays, if the classes expose a suitable interface.

Internal data structures

If you’re debugging or trying to understand how AxisManager works, know that the entire state is stored in the following private data members:

_axes
An odict that maps axis name to an Axis-type object (such as an instance of LabelAxis or IndexAxis).

_assignments
An odict that maps field name to a tuple, where each entry in the tuple gives the name of the axis to which that dimension of the underlying data is associated (or None, for a dimension that is unassociated). For example an array with shape (100, 2, 20000) might have an axis assignment of ('dets', None, 'samps').

_fields
An odict that maps field name to the data object.

Consistency of the internal state requires:

The key under which an Axis-type object in _axes must match that object’s .name attribute. (I.e. _axes[k].name == k.)
The _assignments and _fields odicts should have the same keys, and the dimensionality must agree. (I.e. len(_assignments[k]) == len(_fields[k].shape).)

Tutorial

Here we demonstrate some usage patterns for AxisManager. For these to work, you need these imports:

from sotodlib import core
import numpy as np

Suppose you have an array of detector readings. It’s 2d, with the first axis representing some particular detectors and the second axis representing a time index.

dets = ['det10', 'det11', 'det12']
tod = np.zeros((3, 10000)) + [[10],[11],[12]]

AxisManager is a container that can hold numpy arrays (and also it can hold other AxisManagers, and some other stuff too).

dset = core.AxisManager().wrap('tod', tod)

AxisManagers can also hold scalars, a value is considered scalar if 'np.isscalar' thinks it is a scalar or if it is 'None'.

dset = dset.wrap('scalar', 1.0)

Inspecting:

>>> print(dset)
AxisManager(tod[3,10000], scalar)
>>> print(dset.tod)
[[10. 10. 10. ... 10. 10. 10.]
 [11. 11. 11. ... 11. 11. 11.]
 [12. 12. 12. ... 12. 12. 12.]]
>>> print(dset.scalar)
1.0

The value that AxisManager adds is an ability to relate an axis in one child array to an axis in another child. This time, when we add 'tod', we describe the two dimensions of that array with LabelAxis and IndexAxis objects.

dset = core.AxisManager().wrap('tod', tod, [(0, core.LabelAxis('dets', dets)),
                                            (1, core.IndexAxis('samps'))])

Inspecting:

>>> print(dset)
AxisManager(tod[dets,samps], dets:LabelAxis(3), samps:IndexAxis(10000))

Now if we add other arrays, we can assign their axes to these existing ones:

hwp_angle = np.arange(tod.shape[1]) * 2. / 400 % 360.
dset.wrap('hwp_angle', hwp_angle, [(0, 'samps')])

The output of the wrap call is:

AxisManager(tod[dets,samps], hwp_angle[samps], dets:LabelAxis(3),
  samps:IndexAxis(10000))

To create new arrays in the AxisManager, and have certain dimensions automatically matched to a particular named axis, use the wrap_new method:

events = dset.wrap_new('event_count', shape=('samps', 3), dtype='int')

This object returned by this call is a numpy array of ints with shape (10000, 3). It is wrapped into dset under the name 'event_count'. Its first axis is tied to the 'dets' axis.

We can also embed related AxisManagers within the existing one, to establish a hierarchical structure:

boresight = core.AxisManager(core.IndexAxis('samps'))
for k in ['el', 'az', 'roll']:
    boresight.wrap(k, np.zeros(tod.shape[1]), [(0,'samps')])

dset.wrap('boresight', boresight)

The output of the wrap cal should be:

AxisManager(tod[dets,samps], hwp_angle[samps], boresight*[samps],
  dets:LabelAxis(3), samps:IndexAxis(10000))

Note the boresight entry is marked with a *, indicating that it’s an AxisManager rather than a numpy array.

To slice this object, use the restrict() method. First, let’s restrict in the ‘dets’ axis. Since it’s an Axis of type LabelAxis, the restriction selector must be a list of strings:

dset.restrict('dets', ['det11', 'det12'])

Similarly, restricting in the samps axis:

dset.restrict('samps', (10, 300))

After those two restrictions, inspect the shapes of contained objects:

>>> print(dset.tod.shape)
(2, 290)
>>> print(dset.boresight.az.shape)
(290,)

For debugging, you can write AxisManagers to HDF5 files and then read them back. (This is an experimental feature so don’t rely on this for long term stability!):

>>> dset.save('output.h5', 'my_axismanager/dset')

>>> dset_reloaded = AxisManager.load('output.h5', 'my_axismanager/dset')
>>> dset_reloaded
AxisManager(tod[dets,samps], hwp_angle[samps], boresight*[samps],
  dets:LabelAxis(2), samps:IndexAxis(290))

Numerical arrays are stored as simple HDF5 datasets, so you can also use h5py to load the saved arrays:

>>> import h5py
>>> f = h5py.File('output.h5')
>>> f['my_axismanager/dset/tod'][:]
<HDF5 dataset "tod": shape (2,290), type "<f8">

Standardized Fields

As we develop the SO pipeline we will need to standardize field names that have specific uses within the pipeline so that functions can be written to expect a specific set of fields. Not all AxisManagers will have all these fields by default and many fields are linked to documentation locations where more details can be found. These are meant to prevent naming collisions and more will be added here as the code develops.

dets - the axis for detectors
samps - the axis for samples
timestamps [samps] - the field for UTC timestamps
signal [dets, samps] - the field for detector signal
obs_info - AxisManager of scalars with ObsDb information for the loaded

observation. Details here.
det_info [dets] - AxisManager containing loaded detector metadata
- readout_id - The unique readout ID of the resonator
- det_id - The unique detector ID matched to the resonator.
- wafer - An AxisManager of different parameters related to the hardware
  
  mapping on the UFM itself. Loaded based on det_id field in the
  
  det_info.

SMuRF fields loaded through sotodlib.io.load_file().

bias_lines - the axis for bias lines in a UFM.
status - A SmurfStatus AxisManager containing information from status

frames in the .g3 timestreams. sotodlib.io.load_smurf.SmurfStatus
iir_params - An AxisManager with the readout filter parameters. Used by

sotodlib.tod_ops.filters.iir_filter().
primary [samps] - An AxisManager with SMuRF readout information that is

synced with the timestreams.
biases [bias_lines, samps] - Bias voltage applied to TESes over time.

Pointing information required for sotodlib.coords.

boresight [samps] - AxisManager with boresight pointing in horizon

coordinates. Child fields are az, el, and roll.
focal_plane [dets] - AxisManager with detector position and orientation

relative to boresight pointing.
boresight_equ [samps] - AxisManager with boresight in equitorial

coordinates.

HWP information

hwp_angle [samps] - AxisManager with hwp rotation angle required for

sotodlib.io.g3tsmurf_utils.load_hwp_data
hwpss_model [dets, samps] - the fields of fitted HWP synchronous signal

derived from sotodlib.hwp.get_hwpss

Reference

The class documentation for AxisManager and the basic Axis types should be rendered here.

AxisManager

class sotodlib.core.AxisManager(*args)[source]

A container for numpy arrays and other multi-dimensional data-carrying objects (including other AxisManagers). This object keeps track of which dimensions of each object are concordant, and allows one to slice all hosted data simultaneously.

__init__(*args)[source]

move(name, new_name)[source]: Rename or remove a data field. To delete the field, pass new_name=None.

static concatenate(items, axis=0, other_fields='exact')[source]

Concatenate multiple AxisManagers along the specified axis, which can be an integer (corresponding to the order in items[0]._axes) or the string name of the axis.

This operation is difficult to sanity check so it’s best to use it only in simple, controlled cases! The first item is given significant privilege in defining what fields are relevant. Fields that appear in the first item, but do not share the target axis, will be treated as follows depending on the value of other_fields:

If other_fields=’exact’ will compare entries in all items and if they’re identical will add it. Otherwise will fail with a ValueError.
If other_fields=’fail’, the function will fail with a ValueError.
If other_fields=’first’, the values from the first element of items will be copied into the output.
If other_fields=’drop’, the fields will simply be ignored (and the output will only contain fields that share the target axis).

wrap(name, data, axis_map=None, accept_new=True, accept_merge=True)[source]

Add data into the AxisManager.

Parameters:

name (str) – name of the new data.
data – The data to register. This must be of an acceptable type, i.e. a numpy array or another AxisManager. If scalar (or None) then data will be directly added to _fields with no associated axis.
axis_map – A list that assigns dimensions of data to particular Axes. Each entry in the list must be a tuple with the form (dim, name) or (dim, ax), where dim is the index of the dimension being described, name is a string giving the name of an axis already described in the present object, and ax is an AxisInterface object.

wrap_new(name, shape=None, cls=None, **kwargs)[source]

Create a new object and wrap it, with axes mapped. The shape can include axis names instead of ints, and that will cause the new object to be dimensioned properly and its axes mapped.

Parameters:

name (str) – name of the new data.
shape (tuple of int and std) – shape in the same sense as numpy, except that instead of int it is allowed to pass the name of a managed axis.
cls (callable) – Constructor that should be used to construct the object; it will be called with all kwargs passed to this function, and with the resolved shape as described here. Defaults to numpy.ndarray.

Examples

Construct a 2d array and assign it the name ‘boresight_quat’, with its first axis mapped to the AxisManager tod’s “samps” axis:

>>> tod.wrap_new('boresight_quat', shape=('samps', 4), dtype='float64')

Create a new empty RangesMatrix, carrying a per-det, per-samp flags:

>>> tod.wrap_new('glitch_flags', shape=('dets', 'samps'),
                 cls=so3g.proj.RangesMatrix.zeros)

restrict_axes(axes, in_place=True)[source]

Restrict this AxisManager by intersecting it with a set of Axis definitions.

Parameters:

axes (list or dict of Axis)
in_place (bool) – If in_place == True, the intersection is applied to self. Otherwise, a new object is returned, with data copied out.

Returns:

The restricted AxisManager.

restrict(axis_name, selector, in_place=True)[source]

Restrict the AxisManager by selecting a subset of items in some Axis. The Axis definition and all data fields mapped to that axis will be modified.

Parameters:

axis_name (str) – The name of the Axis.
selector (slice or special) – Selector, in a form understood by the underlying Axis class (see the .restriction method for the Axis).
in_place (bool) – If True, modifications are made to this object. Otherwise, a new object with the restriction applied is returned.

Returns:

The AxisManager with restrictions applied.

static intersection_info(*items)[source]: Given a list of AxisManagers, scan the axes and combine (intersect) any common axes. Returns a dict that maps axis name to restricted Axis object.

merge(*amans)[source]: Merge the data from other AxisMangers into this one. Axes with the same name will be intersected.

save(dest, group=None, overwrite=False, compression=None)[source]

Write this AxisManager data to an HDF5 group. This is an experimental feature primarily intended to assist with debugging. The schema is subject to change, and it’s possible that not all objects supported by AxisManager can be serialized.

Parameters:

dest (str or h5py.Group) – Place to save it (in combination with group).
group (str or None) – Group within the HDF5 file (relative to dest).
overwrite (bool) – If True, remove any existing thing at the specified address before writing there.
compression (str or None) – Compression filter to apply. E.g. ‘gzip’. This string is passed directly to HDF5 dataset routines.

Notes

If dest is a string, it is taken to be an HDF5 filename and is opened in ‘a’ mode. The group, in that case, is the full group name in the file where the data should be written.

If dest is an h5py.Group, the group is the group name in the file relative to dest.

The overwrite argument only matters if group is passed as a string. A RuntimeError is raised if the group address already exists and overwrite==False.

For example, these are equivalent:

# Filename + group address:
axisman.save('test.h5', 'x/y/z')

# Open h5py.File + group address:
with h5py.File('test.h5', 'a') as h:
  axisman.save(h, 'x/y/z')

# Partial group address
with h5py.File('test.h5', 'a') as h:
  g = h.create_group('x/y')
  axisman.save(g, 'z')

When passing a filename, the code probably won’t use a context manager… so if you want that protection, open your own h5py.File as in the 2nd and 3rd example.

classmethod load(src, group=None, fields=None)[source]

Load a saved AxisManager from an HDF5 file and return it. See docs for save() function.

The (src, group) args are combined in the same way as (dest, group) in the save function. Examples:

axisman = AxisManager.load('test.h5', 'x/y/z')

with h5py.File('test.h5', 'r') as h:
  axisman = AxisManager.load(h, 'x/y/z')

If the fields argument is specified, it must be a list of strings indicating what subfields of the stored AxisManager should be extracted. For nested entries, connect fields with “.”. For example fields=['subaman.field1', 'subaman.field2']. When fields is specified, _all_ axes from the AxisManager are included in the result, even if not directly referenced by the requested fields; this behavior is subject to change.

IndexAxis

class sotodlib.core.IndexAxis(name, count=None)[source]

This class manages a simple integer-indexed axis. When intersecting data, the longer one will be simply truncated to match the shorter. Selectors must be slice objects (with stride 1!) or tuples to be passed into slice(), e.g. (0, 1000) or (0, None, 1)..

__init__(name, count=None)[source]

resolve(src, axis_index=None)[source]

Perform a check or promote-and-check of this Axis against a data object.

The promotion step only applies to “unset” Axes, i.e. those here count is None. Not all Axis types will be able to support this. Promotion involves inspection of src and axis_index to fix free parameters in the Axis. If promotion is successful, a new (“set”) Axis is returned. If promotion is attempted and fails then a ValueError is raised.X

The check step involes confirming that the data object described by src (and axis_index) is compatible with the current axis (or with the axis resulting from axis Promotion). Typically that simply involves confirming that src.shape[axis_index] == self.count. If the check fails, a ValueError is raised.

Parameters:

src – The data object to be wrapped (e.g. a numpy array)
axis_index – The index of the data object to test for compatibility.

Returns:

Either self, or the result of promotion.

Return type:

axis

restriction(selector)[source]

Apply selector to the elements of this axis, returning a new Axis of the same type and an array indexing object (a slice or an array of integer indices) that may be used to extract the corresponding elements from a vector.

See class header for acceptable selector objects.

Returns (new_axis, ar_index).

intersection(friend, return_slices=False)[source]

Find the intersection of this Axis and the friend Axis, returning a new Axis of the same type. Optionally, also return array indexing objects that select the common elements from array dimensions corresponding to self and friend, respectively.

See class header for acceptable selector objects.

Returns (new_axis), or (new_axis, ar_index_self, ar_index_friend) if return_slices is True.

OffsetAxis

class sotodlib.core.OffsetAxis(name, count=None, offset=0, origin_tag=None)[source]

This class manages an integer-indexed axis, with an accounting for an integer offset of any single vector relative to some absolute reference point. For example, one vector could could have 100 elements at offset 50, and a second vector could have 100 elements at offset -20. On intersection, the result would have 30 elements at offset 50.

The property origin_tag may be used to identify the absolute reference point. It could be a TOD name (‘obs_2020-12-01’) or a timestamp or whatever.

Selectors must be slice objects (with stride 1!) or tuples to be passed into slice(), e.g. (0, 1000) or (0, None, 1).

__init__(name, count=None, offset=0, origin_tag=None)[source]

resolve(src, axis_index=None)[source]

Perform a check or promote-and-check of this Axis against a data object.

The promotion step only applies to “unset” Axes, i.e. those here count is None. Not all Axis types will be able to support this. Promotion involves inspection of src and axis_index to fix free parameters in the Axis. If promotion is successful, a new (“set”) Axis is returned. If promotion is attempted and fails then a ValueError is raised.X

The check step involes confirming that the data object described by src (and axis_index) is compatible with the current axis (or with the axis resulting from axis Promotion). Typically that simply involves confirming that src.shape[axis_index] == self.count. If the check fails, a ValueError is raised.

Parameters:

src – The data object to be wrapped (e.g. a numpy array)
axis_index – The index of the data object to test for compatibility.

Returns:

Either self, or the result of promotion.

Return type:

axis

restriction(selector)[source]

Apply selector to the elements of this axis, returning a new Axis of the same type and an array indexing object (a slice or an array of integer indices) that may be used to extract the corresponding elements from a vector.

See class header for acceptable selector objects.

Returns (new_axis, ar_index).

intersection(friend, return_slices=False)[source]

Find the intersection of this Axis and the friend Axis, returning a new Axis of the same type. Optionally, also return array indexing objects that select the common elements from array dimensions corresponding to self and friend, respectively.

See class header for acceptable selector objects.

Returns (new_axis), or (new_axis, ar_index_self, ar_index_friend) if return_slices is True.

LabelAxis

class sotodlib.core.LabelAxis(name, vals=None)[source]

This class manages a string-labeled axis, i.e., an axis where each element has been given a unique name. The vector of names can be found in self.vals.

Instantiation with labels that are not strings will raise a TypeError.

On intersection of two vectors, only elements whose names appear in both axes will be preserved.

Selectors should be lists (or arrays) of label strings.

__init__(name, vals=None)[source]

resolve(src, axis_index=None)[source]

Perform a check or promote-and-check of this Axis against a data object.

The promotion step only applies to “unset” Axes, i.e. those here count is None. Not all Axis types will be able to support this. Promotion involves inspection of src and axis_index to fix free parameters in the Axis. If promotion is successful, a new (“set”) Axis is returned. If promotion is attempted and fails then a ValueError is raised.X

The check step involes confirming that the data object described by src (and axis_index) is compatible with the current axis (or with the axis resulting from axis Promotion). Typically that simply involves confirming that src.shape[axis_index] == self.count. If the check fails, a ValueError is raised.

Parameters:

src – The data object to be wrapped (e.g. a numpy array)
axis_index – The index of the data object to test for compatibility.

Returns:

Either self, or the result of promotion.

Return type:

axis

restriction(selector)[source]

Apply selector to the elements of this axis, returning a new Axis of the same type and an array indexing object (a slice or an array of integer indices) that may be used to extract the corresponding elements from a vector.

See class header for acceptable selector objects.

Returns (new_axis, ar_index).

intersection(friend, return_slices=False)[source]

Find the intersection of this Axis and the friend Axis, returning a new Axis of the same type. Optionally, also return array indexing objects that select the common elements from array dimensions corresponding to self and friend, respectively.

See class header for acceptable selector objects.

Returns (new_axis), or (new_axis, ar_index_self, ar_index_friend) if return_slices is True.