From 6f5c79b5f2daed0068213239b51cc8187c08f5df Mon Sep 17 00:00:00 2001 From: Tharindu Jayasinghe Date: Tue, 30 Jun 2026 12:30:58 -0700 Subject: [PATCH] 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