From 4a754c9faf13d705950351b64f915896406f6697 Mon Sep 17 00:00:00 2001 From: Tharindu Jayasinghe Date: Tue, 30 Jun 2026 11:47:06 -0700 Subject: [PATCH 1/6] feat(backend): portable torch backend for GLS + BLS via the array API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Foundations for multi-vendor GPU support (AMD/Intel/Mac + a real CPU path). GLS and BLS gain a portable `torch` backend that runs on any torch device (CUDA/ROCm/MPS/XPU/CPU) through one array-API code path, leaving the NVIDIA fast paths (cufinufft, cupy RawKernel) untouched. - core/_arrayapi.py (new): array-API dispatch + the ops outside the standard — scatter_add, to_host, device/precision resolution, torch device resolution. - core/backend.py: torch detection (torch_available / torch_devices / best_torch_device); register `torch` in available_backends (no torch import). - methods/base.py: vendor-aware resolve_backend — `auto` prefers cufinufft/cupy on CUDA, then torch on other GPUs, else the proven CPU path; parse + validate `torch` / `torch:`. - methods/gls.py: NUFFT-free direct trig-sum path (lombscargle_power_torch), sharing the Zechmeister-Kürster assembly via _assemble_power_parts (real components, so MPS without complex128 works). NUFFT path unchanged. - methods/_bls_core.py: _bls_search rewritten to array-API-standard ops; numpy and torch both flow through array_api_compat namespaces. cupy RawKernel + numba unchanged. - config.py: `torch` backend + orthogonal device/precision selectors (shared mixin on GLS/BLS) + GLSSettings.direct_freq_batch. - pyproject: array-api-compat core dep; [torch] extra; torch in [dev]; mypy/ruff config; version 1.1.0.dev0. Precision policy: precision="auto" is float64 except on MPS (float32); float64 on MPS raises rather than silently downgrading. Output is always numpy float64. Validated on CPU/torch-CPU: GLS torch vs finufft ~2e-11; BLS torch vs astropy ~4e-11 (torch vs numpy bit-identical); period recovery confirmed end-to-end. 151 passed, 8 skipped (CUDA-only). Also fixes two pre-existing numpy-2 stub typing nits (tls/bls geomspace) the toolchain surfaced. Co-Authored-By: Claude Opus 4.8 --- CHANGELOG.md | 20 +++ pyproject.toml | 25 +++- src/cuperiod/core/_arrayapi.py | 186 ++++++++++++++++++++++++++++ src/cuperiod/core/backend.py | 70 ++++++++++- src/cuperiod/core/config.py | 41 +++++-- src/cuperiod/methods/_bls_core.py | 195 ++++++++++++++++-------------- src/cuperiod/methods/base.py | 73 +++++++++-- src/cuperiod/methods/bls.py | 48 ++++++-- src/cuperiod/methods/gls.py | 135 ++++++++++++++++++--- src/cuperiod/methods/tls.py | 2 +- tests/conftest.py | 19 ++- tests/test_arrayapi.py | 138 +++++++++++++++++++++ tests/test_bls.py | 24 +++- tests/test_gls.py | 19 ++- 14 files changed, 846 insertions(+), 149 deletions(-) create mode 100644 src/cuperiod/core/_arrayapi.py create mode 100644 tests/test_arrayapi.py diff --git a/CHANGELOG.md b/CHANGELOG.md index b5f7f70..f8b9852 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,26 @@ All notable changes to cuPeriod are documented here. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +### Added + +- **Multi-vendor GPU support via PyTorch and the Python array API.** GLS and BLS gain a + portable `torch` backend that runs on AMD (ROCm), Intel (XPU), and Apple (MPS) GPUs as + well as a real CPU path — so the accelerated code is no longer NVIDIA-only, and works + even with no GPU at all. Select it with `backend="torch"` (or `"torch:cpu"`, + `"torch:cuda"`, `"torch:mps"`, `"torch:xpu"`); `backend="auto"` now reaches a torch GPU + on non-NVIDIA machines after the cufinufft/cupy fast paths. + - GLS adds a NUFFT-free direct trig-sum path (the portable formulation; cufinufft + remains the NVIDIA fast path). + - BLS runs its vectorized box search through the array-API namespace (the cupy + `RawKernel` remains the NVIDIA fast path). + - New `device` and `precision` settings: `precision="auto"` is float64 everywhere it is + supported and float32 only where the device forces it (Apple MPS cannot do float64); + an explicit `precision="float64"` on MPS raises rather than silently downgrading. +- `array-api-compat` is now a dependency; install the portable accelerator with the + `[torch]` extra (`pip install 'cuperiod[torch]'`). + ## [1.0.0] — 2026-06-29 First public release. diff --git a/pyproject.toml b/pyproject.toml index de713b2..db19aad 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "cuperiod" -version = "1.0.0" +version = "1.1.0.dev0" description = "Optimized, GPU-accelerated periodograms for astronomy" readme = "README.md" requires-python = ">=3.11" @@ -19,6 +19,12 @@ keywords = [ "time-series", "gpu", "cuda", + "rocm", + "amd", + "intel", + "metal", + "array-api", + "pytorch", ] classifiers = [ "Development Status :: 5 - Production/Stable", @@ -36,6 +42,9 @@ dependencies = [ "scipy>=1.10", "astropy>=6.0", "finufft>=2.2", + # Array-API dispatch: the portable (numpy/cupy/torch) compute paths run through one + # standard namespace. Lightweight, pure-Python — always installed. + "array-api-compat>=1.9", "pydantic>=2.5", "pydantic-settings>=2.1", "typer>=0.12", @@ -59,6 +68,11 @@ pandas = ["pandas>=2.0"] # astropy's compiled BoxLeastSquares by ~20x and is auto-selected on the CPU when # present (otherwise BLS falls back to astropy). fast = ["numba>=0.59"] +# Portable accelerator reaching AMD (ROCm), Intel (XPU), Apple (MPS), and a fast CPU +# path. Install the wheel matching your accelerator from https://pytorch.org/ — the +# plain wheel is CPU-only; CUDA/ROCm/XPU builds come from PyTorch's own index, so we do +# not pin a hardware-specific build here. +torch = ["torch>=2.2"] # Documentation toolchain (Sphinx + Furo theme). pandas is included so the # DataFrame-ingestion examples and the autodoc of from_dataframe build cleanly. docs = [ @@ -77,6 +91,8 @@ dev = [ "pandas>=2.0", "hypothesis>=6", "numba>=0.59", + # Exercise the portable torch-CPU path in CI (GPU devices auto-skip). + "torch>=2.2", ] [project.scripts] @@ -112,6 +128,9 @@ src = ["src", "tests"] [tool.ruff.lint] select = ["E", "F", "I", "UP", "B", "SIM", "NPY"] +# UP038 (use `X | Y` in isinstance) is deprecated upstream: the tuple form +# `isinstance(x, (A, B))` is faster at runtime, so we keep it. +ignore = ["UP038"] [tool.ruff.lint.per-file-ignores] # typer's API requires function calls (Option/Argument) in parameter defaults. @@ -140,6 +159,10 @@ module = [ "astropy.*", "scipy", "scipy.*", + "torch", + "torch.*", + "array_api_compat", + "array_api_compat.*", ] ignore_missing_imports = true diff --git a/src/cuperiod/core/_arrayapi.py b/src/cuperiod/core/_arrayapi.py new file mode 100644 index 0000000..62c0555 --- /dev/null +++ b/src/cuperiod/core/_arrayapi.py @@ -0,0 +1,186 @@ +"""Array-API dispatch and the few ops that live outside the standard. + +cuPeriod's portable compute paths are written against the Python array API standard +(via :mod:`array_api_compat`) so a single code path runs on NumPy (CPU), CuPy +(CUDA/ROCm), and PyTorch (CUDA/ROCm/MPS/XPU/CPU). A handful of operations the kernels +need are not in the standard — scatter-add, host transfer, and device/precision +resolution — so they live here as thin, backend-aware shims and the method bodies stay +clean and namespace-generic. + +The vendor-specific fast paths (cufinufft for GLS; the cupy ``RawKernel`` searches for +BLS/PDM/CE) do *not* go through here — they remain the NVIDIA accelerators. This module +is the portable tier that everything else rides on. +""" + +from __future__ import annotations + +from types import ModuleType +from typing import Any, Literal + +import numpy as np + +from cuperiod.core.errors import BackendUnavailableError + +#: User-facing precision selector. +Precision = Literal["auto", "float64", "float32"] + +#: Torch device kinds cuPeriod understands (ROCm reports as ``"cuda"``). +TORCH_DEVICES: tuple[str, ...] = ("cpu", "cuda", "mps", "xpu") + + +def array_namespace(*xs: Any) -> ModuleType: + """Return the array-API namespace for ``xs`` (an ``array_api_compat`` wrapper). + + Works for NumPy, CuPy, and PyTorch arrays. The returned namespace exposes the + standard surface (``astype`` as a function, ``clip`` without ``out=``, + ``remainder``, ``xp.float64`` dtype objects) regardless of the library version. + """ + import array_api_compat + + return array_api_compat.array_namespace(*xs) # type: ignore[no-any-return] + + +def is_torch_array(x: Any) -> bool: + """Whether ``x`` is a torch tensor (``False`` if torch is not installed).""" + try: + import array_api_compat + + return bool(array_api_compat.is_torch_array(x)) + except Exception: + return False + + +def is_cupy_array(x: Any) -> bool: + """Whether ``x`` is a cupy ndarray (``False`` if cupy is not installed).""" + try: + import array_api_compat + + return bool(array_api_compat.is_cupy_array(x)) + except Exception: + return False + + +def device_of(x: Any) -> str: + """Device kind of an array: ``"cpu"``, ``"cuda"``, ``"mps"``, or ``"xpu"``. + + A cupy array is always on ``"cuda"`` (ROCm included); a numpy array on ``"cpu"``; + a torch tensor reports its own ``device.type``. + """ + if is_torch_array(x): + return str(x.device.type) + if is_cupy_array(x): + return "cuda" + return "cpu" + + +def scatter_add(target: Any, index: Any, values: Any) -> None: + """In-place ``target[index] += values`` with repeated indices accumulated. + + ``target`` is 1-D and ``index``/``values`` are 1-D and aligned. ``numpy.add.at`` is + not part of the array-API standard and PyTorch has no equivalent *function*, so this + dispatches per backend: torch uses ``Tensor.index_add_``; numpy/cupy use ``add.at``. + """ + if is_torch_array(target): + target.index_add_(0, index.long(), values) + return + if is_cupy_array(target): + import cupy + + cupy.add.at(target, index, values) + return + np.add.at(target, index, values) + + +def to_host(a: Any) -> np.ndarray: + """Contiguous NumPy ``float64`` copy of a backend array (device → host). + + Preserves cuPeriod's output contract that every periodogram is returned as numpy + float64, regardless of the device/precision it was computed in. + """ + if is_torch_array(a): + return np.ascontiguousarray(a.detach().to("cpu").numpy(), dtype=np.float64) + if is_cupy_array(a): + import cupy + + return np.ascontiguousarray(cupy.asnumpy(a), dtype=np.float64) + return np.ascontiguousarray(np.asarray(a), dtype=np.float64) + + +def resolve_precision(precision: str, device: str) -> str: + """Resolve ``precision`` to a concrete ``"float64"``/``"float32"`` for ``device``. + + float64 is the default everywhere it is supported. Apple MPS cannot represent + float64 (a Metal limitation), so ``"auto"`` becomes float32 there; an explicit + ``"float64"`` request on MPS raises rather than silently downgrading. + """ + if precision == "float64": + if device == "mps": + raise BackendUnavailableError( + "MPS cannot compute in float64; pass precision='float32' for the Apple " + "GPU, or backend='cpu' (or backend='torch:cpu') for full float64." + ) + return "float64" + if precision == "float32": + return "float32" + return "float32" if device == "mps" else "float64" + + +def resolve_torch_device(backend: str, settings_device: str = "auto") -> str: + """Concrete torch device for a (possibly device-qualified) ``backend`` string. + + ``"torch:mps"`` forces ``mps``; bare ``"torch"`` defers to ``settings_device`` + (``"auto"`` → :func:`~cuperiod.core.backend.best_torch_device`). Raises if the + chosen device is not present here. + """ + from cuperiod.core.backend import best_torch_device, torch_devices + + device = backend.split(":", 1)[1] if ":" in backend else settings_device + if device == "auto": + device = best_torch_device() + if device not in torch_devices(): + raise BackendUnavailableError( + f"torch device {device!r} is not available here " + f"(present: {sorted(torch_devices()) or ['none']})" + ) + return device + + +def float_dtype(xp: ModuleType, resolved_precision: str) -> Any: + """The namespace float dtype object for a resolved precision name.""" + return xp.float32 if resolved_precision == "float32" else xp.float64 + + +def int_dtype(xp: ModuleType) -> Any: + """The namespace integer dtype used for index/bin arrays (always int64).""" + return xp.int64 + + +def to_device_array(host: np.ndarray, *, device: str, dtype: Any) -> Any: + """Place a host numpy array onto a torch ``device`` as ``dtype``. + + Used by the portable torch paths. ``dtype`` is a torch dtype object. Returns a torch + tensor; the caller obtains the matching array-API namespace via + :func:`array_namespace`. + """ + import torch + + return torch.as_tensor( + np.ascontiguousarray(host), dtype=dtype, device=torch.device(device) + ) + + +__all__ = [ + "TORCH_DEVICES", + "Precision", + "array_namespace", + "device_of", + "float_dtype", + "int_dtype", + "is_cupy_array", + "is_torch_array", + "resolve_precision", + "resolve_torch_device", + "scatter_add", + "to_device_array", + "to_host", +] diff --git a/src/cuperiod/core/backend.py b/src/cuperiod/core/backend.py index f29dc20..5305901 100644 --- a/src/cuperiod/core/backend.py +++ b/src/cuperiod/core/backend.py @@ -73,17 +73,79 @@ def cuda_available() -> bool: return False +def torch_available() -> bool: + """Whether PyTorch is importable. + + A CPU device always exists, so once torch is installed the portable ``torch`` + backend is never *unavailable* — this is what lets users with no GPU (or an + unsupported GPU) still run the accelerated code paths. Cheap: no import. + """ + return has_module("torch") + + +def torch_devices() -> set[str]: + """Torch device kinds usable here: ``{"cpu"}`` plus any of ``cuda``/``mps``/``xpu``. + + Imports torch to query the runtimes (so call only when a torch path is actually + being taken). ROCm builds report AMD GPUs as ``"cuda"``, so AMD needs no separate + name. Returns an empty set if torch is absent or fails to import. + """ + if not has_module("torch"): + return set() + try: + import torch + except Exception: + return set() + out = {"cpu"} + try: + if torch.cuda.is_available(): + out.add("cuda") + except Exception: + pass + try: + if torch.backends.mps.is_available(): + out.add("mps") + except Exception: + pass + try: + if hasattr(torch, "xpu") and torch.xpu.is_available(): + out.add("xpu") + except Exception: + pass + return out + + +def torch_gpu_available() -> bool: + """Whether torch sees any non-CPU device (CUDA/ROCm, MPS, or XPU).""" + return bool(torch_devices() - {"cpu"}) + + +def best_torch_device() -> str: + """Preferred torch device, in order ``cuda`` → ``xpu`` → ``mps`` → ``cpu``.""" + devices = torch_devices() + for kind in ("cuda", "xpu", "mps"): + if kind in devices: + return kind + return "cpu" + + def available_backends() -> set[str]: """The set of backend names importable in this environment. Returns a union across methods: always includes ``numpy``; adds ``finufft``, - ``astropy``, ``numba`` when importable, and the GPU names ``cufinufft``/``cupy`` - only when a CUDA device is present. + ``astropy``, ``numba``, ``torch`` when importable, and the GPU names + ``cufinufft``/``cupy`` only when a CUDA device is present. + + ``torch`` is added on a cheap import-spec check (no torch import here): a CPU device + always exists, so an installed torch is always a usable backend. Which torch + *devices* are present is answered separately by :func:`torch_devices`. """ out: set[str] = {"numpy"} for name in ("finufft", "astropy", "numba"): if has_module(name): out.add(name) + if torch_available(): + out.add("torch") if cuda_available(): out.add("cupy") if has_module("cufinufft"): @@ -151,8 +213,12 @@ def array_module(a: object) -> ModuleType: __all__ = [ "array_module", "available_backends", + "best_torch_device", "cuda_available", "ensure_cuda_dll_path", "ensure_shared_memory", "has_module", + "torch_available", + "torch_devices", + "torch_gpu_available", ] diff --git a/src/cuperiod/core/config.py b/src/cuperiod/core/config.py index 1d0e5e8..73c3699 100644 --- a/src/cuperiod/core/config.py +++ b/src/cuperiod/core/config.py @@ -26,7 +26,27 @@ def _require_lt(lo: float | None, hi: float | None, lo_name: str, hi_name: str) raise ValueError(f"{lo_name} ({lo}) must be < {hi_name} ({hi})") -class GLSSettings(BaseSettings): +class _DeviceSettings(BaseSettings): + """Device and precision selectors shared by torch-capable methods. + + Orthogonal to ``backend``: ``device`` chooses the torch device when the portable + ``torch`` backend is selected (``"auto"`` picks the best present), and ``precision`` + controls the compute dtype. ``precision="auto"`` is float64 everywhere it is + supported and float32 only where the device forces it (Apple MPS cannot do float64); + an explicit ``"float64"`` on MPS raises rather than silently downgrading. Both + are environment-overridable like every other setting (``CUPERIOD__DEVICE``). + """ + + device: Literal["auto", "cpu", "cuda", "mps", "xpu"] = Field( + default="auto", description="Torch device when the torch backend is used." + ) + precision: Literal["auto", "float64", "float32"] = Field( + default="auto", + description="Compute precision; 'auto' is float64 except on MPS (float32).", + ) + + +class GLSSettings(_DeviceSettings): """Settings for the generalized Lomb-Scargle (GLS) periodogram.""" model_config = SettingsConfigDict(env_prefix="CUPERIOD_GLS_", extra="forbid") @@ -66,18 +86,23 @@ def _check_bounds(self) -> Self: min_detections: int = Field( default=10, ge=3, description="Skip if fewer finite points." ) - backend: Literal["auto", "cpu", "gpu", "finufft", "cufinufft", "astropy"] = Field( - default="auto", description="Compute backend." - ) + backend: Literal[ + "auto", "cpu", "gpu", "finufft", "cufinufft", "torch", "astropy" + ] = Field(default="auto", description="Compute backend.") nufft_eps: float = Field( default=1e-9, gt=0.0, description="NUFFT relative tolerance." ) + direct_freq_batch: int = Field( + default=4096, + ge=1, + description="Frequency chunk for the portable (torch) direct trig-sum path.", + ) downsample_points: int = Field( default=2000, ge=2, description="Stored downsampled-spectrum size." ) -class BLSSettings(BaseSettings): +class BLSSettings(_DeviceSettings): """Settings for the box least squares (BLS) search.""" model_config = SettingsConfigDict(env_prefix="CUPERIOD_BLS_", extra="forbid") @@ -139,9 +164,9 @@ def _check_bounds(self) -> Self: min_detections: int = Field( default=20, ge=3, description="Skip if fewer finite points." ) - backend: Literal["auto", "cpu", "gpu", "numpy", "astropy", "cupy"] = Field( - default="auto", description="Compute backend." - ) + backend: Literal[ + "auto", "cpu", "gpu", "numpy", "astropy", "cupy", "torch" + ] = Field(default="auto", description="Compute backend.") batch_periods: int = Field( default=2048, ge=1, description="Trial periods per vectorized batch." ) diff --git a/src/cuperiod/methods/_bls_core.py b/src/cuperiod/methods/_bls_core.py index 9d665d7..d3ca1e1 100644 --- a/src/cuperiod/methods/_bls_core.py +++ b/src/cuperiod/methods/_bls_core.py @@ -24,8 +24,17 @@ import numpy as np +from cuperiod.core._arrayapi import ( + array_namespace, + resolve_precision, + scatter_add, + to_device_array, + to_host, +) from cuperiod.core._typing import FloatArray +#: Core single-array-library backends (the torch backend is device-qualified, e.g. +#: ``"torch:cpu"``, so backend strings are typed as ``str`` where torch is accepted). BLSBackend = Literal["numpy", "cupy", "numba"] #: Inverse-variance floor: an empty in/out side (ivar sum 0) is skipped, not divided by. @@ -60,11 +69,6 @@ class BLSPower: log_likelihood: FloatArray -def _scatter_add(xp: ModuleType, target: Any, idx: Any, values: Any) -> None: - """``target[idx] += values`` with index repeats accumulated, on numpy/cupy.""" - xp.add.at(target, idx, values) - - def _duration_bins(durations: FloatArray, bin_duration: float) -> list[int]: """Box widths in bins, deduped in first-seen order (astropy's loop order).""" seen: set[int] = set() @@ -118,23 +122,31 @@ def _bls_search( obj_flag: int, batch: int, ) -> dict[str, Any]: - """Vectorized port of astropy ``run_bls`` for one duration-homogeneous call.""" + """Vectorized port of astropy ``run_bls`` for one duration-homogeneous call. + + Array-API generic: ``xp`` is an :mod:`array_api_compat` namespace, so the identical + body runs on numpy (CPU) and torch (CPU/CUDA/ROCm/MPS/XPU). The working float dtype + follows ``periods`` (float64, or float32 on a float32 device); index/bin arrays are + int64. The cupy ``RawKernel`` fast path (NVIDIA) is separate and not this code. + """ + fdtype = periods.dtype + idtype = xp.int64 n_points = int(t.shape[0]) yw = y * ivar - sum_y = float(yw.sum()) - sum_ivar = float(ivar.sum()) - t_min = float(t.min()) + sum_y = float(xp.sum(yw)) + sum_ivar = float(xp.sum(ivar)) + t_min = float(xp.min(t)) tau = t - t_min n_periods = int(periods.shape[0]) out = { - "power": xp.full(n_periods, -np.inf, dtype=np.float64), - "depth": xp.zeros(n_periods, dtype=np.float64), - "depth_err": xp.zeros(n_periods, dtype=np.float64), - "depth_snr": xp.zeros(n_periods, dtype=np.float64), - "duration": xp.zeros(n_periods, dtype=np.float64), - "transit_time": xp.zeros(n_periods, dtype=np.float64), - "log_likelihood": xp.zeros(n_periods, dtype=np.float64), + "power": xp.full(n_periods, -np.inf, dtype=fdtype), + "depth": xp.zeros(n_periods, dtype=fdtype), + "depth_err": xp.zeros(n_periods, dtype=fdtype), + "depth_snr": xp.zeros(n_periods, dtype=fdtype), + "duration": xp.zeros(n_periods, dtype=fdtype), + "transit_time": xp.zeros(n_periods, dtype=fdtype), + "log_likelihood": xp.zeros(n_periods, dtype=fdtype), } if n_periods == 0 or not dur_bins: return out @@ -143,21 +155,21 @@ def _bls_search( stop = min(start + batch, n_periods) pb = periods[start:stop] n_p = int(pb.shape[0]) - rows = xp.arange(n_p) - n_bins = xp.ceil(pb / bin_duration).astype(np.int64) + oversample - - phase_t = xp.mod(tau[None, :], pb[:, None]) - ind = (phase_t / bin_duration).astype(np.int64) + 1 - xp.clip(ind, 0, width - 1, out=ind) - flat = (rows[:, None] * width + ind).ravel() - mean_y = xp.zeros(n_p * width, dtype=np.float64) - mean_ivar = xp.zeros(n_p * width, dtype=np.float64) - yw_b = xp.broadcast_to(yw, (n_p, n_points)).ravel() - ivar_b = xp.broadcast_to(ivar, (n_p, n_points)).ravel() - _scatter_add(xp, mean_y, flat, yw_b) - _scatter_add(xp, mean_ivar, flat, ivar_b) - mean_y = mean_y.reshape(n_p, width) - mean_ivar = mean_ivar.reshape(n_p, width) + rows = xp.arange(n_p, dtype=idtype) + n_bins = xp.astype(xp.ceil(pb / bin_duration), idtype) + oversample + + phase_t = xp.remainder(tau[None, :], pb[:, None]) + ind = xp.astype(phase_t / bin_duration, idtype) + 1 + ind = xp.clip(ind, 0, width - 1) + flat = xp.reshape(rows[:, None] * width + ind, (-1,)) + mean_y = xp.zeros(n_p * width, dtype=fdtype) + mean_ivar = xp.zeros(n_p * width, dtype=fdtype) + yw_b = xp.reshape(xp.broadcast_to(yw, (n_p, n_points)), (-1,)) + ivar_b = xp.reshape(xp.broadcast_to(ivar, (n_p, n_points)), (-1,)) + scatter_add(mean_y, flat, yw_b) + scatter_add(mean_ivar, flat, ivar_b) + mean_y = xp.reshape(mean_y, (n_p, width)) + mean_ivar = xp.reshape(mean_ivar, (n_p, width)) for j in range(oversample): dst = xp.clip(n_bins - oversample + j, 0, width - 1) @@ -168,16 +180,16 @@ def _bls_search( cw = xp.cumsum(mean_ivar, axis=1) del mean_y, mean_ivar - best_obj = xp.full(n_p, -np.inf, dtype=np.float64) - best_n = xp.zeros(n_p, dtype=np.int64) - best_d = xp.zeros(n_p, dtype=np.int64) + best_obj = xp.full(n_p, -np.inf, dtype=fdtype) + best_n = xp.zeros(n_p, dtype=idtype) + best_d = xp.zeros(n_p, dtype=idtype) for kd in dur_bins: if kd >= width: continue y_in = cy[:, kd:] - cy[:, :-kd] ivar_in = cw[:, kd:] - cw[:, :-kd] ivar_out = sum_ivar - ivar_in - cols = xp.arange(y_in.shape[1]) + cols = xp.arange(y_in.shape[1], dtype=idtype) valid = ( (cols[None, :] <= (n_bins[:, None] - kd)) & (ivar_in >= _IVAR_EPS) @@ -188,8 +200,8 @@ def _bls_search( cand_obj = obj[rows, cand_n] improve = cand_obj > best_obj best_obj = xp.where(improve, cand_obj, best_obj) - best_n = xp.where(improve, cand_n.astype(np.int64), best_n) - best_d = xp.where(improve, np.int64(kd), best_d) + best_n = xp.where(improve, xp.astype(cand_n, idtype), best_n) + best_d = xp.where(improve, kd, best_d) finite = xp.isfinite(best_obj) y_in = cy[rows, best_n + best_d] - cy[rows, best_n] @@ -203,15 +215,15 @@ def _bls_search( depth_err = xp.sqrt(1.0 / safe_in + 1.0 / safe_out) depth_snr = depth / depth_err log_like = 0.5 * ivar_in * depth * depth - duration = best_d.astype(np.float64) * bin_duration + duration = xp.astype(best_d, fdtype) * bin_duration transit_time = ( - xp.mod(best_n.astype(np.float64) * bin_duration + 0.5 * duration, pb) + xp.remainder(xp.astype(best_n, fdtype) * bin_duration + 0.5 * duration, pb) + t_min ) power = depth_snr if obj_flag == 0 else log_like sl = slice(start, stop) - zero = xp.zeros(n_p, dtype=np.float64) + zero = xp.zeros(n_p, dtype=fdtype) out["power"][sl] = xp.where(finite, power, -np.inf) out["depth"][sl] = xp.where(finite, depth, zero) out["depth_err"][sl] = xp.where(finite, depth_err, zero) @@ -608,10 +620,11 @@ def bls_power( oversample: int, *, objective: str = "snr", - backend: BLSBackend = "numpy", + backend: str = "numpy", batch: int = DEFAULT_BATCH, + precision: str = "auto", ) -> BLSPower: - """BLS box search over ``periods`` via numpy (CPU) or cupy (GPU). + """BLS box search over ``periods`` via numpy/torch (portable), cupy, or numba. Equivalent to ``BoxLeastSquares(t, y, dy).power(periods, durations, objective=objective, oversample=oversample)`` — same binning and objective. ``y`` @@ -629,10 +642,14 @@ def bls_power( Phase bins per shortest duration. objective : {"snr", "likelihood"}, default "snr" Box objective. - backend : {"numpy", "cupy"}, default "numpy" - CPU reference or GPU kernel. + backend : str, default "numpy" + ``"numpy"`` (array-API CPU reference), ``"torch"`` / ``"torch:"`` + (portable array-API path, any torch device), ``"cupy"`` (NVIDIA RawKernel), or + ``"numba"`` (multicore CPU). batch : int, default 2048 - Trial periods per vectorized batch (numpy backend). + Trial periods per vectorized batch (numpy/torch backends). + precision : {"auto", "float64", "float32"}, default "auto" + Device-side compute precision for the torch backend (float64 except on MPS). Returns ------- @@ -642,7 +659,8 @@ def bls_power( if objective not in ("snr", "likelihood"): raise ValueError("objective must be 'snr' or 'likelihood'") obj_flag = 0 if objective == "snr" else 1 - if backend not in ("numpy", "cupy", "numba"): + is_torch = backend == "torch" or backend.startswith("torch:") + if not is_torch and backend not in ("numpy", "cupy", "numba"): raise ValueError(f"unknown backend {backend!r}") periods_host = np.ascontiguousarray(periods, dtype=np.float64) @@ -655,68 +673,59 @@ def bls_power( max_n_bins = int(np.ceil(float(periods_host.max()) / bin_duration)) + oversample width = max_n_bins + 1 + t_host = np.ascontiguousarray(t, dtype=np.float64) + y_host = np.ascontiguousarray(y, dtype=np.float64) ivar_host = 1.0 / (np.ascontiguousarray(dy, dtype=np.float64) ** 2) + if backend == "cupy": from cuperiod.core.backend import ensure_cuda_dll_path ensure_cuda_dll_path() out = _bls_search_cuda( - np.ascontiguousarray(t, dtype=np.float64), - np.ascontiguousarray(y, dtype=np.float64), - ivar_host, - periods_host, - bin_duration=bin_duration, - dur_bins=dur_bins, - oversample=oversample, - width=width, - obj_flag=obj_flag, + t_host, y_host, ivar_host, periods_host, + bin_duration=bin_duration, dur_bins=dur_bins, + oversample=oversample, width=width, obj_flag=obj_flag, ) - import cupy - - def host(a: Any) -> FloatArray: - return np.asarray(cupy.asnumpy(a), dtype=np.float64) elif backend == "numba": out = _bls_search_numba( - np.ascontiguousarray(t, dtype=np.float64), - np.ascontiguousarray(y, dtype=np.float64), - ivar_host, - periods_host, - bin_duration=bin_duration, - dur_bins=dur_bins, - oversample=oversample, - width=width, - obj_flag=obj_flag, + t_host, y_host, ivar_host, periods_host, + bin_duration=bin_duration, dur_bins=dur_bins, + oversample=oversample, width=width, obj_flag=obj_flag, ) - - def host(a: Any) -> FloatArray: - return np.asarray(a, dtype=np.float64) - else: + elif is_torch: + import torch + + device = backend.split(":", 1)[1] if ":" in backend else "cpu" + tdtype = ( + torch.float32 + if resolve_precision(precision, device) == "float32" + else torch.float64 + ) + t_d = to_device_array(t_host, device=device, dtype=tdtype) + y_d = to_device_array(y_host, device=device, dtype=tdtype) + ivar_d = to_device_array(ivar_host, device=device, dtype=tdtype) + periods_d = to_device_array(periods_host, device=device, dtype=tdtype) out = _bls_search( - np, - np.ascontiguousarray(t, dtype=np.float64), - np.ascontiguousarray(y, dtype=np.float64), - ivar_host, - periods_host, - bin_duration=bin_duration, - dur_bins=dur_bins, - oversample=oversample, - width=width, - obj_flag=obj_flag, - batch=batch, + array_namespace(periods_d), t_d, y_d, ivar_d, periods_d, + bin_duration=bin_duration, dur_bins=dur_bins, + oversample=oversample, width=width, obj_flag=obj_flag, batch=batch, + ) + else: # numpy, through the array-API compat namespace + out = _bls_search( + array_namespace(periods_host), t_host, y_host, ivar_host, periods_host, + bin_duration=bin_duration, dur_bins=dur_bins, + oversample=oversample, width=width, obj_flag=obj_flag, batch=batch, ) - - def host(a: Any) -> FloatArray: - return np.asarray(a, dtype=np.float64) return BLSPower( period=periods_host, - power=host(out["power"]), - depth=host(out["depth"]), - depth_err=host(out["depth_err"]), - depth_snr=host(out["depth_snr"]), - duration=host(out["duration"]), - transit_time=host(out["transit_time"]), - log_likelihood=host(out["log_likelihood"]), + power=to_host(out["power"]), + depth=to_host(out["depth"]), + depth_err=to_host(out["depth_err"]), + depth_snr=to_host(out["depth_snr"]), + duration=to_host(out["duration"]), + transit_time=to_host(out["transit_time"]), + log_likelihood=to_host(out["log_likelihood"]), ) diff --git a/src/cuperiod/methods/base.py b/src/cuperiod/methods/base.py index 7eb291a..653190d 100644 --- a/src/cuperiod/methods/base.py +++ b/src/cuperiod/methods/base.py @@ -17,7 +17,14 @@ from pydantic_settings import BaseSettings -from cuperiod.core.backend import available_backends, cuda_available +from cuperiod.core._arrayapi import TORCH_DEVICES +from cuperiod.core.backend import ( + available_backends, + cuda_available, + torch_available, + torch_devices, + torch_gpu_available, +) from cuperiod.core.columns import Domain from cuperiod.core.errors import BackendUnavailableError, UnknownMethodError from cuperiod.core.grid import GridSpec @@ -47,8 +54,12 @@ class PeriodogramMethod(ABC): settings_cls: ClassVar[type[BaseSettings]] #: Best CPU backend name. cpu_backend: ClassVar[str] - #: GPU backend name, or ``None`` if the method has no GPU path yet. + #: NVIDIA fast-path GPU backend name (cufinufft / cupy), or ``None`` if the method + #: has no CUDA path. gpu_backend: ClassVar[str | None] = None + #: Portable GPU backend name (``"torch"``) reaching AMD/Intel/Mac/CPU, or ``None`` + #: if the method has not been ported to the array-API path yet. + portable_gpu_backend: ClassVar[str | None] = None #: Every backend this method can run. all_backends: ClassVar[tuple[str, ...]] @@ -68,37 +79,49 @@ def coerce_settings(self, settings: BaseSettings | None) -> BaseSettings: def resolve_backend(self, requested: str) -> str: """Resolve ``auto``/``cpu``/``gpu``/concrete to a runnable backend name. + ``"auto"`` prefers the NVIDIA fast-path (``gpu_backend``) when a CUDA device is + present, then the portable ``torch`` backend when torch sees a non-CPU device + (AMD/Intel/Mac), and otherwise the proven CPU path. ``"gpu"`` is the same but + raises when no GPU is available. A concrete ``"torch"`` / ``"torch:"`` + request selects the portable path explicitly. + Parameters ---------- requested : str - ``"auto"`` (GPU when present, else CPU), ``"cpu"``, ``"gpu"``, or a - concrete backend name belonging to this method. + ``"auto"``, ``"cpu"``, ``"gpu"``, ``"torch"``, ``"torch:cpu|cuda|mps|xpu"``, + or a concrete backend name belonging to this method. Returns ------- str - A concrete, available backend name. + A concrete, available backend name (``"torch:"`` kept as given). Raises ------ BackendUnavailableError - If GPU was requested but is unavailable, or a named backend is unknown to - this method or not importable here. + If GPU/torch was requested but is unavailable, or a named backend is unknown + to this method or not importable here. """ available = available_backends() if requested == "auto": if self.gpu_backend is not None and cuda_available(): return self.gpu_backend + if self.portable_gpu_backend is not None and torch_gpu_available(): + return self.portable_gpu_backend return self.cpu_backend if requested == "cpu": return self.cpu_backend if requested == "gpu": if self.gpu_backend is not None and cuda_available(): return self.gpu_backend + if self.portable_gpu_backend is not None and torch_gpu_available(): + return self.portable_gpu_backend raise BackendUnavailableError( - f"{self.name}: GPU backend unavailable (need the [gpu] extra and a " - "CUDA device)" + f"{self.name}: no GPU backend available (need the [gpu] extra and a " + "CUDA device, or the [torch] extra and a CUDA/ROCm/MPS/XPU device)" ) + if requested == "torch" or requested.startswith("torch:"): + return self._resolve_torch(requested) if requested not in self.all_backends: raise BackendUnavailableError( f"{self.name}: unknown backend {requested!r}; " @@ -116,9 +139,37 @@ def resolve_backend(self, requested: str) -> str: ) return requested + def _resolve_torch(self, requested: str) -> str: + """Validate a concrete ``torch``/``torch:`` request for this method.""" + if self.portable_gpu_backend != "torch": + raise BackendUnavailableError( + f"{self.name}: no portable 'torch' backend for this method" + ) + if not torch_available(): + raise BackendUnavailableError( + f"{self.name}: backend 'torch' needs the [torch] extra (pip install " + "'cuperiod[torch]')" + ) + if ":" in requested: + device = requested.split(":", 1)[1] + if device not in TORCH_DEVICES: + raise BackendUnavailableError( + f"{self.name}: unknown torch device {device!r}; " + f"choose from {TORCH_DEVICES}" + ) + if device not in torch_devices(): + raise BackendUnavailableError( + f"{self.name}: torch device {device!r} is not available here" + ) + return requested + def is_gpu_backend(self, backend: str) -> bool: - """Whether ``backend`` is this method's GPU backend.""" - return backend == self.gpu_backend + """Whether ``backend`` runs on a GPU (NVIDIA fast-path or non-CPU torch).""" + if backend == self.gpu_backend: + return True + if backend == "torch" or backend.startswith("torch:"): + return backend != "torch:cpu" + return False # -- compute ----------------------------------------------------------------- @abstractmethod diff --git a/src/cuperiod/methods/bls.py b/src/cuperiod/methods/bls.py index 3e2dfbd..abb308c 100644 --- a/src/cuperiod/methods/bls.py +++ b/src/cuperiod/methods/bls.py @@ -17,6 +17,7 @@ import numpy as np +from cuperiod.core._arrayapi import resolve_torch_device from cuperiod.core._typing import FloatArray from cuperiod.core.columns import Domain from cuperiod.core.config import BLSSettings @@ -24,7 +25,7 @@ from cuperiod.core.grid import GridSpec from cuperiod.core.lightcurve import LightCurve, MultiBandLightCurve from cuperiod.core.result import Periodogram -from cuperiod.methods._bls_core import BLSBackend, bls_power +from cuperiod.methods._bls_core import bls_power from cuperiod.methods.base import PeriodogramMethod, register #: Per-period fields shared by every backend (astropy + the in-house search). @@ -41,7 +42,7 @@ def _segment_durations(p_lo: float, settings: BLSSettings) -> FloatArray: d_lo = max(settings.duration_min_frac * p_lo, settings.min_duration_days) if d_lo >= d_hi: return np.asarray([d_hi], dtype=np.float64) - return np.geomspace(d_lo, d_hi, settings.n_durations) + return np.asarray(np.geomspace(d_lo, d_hi, settings.n_durations), dtype=np.float64) def _max_period(baseline: float, settings: BLSSettings) -> float: @@ -63,13 +64,14 @@ def _segment_grids( p_hi = min(p_lo * settings.segment_factor, max_period) freq = np.arange(1.0 / p_hi, 1.0 / p_lo, df) if freq.size: - grids.append((1.0 / freq[::-1], _segment_durations(p_lo, settings))) + periods = np.asarray(1.0 / freq[::-1], dtype=np.float64) + grids.append((periods, _segment_durations(p_lo, settings))) p_lo = p_hi return grids def _segment_power( - backend: BLSBackend | Literal["astropy"], + backend: str, jd: FloatArray, flux: FloatArray, err: FloatArray, @@ -101,6 +103,7 @@ def _segment_power( objective=settings.objective, backend=backend, batch=settings.batch_periods, + precision=settings.precision, ) return {name: getattr(power, name) for name in _SEGMENT_FIELDS} @@ -152,15 +155,32 @@ class BLSMethod(PeriodogramMethod): settings_cls: ClassVar[type] = BLSSettings cpu_backend: ClassVar[str] = "astropy" gpu_backend: ClassVar[str | None] = "cupy" - all_backends: ClassVar[tuple[str, ...]] = ("numba", "numpy", "astropy", "cupy") + portable_gpu_backend: ClassVar[str | None] = "torch" + all_backends: ClassVar[tuple[str, ...]] = ( + "numba", "numpy", "astropy", "cupy", "torch", + ) def resolve_backend(self, requested: str) -> str: - """Prefer the multicore numba box search on the CPU when it is installed.""" - from cuperiod.core.backend import available_backends, cuda_available + """Prefer cupy on NVIDIA, then torch on other GPUs, else the numba CPU search. + + Keeps BLS's CPU preference (multicore numba when installed, else astropy) for + ``cpu``/``auto``, while ``auto`` still reaches a GPU: the cupy kernel on CUDA, + then the portable torch path on AMD/Intel/Mac. Concrete ``torch``/``torch:*`` + requests are validated by the base method. + """ + from cuperiod.core.backend import ( + available_backends, + cuda_available, + torch_gpu_available, + ) - if requested == "auto" and self.gpu_backend is not None and cuda_available(): - return self.gpu_backend - if requested in ("cpu", "auto"): + if requested == "auto": + if self.gpu_backend is not None and cuda_available(): + return self.gpu_backend + if self.portable_gpu_backend is not None and torch_gpu_available(): + return self.portable_gpu_backend + return "numba" if "numba" in available_backends() else "astropy" + if requested == "cpu": return "numba" if "numba" in available_backends() else "astropy" return super().resolve_backend(requested) @@ -199,7 +219,11 @@ def power( # type: ignore[override] jd = finite.time flux = finite.value err = finite.error if finite.error is not None else np.ones_like(flux) - bck: BLSBackend | Literal["astropy"] = backend # type: ignore[assignment] + if backend == "torch" or backend.startswith("torch:"): + # Fully-qualify the device so both dispatch and the recorded backend are + # concrete (e.g. "torch:cpu"); settings.device picks the device for bare + # "torch"/"auto"/"gpu". + backend = f"torch:{resolve_torch_device(backend, settings.device)}" if grid.meta.get("segmented", False): segments = _segment_grids(finite.baseline, settings) @@ -214,7 +238,7 @@ def power( # type: ignore[override] chunks: dict[str, list[FloatArray]] = {name: [] for name in _SEGMENT_FIELDS} for periods, durations in segments: - seg = _segment_power(bck, jd, flux, err, periods, durations, settings) + seg = _segment_power(backend, jd, flux, err, periods, durations, settings) for name in _SEGMENT_FIELDS: chunks[name].append(seg[name]) return _assemble(chunks, n, finite.baseline, backend, finite.meta) diff --git a/src/cuperiod/methods/gls.py b/src/cuperiod/methods/gls.py index 6fab01d..f92dd02 100644 --- a/src/cuperiod/methods/gls.py +++ b/src/cuperiod/methods/gls.py @@ -32,6 +32,13 @@ import numpy as np +from cuperiod.core._arrayapi import ( + array_namespace, + resolve_precision, + resolve_torch_device, + to_device_array, + to_host, +) from cuperiod.core._typing import FloatArray from cuperiod.core.backend import array_module, ensure_cuda_dll_path from cuperiod.core.config import GLSSettings @@ -84,19 +91,19 @@ def _trig_sums( raise ValueError(f"unknown NUFFT backend {backend!r}") -def _assemble_power( - sw: Any, swy: Any, sw2: Any, y_mean: float, yy: float, fit_mean: bool +def _assemble_power_parts( + xp: Any, + c: Any, s: Any, yc: Any, ys: Any, c2: Any, s2: Any, + y_mean: float, yy: float, fit_mean: bool, ) -> Any: - """Zechmeister-Kürster generalized-LS power (standard norm) from three trig sums. + """Zechmeister-Kürster generalized-LS power from the real trig-sum components. - Works on numpy or cupy arrays (the GPU path keeps everything on device); - ``y_mean``/``yy`` are host scalars. Degenerate frequencies map to 0. + ``c``/``s`` are the cos/sin sums of the weights, ``yc``/``ys`` of the weighted data, + and ``c2``/``s2`` the doubled-frequency cos/sin sums of the weights. Uses only + array-API-standard ops, so it runs unchanged on numpy, cupy, and torch (any device) + — no complex dtype is needed, so the Apple-MPS path (no complex128) works. + Degenerate frequencies map to 0. """ - xp = array_module(sw) - c, s = sw.real, sw.imag - yc, ys = swy.real, swy.imag - c2, s2 = sw2.real, sw2.imag - cc = 0.5 * (1.0 + c2) ss = 0.5 * (1.0 - c2) cs = 0.5 * s2 @@ -108,12 +115,25 @@ def _assemble_power( cs = cs - c * s denom = cc * ss - cs * cs + power = (ss * yc * yc + cc * ys * ys - 2.0 * cs * yc * ys) / (yy * denom) + return xp.where(xp.isfinite(power), power, xp.zeros_like(power)) + + +def _assemble_power( + sw: Any, swy: Any, sw2: Any, y_mean: float, yy: float, fit_mean: bool +) -> Any: + """GLS power from three *complex* trig sums (the NUFFT path). + + Splits each complex sum into its real/imaginary parts and delegates the Zechmeister- + Kürster math to :func:`_assemble_power_parts`. Works on numpy or cupy arrays (GPU + path keeps everything on device); ``y_mean``/``yy`` are host scalars. + """ + xp = array_module(sw) + parts = (sw.real, sw.imag, swy.real, swy.imag, sw2.real, sw2.imag) if xp is np: with np.errstate(divide="ignore", invalid="ignore"): - power = (ss * yc * yc + cc * ys * ys - 2.0 * cs * yc * ys) / (yy * denom) - else: - power = (ss * yc * yc + cc * ys * ys - 2.0 * cs * yc * ys) / (yy * denom) - return xp.nan_to_num(power, nan=0.0, posinf=0.0, neginf=0.0) + return _assemble_power_parts(xp, *parts, y_mean, yy, fit_mean) + return _assemble_power_parts(xp, *parts, y_mean, yy, fit_mean) def _prep( @@ -188,6 +208,72 @@ def lombscargle_power( ) +def _trig_sums_direct( + xp: Any, tau: Any, strengths: Any, f0: float, df: float, nf: int, *, freq_batch: int +) -> tuple[Any, Any]: + """Direct (NUFFT-free) trig sums on the uniform grid ``f_k = f0 + k*df``. + + Returns ``(cos_sum, sin_sum)`` real arrays of length ``nf`` with + ``cos_sum[k] = sum_j strengths_j cos(2*pi*f_k*tau_j)`` and ``sin_sum`` the matching + ``+sin`` sum (the ``isign=+1`` convention of the NUFFT path). This is ``O(N*nf)`` + rather than the NUFFT's ``O(nf log nf)``, but is pure array-API and so runs on any + backend/device — the portable GLS path for AMD/Intel/Mac. Batched over frequency to + bound the transient ``(chunk, N)`` angle matrix. + """ + two_pi = 2.0 * float(np.pi) + fdtype = tau.dtype + freqs = f0 + df * xp.arange(nf, dtype=fdtype) + cos_sum = xp.empty(nf, dtype=fdtype) + sin_sum = xp.empty(nf, dtype=fdtype) + strength_row = strengths[None, :] + for start in range(0, nf, freq_batch): + stop = min(start + freq_batch, nf) + ang = (two_pi * freqs[start:stop])[:, None] * tau[None, :] + cos_sum[start:stop] = xp.sum(strength_row * xp.cos(ang), axis=1) + sin_sum[start:stop] = xp.sum(strength_row * xp.sin(ang), axis=1) + return cos_sum, sin_sum + + +def lombscargle_power_torch( + t: FloatArray, + y: FloatArray, + dy: FloatArray | None, + f0: float, + df: float, + nf: int, + *, + fit_mean: bool = True, + device: str = "cpu", + precision: str = "auto", + freq_batch: int = 4096, +) -> FloatArray: + """GLS power on ``f0 + df*arange(nf)`` via the portable torch direct trig-sum path. + + Numerically matches :func:`lombscargle_power` (to the working precision): same + Zechmeister-Kürster assembly, just NUFFT-free trig sums so it runs on any torch + device (CUDA/ROCm/MPS/XPU/CPU). ``precision="auto"`` is float64 except on MPS, where + float64 is impossible and float32 is used. Returns numpy float64. + """ + if nf <= 0: + return np.zeros(0, dtype=np.float64) + import torch + + tau, w, y, y_mean, yy = _prep(t, y, dy) + prec = resolve_precision(precision, device) + tdtype = torch.float32 if prec == "float32" else torch.float64 + tau_d = to_device_array(tau, device=device, dtype=tdtype) + w_d = to_device_array(w, device=device, dtype=tdtype) + wy_d = to_device_array(w * y, device=device, dtype=tdtype) + xp = array_namespace(tau_d) + c, s = _trig_sums_direct(xp, tau_d, w_d, f0, df, nf, freq_batch=freq_batch) + yc, ys = _trig_sums_direct(xp, tau_d, wy_d, f0, df, nf, freq_batch=freq_batch) + c2, s2 = _trig_sums_direct( + xp, tau_d, w_d, 2.0 * f0, 2.0 * df, nf, freq_batch=freq_batch + ) + power = _assemble_power_parts(xp, c, s, yc, ys, c2, s2, y_mean, yy, fit_mean) + return to_host(power) + + class CufinufftGLS: """Plan-reusing GPU Lomb-Scargle for batch processing. @@ -321,7 +407,10 @@ class GLSMethod(PeriodogramMethod): settings_cls: ClassVar[type] = GLSSettings cpu_backend: ClassVar[str] = "finufft" gpu_backend: ClassVar[str | None] = "cufinufft" - all_backends: ClassVar[tuple[str, ...]] = ("finufft", "cufinufft", "astropy") + portable_gpu_backend: ClassVar[str | None] = "torch" + all_backends: ClassVar[tuple[str, ...]] = ( + "finufft", "cufinufft", "torch", "astropy", + ) def default_grid(self, lc: LightCurve, settings: GLSSettings) -> GridSpec: # type: ignore[override] finite = lc.finite() @@ -371,15 +460,26 @@ def power( # type: ignore[override] ls.power(frequency, normalization="standard"), dtype=np.float64 ) else: - actual_backend = backend f0, df, nf = grid.uniform_frequency_params() frequency = f0 + df * np.arange(nf, dtype=np.float64) - if backend == "cufinufft" and engine is not None: + if backend == "torch" or backend.startswith("torch:"): + device = resolve_torch_device(backend, settings.device) + actual_backend = f"torch:{device}" + power = lombscargle_power_torch( + finite.time, finite.value, finite.error, f0, df, nf, + fit_mean=settings.fit_mean, + device=device, + precision=settings.precision, + freq_batch=settings.direct_freq_batch, + ) + elif backend == "cufinufft" and engine is not None: + actual_backend = backend power = engine.power( # type: ignore[attr-defined] finite.time, finite.value, finite.error, f0, df, nf, fit_mean=settings.fit_mean, ) else: + actual_backend = backend power = lombscargle_power( finite.time, finite.value, finite.error, f0, df, nf, fit_mean=settings.fit_mean, @@ -435,4 +535,5 @@ def estimate_device_bytes(self, n_points: int) -> int: "GLSMethod", "NufftBackend", "lombscargle_power", + "lombscargle_power_torch", ] diff --git a/src/cuperiod/methods/tls.py b/src/cuperiod/methods/tls.py index 04143c8..674e637 100644 --- a/src/cuperiod/methods/tls.py +++ b/src/cuperiod/methods/tls.py @@ -95,7 +95,7 @@ def _period_grid(baseline: float, settings: TLSSettings) -> FloatArray: freq = np.arange(1.0 / max_period, 1.0 / settings.min_period_days, df) if freq.size == 0: return np.zeros(0, dtype=np.float64) - return (1.0 / freq[::-1]).copy() + return np.ascontiguousarray(1.0 / freq[::-1], dtype=np.float64) def _matched_filter( diff --git a/tests/conftest.py b/tests/conftest.py index 8815a35..e00105c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2,9 +2,17 @@ from __future__ import annotations -import importlib.util +import os -import pytest +# Windows: torch and numpy/scipy (MKL) each vendor an OpenMP runtime; importing torch +# after numpy aborts with "OMP: Error #15" unless the duplicate load is allowed. We set +# it in the test process (before torch is imported) so the torch-CPU tests run; it is +# harmless elsewhere. The library itself does not set this — see the install docs. +os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE") + +import importlib.util # noqa: E402 + +import pytest # noqa: E402 def _importable(name: str) -> bool: @@ -19,3 +27,10 @@ def _importable(name: str) -> bool: not (_importable("cupy") and _importable("cufinufft")), reason="GPU stack (cupy/cufinufft) not installed", ) + +#: Skip a test unless PyTorch (the ``[torch]`` extra) is importable. The torch-CPU path +#: always runs once torch is installed; GPU-device tests add their own skips in PR2. +requires_torch = pytest.mark.skipif( + not _importable("torch"), + reason="torch (the [torch] extra) is not installed", +) diff --git a/tests/test_arrayapi.py b/tests/test_arrayapi.py new file mode 100644 index 0000000..56d7e6b --- /dev/null +++ b/tests/test_arrayapi.py @@ -0,0 +1,138 @@ +"""Array-API shim, torch device detection, and vendor-aware backend resolution. + +These are the new correctness-critical primitives behind the portable (torch) compute +paths. The pure-shim and numpy tests run everywhere; torch-specific tests skip without +the ``[torch]`` extra. No GPU is required (the torch path runs on the CPU device). +""" + +from __future__ import annotations + +import numpy as np +import pytest + +from conftest import requires_torch +from cuperiod.core._arrayapi import ( + array_namespace, + device_of, + is_torch_array, + resolve_precision, + resolve_torch_device, + scatter_add, + to_host, +) +from cuperiod.core.backend import ( + available_backends, + best_torch_device, + torch_available, + torch_devices, +) +from cuperiod.core.errors import BackendUnavailableError +from cuperiod.methods.bls import BLSMethod +from cuperiod.methods.gls import GLSMethod + +# --- shim primitives --------------------------------------------------------- + + +def test_scatter_add_numpy_accumulates_repeats() -> None: + target = np.zeros(4) + scatter_add(target, np.array([0, 0, 1, 3]), np.array([1.0, 2.0, 3.0, 4.0])) + assert target.tolist() == [3.0, 3.0, 0.0, 4.0] + + +def test_to_host_from_numpy_is_float64() -> None: + out = to_host(np.arange(3, dtype=np.float32)) + assert out.dtype == np.float64 + assert out.tolist() == [0.0, 1.0, 2.0] + + +def test_device_of_numpy_is_cpu() -> None: + assert device_of(np.zeros(3)) == "cpu" + assert not is_torch_array(np.zeros(3)) + + +def test_resolve_precision_defaults() -> None: + assert resolve_precision("auto", "cpu") == "float64" + assert resolve_precision("auto", "cuda") == "float64" + assert resolve_precision("auto", "xpu") == "float64" + # MPS cannot do float64, so auto downgrades there (and only there). + assert resolve_precision("auto", "mps") == "float32" + assert resolve_precision("float32", "cpu") == "float32" + assert resolve_precision("float64", "cuda") == "float64" + + +def test_resolve_precision_float64_on_mps_raises() -> None: + with pytest.raises(BackendUnavailableError): + resolve_precision("float64", "mps") + + +# --- torch detection --------------------------------------------------------- + + +def test_torch_available_returns_bool() -> None: + assert isinstance(torch_available(), bool) + + +@requires_torch +def test_torch_registers_as_backend() -> None: + assert "torch" in available_backends() + + +@requires_torch +def test_torch_devices_include_cpu() -> None: + devices = torch_devices() + assert "cpu" in devices + assert best_torch_device() in devices + + +@requires_torch +def test_array_namespace_and_scatter_add_torch() -> None: + import torch + + xp = array_namespace(torch.zeros(3)) + assert "torch" in xp.__name__ + + target = torch.zeros(4, dtype=torch.float64) + values = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float64) + scatter_add(target, torch.tensor([0, 0, 1, 3]), values) + assert target.tolist() == [3.0, 3.0, 0.0, 4.0] + + +@requires_torch +def test_device_of_and_to_host_torch() -> None: + import torch + + t = torch.arange(3, dtype=torch.float32) + assert device_of(t) == "cpu" + assert is_torch_array(t) + out = to_host(t) + assert out.dtype == np.float64 + assert out.tolist() == [0.0, 1.0, 2.0] + + +# --- backend resolution ------------------------------------------------------ + + +@requires_torch +def test_resolve_torch_device_explicit_and_default() -> None: + assert resolve_torch_device("torch:cpu", "auto") == "cpu" + assert resolve_torch_device("torch", "cpu") == "cpu" + + +@requires_torch +def test_gls_resolves_torch_backends() -> None: + method = GLSMethod() + assert method.resolve_backend("torch") == "torch" + assert method.resolve_backend("torch:cpu") == "torch:cpu" + assert method.is_gpu_backend("torch:cuda") + assert not method.is_gpu_backend("torch:cpu") + + +@requires_torch +def test_bls_resolves_torch_backend() -> None: + assert BLSMethod().resolve_backend("torch:cpu") == "torch:cpu" + + +@requires_torch +def test_resolve_unknown_torch_device_raises() -> None: + with pytest.raises(BackendUnavailableError): + GLSMethod().resolve_backend("torch:bogus") diff --git a/tests/test_bls.py b/tests/test_bls.py index 7928f6d..d8dd4a7 100644 --- a/tests/test_bls.py +++ b/tests/test_bls.py @@ -6,7 +6,7 @@ import pytest import cuperiod as cup -from conftest import requires_gpu +from conftest import requires_gpu, requires_torch from cuperiod.methods._bls_core import bls_power from cuperiod.methods.bls import BLSMethod from synth import synthetic_eclipser @@ -103,3 +103,25 @@ def test_bls_gpu_matches_numpy() -> None: gpu = method.power(grid, lc, settings, "cupy") finite = np.isfinite(cpu.power) & np.isfinite(gpu.power) assert float(np.max(np.abs(cpu.power[finite] - gpu.power[finite]))) < 1e-7 + + +@requires_torch +def test_bls_torch_cpu_matches_numpy() -> None: + # The torch box search shares the array-API body with numpy: bit-for-bit close. + grid, lc, settings = _grid_and_lc() + method = BLSMethod() + ref = method.power(grid, lc, settings, "numpy") + tor = method.power(grid, lc, settings, "torch:cpu") + assert tor.backend == "torch:cpu" + finite = np.isfinite(ref.power) & np.isfinite(tor.power) + assert float(np.max(np.abs(ref.power[finite] - tor.power[finite]))) < 1e-7 + assert int(np.argmax(ref.power)) == int(np.argmax(tor.power)) + + +@requires_torch +def test_bls_torch_recovers_period_and_depth() -> None: + t, flux, err = synthetic_eclipser(period=2.5, depth=0.05) + pg = cup.periodogram((t, flux, err), "BLS", domain=cup.Domain.FLUX, backend="torch") + peak = pg.best_periods(1, alias_diverse=True)[0] + assert peak.period == pytest.approx(2.5, rel=3e-3) + assert peak.extra["depth"] == pytest.approx(0.05, abs=0.01) diff --git a/tests/test_gls.py b/tests/test_gls.py index bdb49c1..b0455bc 100644 --- a/tests/test_gls.py +++ b/tests/test_gls.py @@ -7,7 +7,7 @@ from astropy.timeseries import LombScargle import cuperiod as cup -from conftest import requires_gpu +from conftest import requires_gpu, requires_torch from cuperiod.methods.gls import GLSMethod, lombscargle_power from synth import synthetic_sine @@ -87,3 +87,20 @@ def test_gls_gpu_matches_cpu() -> None: cpu = cup.periodogram((t, mag, err), "GLS", backend="finufft") gpu = cup.periodogram((t, mag, err), "GLS", backend="cufinufft") assert float(np.max(np.abs(cpu.power - gpu.power))) < 1e-7 + + +@requires_torch +def test_gls_torch_cpu_matches_finufft() -> None: + # The portable direct trig-sum path must match the NUFFT path in float64. + t, mag, err = synthetic_sine() + fi = cup.periodogram((t, mag, err), "GLS", backend="finufft") + to = cup.periodogram((t, mag, err), "GLS", backend="torch:cpu") + assert to.backend == "torch:cpu" + assert float(np.max(np.abs(fi.power - to.power))) < 1e-6 + + +@requires_torch +def test_gls_torch_recovers_period() -> None: + t, mag, err = synthetic_sine(period=0.6234) + pg = cup.periodogram((t, mag, err), "GLS", backend="torch") + assert pg.best_period() == pytest.approx(0.6234, rel=1e-3) From 36e9904d72e74867827de8910261b37348842036 Mon Sep 17 00:00:00 2001 From: Tharindu Jayasinghe Date: Tue, 30 Jun 2026 12:10:05 -0700 Subject: [PATCH 2/6] bench: time the portable torch backend for GLS + BLS MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fold the torch backend into the permanent benchmark suite so the cross-vendor path is tracked alongside the CPU and CUDA paths. - benchmark.py: bench_single + both scaling sweeps now time backend="torch" (device recorded in torch_backend) for methods that support it (GLS, BLS). New safe_best_time() records NaN for a backend that's unavailable on the host (no CUDA GPU, or no torch) instead of aborting the sweep — so the suite also runs CPU-only, where torch-CPU vs finufft/numba is now a meaningful comparison. Trial-period count (BLS/TLS) read from the CPU run (backend-independent). - make_report.py: §4 performance table surfaces torch_s / torch_backend; the numeric formatters are NaN-safe (blank cell when a backend is absent). - README: note the torch column in the performance benchmark. Verified on CPU-only: GLS cpu(finufft)=19ms, gpu=— (graceful), torch:cpu=81ms. Co-Authored-By: Claude Opus 4.8 --- benchmarks/README.md | 6 ++++ benchmarks/benchmark.py | 74 ++++++++++++++++++++++++++++++++------- benchmarks/make_report.py | 24 ++++++++----- 3 files changed, 82 insertions(+), 22 deletions(-) diff --git a/benchmarks/README.md b/benchmarks/README.md index 3371b79..43e8299 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -18,6 +18,12 @@ The rendered results are in **[REPORT.md](REPORT.md)** with figures in `figures/ | Period recovery | all 7 | VSX literature period | find the real period? | | Performance | all 7 | astropy / PyAstronomy | how much faster, and how does it scale? | +The performance benchmark also times the **portable `torch` backend** (`backend="torch"`; +the resolved device is shown in the `torch_backend` column) for the ported methods — GLS +and BLS — alongside the CPU and CUDA paths, so the cross-vendor path (AMD/Intel/Mac/CPU) is +tracked. A backend absent on the host (no CUDA GPU, or no torch) is recorded blank rather +than failing the sweep. + ## Data * **`dataset/light_curves.parquet`** — 72 real ASAS-SN g-band light curves diff --git a/benchmarks/benchmark.py b/benchmarks/benchmark.py index 79ffd7a..3517ee9 100644 --- a/benchmarks/benchmark.py +++ b/benchmarks/benchmark.py @@ -26,6 +26,9 @@ import cuperiod as cup # noqa: E402 from _common import RESULTS, load_dataset # noqa: E402 +from cuperiod.core.backend import torch_available # noqa: E402 +from cuperiod.core.errors import BackendUnavailableError # noqa: E402 +from cuperiod.methods.base import get_method # noqa: E402 # bounded period window for the box methods (BLS/TLS) — independent of the star, # so the trial-period count can never blow up on a short-period target. @@ -60,6 +63,23 @@ def best_time(fn, repeat=3): return min(_timed(fn) for _ in range(repeat)) +def safe_best_time(fn, repeat=3): + """:func:`best_time`, but a backend unavailable here records NaN instead of raising. + + Lets the sweep run on machines missing a backend — no CUDA GPU (the ``gpu`` column), + or no torch / no torch GPU device (the ``torch`` column) — leaving that cell blank. + """ + try: + return best_time(fn, repeat) + except BackendUnavailableError: + return float("nan") + + +def supports_torch(method): + """Whether ``method`` has the portable torch backend and torch is importable.""" + return torch_available() and "torch" in get_method(method).all_backends + + def _timed(fn): t0 = time.perf_counter(); fn(); return time.perf_counter() - t0 @@ -91,13 +111,23 @@ def bench_single(t, y, e): for m, st in FREQ_SETTINGS.items(): be = cup.periodogram((t, y, e), m, backend="cpu", grid=grid, settings=st()).backend tc = best_time(lambda: cup.periodogram((t, y, e), m, backend="cpu", grid=grid, settings=st())) - tg = best_time(lambda: cup.periodogram((t, y, e), m, backend="gpu", grid=grid, settings=st())) + tg = safe_best_time(lambda: cup.periodogram((t, y, e), m, backend="gpu", grid=grid, settings=st())) + has_torch = supports_torch(m) + tt = (safe_best_time(lambda: cup.periodogram((t, y, e), m, backend="torch", grid=grid, settings=st())) + if has_torch else np.nan) + tbe = (cup.periodogram((t, y, e), m, backend="torch", grid=grid, settings=st()).backend + if has_torch and np.isfinite(tt) else "—") tr = best_time(reftime[m], repeat=1) if m in reftime else np.nan rows.append(dict(method=m, n_grid=grid.size, cpu_backend=be, cpu_s=tc, gpu_s=tg, - ref_s=tr, ref=refname.get(m, "—"), gpu_speedup=tc / tg, - gpu_vs_ref=(tr / tg if np.isfinite(tr) else np.nan), + torch_s=tt, torch_backend=tbe, + ref_s=tr, ref=refname.get(m, "—"), + gpu_speedup=(tc / tg if np.isfinite(tg) else np.nan), + torch_speedup=(tc / tt if np.isfinite(tt) else np.nan), + gpu_vs_ref=(tr / tg if np.isfinite(tr) and np.isfinite(tg) else np.nan), cpu_vs_ref=(tr / tc if np.isfinite(tr) else np.nan))) - print(f" {m:12s} cpu({be})={tc:.3f}s gpu={tg:.4f}s ({tc/tg:.0f}x)", flush=True) + gstr = f"gpu={tg:.4f}s ({tc/tg:.0f}x)" if np.isfinite(tg) else "gpu=—" + tstr = f" torch({tbe})={tt:.3f}s" if np.isfinite(tt) else "" + print(f" {m:12s} cpu({be})={tc:.3f}s {gstr}{tstr}", flush=True) # box methods on a fixed, bounded period window. cpu_s is cuPeriod's *default* # CPU backend: "cpu" resolves to the multicore numba box search for BLS (or # astropy if numba is not installed), to numpy for TLS. @@ -107,8 +137,13 @@ def bench_single(t, y, e): cup.TLSSettings(min_period_days=BOX_PMIN, max_period_days=BOX_PMAX)) be = cup.periodogram((t, y, e), m, backend="cpu", settings=S).backend tc = best_time(lambda: cup.periodogram((t, y, e), m, backend="cpu", settings=S), repeat=2) - tg = best_time(lambda: cup.periodogram((t, y, e), m, backend="gpu", settings=S), repeat=2) - ng = cup.periodogram((t, y, e), m, backend="gpu", settings=S).power.size + tg = safe_best_time(lambda: cup.periodogram((t, y, e), m, backend="gpu", settings=S), repeat=2) + ng = cup.periodogram((t, y, e), m, backend="cpu", settings=S).power.size # backend-independent + has_torch = supports_torch(m) + tt = (safe_best_time(lambda: cup.periodogram((t, y, e), m, backend="torch", settings=S), repeat=2) + if has_torch else np.nan) + tbe = (cup.periodogram((t, y, e), m, backend="torch", settings=S).backend + if has_torch and np.isfinite(tt) else "—") # BLS reference = astropy's compiled BoxLeastSquares; also time the pure-numpy # GPU-parity reference port to document it is not the product path. ref_s = (best_time(lambda: cup.periodogram((t, y, e), "BLS", backend="astropy", settings=S), repeat=1) @@ -117,13 +152,18 @@ def bench_single(t, y, e): port_s = (best_time(lambda: cup.periodogram((t, y, e), "BLS", backend="numpy", settings=S), repeat=1) if m == "BLS" else np.nan) rows.append(dict(method=m, n_grid=int(ng), cpu_backend=be, cpu_s=tc, gpu_s=tg, - ref_s=ref_s, ref=ref, gpu_speedup=tc / tg, - gpu_vs_ref=(ref_s / tg if np.isfinite(ref_s) else np.nan), + torch_s=tt, torch_backend=tbe, + ref_s=ref_s, ref=ref, + gpu_speedup=(tc / tg if np.isfinite(tg) else np.nan), + torch_speedup=(tc / tt if np.isfinite(tt) else np.nan), + gpu_vs_ref=(ref_s / tg if np.isfinite(ref_s) and np.isfinite(tg) else np.nan), cpu_vs_ref=(ref_s / tc if np.isfinite(ref_s) else np.nan), cpu_port_s=port_s)) + gstr = f"gpu={tg:.4f}s (gpu {tc/tg:.0f}x)" if np.isfinite(tg) else "gpu=—" + tstr = f" torch({tbe})={tt:.3f}s" if np.isfinite(tt) else "" extra = (f" [vs astropy {ref_s/tc:.0f}x faster; numpy-port {port_s:.1f}s]" if m == "BLS" else "") - print(f" {m:12s} cpu({be})={tc:.3f}s gpu={tg:.4f}s (gpu {tc/tg:.0f}x){extra} [{ng:,} periods]", + print(f" {m:12s} cpu({be})={tc:.3f}s {gstr}{tstr}{extra} [{ng:,} periods]", flush=True) df = pd.DataFrame(rows) df.to_parquet(RESULTS / "bench_single.parquet", index=False) @@ -138,8 +178,12 @@ def bench_scaling_npoints(t, y, e): for m in ["GLS", "PDM", "MHAOV"]: st = FREQ_SETTINGS[m] tc = best_time(lambda: cup.periodogram((tt, yy, ee), m, backend="cpu", grid=grid, settings=st()), repeat=1) - tg = best_time(lambda: cup.periodogram((tt, yy, ee), m, backend="gpu", grid=grid, settings=st()), repeat=1) - rows.append(dict(axis="npoints", method=m, n=n, cpu_s=tc, gpu_s=tg, speedup=tc / tg)) + tg = safe_best_time(lambda: cup.periodogram((tt, yy, ee), m, backend="gpu", grid=grid, settings=st()), repeat=1) + ttor = (safe_best_time(lambda: cup.periodogram((tt, yy, ee), m, backend="torch", grid=grid, settings=st()), repeat=1) + if supports_torch(m) else np.nan) + rows.append(dict(axis="npoints", method=m, n=n, cpu_s=tc, gpu_s=tg, torch_s=ttor, + speedup=(tc / tg if np.isfinite(tg) else np.nan), + torch_speedup=(tc / ttor if np.isfinite(ttor) else np.nan))) print(f" N={n:>6}: done", flush=True) df = pd.DataFrame(rows) df.to_parquet(RESULTS / "bench_npoints.parquet", index=False) @@ -153,8 +197,12 @@ def bench_scaling_grid(t, y, e): for m in ["GLS", "PDM", "MHAOV"]: st = FREQ_SETTINGS[m] tc = best_time(lambda: cup.periodogram((t, y, e), m, backend="cpu", grid=grid, settings=st()), repeat=1) - tg = best_time(lambda: cup.periodogram((t, y, e), m, backend="gpu", grid=grid, settings=st()), repeat=1) - rows.append(dict(axis="grid", method=m, n=n, cpu_s=tc, gpu_s=tg, speedup=tc / tg)) + tg = safe_best_time(lambda: cup.periodogram((t, y, e), m, backend="gpu", grid=grid, settings=st()), repeat=1) + ttor = (safe_best_time(lambda: cup.periodogram((t, y, e), m, backend="torch", grid=grid, settings=st()), repeat=1) + if supports_torch(m) else np.nan) + rows.append(dict(axis="grid", method=m, n=n, cpu_s=tc, gpu_s=tg, torch_s=ttor, + speedup=(tc / tg if np.isfinite(tg) else np.nan), + torch_speedup=(tc / ttor if np.isfinite(ttor) else np.nan))) print(f" grid={n:>7}: done", flush=True) df = pd.DataFrame(rows) df.to_parquet(RESULTS / "bench_grid.parquet", index=False) diff --git a/benchmarks/make_report.py b/benchmarks/make_report.py index 8ab2be8..f9c7001 100644 --- a/benchmarks/make_report.py +++ b/benchmarks/make_report.py @@ -432,21 +432,27 @@ def main(): f"({b.cpu_s*1e3:.0f} ms vs {b.ref_s:.1f} s on this light curve), matching it " f"to floating-point{par}. The GPU then adds another {b.gpu_speedup:.0f}× " f"({b.ref_s/b.gpu_s:.0f}× over astropy).\n") - cols = ["method", "cpu_backend", "cpu_s", "gpu_s", "ref", "ref_s", "cpu_vs_ref", "gpu_speedup"] + cols = ["method", "cpu_backend", "cpu_s", "gpu_s", "torch_s", "torch_backend", + "ref", "ref_s", "cpu_vs_ref", "gpu_speedup"] cols = [c for c in cols if c in s.columns] + nan_dash = lambda fmt: (lambda v: ("—" if not np.isfinite(v) else fmt(v))) L.append(md_table(s, cols, { - "cpu_s": lambda v: f"{v:.3f}", "gpu_s": lambda v: f"{v:.4f}", - "gpu_speedup": lambda v: f"{v:.0f}x", - "cpu_vs_ref": lambda v: ("—" if not np.isfinite(v) else f"{v:.0f}x"), - "ref_s": lambda v: ("—" if not np.isfinite(v) else f"{v:.2f}"), + "cpu_s": lambda v: f"{v:.3f}", + "gpu_s": nan_dash(lambda v: f"{v:.4f}"), + "torch_s": nan_dash(lambda v: f"{v:.3f}"), + "gpu_speedup": nan_dash(lambda v: f"{v:.0f}x"), + "cpu_vs_ref": nan_dash(lambda v: f"{v:.0f}x"), + "ref_s": nan_dash(lambda v: f"{v:.2f}"), "method": ml})) L.append("\n*cpu_backend* = what `backend=\"cpu\"` resolves to — the fast default a user " "gets: finufft (GLS), the multicore numba box search (BLS), numpy (the rest). " "*ref* = the established external tool; *cpu_vs_ref* = how much faster cuPeriod's " - "CPU is than that tool; *gpu_speedup* = GPU over cuPeriod's CPU. cuPeriod's CPU " - "path already beats every reference tool it has (GLS, PDM, BLS) — so the GPU's " - "marginal gain is small where the CPU is already fast (BLS, GLS) and large where " - "it is not (PDM, MHAOV, TLS).\n") + "CPU is than that tool; *gpu_speedup* = GPU over cuPeriod's CPU. *torch_s* = the " + "portable PyTorch backend (device shown in *torch_backend*: cpu/cuda/mps/xpu) — " + "the cross-vendor path that also runs on AMD/Intel/Mac; blank for methods not yet " + "ported to it. cuPeriod's CPU path already beats every reference tool it has " + "(GLS, PDM, BLS) — so the GPU's marginal gain is small where the CPU is already " + "fast (BLS, GLS) and large where it is not (PDM, MHAOV, TLS).\n") if len(bls) and "cpu_port_s" in bls and np.isfinite(bls.cpu_port_s.iloc[0]): b = bls.iloc[0] L.append(f"\n> The pure-`numpy` BLS backend shares one array-module-generic source " From 6f5c79b5f2daed0068213239b51cc8187c08f5df Mon Sep 17 00:00:00 2001 From: Tharindu Jayasinghe Date: Tue, 30 Jun 2026 12:30:58 -0700 Subject: [PATCH 3/6] feat(methods): port PDM, CE, String-Length, MHAOV, TLS to the torch backend Extends the portable array-API path from GLS/BLS (PR1) to the remaining five methods, so all seven now run on any torch device (CUDA/ROCm/MPS/XPU/CPU) with the NVIDIA fast paths untouched. - CE / PDM: vectorized histogram/box kernels re-spelled to array-API ops (scatter_add shim, xp.astype/remainder/clip, dtype-following-periods); numpy + torch flow through array_api_compat namespaces. cupy RawKernels unchanged. - String-Length: sort/gather re-spelled with fancy indexing + slicing (no take_along_axis/diff, which aren't in every namespace). - MHAOV: einsum + linalg.solve are available in the compat namespaces, so the torch path is a thin add; raw numpy/cupy keep einsum. - TLS: matched filter re-spelled (notably explicit float dtype on every zeros(), since torch defaults to float32; clip instead of minimum-with-scalar, which torch rejects). - config: device/precision mixin + "torch" backend on all five settings. - New tests/test_torch_bonus.py: torch:cpu vs numpy parity + period recovery for all five (harmonic-aware, since String-Length/MHAOV legitimately lock onto a harmonic identically on both backends). Validated on torch:cpu: CE 1.3e-15, String-Length 2.8e-14, PDM 0.0, TLS 0.0 (bit-identical); MHAOV 2.7e-10 relative (linalg.solve rounding, peak unaffected). 160 passed, 8 skipped (CUDA-only); ruff + mypy clean. Co-Authored-By: Claude Opus 4.8 --- CHANGELOG.md | 18 ++-- src/cuperiod/core/config.py | 20 ++-- src/cuperiod/methods/conditional_entropy.py | 73 +++++++++---- src/cuperiod/methods/mhaov.py | 58 ++++++++--- src/cuperiod/methods/pdm.py | 87 +++++++++++----- src/cuperiod/methods/string_length.py | 54 +++++++--- src/cuperiod/methods/tls.py | 110 +++++++++++++------- tests/test_torch_bonus.py | 66 ++++++++++++ 8 files changed, 357 insertions(+), 129 deletions(-) create mode 100644 tests/test_torch_bonus.py diff --git a/CHANGELOG.md b/CHANGELOG.md index f8b9852..acd6b9f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,16 +8,18 @@ All notable changes to cuPeriod are documented here. The format is based on ### Added -- **Multi-vendor GPU support via PyTorch and the Python array API.** GLS and BLS gain a - portable `torch` backend that runs on AMD (ROCm), Intel (XPU), and Apple (MPS) GPUs as - well as a real CPU path — so the accelerated code is no longer NVIDIA-only, and works - even with no GPU at all. Select it with `backend="torch"` (or `"torch:cpu"`, - `"torch:cuda"`, `"torch:mps"`, `"torch:xpu"`); `backend="auto"` now reaches a torch GPU - on non-NVIDIA machines after the cufinufft/cupy fast paths. +- **Multi-vendor GPU support via PyTorch and the Python array API.** All seven + period-search methods (GLS, BLS, PDM, CE, String-Length, MHAOV, TLS) gain a portable + `torch` backend that runs on AMD (ROCm), Intel (XPU), and Apple (MPS) GPUs as well as a + real CPU path — so the accelerated code is no longer NVIDIA-only, and works even with no + GPU at all. Select it with `backend="torch"` (or `"torch:cpu"`, `"torch:cuda"`, + `"torch:mps"`, `"torch:xpu"`); `backend="auto"` now reaches a torch GPU on non-NVIDIA + machines after the cufinufft/cupy fast paths. - GLS adds a NUFFT-free direct trig-sum path (the portable formulation; cufinufft remains the NVIDIA fast path). - - BLS runs its vectorized box search through the array-API namespace (the cupy - `RawKernel` remains the NVIDIA fast path). + - BLS, PDM, CE, String-Length, MHAOV, and TLS run their vectorized kernels through the + array-API namespace; the cupy `RawKernel`s (BLS/PDM/CE/TLS), numba (BLS), and finufft + (GLS) remain the fast paths where present. - New `device` and `precision` settings: `precision="auto"` is float64 everywhere it is supported and float32 only where the device forces it (Apple MPS cannot do float64); an explicit `precision="float64"` on MPS raises rather than silently downgrading. diff --git a/src/cuperiod/core/config.py b/src/cuperiod/core/config.py index 73c3699..0cd78d9 100644 --- a/src/cuperiod/core/config.py +++ b/src/cuperiod/core/config.py @@ -175,7 +175,7 @@ def _check_bounds(self) -> Self: ) -class PDMSettings(BaseSettings): +class PDMSettings(_DeviceSettings): """Settings for phase dispersion minimization (PDM).""" model_config = SettingsConfigDict(env_prefix="CUPERIOD_PDM_", extra="forbid") @@ -213,7 +213,7 @@ def _check_bounds(self) -> Self: min_detections: int = Field( default=20, ge=3, description="Skip if fewer finite points." ) - backend: Literal["auto", "cpu", "gpu", "numpy", "cupy"] = Field( + backend: Literal["auto", "cpu", "gpu", "numpy", "cupy", "torch"] = Field( default="auto", description="Compute backend." ) batch_periods: int = Field( @@ -224,7 +224,7 @@ def _check_bounds(self) -> Self: ) -class MHAOVSettings(BaseSettings): +class MHAOVSettings(_DeviceSettings): """Settings for the multiharmonic Analysis of Variance (MHAOV) periodogram.""" model_config = SettingsConfigDict(env_prefix="CUPERIOD_MHAOV_", extra="forbid") @@ -261,7 +261,7 @@ def _check_bounds(self) -> Self: min_detections: int = Field( default=20, ge=5, description="Skip if fewer finite points (need > 2H+1)." ) - backend: Literal["auto", "cpu", "gpu", "numpy", "cupy"] = Field( + backend: Literal["auto", "cpu", "gpu", "numpy", "cupy", "torch"] = Field( default="auto", description="Compute backend." ) batch_periods: int = Field( @@ -272,7 +272,7 @@ def _check_bounds(self) -> Self: ) -class CESettings(BaseSettings): +class CESettings(_DeviceSettings): """Settings for the conditional-entropy (CE) period search.""" model_config = SettingsConfigDict(env_prefix="CUPERIOD_CE_", extra="forbid") @@ -308,7 +308,7 @@ def _check_bounds(self) -> Self: min_detections: int = Field( default=20, ge=3, description="Skip if fewer finite points." ) - backend: Literal["auto", "cpu", "gpu", "numpy", "cupy"] = Field( + backend: Literal["auto", "cpu", "gpu", "numpy", "cupy", "torch"] = Field( default="auto", description="Compute backend." ) batch_periods: int = Field( @@ -319,7 +319,7 @@ def _check_bounds(self) -> Self: ) -class StringLengthSettings(BaseSettings): +class StringLengthSettings(_DeviceSettings): """Settings for the string-length (Lafler-Kinman / Dworetsky) period search.""" model_config = SettingsConfigDict(env_prefix="CUPERIOD_SL_", extra="forbid") @@ -353,7 +353,7 @@ def _check_bounds(self) -> Self: min_detections: int = Field( default=20, ge=3, description="Skip if fewer finite points." ) - backend: Literal["auto", "cpu", "gpu", "numpy", "cupy"] = Field( + backend: Literal["auto", "cpu", "gpu", "numpy", "cupy", "torch"] = Field( default="auto", description="Compute backend." ) batch_periods: int = Field( @@ -364,7 +364,7 @@ def _check_bounds(self) -> Self: ) -class TLSSettings(BaseSettings): +class TLSSettings(_DeviceSettings): """Settings for the transit least squares (TLS) search.""" model_config = SettingsConfigDict(env_prefix="CUPERIOD_TLS_", extra="forbid") @@ -425,7 +425,7 @@ def _check_bounds(self) -> Self: min_detections: int = Field( default=20, ge=3, description="Skip if fewer finite points." ) - backend: Literal["auto", "cpu", "gpu", "numpy", "cupy"] = Field( + backend: Literal["auto", "cpu", "gpu", "numpy", "cupy", "torch"] = Field( default="auto", description="Compute backend." ) period_batch: int = Field( diff --git a/src/cuperiod/methods/conditional_entropy.py b/src/cuperiod/methods/conditional_entropy.py index 3e9fa8d..f581039 100644 --- a/src/cuperiod/methods/conditional_entropy.py +++ b/src/cuperiod/methods/conditional_entropy.py @@ -17,6 +17,14 @@ import numpy as np +from cuperiod.core._arrayapi import ( + array_namespace, + resolve_precision, + resolve_torch_device, + scatter_add, + to_device_array, + to_host, +) from cuperiod.core._typing import FloatArray, IntArray from cuperiod.core.backend import ensure_cuda_dll_path from cuperiod.core.config import CESettings @@ -46,33 +54,39 @@ def _entropy_batch( n_mag: int, batch: int, ) -> Any: - """Conditional entropy H(m|phase) for each trial period, vectorized over periods.""" + """Conditional entropy H(m|phase) for each trial period, vectorized over periods. + + Array-API generic (numpy/cupy/torch via array_api_compat). The float dtype follows + ``periods`` (float64, or float32 on a float32 device); index/bin arrays are int64. + The cupy ``RawKernel`` fast path (NVIDIA) is separate and not this code. + """ + fdtype = periods.dtype + idtype = xp.int64 n_points = int(tau.shape[0]) n_periods = int(periods.shape[0]) n_cells = n_phase * n_mag - entropy = xp.empty(n_periods, dtype=np.float64) + entropy = xp.empty(n_periods, dtype=fdtype) for start in range(0, n_periods, batch): stop = min(start + batch, n_periods) pb = periods[start:stop] n_p = int(pb.shape[0]) - rows = xp.arange(n_p) - phase = xp.mod(tau[None, :] / pb[:, None], 1.0) - phase_bin = (phase * n_phase).astype(np.int64) - xp.clip(phase_bin, 0, n_phase - 1, out=phase_bin) + rows = xp.arange(n_p, dtype=idtype) + phase = xp.remainder(tau[None, :] / pb[:, None], 1.0) + phase_bin = xp.clip(xp.astype(phase * n_phase, idtype), 0, n_phase - 1) cell = phase_bin * n_mag + mag_bin[None, :] # (P, N) in [0, n_cells) - flat = (rows[:, None] * n_cells + cell).ravel() - count = xp.zeros(n_p * n_cells, dtype=np.float64) - xp.add.at(count, flat, xp.broadcast_to(xp.ones(1), (n_p, n_points)).ravel()) - count = count.reshape(n_p, n_phase, n_mag) + flat = xp.reshape(rows[:, None] * n_cells + cell, (-1,)) + count = xp.zeros(n_p * n_cells, dtype=fdtype) + scatter_add(count, flat, xp.ones(n_p * n_points, dtype=fdtype)) + count = xp.reshape(count, (n_p, n_phase, n_mag)) - phase_total = count.sum(axis=2, keepdims=True) # (P, n_phase, 1) + phase_total = xp.sum(count, axis=2, keepdims=True) # (P, n_phase, 1) mask = count > 0.0 safe_count = xp.where(mask, count, 1.0) safe_total = xp.where(phase_total > 0.0, phase_total, 1.0) term = xp.where( mask, count * (xp.log(safe_total) - xp.log(safe_count)), 0.0 ) - entropy[start:stop] = term.sum(axis=(1, 2)) / n_points + entropy[start:stop] = xp.sum(term, axis=(1, 2)) / n_points return entropy @@ -186,8 +200,9 @@ def conditional_entropy( *, n_phase_bins: int = 10, n_mag_bins: int = 10, - backend: CEBackend = "numpy", + backend: str = "numpy", batch: int = DEFAULT_BATCH, + precision: str = "auto", ) -> FloatArray: """Conditional entropy for each trial period (minimized at the true period). @@ -227,15 +242,25 @@ def conditional_entropy( return _ce_cuda( tau, mag_bin, periods_host, n_phase=n_phase_bins, n_mag=n_mag_bins ) + if backend == "torch" or backend.startswith("torch:"): + import torch + + device = backend.split(":", 1)[1] if ":" in backend else "cpu" + fdt = (torch.float32 + if resolve_precision(precision, device) == "float32" else torch.float64) + tau_d = to_device_array(tau, device=device, dtype=fdt) + mag_d = to_device_array(mag_bin, device=device, dtype=torch.int64) + per_d = to_device_array(periods_host, device=device, dtype=fdt) + return to_host(_entropy_batch( + array_namespace(per_d), tau_d, mag_d, per_d, + n_phase=n_phase_bins, n_mag=n_mag_bins, batch=batch, + )) if backend != "numpy": raise ValueError(f"unknown backend {backend!r}") - return np.asarray( - _entropy_batch( - np, tau, mag_bin, periods_host, - n_phase=n_phase_bins, n_mag=n_mag_bins, batch=batch, - ), - dtype=np.float64, - ) + return to_host(_entropy_batch( + array_namespace(periods_host), tau, mag_bin, periods_host, + n_phase=n_phase_bins, n_mag=n_mag_bins, batch=batch, + )) class ConditionalEntropyMethod(PeriodogramMethod): @@ -247,7 +272,8 @@ class ConditionalEntropyMethod(PeriodogramMethod): settings_cls: ClassVar[type] = CESettings cpu_backend: ClassVar[str] = "numpy" gpu_backend: ClassVar[str | None] = "cupy" - all_backends: ClassVar[tuple[str, ...]] = ("numpy", "cupy") + portable_gpu_backend: ClassVar[str | None] = "torch" + all_backends: ClassVar[tuple[str, ...]] = ("numpy", "cupy", "torch") def default_grid(self, lc: LightCurve, settings: CESettings) -> GridSpec: # type: ignore[override] finite = lc.finite() @@ -281,10 +307,13 @@ def power( # type: ignore[override] if finite.baseline <= 0.0: raise InsufficientDataError("CE: no usable time baseline") periods = grid.period + if backend == "torch" or backend.startswith("torch:"): + backend = f"torch:{resolve_torch_device(backend, settings.device)}" entropy = conditional_entropy( finite.time, finite.value, periods, n_phase_bins=settings.n_phase_bins, n_mag_bins=settings.n_mag_bins, - backend=backend, batch=settings.batch_periods, # type: ignore[arg-type] + backend=backend, batch=settings.batch_periods, + precision=settings.precision, ) return Periodogram.from_spectrum( method="CE", diff --git a/src/cuperiod/methods/mhaov.py b/src/cuperiod/methods/mhaov.py index a4b6501..0c7d850 100644 --- a/src/cuperiod/methods/mhaov.py +++ b/src/cuperiod/methods/mhaov.py @@ -25,6 +25,13 @@ import numpy as np +from cuperiod.core._arrayapi import ( + array_namespace, + resolve_precision, + resolve_torch_device, + to_device_array, + to_host, +) from cuperiod.core._typing import FloatArray from cuperiod.core.backend import ensure_cuda_dll_path from cuperiod.core.config import MHAOVSettings @@ -51,7 +58,7 @@ def _design(xp: ModuleType, angle: Any, n_harmonics: int) -> Any: """Trig-polynomial design tensor ``(F, N, 2H+1)`` = [1, cos kθ, sin kθ].""" n_freq, n_points = angle.shape d = 2 * n_harmonics + 1 - design = xp.empty((n_freq, n_points, d), dtype=np.float64) + design = xp.empty((n_freq, n_points, d), dtype=angle.dtype) design[:, :, 0] = 1.0 for k in range(1, n_harmonics + 1): design[:, :, 2 * k - 1] = xp.cos(k * angle) @@ -78,9 +85,10 @@ def _model_ss_batch( """ d = 2 * n_harmonics + 1 n_freq = int(frequencies.shape[0]) - eye = xp.eye(d, dtype=np.float64) * _RIDGE - out = xp.empty(n_freq, dtype=np.float64) - two_pi = 2.0 * np.pi + fdtype = frequencies.dtype + eye = xp.eye(d, dtype=fdtype) * _RIDGE + out = xp.empty(n_freq, dtype=fdtype) + two_pi = 2.0 * float(np.pi) for start in range(0, n_freq, batch): stop = min(start + batch, n_freq) @@ -92,7 +100,7 @@ def _model_ss_batch( # numpy 2.x batched solve treats a 2-D RHS as matrices, so add a trailing # singleton to keep it a per-frequency vector solve. beta = xp.linalg.solve(gram, proj[..., None])[..., 0] - model_ss = (beta * proj).sum(axis=1) - n_points * y_mean * y_mean + model_ss = xp.sum(beta * proj, axis=1) - n_points * y_mean * y_mean out[start:stop] = xp.clip(model_ss, 0.0, total_ss) return out @@ -103,8 +111,9 @@ def _compute_model_ss( frequencies: FloatArray, *, n_harmonics: int, - backend: MHAOVBackend, + backend: str, batch: int, + precision: str = "auto", ) -> tuple[FloatArray, float, int]: """Host-side regression SS per frequency, plus ``(total_ss, n)`` for one band. @@ -133,6 +142,21 @@ def _compute_model_ss( n_points=n, batch=batch, ) return np.asarray(cp.asnumpy(out), dtype=np.float64), total_ss, n + if backend == "torch" or backend.startswith("torch:"): + import torch + + device = backend.split(":", 1)[1] if ":" in backend else "cpu" + fdt = (torch.float32 + if resolve_precision(precision, device) == "float32" else torch.float64) + out = _model_ss_batch( + array_namespace(to_device_array(freqs, device=device, dtype=fdt)), + to_device_array(tau, device=device, dtype=fdt), + to_device_array(y, device=device, dtype=fdt), + to_device_array(freqs, device=device, dtype=fdt), + n_harmonics=n_harmonics, total_ss=total_ss, y_mean=y_mean, + n_points=n, batch=batch, + ) + return to_host(out), total_ss, n if backend != "numpy": raise ValueError(f"unknown backend {backend!r}") out = _model_ss_batch( @@ -149,8 +173,9 @@ def aov_power( frequencies: FloatArray, *, n_harmonics: int = 3, - backend: MHAOVBackend = "numpy", + backend: str = "numpy", batch: int = DEFAULT_BATCH, + precision: str = "auto", ) -> FloatArray: """Multiharmonic AOV statistic for each trial frequency. @@ -174,7 +199,8 @@ def aov_power( constant signal or when there are too few points. """ model_ss, total_ss, n = _compute_model_ss( - t, y, frequencies, n_harmonics=n_harmonics, backend=backend, batch=batch + t, y, frequencies, n_harmonics=n_harmonics, backend=backend, batch=batch, + precision=precision, ) if total_ss <= 0.0: return model_ss @@ -189,8 +215,9 @@ def aov_multiband_power( bands: list[tuple[FloatArray, FloatArray]], *, n_harmonics: int = 3, - backend: MHAOVBackend = "numpy", + backend: str = "numpy", batch: int = DEFAULT_BATCH, + precision: str = "auto", ) -> FloatArray: """Pooled multiband AOV F-statistic at a shared frequency, per-band amplitudes. @@ -222,7 +249,8 @@ def aov_multiband_power( n_used_bands = 0 for t, y in bands: model_ss, total_ss, n = _compute_model_ss( - t, y, freqs, n_harmonics=n_harmonics, backend=backend, batch=batch + t, y, freqs, n_harmonics=n_harmonics, backend=backend, batch=batch, + precision=precision, ) if total_ss <= 0.0: continue @@ -248,7 +276,8 @@ class MHAOVMethod(PeriodogramMethod): settings_cls: ClassVar[type] = MHAOVSettings cpu_backend: ClassVar[str] = "numpy" gpu_backend: ClassVar[str | None] = "cupy" - all_backends: ClassVar[tuple[str, ...]] = ("numpy", "cupy") + portable_gpu_backend: ClassVar[str | None] = "torch" + all_backends: ClassVar[tuple[str, ...]] = ("numpy", "cupy", "torch") def default_grid(self, lc: LightCurve, settings: MHAOVSettings) -> GridSpec: # type: ignore[override] finite = lc.finite() @@ -283,10 +312,13 @@ def power( # type: ignore[override] if finite.baseline <= 0.0: raise InsufficientDataError("MHAOV: no usable time baseline") frequency = grid.frequency + if backend == "torch" or backend.startswith("torch:"): + backend = f"torch:{resolve_torch_device(backend, settings.device)}" power = aov_power( finite.time, finite.value, frequency, n_harmonics=settings.n_harmonics, - backend=backend, batch=settings.batch_periods, # type: ignore[arg-type] + backend=backend, batch=settings.batch_periods, + precision=settings.precision, ) return Periodogram.from_spectrum( method="MHAOV", @@ -308,6 +340,8 @@ def multiband_power( # type: ignore[override] ) -> Periodogram: from cuperiod.multiband.mhaov_mb import mhaov_multiband_power + if backend == "torch" or backend.startswith("torch:"): + backend = f"torch:{resolve_torch_device(backend, settings.device)}" return mhaov_multiband_power(grid, mblc, settings, backend) def estimate_device_bytes(self, n_points: int) -> int: diff --git a/src/cuperiod/methods/pdm.py b/src/cuperiod/methods/pdm.py index 0720585..3580947 100644 --- a/src/cuperiod/methods/pdm.py +++ b/src/cuperiod/methods/pdm.py @@ -22,6 +22,14 @@ import numpy as np +from cuperiod.core._arrayapi import ( + array_namespace, + resolve_precision, + resolve_torch_device, + scatter_add, + to_device_array, + to_host, +) from cuperiod.core._typing import FloatArray from cuperiod.core.backend import ensure_cuda_dll_path from cuperiod.core.config import PDMSettings @@ -58,43 +66,47 @@ def _theta_batch( ``s^2 = sum_j SSD_j / (N*n_covers - n_nonempty)`` with ``SSD_j`` the within-bin sum of squared deviations across all covers, and ``Theta = s^2 / sigma^2``. """ + fdtype = periods.dtype + idtype = xp.int64 + inf = float("inf") n_points = int(tau.shape[0]) n_periods = int(periods.shape[0]) n_global = n_bins * n_covers cover_step = 1.0 / (n_bins * n_covers) - theta = xp.empty(n_periods, dtype=np.float64) + theta = xp.empty(n_periods, dtype=fdtype) for start in range(0, n_periods, batch): stop = min(start + batch, n_periods) pb = periods[start:stop] n_p = int(pb.shape[0]) - rows = xp.arange(n_p) - phase = xp.mod(tau[None, :] / pb[:, None], 1.0) # (P, N) in [0, 1) - - count = xp.zeros(n_p * n_global, dtype=np.float64) - ysum = xp.zeros(n_p * n_global, dtype=np.float64) - ysq = xp.zeros(n_p * n_global, dtype=np.float64) - ones = xp.broadcast_to(xp.ones(1), (n_p, n_points)).ravel() - y_b = xp.broadcast_to(y, (n_p, n_points)).ravel() - y2_b = xp.broadcast_to(y2, (n_p, n_points)).ravel() + rows = xp.arange(n_p, dtype=idtype) + phase = xp.remainder(tau[None, :] / pb[:, None], 1.0) # (P, N) in [0, 1) + + count = xp.zeros(n_p * n_global, dtype=fdtype) + ysum = xp.zeros(n_p * n_global, dtype=fdtype) + ysq = xp.zeros(n_p * n_global, dtype=fdtype) + ones = xp.ones(n_p * n_points, dtype=fdtype) + y_b = xp.reshape(xp.broadcast_to(y, (n_p, n_points)), (-1,)) + y2_b = xp.reshape(xp.broadcast_to(y2, (n_p, n_points)), (-1,)) for cover in range(n_covers): offset = cover * cover_step - b = (xp.mod(phase + offset, 1.0) * n_bins).astype(np.int64) - xp.clip(b, 0, n_bins - 1, out=b) - flat = (rows[:, None] * n_global + (b + cover * n_bins)).ravel() - xp.add.at(count, flat, ones) - xp.add.at(ysum, flat, y_b) - xp.add.at(ysq, flat, y2_b) - - count = count.reshape(n_p, n_global) - ysum = ysum.reshape(n_p, n_global) - ysq = ysq.reshape(n_p, n_global) - safe = xp.where(count > 0.0, count, 1.0) - ssd = xp.where(count > 0.0, ysq - ysum * ysum / safe, 0.0) - nonempty = (count > 0.0).sum(axis=1) + b = xp.astype(xp.remainder(phase + offset, 1.0) * n_bins, idtype) + b = xp.clip(b, 0, n_bins - 1) + flat = xp.reshape(rows[:, None] * n_global + (b + cover * n_bins), (-1,)) + scatter_add(count, flat, ones) + scatter_add(ysum, flat, y_b) + scatter_add(ysq, flat, y2_b) + + count = xp.reshape(count, (n_p, n_global)) + ysum = xp.reshape(ysum, (n_p, n_global)) + ysq = xp.reshape(ysq, (n_p, n_global)) + mask = count > 0.0 + safe = xp.where(mask, count, 1.0) + ssd = xp.where(mask, ysq - ysum * ysum / safe, 0.0) + nonempty = xp.sum(xp.astype(mask, fdtype), axis=1) den = float(n_points * n_covers) - nonempty safe_den = xp.where(den > 0.0, den, 1.0) - s2 = xp.where(den > 0.0, ssd.sum(axis=1) / safe_den, xp.inf) + s2 = xp.where(den > 0.0, xp.sum(ssd, axis=1) / safe_den, inf) theta[start:stop] = s2 / sigma2 return theta @@ -221,8 +233,9 @@ def pdm_theta( *, n_bins: int = 10, n_covers: int = 3, - backend: PDMBackend = "numpy", + backend: str = "numpy", batch: int = DEFAULT_BATCH, + precision: str = "auto", ) -> FloatArray: """PDM Theta statistic for each trial period. @@ -264,10 +277,24 @@ def pdm_theta( return _pdm_cuda( tau, y, periods_host, n_bins=n_bins, n_covers=n_covers, sigma2=sigma2 ) + if backend == "torch" or backend.startswith("torch:"): + import torch + + device = backend.split(":", 1)[1] if ":" in backend else "cpu" + fdt = (torch.float32 + if resolve_precision(precision, device) == "float32" else torch.float64) + tau_d = to_device_array(tau, device=device, dtype=fdt) + y_d = to_device_array(y, device=device, dtype=fdt) + y2_d = to_device_array(y * y, device=device, dtype=fdt) + per_d = to_device_array(periods_host, device=device, dtype=fdt) + return to_host(_theta_batch( + array_namespace(per_d), tau_d, y_d, y2_d, per_d, + n_bins=n_bins, n_covers=n_covers, sigma2=sigma2, batch=batch, + )) if backend != "numpy": raise ValueError(f"unknown backend {backend!r}") theta = _theta_batch( - np, tau, y, y * y, periods_host, + array_namespace(periods_host), tau, y, y * y, periods_host, n_bins=n_bins, n_covers=n_covers, sigma2=sigma2, batch=batch, ) return np.asarray(theta, dtype=np.float64) @@ -282,7 +309,8 @@ class PDMMethod(PeriodogramMethod): settings_cls: ClassVar[type] = PDMSettings cpu_backend: ClassVar[str] = "numpy" gpu_backend: ClassVar[str | None] = "cupy" - all_backends: ClassVar[tuple[str, ...]] = ("numpy", "cupy") + portable_gpu_backend: ClassVar[str | None] = "torch" + all_backends: ClassVar[tuple[str, ...]] = ("numpy", "cupy", "torch") def default_grid(self, lc: LightCurve, settings: PDMSettings) -> GridSpec: # type: ignore[override] finite = lc.finite() @@ -316,14 +344,17 @@ def power( # type: ignore[override] if finite.baseline <= 0.0: raise InsufficientDataError("PDM: no usable time baseline") periods = grid.period + if backend == "torch" or backend.startswith("torch:"): + backend = f"torch:{resolve_torch_device(backend, settings.device)}" theta = pdm_theta( finite.time, finite.value, periods, n_bins=settings.n_bins, n_covers=settings.n_covers, - backend=backend, # type: ignore[arg-type] + backend=backend, batch=settings.batch_periods, + precision=settings.precision, ) return Periodogram.from_spectrum( method="PDM", diff --git a/src/cuperiod/methods/string_length.py b/src/cuperiod/methods/string_length.py index 374fa96..a7749fe 100644 --- a/src/cuperiod/methods/string_length.py +++ b/src/cuperiod/methods/string_length.py @@ -17,6 +17,13 @@ import numpy as np +from cuperiod.core._arrayapi import ( + array_namespace, + resolve_precision, + resolve_torch_device, + to_device_array, + to_host, +) from cuperiod.core._typing import FloatArray from cuperiod.core.backend import ensure_cuda_dll_path from cuperiod.core.config import StringLengthSettings @@ -39,21 +46,27 @@ def _length_batch( xp: ModuleType, tau: Any, m_scaled: Any, periods: Any, *, batch: int ) -> Any: - """Total folded string length for each trial period, vectorized over periods.""" - n_points = int(tau.shape[0]) + """Total folded string length for each trial period, vectorized over periods. + + Array-API generic (numpy/cupy/torch). Uses only fancy indexing and slicing for the + per-row sort/gather and the consecutive differences, avoiding ``take_along_axis`` / + ``diff`` (not in every namespace); float dtype follows ``periods``. + """ + idtype = xp.int64 n_periods = int(periods.shape[0]) - length = xp.empty(n_periods, dtype=np.float64) + length = xp.empty(n_periods, dtype=periods.dtype) for start in range(0, n_periods, batch): stop = min(start + batch, n_periods) pb = periods[start:stop] n_p = int(pb.shape[0]) - phase = xp.mod(tau[None, :] / pb[:, None], 1.0) # (P, N) + phase = xp.remainder(tau[None, :] / pb[:, None], 1.0) # (P, N) order = xp.argsort(phase, axis=1) - ph = xp.take_along_axis(phase, order, axis=1) - mm = xp.take_along_axis(xp.broadcast_to(m_scaled, (n_p, n_points)), order, 1) - dphi = xp.diff(ph, axis=1) - dmag = xp.diff(mm, axis=1) - total = xp.sqrt(dphi * dphi + dmag * dmag).sum(axis=1) + rows = xp.arange(n_p, dtype=idtype)[:, None] + ph = phase[rows, order] # phase sorted per row + mm = m_scaled[order] # magnitudes gathered in the same order + dphi = ph[:, 1:] - ph[:, :-1] + dmag = mm[:, 1:] - mm[:, :-1] + total = xp.sum(xp.sqrt(dphi * dphi + dmag * dmag), axis=1) wrap_phi = (ph[:, 0] + 1.0) - ph[:, -1] wrap_mag = mm[:, 0] - mm[:, -1] total = total + xp.sqrt(wrap_phi * wrap_phi + wrap_mag * wrap_mag) @@ -66,8 +79,9 @@ def string_length( y: FloatArray, periods: FloatArray, *, - backend: SLBackend = "numpy", + backend: str = "numpy", batch: int = DEFAULT_BATCH, + precision: str = "auto", ) -> FloatArray: """String length for each trial period (minimized at the true period). @@ -105,6 +119,18 @@ def string_length( batch=batch, ) return np.asarray(cp.asnumpy(length), dtype=np.float64) + if backend == "torch" or backend.startswith("torch:"): + import torch + + device = backend.split(":", 1)[1] if ":" in backend else "cpu" + fdt = (torch.float32 + if resolve_precision(precision, device) == "float32" else torch.float64) + tau_d = to_device_array(tau, device=device, dtype=fdt) + m_d = to_device_array(m_scaled, device=device, dtype=fdt) + per_d = to_device_array(periods_host, device=device, dtype=fdt) + return to_host( + _length_batch(array_namespace(per_d), tau_d, m_d, per_d, batch=batch) + ) if backend != "numpy": raise ValueError(f"unknown backend {backend!r}") return np.asarray( @@ -121,7 +147,8 @@ class StringLengthMethod(PeriodogramMethod): settings_cls: ClassVar[type] = StringLengthSettings cpu_backend: ClassVar[str] = "numpy" gpu_backend: ClassVar[str | None] = "cupy" - all_backends: ClassVar[tuple[str, ...]] = ("numpy", "cupy") + portable_gpu_backend: ClassVar[str | None] = "torch" + all_backends: ClassVar[tuple[str, ...]] = ("numpy", "cupy", "torch") def default_grid(self, lc: LightCurve, settings: StringLengthSettings) -> GridSpec: # type: ignore[override] finite = lc.finite() @@ -156,9 +183,12 @@ def power( # type: ignore[override] if finite.baseline <= 0.0: raise InsufficientDataError("string-length: no usable time baseline") periods = grid.period + if backend == "torch" or backend.startswith("torch:"): + backend = f"torch:{resolve_torch_device(backend, settings.device)}" length = string_length( finite.time, finite.value, periods, - backend=backend, batch=settings.batch_periods, # type: ignore[arg-type] + backend=backend, batch=settings.batch_periods, + precision=settings.precision, ) return Periodogram.from_spectrum( method="STRINGLENGTH", diff --git a/src/cuperiod/methods/tls.py b/src/cuperiod/methods/tls.py index 674e637..74a7bec 100644 --- a/src/cuperiod/methods/tls.py +++ b/src/cuperiod/methods/tls.py @@ -24,6 +24,14 @@ import numpy as np +from cuperiod.core._arrayapi import ( + array_namespace, + resolve_precision, + resolve_torch_device, + scatter_add, + to_device_array, + to_host, +) from cuperiod.core._typing import FloatArray from cuperiod.core.backend import ensure_cuda_dll_path from cuperiod.core.columns import Domain @@ -112,48 +120,53 @@ def _matched_filter( ) -> dict[str, Any]: """Per-period best matched-filter signal residue and transit parameters. - Device-agnostic: ``xp`` is numpy (CPU) or cupy (GPU). The phase correlation is a - short width-loop of shifted slices (the template width is small), which runs on - either backend without ``sliding_window_view`` (absent in cupy). Templates stay on - the host; their per-bin scalars broadcast against device arrays. + Array-API generic: ``xp`` is an array_api_compat namespace (numpy on CPU, torch on + any device). The phase correlation is a short width-loop of shifted slices (the + template width is small), which needs no ``sliding_window_view``. Templates stay on + the host; their per-bin scalars broadcast against the device arrays. The cupy + ``RawKernel`` fast path (NVIDIA) is separate. Float dtype follows ``periods``. """ + fdtype = periods.dtype + idtype = xp.int64 n_periods = int(periods.shape[0]) out = { - "sr": xp.zeros(n_periods), - "depth": xp.zeros(n_periods), - "duration": xp.zeros(n_periods), - "t0": xp.zeros(n_periods), + "sr": xp.zeros(n_periods, dtype=fdtype), + "depth": xp.zeros(n_periods, dtype=fdtype), + "duration": xp.zeros(n_periods, dtype=fdtype), + "t0": xp.zeros(n_periods, dtype=fdtype), } n_points = int(tau.shape[0]) for start in range(0, n_periods, period_batch): stop = min(start + period_batch, n_periods) pb = periods[start:stop] n_p = int(pb.shape[0]) - rows_p = xp.arange(n_p) - phase = xp.mod(tau[None, :] / pb[:, None], 1.0) - bin_idx = xp.minimum((phase * n_bins).astype(xp.int64), n_bins - 1) - flat = (rows_p[:, None] * n_bins + bin_idx).ravel() - a_flat = xp.zeros(n_p * n_bins) # sum w*y' per bin - b_flat = xp.zeros(n_p * n_bins) # sum w per bin - xp.add.at(a_flat, flat, xp.broadcast_to(yw, (n_p, n_points)).ravel()) - xp.add.at(b_flat, flat, xp.broadcast_to(w, (n_p, n_points)).ravel()) - a = a_flat.reshape(n_p, n_bins) - b = b_flat.reshape(n_p, n_bins) - - best_sr = xp.zeros(n_p) - best_depth = xp.zeros(n_p) - best_start = xp.zeros(n_p, dtype=xp.int64) - best_width = xp.zeros(n_p, dtype=xp.int64) + rows_p = xp.arange(n_p, dtype=idtype) + phase = xp.remainder(tau[None, :] / pb[:, None], 1.0) + bin_idx = xp.clip(xp.astype(phase * n_bins, idtype), 0, n_bins - 1) + flat = xp.reshape(rows_p[:, None] * n_bins + bin_idx, (-1,)) + a_flat = xp.zeros(n_p * n_bins, dtype=fdtype) # sum w*y' per bin + b_flat = xp.zeros(n_p * n_bins, dtype=fdtype) # sum w per bin + yw_b = xp.reshape(xp.broadcast_to(yw, (n_p, n_points)), (-1,)) + w_b = xp.reshape(xp.broadcast_to(w, (n_p, n_points)), (-1,)) + scatter_add(a_flat, flat, yw_b) + scatter_add(b_flat, flat, w_b) + a = xp.reshape(a_flat, (n_p, n_bins)) + b = xp.reshape(b_flat, (n_p, n_bins)) + + best_sr = xp.zeros(n_p, dtype=fdtype) + best_depth = xp.zeros(n_p, dtype=fdtype) + best_start = xp.zeros(n_p, dtype=idtype) + best_width = xp.zeros(n_p, dtype=idtype) for width in dur_bins: g = templates[width] - a_ext = xp.concatenate([a, a[:, : width - 1]], axis=1) - b_ext = xp.concatenate([b, b[:, : width - 1]], axis=1) - num = xp.zeros((n_p, n_bins)) - den = xp.zeros((n_p, n_bins)) + a_ext = xp.concat([a, a[:, : width - 1]], axis=1) + b_ext = xp.concat([b, b[:, : width - 1]], axis=1) + num = xp.zeros((n_p, n_bins), dtype=fdtype) + den = xp.zeros((n_p, n_bins), dtype=fdtype) for k in range(width): # correlate the folded data with the template gk = float(g[k]) - num += gk * a_ext[:, k : k + n_bins] - den += (gk * gk) * b_ext[:, k : k + n_bins] + num = num + gk * a_ext[:, k : k + n_bins] + den = den + (gk * gk) * b_ext[:, k : k + n_bins] safe_den = xp.where(den > _W_EPS, den, 1.0) # a dip means the in-transit weighted residual is negative -> num < 0 sr = xp.where((den > _W_EPS) & (num < 0.0), num * num / safe_den, 0.0) @@ -165,14 +178,16 @@ def _matched_filter( depth = -num[rows_p, s_best] / xp.where(den_best > _W_EPS, den_best, 1.0) best_depth = xp.where(improve, depth, best_depth) best_start = xp.where(improve, s_best, best_start) - best_width = xp.where(improve, xp.int64(width), best_width) + best_width = xp.where(improve, width, best_width) sl = slice(start, stop) out["sr"][sl] = best_sr out["depth"][sl] = best_depth - out["duration"][sl] = best_width.astype(xp.float64) / n_bins * pb - centre = (best_start.astype(xp.float64) + best_width / 2.0) / n_bins - out["t0"][sl] = xp.mod(centre, 1.0) * pb + out["duration"][sl] = xp.astype(best_width, fdtype) / n_bins * pb + centre = ( + xp.astype(best_start, fdtype) + xp.astype(best_width, fdtype) / 2.0 + ) / n_bins + out["t0"][sl] = xp.remainder(centre, 1.0) * pb return out @@ -337,7 +352,7 @@ def tls_power( periods: FloatArray, *, settings: TLSSettings, - backend: TLSBackend = "numpy", + backend: str = "numpy", ) -> dict[str, FloatArray]: """TLS matched-filter search over ``periods`` (flux input; a transit is a dip). @@ -391,9 +406,27 @@ def tls_power( tau, yw, w, periods_host, n_bins=n_bins, dur_bins=dur_bins, templates=templates, ) + elif backend == "torch" or backend.startswith("torch:"): + import torch + + device = backend.split(":", 1)[1] if ":" in backend else "cpu" + fdt = (torch.float32 + if resolve_precision(settings.precision, device) == "float32" + else torch.float64) + per_d = to_device_array(periods_host, device=device, dtype=fdt) + dev_res = _matched_filter( + array_namespace(per_d), + to_device_array(tau, device=device, dtype=fdt), + to_device_array(yw, device=device, dtype=fdt), + to_device_array(w, device=device, dtype=fdt), + per_d, + n_bins=n_bins, dur_bins=dur_bins, templates=templates, + period_batch=settings.period_batch, + ) + res = {k: to_host(v) for k, v in dev_res.items()} elif backend == "numpy": res = _matched_filter( - np, tau, yw, w, periods_host, + array_namespace(periods_host), tau, yw, w, periods_host, n_bins=n_bins, dur_bins=dur_bins, templates=templates, period_batch=settings.period_batch, ) @@ -416,7 +449,8 @@ class TLSMethod(PeriodogramMethod): settings_cls: ClassVar[type] = TLSSettings cpu_backend: ClassVar[str] = "numpy" gpu_backend: ClassVar[str | None] = "cupy" - all_backends: ClassVar[tuple[str, ...]] = ("numpy", "cupy") + portable_gpu_backend: ClassVar[str | None] = "torch" + all_backends: ClassVar[tuple[str, ...]] = ("numpy", "cupy", "torch") def default_grid(self, lc: LightCurve, settings: TLSSettings) -> GridSpec: # type: ignore[override] finite = lc.finite() @@ -444,9 +478,11 @@ def power( # type: ignore[override] if finite.baseline <= 0.0: raise InsufficientDataError("TLS: no usable time baseline") periods = grid.period + if backend == "torch" or backend.startswith("torch:"): + backend = f"torch:{resolve_torch_device(backend, settings.device)}" res = tls_power( finite.time, finite.value, finite.error, periods, - settings=settings, backend=backend, # type: ignore[arg-type] + settings=settings, backend=backend, ) extras = { "sr": res["sr"], diff --git a/tests/test_torch_bonus.py b/tests/test_torch_bonus.py new file mode 100644 index 0000000..b1eecb6 --- /dev/null +++ b/tests/test_torch_bonus.py @@ -0,0 +1,66 @@ +"""Torch-CPU parity for the bonus methods (CE, String-Length, MHAOV, PDM, TLS). + +Each method's portable torch path must match its numpy path (both float64 on CPU) and +pick the same best period. MHAOV goes through ``linalg.solve`` so its agreement is +relative (peak-normalized), not bit-exact; the rest are near-identical. GPU devices +are validated separately in PR2; these run on the CPU torch device. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +import cuperiod as cup +from conftest import requires_torch +from synth import synthetic_eclipser, synthetic_sine + +FREQ_METHODS = ["CE", "STRINGLENGTH", "MHAOV", "PDM"] + + +@requires_torch +@pytest.mark.parametrize("method", FREQ_METHODS) +def test_bonus_torch_matches_numpy(method: str) -> None: + t, mag, err = synthetic_sine(period=0.6234) + lc = cup.LightCurve.from_arrays(t, mag, err) + m = cup.get_method(method) + settings = m.settings_cls() + grid = m.default_grid(lc, settings) + ref = m.power(grid, lc, settings, "numpy") + tor = m.power(grid, lc, settings, "torch:cpu") + assert tor.backend == "torch:cpu" + finite = np.isfinite(ref.power) & np.isfinite(tor.power) + scale = max(float(np.max(np.abs(ref.power[finite]))), 1e-30) + max_abs = float(np.max(np.abs(ref.power[finite] - tor.power[finite]))) + assert max_abs / scale < 1e-6 + # Same best period off the identical grid (whichever sense the method is). + assert ref.best_period() == pytest.approx(tor.best_period(), rel=1e-9) + + +@requires_torch +def test_tls_torch_matches_numpy() -> None: + t, flux, err = synthetic_eclipser(period=2.5, depth=0.05) + lc = cup.LightCurve.from_arrays(t, flux, err, domain=cup.Domain.FLUX) + m = cup.get_method("TLS") + settings = m.settings_cls(min_period_days=1.5, max_period_days=4.0) + grid = m.default_grid(lc, settings) + ref = m.power(grid, lc, settings, "numpy") + tor = m.power(grid, lc, settings, "torch:cpu") + assert tor.backend == "torch:cpu" + finite = np.isfinite(ref.power) & np.isfinite(tor.power) + scale = max(float(np.max(np.abs(ref.power[finite]))), 1e-30) + assert float(np.max(np.abs(ref.power[finite] - tor.power[finite]))) / scale < 1e-6 + assert int(np.argmax(ref.power)) == int(np.argmax(tor.power)) + + +@requires_torch +@pytest.mark.parametrize("method", FREQ_METHODS) +def test_bonus_torch_recovers_period(method: str) -> None: + # The portable torch path recovers the period (or a small-integer harmonic — + # String-Length and MHAOV legitimately lock onto 3P here, identically to numpy). + t, mag, err = synthetic_sine(period=0.6234) + pg = cup.periodogram((t, mag, err), method, backend="torch") + assert pg.backend == "torch:cpu" + ratio = pg.best_period() / 0.6234 + harmonics = (1.0, 0.5, 2.0, 1 / 3, 3.0, 2 / 3, 1.5) + assert min(abs(ratio / r - 1.0) for r in harmonics) < 5e-3 From 36a28a9e030a9c41a4be13dbc6f9a2246472e84f Mon Sep 17 00:00:00 2001 From: Tharindu Jayasinghe Date: Tue, 30 Jun 2026 12:40:31 -0700 Subject: [PATCH 4/6] feat(cli,docs): `cuperiod doctor` + PyTorch install/backend docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR4 of the multi-vendor GPU effort — the install UX for the torch backend. - cli: new `cuperiod doctor` command — reports installed backends, the NVIDIA CUDA fast paths, the portable torch backend and each device it sees (CUDA/ROCm/MPS/XPU/CPU) with the precision each uses, and what `backend="auto"` resolves to per method. It sets KMP_DUPLICATE_LIB_OK for its own read-only device probe only (it does no numerics), so it can't OMP-abort on Windows; the library still never sets it for compute paths. - docs/installation: `[torch]` extra, per-accelerator PyTorch wheel guidance (the plain wheel is CPU-only), the Apple-MPS float32 note, and a Windows OpenMP-clash warning with the KMP_DUPLICATE_LIB_OK workaround. - docs/guide/backends: torch in the selector table; a "portable PyTorch backend" section covering devices, the device/precision settings, and the fast-on-NVIDIA / portable-everywhere split. - docs/guide/cli: document `doctor`. Verified: `cuperiod doctor` runs clean on a CPU-only Windows box without any external env (torch:cpu -> float64; auto -> finufft/numba/numpy). Docs build with `-W`. 161 passed, 8 skipped; ruff + mypy clean. Co-Authored-By: Claude Opus 4.8 --- docs/guide/backends.md | 42 +++++++++++++++++++++-- docs/guide/cli.md | 15 ++++++++- docs/installation.md | 51 ++++++++++++++++++++++++++++ src/cuperiod/cli/app.py | 75 ++++++++++++++++++++++++++++++++++++++++- tests/test_cli.py | 7 ++++ 5 files changed, 185 insertions(+), 5 deletions(-) diff --git a/docs/guide/backends.md b/docs/guide/backends.md index 5d0f465..baa875f 100644 --- a/docs/guide/backends.md +++ b/docs/guide/backends.md @@ -10,13 +10,14 @@ floating-point round-off regardless of which one runs. pg = cup.periodogram(lc, "GLS", backend="auto") # default ``` -The four selectors: +The selectors: | `backend=` | Meaning | | --- | --- | -| `"auto"` *(default)* | Use the GPU when a CUDA device and the `[gpu]` extra are present; otherwise the best CPU backend. The same code runs on any machine. | +| `"auto"` *(default)* | The NVIDIA CUDA fast path when a CUDA device and the `[gpu]` extra are present; then the portable **torch** backend on another GPU (AMD/Intel/Mac) when the `[torch]` extra sees one; otherwise the best CPU backend. The same code runs on any machine. | | `"cpu"` | Force the best CPU backend for this method. | -| `"gpu"` | Force the GPU backend. Raises {exc}`~cuperiod.BackendUnavailableError` if no GPU is usable. | +| `"gpu"` | Force a GPU backend — the CUDA fast path, or torch on a non-NVIDIA GPU. Raises {exc}`~cuperiod.BackendUnavailableError` if no GPU is usable. | +| `"torch"` / `"torch:"` | Force the portable PyTorch backend; `` is `cpu`, `cuda`, `mps`, or `xpu` (bare `"torch"` picks the best present). Needs the `[torch]` extra. | | a concrete name | Force a specific implementation, e.g. `"finufft"`, `"astropy"`, `"numpy"`, `"cupy"`, `"cufinufft"`, `"numba"`. | ## What each method can run @@ -64,6 +65,41 @@ array-module-generic source with the CUDA kernel so the two validate to floating but it is slow (it trades memory traffic for the parallelism that makes the GPU fast). For CPU BLS use `numba` (the default with `[fast]`) or `astropy`, **not** `numpy`. +In addition, **every method has a portable `torch` backend** (the `[torch]` extra) that +runs the same array-API code on any torch device — see below. + +## The portable PyTorch backend + +`backend="torch"` runs a method through the [array API](https://data-apis.org/array-api/) +on whichever torch device you have — CUDA, AMD **ROCm**, Apple **MPS**, Intel **XPU**, or +**CPU** — so the accelerated code is no longer NVIDIA-only and works even with no GPU at +all. The NVIDIA fast paths (cufinufft for GLS, the cupy kernels for BLS/PDM/CE/TLS) are +untouched and remain what `"auto"`/`"gpu"` pick on CUDA; torch is the cross-vendor path +for everything else. + +```python +pg = cup.periodogram(lc, "GLS", backend="torch") # best torch device +pg = cup.periodogram(lc, "BLS", backend="torch:xpu") # pin the Intel GPU +pg = cup.periodogram(lc, "PDM", backend="torch:cpu", settings=cup.PDMSettings(device="cpu")) +``` + +Two orthogonal settings tune it (both environment-overridable, e.g. `CUPERIOD_GLS_DEVICE`): + +- **`device`** — `"auto"` (default; the best present), `"cpu"`, `"cuda"`, `"mps"`, `"xpu"`. + A `"torch:"` backend string overrides it. +- **`precision`** — `"auto"` (default) is float64 everywhere it is supported and float32 + only where the device forces it (Apple MPS; some Intel GPUs). `"float64"` and `"float32"` + force it; `precision="float64"` on MPS raises rather than silently downgrading. Results + are always returned as float64 numpy arrays regardless of the device precision. + +:::{note} +The portable path is correct everywhere (it matches the CPU reference to round-off — float +methods bit-for-bit, MHAOV to ~1e-10 from its linear solve) but it is **not** the speed +champion on CPU, where finufft (GLS) and the numba box search (BLS) are faster. Its value +is reaching GPUs the CUDA fast paths can't — AMD, Intel, and Apple. On those devices its +massively-parallel evaluation is the win. Run `cuperiod doctor` to see what you have. +::: + Inspect the live picture for your install: ```python diff --git a/docs/guide/cli.md b/docs/guide/cli.md index 50792dc..1ca9ab6 100644 --- a/docs/guide/cli.md +++ b/docs/guide/cli.md @@ -8,7 +8,8 @@ Python API, so the CLI and library share one code path and give identical result cuperiod run one light curve, one or more methods → prints the best periods cuperiod batch many light curves with CPU or GPU workers → Parquet/CSV cuperiod methods list registered methods and their backends -cuperiod gpu-info show the GPU and suggested worker counts +cuperiod gpu-info show the CUDA GPU and suggested worker counts +cuperiod doctor diagnose backends, torch devices, and the precision each uses cuperiod grid-info show a method's trial grid for a light curve (no compute) ``` @@ -102,6 +103,18 @@ cuperiod gpu-info Shows the CUDA device (name, free/total memory, MPS status) and the suggested batch worker count per GPU-capable method. If no GPU is present, it says so and exits cleanly. +## `doctor` — full environment diagnosis + +```bash +cuperiod doctor +``` + +A one-stop "will the accelerated paths run here, and on what?" report: which backends are +installed, the NVIDIA CUDA fast paths, the portable **torch** backend and each device it +sees (CUDA/ROCm/MPS/XPU/CPU) with the precision it would use, and what `backend="auto"` +resolves to for every method. Reach for it first when a GPU isn't being picked up or you're +unsure which build of PyTorch you have. + ## `grid-info` — inspect a grid ```bash diff --git a/docs/installation.md b/docs/installation.md index 96bc6d7..467de5f 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -18,6 +18,7 @@ line, and batch processing over a process pool. | Extra | Command | Adds | | --- | --- | --- | | **gpu** | `pip install "cuperiod[gpu]"` | CUDA-12 GPU backends (`cupy-cuda12x`, `cufinufft`, the NVIDIA runtime wheels) | +| **torch** | `pip install "cuperiod[torch]"` | the portable **PyTorch** backend — runs every method on AMD (ROCm), Intel (XPU), Apple (MPS), and a CPU path | | **fast** | `pip install "cuperiod[fast]"` | a multicore `numba` box search — BLS's CPU default, ~20× faster than astropy's compiled `BoxLeastSquares` | | **pandas** | `pip install "cuperiod[pandas]"` | pandas `DataFrame` ingestion | @@ -56,6 +57,49 @@ first used, so no manual `PATH` editing is needed. cupy may print a benign system toolkit — this is harmless. ::: +## Portable GPU backend (PyTorch) + +The `[torch]` extra adds a **PyTorch backend** that runs every method beyond NVIDIA — on +AMD (ROCm), Intel (XPU), and Apple-Silicon (MPS) GPUs, and on a CPU path everywhere (handy +even with no GPU at all): + +```bash +pip install "cuperiod[torch]" +``` + +The plain wheel above is **CPU-only**. To use a GPU, install the PyTorch build matching +your accelerator from the [official index](https://pytorch.org/get-started/locally/) — we +deliberately don't pin a hardware-specific wheel: + +| Hardware | PyTorch build | +| --- | --- | +| NVIDIA (CUDA) | the CUDA wheel — or just use the `[gpu]` extra's faster cufinufft/cupy paths | +| AMD (ROCm) | the ROCm wheel (Linux only) | +| Intel (XPU) | the XPU wheel (`torch.xpu`) | +| Apple Silicon | the standard macOS wheel (MPS is built in) | +| CPU only | the default wheel | + +Select it with `backend="torch"` (or `"torch:cpu"`, `"torch:cuda"`, `"torch:mps"`, +`"torch:xpu"`); `backend="auto"` reaches a torch GPU on non-NVIDIA machines after the +cufinufft/cupy fast paths. See {doc}`guide/backends`. + +:::{note} +**Apple MPS** cannot compute in float64 (a Metal limitation), so the Mac-GPU path uses +float32. `precision="auto"` (the default) keeps float64 everywhere it is supported and +drops to float32 only where the device forces it (MPS, and some Intel GPUs); an explicit +`precision="float64"` on MPS raises rather than silently downgrading. +::: + +:::{warning} +**Windows OpenMP clash.** PyTorch and NumPy/SciPy (MKL) each ship an OpenMP runtime, and +importing torch after numpy can abort with *"OMP: Error #15 … libiomp5md.dll already +initialized."* cuPeriod does **not** set a workaround for you — it can silently affect +numerical results. If you hit this running a torch workload on Windows, set +`KMP_DUPLICATE_LIB_OK=TRUE` in your environment, or install torch and numpy builds that +share one OpenMP runtime. (`cuperiod doctor` sets it only for its own read-only device +probe.) +::: + ## Verifying the install List the registered methods and the backends available in your environment: @@ -70,6 +114,13 @@ Check whether a GPU is visible and how many batch workers it would suggest: cuperiod gpu-info ``` +For the full picture — every installed backend, the available torch devices and the +precision each will use, and what `backend="auto"` resolves to for every method — run: + +```bash +cuperiod doctor +``` + If no CUDA device is present, `gpu-info` says so and exits cleanly — the CPU paths still work. From Python: diff --git a/src/cuperiod/cli/app.py b/src/cuperiod/cli/app.py index e4cbc96..ddf7730 100644 --- a/src/cuperiod/cli/app.py +++ b/src/cuperiod/cli/app.py @@ -6,7 +6,8 @@ * ``run`` — one light curve, one or more methods; prints the N best periods. * ``batch`` — many light curves with CPU or GPU workers, written to Parquet/CSV. * ``methods`` — list registered methods and their backends. -* ``gpu-info`` — show the GPU and suggested worker counts. +* ``gpu-info`` — show the CUDA GPU and suggested worker counts. +* ``doctor`` — diagnose available backends, torch devices, and the precision each uses. * ``grid-info`` — show a method's trial grid for a light curve without computing it. """ @@ -193,6 +194,78 @@ def gpu_info_cmd() -> None: typer.echo(f" {m.name}: {suggest_gpu_workers(m.name)}") +@app.command() +def doctor() -> None: + """Diagnose available backends, devices, and the precision each will use. + + A one-stop "will the accelerated paths run here, and on what?" check: the installed + backends, the NVIDIA CUDA fast paths, the portable torch backend and its devices + (CUDA/ROCm/MPS/XPU/CPU), and what ``backend="auto"`` resolves to per method. + """ + import os + import platform + import sys + from importlib.metadata import PackageNotFoundError + from importlib.metadata import version as _pkg_version + + # This probe only enumerates devices (no numerics), so on Windows it allows the + # torch + numpy/MKL OpenMP duplicate so importing torch can't abort it. The library + # never sets this for compute paths (see the install docs) — a torch workload on a + # conflicting Windows env should set KMP_DUPLICATE_LIB_OK itself. + os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE") + + from cuperiod.core._arrayapi import resolve_precision + from cuperiod.core.backend import ( + available_backends, + best_torch_device, + torch_available, + torch_devices, + ) + + try: + ver = _pkg_version("cuperiod") + except PackageNotFoundError: # pragma: no cover - editable/source runs + ver = "?" + typer.echo( + f"cuPeriod {ver} | Python {sys.version.split()[0]} | " + f"{platform.system()} {platform.machine()}" + ) + + avail = available_backends() + typer.echo("\nbackends installed:") + for name in ("numpy", "finufft", "numba", "astropy", "cupy", "cufinufft", "torch"): + typer.echo(f" {'OK' if name in avail else '--':>2} {name}") + + typer.echo("\nNVIDIA CUDA fast paths (cufinufft, cupy kernels):") + info = _gpu_info() + typer.echo( + f" {info}" if info is not None + else " no CUDA device (needs the [gpu] extra and an NVIDIA GPU)" + ) + + typer.echo("\nportable torch backend (AMD/Intel/Mac/CPU):") + if not torch_available(): + typer.echo(" torch not installed — `pip install 'cuperiod[torch]'`") + else: + devices = torch_devices() + best = best_torch_device() + for d in ("cuda", "xpu", "mps", "cpu"): + if d in devices: + tag = " <- best (used by backend='auto')" if d == best else "" + prec = resolve_precision("auto", d) + typer.echo(f" OK torch:{d:4s} precision auto -> {prec}{tag}") + else: + typer.echo(f" -- torch:{d}") + + typer.echo("\nbackend='auto' resolves to:") + for m in list_methods(): + try: + resolved = get_method(m.name).resolve_backend("auto") + except Exception as exc: # pragma: no cover - defensive + resolved = f"error: {exc}" + typer.echo(f" {m.name:14s} {resolved}") + + @app.command(name="grid-info") def grid_info_cmd( path: Path = typer.Argument(..., help="Light-curve file."), diff --git a/tests/test_cli.py b/tests/test_cli.py index 6957a32..bcf6382 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -82,6 +82,13 @@ def test_gpu_info_command() -> None: assert result.exit_code == 0 +def test_doctor_command() -> None: + result = runner.invoke(app, ["doctor"]) + assert result.exit_code == 0 + assert "backends installed" in result.stdout + assert "backend='auto' resolves to" in result.stdout + + def test_grid_info_command(tmp_path: Path) -> None: csv = _write_csv(tmp_path / "star.csv") result = runner.invoke(app, ["grid-info", str(csv), "--method", "GLS"]) From c2626bc0dd222f0f644b564fc79f78b3e812468d Mon Sep 17 00:00:00 2001 From: Tharindu Jayasinghe Date: Tue, 30 Jun 2026 14:41:41 -0700 Subject: [PATCH 5/6] fix(torch): resolve PR5 review findings (float32 correctness) The PR5 local adversarial review surfaced three float32-only bugs in the portable torch backend, all masked because every torch test ran float64-on-CPU: - BLS cast absolute BJD time to float32 *before* subtracting t_min, destroying sub-0.25-day timing on real light curves (silent wrong periodogram; this is the default path on Apple MPS). Subtract the origin in float64 before the device cast and restore the absolute transit_time on the host (mirrors gls._prep). - MHAOV multiband dropped `precision`, silently ignoring it (and, on MPS, downcasting to float32 instead of raising). Forward precision to aov_multiband_power. - MHAOV's fixed 1e-10 diagonal ridge underflowed the ~N Gram diagonal in float32, so linalg.solve raised on a singular matrix at degenerate frequencies. Scale the ridge by eps(dtype)*n_points. Add float32 regression tests (three proven to fail pre-fix) and the requires_torch_gpu marker. Document the remaining hardware-gated limitations (XPU fp64 probe, torch-GPU VRAM auto-sizing, GPU argmax tie-breaks of extras) in the changelog. NVIDIA fast paths verified unchanged; 165 passed, ruff+mypy clean. Co-Authored-By: Claude Opus 4.8 --- CHANGELOG.md | 15 +++++++++++ src/cuperiod/core/_arrayapi.py | 4 +++ src/cuperiod/methods/_bls_core.py | 16 +++++++++-- src/cuperiod/methods/mhaov.py | 14 +++++++--- src/cuperiod/multiband/mhaov_mb.py | 1 + tests/conftest.py | 22 ++++++++++++++- tests/test_bls.py | 26 ++++++++++++++++++ tests/test_gls.py | 13 +++++++++ tests/test_torch_bonus.py | 43 ++++++++++++++++++++++++++++++ 9 files changed, 148 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index acd6b9f..64da595 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,21 @@ All notable changes to cuPeriod are documented here. The format is based on - `array-api-compat` is now a dependency; install the portable accelerator with the `[torch]` extra (`pip install 'cuperiod[torch]'`). +### Known limitations + +- **Non-NVIDIA GPU numerics are written-to-spec and CPU-validated, not yet hardware- + verified.** The torch CUDA/ROCm/MPS/XPU paths share the array-API body that is parity- + tested on the CPU torch device; on-device parity self-skips (`requires_torch_gpu`) until + such hardware is available. +- **No fp64 capability probe on Intel XPU.** `precision="auto"` resolves to float64 on an + XPU; a device without native float64 will error at compute time rather than falling back + to float32. Pass `precision="float32"` explicitly on such a device. +- **The torch GPU path does not auto-shrink to small VRAM.** A large period×bin grid on a + small consumer GPU can raise an out-of-memory error; reduce `batch_periods`. +- **Tie-broken best-fit *extras* may differ across devices.** Where an `argmax` lands on an + exact tie (e.g. BLS `transit_time`, TLS `t0`/`duration` at non-transit periods), the + chosen index is device-dependent; the periodogram power and best period are unaffected. + ## [1.0.0] — 2026-06-29 First public release. diff --git a/src/cuperiod/core/_arrayapi.py b/src/cuperiod/core/_arrayapi.py index 62c0555..74fc32f 100644 --- a/src/cuperiod/core/_arrayapi.py +++ b/src/cuperiod/core/_arrayapi.py @@ -79,6 +79,10 @@ def scatter_add(target: Any, index: Any, values: Any) -> None: ``target`` is 1-D and ``index``/``values`` are 1-D and aligned. ``numpy.add.at`` is not part of the array-API standard and PyTorch has no equivalent *function*, so this dispatches per backend: torch uses ``Tensor.index_add_``; numpy/cupy use ``add.at``. + + ``values.dtype`` must equal ``target.dtype``: torch ``index_add_`` rejects a + mismatch (numpy/cupy would silently cast), so callers build both at the same + working float dtype. ``index`` may be any integer dtype (coerced to int64). """ if is_torch_array(target): target.index_add_(0, index.long(), values) diff --git a/src/cuperiod/methods/_bls_core.py b/src/cuperiod/methods/_bls_core.py index d3ca1e1..f6e75da 100644 --- a/src/cuperiod/methods/_bls_core.py +++ b/src/cuperiod/methods/_bls_core.py @@ -677,6 +677,12 @@ def bls_power( y_host = np.ascontiguousarray(y, dtype=np.float64) ivar_host = 1.0 / (np.ascontiguousarray(dy, dtype=np.float64) ** 2) + # Time origin subtracted before any (possibly float32) device cast; restored + # into the absolute ``transit_time`` on the host below. Stays 0.0 for the float64 + # host paths (numpy/numba/cupy), which build the absolute time directly. See the + # torch branch. + t_ref = 0.0 + if backend == "cupy": from cuperiod.core.backend import ensure_cuda_dll_path @@ -701,7 +707,13 @@ def bls_power( if resolve_precision(precision, device) == "float32" else torch.float64 ) - t_d = to_device_array(t_host, device=device, dtype=tdtype) + # Subtract the time origin in float64 *before* the device cast. Absolute BJDs + # (~2.458e6) lose all sub-0.25-day timing when cast to float32 — the default + # precision on Apple MPS — so a raw cast would silently corrupt the phase fold. + # The small t-min-relative tau survives float32; ``transit_time`` is shifted + # back to absolute on the host (float64) at the return. Mirrors ``gls._prep``. + t_ref = float(np.min(t_host)) + t_d = to_device_array(t_host - t_ref, device=device, dtype=tdtype) y_d = to_device_array(y_host, device=device, dtype=tdtype) ivar_d = to_device_array(ivar_host, device=device, dtype=tdtype) periods_d = to_device_array(periods_host, device=device, dtype=tdtype) @@ -724,7 +736,7 @@ def bls_power( depth_err=to_host(out["depth_err"]), depth_snr=to_host(out["depth_snr"]), duration=to_host(out["duration"]), - transit_time=to_host(out["transit_time"]), + transit_time=to_host(out["transit_time"]) + t_ref, log_likelihood=to_host(out["log_likelihood"]), ) diff --git a/src/cuperiod/methods/mhaov.py b/src/cuperiod/methods/mhaov.py index 0c7d850..8e3bceb 100644 --- a/src/cuperiod/methods/mhaov.py +++ b/src/cuperiod/methods/mhaov.py @@ -50,8 +50,15 @@ #: Trial frequencies per vectorized batch (bounds the (F, N, 2H+1) design tensor). DEFAULT_BATCH: Final = 512 -#: Diagonal ridge to keep the normal equations solvable at degenerate frequencies. -_RIDGE: Final = 1e-10 +#: Diagonal ridge that keeps the harmonic normal equations solvable at degenerate +#: frequencies (f→0, where the cosine columns collapse onto the constant column). It is +#: applied as ``_RIDGE_EPS · eps(dtype) · n_points``: scaling by the working precision's +#: machine epsilon and the Gram diagonal magnitude (≈ ``n_points``, from the all-ones +#: constant column) makes it representable in float32. A fixed absolute 1e-10 underflows +#: against the ~N-sized diagonal on a float32 device (1e-10 ≪ eps_f32·N), leaving the +#: matrix singular so ``linalg.solve`` raises. In float64 this reproduces the previous +#: ~1e-10 ridge to within rounding, so float64 results are unchanged. +_RIDGE_EPS: Final = 1.0e3 def _design(xp: ModuleType, angle: Any, n_harmonics: int) -> Any: @@ -86,7 +93,8 @@ def _model_ss_batch( d = 2 * n_harmonics + 1 n_freq = int(frequencies.shape[0]) fdtype = frequencies.dtype - eye = xp.eye(d, dtype=fdtype) * _RIDGE + ridge = _RIDGE_EPS * float(xp.finfo(fdtype).eps) * float(n_points) + eye = xp.eye(d, dtype=fdtype) * ridge out = xp.empty(n_freq, dtype=fdtype) two_pi = 2.0 * float(np.pi) diff --git a/src/cuperiod/multiband/mhaov_mb.py b/src/cuperiod/multiband/mhaov_mb.py index ccf192a..e8ce43a 100644 --- a/src/cuperiod/multiband/mhaov_mb.py +++ b/src/cuperiod/multiband/mhaov_mb.py @@ -70,6 +70,7 @@ def mhaov_multiband_power( n_harmonics=settings.n_harmonics, backend=backend, # type: ignore[arg-type] batch=settings.batch_periods, + precision=settings.precision, ) n_total = sum(t.size for t, _ in bands) return Periodogram.from_spectrum( diff --git a/tests/conftest.py b/tests/conftest.py index e00105c..8eecb96 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -29,8 +29,28 @@ def _importable(name: str) -> bool: ) #: Skip a test unless PyTorch (the ``[torch]`` extra) is importable. The torch-CPU path -#: always runs once torch is installed; GPU-device tests add their own skips in PR2. +#: always runs once torch is installed; GPU-device parity adds the marker below. requires_torch = pytest.mark.skipif( not _importable("torch"), reason="torch (the [torch] extra) is not installed", ) + + +def _torch_gpu_available() -> bool: + if not _importable("torch"): + return False + try: + from cuperiod.core.backend import torch_gpu_available + + return torch_gpu_available() + except Exception: + return False + + +#: Skip a test unless torch sees a non-CPU device (CUDA/ROCm, MPS, or XPU). The portable +#: torch GPU numeric paths are written to spec and CPU-validated; on-hardware parity +#: self-skips until such a device is available (none on the current dev machine/CI). +requires_torch_gpu = pytest.mark.skipif( + not _torch_gpu_available(), + reason="no non-CPU torch device (CUDA/ROCm/MPS/XPU) available", +) diff --git a/tests/test_bls.py b/tests/test_bls.py index d8dd4a7..b413071 100644 --- a/tests/test_bls.py +++ b/tests/test_bls.py @@ -125,3 +125,29 @@ def test_bls_torch_recovers_period_and_depth() -> None: peak = pg.best_periods(1, alias_diverse=True)[0] assert peak.period == pytest.approx(2.5, rel=3e-3) assert peak.extra["depth"] == pytest.approx(0.05, abs=0.01) + + +@requires_torch +def test_bls_torch_float32_tracks_float64_on_bjd_times() -> None: + # Regression: the absolute time origin must be subtracted in float64 *before* the + # float32 device cast. Synthetic times are real-scale BJD (~2.458e6), where float32 + # spacing is ~0.25 d — coarser than a transit bin — so a raw cast destroys the phase + # fold and float32 power diverges by ~order-1 from float64. float32 is the default + # precision on Apple MPS, so this is the out-of-the-box Mac path. After the fix the + # only difference is float32 round-off. + t, flux, err = synthetic_eclipser(period=2.5, depth=0.05) + assert float(t.min()) > 2.4e6 # the catastrophic-cancellation regime + common = dict(domain=cup.Domain.FLUX, backend="torch:cpu") + p64 = cup.periodogram( + (t, flux, err), "BLS", settings=cup.BLSSettings(precision="float64"), **common + ) + p32 = cup.periodogram( + (t, flux, err), "BLS", settings=cup.BLSSettings(precision="float32"), **common + ) + assert p32.backend == "torch:cpu" + finite = np.isfinite(p64.power) & np.isfinite(p32.power) + scale = max(float(np.max(np.abs(p64.power[finite]))), 1e-30) + rel = float(np.max(np.abs(p64.power[finite] - p32.power[finite]))) / scale + assert rel < 5e-2 # pre-fix the lost timing made this ~order-1 + best = p32.best_periods(1, alias_diverse=True)[0] + assert best.period == pytest.approx(2.5, rel=5e-3) diff --git a/tests/test_gls.py b/tests/test_gls.py index b0455bc..de4b9d7 100644 --- a/tests/test_gls.py +++ b/tests/test_gls.py @@ -104,3 +104,16 @@ def test_gls_torch_recovers_period() -> None: t, mag, err = synthetic_sine(period=0.6234) pg = cup.periodogram((t, mag, err), "GLS", backend="torch") assert pg.best_period() == pytest.approx(0.6234, rel=1e-3) + + +@requires_torch +def test_gls_torch_float32_recovers_period() -> None: + # float32 is the forced precision on Apple MPS; the direct trig-sum's power drifts + # but the peak must stay robust (the float32 path previously had no test coverage). + t, mag, err = synthetic_sine(period=0.6234) + pg = cup.periodogram( + (t, mag, err), "GLS", backend="torch:cpu", + settings=cup.GLSSettings(precision="float32"), + ) + assert pg.backend == "torch:cpu" + assert pg.best_period() == pytest.approx(0.6234, rel=1e-3) diff --git a/tests/test_torch_bonus.py b/tests/test_torch_bonus.py index b1eecb6..0d85cc5 100644 --- a/tests/test_torch_bonus.py +++ b/tests/test_torch_bonus.py @@ -13,6 +13,7 @@ import cuperiod as cup from conftest import requires_torch +from cuperiod.methods.mhaov import aov_power from synth import synthetic_eclipser, synthetic_sine FREQ_METHODS = ["CE", "STRINGLENGTH", "MHAOV", "PDM"] @@ -64,3 +65,45 @@ def test_bonus_torch_recovers_period(method: str) -> None: ratio = pg.best_period() / 0.6234 harmonics = (1.0, 0.5, 2.0, 1 / 3, 3.0, 2 / 3, 1.5) assert min(abs(ratio / r - 1.0) for r in harmonics) < 5e-3 + + +@requires_torch +def test_mhaov_multiband_torch_honours_precision() -> None: + # Regression: the multiband wrapper must forward ``precision`` to the torch compute. + # It previously dropped it (defaulting to auto→float64), so an explicit float32 was + # silently ignored — bit-identical to float64 (and on MPS this turned the mandatory + # float64 raise into a silent float32 downcast). After the fix the two precisions + # produce different bits, while both still recover the period. + tg, mg, eg = synthetic_sine(n=300, period=0.6234, seed=2) + tr, mr, er = synthetic_sine(n=300, period=0.6234, amp=0.3, seed=3) + mb = cup.MultiBandLightCurve.from_light_curves( + { + "g": cup.LightCurve.from_arrays(tg, mg, eg), + "r": cup.LightCurve.from_arrays(tr, mr + 1.0, er), + } + ) + m = cup.get_method("MHAOV") + p64 = cup.periodogram( + mb, "MHAOV", backend="torch:cpu", settings=m.settings_cls(precision="float64") + ) + p32 = cup.periodogram( + mb, "MHAOV", backend="torch:cpu", settings=m.settings_cls(precision="float32") + ) + assert p32.backend == "torch:cpu" + assert not np.array_equal(p64.power, p32.power) # float32 was actually used + assert p64.best_period() == pytest.approx(p32.best_period(), rel=1e-2) + + +@requires_torch +def test_mhaov_torch_float32_survives_degenerate_frequency() -> None: + # Regression: at f→0 the harmonic basis collapses onto the constant column; the + # diagonal ridge must survive the float32 cast. A fixed absolute 1e-10 underflowed + # against the ~N-sized Gram diagonal in float32, leaving the matrix singular so + # ``linalg.solve`` raised _LinAlgError. The ridge is now eps(dtype)·N-scaled. + t, mag, err = synthetic_sine(n=400, period=0.6234) + freqs = np.array([1e-8, 1e-6, 1.0 / 0.6234]) # first is near-degenerate + power = aov_power( + t, mag, freqs, n_harmonics=3, backend="torch:cpu", precision="float32" + ) + assert np.all(np.isfinite(power)) # no crash, no NaN/inf + assert power[-1] > power[0] # the real signal beats the degenerate frequency From 6e7ea3919c9fdb319a8b8a40cf59e2ffbc0a0a44 Mon Sep 17 00:00:00 2001 From: Tharindu Jayasinghe Date: Tue, 30 Jun 2026 14:49:59 -0700 Subject: [PATCH 6/6] fix(string-length): stable phase sort across all backends MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The portable String-Length body sorted folded phases with whichever argsort the backend namespace provided: the numpy path used raw numpy.argsort (quicksort, unstable) while the torch path used the array-API argsort (stable). On exactly-tied phases the two ordered points differently, so the string length diverged — and because quicksort's tie order is platform-dependent, the numpy-vs-torch parity test passed locally but failed in CI across OS/Python. Route the numpy and cupy paths through the array-API namespace too, so every backend uses the standard stable sort and ties keep a backend- and platform-independent order. Add a regression test that forces heavy phase ties (deterministically failing pre-fix). Surfaced by the PR5 review (the argsort tie divergence) and confirmed by CI. Co-Authored-By: Claude Opus 4.8 --- src/cuperiod/methods/string_length.py | 14 ++++++++++++-- tests/test_torch_bonus.py | 19 +++++++++++++++++++ 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/src/cuperiod/methods/string_length.py b/src/cuperiod/methods/string_length.py index a7749fe..56df71d 100644 --- a/src/cuperiod/methods/string_length.py +++ b/src/cuperiod/methods/string_length.py @@ -51,6 +51,12 @@ def _length_batch( Array-API generic (numpy/cupy/torch). Uses only fancy indexing and slicing for the per-row sort/gather and the consecutive differences, avoiding ``take_along_axis`` / ``diff`` (not in every namespace); float dtype follows ``periods``. + + The phase sort must be **stable** so equal phases keep a backend-independent + order: the string length depends on neighbour pairing, and an unstable sort breaks + ties differently across numpy/torch (even across platforms), drifting the result. + All backends therefore run through the array-API namespace, whose ``argsort`` + defaults to ``stable=True`` — never raw ``numpy``/``cupy`` (quicksort, unstable). """ idtype = xp.int64 n_periods = int(periods.shape[0]) @@ -114,8 +120,9 @@ def string_length( ensure_cuda_dll_path() import cupy as cp + per_cp = cp.asarray(periods_host) length = _length_batch( - cp, cp.asarray(tau), cp.asarray(m_scaled), cp.asarray(periods_host), + array_namespace(per_cp), cp.asarray(tau), cp.asarray(m_scaled), per_cp, batch=batch, ) return np.asarray(cp.asnumpy(length), dtype=np.float64) @@ -134,7 +141,10 @@ def string_length( if backend != "numpy": raise ValueError(f"unknown backend {backend!r}") return np.asarray( - _length_batch(np, tau, m_scaled, periods_host, batch=batch), dtype=np.float64 + _length_batch( + array_namespace(periods_host), tau, m_scaled, periods_host, batch=batch + ), + dtype=np.float64, ) diff --git a/tests/test_torch_bonus.py b/tests/test_torch_bonus.py index 0d85cc5..0fc69ec 100644 --- a/tests/test_torch_bonus.py +++ b/tests/test_torch_bonus.py @@ -67,6 +67,25 @@ def test_bonus_torch_recovers_period(method: str) -> None: assert min(abs(ratio / r - 1.0) for r in harmonics) < 5e-3 +@requires_torch +def test_string_length_torch_matches_numpy_with_tied_phases() -> None: + # Regression: equal folded phases must sort *stably* and identically on numpy and + # torch. An unstable sort (raw numpy quicksort) breaks ties differently from the + # array-API stable sort the torch path uses — and quicksort's tie order is + # platform-dependent, so this passed locally but failed in CI. Force heavy ties + # (times on a period/8 lattice fold to only 8 distinct phases) to pin it everywhere. + from cuperiod.methods.string_length import string_length + + rng = np.random.default_rng(0) + period = 2.0 + t = 2458000.0 + (rng.integers(0, 400, 400) * (period / 8)).astype(float) + y = rng.normal(0.0, 1.0, t.size) + periods = np.linspace(1.5, 3.0, 200) + ref = string_length(t, y, periods, backend="numpy") + tor = string_length(t, y, periods, backend="torch:cpu") + assert float(np.max(np.abs(ref - tor))) < 1e-9 + + @requires_torch def test_mhaov_multiband_torch_honours_precision() -> None: # Regression: the multiband wrapper must forward ``precision`` to the torch compute.