From 014e009f3a85fe25c5e5411171f870322d495b06 Mon Sep 17 00:00:00 2001 From: jaimerzp Date: Mon, 2 Mar 2026 11:44:32 +0000 Subject: [PATCH 1/5] options for inv_mat --- heracles/twopoint.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/heracles/twopoint.py b/heracles/twopoint.py index d53c792..1e692c5 100644 --- a/heracles/twopoint.py +++ b/heracles/twopoint.py @@ -404,7 +404,7 @@ def mixing_matrices( def invert_mixing_matrix( M, - rtol: float = 1e-5, + rcond: float = 1e-5, progress: Progress | None = None, ): """ @@ -412,7 +412,7 @@ def invert_mixing_matrix( Args: M: Mixing matrix (mapping of keys -> Result objects) - rtol: relative tolerance for pseudo-inverse + rcond: relative tolerance for pseudo-inverse progress: optional progress reporter Returns: @@ -432,14 +432,21 @@ def invert_mixing_matrix( s1, s2 = value.spin *_, _n, _m = _M.shape + if isinstance(rcond, Mapping): + if key not in rcond: + raise KeyError(f"Missing rcond value for wm key: {key}") + _rcond = rcond[key] + else: + _rcond = rcond + with progress.task(f"invert {key}"): if (s1 != 0) and (s2 != 0): # Cl^EE+Cl^BB and Cl^EE-Cl^BB transformation # makes the mixing matrix block-diagonal M_p = _M[0] + _M[1] M_m = _M[0] - _M[1] - inv_M_p = np.linalg.pinv(M_p, rcond=rtol) - inv_M_m = np.linalg.pinv(M_m, rcond=rtol) + inv_M_p = np.linalg.pinv(M_p, rcond=_rcond) + inv_M_m = np.linalg.pinv(M_m, rcond=_rcond) _inv_m = np.vstack( ( np.hstack(((inv_M_p + inv_M_m) / 2, (inv_M_p - inv_M_m) / 2)), @@ -448,10 +455,10 @@ def invert_mixing_matrix( ) _inv_M_EEEE = _inv_m[:_m, :_n] _inv_M_EEBB = _inv_m[_m:, :_n] - _inv_M_EBEB = np.linalg.pinv(_M[2], rcond=rtol) + _inv_M_EBEB = np.linalg.pinv(_M[2], rcond=_rcond) _inv_M = np.array([_inv_M_EEEE, _inv_M_EEBB, _inv_M_EBEB]) else: - _inv_M = np.linalg.pinv(_M, rcond=rtol) + _inv_M = np.linalg.pinv(_M, rcond=_rcond) inv_M[key] = replace(M[key], array=_inv_M) return inv_M From 6129247ff94a12a3cc39214f5d0992cc7f3c8219 Mon Sep 17 00:00:00 2001 From: jaimerzp Date: Mon, 2 Mar 2026 11:47:59 +0000 Subject: [PATCH 2/5] tests --- heracles/twopoint.py | 3 ++- tests/test_twopoint.py | 45 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/heracles/twopoint.py b/heracles/twopoint.py index 1e692c5..5f7bd97 100644 --- a/heracles/twopoint.py +++ b/heracles/twopoint.py @@ -22,6 +22,7 @@ import logging import time +from collections.abc import Mapping from datetime import timedelta from itertools import combinations_with_replacement, product from typing import TYPE_CHECKING, Any @@ -39,7 +40,7 @@ from dataclasses import replace if TYPE_CHECKING: - from collections.abc import Mapping, MutableMapping + from collections.abc import MutableMapping from numpy.typing import ArrayLike, NDArray diff --git a/tests/test_twopoint.py b/tests/test_twopoint.py index 4863a7d..9794459 100644 --- a/tests/test_twopoint.py +++ b/tests/test_twopoint.py @@ -411,3 +411,48 @@ 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_invert_mixing_matrix_accepts_per_key_rcond(): + from heracles.twopoint import Result, invert_mixing_matrix + + key_spin0 = ("POS", "POS", 0, 0) + key_spin2 = ("SHE", "SHE", 0, 0) + + mms = { + key_spin0: Result(np.eye(2), spin=(0, 0), axis=(0,)), + key_spin2: Result( + np.array([np.eye(2), np.zeros((2, 2)), np.eye(2)]), + spin=(2, 2), + axis=(1,), + ), + } + rcond = { + key_spin0: 1e-2, + key_spin2: 1e-4, + } + + orig_pinv = np.linalg.pinv + pinv_calls = [] + + def spy_pinv(a, rcond): + pinv_calls.append(rcond) + return orig_pinv(a, rcond=rcond) + + with patch("numpy.linalg.pinv", side_effect=spy_pinv): + inv_mms = invert_mixing_matrix(mms, rcond=rcond) + + assert inv_mms.keys() == mms.keys() + assert pinv_calls == [1e-2, 1e-4, 1e-4, 1e-4] + + +def test_invert_mixing_matrix_errors_if_rcond_mapping_missing_key(): + from heracles.twopoint import Result, invert_mixing_matrix + + key = ("POS", "POS", 0, 0) + mms = { + key: Result(np.eye(2), spin=(0, 0), axis=(0,)), + } + + with pytest.raises(KeyError, match="Missing rcond value"): + invert_mixing_matrix(mms, rcond={}) From 183228a6e1cf1635772812688357e99b2a530000 Mon Sep 17 00:00:00 2001 From: jaimerzp Date: Mon, 2 Mar 2026 11:57:35 +0000 Subject: [PATCH 3/5] compact testing --- tests/test_twopoint.py | 50 +++--------------------------------------- 1 file changed, 3 insertions(+), 47 deletions(-) diff --git a/tests/test_twopoint.py b/tests/test_twopoint.py index 9794459..3b8b6d1 100644 --- a/tests/test_twopoint.py +++ b/tests/test_twopoint.py @@ -368,7 +368,7 @@ def test_inverting_mixing_matrices(): ("WHT", "WHT", 0, 0): Result(cl, spin=(0, 0), axis=(0,)), } mms = mixing_matrices(fields, cls2, l1max=10, l2max=20) - inv_mms = invert_mixing_matrix(mms) + inv_mms = invert_mixing_matrix(mms, rcond=1e-2) # test for correct shape for key in mms.keys(): @@ -387,7 +387,8 @@ def test_inverting_mixing_matrices(): _m = np.ones_like(mms[key].array) mms[key] = Result(_m, spin=mms[key].spin, axis=mms[key].axis, ell=mms[key].ell) - inv_mms = invert_mixing_matrix(mms) + inv_mms = invert_mixing_matrix(mms, + rcond={("POS", "POS", 0, 0): 1e-2, ("POS", "SHE", 0, 0): 1e-3, ("SHE", "SHE", 0, 0): 1e-4}) assert inv_mms.keys() == mms.keys() for key in mms: inv_mm = inv_mms[key].array @@ -411,48 +412,3 @@ 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_invert_mixing_matrix_accepts_per_key_rcond(): - from heracles.twopoint import Result, invert_mixing_matrix - - key_spin0 = ("POS", "POS", 0, 0) - key_spin2 = ("SHE", "SHE", 0, 0) - - mms = { - key_spin0: Result(np.eye(2), spin=(0, 0), axis=(0,)), - key_spin2: Result( - np.array([np.eye(2), np.zeros((2, 2)), np.eye(2)]), - spin=(2, 2), - axis=(1,), - ), - } - rcond = { - key_spin0: 1e-2, - key_spin2: 1e-4, - } - - orig_pinv = np.linalg.pinv - pinv_calls = [] - - def spy_pinv(a, rcond): - pinv_calls.append(rcond) - return orig_pinv(a, rcond=rcond) - - with patch("numpy.linalg.pinv", side_effect=spy_pinv): - inv_mms = invert_mixing_matrix(mms, rcond=rcond) - - assert inv_mms.keys() == mms.keys() - assert pinv_calls == [1e-2, 1e-4, 1e-4, 1e-4] - - -def test_invert_mixing_matrix_errors_if_rcond_mapping_missing_key(): - from heracles.twopoint import Result, invert_mixing_matrix - - key = ("POS", "POS", 0, 0) - mms = { - key: Result(np.eye(2), spin=(0, 0), axis=(0,)), - } - - with pytest.raises(KeyError, match="Missing rcond value"): - invert_mixing_matrix(mms, rcond={}) From 74c5b11775818baccb1683e1ea8fe77411d28eeb Mon Sep 17 00:00:00 2001 From: jaimerzp Date: Mon, 2 Mar 2026 11:59:13 +0000 Subject: [PATCH 4/5] ruff --- tests/test_twopoint.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/test_twopoint.py b/tests/test_twopoint.py index 3b8b6d1..802771a 100644 --- a/tests/test_twopoint.py +++ b/tests/test_twopoint.py @@ -387,8 +387,14 @@ def test_inverting_mixing_matrices(): _m = np.ones_like(mms[key].array) mms[key] = Result(_m, spin=mms[key].spin, axis=mms[key].axis, ell=mms[key].ell) - inv_mms = invert_mixing_matrix(mms, - rcond={("POS", "POS", 0, 0): 1e-2, ("POS", "SHE", 0, 0): 1e-3, ("SHE", "SHE", 0, 0): 1e-4}) + inv_mms = invert_mixing_matrix( + mms, + rcond={ + ("POS", "POS", 0, 0): 1e-2, + ("POS", "SHE", 0, 0): 1e-3, + ("SHE", "SHE", 0, 0): 1e-4, + }, + ) assert inv_mms.keys() == mms.keys() for key in mms: inv_mm = inv_mms[key].array From b41669fd96e32cd2ba59aa291b7fa3837171d04b Mon Sep 17 00:00:00 2001 From: jaimerzp Date: Mon, 9 Mar 2026 15:52:22 +0000 Subject: [PATCH 5/5] 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