warn or raise error when merge/wrap inconsistent axes

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.

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']