mass-screen-rs 0.1.0__tar.gz

mass_screen_rs-0.1.0/.cargo/config.toml

@@ -0,0 +1,12 @@
+# See https://pyo3.rs/latest/building-and-distribution.html#macos
+[target.x86_64-apple-darwin]
+rustflags = [
+  "-C", "link-arg=-undefined",
+  "-C", "link-arg=dynamic_lookup",
+]
+
+[target.aarch64-apple-darwin]
+rustflags = [
+  "-C", "link-arg=-undefined",
+  "-C", "link-arg=dynamic_lookup",
+]

mass_screen_rs-0.1.0/Cargo.lock

@@ -0,0 +1,280 @@
+# This file is automatically @generated by Cargo.
+# It is not intended for manual editing.
+version = 4
+
+[[package]]
+name = "autocfg"
+version = "1.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53"
+
+[[package]]
+name = "crossbeam-deque"
+version = "0.8.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "622f3fc73690be383c7214310406f28a90e6edeadc3cea882f9d71e495b9711a"
+dependencies = [
+ "crossbeam-epoch",
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "crossbeam-epoch"
+version = "0.9.21"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dc74980687109a3b14c72fd458107bf0baa1da1a1a805e178d15501ba9b86d9d"
+dependencies = [
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "crossbeam-utils"
+version = "0.8.23"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a31eee39dddec8330830986fcd7625edb5a24ec90ea038215273bbc3adb08ac6"
+
+[[package]]
+name = "either"
+version = "1.19.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0e9c71c2167ca323c882b99918929403426e2373ea17242ff5653e0d5e1058be"
+
+[[package]]
+name = "heck"
+version = "0.5.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea"
+
+[[package]]
+name = "libc"
+version = "0.2.191"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "65cd87c9a6d60536580b316ad985726104730a5151c99845f2490428ef91240a"
+
+[[package]]
+name = "mass-screen-rs"
+version = "0.1.0"
+dependencies = [
+ "numpy",
+ "pyo3",
+ "rayon",
+]
+
+[[package]]
+name = "matrixmultiply"
+version = "0.3.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3f607c237553f086e7043417a51df26b2eb899d3caff94e6a67592ff992fedc7"
+dependencies = [
+ "autocfg",
+ "rawpointer",
+]
+
+[[package]]
+name = "ndarray"
+version = "0.17.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d"
+dependencies = [
+ "matrixmultiply",
+ "num-complex",
+ "num-integer",
+ "num-traits",
+ "portable-atomic",
+ "portable-atomic-util",
+ "rawpointer",
+]
+
+[[package]]
+name = "num-complex"
+version = "0.4.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495"
+dependencies = [
+ "num-traits",
+]
+
+[[package]]
+name = "num-integer"
+version = "0.1.47"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7ce2d95d4b3734dc35aa2f45e1aa22cd416814592a4f9d9205e11affd5b8e10b"
+dependencies = [
+ "num-traits",
+]
+
+[[package]]
+name = "num-traits"
+version = "0.2.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
+dependencies = [
+ "autocfg",
+]
+
+[[package]]
+name = "numpy"
+version = "0.30.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "aa0b9111761885449e894981fa8051dcf69e0bd114aecbd4bcdd830b6be4ed92"
+dependencies = [
+ "libc",
+ "ndarray",
+ "num-complex",
+ "num-integer",
+ "num-traits",
+ "pyo3",
+ "pyo3-build-config",
+ "rustc-hash",
+]
+
+[[package]]
+name = "once_cell"
+version = "1.21.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50"
+
+[[package]]
+name = "portable-atomic"
+version = "1.15.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85"
+
+[[package]]
+name = "portable-atomic-util"
+version = "0.2.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "10ab3eb7f3becc3a1cbc4f2c6f20267996cfc1a6467a873763411b136a122715"
+dependencies = [
+ "portable-atomic",
+]
+
+[[package]]
+name = "proc-macro2"
+version = "1.0.107"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9"
+dependencies = [
+ "unicode-ident",
+]
+
+[[package]]
+name = "pyo3"
+version = "0.30.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "229ca13b8a15bb0d1487b81ab975bfcd97513416499a8312011b6036e42e0de4"
+dependencies = [
+ "libc",
+ "once_cell",
+ "portable-atomic",
+ "pyo3-build-config",
+ "pyo3-ffi",
+ "pyo3-macros",
+]
+
+[[package]]
+name = "pyo3-build-config"
+version = "0.30.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "70130e2b6dd1ee589a8788026a429561ca5c40fa7a2fa7914c2282eddae24700"
+dependencies = [
+ "target-lexicon",
+]
+
+[[package]]
+name = "pyo3-ffi"
+version = "0.30.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5e7ff73c555c286d212d4ce7e5fbda845a973c4f61bd5f160f5f2c3c2a04d1c4"
+dependencies = [
+ "libc",
+ "pyo3-build-config",
+]
+
+[[package]]
+name = "pyo3-macros"
+version = "0.30.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c46d2e5862c7bb8207e3b639f713dda3c499933444784154c11d013c5a691b60"
+dependencies = [
+ "proc-macro2",
+ "pyo3-macros-backend",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "pyo3-macros-backend"
+version = "0.30.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "38b89f066f2ae72a2aef52d1537ca2fcd104440ebfadee06e97aa90738a25b9b"
+dependencies = [
+ "heck",
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
+[[package]]
+name = "quote"
+version = "1.0.47"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001"
+dependencies = [
+ "proc-macro2",
+]
+
+[[package]]
+name = "rawpointer"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"
+
+[[package]]
+name = "rayon"
+version = "1.12.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d"
+dependencies = [
+ "either",
+ "rayon-core",
+]
+
+[[package]]
+name = "rayon-core"
+version = "1.13.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91"
+dependencies = [
+ "crossbeam-deque",
+ "crossbeam-utils",
+]
+
+[[package]]
+name = "rustc-hash"
+version = "2.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d"
+
+[[package]]
+name = "syn"
+version = "3.0.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d78c8dee4c7bf0e14673097256fed6142ce9d3b85a408189d07482442145823b"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "unicode-ident",
+]
+
+[[package]]
+name = "target-lexicon"
+version = "0.13.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca"
+
+[[package]]
+name = "unicode-ident"
+version = "1.0.27"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a2c754d6c33795a1c324727428e5a7dedb5b06195f9890bdbcba760d3e246563"

mass_screen_rs-0.1.0/Cargo.toml

@@ -0,0 +1,22 @@
+[package]
+name = "mass-screen-rs"
+version = "0.1.0"
+edition = "2021"
+description = "Rust backend for the MUSE data-driven response mask (mass screen)"
+license = "BSD-3-Clause"
+repository = "https://github.com/LM-SAL/mass-screen-rs"
+readme = "README.md"
+
+[lib]
+name = "mass_screen_rs"
+crate-type = ["cdylib"]
+
+[dependencies]
+numpy = "0.30"
+pyo3 = { version = "0.30", features = ["extension-module"] }
+rayon = "1"
+
+[profile.release]
+lto = true
+codegen-units = 1
+strip = "symbols"

mass_screen_rs-0.1.0/LICENSE

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2026, The MUSE Instrument Team @ LMSAL
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

mass_screen_rs-0.1.0/PKG-INFO

@@ -0,0 +1,62 @@
+Metadata-Version: 2.4
+Name: mass-screen-rs
+Version: 0.1.0
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Rust
+Classifier: Topic :: Scientific/Engineering :: Astronomy
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Dist: numpy>=2.0.0
+Requires-Dist: mass-screen-rs[tests,docs] ; extra == 'dev'
+Requires-Dist: sphinx>=8 ; extra == 'docs'
+Requires-Dist: furo ; extra == 'docs'
+Requires-Dist: sphinx-copybutton ; extra == 'docs'
+Requires-Dist: pytest ; extra == 'tests'
+Provides-Extra: dev
+Provides-Extra: docs
+Provides-Extra: tests
+License-File: LICENSE
+Summary: Rust backend for the MUSE data-driven response mask (mass screen)
+Author-email: Nabil Freij <nabil.freij@gmail.com>
+License-Expression: BSD-3-Clause
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
+Project-URL: Documentation, https://mass-screen-rs.readthedocs.io
+Project-URL: Issues, https://github.com/LM-SAL/mass-screen-rs/issues
+Project-URL: Source, https://github.com/LM-SAL/mass-screen-rs
+
+# mass-screen-rs
+
+The MUSE data-driven response mask ("mass screen") with a Rust core and a Python API.
+
+`mass_screen(response, spectra, threshold=..., fill_value=...)` takes the response matrix
+`[n, m]` (spectral pixel × VDEM component) and the spectra `[n, b]` (spectral pixel × image
+pixel) and returns a float32 array `[b, m]` with 1.0 where a component can produce at least
+`threshold` of signal for that pixel and NaN where it cannot, after replacing spectra `<= 0` by
+`fill_value`. It computes the same mask as the PyTorch reference in `muse.inversions` and
+replaces the C++ kernel muse used to build at install time. The tests check it against a recorded
+PyTorch mask.
+
+## Install
+
+```bash
+pip install mass-screen-rs
+```
+
+Wheels are published for Linux (x86-64), macOS (arm64) and Windows, Python 3.12 to 3.15.
+
+## Development
+
+```bash
+pip install maturin pytest
+maturin develop --release        # build into the active environment
+cargo test                       # Rust unit tests
+pytest --pyargs mass_screen_rs   # Python tests, including parity with the recorded PyTorch mask
+prek run --all-files             # lint hooks (ruff, docformatter, cargo fmt, clippy, cargo test)
+```
+
+With `torch` installed, the parity test also checks a fresh random case against PyTorch directly,
+and `python -m mass_screen_rs.tests.make_torch_reference` re-records the reference.
+
+Releases: bump `version` in `Cargo.toml` (the only place it is set), merge, then tag `vX.Y.Z` on
+`main`; the workflow builds the wheels and the sdist and uploads them to PyPI.
+

mass_screen_rs-0.1.0/README.md

@@ -0,0 +1,35 @@
+# mass-screen-rs
+
+The MUSE data-driven response mask ("mass screen") with a Rust core and a Python API.
+
+`mass_screen(response, spectra, threshold=..., fill_value=...)` takes the response matrix
+`[n, m]` (spectral pixel × VDEM component) and the spectra `[n, b]` (spectral pixel × image
+pixel) and returns a float32 array `[b, m]` with 1.0 where a component can produce at least
+`threshold` of signal for that pixel and NaN where it cannot, after replacing spectra `<= 0` by
+`fill_value`. It computes the same mask as the PyTorch reference in `muse.inversions` and
+replaces the C++ kernel muse used to build at install time. The tests check it against a recorded
+PyTorch mask.
+
+## Install
+
+```bash
+pip install mass-screen-rs
+```
+
+Wheels are published for Linux (x86-64), macOS (arm64) and Windows, Python 3.12 to 3.15.
+
+## Development
+
+```bash
+pip install maturin pytest
+maturin develop --release        # build into the active environment
+cargo test                       # Rust unit tests
+pytest --pyargs mass_screen_rs   # Python tests, including parity with the recorded PyTorch mask
+prek run --all-files             # lint hooks (ruff, docformatter, cargo fmt, clippy, cargo test)
+```
+
+With `torch` installed, the parity test also checks a fresh random case against PyTorch directly,
+and `python -m mass_screen_rs.tests.make_torch_reference` re-records the reference.
+
+Releases: bump `version` in `Cargo.toml` (the only place it is set), merge, then tag `vX.Y.Z` on
+`main`; the workflow builds the wheels and the sdist and uploads them to PyPI.

mass_screen_rs-0.1.0/pyproject.toml

@@ -0,0 +1,78 @@
+[build-system]
+requires = [
+  "maturin>=1.8,<2.0",
+]
+build-backend = "maturin"
+
+[project]
+name = "mass-screen-rs"
+# The version lives in Cargo.toml only; maturin reads it from there.
+dynamic = ["version"]
+description = "Rust backend for the MUSE data-driven response mask (mass screen)"
+readme = "README.md"
+license = "BSD-3-Clause"
+license-files = ["LICENSE"]
+authors = [
+  { name = "Nabil Freij", email = "nabil.freij@gmail.com" },
+]
+requires-python = ">=3.12"
+dependencies = [
+  "numpy>=2.0.0",
+]
+classifiers = [
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Rust",
+  "Topic :: Scientific/Engineering :: Astronomy",
+  "Topic :: Scientific/Engineering :: Physics",
+]
+
+[project.optional-dependencies]
+docs = [
+  "sphinx>=8",
+  "furo",
+  "sphinx-copybutton",
+]
+tests = [
+  "pytest",
+]
+dev = [
+  "mass-screen-rs[tests,docs]",
+]
+
+[project.urls]
+Source = "https://github.com/LM-SAL/mass-screen-rs"
+Issues = "https://github.com/LM-SAL/mass-screen-rs/issues"
+Documentation = "https://mass-screen-rs.readthedocs.io"
+
+[tool.maturin]
+python-source = "python"
+module-name = "mass_screen_rs._mass_screen_rs"
+exclude = [
+  ".github/**",
+  "docs/**",
+  ".readthedocs.yaml",
+  ".gitignore",
+  ".pre-commit-config.yaml",
+  ".clippy.toml",
+  "rustfmt.toml",
+  "ruff.toml",
+  "tox.ini",
+]
+
+[tool.pytest.ini_options]
+testpaths = ["python/mass_screen_rs/tests"]
+
+[tool.cibuildwheel]
+build-frontend = "build"
+skip = ["pp*", "*-win32", "*-manylinux_i686"]
+test-extras = ["tests"]
+test-command = "pytest --pyargs mass_screen_rs"
+
+[tool.docformatter]
+recursive = true
+wrap-summaries = 100
+wrap-descriptions = 100
+pre-summary-newline = true
+close-quotes-on-newline = true
+make-summary-multi-line = true
+black = true

mass_screen_rs-0.1.0/python/mass_screen_rs/__init__.py

@@ -0,0 +1,11 @@
+"""
+The MUSE data-driven response mask (mass screen), in Rust.
+"""
+
+from importlib.metadata import version as _get_version
+
+from .screen import mass_screen
+
+__version__ = _get_version("mass-screen-rs")
+
+__all__ = ["__version__", "mass_screen"]

mass_screen_rs-0.1.0/python/mass_screen_rs/screen.py

@@ -0,0 +1,55 @@
+"""
+Python API for the mass-screen-rs Rust backend.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from ._mass_screen_rs import mass_screen as _mass_screen
+
+if TYPE_CHECKING:
+    from numpy.typing import ArrayLike, NDArray
+
+__all__ = ["mass_screen"]
+
+
+def mass_screen(
+    response: ArrayLike, spectra: ArrayLike, *, threshold: float, fill_value: float
+) -> NDArray[np.float32]:
+    """
+    Mask the VDEM components a set of spectra cannot constrain.
+
+    Component ``m`` is kept for image pixel ``b`` when
+    ``min_n(spectra[n, b] / response[n, m]) * max_n(response[n, m]) >= threshold``,
+    after replacing ``spectra <= 0`` by ``fill_value``. A component whose response is
+    nowhere positive is always masked.
+
+    Parameters
+    ----------
+    response : array_like
+        Response matrix of shape ``[n, m]``: spectral pixel × VDEM component.
+    spectra : array_like
+        Spectra of shape ``[n, b]``: spectral pixel × image pixel.
+    threshold : float
+        Smallest signal, in the units of ``spectra``, a component must be able to produce.
+    fill_value : float
+        Replaces the ``spectra`` values that are ``<= 0``; must be ``>= 0``.
+
+    Returns
+    -------
+    numpy.ndarray
+        float32 array of shape ``[b, m]``: 1.0 where a component is kept, NaN where masked.
+
+    Notes
+    -----
+    The kernel runs on all cores through rayon; ``RAYON_NUM_THREADS`` caps it.
+    """
+    response = np.ascontiguousarray(response, dtype=np.float32)
+    spectra = np.ascontiguousarray(spectra, dtype=np.float32)
+    if (response.ndim, spectra.ndim) != (2, 2):
+        msg = f"response and spectra must be 2-D, got shapes {response.shape} and {spectra.shape}"
+        raise ValueError(msg)
+    return _mass_screen(response, spectra, float(threshold), float(fill_value))

mass_screen_rs-0.1.0/python/mass_screen_rs/tests/make_torch_reference.py

@@ -0,0 +1,63 @@
+"""
+Record the PyTorch reference mask for the cases ``test_torch_parity.py`` checks.
+
+Run with torch installed: ``python -m mass_screen_rs.tests.make_torch_reference``.
+"""
+
+from pathlib import Path
+
+import numpy as np
+
+DATA = Path(__file__).parent / "data" / "torch_reference.npz"
+# name: (n, m, b, density, threshold, fill_value, seed). The first case spans several column blocks
+# and pixel tiles of the kernel; the second is small and dense.
+CASES = {
+    "blocks_and_tiles": (24, 2 * 8192 + 5, 3 * 64 + 9, 0.03, 2.24, 0.25, 1),
+    "small_dense": (40, 300, 50, 0.3, 0.5, 1.0, 2),
+}
+
+
+def torch_mask(response, spectra, threshold, fill_value):
+    """
+    Return the mask of ``muse.inversions.inversions._mass_screen_eager``, computed the same way.
+    """
+    import torch  # noqa: PLC0415 - only the recording and the live test need torch
+
+    rg = torch.from_numpy(response)
+    ys = torch.where(torch.from_numpy(spectra) <= 0, fill_value, torch.from_numpy(spectra))
+    max_bnd = torch.amin(ys[:, :, None] / rg[:, None, :], dim=0, keepdims=True)
+    keep = torch.any(rg[:, None, :] * max_bnd >= threshold, dim=0)
+    return torch.where(keep, 1, float("nan")).numpy()
+
+
+def make_case(n, m, b, density, seed):
+    """
+    Build a random case.
+
+    The response is sparse with blocks of zero and of tiny components; about a quarter of the
+    spectra are ``<= 0``.
+    """
+    rng = np.random.default_rng(seed)
+    response = np.where(rng.random((n, m)) < density, rng.random((n, m)), 0).astype(np.float32)
+    response[:, : m // 10] = 0
+    response[:, m // 10 : m // 5] *= np.float32(1e-10)
+    spectra = (rng.random((n, b)) * 4 - 1).astype(np.float32)
+    return response, spectra
+
+
+def main():
+    arrays = {}
+    for name, (n, m, b, density, threshold, fill_value, seed) in CASES.items():
+        response, spectra = make_case(n, m, b, density, seed)
+        kept = torch_mask(response, spectra, threshold, fill_value) == 1
+        arrays[f"{name}/response"] = response
+        arrays[f"{name}/spectra"] = spectra
+        arrays[f"{name}/kept"] = kept
+        arrays[f"{name}/params"] = np.array([threshold, fill_value], dtype=np.float32)
+        print(f"{name}: {response.shape}, {spectra.shape}, kept {kept.mean() * 100:.1f} %")  # noqa: T201
+    np.savez_compressed(DATA, **arrays)
+    print(f"wrote {DATA} ({DATA.stat().st_size / 1024:.0f} KiB)")  # noqa: T201
+
+
+if __name__ == "__main__":
+    main()

mass_screen_rs-0.1.0/python/mass_screen_rs/tests/test_mass_screen.py

@@ -0,0 +1,95 @@
+"""
+Tests for mass_screen.
+"""
+
+import numpy as np
+import pytest
+
+from mass_screen_rs import mass_screen
+
+
+def reference(response, spectra, threshold, fill_value):
+    """
+    Evaluate the mask rule in numpy, one cell at a time.
+    """
+    response = np.asarray(response, dtype=np.float32)
+    spectra = np.asarray(spectra, dtype=np.float32)
+    spectra = np.where(spectra <= 0, np.float32(fill_value), spectra)
+    col_max = response.max(axis=0)
+    scale = np.divide(np.float32(threshold), col_max, out=np.zeros_like(col_max), where=col_max > 0)
+    rhs = scale * response  # [n, m], same float32 arithmetic as the kernel
+    kept = np.all(spectra[:, :, None] >= rhs[:, None, :], axis=0) & (col_max > 0)
+    return np.where(kept, np.float32(1), np.float32(np.nan))
+
+
+def random_case(rng, n, m, b, density=0.05):
+    response = np.where(rng.random((n, m)) < density, rng.random((n, m)), 0).astype(np.float32)
+    spectra = (rng.random((n, b)) * 4 - 1).astype(np.float32)  # about a quarter <= 0
+    return response, spectra
+
+
+def test_muse_example():
+    response = np.array([[1, 0.5, 0.1], [0.2, 1, 0.3]], dtype=np.float32)
+    spectra = np.array([[3, 1], [4, 5]], dtype=np.float32)
+    keep = mass_screen(response, spectra, threshold=2.24, fill_value=1)
+    assert keep.dtype == np.float32
+    np.testing.assert_array_equal(keep, [[1, 1, 1], [np.nan, np.nan, 1]])
+
+
+def test_zero_response_component_is_masked():
+    response = np.array([[1, 0, 0.1], [0.5, 0, 0.3]], dtype=np.float32)
+    spectra = np.array([[3, 4], [4, 5]], dtype=np.float32)
+    keep = mass_screen(response, spectra, threshold=2.24, fill_value=1)
+    np.testing.assert_array_equal(keep[:, 1], [np.nan, np.nan])
+    np.testing.assert_array_equal(keep[:, [0, 2]], 1)
+
+
+def test_fill_value_replaces_non_positive_spectra():
+    response = np.array([[1.0], [1.0]], dtype=np.float32)
+    spectra = np.array([[-5.0, 0.0], [3.0, 3.0]], dtype=np.float32)
+    kept = mass_screen(response, spectra, threshold=2.0, fill_value=2.5)
+    masked = mass_screen(response, spectra, threshold=2.0, fill_value=0.25)
+    np.testing.assert_array_equal(kept, 1)
+    np.testing.assert_array_equal(masked, np.nan)
+
+
+@pytest.mark.parametrize(
+    ("n", "m", "b"), [(5, 7, 3), (40, 300, 50), (12, 8192 * 2 + 5, 64 * 3 + 9)]
+)
+def test_matches_reference(n, m, b):
+    rng = np.random.default_rng(n * m + b)
+    response, spectra = random_case(rng, n, m, b)
+    keep = mass_screen(response, spectra, threshold=2.24, fill_value=0.25)
+    assert keep.shape == (b, m)
+    np.testing.assert_array_equal(keep, reference(response, spectra, 2.24, 0.25))
+
+
+def test_accepts_float64_and_non_contiguous_input():
+    rng = np.random.default_rng(0)
+    response, spectra = random_case(rng, 10, 20, 6)
+    keep = mass_screen(
+        response.astype(np.float64), spectra.T.copy().T, threshold=2.24, fill_value=1
+    )
+    np.testing.assert_array_equal(keep, reference(response, spectra, 2.24, 1))
+
+
+@pytest.mark.parametrize(("m", "b"), [(0, 3), (4, 0)])
+def test_empty(m, b):
+    keep = mass_screen(
+        np.ones((2, m), np.float32), np.ones((2, b), np.float32), threshold=2.24, fill_value=1
+    )
+    assert keep.shape == (b, m)
+
+
+@pytest.mark.parametrize(
+    ("response", "spectra", "kwargs", "match"),
+    [
+        (np.ones((2, 3)), np.ones((3, 2)), {"threshold": 1, "fill_value": 1}, "rows"),
+        (np.ones((2, 3)), np.ones((2, 2)), {"threshold": 1, "fill_value": -1}, "fill_value"),
+        (np.ones((2, 3)), np.ones((2, 2)), {"threshold": np.nan, "fill_value": 1}, "threshold"),
+        (np.ones(3), np.ones((2, 2)), {"threshold": 1, "fill_value": 1}, "2-D"),
+    ],
+)
+def test_rejects_bad_input(response, spectra, kwargs, match):
+    with pytest.raises(ValueError, match=match):
+        mass_screen(response, spectra, **kwargs)

mass_screen_rs-0.1.0/python/mass_screen_rs/tests/test_torch_parity.py

@@ -0,0 +1,31 @@
+"""
+The kernel reproduces the PyTorch reference mask of ``muse.inversions``.
+"""
+
+import numpy as np
+import pytest
+
+from mass_screen_rs import mass_screen
+from mass_screen_rs.tests.make_torch_reference import CASES, DATA, make_case, torch_mask
+
+
+@pytest.mark.parametrize("name", list(CASES))
+def test_matches_recorded_torch_mask(name):
+    with np.load(DATA) as data:
+        response, spectra = data[f"{name}/response"], data[f"{name}/spectra"]
+        kept, (threshold, fill_value) = data[f"{name}/kept"], data[f"{name}/params"]
+
+    keep = mass_screen(response, spectra, threshold=threshold, fill_value=fill_value)
+
+    assert keep.shape == kept.shape
+    np.testing.assert_array_equal(keep == 1, kept)
+    np.testing.assert_array_equal(np.isnan(keep), ~kept)
+
+
+def test_matches_torch_live():
+    pytest.importorskip("torch")
+    response, spectra = make_case(30, 3000, 70, 0.05, seed=3)
+
+    keep = mass_screen(response, spectra, threshold=2.24, fill_value=0.25)
+
+    np.testing.assert_array_equal(keep, torch_mask(response, spectra, 2.24, 0.25))

mass_screen_rs-0.1.0/src/api.rs

@@ -0,0 +1,50 @@
+//! Python entry point: validation, the output array, and the GIL release around the kernel.
+use numpy::{PyArray2, PyArrayMethods, PyReadonlyArray2, PyUntypedArrayMethods};
+use pyo3::exceptions::PyValueError;
+use pyo3::prelude::{pyfunction, Bound, PyResult, Python};
+
+use crate::screen::mass_screen_core;
+
+/// Mask the VDEM components a set of spectra cannot constrain.
+///
+/// `response` is `[n, m]` and `spectra` is `[n, b]`, both C-contiguous float32. Returns a float32
+/// array of shape `[b, m]` holding 1.0 where a component is kept and NaN where it is masked:
+/// component `m` is kept for pixel `b` when `min_n(spectra[n, b] / response[n, m]) *
+/// max_n(response[n, m]) >= threshold`, after replacing `spectra <= 0` by `fill_value`. A
+/// component whose response is nowhere positive is always masked.
+#[pyfunction]
+#[pyo3(signature = (response, spectra, threshold, fill_value))]
+pub(crate) fn mass_screen<'py>(
+    py: Python<'py>,
+    response: PyReadonlyArray2<'py, f32>,
+    spectra: PyReadonlyArray2<'py, f32>,
+    threshold: f32,
+    fill_value: f32,
+) -> PyResult<Bound<'py, PyArray2<f32>>> {
+    if !threshold.is_finite() {
+        return Err(PyValueError::new_err("threshold must be finite"));
+    }
+    if !fill_value.is_finite() || fill_value < 0.0 {
+        return Err(PyValueError::new_err("fill_value must be finite and >= 0"));
+    }
+    let [n, m] = [response.shape()[0], response.shape()[1]];
+    let [n_ys, b] = [spectra.shape()[0], spectra.shape()[1]];
+    if n_ys != n {
+        return Err(PyValueError::new_err(format!(
+            "response has {n} rows but spectra has {n_ys}; both are indexed by spectral pixel"
+        )));
+    }
+    let rg = response
+        .as_slice()
+        .map_err(|_| PyValueError::new_err("response must be a C-contiguous float32 array"))?;
+    let ys = spectra
+        .as_slice()
+        .map_err(|_| PyValueError::new_err("spectra must be a C-contiguous float32 array"))?;
+    // SAFETY: the array is uninitialised until mass_screen_core has written every element,
+    // and nothing reads it before that.
+    let out = unsafe { PyArray2::<f32>::new(py, [b, m], false) };
+    let keep = unsafe { out.as_slice_mut() }
+        .map_err(|_| PyValueError::new_err("could not allocate the mask"))?;
+    py.detach(|| mass_screen_core(rg, ys, keep, n, m, b, threshold, fill_value));
+    Ok(out)
+}

mass_screen_rs-0.1.0/src/lib.rs

@@ -0,0 +1,31 @@
+//! Data-driven response mask for the MUSE inversions.
+//!
+//! For a response matrix `rg` of shape `[n, m]` (spectral pixel × VDEM component) and spectra `ys`
+//! of shape `[n, b]` (spectral pixel × image pixel), the mask is
+//!
+//! ```text
+//! keep[b, m] = 1    if  min_n(ys[n, b] / rg[n, m]) * max_n(rg[n, m]) >= threshold
+//!            = NaN  otherwise,
+//! ```
+//!
+//! with `ys <= 0` replaced by `fill_value` first. Rewritten as `ys[n, b] >= rhs[n, m]` for every
+//! `n`, with `rhs = threshold / max_n(rg) * rg`, only the entries with `rhs > 0` can mask a cell,
+//! and a MUSE response is a few percent non-zero. The kernel keeps those entries per row, orders the
+//! rows so the ones that mask most run first, and updates the flags of 64 image pixels per entry
+//! with one vector compare.
+mod api;
+mod screen;
+
+#[cfg(test)]
+mod tests;
+
+use pyo3::prelude::{pymodule, Bound, PyModule, PyResult, Python};
+use pyo3::types::PyModuleMethods;
+use pyo3::wrap_pyfunction;
+
+#[pymodule]
+#[pyo3(name = "_mass_screen_rs")]
+fn mass_screen_rs(_py: Python<'_>, module: &Bound<'_, PyModule>) -> PyResult<()> {
+    module.add_function(wrap_pyfunction!(api::mass_screen, module)?)?;
+    Ok(())
+}

mass_screen_rs-0.1.0/src/screen.rs

@@ -0,0 +1,180 @@
+//! The mask kernel: compressed positive entries per row, flags laid out per column for 64 pixels.
+use rayon::prelude::*;
+use std::slice;
+
+/// Image pixels per tile: a column's flags are one 64-byte line, one f32 compare covers them.
+pub(crate) const BT: usize = 64;
+/// Columns per block: the fail block is MBLK * BT = 512 KiB, which stays in L2 across all rows.
+pub(crate) const MBLK: usize = 8192;
+
+pub(crate) fn column_max(rg: &[f32], m: usize) -> Vec<f32> {
+    let merge = |mut a: Vec<f32>, b: &[f32]| {
+        for (x, &y) in a.iter_mut().zip(b) {
+            if y > *x {
+                *x = y;
+            }
+        }
+        a
+    };
+    // A partial maximum per chunk of rows, not per row: each partial is an m-float vector.
+    let rows_per_chunk = (rg.len() / m)
+        .div_ceil(4 * rayon::current_num_threads())
+        .max(1);
+    rg.par_chunks(m * rows_per_chunk)
+        .map(|chunk| chunk.chunks(m).fold(vec![f32::NEG_INFINITY; m], merge))
+        .reduce(|| vec![f32::NEG_INFINITY; m], |a, b| merge(a, &b))
+}
+
+/// One response row as (column, rhs) pairs for the entries with rhs > 0, plus block offsets.
+struct CsrRow {
+    idx: Vec<u32>,
+    val: Vec<f32>,
+    /// `blk[q]..blk[q + 1]` is the slice of `idx`/`val` inside column block `q`.
+    blk: Vec<u32>,
+    /// Sum of `val`: rows with large rhs mask cells earliest, so they run first.
+    score: f64,
+    max_val: f32,
+}
+
+fn csr_row(rg_row: &[f32], scale: &[f32], nblocks: usize) -> CsrRow {
+    let mut idx = Vec::new();
+    let mut val = Vec::new();
+    let mut score = 0.0f64;
+    let mut max_val = 0.0f32;
+    for (j, (&r, &s)) in rg_row.iter().zip(scale).enumerate() {
+        let v = s * r;
+        if v > 0.0 {
+            idx.push(j as u32);
+            val.push(v);
+            score += f64::from(v);
+            max_val = max_val.max(v);
+        }
+    }
+    let blk = (0..=nblocks)
+        .map(|q| idx.partition_point(|&j| (j as usize) < q * MBLK) as u32)
+        .collect();
+    CsrRow {
+        idx,
+        val,
+        blk,
+        score,
+        max_val,
+    }
+}
+
+struct SyncPtr(*mut f32);
+unsafe impl Sync for SyncPtr {}
+unsafe impl Send for SyncPtr {}
+impl SyncPtr {
+    // A method keeps the closure capturing the wrapper, not the bare pointer.
+    unsafe fn add(&self, off: usize) -> *mut f32 {
+        self.0.add(off)
+    }
+}
+
+/// Writes the mask into `keep` (`[b, m]`, fully overwritten). `rg` is `[n, m]` and `ys` is
+/// `[n, b]`, all C order. Columns with no positive response are masked. Needs `fill_value >= 0`,
+/// so that every screened value is non-negative and the zero entries of `rg` can be skipped.
+#[allow(clippy::too_many_arguments)]
+pub(crate) fn mass_screen_core(
+    rg: &[f32],
+    ys: &[f32],
+    keep: &mut [f32],
+    n: usize,
+    m: usize,
+    b: usize,
+    threshold: f32,
+    fill_value: f32,
+) {
+    assert_eq!(rg.len(), n * m);
+    assert_eq!(ys.len(), n * b);
+    assert_eq!(keep.len(), b * m);
+    assert!(fill_value >= 0.0);
+    if b == 0 || m == 0 {
+        return;
+    }
+    let col_max = column_max(rg, m);
+    let dead_cols: Vec<u32> = (0..m)
+        .filter(|&j| col_max[j].is_nan() || col_max[j] <= 0.0)
+        .map(|j| j as u32)
+        .collect();
+    let scale: Vec<f32> = col_max
+        .iter()
+        .map(|&mx| if mx > 0.0 { threshold / mx } else { 0.0 })
+        .collect();
+    let nblocks = m.div_ceil(MBLK);
+    let ntiles = b.div_ceil(BT);
+    let rows: Vec<CsrRow> = rg
+        .par_chunks(m)
+        .map(|r| csr_row(r, &scale, nblocks))
+        .collect();
+    let mut order: Vec<usize> = (0..n).filter(|&i| !rows[i].idx.is_empty()).collect();
+    order.sort_by(|&i, &j| rows[j].score.total_cmp(&rows[i].score));
+    let keep_ptr = SyncPtr(keep.as_mut_ptr());
+
+    // One flag block per worker, reused across its tasks: a fresh zeroed block per task costs
+    // a page fault per 4 KiB and showed as a 50 % slowdown on 32 threads.
+    (0..ntiles * nblocks).into_par_iter().for_each_init(
+        || vec![0u8; MBLK * BT],
+        |fail, task| {
+            let (t, q) = (task / nblocks, task % nblocks);
+            let b0 = t * BT;
+            let bt = BT.min(b - b0);
+            let m0 = q * MBLK;
+            let mb = MBLK.min(m - m0);
+            // fail[j * BT + k] != 0 once column m0 + j is masked for image pixel b0 + k.
+            let fail = &mut fail[..mb * BT];
+            fail.fill(0);
+            let dead_lo = dead_cols.partition_point(|&j| (j as usize) < m0);
+            let dead_hi = dead_cols.partition_point(|&j| (j as usize) < m0 + mb);
+            for &j in &dead_cols[dead_lo..dead_hi] {
+                fail[(j as usize - m0) * BT..][..BT].fill(1);
+            }
+            let mut ybuf = [f32::INFINITY; BT]; // lanes past bt never fail
+            let mut since_recount = 0usize;
+            for &i in &order {
+                let row = &rows[i];
+                let (lo, hi) = (row.blk[q] as usize, row.blk[q + 1] as usize);
+                if lo == hi {
+                    continue;
+                }
+                let mut ymin = f32::INFINITY;
+                for (yb, &y) in ybuf[..bt].iter_mut().zip(&ys[i * b + b0..i * b + b0 + bt]) {
+                    *yb = if y <= 0.0 { fill_value } else { y };
+                    ymin = ymin.min(*yb);
+                }
+                if ymin >= row.max_val {
+                    continue; // no pixel of this tile is below any rhs of this row
+                }
+                for (&j, &v) in row.idx[lo..hi].iter().zip(&row.val[lo..hi]) {
+                    let f = &mut fail[(j as usize - m0) * BT..][..BT];
+                    for (f, &y) in f.iter_mut().zip(&ybuf) {
+                        *f |= u8::from(y < v);
+                    }
+                }
+                // Stop once every column of the block has failed for every pixel of the tile. The
+                // check costs one pass over fail, so it runs once per ~4 passes' worth of updates.
+                since_recount += hi - lo;
+                if since_recount >= 4 * mb {
+                    since_recount = 0;
+                    let mut all_failed = [1u8; BT];
+                    for col in fail.chunks_exact(BT) {
+                        for (a, &f) in all_failed.iter_mut().zip(col) {
+                            *a &= f;
+                        }
+                    }
+                    if all_failed[..bt].iter().all(|&a| a == 1) {
+                        break;
+                    }
+                }
+            }
+            for k in 0..bt {
+                // SAFETY: tasks write disjoint rectangles [b0 + k, m0..m0 + mb] of keep.
+                let out = unsafe { slice::from_raw_parts_mut(keep_ptr.add((b0 + k) * m + m0), mb) };
+                for (j, o) in out.iter_mut().enumerate() {
+                    *o = if fail[j * BT + k] != 0 { f32::NAN } else { 1.0 };
+                }
+            }
+        },
+    );
+}

mass_screen_rs-0.1.0/src/tests/mod.rs

@@ -0,0 +1 @@
+mod screen;

mass_screen_rs-0.1.0/src/tests/screen.rs

@@ -0,0 +1,65 @@
+use crate::screen::{column_max, mass_screen_core, BT, MBLK};
+
+fn run(rg: &[f32], ys: &[f32], n: usize, m: usize, b: usize, fill: f32) -> Vec<f32> {
+    let mut keep = vec![0.0; b * m];
+    mass_screen_core(rg, ys, &mut keep, n, m, b, 2.24, fill);
+    keep
+}
+
+fn same(a: &[f32], b: &[f32]) -> bool {
+    a.len() == b.len()
+        && a.iter()
+            .zip(b)
+            .all(|(x, y)| x == y || (x.is_nan() && y.is_nan()))
+}
+
+#[test]
+fn muse_example() {
+    // The case muse/inversions/tests/test_mass_screen.py checks.
+    let rg = [1.0, 0.5, 0.1, 0.2, 1.0, 0.3];
+    let ys = [3.0, 1.0, 4.0, 5.0];
+    let keep = run(&rg, &ys, 2, 3, 2, 1.0);
+    assert!(same(&keep, &[1.0, 1.0, 1.0, f32::NAN, f32::NAN, 1.0]));
+}
+
+#[test]
+fn zero_column_is_masked_and_fill_value_applies() {
+    // Column 1 has no response; ys[0, 1] <= 0 is replaced by the fill value 1.0.
+    let rg = [1.0, 0.0, 0.1, 0.5, 0.0, 0.3];
+    let ys = [3.0, -1.0, 4.0, 5.0];
+    let keep = run(&rg, &ys, 2, 3, 2, 1.0);
+    assert!(same(&keep, &[1.0, f32::NAN, 1.0, f32::NAN, f32::NAN, 1.0]));
+}
+
+#[test]
+fn many_tiles_and_blocks_match_the_dense_rule() {
+    // Exercises partial tiles (b % BT != 0) and several column blocks.
+    let (n, m, b) = (7, 2 * MBLK + 5, 3 * BT + 9);
+    let mut seed = 12345u64;
+    let mut next = || {
+        seed = seed
+            .wrapping_mul(6364136223846793005)
+            .wrapping_add(1442695040888963407);
+        (seed >> 40) as f32 / (1u64 << 24) as f32
+    };
+    let rg: Vec<f32> = (0..n * m)
+        .map(|_| if next() < 0.05 { next() } else { 0.0 })
+        .collect();
+    let ys: Vec<f32> = (0..n * b).map(|_| next() * 4.0 - 1.0).collect();
+    let keep = run(&rg, &ys, n, m, b, 0.25);
+    let col_max = column_max(&rg, m);
+    for bi in 0..b {
+        for j in 0..m {
+            let kept = col_max[j] > 0.0
+                && (0..n).all(|i| {
+                    let y = if ys[i * b + bi] <= 0.0 {
+                        0.25
+                    } else {
+                        ys[i * b + bi]
+                    };
+                    y >= 2.24 / col_max[j] * rg[i * m + j]
+                });
+            assert_eq!(keep[bi * m + j] == 1.0, kept, "cell ({bi}, {j})");
+        }
+    }
+}