diff --git a/CHANGELOG.md b/CHANGELOG.md index b5f7f70..64da595 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,43 @@ 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.** 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, 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. +- `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/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 " 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/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/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/src/cuperiod/core/_arrayapi.py b/src/cuperiod/core/_arrayapi.py new file mode 100644 index 0000000..74fc32f --- /dev/null +++ b/src/cuperiod/core/_arrayapi.py @@ -0,0 +1,190 @@ +"""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``. + + ``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) + 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..0cd78d9 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." ) @@ -150,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") @@ -188,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( @@ -199,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") @@ -236,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( @@ -247,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") @@ -283,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( @@ -294,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") @@ -328,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( @@ -339,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") @@ -400,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/_bls_core.py b/src/cuperiod/methods/_bls_core.py index 9d665d7..f6e75da 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,71 @@ 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) + + # 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 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 + ) + # 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) 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"]) + t_ref, + 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/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/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/mhaov.py b/src/cuperiod/methods/mhaov.py index a4b6501..8e3bceb 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 @@ -43,15 +50,22 @@ #: 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: """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 +92,11 @@ 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 + 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) for start in range(0, n_freq, batch): stop = min(start + batch, n_freq) @@ -92,7 +108,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 +119,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 +150,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 +181,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 +207,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 +223,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 +257,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 +284,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 +320,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 +348,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..56df71d 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,33 @@ 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``. + + 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]) - 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 +85,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). @@ -100,15 +120,31 @@ 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) + 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( - _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, ) @@ -121,7 +157,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 +193,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 04143c8..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 @@ -95,7 +103,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( @@ -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/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 8815a35..8eecb96 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,30 @@ 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 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_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..b413071 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,51 @@ 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) + + +@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_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"]) diff --git a/tests/test_gls.py b/tests/test_gls.py index bdb49c1..de4b9d7 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,33 @@ 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) + + +@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 new file mode 100644 index 0000000..0fc69ec --- /dev/null +++ b/tests/test_torch_bonus.py @@ -0,0 +1,128 @@ +"""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 cuperiod.methods.mhaov import aov_power +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 + + +@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. + # 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