gh 319: Truncated

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",

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 (

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,
+    )

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