From 8f3eabe146b089aa7e349dff70034871cd256d0e Mon Sep 17 00:00:00 2001 From: Florian Pfaff <6773539+FlorianPfaff@users.noreply.github.com> Date: Mon, 3 Aug 2026 09:53:23 +0200 Subject: [PATCH] Reject masked assignment inputs --- src/pyrecest/utils/assignment.py | 31 +++++++ tests/utils/test_assignment_masked_costs.py | 95 +++++++++++++++++++++ 2 files changed, 126 insertions(+) create mode 100644 tests/utils/test_assignment_masked_costs.py diff --git a/src/pyrecest/utils/assignment.py b/src/pyrecest/utils/assignment.py index 188a70ef02..537c0a628a 100644 --- a/src/pyrecest/utils/assignment.py +++ b/src/pyrecest/utils/assignment.py @@ -33,6 +33,31 @@ _INVALID_SCALAR_TYPES = _BOOLEAN_TYPES + _TEXT_TYPES + _TEMPORAL_TYPES +def _contains_masked_values(value, active_ids: set[int] | None = None) -> bool: + """Return whether a nested input contains genuinely masked NumPy values.""" + if _np.ma.is_masked(value): + return True + if isinstance(value, _np.ndarray): + if value.dtype != object: + return False + items = value.reshape(-1) + elif isinstance(value, (list, tuple)): + items = value + else: + return False + + if active_ids is None: + active_ids = set() + value_id = id(value) + if value_id in active_ids: + return False + active_ids.add(value_id) + try: + return any(_contains_masked_values(item, active_ids) for item in items) + finally: + active_ids.remove(value_id) + + @dataclass(frozen=True) class _MurtySubproblem: """Internal Murty subproblem descriptor.""" @@ -43,6 +68,8 @@ class _MurtySubproblem: def _validate_assignment_count(k: int) -> int: + if _contains_masked_values(k): + raise ValueError("k must be an integer") if isinstance(k, _INVALID_SCALAR_TYPES): raise ValueError("k must be an integer") if isinstance(k, Integral): @@ -163,6 +190,8 @@ def _contains_complex_values(value) -> bool: def _coerce_cost_matrix(cost_matrix): + if _contains_masked_values(cost_matrix): + raise ValueError("cost_matrix must not contain masked values") if _contains_boolean_values(cost_matrix): raise ValueError("cost_matrix must be numeric, not boolean") if _contains_text_values(cost_matrix) or _contains_temporal_values(cost_matrix): @@ -207,6 +236,8 @@ def _coerce_cost_matrix(cost_matrix): def _coerce_non_assignment_costs(costs, size: int, name: str): if costs is None: return _zeros(size, dtype=float) + if _contains_masked_values(costs): + raise ValueError(f"{name} must not contain masked values") if _contains_temporal_values(costs): raise ValueError(f"{name} must be numeric and finite") diff --git a/tests/utils/test_assignment_masked_costs.py b/tests/utils/test_assignment_masked_costs.py new file mode 100644 index 0000000000..9894255efc --- /dev/null +++ b/tests/utils/test_assignment_masked_costs.py @@ -0,0 +1,95 @@ +import unittest + +import numpy as np +import pyrecest.backend +from pyrecest.utils import ( + min_cost_max_cardinality_assignment, + murty_k_best_assignments, +) + + +class AssignmentMaskedCostValidationTest(unittest.TestCase): + @staticmethod + def _solvers(): + return ( + ("murty", lambda matrix: murty_k_best_assignments(matrix, k=1)), + ("max_cardinality", min_cost_max_cardinality_assignment), + ) + + @unittest.skipIf( + pyrecest.backend.__backend_name__ == "jax", # pylint: disable=no-member + reason="Not supported on the JAX backend", + ) + def test_masked_cost_matrix_entries_are_rejected(self): + invalid_matrices = ( + np.ma.array([[0.0, 1.0]], mask=[[False, True]]), + np.array([[0.0, np.ma.masked]], dtype=object), + ) + + for solver_name, solver in self._solvers(): + for matrix in invalid_matrices: + with self.subTest(solver=solver_name, matrix_type=type(matrix).__name__): + with self.assertRaisesRegex( + ValueError, "cost_matrix must not contain masked values" + ): + solver(matrix) + + @unittest.skipIf( + pyrecest.backend.__backend_name__ == "jax", # pylint: disable=no-member + reason="Not supported on the JAX backend", + ) + def test_masked_non_assignment_costs_are_rejected(self): + matrix = np.array([[1.0, 2.0]]) + invalid_costs = ( + {"row_non_assignment_costs": np.ma.array([0.5], mask=[True])}, + { + "col_non_assignment_costs": np.array( + [0.5, np.ma.masked], dtype=object + ) + }, + ) + + for kwargs in invalid_costs: + name = next(iter(kwargs)) + with self.subTest(name=name): + with self.assertRaisesRegex( + ValueError, f"{name} must not contain masked values" + ): + murty_k_best_assignments(matrix, k=1, **kwargs) + + @unittest.skipIf( + pyrecest.backend.__backend_name__ == "jax", # pylint: disable=no-member + reason="Not supported on the JAX backend", + ) + def test_masked_assignment_count_is_rejected(self): + for k in (np.ma.masked, np.ma.array(1, mask=True)): + with self.subTest(k=repr(k)): + with self.assertRaisesRegex(ValueError, "k must be an integer"): + murty_k_best_assignments(np.array([[1.0]]), k=k) + + @unittest.skipIf( + pyrecest.backend.__backend_name__ == "jax", # pylint: disable=no-member + reason="Not supported on the JAX backend", + ) + def test_fully_unmasked_masked_arrays_remain_supported(self): + matrix = np.ma.array([[1.0, 2.0]], mask=False) + for solver_name, solver in self._solvers(): + with self.subTest(solver=solver_name): + result = solver(matrix) + if solver_name == "murty": + self.assertEqual(len(result), 1) + else: + np.testing.assert_array_equal(result["assignment"], np.array([0])) + + solutions = murty_k_best_assignments( + matrix, + k=np.ma.array(1, mask=False), + row_non_assignment_costs=np.ma.array([3.0], mask=False), + col_non_assignment_costs=np.ma.array([0.0, 0.0], mask=False), + ) + self.assertEqual(len(solutions), 1) + np.testing.assert_array_equal(solutions[0]["assignment"], np.array([0])) + + +if __name__ == "__main__": + unittest.main()