Single shared mask-condition grammar (masks.py + cosmo_inference)

cosmo_inference/scripts/masking.py

@@ -8,6 +8,8 @@
 import numpy as np
 import yaml
 
+from sp_validation.masks import apply_condition
+
 # -------------------------
 # Spatially-structured cuts: these define the survey footprint.
 # All other cuts (FLAGS, mag, SNR, shape measurement, PSF ellipticity,
@@ -30,29 +32,6 @@
 # Masking logic
 
 
-def apply_condition(array, kind, value):
-    """
-    Apply a logical condition to a NumPy array and return a boolean mask, based
-    on the "kind" key in the mask config YAML file.
-    """
-    if kind == "equal":
-        return array == value
-    elif kind == "not_equal":
-        return array != value
-    elif kind == "greater_equal":
-        return array >= value
-    elif kind == "greater":
-        return array > value
-    elif kind == "less_equal":
-        return array <= value
-    elif kind == "less":
-        return array < value
-    elif kind == "range":
-        return (array >= value[0]) & (array <= value[1])
-    else:
-        raise ValueError(f"Unknown kind: {kind}")
-
-
 def apply_masks(data, data_ext, mask_config, footprint_only=False):
     """
     Construct a boolean mask selecting galaxies that satisfy all

pyproject.toml

@@ -44,7 +44,6 @@ dependencies = [
     "jupytext>=1.15",
     "lenspack",
     "lmfit",
-    "numexpr",
     "numpy>=2.0",
     "opencv-python-headless",
     "pyccl",

src/sp_validation/masks.py

@@ -8,11 +8,53 @@
 """
 
 import healsparse as hsp
-import numexpr as ne
 import numpy as np
 from astropy.io import fits
 from scipy import stats
 
+_KIND_ALIASES = {"smaller_equal": "less_equal"}
+
+_KIND_OPS = {
+    "equal": lambda a, v: a == v,
+    "not_equal": lambda a, v: a != v,
+    "greater": lambda a, v: a > v,
+    "greater_equal": lambda a, v: a >= v,
+    "less": lambda a, v: a < v,
+    "less_equal": lambda a, v: a <= v,
+    "range": lambda a, v: (a >= v[0]) & (a <= v[1]),
+}
+
+
+def apply_condition(array, kind, value):
+    """Apply Condition.
+
+    Evaluate one mask condition (``kind``/``value``, as specified in mask
+    config YAML files) against an array and return a boolean mask. Single
+    shared grammar for the object-selection ``Mask`` class and the
+    cosmo_inference footprint builder (see issue #181).
+
+    Parameters
+    ----------
+    array : numpy.ndarray
+        input data
+    kind : str
+        operation type, one of "equal", "not_equal", "greater",
+        "greater_equal", "less", "less_equal" (alias "smaller_equal"),
+        "range"
+    value : float or list
+        value(s) to be used in mask operation; two-element list for "range"
+
+    Returns
+    -------
+    numpy.ndarray
+        boolean mask
+
+    """
+    kind = _KIND_ALIASES.get(kind, kind)
+    if kind not in _KIND_OPS:
+        raise ValueError(f"Unknown mask condition kind: {kind!r}")
+    return _KIND_OPS[kind](array, value)
+
 
 def correlation_matrix(masks, confidence_level=0.9):
 
@@ -76,8 +118,7 @@ class Mask:
     label : str
         mask label
     kind : str
-        operation type, allowed are "equal", "not_equal, ""greater_equal",
-        "smaller_equal", "range"
+        operation type; see :func:`apply_condition` for the allowed values
     value : float or list
         value(s) to be used in mask operation
     dat : numpy.ndarray, optional
@@ -123,40 +164,7 @@ def from_list(cls, masks, label="combined", verbose=False):
 
     def apply(self, dat):
 
-        # Get column
-        col_data = dat[self._col_name]
-
-        if self._kind == "equal":
-            self._mask = ne.evaluate(
-                "col_data == value",
-                local_dict={"col_data": col_data, "value": self._value},
-            )
-        elif self._kind == "not_equal":
-            self._mask = ne.evaluate(
-                "col_data != value",
-                local_dict={"col_data": col_data, "value": self._value},
-            )
-        elif self._kind == "greater_equal":
-            self._mask = ne.evaluate(
-                "col_data >= value",
-                local_dict={"col_data": col_data, "value": self._value},
-            )
-        elif self._kind == "smaller_equal":
-            self._mask = ne.evaluate(
-                "col_data <= value",
-                local_dict={"col_data": col_data, "value": self._value},
-            )
-        elif self._kind == "range":
-            self._mask = ne.evaluate(
-                "(col_data >= low) & (col_data <= high)",
-                local_dict={
-                    "col_data": col_data,
-                    "low": self._value[0],
-                    "high": self._value[1],
-                },
-            )
-        else:
-            raise ValueError(f"Invalid kind {self._kind}")
+        self._mask = apply_condition(dat[self._col_name], self._kind, self._value)
 
     def to_bool(self, hsp_mask):
 

src/sp_validation/tests/test_masks.py

@@ -0,0 +1,63 @@
+"""TESTS FOR THE SHARED MASK-CONDITION GRAMMAR.
+
+Covers ``sp_validation.masks.apply_condition`` — the single ``kind``/``value``
+evaluator shared by the object-selection ``Mask`` class and the
+cosmo_inference footprint builder (issue #181) — plus an equivalence check
+against ``Mask.apply()`` itself.
+
+:Author: cdaley
+
+"""
+
+import numpy as np
+import numpy.testing as npt
+import pytest
+
+from sp_validation.masks import Mask, apply_condition
+
+pytestmark = pytest.mark.fast
+
+
+_ARRAY = np.array([1, 2, 3, 4, 5])
+
+
+@pytest.mark.parametrize(
+    "kind, value, expected",
+    [
+        ("equal", 3, [False, False, True, False, False]),
+        ("not_equal", 3, [True, True, False, True, True]),
+        ("greater", 3, [False, False, False, True, True]),
+        ("greater_equal", 3, [False, False, True, True, True]),
+        ("less", 3, [True, True, False, False, False]),
+        ("less_equal", 3, [True, True, True, False, False]),
+        ("range", [2, 4], [False, True, True, True, False]),
+    ],
+)
+def test_apply_condition_kinds(kind, value, expected):
+
+    npt.assert_array_equal(apply_condition(_ARRAY, kind, value), expected)
+
+
+def test_smaller_equal_alias_matches_less_equal():
+
+    npt.assert_array_equal(
+        apply_condition(_ARRAY, "smaller_equal", 3),
+        apply_condition(_ARRAY, "less_equal", 3),
+    )
+
+
+def test_unknown_kind_raises():
+
+    with pytest.raises(ValueError):
+        apply_condition(_ARRAY, "not_a_real_kind", 3)
+
+
+def test_mask_apply_matches_apply_condition():
+
+    dat = np.array([(1,), (2,), (3,), (4,), (5,)], dtype=[("col", "i8")])
+
+    my_mask = Mask("col", "test_mask", kind="greater_equal", value=3, dat=dat)
+
+    npt.assert_array_equal(
+        my_mask._mask, apply_condition(dat["col"], "greater_equal", 3)
+    )