gh-385: Non square mms (#384)

Author: JaimeRZP

heracles/twopoint.py

@@ -31,7 +31,7 @@
 
 from .core import TocDict, toc_match, update_metadata
 from .progress import NoProgress, Progress
-from .result import Result, binned
+from .result import Result, binned, get_result_array
 
 try:
     from copy import replace
@@ -461,7 +461,25 @@ def invert_mixing_matrix(
             else:
                 _inv_M = np.linalg.pinv(_M, rcond=_rcond)
 
-            inv_M[key] = replace(M[key], array=_inv_M)
+            if _n != _m:
+                # A pseudo-inverse swaps output/input ell dimensions. Build
+                # angular arrays from the inverse output size (row dimension).
+                axis = value.axis[0]
+                out_size = _inv_M.shape[axis]
+                ell = np.arange(out_size)
+                lower = np.arange(out_size)
+                upper = np.arange(1, out_size + 1)
+                weight = np.ones(out_size)
+                inv_M[key] = replace(
+                    M[key],
+                    array=_inv_M,
+                    ell=ell,
+                    lower=lower,
+                    upper=upper,
+                    weight=weight,
+                )
+            else:
+                inv_M[key] = replace(M[key], array=_inv_M)
     return inv_M
 
 
@@ -480,6 +498,11 @@ def apply_mixing_matrix(d, M):
         s1, s2 = d[key].spin
         _d = np.atleast_2d(d[key].array)
         _M = M[key].array
+        # Output angular metadata follows the matrix output ell axis.
+        _ell = get_result_array(M[key], "ell")[0]
+        _lower = get_result_array(M[key], "lower")[0]
+        _upper = get_result_array(M[key], "upper")[0]
+        _weight = get_result_array(M[key], "weight")[0]
         if (s1 != 0) and (s2 != 0):
             _corr_d_EE = _M[0] @ _d[0, 0] + _M[1] @ _d[1, 1]
             _corr_d_BB = _M[1] @ _d[0, 0] + _M[0] @ _d[1, 1]
@@ -492,5 +515,12 @@ def apply_mixing_matrix(d, M):
                 _corr_d.append(_M @ cl)
             _corr_d = np.squeeze(_corr_d)
         _corr_d = np.array(list(_corr_d), dtype=dtype)
-        corr_d[key] = replace(d[key], array=_corr_d)
+        corr_d[key] = replace(
+            d[key],
+            array=_corr_d,
+            ell=_ell,
+            lower=_lower,
+            upper=_upper,
+            weight=_weight,
+        )
     return corr_d

tests/test_twopoint.py

@@ -418,3 +418,85 @@ def test_inverting_mixing_matrices():
         _mixed_cls = np.sum(mixed_cls[key].array)
         print(key, _mixed_cls)
         np.testing.assert_allclose(_mixed_cls, (n + 1) * 1.0)
+
+
+def test_nonsquare_mixing_matrix_metadata_binning_and_io(tmp_path):
+    import heracles
+
+    from heracles.twopoint import Result, apply_mixing_matrix, invert_mixing_matrix
+
+    key = ("POS", "POS", 0, 0)
+
+    in_ell = np.arange(8)
+    out_ell = np.array([2.0, 4.0, 6.0, 8.0, 10.0])
+    out_lower = np.array([1.5, 3.5, 5.5, 7.5, 9.5])
+    out_upper = np.array([2.5, 4.5, 6.5, 8.5, 10.5])
+
+    cls_in = {
+        key: Result(
+            np.linspace(1.0, 8.0, in_ell.size),
+            spin=(0, 0),
+            axis=(0,),
+            ell=in_ell,
+        )
+    }
+
+    mm = {
+        key: Result(
+            np.arange(out_ell.size * in_ell.size, dtype=float).reshape(
+                out_ell.size, in_ell.size
+            ),
+            spin=(0, 0),
+            axis=(0,),
+            ell=out_ell,
+            lower=out_lower,
+            upper=out_upper,
+        )
+    }
+
+    mixed = apply_mixing_matrix(cls_in, mm)
+    np.testing.assert_array_equal(mixed[key].ell, out_ell)
+    np.testing.assert_array_equal(mixed[key].lower, out_lower)
+    np.testing.assert_array_equal(mixed[key].upper, out_upper)
+
+    bins = np.array([2.0, 5.0, 11.0])
+    mixed_binned = heracles.binned(mixed, bins)
+    assert mixed_binned[key].shape == (2,)
+    np.testing.assert_array_equal(mixed_binned[key].lower, bins[:-1])
+    np.testing.assert_array_equal(mixed_binned[key].upper, bins[1:])
+
+    mixed_path = tmp_path / "mixed_nonsquare.fits"
+    heracles.write(mixed_path, mixed_binned)
+    mixed_read = heracles.read(mixed_path)
+    np.testing.assert_array_equal(mixed_read[key], mixed_binned[key])
+    np.testing.assert_array_equal(mixed_read[key].ell, mixed_binned[key].ell)
+
+    inv_mm = invert_mixing_matrix(mm, rcond=1e-12)
+    assert inv_mm[key].shape == (in_ell.size, out_ell.size)
+    np.testing.assert_array_equal(inv_mm[key].ell, np.arange(in_ell.size))
+    np.testing.assert_array_equal(inv_mm[key].lower, np.arange(in_ell.size))
+    np.testing.assert_array_equal(inv_mm[key].upper, np.arange(1, in_ell.size + 1))
+
+    cls_out = {
+        key: Result(
+            np.linspace(1.0, out_ell.size, out_ell.size),
+            spin=(0, 0),
+            axis=(0,),
+            ell=out_ell,
+        )
+    }
+
+    unmixed = apply_mixing_matrix(cls_out, inv_mm)
+    np.testing.assert_array_equal(unmixed[key].ell, np.arange(in_ell.size))
+    np.testing.assert_array_equal(unmixed[key].lower, np.arange(in_ell.size))
+    np.testing.assert_array_equal(unmixed[key].upper, np.arange(1, in_ell.size + 1))
+
+    inv_bins = np.array([0.0, 4.0, 8.0])
+    unmixed_binned = heracles.binned(unmixed, inv_bins)
+    assert unmixed_binned[key].shape == (2,)
+
+    unmixed_path = tmp_path / "unmixed_nonsquare.fits"
+    heracles.write(unmixed_path, unmixed_binned)
+    unmixed_read = heracles.read(unmixed_path)
+    np.testing.assert_array_equal(unmixed_read[key], unmixed_binned[key])
+    np.testing.assert_array_equal(unmixed_read[key].ell, unmixed_binned[key].ell)