From be3c601618a80398d87ea5987a64957a2f12cd27 Mon Sep 17 00:00:00 2001 From: jaimerzp Date: Mon, 9 Mar 2026 15:54:45 +0000 Subject: [PATCH 1/3] working --- heracles/twopoint.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/heracles/twopoint.py b/heracles/twopoint.py index 5f7bd97..401b182 100644 --- a/heracles/twopoint.py +++ b/heracles/twopoint.py @@ -465,7 +465,7 @@ def invert_mixing_matrix( return inv_M -def apply_mixing_matrix(d, M, lmax=None): +def apply_mixing_matrix(d, M): """ Apply mixing matrix to the data Cl. Args: @@ -479,7 +479,6 @@ def apply_mixing_matrix(d, M, lmax=None): if lmax is None: *_, lmax = d[key].shape dtype = d[key].array.dtype - ell_mask = M[key].ell s1, s2 = d[key].spin _d = np.atleast_2d(d[key].array) _M = M[key].array @@ -495,7 +494,5 @@ def apply_mixing_matrix(d, M, lmax=None): _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, ell=ell_mask) - # truncate - corr_d = binned(corr_d, np.arange(0, lmax + 1)) + corr_d[key] = replace(d[key], array=_corr_d) return corr_d From a845afadee6c98da9d99307b284e1030d344f235 Mon Sep 17 00:00:00 2001 From: jaimerzp Date: Mon, 9 Mar 2026 15:56:18 +0000 Subject: [PATCH 2/3] ruff --- heracles/twopoint.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/heracles/twopoint.py b/heracles/twopoint.py index 401b182..637bad5 100644 --- a/heracles/twopoint.py +++ b/heracles/twopoint.py @@ -476,8 +476,6 @@ def apply_mixing_matrix(d, M): """ corr_d = {} for key in d.keys(): - if lmax is None: - *_, lmax = d[key].shape dtype = d[key].array.dtype s1, s2 = d[key].spin _d = np.atleast_2d(d[key].array) From 5f1f55db505b4933272a617c918acc7e98252b36 Mon Sep 17 00:00:00 2001 From: jaimerzp Date: Tue, 10 Mar 2026 11:54:37 +0000 Subject: [PATCH 3/3] bug --- heracles/twopoint.py | 36 +++++++++++++++++-- tests/test_twopoint.py | 82 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 115 insertions(+), 3 deletions(-) diff --git a/heracles/twopoint.py b/heracles/twopoint.py index 637bad5..1472095 100644 --- a/heracles/twopoint.py +++ b/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 diff --git a/tests/test_twopoint.py b/tests/test_twopoint.py index 802771a..50d59bc 100644 --- a/tests/test_twopoint.py +++ b/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)