diff --git a/examples/unmixing.ipynb b/examples/unmixing.ipynb index 88111960..aff821f4 100644 --- a/examples/unmixing.ipynb +++ b/examples/unmixing.ipynb @@ -46,7 +46,7 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, "id": "4e85327b", "metadata": {}, "outputs": [], @@ -57,12 +57,29 @@ "# In the future it will not be necessary to trim the mask cls\n", "# to the same lmax\n", "lmax = 1500\n", - "for key in mask_cls.keys():\n", - " mask_cls[key] = heracles.Result(\n", - " mask_cls[key].array[: lmax + 1], axis=mask_cls[key].axis\n", - " )" + "\n", + "# for key in mask_cls.keys():\n", + "# mask_cls[key] = heracles.Result(\n", + "# mask_cls[key].array[: lmax + 1], axis=mask_cls[key].axis\n", + "# )" ] }, + { + "cell_type": "code", + "execution_count": 18, + "id": "1a422093", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "35ba5bad", + "metadata": {}, + "outputs": [], + "source": [] + }, { "cell_type": "markdown", "id": "a7868723", diff --git a/heracles/__init__.py b/heracles/__init__.py index 6e40d2fd..ee0ffbed 100644 --- a/heracles/__init__.py +++ b/heracles/__init__.py @@ -65,6 +65,7 @@ # result "Result", "binned", + "truncated", # twopoint "angular_power_spectra", "debias_cls", @@ -137,6 +138,7 @@ from .result import ( Result, binned, + truncated, ) from .twopoint import ( diff --git a/heracles/result.py b/heracles/result.py index 951434b7..f3f76e74 100644 --- a/heracles/result.py +++ b/heracles/result.py @@ -236,3 +236,62 @@ def norm(a, b): upper=binned_upper, weight=binned_weight, ) + + +def truncated(result, ell_max): + """ + Truncate result arrays at given maximum ell values. + """ + + if isinstance(result, Mapping): + return {key: truncated(value, ell_max) for key, value in result.items()} + + ells = get_result_array(result, "ell") + axes = normalize_result_axis(getattr(result, "axis", None), result, ells) + + if not isinstance(ell_max, tuple): + ell_max = (ell_max,) * len(axes) + if len(ell_max) != len(axes): + raise ValueError("result and ell_max have different number of ell axes") + + md = {} + if result.dtype.metadata: + md.update(result.dtype.metadata) + dt = np.dtype(float, metadata=md) + + out = np.copy(result).view(dt) + result_weight = get_result_array(result, "weight") + + truncated_ell = () + truncated_weight = () + + for axis, ell, w, maxval in zip(axes, ells, result_weight, ell_max): + mask = ell <= maxval + n = np.count_nonzero(mask) + + ell_trunc = ell[mask] + w_trunc = w[mask] + + shape = out.shape[:axis] + (n,) + out.shape[axis + 1 :] + tmp = np.empty(shape, dtype=dt) + + for before in np.ndindex(shape[:axis]): + for after in np.ndindex(shape[axis + 1 :]): + k_in = (*before, mask, *after) + k_out = (*before, slice(None), *after) + tmp[k_out] = out[k_in] + + out = tmp + truncated_ell += (ell_trunc,) + truncated_weight += (w_trunc,) + + if len(axes) == 1: + truncated_ell = truncated_ell[0] + truncated_weight = truncated_weight[0] + + return Result( + out, + ell=truncated_ell, + axis=axes, + weight=truncated_weight, + ) diff --git a/tests/test_result.py b/tests/test_result.py index 3ed34047..b3a0e939 100644 --- a/tests/test_result.py +++ b/tests/test_result.py @@ -200,3 +200,70 @@ def test_binned_metadata(): binned = heracles.binned(result, np.array([0, 1, 2])) assert binned.dtype.metadata == md + + +def trunc1(data, ell_max, axis): + """truncate data over a single axis""" + ell = np.arange(data.shape[axis]) + mask = ell <= ell_max + out = np.take(data, np.where(mask)[0], axis=axis) + return out, ell[mask], np.ones_like(ell[mask]) + + +@pytest.mark.parametrize("ndim,axis", [(1, 0), (2, 0), (3, 1)]) +def test_truncated(ndim, axis, rng): + shape = rng.integers(5, 50, ndim) + lmax = shape[axis] - 1 + ell_max = rng.integers(0, lmax) + + data = heracles.Result(rng.standard_normal(shape), axis=axis) + result = heracles.truncated(data, ell_max) + + trunc_data, trunc_ell, trunc_weight = trunc1(data, ell_max, axis) + + np.testing.assert_array_almost_equal(result, trunc_data) + np.testing.assert_array_equal(result.ell, trunc_ell) + np.testing.assert_array_equal(result.weight, trunc_weight) + + +def test_truncated_2d(rng): + ndim = 3 + axes = (0, 2) + shape = rng.integers(5, 50, ndim) + data = heracles.Result(rng.standard_normal(shape), axis=axes) + + ell_max = tuple(rng.integers(0, shape[ax]) for ax in axes) + result = heracles.truncated(data, ell_max) + + trunc = data.array + for i, axis in enumerate(axes): + trunc, trunc_ell, trunc_weight = trunc1(trunc, ell_max[i], axis) + np.testing.assert_array_equal(result.ell[i], trunc_ell) + np.testing.assert_array_equal(result.weight[i], trunc_weight) + np.testing.assert_array_almost_equal(result, trunc) + + +def test_truncated_mapping(): + result = { + object(): object(), + object(): object(), + object(): object(), + } + ell_max = object() + + with patch("heracles.result.truncated") as mock: + out = heracles.truncated(result, ell_max) + assert mock.call_count == len(out) == len(result) + + for i, key in enumerate(result): + assert mock.call_args_list[i] == call(result[key], ell_max) + assert out[key] is mock.return_value + + +def test_truncated_metadata(): + md = {"test": object()} + result = np.zeros(3, dtype=np.dtype(float, metadata=md)) + assert result.dtype.metadata == md + + truncated = heracles.truncated(result, 1) + assert truncated.dtype.metadata == md