diff --git a/sotodlib/core/axisman.py b/sotodlib/core/axisman.py index 06792a451..a30b57716 100644 --- a/sotodlib/core/axisman.py +++ b/sotodlib/core/axisman.py @@ -1,3 +1,4 @@ +import logging import numpy as np from collections import OrderedDict as odict @@ -12,6 +13,8 @@ from .util import get_coindices +logger = logging.getLogger(__name__) + class AxisInterface: """Abstract base class for axes managed by AxisManager.""" @@ -635,7 +638,7 @@ def concatenate(items, axis=0, other_fields='exact'): # Add and remove data while maintaining internal consistency. def wrap(self, name, data, axis_map=None, - overwrite=False, restrict_in_place=False): + overwrite=False, restrict_in_place=False, on_mismatch='warn'): """Add data into the AxisManager. Arguments: @@ -662,6 +665,12 @@ def wrap(self, name, data, axis_map=None, first. This can be much faster, if there's no need to preserve the wrapped item. + on_mismatch (str): Policy when an axis of the new data + has the same name as, but different coverage than, an + existing axis: 'intersect', 'warn' (intersect, but log a + warning; default) or 'error' (raise ValueError, before + modifying anything). + """ if overwrite and (name in self._fields): self.move(name, None) @@ -704,7 +713,8 @@ def wrap(self, name, data, axis_map=None, assign[index] = axis.name helper._fields[name] = data helper._assignments[name] = assign - return self.merge(helper, restrict_in_place=restrict_in_place) + return self.merge(helper, restrict_in_place=restrict_in_place, + on_mismatch=on_mismatch) def wrap_new(self, name, shape=None, cls=None, **kwargs): """Create a new object and wrap it, with axes mapped. The shape can @@ -1042,9 +1052,12 @@ def intersection_info(*items): ax, False) return axes_out - def merge(self, *amans, restrict_in_place=False): + def merge(self, *amans, restrict_in_place=False, on_mismatch='warn'): """Merge the data from other AxisMangers into this one. Axes with the - same name will be intersected. + same name will be intersected. The on_mismatch argument sets the policy + for axes that must be truncated (intersected) because amans share an + axis name but not its coverage: 'intersect', 'warn' (log a warning) + or 'error' (raise ValueError before modifying anything). If restrict_in_place=True, then the amans may be modified as they are added to the output objcet. When that arg is False, @@ -1064,6 +1077,28 @@ def merge(self, *amans, restrict_in_place=False): # Get the intersected axis descriptions. axes_out = self.intersection_info(self, *amans) + + # Apply the axis-mismatch policy before modifying anything. + assert on_mismatch in ('intersect', 'warn', 'error') + if on_mismatch != 'intersect': + report = [] + for src, aman in [('self', self)] + [ + ('merged[%i]' % i, a) for i, a in enumerate(amans)]: + for ax in aman._axes.values(): + if ax.count is None or ax.name not in axes_out: + continue + new_ax = axes_out[ax.name] + if new_ax.count != ax.count: + report.append( + "axis '%s' of %s changes %s -> %s." % ( + ax.name, src, repr(ax), repr(new_ax))) + if report: + msg = ('Axis mismatch on merge/wrap causes truncation: ' + + '; '.join(report)) + if on_mismatch == 'error': + raise ValueError(msg) + logger.warning(msg) + # Reduce the data in self, update our axes. self.restrict_axes(axes_out) # Import the other ones. diff --git a/tests/test_core.py b/tests/test_core.py index c31b0fd49..ad7d47736 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -489,6 +489,78 @@ def test_410_merge(self): self.assertEqual(coparent.shape, (4, n)) + def test_420_axis_mismatch_policy(self): + # wrap/merge intersect axes that share a name but not coverage; + # the on_mismatch policy controls whether that is silent + # ('intersect'), logged ('warn', default) or a ValueError ('error'). + logger = core.axisman.logger + dets = ['det0', 'det1', 'det2'] + n, ofs = 1000, 0 + + def make_tod(): + tod = core.AxisManager( + core.LabelAxis('dets', dets), + core.OffsetAxis('samps', n, ofs)) + tod.wrap('sig', np.zeros((len(dets), n)), + [(0, 'dets'), (1, 'samps')]) + return tod + + # Default policy: shorter incoming axis logs a warning, then + # truncates exactly as before. + tod = make_tod() + with self.assertLogs(logger, 'WARNING'): + tod.wrap('x', np.zeros(n - 100), + [(0, core.OffsetAxis('samps', n - 100, ofs))]) + self.assertEqual(tod.samps.count, n - 100) + self.assertEqual(tod.sig.shape, (3, n - 100)) + + # Longer incoming axis warns too (the incoming data is cut). + tod = make_tod() + with self.assertLogs(logger, 'WARNING'): + tod.wrap('x', np.zeros(n + 100), + [(0, core.OffsetAxis('samps', n + 100, ofs))]) + self.assertEqual(len(tod.x), n) + + # Same count but different offset warns. + tod = make_tod() + with self.assertLogs(logger, 'WARNING'): + tod.wrap('x', np.zeros(n), + [(0, core.OffsetAxis('samps', n, ofs + 10))]) + self.assertEqual(tod.samps.count, n - 10) + + # LabelAxis truncation (dets subset via child) warns. + tod = make_tod() + child = core.AxisManager(core.LabelAxis('dets', dets[:2])) + with self.assertLogs(logger, 'WARNING'): + tod.wrap('child', child) + self.assertEqual(tod.dets.count, 2) + + # No error with matching axes. + tod = make_tod() + tod.wrap('x', np.zeros(n), + [(0, core.OffsetAxis('samps', n, ofs))], + on_mismatch='error') + + # on_mismatch='error' raises, leaving the target untouched. + tod = make_tod() + with self.assertRaises(ValueError): + tod.wrap('x', np.zeros(n - 100), + [(0, core.OffsetAxis('samps', n - 100, ofs))], + on_mismatch='error') + self.assertEqual(tod.samps.count, n) + self.assertEqual(tod.sig.shape, (3, n)) + self.assertNotIn('x', tod) + + # Through merge() directly. + tod = make_tod() + _tod = core.AxisManager( + core.LabelAxis('dets', dets + ['det3']), + core.OffsetAxis('samps', n, ofs - n // 2)) + with self.assertRaises(ValueError): + tod.merge(_tod, on_mismatch='error') + self.assertEqual(tod.shape, (3, n)) + + def test_500_io(self): # Test save/load HDF5 dets = ['det0', 'det1', 'det2']