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.

Data access under an AxisManager is done based on field names. For example:

>>> print(dset.boresight.az)
[0. 0. 0. ... 0. 0. 0.]

Advanced data access is possible by a path like syntax. This is especially useful when data access is dynamic and the field name is not known in advance. For example:

>>> print(dset["boresight.az"])
[0. 0. 0. ... 0. 0. 0.]

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">

To save data with flacarray compression, pass the “encodings” argument, and use a nested dict to identify the field you want to save and to specify the storage details. For example, suppose container dset has field dset.thermometers.diode1 that is a float array that can be safely compressed with precision 1e-4, and field dset.thermometers.flags is an integer array. Then sensible encodings request is:

 >>> encodings = {
      'thermometers': {
         'diode1': {
           'type': 'flacarray',
           'args': {
             'quanta': 1e-4
           }
         },
         'flags': {
           'type': 'flacarray'
         }
      }
    }
>>> dset.save('output.h5', encodings=encodings)

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.

Example usage:

  1. aman.move('hwp_angle', None)

    Deletes the field hwp_angle from aman.

  2. aman.move('hwp_angle', 'angle')

    Renames the field hwp_angle to angle.

  3. aman.move('preprocess.t2p.t2p_stats', None)

    Deletes the field t2p_stats from the sub-AxisManager aman.preprocess.t2p.

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, overwrite=False, restrict_in_place=False)[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.

  • overwrite (bool) – If True then will write over existing data in field name if present.

  • restrict_in_place (bool) – If True, then a wrapped AxisManager may be modified and added, without a copy first. This can be much faster, if there’s no need to preserve the wrapped item.

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.

reindex_axis(axis, indexes, in_place=True)[source]

Reindexes all data that is assigned to a specified axis with a new list/array of indexes. This is particularly useful if the number of detectors between the meta and obs data don’t match. This function will recursively delve through all AxisManagers in aman and will reindex every data array that is found assigned to an axis matching the specified axis.

Parameters:

  • axis (str) – The name of the axis in the aman to reindex.

  • indexes (int array) – an array of ints with length equal to the length of the new array and values equal to the idxs of the values in the data to be reindexed. Indexes that should be left as nan in the new array should be set to -1 or nan.

  • example (For) – data = [1,3,5], indexes = [0, -1, 2, 1] would result in new_data = [1, nan, 5, 3]

  • in_place (bool) – If in_place == True, the intersection is

  • Otherwise (applied to self.)

  • returned (a new object is)

:param : :param with data copied out.:

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, restrict_in_place=False)[source]

Merge the data from other AxisMangers into this one. Axes with the same name will be intersected.

If restrict_in_place=True, then the amans may be modified as they are added to the output objcet. When that arg is False, the incoming amans are all copied, even if no modifications are needed.

save(dest, group=None, overwrite=False, compression=None, encodings=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.

  • encodings (dict or None) – Special instructions for encoding / compression. See notes.

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.

The encodings dict is used to support flacarray compression of fields. Suppose member ‘quant_field’ should be FLAC-compressed, with precision 0.1. Then pass:

encodings={
  'quant_field': {
    'type': 'flacarray':
    'args': {
      'quanta': 0.1
    }
  }
}

For fields in sub-axismanagers, nest the specification; e.g.:

encodings={
  'subaman': {
     'quant_field': ...
  }
}

Only multi-dimensional int and float arrays may be so compressed. Unconsumed encoding information (e.g. specifying compression for a non-existent field or a field that is not an array) will cause an exception to be raised.

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.