makewfs 2.1.0__tar.gz → 2.1.1__tar.gz

{makewfs-2.1.0 → makewfs-2.1.1}/AGENTS.md

@@ -255,15 +255,18 @@ MPLBACKEND=Agg python examples/gallery.py
   fixed shapes and never synchronize with the host (no `scalar`, `.item()`,
   `to_host`, or data-dependent shapes). A failed capture silently falls back to
   eager execution and records why in `engine._graph.failure`; check it after
-  changing that method.
+  changing that method. Anything the graph uses must outlive it: its FFTs run
+  on the engine's own cuFFT plans (`plans=` in `fft_axis`), never CuPy's global
+  plan cache, which evicts and frees plans the graph still references.
 - The compiled SH kernel reads the OPD as float64; `_CompiledShackHartmannExecutor.render`
   widens a float32 lenslet-grid resample (exact). Its static thread and
   shared-memory checks cannot see register pressure, so construction also
   checks the compiled kernel's `max_threads_per_block` and falls back.
-- `ArrayBackend.next_fast_length` uses SciPy on the CPU and CuPy on the GPU,
-  which disagree when SciPy picks a factor of 11 (for example 33, 66, 99). The
-  pyramid's FFT size, and therefore its result, can then differ between
-  devices; parity tests use geometries where they agree.
+- Grid sizes set the physics, so they must never depend on the backend.
+  `backend.next_fast_length` is one 7-smooth rule for CPU and GPU (it equals
+  CuPy's `next_fast_len`); do not call `scipy.fft.next_fast_len` or
+  `cupyx.scipy.fft.next_fast_len` for a size. SciPy's admits 11, which made the
+  pyramid pad differently on the two devices before 2.1.1.
 
 ## Documentation and examples
 

{makewfs-2.1.0 → makewfs-2.1.1}/CHANGELOG.md

@@ -4,6 +4,39 @@ All notable changes to `makewfs` are documented here.
 
 ## [Unreleased]
 
+## [2.1.1] - 2026-10-07
+
+### Fixed
+
+- **CPU and GPU pyramids used different propagation grids for some
+  geometries** ([#15](https://github.com/jacotay7/makewfs/issues/15)). The FFT
+  size came from `scipy.fft.next_fast_len` on the CPU and
+  `cupyx.scipy.fft.next_fast_len` on the GPU, and SciPy also admits a factor
+  of 11. Whenever it picked one (requests such as 33, 66 or 99 samples), the
+  CPU padded to 33/66/99 while the GPU padded to 35/70/100, and the photon
+  rates differed by up to 46% (24-pixel pupil, 9-pixel separation, 12-point
+  modulation at 3 lambda/D). Both backends now use one rule: the smallest
+  length with no prime factor above 7, which equals CuPy's choice, so **GPU
+  results are unchanged** and CPU and GPU now agree to rounding.
+  - **Behaviour change on the CPU** for every pyramid whose old SciPy size had a
+    factor of 11: the grid grows by at most 9.1% and the rate changes. At the
+    default `fft_oversampling = 2` this affects detector widths
+    (`pixels_across_pupil + pupil_separation_pixels + 2 * detector_margin_pixels`)
+    of 11, 22, 33, 38, 43, 44, 55, 65, 66, 76, 77, 82, 88, 99, 109, 110,
+    113-115, 121, 129-132, 136, 137, 148, 151-154, 163-165, 176, 181, 197, 198,
+    217-220, 226-231, 241, 242, 246, 247, 263-269, 271-275 and 295-297 pixels
+    (up to 300). Overall, 1,342 of the requests 1-4096 change size.
+  - **Unaffected:** the shipped example and benchmark pyramids (grids of 288,
+    108, 160 and 216 samples) and every Shack-Hartmann configuration, whose FFT
+    sizes follow from the configured sampling.
+  - Frames now record the grid as `wfs_pyramid_fft_size_px`.
+- **The 2.1.0 GPU pyramid's CUDA graph could replay freed cuFFT memory.** It
+  used plans from CuPy's global plan cache, which evicts them once other
+  transforms fill it (for example many sensors or user FFTs in one process),
+  and replay then read freed memory (`cudaErrorIllegalAddress`). The pyramid
+  now owns its plans, built exactly as CuPy builds them, so results are
+  unchanged.
+
 ## [2.1.0] - 2026-10-07
 
 ### Performance

{makewfs-2.1.0 → makewfs-2.1.1}/CITATION.cff

@@ -5,7 +5,7 @@ type: software
 authors:
   - family-names: Taylor
     given-names: Jacob
-version: 2.1.0
+version: 2.1.1
 date-released: 2026-10-07
 license: MIT
 repository-code: "https://github.com/jacotay7/makewfs"

{makewfs-2.1.0 → makewfs-2.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: makewfs
-Version: 2.1.0
+Version: 2.1.1
 Summary: Configuration-driven adaptive-optics wavefront sensor image simulation.
 Project-URL: Homepage, https://github.com/jacotay7/makewfs
 Project-URL: Documentation, https://jacotay7.github.io/makewfs/

{makewfs-2.1.0 → makewfs-2.1.1}/docs/configuration.md

@@ -116,7 +116,7 @@ kernel file hash and all normalized states are included in frame provenance.
 | `detector.roi` | Optional full-detector ROI table with non-negative `left_px`/`top_px` and positive `width_px`/`height_px`. Its shape must match the optical mosaic. |
 | `numerics.dtype` | Optical real precision, `"float32"` or `"float64"` (default `"float64"`). |
 | `numerics.device` | Execution device, `"cpu"` (default) or `"gpu"`. GPU requires the `makewfs[gpu]` extra and a GPU-capable `getframes`; optical, truth, and ADU arrays stay device-resident. |
-| `numerics.fft_oversampling` | Positive FFT integration oversampling (default `2`). Also scales the pyramid propagation grid so diffraction beyond the detector crop is discarded instead of wrapping onto the pupil rims. |
+| `numerics.fft_oversampling` | Positive FFT integration oversampling (default `2`). Also scales the pyramid propagation grid so diffraction beyond the detector crop is discarded instead of wrapping onto the pupil rims. The pyramid grid side is the smallest length with no prime factor above 7 that is at least `fft_oversampling` times the detector width; the rule is the same on CPU and GPU, and frames record it as `wfs_pyramid_fft_size_px`. |
 | `numerics.fft_workers` | Positive `scipy.fft` worker count (default `1`). |
 | `numerics.pupil_samples_per_lenslet` | Optional integer ≥4: samples per lenslet pitch of the SH pupil grid that holds the amplitude mask, a custom mask, and lenslet illumination. Default `max(8, ceil(max(input.shape) / lenslets_across_pupil))`. Each lenslet field is then propagated on an integer refinement of this grid that is at least as fine as the widest detector window in lenslet `lambda/d` (see [Lenslet-field sampling](shack-hartmann.md#lenslet-field-sampling)); a value already that fine is used unchanged. |
 | `numerics.pupil_supersampling` | Positive analytic pupil boundary sub-sampling factor (default `1`). |

{makewfs-2.1.0 → makewfs-2.1.1}/src/makewfs/__about__.py

@@ -1,3 +1,3 @@
 """Package version metadata."""
 
-__version__ = "2.1.0"
+__version__ = "2.1.1"

{makewfs-2.1.0 → makewfs-2.1.1}/src/makewfs/api.py

@@ -58,6 +58,7 @@ class WavefrontSensor:
             seed=None,
             source_states=self.engine.source_states,
             file_digests=self.engine.file_digests,
+            pyramid_fft_size=getattr(self.engine, "nfft", None),
         )
         self._rms_weights = pupil_weights(self._input_pupil_intensity(), backend=self.backend)
 

{makewfs-2.1.0 → makewfs-2.1.1}/src/makewfs/backend.py

@@ -255,19 +255,53 @@ class ArrayBackend:
         norm: str,
         inverse: bool = False,
         workers: int = 1,
+        plans: dict[Any, Any] | None = None,
     ) -> Any:
         """One-dimensional FFT along ``axis`` of an array the caller owns.
 
         The pruned two-dimensional transforms are built from these passes. On
         the CPU the input may be overwritten, so pass only a temporary.
+
+        On a GPU, ``plans`` is a caller-owned cache of cuFFT plans. CuPy's own
+        plan cache is global and evicts plans when other transforms run, which
+        frees memory a captured CUDA graph still uses; a graph's owner passes
+        its own dict so the plans live exactly as long as the graph. The plan
+        is the same ``Plan1d`` CuPy would build, so values are unchanged.
         """
         if self.is_cpu:
             from scipy import fft
 
             transform = fft.ifft if inverse else fft.fft
             return transform(array, axis=axis, norm=norm, workers=workers, overwrite_x=True)
+        return self._device_fft_axis(  # pragma: no cover - optional CUDA execution
+            array, axis=axis, norm=norm, inverse=inverse, plans=plans
+        )
+
+    def _device_fft_axis(  # pragma: no cover - optional CUDA execution
+        self,
+        array: Any,
+        *,
+        axis: int,
+        norm: str,
+        inverse: bool,
+        plans: dict[Any, Any] | None,
+    ) -> Any:
         transform = self.xp.fft.ifft if inverse else self.xp.fft.fft
-        return transform(array, axis=axis, norm=norm)
+        if plans is None:
+            return transform(array, axis=axis, norm=norm)
+        # CuPy transforms ``axis`` as the last axis of a contiguous copy, with a
+        # ``Plan1d`` keyed by length, type and batch; build that same plan.
+        length = int(array.shape[axis])
+        batch = int(array.size) // length
+        key = (length, array.dtype.str, batch)
+        plan = plans.get(key)
+        if plan is None:
+            from cupy.cuda import cufft
+
+            kind = cufft.CUFFT_C2C if array.dtype == np.complex64 else cufft.CUFFT_Z2Z
+            plan = plans[key] = cufft.Plan1d(length, kind, batch)
+        with plan:
+            return transform(array, axis=axis, norm=norm)
 
     def pruned_fft2(
         self,
@@ -280,6 +314,7 @@ class ArrayBackend:
         inverse: bool = False,
         axes: tuple[int, int] = (-2, -1),
         workers: int = 1,
+        plans: dict[Any, Any] | None = None,
     ) -> Any:
         """Unitary square 2-D FFT that skips zero input and unwanted output lines.
 
@@ -307,7 +342,7 @@ class ArrayBackend:
         Keep ``axes`` equal to the pass order of the transform being replaced,
         because rounding depends on it: SciPy's ``fft2`` runs the listed axes in
         order when ``overwrite_x=True`` but the last axis first when it
-        allocates its output.
+        allocates its output. ``plans`` is passed to :meth:`fft_axis`.
         """
         if (output_start is None) != (output_length is None):
             raise ValueError("output_start and output_length go together")
@@ -317,7 +352,12 @@ class ArrayBackend:
         if input_start is not None:
             spectrum = self._embed(spectrum, input_start, size=size, axis=first_axis)
         spectrum = self.fft_axis(
-            spectrum, axis=first_axis, norm=first_norm, inverse=inverse, workers=workers
+            spectrum,
+            axis=first_axis,
+            norm=first_norm,
+            inverse=inverse,
+            workers=workers,
+            plans=plans,
         )
         if output_start is not None:
             assert output_length is not None
@@ -325,7 +365,12 @@ class ArrayBackend:
         if input_start is not None:
             spectrum = self._embed(spectrum, input_start, size=size, axis=second_axis)
         spectrum = self.fft_axis(
-            spectrum, axis=second_axis, norm=second_norm, inverse=inverse, workers=workers
+            spectrum,
+            axis=second_axis,
+            norm=second_norm,
+            inverse=inverse,
+            workers=workers,
+            plans=plans,
         )
         if output_start is not None:
             assert output_length is not None
@@ -412,13 +457,12 @@ class ArrayBackend:
         return gaussian_filter(array, sigma=sigma, mode="constant")
 
     def next_fast_length(self, value: int) -> int:
-        """Return an FFT-friendly length using this backend's implementation."""
-        if self.is_cpu:
-            from scipy.fft import next_fast_len
-        else:  # pragma: no cover - GPU optional
-            from cupyx.scipy.fft import next_fast_len
+        """Return the shared FFT length for ``value`` (see :func:`next_fast_length`).
 
-        return int(next_fast_len(value))
+        Deliberately independent of the backend: a grid size sets the physics
+        (padding, aliasing), so the CPU and the GPU must choose the same one.
+        """
+        return next_fast_length(value)
 
     def scalar(self, value: Any) -> float:
         """Extract one host scalar at an explicit metadata/geometry boundary."""
@@ -523,9 +567,31 @@ def centered_fft_intensity(
     )
 
 
+_FAST_FFT_FACTORS = (2, 3, 5, 7)
+
+
 def next_fast_length(value: int) -> int:
-    """Return a convenient CPU FFT length."""
-    return cpu_backend().next_fast_length(value)
+    """Return the smallest length ``>= value`` with no prime factor above 7.
+
+    One rule for every backend, so a configuration propagates on the same grid
+    on the CPU and the GPU. cuFFT has optimized kernels for radices 2, 3, 5
+    and 7, and SciPy's pocketfft handles them efficiently too. SciPy's own
+    ``next_fast_len`` also admits 11, which CuPy's does not; before 2.1.1 the
+    pyramid therefore padded differently on the two devices whenever SciPy
+    picked a factor of 11 (for example 33, 66 or 99 against 35, 70 or 100).
+    This rule equals CuPy's ``next_fast_len``.
+    """
+    if value < 1:
+        raise ValueError("FFT length must be positive")
+    length = int(value)
+    while True:
+        remainder = length
+        for factor in _FAST_FFT_FACTORS:
+            while remainder % factor == 0:
+                remainder //= factor
+        if remainder == 1:
+            return length
+        length += 1
 
 
 __all__ = [

{makewfs-2.1.0 → makewfs-2.1.1}/src/makewfs/provenance.py

@@ -52,12 +52,15 @@ def metadata(
     seed: int | None,
     source_states: tuple[SourceState, ...] | None = None,
     file_digests: dict[str, str] | None = None,
+    pyramid_fft_size: int | None = None,
 ) -> dict[str, Any]:
     """Build serializable metadata for an ideal or detector frame.
 
     ``opd_rms_m`` is the input OPD's pupil-weighted, piston-removed RMS and
     ``opd_rms_unweighted_m`` its unweighted whole-grid RMS with piston
     included, both already reduced by the caller (aocore CONVENTIONS 4.1).
+    ``pyramid_fft_size`` is the side of the pyramid's square propagation grid,
+    recorded as ``wfs_pyramid_fft_size_px`` because it sets the physics.
     """
     states = iter_source_states(config) if source_states is None else source_states
     result: dict[str, Any] = {
@@ -97,6 +100,8 @@ def metadata(
             "lower_left",
             "lower_right",
         ]
+        if pyramid_fft_size is not None:
+            result["wfs_pyramid_fft_size_px"] = int(pyramid_fft_size)
     return result
 
 

{makewfs-2.1.0 → makewfs-2.1.1}/src/makewfs/sensors/pyramid.py

@@ -50,7 +50,8 @@ class PyramidEngine(SensorEngine):
         self.output_shape = (pixels + separation + 2 * margin, pixels + separation + 2 * margin)
         # Oversample the propagation grid so the diffraction halo lands outside
         # the detector crop instead of wrapping onto the pupil rims; cropped
-        # flux is reported as captured rate, never renormalized.
+        # flux is reported as captured rate, never renormalized. The length
+        # rule is shared by every backend, so CPU and GPU use the same grid.
         self.nfft = self.backend.next_fast_length(
             config.numerics.fft_oversampling * max(self.output_shape)
         )
@@ -122,6 +123,9 @@ class PyramidEngine(SensorEngine):
         self._graph = (
             None if self.backend.is_cpu else _CudaGraphReplay(self.backend.xp, self._propagate)
         )
+        # The graph replays cuFFT work, so its plans must outlive it rather than
+        # live in CuPy's global, evicting plan cache (see ``fft_axis``).
+        self._fft_plans: dict[Any, Any] = {}
 
     def _make_pyramid_mask(self) -> NDArray[Any]:
         """Build four signed focal-plane ramps that separate the pupils."""
@@ -243,6 +247,7 @@ class PyramidEngine(SensorEngine):
                 size=self.nfft,
                 input_start=self._pupil_start,
                 workers=self.config.numerics.fft_workers,
+                plans=self._fft_plans,
             )
             focal *= self._unshifted_mask
             exit_pupil = self.backend.pruned_fft2(
@@ -252,6 +257,7 @@ class PyramidEngine(SensorEngine):
                 output_length=self._base_output_shape[0],
                 inverse=True,
                 workers=self.config.numerics.fft_workers,
+                plans=self._fft_plans,
             )
             cropped = self.backend.abs(exit_pupil) ** 2
             mosaic = self.backend.mean(cropped, axis=0)

{makewfs-2.1.0 → makewfs-2.1.1}/tests/test_gpu_backend.py

@@ -492,3 +492,55 @@ def test_pyramid_cuda_graph_replay_equals_eager_and_owns_outputs(broadband: bool
     assert float(second.captured_rate_per_s) == float(eager.captured_rate_per_s)
     expected_cube = (3, *engine.output_shape) if broadband else (1, *engine.output_shape)
     assert second.spectral_photon_rate.shape == expected_cube
+
+
+@pytest.mark.gpu
+def test_pyramid_geometry_with_former_factor_eleven_grid_matches_cpu() -> None:
+    # 24 + 9 pixels at 2x oversampling requests 66 samples: SciPy used to pick
+    # 66 (2 x 3 x 11) on the CPU and CuPy 70 on the GPU, which changed the
+    # rate by up to 46% for this modulated geometry. Both now use 70.
+    cupy = _cupy()
+    config = _config("pyramid_minimal.toml")
+    assert config.pyramid is not None
+    config = replace(
+        config,
+        numerics=replace(config.numerics, dtype="float64"),
+        pyramid=replace(
+            config.pyramid,
+            pixels_across_pupil=24,
+            pupil_separation_pixels=9,
+            modulation_radius_lambda_over_d=3.0,
+            modulation_samples=12,
+        ),
+    )
+    cpu = WavefrontSensor(config)
+    gpu = WavefrontSensor(replace(config, numerics=replace(config.numerics, device="gpu")))
+    assert cpu.engine.nfft == gpu.engine.nfft == 70
+    opd = np.random.default_rng(7).normal(0.0, 5.0e-8, config.input.shape)
+
+    expected = cpu.photon_rate(opd)
+    actual = cupy.asnumpy(gpu.photon_rate(cupy.asarray(opd)))
+    np.testing.assert_allclose(actual, expected, rtol=1e-10, atol=1e-10 * expected.max())
+    cpu_size = cpu.expose(opd, seed=1).metadata["wfs_pyramid_fft_size_px"]
+    gpu_size = gpu.expose(cupy.asarray(opd), seed=1).metadata["wfs_pyramid_fft_size_px"]
+    assert cpu_size == gpu_size == 70
+
+
+@pytest.mark.gpu
+def test_pyramid_cuda_graph_survives_cupy_plan_cache_churn() -> None:
+    # CuPy's plan cache is global and evicts plans when other transforms run;
+    # the captured graph must not depend on it (2.1.0 replayed freed memory).
+    cupy = _cupy()
+    sensor = WavefrontSensor(_gpu_config("pyramid_minimal.toml"))
+    opd = cupy.asarray(np.random.default_rng(3).normal(0.0, 5.0e-8, sensor.config.input.shape))
+    first = sensor.photon_rate(opd).copy()
+    assert sensor.engine._graph._graph is not None
+    for length in range(20, 60):
+        cupy.fft.fft(cupy.ones((3, length), dtype=cupy.complex64), axis=-1)
+    cupy.fft.config.get_plan_cache().clear()
+    import gc
+
+    gc.collect()
+    again = sensor.photon_rate(opd)
+    cupy.cuda.Device().synchronize()
+    assert bool(cupy.array_equal(first, again))

{makewfs-2.1.0 → makewfs-2.1.1}/tests/test_numerics.py

@@ -7,11 +7,13 @@ import numpy as np
 import pytest
 
 from makewfs.backend import (
+    ArrayBackend,
     centered_fft2,
     centered_fft_intensity,
     centered_ifft2,
     complex_dtype,
     cpu_backend,
+    next_fast_length,
     real_dtype,
 )
 from makewfs.config import SourceConfig, SpiderConfig, TelescopeConfig
@@ -372,3 +374,31 @@ def test_magnitude_rate_follows_pogson_scaling() -> None:
     faint = replace(bright, magnitude=12.0)
     ratio = source_rate_per_s(bright, telescope) / source_rate_per_s(faint, telescope)
     assert np.isclose(ratio, 10.0**0.8, rtol=1e-12)
+
+
+def _seven_smooth(value: int) -> bool:
+    for factor in (2, 3, 5, 7):
+        while value % factor == 0:
+            value //= factor
+    return value == 1
+
+
+def test_fft_length_rule_is_minimal_seven_smooth_and_backend_independent() -> None:
+    gpu_named = ArrayBackend(np, name="cupy")  # only the name differs from the CPU backend
+    for value in range(1, 3001):
+        length = next_fast_length(value)
+        assert length >= value
+        assert _seven_smooth(length)
+        assert not any(_seven_smooth(candidate) for candidate in range(value, length))
+        assert cpu_backend().next_fast_length(value) == length
+        assert gpu_named.next_fast_length(value) == length
+    # Lengths SciPy used to pick on the CPU while CuPy picked another.
+    assert [next_fast_length(value) for value in (33, 66, 99)] == [35, 70, 100]
+    with pytest.raises(ValueError, match="positive"):
+        next_fast_length(0)
+
+
+def test_fft_length_rule_matches_cupy_next_fast_len() -> None:
+    cupy_fft = pytest.importorskip("cupyx.scipy.fft")
+    for value in range(1, 3001):
+        assert next_fast_length(value) == int(cupy_fft.next_fast_len(value))

{makewfs-2.1.0 → makewfs-2.1.1}/tests/test_pyramid.py

@@ -271,3 +271,21 @@ def test_unshifted_pruned_propagation_matches_centered_reference(
     np.testing.assert_allclose(
         result.photon_rate, expected, rtol=tolerance, atol=tolerance * expected.max()
     )
+
+
+@pytest.mark.parametrize(
+    ("path", "size"),
+    [
+        (CONFIG, 288),
+        (CONFIG.parents[2] / "benchmarks" / "configs" / "pyramid_40_float32.toml", 108),
+        (CONFIG.parents[2] / "benchmarks" / "configs" / "pyramid_60_mod8_float32.toml", 160),
+        (CONFIG.parents[2] / "benchmarks" / "configs" / "pyramid_80_mod32_float64.toml", 216),
+    ],
+)
+def test_pyramid_fft_size_is_recorded_and_unchanged_for_shipped_configs(
+    path: Path, size: int
+) -> None:
+    sensor = WavefrontSensor.from_toml(path)
+    assert sensor.engine.nfft == size
+    frame = sensor.expose(np.zeros(sensor.config.input.shape), seed=1)
+    assert frame.metadata["wfs_pyramid_fft_size_px"] == size