torchflat 0.8.0__tar.gz

torchflat-0.8.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Omar Khan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

torchflat-0.8.0/PKG-INFO

@@ -0,0 +1,234 @@
+Metadata-Version: 2.4
+Name: torchflat
+Version: 0.8.0
+Summary: GPU-accelerated photometric preprocessing with UMI detrending for exoplanet transit searches
+Author: Omar Khan
+License: MIT
+Project-URL: Homepage, https://github.com/omarkhan2217/TorchFlat
+Project-URL: Repository, https://github.com/omarkhan2217/TorchFlat
+Keywords: astronomy,exoplanets,transits,TESS,GPU,PyTorch,detrending,UMI
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Astronomy
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.1.0
+Requires-Dist: numpy>=1.24.0
+Requires-Dist: scipy>=1.10.0
+Provides-Extra: test
+Requires-Dist: pytest>=7.0; extra == "test"
+Requires-Dist: pytest-benchmark; extra == "test"
+Requires-Dist: wotan; extra == "test"
+Provides-Extra: dev
+Requires-Dist: torchflat[test]; extra == "dev"
+Requires-Dist: astropy; extra == "dev"
+Dynamic: license-file
+
+# TorchFlat
+
+**GPU-native photometric preprocessing pipeline for exoplanet transit searches.**
+
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
+
+TorchFlat replaces the standard CPU preprocessing workflow (quality filtering, gap handling, sigma clipping, detrending, normalization, and windowing) with a GPU-accelerated pipeline. It uses **UMI** (Unified Median Iterative), a novel asymmetric robust location estimator implemented as a fused HIP/CUDA kernel, to detrend light curves faster and more accurately than existing methods.
+
+## Performance
+
+Benchmarked on AMD Radeon RX 9060 XT (16 GB VRAM) with real TESS sector 6 data (19,618 stars):
+
+| Method | Rate | Full Sector | Speedup |
+|--------|------|-------------|---------|
+| Celix wotan 12-worker | 4.2 stars/sec | ~78 min | baseline |
+| TorchFlat v0.5.0 (hybrid) | 59.3 stars/sec | ~5.5 min | 14.2x |
+| **TorchFlat v0.8.0 + UMI kernel** | **139 stars/sec** | **~2.4 min** | **33x** |
+
+### Transit Depth Recovery Accuracy
+
+Injection-recovery test on 200 real TESS stars, median per-star error (lower = better):
+
+| Depth | wotan biweight | TorchFlat UMI | Winner |
+|-------|---------------|---------------|--------|
+| 0.1% (super-Earths) | 19.8% | **14.7%** | TorchFlat |
+| 0.3% (sub-Neptunes) | 8.4% | **3.7%** | TorchFlat |
+| 0.5% (Neptunes) | 1.5% | **1.5%** | tie |
+| 1.0% (hot Jupiters) | **0.4%** | 0.7% | wotan |
+| 5.0% (deep transits) | **0.0%** | 0.1% | both perfect |
+
+TorchFlat is more accurate at the transit depths where most detectable planets live (0.1-0.5%).
+
+Validated on 3 TESS sectors (6, 7, 12), 42 confirmed planets, and 1000-star train/test split. Results in `results/`.
+
+## The UMI Algorithm
+
+UMI (Unified Median Iterative) is a three-phase robust location estimator:
+
+1. **Quickselect median** -- exact median via O(n) selection algorithm, computed per-thread on GPU
+2. **Upper-RMS scale** -- RMS of points above the median only. Transit dips never contaminate the scale estimate, giving a tighter and more accurate noise measurement than standard MAD
+3. **Asymmetric bisquare iterations** -- weighted location refinement where downward deviations (transit dips) are penalized 1.5x more than upward ones
+
+The asymmetric weight function exploits the fact that **transits are always below the continuum**. Standard biweight treats dips and spikes equally. UMI penalizes dips more aggressively, so the trend stays above the transit and transit depth is preserved.
+
+All three phases run in a single fused GPU kernel call -- median, upper-RMS, and 5 iterations happen per-thread with zero global memory traffic between steps.
+
+When the GPU kernel is not available (no ROCm/CUDA toolkit), UMI falls back to a pure-PyTorch path using torch.sort for median + upper-RMS scale.
+
+## Installation
+
+```bash
+git clone https://github.com/omarkhan2217/TorchFlat.git
+cd TorchFlat
+pip install -e .
+```
+
+**Requirements:** PyTorch >= 2.1.0, NumPy >= 1.24.0, Numba >= 0.57.0, SciPy >= 1.10.0
+
+Works with both **NVIDIA CUDA** and **AMD ROCm** (via PyTorch's unified CUDA API). The UMI kernel compiles automatically on first use via JIT (requires ROCm SDK or CUDA toolkit).
+
+## Quick Start
+
+### Process a TESS sector
+
+```python
+import numpy as np
+import torchflat
+
+star_data = [
+    {
+        "time": np.load("star_001_time.npy"),
+        "pdcsap_flux": np.load("star_001_pdcsap.npy"),
+        "sap_flux": np.load("star_001_sap.npy"),
+        "quality": np.load("star_001_quality.npy"),
+    }
+    # ... for each star in the sector
+]
+
+results, skipped = torchflat.preprocess_sector(
+    star_data,
+    device="cuda",
+)
+
+for i, result in enumerate(results):
+    if not result:
+        continue
+    windows = result["windows_2048"]
+    trend = result["trend"]
+```
+
+### Standalone UMI detrending
+
+```python
+import torch
+from torchflat import umi_detrend
+
+# flux, time, valid_mask, segment_id are [B, L] tensors on GPU
+detrended, trend = umi_detrend(
+    flux, time, valid_mask, segment_id,
+    window_length_days=0.5,
+    asymmetry=2.0,       # 2.0=best accuracy, 1.0=variable stars, 1.5=mixed
+)
+```
+
+## Architecture
+
+TorchFlat implements two processing tracks:
+
+- **Track A (Transit Search):** Quality filter > gap interpolation > sigma clipping > UMI detrending > normalization > multi-scale window extraction
+- **Track B (Anomaly Detection):** Quality filter > gap interpolation > conservative clipping > FFT highpass filter > MAD normalization > fixed-length padding
+
+### UMI kernel
+
+The direct HIP/CUDA kernel (`torchflat/csrc/umi_kernel.cu`) runs one thread per (star, window position) pair. Each thread reads directly from the raw `[B, L]` flux array, no unfold or tensor copies needed:
+
+1. Reads W values from raw flux, checks segment validity inline
+2. Quickselect for exact median (O(n))
+3. Absolute deviations + quickselect for exact MAD
+4. 5 asymmetric bisquare iterations
+5. Writes the final location estimate
+
+VRAM usage: 319 MB for a 50-star batch. The kernel compiles via JIT on first import and is cached for subsequent runs.
+
+## Validation
+
+All validation results are saved as JSON in `results/`:
+
+| Validation | Result | File |
+|-----------|--------|------|
+| Asymmetry train/test split | optimal=1.5, generalizes across held-out stars | `asymmetry_validation.json` |
+| Known planet recovery | TorchFlat wins 24/41 (59%) confirmed planets | `known_planet_recovery.json` |
+| Multi-sector consistency | UMI wins 10/15 depth-sector combos across sectors 6,7,12 | `multisector_validation.json` |
+
+135/135 unit tests passing.
+
+## Benchmarks
+
+```bash
+# Full sector speed benchmark
+python benchmarks/bench_real_tess.py --data-dir /path/to/fits/sector_6 --n-stars 19618
+
+# Asymmetry parameter validation
+python benchmarks/validate_asymmetry.py
+
+# Known planet recovery
+python benchmarks/validate_known_planets.py
+
+# Multi-sector validation
+python benchmarks/validate_multisector.py
+```
+
+**Note:** Set `$env:TORCHFLAT_NO_KERNEL = "0"` (PowerShell) or `export TORCHFLAT_NO_KERNEL=0` (bash) to enable the UMI kernel.
+
+## API Reference
+
+### Main Entry Points
+
+- **`torchflat.preprocess_sector(star_data, ...)`** -- Full pipeline (Track A + Track B).
+- **`torchflat.preprocess_track_a(times, fluxes, qualities, ...)`** -- Track A only.
+- **`torchflat.preprocess_track_b(times, sap_fluxes, qualities, ...)`** -- Track B only.
+- **`torchflat.umi_detrend(flux, time, valid_mask, segment_id, ...)`** -- Standalone UMI kernel.
+
+### Key Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `device` | `"cuda"` | Torch device |
+| `window_length_days` | `0.5` | Sliding window width (days) |
+| `asymmetry` | `2.0` | Dip penalty: 2.0 (quiet stars), 1.5 (mixed), 1.0 (variable stars) |
+| `n_iter` | `5` | Number of bisquare iterations |
+| `cval` | `5.0` | Rejection threshold in MAD units |
+| `skip_track_b` | `False` | Skip Track B (FFT highpass) |
+| `window_scales` | 4 scales | `[(256,128), (512,256), (2048,512), (8192,2048)]` |
+| `dtype` | `float32` | Computation precision |
+
+## Development
+
+```bash
+git clone https://github.com/omarkhan2217/TorchFlat.git
+cd TorchFlat
+pip install -e ".[dev]"
+pytest tests/ -v
+```
+
+## Citation
+
+If you use TorchFlat in your research, please cite:
+
+```bibtex
+@software{torchflat,
+  author = {Khan, Omar},
+  title = {TorchFlat: GPU-Accelerated Photometric Preprocessing with UMI Detrending},
+  year = {2026},
+  url = {https://github.com/omarkhan2217/TorchFlat}
+}
+```
+
+## License
+
+MIT License. See [LICENSE](LICENSE) for details.

torchflat-0.8.0/README.md

@@ -0,0 +1,201 @@
+# TorchFlat
+
+**GPU-native photometric preprocessing pipeline for exoplanet transit searches.**
+
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
+[![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
+
+TorchFlat replaces the standard CPU preprocessing workflow (quality filtering, gap handling, sigma clipping, detrending, normalization, and windowing) with a GPU-accelerated pipeline. It uses **UMI** (Unified Median Iterative), a novel asymmetric robust location estimator implemented as a fused HIP/CUDA kernel, to detrend light curves faster and more accurately than existing methods.
+
+## Performance
+
+Benchmarked on AMD Radeon RX 9060 XT (16 GB VRAM) with real TESS sector 6 data (19,618 stars):
+
+| Method | Rate | Full Sector | Speedup |
+|--------|------|-------------|---------|
+| Celix wotan 12-worker | 4.2 stars/sec | ~78 min | baseline |
+| TorchFlat v0.5.0 (hybrid) | 59.3 stars/sec | ~5.5 min | 14.2x |
+| **TorchFlat v0.8.0 + UMI kernel** | **139 stars/sec** | **~2.4 min** | **33x** |
+
+### Transit Depth Recovery Accuracy
+
+Injection-recovery test on 200 real TESS stars, median per-star error (lower = better):
+
+| Depth | wotan biweight | TorchFlat UMI | Winner |
+|-------|---------------|---------------|--------|
+| 0.1% (super-Earths) | 19.8% | **14.7%** | TorchFlat |
+| 0.3% (sub-Neptunes) | 8.4% | **3.7%** | TorchFlat |
+| 0.5% (Neptunes) | 1.5% | **1.5%** | tie |
+| 1.0% (hot Jupiters) | **0.4%** | 0.7% | wotan |
+| 5.0% (deep transits) | **0.0%** | 0.1% | both perfect |
+
+TorchFlat is more accurate at the transit depths where most detectable planets live (0.1-0.5%).
+
+Validated on 3 TESS sectors (6, 7, 12), 42 confirmed planets, and 1000-star train/test split. Results in `results/`.
+
+## The UMI Algorithm
+
+UMI (Unified Median Iterative) is a three-phase robust location estimator:
+
+1. **Quickselect median** -- exact median via O(n) selection algorithm, computed per-thread on GPU
+2. **Upper-RMS scale** -- RMS of points above the median only. Transit dips never contaminate the scale estimate, giving a tighter and more accurate noise measurement than standard MAD
+3. **Asymmetric bisquare iterations** -- weighted location refinement where downward deviations (transit dips) are penalized 1.5x more than upward ones
+
+The asymmetric weight function exploits the fact that **transits are always below the continuum**. Standard biweight treats dips and spikes equally. UMI penalizes dips more aggressively, so the trend stays above the transit and transit depth is preserved.
+
+All three phases run in a single fused GPU kernel call -- median, upper-RMS, and 5 iterations happen per-thread with zero global memory traffic between steps.
+
+When the GPU kernel is not available (no ROCm/CUDA toolkit), UMI falls back to a pure-PyTorch path using torch.sort for median + upper-RMS scale.
+
+## Installation
+
+```bash
+git clone https://github.com/omarkhan2217/TorchFlat.git
+cd TorchFlat
+pip install -e .
+```
+
+**Requirements:** PyTorch >= 2.1.0, NumPy >= 1.24.0, Numba >= 0.57.0, SciPy >= 1.10.0
+
+Works with both **NVIDIA CUDA** and **AMD ROCm** (via PyTorch's unified CUDA API). The UMI kernel compiles automatically on first use via JIT (requires ROCm SDK or CUDA toolkit).
+
+## Quick Start
+
+### Process a TESS sector
+
+```python
+import numpy as np
+import torchflat
+
+star_data = [
+    {
+        "time": np.load("star_001_time.npy"),
+        "pdcsap_flux": np.load("star_001_pdcsap.npy"),
+        "sap_flux": np.load("star_001_sap.npy"),
+        "quality": np.load("star_001_quality.npy"),
+    }
+    # ... for each star in the sector
+]
+
+results, skipped = torchflat.preprocess_sector(
+    star_data,
+    device="cuda",
+)
+
+for i, result in enumerate(results):
+    if not result:
+        continue
+    windows = result["windows_2048"]
+    trend = result["trend"]
+```
+
+### Standalone UMI detrending
+
+```python
+import torch
+from torchflat import umi_detrend
+
+# flux, time, valid_mask, segment_id are [B, L] tensors on GPU
+detrended, trend = umi_detrend(
+    flux, time, valid_mask, segment_id,
+    window_length_days=0.5,
+    asymmetry=2.0,       # 2.0=best accuracy, 1.0=variable stars, 1.5=mixed
+)
+```
+
+## Architecture
+
+TorchFlat implements two processing tracks:
+
+- **Track A (Transit Search):** Quality filter > gap interpolation > sigma clipping > UMI detrending > normalization > multi-scale window extraction
+- **Track B (Anomaly Detection):** Quality filter > gap interpolation > conservative clipping > FFT highpass filter > MAD normalization > fixed-length padding
+
+### UMI kernel
+
+The direct HIP/CUDA kernel (`torchflat/csrc/umi_kernel.cu`) runs one thread per (star, window position) pair. Each thread reads directly from the raw `[B, L]` flux array, no unfold or tensor copies needed:
+
+1. Reads W values from raw flux, checks segment validity inline
+2. Quickselect for exact median (O(n))
+3. Absolute deviations + quickselect for exact MAD
+4. 5 asymmetric bisquare iterations
+5. Writes the final location estimate
+
+VRAM usage: 319 MB for a 50-star batch. The kernel compiles via JIT on first import and is cached for subsequent runs.
+
+## Validation
+
+All validation results are saved as JSON in `results/`:
+
+| Validation | Result | File |
+|-----------|--------|------|
+| Asymmetry train/test split | optimal=1.5, generalizes across held-out stars | `asymmetry_validation.json` |
+| Known planet recovery | TorchFlat wins 24/41 (59%) confirmed planets | `known_planet_recovery.json` |
+| Multi-sector consistency | UMI wins 10/15 depth-sector combos across sectors 6,7,12 | `multisector_validation.json` |
+
+135/135 unit tests passing.
+
+## Benchmarks
+
+```bash
+# Full sector speed benchmark
+python benchmarks/bench_real_tess.py --data-dir /path/to/fits/sector_6 --n-stars 19618
+
+# Asymmetry parameter validation
+python benchmarks/validate_asymmetry.py
+
+# Known planet recovery
+python benchmarks/validate_known_planets.py
+
+# Multi-sector validation
+python benchmarks/validate_multisector.py
+```
+
+**Note:** Set `$env:TORCHFLAT_NO_KERNEL = "0"` (PowerShell) or `export TORCHFLAT_NO_KERNEL=0` (bash) to enable the UMI kernel.
+
+## API Reference
+
+### Main Entry Points
+
+- **`torchflat.preprocess_sector(star_data, ...)`** -- Full pipeline (Track A + Track B).
+- **`torchflat.preprocess_track_a(times, fluxes, qualities, ...)`** -- Track A only.
+- **`torchflat.preprocess_track_b(times, sap_fluxes, qualities, ...)`** -- Track B only.
+- **`torchflat.umi_detrend(flux, time, valid_mask, segment_id, ...)`** -- Standalone UMI kernel.
+
+### Key Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `device` | `"cuda"` | Torch device |
+| `window_length_days` | `0.5` | Sliding window width (days) |
+| `asymmetry` | `2.0` | Dip penalty: 2.0 (quiet stars), 1.5 (mixed), 1.0 (variable stars) |
+| `n_iter` | `5` | Number of bisquare iterations |
+| `cval` | `5.0` | Rejection threshold in MAD units |
+| `skip_track_b` | `False` | Skip Track B (FFT highpass) |
+| `window_scales` | 4 scales | `[(256,128), (512,256), (2048,512), (8192,2048)]` |
+| `dtype` | `float32` | Computation precision |
+
+## Development
+
+```bash
+git clone https://github.com/omarkhan2217/TorchFlat.git
+cd TorchFlat
+pip install -e ".[dev]"
+pytest tests/ -v
+```
+
+## Citation
+
+If you use TorchFlat in your research, please cite:
+
+```bibtex
+@software{torchflat,
+  author = {Khan, Omar},
+  title = {TorchFlat: GPU-Accelerated Photometric Preprocessing with UMI Detrending},
+  year = {2026},
+  url = {https://github.com/omarkhan2217/TorchFlat}
+}
+```
+
+## License
+
+MIT License. See [LICENSE](LICENSE) for details.

torchflat-0.8.0/pyproject.toml

@@ -0,0 +1,56 @@
+[build-system]
+requires = ["setuptools>=68.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "torchflat"
+version = "0.8.0"
+description = "GPU-accelerated photometric preprocessing with UMI detrending for exoplanet transit searches"
+readme = "README.md"
+license = {text = "MIT"}
+authors = [{name = "Omar Khan"}]
+requires-python = ">=3.10"
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering :: Astronomy",
+]
+keywords = ["astronomy", "exoplanets", "transits", "TESS", "GPU", "PyTorch", "detrending", "UMI"]
+dependencies = [
+    "torch>=2.1.0",
+    "numpy>=1.24.0",
+    "scipy>=1.10.0",
+]
+
+[project.scripts]
+torchflat = "torchflat.cli:main"
+
+[project.urls]
+Homepage = "https://github.com/omarkhan2217/TorchFlat"
+Repository = "https://github.com/omarkhan2217/TorchFlat"
+
+[project.optional-dependencies]
+test = [
+    "pytest>=7.0",
+    "pytest-benchmark",
+    "wotan",
+]
+dev = [
+    "torchflat[test]",
+    "astropy",
+]
+
+[tool.setuptools.packages.find]
+include = ["torchflat*"]
+
+[tool.setuptools.package-data]
+torchflat = ["csrc/*.cu", "csrc/*.cpp", "csrc/build/*.cpp"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

torchflat-0.8.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

torchflat-0.8.0/tests/test_batching.py

@@ -0,0 +1,143 @@
+"""Tests for torchflat.batching."""
+
+import numpy as np
+import pytest
+import torch
+
+from torchflat.batching import (
+    assemble_batch,
+    bucket_stars,
+    compute_max_batch,
+    cpu_prescan,
+    estimate_peak_vram,
+)
+
+CADENCE = 2.0 / 1440.0
+
+
+def _make_star(n_points: int = 18000, seed: int = 0):
+    """Create a clean synthetic star as numpy arrays."""
+    rng = np.random.default_rng(seed)
+    time = np.arange(n_points, dtype=np.float64) * CADENCE
+    flux = (1.0 + rng.normal(0, 0.001, n_points)).astype(np.float32)
+    quality = np.zeros(n_points, dtype=np.int32)
+    return time, flux, quality
+
+
+class TestCpuPrescan:
+
+    def test_clean_star(self):
+        t, f, q = _make_star(18000)
+        results = cpu_prescan([t], [f], [q])
+        assert len(results) == 1
+        r = results[0]
+        assert r["index"] == 0
+        assert r["n_valid"] == 18000
+        assert r["degenerate"] is False
+        assert r["post_filter_length"] >= r["n_valid"]
+
+    def test_degenerate_too_few_points(self):
+        t, f, q = _make_star(50)
+        results = cpu_prescan([t], [f], [q])
+        assert results[0]["degenerate"] is True
+        assert results[0]["degenerate_reason"] == "too_few_valid_points"
+
+    def test_degenerate_segment_too_short(self):
+        # Star with many small segments, none long enough for biweight window
+        n = 5000
+        t = np.arange(n, dtype=np.float64) * CADENCE
+        # Insert large gaps every 100 points -> segments of ~100, window=360
+        for i in range(100, n, 100):
+            t[i:] += 20 * CADENCE
+        f = np.ones(n, dtype=np.float32)
+        q = np.zeros(n, dtype=np.int32)
+        results = cpu_prescan([t], [f], [q], window_samples=360)
+        assert results[0]["degenerate"] is True
+        assert results[0]["degenerate_reason"] == "segment_too_short"
+
+    def test_multiple_stars(self):
+        stars = [_make_star(n, seed=i) for i, n in enumerate([18000, 50, 15000])]
+        ts, fs, qs = zip(*stars)
+        results = cpu_prescan(list(ts), list(fs), list(qs))
+        assert len(results) == 3
+        assert results[0]["degenerate"] is False
+        assert results[1]["degenerate"] is True
+        assert results[2]["degenerate"] is False
+
+
+class TestBucketStars:
+
+    def test_bucketing_correctness(self):
+        prescan = [
+            {"index": 0, "post_filter_length": 14500, "degenerate": False},
+            {"index": 1, "post_filter_length": 15500, "degenerate": False},
+            {"index": 2, "post_filter_length": 14800, "degenerate": False},
+            {"index": 3, "post_filter_length": 50, "degenerate": True, "degenerate_reason": "x"},
+        ]
+        buckets = bucket_stars(prescan, bucket_width=1000)
+        # Star 3 is degenerate -> excluded
+        total_stars = sum(len(b["star_indices"]) for b in buckets)
+        assert total_stars == 3
+
+    def test_bucket_width(self):
+        prescan = [
+            {"index": i, "post_filter_length": 14000 + i * 200, "degenerate": False}
+            for i in range(10)
+        ]
+        buckets = bucket_stars(prescan, bucket_width=1000)
+        for b in buckets:
+            # All stars in bucket should fit within pad_length
+            for idx in b["star_indices"]:
+                pfl = prescan[idx]["post_filter_length"]
+                assert pfl <= b["pad_length"]
+
+
+class TestAssembleBatch:
+
+    def test_padding_no_data_leakage(self):
+        t1, f1, q1 = _make_star(100, seed=1)
+        t2, f2, q2 = _make_star(80, seed=2)
+        batch = assemble_batch([0, 1], [t1, t2], [f1, f2], [q1, q2], 120, torch.device("cpu"))
+        # Padded positions should be zero
+        assert (batch["flux"][0, 100:] == 0).all()
+        assert (batch["flux"][1, 80:] == 0).all()
+
+    def test_padding_mask(self):
+        t, f, q = _make_star(100)
+        batch = assemble_batch([0], [t], [f], [q], 150, torch.device("cpu"))
+        assert batch["valid_mask"][0, :100].all()
+        assert not batch["valid_mask"][0, 100:].any()
+
+    def test_lengths(self):
+        stars = [_make_star(n) for n in [100, 200, 150]]
+        ts, fs, qs = zip(*stars)
+        batch = assemble_batch([0, 1, 2], list(ts), list(fs), list(qs), 250, torch.device("cpu"))
+        assert batch["lengths"].tolist() == [100, 200, 150]
+
+
+class TestVramEstimation:
+
+    def test_monotonic(self):
+        v1 = estimate_peak_vram(15000, 360)
+        v2 = estimate_peak_vram(20000, 360)
+        assert v2 > v1
+
+    def test_reasonable_range(self):
+        v = estimate_peak_vram(20000, 360)
+        # Should be between 100MB and 500MB per star
+        assert 100 * 1024**2 < v < 500 * 1024**2
+
+
+class TestComputeMaxBatch:
+
+    def test_override(self):
+        assert compute_max_batch(20000, max_batch_override=10) == 10
+
+    def test_budget(self):
+        mb_12 = compute_max_batch(20000, vram_budget_gb=12.0)
+        mb_6 = compute_max_batch(20000, vram_budget_gb=6.0)
+        assert mb_12 > mb_6
+        assert mb_6 >= 1
+
+    def test_cpu_fallback(self):
+        assert compute_max_batch(20000, device=torch.device("cpu")) == 1