From ffbbd03b274132eaf7ceb8fa99251acbec16ddf0 Mon Sep 17 00:00:00 2001 From: Jacob Taylor Date: Tue, 6 Oct 2026 19:08:24 -1000 Subject: [PATCH 1/2] Weight wfs_input_opd_rms_m by the pupil and remove piston (2.0.0) wfs_input_opd_rms_m now follows aocore CONVENTIONS 4.1: the RMS of the input OPD in metres weighted by the intensity of the pupil the optics use, on the input grid, with the intensity-weighted mean removed. An analytic pupil is evaluated on input.shape with the configured supersampling; a custom mask is area-averaged from the engine's pupil grid. The 1.x quantity, the unweighted whole-grid RMS with piston kept, is still recorded as wfs_input_opd_rms_unweighted_m. Both are reduced on the device and cross in the existing single batched transfer. Release 2.0.0: this changes the meaning of an existing metadata key. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 11 ++ CHANGELOG.md | 43 ++++++ CITATION.cff | 2 +- docs/concepts.md | 29 ++++ docs/stability.md | 17 +++ examples/closed_loop_injection.py | 13 +- examples/showcase.py | 5 +- src/makewfs/__about__.py | 2 +- src/makewfs/api.py | 70 +++++++-- src/makewfs/provenance.py | 26 ++-- src/makewfs/sampling.py | 57 +++++++ src/makewfs/sensors/base.py | 6 + src/makewfs/sensors/pyramid.py | 1 + src/makewfs/sensors/shack_hartmann.py | 4 +- src/makewfs/wavefront.py | 77 ++++++++++ tests/test_gpu_backend.py | 37 +++++ tests/test_input_rms.py | 211 ++++++++++++++++++++++++++ tests/test_provenance.py | 30 ---- 18 files changed, 573 insertions(+), 68 deletions(-) create mode 100644 tests/test_input_rms.py diff --git a/AGENTS.md b/AGENTS.md index 0414f4a..a3d2b34 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -87,6 +87,14 @@ write a failing integration test/design note, use the conditional gates in - Intensities, not fields, are summed over incoherent wavelengths, modulation points, finite-source samples, and sodium slices. - Cropping reports lost flux; it does not renormalize it away. +- Wavefront metrics follow aocore CONVENTIONS 4.1. Frame metadata + `wfs_input_opd_rms_m` is the pupil-intensity-weighted, piston-removed RMS of + the input OPD on the input grid (weights from `WavefrontSensor._rms_weights`: + the analytic pupil evaluated on `input.shape`, or a custom mask area-averaged + from the engine's `configured_pupil`); `wfs_input_opd_rms_unweighted_m` is the + whole-grid quadratic mean with piston kept. Both are reduced on the device and + cross in the one `backend.scalars` batch. Any other RMS-like key must say its + variant in its name (`_unweighted`, `_tiptilt_removed`). - The intended top-level API is `load_config`, `WavefrontSensor`, and `simulate`. Keep other implementation objects out of `makewfs.__init__` unless an API review explicitly accepts them. @@ -155,6 +163,9 @@ Follow the target layout in `ROADMAP.md`: FFT workers and `ndimage` helpers, which `aocore.Backend` does not, and its centred FFTs keep the `fftshift` convention. `sampling.block_sum` wraps `aocore.block_sum` but keeps its own factor-two fast path. + `sampling.area_rebin` is exact-overlap area averaging between grids of the + same extent and any shape ratio, which `aocore.block_sum` (integer factors + only) does not cover. - `sensors/` contains deterministic ideal optical engines and no camera noise. `_shack_hartmann_cuda.py` is a private first-use-JIT execution plan for exact compatible CUDA geometries; `shack_hartmann.py` remains the readable physics diff --git a/CHANGELOG.md b/CHANGELOG.md index ab9f4d3..bbcb73d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,49 @@ All notable changes to `makewfs` are documented here. ## [Unreleased] +## [2.0.0] - 2026-10-07 + +### Breaking + +- **`wfs_input_opd_rms_m` is now the pupil-weighted, piston-removed RMS.** + It follows aocore CONVENTIONS.md 4.1: the RMS of the input OPD in metres, + weighted by the intensity of the pupil the optics use and with the + intensity-weighted mean (piston) removed. Before 2.0 it was the quadratic + mean over the whole input grid, with every pixel counted equally, pixels + outside the pupil included and piston kept. For the same wavefront the new + value is usually smaller; a pure piston now reports 0. The weights are the + configured pupil's intensity (the amplitude the engines use, squared) on the + input grid: an analytic pupil is evaluated on `input.shape` with the + configured `numerics.pupil_supersampling`, the same map + `WavefrontSensor.pupil_illumination()` returns; a custom mask, which exists + only on the engine's pupil grid, is area-averaged onto the input grid. + Phase input is converted to OPD first, and `expose_integrated` reports the + RMS of the mean OPD, as before. The value is still reduced on the device and + crosses to the host in the same single batched transfer as the captured + photon rate. A pupil with no transmission on the input grid is now rejected + when the sensor is built. +- **`makewfs.provenance.metadata`**, an internal helper, takes the two + reduced RMS values as required `opd_rms_m` and `opd_rms_unweighted_m` + arguments and no longer accepts `opd_m`. + +See "Migrating to 2.0" in the stability guide. + +### Added + +- **`wfs_input_opd_rms_unweighted_m`** frame metadata keeps the 1.x quantity, + the unweighted RMS over the whole input grid with piston included, so no + information is lost. Read it wherever the old number is still wanted. +- Both sensor engines expose `configured_pupil`, the pupil amplitude on their + own pupil grid, and `makewfs.sampling.area_rebin` area-averages a map onto + another grid of the same extent, for any shape ratio. + +### Changed + +- The `closed_loop_injection.py` and `showcase.py` examples take the residual + RMS they plot from each frame's `wfs_input_opd_rms_m` instead of a + whole-grid `np.std`, so the numbers they print are pupil-weighted. The + checked-in `makewfs_showcase.webp` was not regenerated. + ## [1.2.0] - 2026-10-07 - **Fixed: phase input was reported as metres.** With diff --git a/CITATION.cff b/CITATION.cff index d421bb2..11fdd94 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -5,7 +5,7 @@ type: software authors: - family-names: Taylor given-names: Jacob -version: 1.2.0 +version: 2.0.0 date-released: 2026-10-07 license: MIT repository-code: "https://github.com/jacotay7/makewfs" diff --git a/docs/concepts.md b/docs/concepts.md index 423f6b3..12ca006 100644 --- a/docs/concepts.md +++ b/docs/concepts.md @@ -31,6 +31,35 @@ in single precision. When an input grid does not divide into the configured lenslets, OPD is resampled on physical coordinates. Wrapped phase is never interpolated. +## Input wavefront RMS in frame metadata + +Every frame records two RMS values of the input wavefront, both in OPD metres +(phase input is converted at `input.reference_wavelength_m` first). Both +describe the input array as given, without `input.static_opd_path`; for +`expose_integrated` they describe the mean of the temporal samples. + +| Key | Definition | +| --- | --- | +| `wfs_input_opd_rms_m` | Pupil-weighted, piston-removed RMS, the `rms` of aocore CONVENTIONS 4.1: `sqrt(sum a^2 (opd - _a)^2 / sum a^2)`, where `_a` is the intensity-weighted mean. | +| `wfs_input_opd_rms_unweighted_m` | `sqrt(mean(opd^2))` over every pixel of the input grid, with piston included and pixels outside the pupil counted. This was `wfs_input_opd_rms_m` before 2.0. | + +The weight `a^2` is the intensity of the pupil the optics use, on the input +grid. The engines use the configured pupil as a field amplitude, so the +intensity is its square. An analytic pupil (diameter, obscuration, spiders, +segment gaps, rotation) is evaluated on `input.shape` with the same +`numerics.pupil_supersampling`; this is the map +`WavefrontSensor.pupil_illumination()` returns, computed in float64. A +`telescope.custom_mask_path` mask exists only on the engine's own pupil grid +(`lenslets_across_pupil` times the pupil samples per lenslet for a +Shack-Hartmann, `pixels_across_pupil` for a pyramid), so its intensity is +area-averaged onto the input grid: each input +pixel takes the mean over the mask cells it overlaps, which is exact when the +two shapes match. A piston, or OPD outside the pupil, therefore changes only +the unweighted value. + +Both values are reduced on the selected device and cross to the host together +with the captured photon rate in one transfer. + ## Closed-loop use The package intentionally stops at the detector image. A downstream controller diff --git a/docs/stability.md b/docs/stability.md index 68afb9e..236bb84 100644 --- a/docs/stability.md +++ b/docs/stability.md @@ -25,3 +25,20 @@ The physical per-frame contract is a wavefront array plus configuration (and an optional detector seed). `out` controls storage lifetime only and cannot change the simulated result. Atmosphere, reconstruction, controllers, and detector physics remain outside the package boundary. + +## Migrating to 2.0 + +makewfs 2.0 changes the meaning of one frame-metadata key; the API and the +configuration schema are unchanged. + +- `wfs_input_opd_rms_m` is now the pupil-weighted, piston-removed RMS of + aocore CONVENTIONS 4.1 (see + [Concepts](concepts.md#input-wavefront-rms-in-frame-metadata)). For the same + wavefront it is usually smaller than before: it no longer includes piston or + any OPD outside the pupil. +- The old quantity, the unweighted RMS over the whole input grid with piston + included, is still recorded as `wfs_input_opd_rms_unweighted_m`. Code that + needs the 1.x number reads that key instead. +- `makewfs.provenance.metadata`, an internal helper, now takes both RMS values + as required `opd_rms_m` and `opd_rms_unweighted_m` arguments and no longer + accepts `opd_m`. diff --git a/examples/closed_loop_injection.py b/examples/closed_loop_injection.py index e998567..876b127 100644 --- a/examples/closed_loop_injection.py +++ b/examples/closed_loop_injection.py @@ -5,7 +5,9 @@ next *residual* wavefront back to ``WavefrontSensor.expose``. The residual is a low-order aberration (defocus + astigmatism) so the Shack-Hartmann spots shift differently across the pupil and the relaxation toward the flat reference is -visible, and a convergence panel tracks the residual RMS. +visible, and a convergence panel tracks the residual RMS that each frame +records in its ``wfs_input_opd_rms_m`` metadata: pupil-weighted with piston +removed, as aocore CONVENTIONS 4.1 defines ``rms``. """ from __future__ import annotations @@ -19,7 +21,7 @@ def _low_order_residual(shape: tuple[int, int], amplitude_m: float) -> np.ndarray: - """Return a defocus + astigmatism OPD normalized to ``amplitude_m`` RMS.""" + """Return a defocus + astigmatism OPD with ``amplitude_m`` RMS over the grid.""" yy, xx = np.mgrid[: shape[0], : shape[1]] x = (xx - (shape[1] - 1) / 2) / (shape[1] / 2) y = (yy - (shape[0] - 1) / 2) / (shape[0] / 2) @@ -56,8 +58,9 @@ def main() -> None: residual_rms_nm: list[float] = [] for step in range(args.steps): residual = _low_order_residual(shape, amplitude_m=200e-9) * (0.6**step) - residual_rms_nm.append(float(np.std(residual) * 1e9)) - frames.append(np.asarray(sensor.expose(residual, seed=0)).astype(np.float64)) + frame = sensor.expose(residual, seed=0) + residual_rms_nm.append(frame.metadata["wfs_input_opd_rms_m"] * 1e9) + frames.append(np.asarray(frame).astype(np.float64)) columns = min(4, args.steps) shown = sorted({0, 1, args.steps // 2, args.steps - 1})[:columns] @@ -91,7 +94,7 @@ def main() -> None: convergence.plot(range(args.steps), residual_rms_nm, "o-") convergence.set_title("external loop convergence") convergence.set_xlabel("loop step") - convergence.set_ylabel("residual OPD RMS (nm)") + convergence.set_ylabel("residual OPD RMS over the pupil (nm)") convergence.grid(True, alpha=0.3) frame_rms = [float(np.sqrt(np.mean(diff**2))) for diff in diffs] diff --git a/examples/showcase.py b/examples/showcase.py index 729d800..d94161b 100644 --- a/examples/showcase.py +++ b/examples/showcase.py @@ -263,7 +263,6 @@ def benchmark_panel(panel: Panel, sync: Callable[[], None]) -> float: def collect_frames(panel: Panel) -> np.ndarray: """``(N_FRAMES, *output_shape)`` detector ADU on the host, plus the OPD rms.""" import getframes - import pyturb sensor = panel.build_sensor() atmosphere = panel.build_atmosphere() @@ -271,8 +270,10 @@ def collect_frames(panel: Panel) -> np.ndarray: rms = [] for index in range(N_FRAMES): residual = panel.residual(atmosphere, index) - rms.append(float(np.std(pyturb.to_numpy(residual))) * 1e9) frame = sensor.expose(residual, seed=index) + # Pupil-weighted, piston-removed (aocore CONVENTIONS 4.1), reduced on + # the sensor's device. + rms.append(frame.metadata["wfs_input_opd_rms_m"] * 1e9) out[index] = getframes.to_numpy(frame.data).astype(np.float32) panel.badge_extra = f"residual {np.mean(rms):.0f} nm rms" return np.clip(out - panel.bias_adu, 0.0, None) diff --git a/src/makewfs/__about__.py b/src/makewfs/__about__.py index fc31d78..eab05f7 100644 --- a/src/makewfs/__about__.py +++ b/src/makewfs/__about__.py @@ -1,3 +1,3 @@ """Package version metadata.""" -__version__ = "1.2.0" +__version__ = "2.0.0" diff --git a/src/makewfs/api.py b/src/makewfs/api.py index 9aa05a9..fa334a7 100644 --- a/src/makewfs/api.py +++ b/src/makewfs/api.py @@ -16,10 +16,11 @@ from .detector import DetectorAdapter from .provenance import metadata as build_metadata from .pupil import make_pupil +from .sampling import area_rebin from .sensors.base import OpticalResult, SensorEngine from .sensors.pyramid import PyramidEngine from .sensors.shack_hartmann import ShackHartmannEngine -from .wavefront import iter_phase_samples +from .wavefront import grid_rms, iter_phase_samples, pupil_rms, pupil_weights class WavefrontSensor: @@ -53,10 +54,12 @@ def __init__(self, config: WFSConfig, *, _backend: ArrayBackend | None = None) - launched_rate=0.0, captured_rate=0.0, opd_rms_m=0.0, + opd_rms_unweighted_m=0.0, seed=None, source_states=self.engine.source_states, file_digests=self.engine.file_digests, ) + self._rms_weights = pupil_weights(self._input_pupil_intensity(), backend=self.backend) @classmethod def from_toml(cls, path: str | Path) -> WavefrontSensor: @@ -66,9 +69,42 @@ def from_toml(cls, path: str | Path) -> WavefrontSensor: def _render(self, wavefront: ArrayLike) -> OpticalResult: return self.engine.render(cast(NDArray[np.float64], wavefront)) - def _opd_rms(self, opd: Any) -> Any: - """Reduce OPD RMS on-device before the batched metadata crossing.""" - return self.backend.sqrt(self.backend.mean(opd**2)) + def _input_pupil_intensity(self) -> Any: + """Return the pupil intensity the optics use, on the input grid. + + An analytic pupil is evaluated on ``input.shape`` from the same + telescope model and ``numerics.pupil_supersampling`` the engine uses on + its own grid, as :meth:`pupil_illumination` does; both grids span + ``input.grid_extent_m``. A custom mask exists only on the engine's + configured pupil grid, so its intensity is area-averaged from there + onto the input grid, which is exact when the two shapes match. The + engines use the pupil as a field amplitude, so the intensity is its + square. + """ + if self.config.telescope.custom_mask_path is None: + amplitude = make_pupil( + self.config.telescope, + self.config.input.shape, + self.config.input.grid_extent_m, + supersampling=self.config.numerics.pupil_supersampling, + backend=self.backend, + dtype=np.float64, + ) + return amplitude * amplitude + amplitude = self.backend.asarray(self.engine.configured_pupil, dtype=np.float64) + return area_rebin(amplitude * amplitude, self.config.input.shape, backend=self.backend) + + def _opd_rms(self, opd: Any) -> tuple[Any, Any]: + """Reduce both input-OPD RMS values on the device. + + Returns the pupil-weighted, piston-removed RMS and the unweighted + whole-grid RMS as device scalars, for the one batched metadata + crossing in ``backend.scalars``. + """ + return ( + pupil_rms(opd, self._rms_weights, backend=self.backend), + grid_rms(opd, backend=self.backend), + ) def _frame_metadata( self, @@ -76,6 +112,7 @@ def _frame_metadata( launched_rate: float, captured_rate: float, opd_rms_m: float, + opd_rms_unweighted_m: float, seed: int | None, ) -> dict[str, Any]: """Copy cached static provenance and fill the per-frame values.""" @@ -85,6 +122,7 @@ def _frame_metadata( "wfs_launched_photons_s": float(launched_rate), "wfs_captured_photons_s": float(captured_rate), "wfs_input_opd_rms_m": float(opd_rms_m), + "wfs_input_opd_rms_unweighted_m": float(opd_rms_unweighted_m), "wfs_seed": seed if seed is not None else "internal", } ) @@ -182,18 +220,26 @@ def expose( seed: int | None = None, out: Any | None = None, ) -> Any: - """Render one wavefront into optional caller-owned detector storage.""" + """Render one wavefront into optional caller-owned detector storage. + + Besides provenance and timings, the frame metadata records the input + wavefront's ``wfs_input_opd_rms_m``, its RMS in OPD metres weighted by + the pupil intensity with piston removed (aocore CONVENTIONS 4.1), and + ``wfs_input_opd_rms_unweighted_m``, its RMS over the whole input grid + with piston kept. + """ total_start = perf_counter() optical_start = total_start result = self._render(wavefront) - captured_rate, opd_rms = self.backend.scalars( - result.captured_rate_per_s, self._opd_rms(result.opd_m) + captured_rate, opd_rms, opd_rms_unweighted = self.backend.scalars( + result.captured_rate_per_s, *self._opd_rms(result.opd_m) ) optical_elapsed = perf_counter() - optical_start frame_metadata = self._frame_metadata( launched_rate=result.launched_rate_per_s, captured_rate=captured_rate, opd_rms_m=opd_rms, + opd_rms_unweighted_m=opd_rms_unweighted, seed=seed, ) detector_start = perf_counter() @@ -244,13 +290,14 @@ def expose_integrated( if not samples: raise ValueError("phase_samples must contain at least one sample") result = batched(samples) - captured_rate, opd_rms = self.backend.scalars( - result.captured_rate_per_s, self._opd_rms(result.opd_m) + captured_rate, opd_rms, opd_rms_unweighted = self.backend.scalars( + result.captured_rate_per_s, *self._opd_rms(result.opd_m) ) frame_metadata = self._frame_metadata( launched_rate=result.launched_rate_per_s, captured_rate=captured_rate, opd_rms_m=opd_rms, + opd_rms_unweighted_m=opd_rms_unweighted, seed=seed, ) frame_metadata["wfs_temporal_samples"] = len(samples) @@ -305,13 +352,14 @@ def expose_integrated( None if spectral_rate_sum is None else spectral_rate_sum / sample_count ) average_opd = opd_sum / sample_count - captured_rate, opd_rms = self.backend.scalars( - captured / sample_count, self._opd_rms(average_opd) + captured_rate, opd_rms, opd_rms_unweighted = self.backend.scalars( + captured / sample_count, *self._opd_rms(average_opd) ) frame_metadata = self._frame_metadata( launched_rate=launched, captured_rate=captured_rate, opd_rms_m=opd_rms, + opd_rms_unweighted_m=opd_rms_unweighted, seed=seed, ) frame_metadata["wfs_temporal_samples"] = sample_count diff --git a/src/makewfs/provenance.py b/src/makewfs/provenance.py index 28d8eab..1ff462e 100644 --- a/src/makewfs/provenance.py +++ b/src/makewfs/provenance.py @@ -7,9 +7,6 @@ from pathlib import Path from typing import Any -import numpy as np -from numpy.typing import NDArray - from .config import WFSConfig from .source import SourceState, iter_source_states @@ -44,28 +41,24 @@ def referenced_file_digests(config: WFSConfig) -> dict[str, str]: } -def _opd_rms(opd_m: NDArray[Any] | None, opd_rms_m: float | None) -> float: - """Resolve a host OPD RMS, accepting a backend-reduced value.""" - if opd_rms_m is not None: - return float(opd_rms_m) - if opd_m is None: - raise ValueError("metadata requires opd_m or opd_rms_m") - return float(np.sqrt(np.mean(np.asarray(opd_m) ** 2))) - - def metadata( config: WFSConfig, *, sensor_kind: str, launched_rate: float, captured_rate: float, - opd_m: NDArray[Any] | None = None, - opd_rms_m: float | None = None, + opd_rms_m: float, + opd_rms_unweighted_m: float, seed: int | None, source_states: tuple[SourceState, ...] | None = None, file_digests: dict[str, str] | None = None, ) -> dict[str, Any]: - """Build serializable metadata for an ideal or detector frame.""" + """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). + """ states = iter_source_states(config) if source_states is None else source_states result: dict[str, Any] = { "frame_type": "wfs", @@ -74,7 +67,8 @@ def metadata( "wfs_wavelength_m": config.sensor.wavelength_m, "wfs_launched_photons_s": float(launched_rate), "wfs_captured_photons_s": float(captured_rate), - "wfs_input_opd_rms_m": _opd_rms(opd_m, opd_rms_m), + "wfs_input_opd_rms_m": float(opd_rms_m), + "wfs_input_opd_rms_unweighted_m": float(opd_rms_unweighted_m), "wfs_seed": seed if seed is not None else "internal", "wfs_source_kind": config.source.kind, "wfs_source_state_count": len(states), diff --git a/src/makewfs/sampling.py b/src/makewfs/sampling.py index 2245516..a24610e 100644 --- a/src/makewfs/sampling.py +++ b/src/makewfs/sampling.py @@ -172,6 +172,62 @@ def block_sum( return cast(NDArray[Any], _aocore_block_sum(array, factor)) +def _area_overlap(target: int, source: int, *, backend: ArrayBackend, dtype: Any) -> NDArray[Any]: + """Return the ``(target, source)`` fractions of each target cell's length. + + Both grids tile the same interval. Edges are compared in integer units of + ``1 / (target * source)``, so the overlaps are exact and every row sums to + one. + """ + rows = backend.arange(target)[:, None] + columns = backend.arange(source)[None, :] + upper = backend.where( + (rows + 1) * source < (columns + 1) * target, + (rows + 1) * source, + (columns + 1) * target, + ) + lower = backend.where(rows * source > columns * target, rows * source, columns * target) + overlap = backend.where(upper > lower, upper - lower, 0) + return cast(NDArray[Any], backend.asarray(overlap, dtype=dtype) / source) + + +def area_rebin( + array: NDArray[Any], + shape: tuple[int, int], + *, + backend: ArrayBackend | None = None, +) -> NDArray[Any]: + """Area-average a map onto another grid spanning the same extent. + + Each target pixel takes the mean of the source map over its own area, + weighting every source pixel by the exact fraction it overlaps, so any + integer or non-integer ratio, up or down, is handled without + interpolation. A uniform map stays uniform and the area integral of the map + is preserved. Equal shapes return the map unchanged. + + Parameters + ---------- + array + Two-dimensional ``(y, x)`` map on the source grid. + shape + Target ``(height, width)``. + backend + Array backend holding ``array``. + + Returns + ------- + numpy.ndarray or cupy.ndarray + The area-averaged map on the target grid, in ``array``'s dtype. + """ + resolved = backend or cpu_backend() + height, width = array.shape + if (height, width) == tuple(shape): + return array + rows = _area_overlap(shape[0], height, backend=resolved, dtype=array.dtype) + columns = _area_overlap(shape[1], width, backend=resolved, dtype=array.dtype) + return cast(NDArray[Any], resolved.matmul(resolved.matmul(rows, array), columns.T)) + + @dataclass class _SpotPropagationPlan: """Cached backend-resident geometry for repeated spot propagation. @@ -439,6 +495,7 @@ def spot_intensity( __all__ = [ + "area_rebin", "block_sum", "crop_center", "lenslet_field_upsampling", diff --git a/src/makewfs/sensors/base.py b/src/makewfs/sensors/base.py index 7c42ad6..3e107a5 100644 --- a/src/makewfs/sensors/base.py +++ b/src/makewfs/sensors/base.py @@ -33,6 +33,12 @@ class SensorEngine: source_states: tuple[SourceState, ...] file_digests: dict[str, str] backend: ArrayBackend + configured_pupil: Any + """Pupil amplitude on the engine's configured pupil grid. + + The grid spans ``input.grid_extent_m`` on each axis. A custom mask is + defined on this grid, so it is the only place that mask exists. + """ @staticmethod def resolve_backend(backend: ArrayBackend | None) -> ArrayBackend: diff --git a/src/makewfs/sensors/pyramid.py b/src/makewfs/sensors/pyramid.py index d20d2b7..3d8569d 100644 --- a/src/makewfs/sensors/pyramid.py +++ b/src/makewfs/sensors/pyramid.py @@ -62,6 +62,7 @@ def __init__(self, config: WFSConfig, *, backend: ArrayBackend | None = None) -> backend=self.backend, dtype=self._real_dtype, ) + self.configured_pupil = self.pupil self.wavefront = WavefrontInput( config, load_static_opd(config), diff --git a/src/makewfs/sensors/shack_hartmann.py b/src/makewfs/sensors/shack_hartmann.py index e17b129..6246d09 100644 --- a/src/makewfs/sensors/shack_hartmann.py +++ b/src/makewfs/sensors/shack_hartmann.py @@ -80,7 +80,7 @@ def __init__(self, config: WFSConfig, *, backend: ArrayBackend | None = None) -> backend=self.backend, dtype=self._real_dtype, ) - self._configured_pupil = pupil + self.configured_pupil = pupil if self.field_upsampling > 1: # Hold each configured pupil cell's area-weighted transmission # constant over its refined sub-cells, so the illuminated area and @@ -182,7 +182,7 @@ def __init__(self, config: WFSConfig, *, backend: ArrayBackend | None = None) -> # Spiders, segment gaps and custom masks block light the analytic # annulus of the magnitude normalization would otherwise count. self.clear_aperture_fraction = clear_aperture_fraction( - self._configured_pupil, + self.configured_pupil, config.telescope, self.pupil_shape, config.input.grid_extent_m, diff --git a/src/makewfs/wavefront.py b/src/makewfs/wavefront.py index 0132623..0b2bfbc 100644 --- a/src/makewfs/wavefront.py +++ b/src/makewfs/wavefront.py @@ -194,6 +194,80 @@ def _resampling_coordinates( return backend.stack((yy, xx), axis=0) +def pupil_weights(intensity: NDArray[Any], *, backend: ArrayBackend | None = None) -> Any: + """Normalize a pupil intensity map into RMS weights that sum to one. + + Parameters + ---------- + intensity + Non-negative pupil intensity (amplitude squared) on the OPD grid. + backend + Array backend holding ``intensity``. + + Returns + ------- + numpy.ndarray or cupy.ndarray + Float64 weights on the same backend. + + Raises + ------ + ValueError + If the pupil transmits nothing on this grid. + """ + resolved = backend or cpu_backend() + weights = resolved.asarray(intensity, dtype=np.float64) + total = resolved.sum(weights) + if not resolved.scalar(total) > 0.0: + raise ValueError("pupil has no illuminated pixels on the input grid") + return weights / total + + +def pupil_rms(opd: Any, weights: Any, *, backend: ArrayBackend | None = None) -> Any: + """Pupil-weighted, piston-removed OPD RMS, left on the device. + + Implements ``rms`` of aocore CONVENTIONS 4.1, + ``sqrt(sum a^2 (opd - _a)^2 / sum a^2)``, where ``weights`` is the + intensity ``a^2`` already normalized by :func:`pupil_weights`. The piston is + subtracted before squaring rather than taken from `` - ^2``, so + a large piston does not cancel away the residual's precision and a + piston-only input gives zero to rounding. + + ``aocore.rms`` defines the same quantity but reduces on the host with + NumPy; this stays on the selected backend so frame metadata needs no extra + device-to-host copy. The tests check it with ``aocore.conformance.check_rms``. + + Parameters + ---------- + opd + OPD in metres on the same grid as ``weights``. + weights + Normalized pupil intensity weights. + backend + Array backend holding both arrays. + + Returns + ------- + numpy.ndarray or cupy.ndarray + A zero-dimensional float64 array; no host synchronization happens here. + """ + resolved = backend or cpu_backend() + values = resolved.asarray(opd, dtype=np.float64) + residual = values - resolved.sum(weights * values) + return resolved.sqrt(resolved.sum(weights * residual * residual)) + + +def grid_rms(opd: Any, *, backend: ArrayBackend | None = None) -> Any: + """Unweighted OPD RMS over the whole grid, piston included, on the device. + + This is ``rms_unweighted`` in the sense of aocore CONVENTIONS 4.1: every + grid pixel counts equally, including pixels outside the pupil, and the mean + is not removed. + """ + resolved = backend or cpu_backend() + values = resolved.asarray(opd, dtype=np.float64) + return resolved.sqrt(resolved.mean(values * values)) + + def iter_phase_samples( value: ArrayLike | Iterable[ArrayLike], shape: tuple[int, int], @@ -221,7 +295,10 @@ def iter_phase_samples( __all__ = [ "WavefrontInput", "_coordinates", + "grid_rms", "iter_phase_samples", "load_static_opd", + "pupil_rms", + "pupil_weights", "resample_opd", ] diff --git a/tests/test_gpu_backend.py b/tests/test_gpu_backend.py index 9d6c536..3cc73a7 100644 --- a/tests/test_gpu_backend.py +++ b/tests/test_gpu_backend.py @@ -341,3 +341,40 @@ def test_compiled_sh_records_feature_fallback_once() -> None: assert not engine._compiled_executors assert engine._compiled_executor_rejections == {1: "continuous optical blur"} assert bool(cupy.array_equal(first.photon_rate, second.photon_rate)) + + +@pytest.mark.gpu +@pytest.mark.parametrize("name", ["shack_hartmann_minimal.toml", "pyramid_minimal.toml"]) +@pytest.mark.parametrize("custom_mask", [False, True]) +def test_input_rms_metadata_matches_cpu(name: str, custom_mask: bool, tmp_path: Path) -> None: + # Both RMS keys are reduced on the device and cross with the captured rate + # in one batch; the values must match the CPU reduction. + cupy = _cupy() + config = _config(name) + if custom_mask: + shape = tuple(WavefrontSensor(config).engine.configured_pupil.shape) + mask = np.zeros(shape) + mask[shape[0] // 8 : shape[0] // 2, shape[1] // 4 : -shape[1] // 8] = 1.0 + mask[shape[0] // 2 : shape[0] // 2 + 3, shape[1] // 4 : -shape[1] // 8] = 0.4 + path = tmp_path / "mask.npy" + np.save(path, mask) + config = replace(config, telescope=replace(config.telescope, custom_mask_path=str(path))) + gpu_config = replace(config, numerics=replace(config.numerics, device="gpu")) + cpu_sensor = WavefrontSensor(config) + gpu_sensor = WavefrontSensor(gpu_config) + rng = np.random.default_rng(11) + height, width = config.input.shape + x = np.broadcast_to(np.arange(width) - (width - 1) / 2, (height, width)) + opd = 2.0e-9 * x + rng.normal(0.0, 2.0e-8, (height, width)) + 1.5e-6 + + cpu = cpu_sensor.expose(opd, seed=3).metadata + gpu = gpu_sensor.expose(cupy.asarray(opd), seed=3).metadata + for key in ("wfs_input_opd_rms_m", "wfs_input_opd_rms_unweighted_m"): + assert isinstance(gpu[key], float) + assert gpu[key] == pytest.approx(cpu[key], rel=1e-12) + + samples = np.stack([opd, 0.5 * opd - 1.0e-6]) + cpu_integrated = cpu_sensor.expose_integrated(samples, seed=3).metadata + gpu_integrated = gpu_sensor.expose_integrated(cupy.asarray(samples), seed=3).metadata + for key in ("wfs_input_opd_rms_m", "wfs_input_opd_rms_unweighted_m"): + assert gpu_integrated[key] == pytest.approx(cpu_integrated[key], rel=1e-12) diff --git a/tests/test_input_rms.py b/tests/test_input_rms.py new file mode 100644 index 0000000..fde893a --- /dev/null +++ b/tests/test_input_rms.py @@ -0,0 +1,211 @@ +"""Input-OPD RMS frame metadata (aocore CONVENTIONS 4.1).""" + +from __future__ import annotations + +import math +from dataclasses import replace +from pathlib import Path + +import numpy as np +import pytest +from aocore import rms as aocore_rms +from aocore.conformance import check_rms + +from makewfs import WavefrontSensor, load_config +from makewfs.pupil import make_pupil +from makewfs.sampling import area_rebin +from makewfs.wavefront import grid_rms, pupil_rms, pupil_weights + +CONFIGS = Path(__file__).parents[1] / "examples" / "configs" +SENSORS = ["shack_hartmann_minimal.toml", "pyramid_minimal.toml"] + + +def _sensor(name: str) -> WavefrontSensor: + return WavefrontSensor.from_toml(CONFIGS / name) + + +def _x_m(sensor: WavefrontSensor) -> np.ndarray: + """Centred physical x coordinate of every input pixel, in metres.""" + height, width = sensor.config.input.shape + extent = sensor.config.input.grid_extent_m + x = (np.arange(width) - (width - 1) / 2) * extent / width + return np.broadcast_to(x, (height, width)).copy() + + +def _pupil(sensor: WavefrontSensor) -> np.ndarray: + """The configured pupil amplitude on the input grid, in float64.""" + config = sensor.config + return make_pupil( + config.telescope, + config.input.shape, + config.input.grid_extent_m, + supersampling=config.numerics.pupil_supersampling, + ) + + +def test_weighted_rms_helper_passes_the_aocore_rms_check() -> None: + def makewfs_rms(opd: np.ndarray, amplitude: np.ndarray) -> float: + return float(pupil_rms(opd, pupil_weights(amplitude**2))) + + check_rms(makewfs_rms, rel_tol=1e-12) + + +@pytest.mark.parametrize("name", SENSORS) +def test_piston_only_input_has_zero_rms(name: str) -> None: + sensor = _sensor(name) + piston = np.full(sensor.config.input.shape, 2.5e-6) + metadata = sensor.expose(piston, seed=1).metadata + assert metadata["wfs_input_opd_rms_m"] == pytest.approx(0.0, abs=1e-20) + assert metadata["wfs_input_opd_rms_unweighted_m"] == pytest.approx(2.5e-6, rel=1e-12) + + +@pytest.mark.parametrize("name", SENSORS) +def test_tilt_matches_the_analytic_pupil_weighted_rms(name: str) -> None: + # On an annulus of outer radius R and obscuration ratio e, = 0 and + # = R^2 (1 + e^2) / 4, so a slope s gives s R sqrt(1 + e^2) / 2. A + # piston is added to show that it is removed. + sensor = _sensor(name) + telescope = sensor.config.telescope + assert not telescope.spiders and telescope.custom_mask_path is None + slope = 1.0e-7 + opd = slope * _x_m(sensor) + 3.0e-6 + radius = telescope.pupil_diameter_m / 2 + analytic = slope * radius * math.sqrt(1 + telescope.central_obscuration_ratio**2) / 2 + + metadata = sensor.expose(opd, seed=1).metadata + # The sampled 128-pixel pupil reproduces the continuous annulus to 0.1 %. + assert metadata["wfs_input_opd_rms_m"] == pytest.approx(analytic, rel=2e-3) + # It is exactly CONVENTIONS 4.1 on the pupil the sensor uses. + assert metadata["wfs_input_opd_rms_m"] == pytest.approx( + aocore_rms(opd, _pupil(sensor)), rel=1e-12 + ) + assert metadata["wfs_input_opd_rms_unweighted_m"] == pytest.approx( + math.sqrt(float(np.mean(opd**2))), rel=1e-12 + ) + + +@pytest.mark.parametrize("name", SENSORS) +def test_only_the_unweighted_rms_sees_opd_outside_the_pupil(name: str) -> None: + sensor = _sensor(name) + inside = 4.0e-8 * _x_m(sensor) + outside = _pupil(sensor) == 0.0 + assert outside.any() + rng = np.random.default_rng(3) + polluted = inside.copy() + polluted[outside] = rng.normal(0.0, 1.0e-5, int(outside.sum())) + + clean = sensor.expose(inside, seed=1).metadata + dirty = sensor.expose(polluted, seed=1).metadata + assert dirty["wfs_input_opd_rms_m"] == clean["wfs_input_opd_rms_m"] + assert dirty["wfs_input_opd_rms_unweighted_m"] > 10 * clean["wfs_input_opd_rms_unweighted_m"] + + +@pytest.mark.parametrize("name", SENSORS) +def test_integrated_exposure_reports_the_rms_of_the_mean_opd(name: str) -> None: + sensor = _sensor(name) + tilt = 1.0e-7 * _x_m(sensor) + samples = np.stack([tilt + 1.0e-6, 3.0 * tilt - 1.0e-6]) + metadata = sensor.expose_integrated(samples, seed=1).metadata + mean_opd = samples.mean(axis=0) + assert metadata["wfs_input_opd_rms_m"] == pytest.approx( + aocore_rms(mean_opd, _pupil(sensor)), rel=1e-12 + ) + assert metadata["wfs_input_opd_rms_unweighted_m"] == pytest.approx( + math.sqrt(float(np.mean(mean_opd**2))), rel=1e-12 + ) + + +@pytest.mark.parametrize("name", SENSORS) +def test_custom_mask_weights_are_area_averaged_onto_the_input_grid( + name: str, tmp_path: Path +) -> None: + # A custom mask exists only on the engine's pupil grid, so each input pixel + # takes the mean intensity of the mask cells it covers. The grid is 128 x 128 + # for the Shack-Hartmann, the same as the input, and 64 x 64 for the + # pyramid, half of it. + config = load_config(CONFIGS / name) + probe = WavefrontSensor(config) + engine_shape = tuple(probe.engine.configured_pupil.shape) + mask = np.zeros(engine_shape) + mask[8:40, 16:56] = 1.0 + mask[40:44, 16:56] = 0.5 + path = tmp_path / "mask.npy" + np.save(path, mask) + sensor = WavefrontSensor( + replace(config, telescope=replace(config.telescope, custom_mask_path=str(path))) + ) + factor = config.input.shape[0] // engine_shape[0] + assert factor * engine_shape[0] == config.input.shape[0] + amplitude_on_input = np.sqrt(np.kron(mask**2, np.ones((factor, factor)))) + + rng = np.random.default_rng(5) + opd = rng.normal(0.0, 5.0e-8, config.input.shape) + 1.0e-6 + metadata = sensor.expose(opd, seed=1).metadata + assert metadata["wfs_input_opd_rms_m"] == pytest.approx( + aocore_rms(opd, amplitude_on_input), rel=1e-12 + ) + + +@pytest.mark.parametrize("name", SENSORS) +def test_phase_input_rms_is_reported_in_opd_metres(name: str) -> None: + # Phase input is converted at the reference wavelength before anything + # downstream sees it, including both input-RMS keys and OpticalResult. + opd_sensor = _sensor(name) + wavelength = 1.0e-6 + phase_input = replace( + opd_sensor.config.input, + quantity="phase", + unit="rad", + reference_wavelength_m=wavelength, + ) + phase_sensor = WavefrontSensor(replace(opd_sensor.config, input=phase_input)) + opd = 3.0e-8 * _x_m(opd_sensor) / opd_sensor.config.input.grid_extent_m + 1.0e-8 + phase = 2.0 * np.pi * opd / wavelength + + weighted = aocore_rms(opd, _pupil(opd_sensor)) + unweighted = math.sqrt(float(np.mean(opd**2))) + for sensor, wavefront in ((opd_sensor, opd), (phase_sensor, phase)): + metadata = sensor.expose(wavefront, seed=1).metadata + assert metadata["wfs_input_opd_rms_m"] == pytest.approx(weighted, rel=1e-12) + assert metadata["wfs_input_opd_rms_unweighted_m"] == pytest.approx(unweighted, rel=1e-12) + np.testing.assert_allclose(phase_sensor._render(phase).opd_m, opd, rtol=1e-12) + np.testing.assert_allclose( + phase_sensor.photon_rate(phase), opd_sensor.photon_rate(opd), rtol=1e-6, atol=0.0 + ) + + +def test_empty_pupil_on_the_input_grid_is_rejected() -> None: + with pytest.raises(ValueError, match="no illuminated pixels"): + pupil_weights(np.zeros((4, 4))) + + +def test_grid_rms_keeps_piston_and_counts_every_pixel() -> None: + opd = np.array([[1.0, 1.0], [1.0, 5.0]]) + assert float(grid_rms(opd)) == pytest.approx(math.sqrt(28.0 / 4.0)) + + +@pytest.mark.parametrize( + ("source", "target"), + [((64, 64), (128, 128)), ((128, 128), (64, 64)), ((40, 48), (64, 30)), ((7, 7), (7, 7))], +) +def test_area_rebin_preserves_uniform_maps_and_the_area_integral( + source: tuple[int, int], target: tuple[int, int] +) -> None: + rng = np.random.default_rng(7) + array = rng.uniform(0.0, 1.0, source) + rebinned = area_rebin(array, target) + assert rebinned.shape == target + source_area = 1.0 / (source[0] * source[1]) + target_area = 1.0 / (target[0] * target[1]) + assert float(rebinned.sum()) * target_area == pytest.approx( + float(array.sum()) * source_area, rel=1e-12 + ) + np.testing.assert_allclose(area_rebin(np.ones(source), target), 1.0, rtol=1e-12) + + +def test_area_rebin_downsampling_is_a_block_mean() -> None: + array = np.arange(36.0).reshape(6, 6) + expected = array.reshape(3, 2, 3, 2).mean(axis=(1, 3)) + np.testing.assert_allclose(area_rebin(array, (3, 3)), expected, rtol=1e-12) + # A non-integer ratio splits a source cell between two target cells. + np.testing.assert_allclose(area_rebin(np.array([[0.0, 3.0, 6.0]]), (1, 2)), [[1.0, 5.0]]) diff --git a/tests/test_provenance.py b/tests/test_provenance.py index b8129b9..ddbb3ca 100644 --- a/tests/test_provenance.py +++ b/tests/test_provenance.py @@ -4,7 +4,6 @@ from pathlib import Path import numpy as np -import pytest from makewfs import WavefrontSensor, load_config @@ -51,32 +50,3 @@ def test_optical_blur_kernel_digest_is_recorded(tmp_path: Path) -> None: np.zeros(config.input.shape), seed=4 ) assert len(frame.metadata["wfs_shack_hartmann_optical_blur_kernel_sha256"]) == 16 - - -@pytest.mark.parametrize("name", ["shack_hartmann_minimal.toml", "pyramid_minimal.toml"]) -def test_phase_input_reports_its_opd_in_metres(name: str) -> None: - # Phase input is converted at the reference wavelength before anything - # downstream sees it, including the input-RMS metadata and OpticalResult. - opd_sensor = WavefrontSensor.from_toml(CONFIGS / name) - wavelength = 1.0e-6 - phase_input = replace( - opd_sensor.config.input, - quantity="phase", - unit="rad", - reference_wavelength_m=wavelength, - ) - phase_sensor = WavefrontSensor(replace(opd_sensor.config, input=phase_input)) - shape = opd_sensor.config.input.shape - _, xx = np.indices(shape, dtype=np.float64) - opd = 30e-9 * (xx - (shape[1] - 1) / 2) / shape[1] + 10e-9 - phase = 2.0 * np.pi * opd / wavelength - - expected = float(np.sqrt(np.mean(opd**2))) - from_opd = opd_sensor.expose(opd, seed=1).metadata["wfs_input_opd_rms_m"] - from_phase = phase_sensor.expose(phase, seed=1).metadata["wfs_input_opd_rms_m"] - assert from_opd == pytest.approx(expected, rel=1e-12) - assert from_phase == pytest.approx(expected, rel=1e-12) - np.testing.assert_allclose(phase_sensor._render(phase).opd_m, opd, rtol=1e-12) - np.testing.assert_allclose( - phase_sensor.photon_rate(phase), opd_sensor.photon_rate(opd), rtol=1e-6, atol=0.0 - ) From 3d14ac9ac3c4c96733f8114ff20ce5d3204dc51a Mon Sep 17 00:00:00 2001 From: Jacob Taylor Date: Tue, 6 Oct 2026 19:17:16 -1000 Subject: [PATCH 2/2] Adopt aocore 0.1.3: plain block_sum, device coordinates Require aocore>=0.1.3,<0.2. block_sum delegates every factor to aocore.block_sum, which measured at least as fast as the local factor-two shortcut on the Shack-Hartmann spot stacks on CPU and GPU. Pixel-centre coordinates are built on the selected device through ArrayBackend.centered_coordinates. The input-RMS tests check the unweighted key against aocore.rms_unweighted; the per-frame reduction stays local because the aocore metrics return host floats. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 12 ++++++++---- CHANGELOG.md | 16 ++++++++++++++++ pyproject.toml | 2 +- src/makewfs/backend.py | 13 +++++++++++++ src/makewfs/sampling.py | 28 ++++++++-------------------- src/makewfs/wavefront.py | 16 +++++++++------- tests/test_input_rms.py | 8 +++++--- 7 files changed, 60 insertions(+), 35 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index a3d2b34..83232ed 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -93,8 +93,10 @@ write a failing integration test/design note, use the conditional gates in the analytic pupil evaluated on `input.shape`, or a custom mask area-averaged from the engine's `configured_pupil`); `wfs_input_opd_rms_unweighted_m` is the whole-grid quadratic mean with piston kept. Both are reduced on the device and - cross in the one `backend.scalars` batch. Any other RMS-like key must say its - variant in its name (`_unweighted`, `_tiptilt_removed`). + cross in the one `backend.scalars` batch, so do not swap in `aocore.rms` or + `aocore.rms_unweighted` there: they return host floats and would add a + synchronization each. Any other RMS-like key must say its variant in its + name (`_unweighted`, `_tiptilt_removed`). - The intended top-level API is `load_config`, `WavefrontSensor`, and `simulate`. Keep other implementation objects out of `makewfs.__init__` unless an API review explicitly accepts them. @@ -161,8 +163,10 @@ Follow the target layout in `ROADMAP.md`: rotation and rectangular grids on the selected backend, none of which `aocore.Pupil` models. `backend.ArrayBackend` takes a dtype per array, explicit FFT workers and `ndimage` helpers, which `aocore.Backend` does not, and its - centred FFTs keep the `fftshift` convention. `sampling.block_sum` wraps - `aocore.block_sum` but keeps its own factor-two fast path. + centred FFTs keep the `fftshift` convention; `ArrayBackend.centered_coordinates` + builds aocore's coordinates directly on the device. `sampling.block_sum` is a + thin wrapper over `aocore.block_sum` (no local fast path since aocore 0.1.3, + which measured at least as fast at the SH call sites). `sampling.area_rebin` is exact-overlap area averaging between grids of the same extent and any shape ratio, which `aocore.block_sum` (integer factors only) does not cover. diff --git a/CHANGELOG.md b/CHANGELOG.md index bbcb73d..57efa08 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -46,6 +46,22 @@ See "Migrating to 2.0" in the stability guide. RMS they plot from each frame's `wfs_input_opd_rms_m` instead of a whole-grid `np.std`, so the numbers they print are pupil-weighted. The checked-in `makewfs_showcase.webp` was not regenerated. +- **Requires `aocore>=0.1.3,<0.2`.** `makewfs.sampling.block_sum` drops its + own factor-two shortcut and delegates every factor to `aocore.block_sum`, + whose strided CPU adds and single CuPy kernel measured at least as fast on + the spot stacks the Shack-Hartmann actually bins (for example + `(400, 16, 16)` float32: CPU 112 vs 173 us, Quadro P620 65 vs 94 us; + `(3600, 12, 12)` float64: CPU 0.90 vs 1.27 ms, GPU 98 vs 173 us). The + additions happen in a different order, so Shack-Hartmann float32 images + differ from 1.2.0 by float32 rounding (at most about 1e-7 relative) and + float64 ones by about 2e-16; pyramid images are bit-for-bit unchanged. + Pixel-centre coordinates for the input, pupil and DFT detector grids are + now built on the selected device by `aocore.centered_coordinates` + (`ArrayBackend.centered_coordinates`), with identical values. +- The input-RMS tests check the unweighted key against + `aocore.rms_unweighted`. The per-frame path keeps its own device reduction, + because the aocore functions return host floats and would each add a + synchronization. ## [1.2.0] - 2026-10-07 diff --git a/pyproject.toml b/pyproject.toml index e40e49f..83c73b4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Typing :: Typed", ] dependencies = [ - "aocore>=0.1.2,<0.2", + "aocore>=0.1.3,<0.2", "numpy>=1.23", "scipy>=1.10", "getframes>=2.2.0", diff --git a/src/makewfs/backend.py b/src/makewfs/backend.py index a7c64e0..36447e9 100644 --- a/src/makewfs/backend.py +++ b/src/makewfs/backend.py @@ -21,6 +21,7 @@ from typing import Any, cast import numpy as np +from aocore import centered_coordinates as _aocore_centered_coordinates from numpy.typing import NDArray @@ -48,6 +49,18 @@ def asarray(self, value: Any, *, dtype: Any | None = None) -> Any: """Convert a value using this backend's array namespace.""" return self.xp.asarray(value, dtype=dtype) + def centered_coordinates(self, n: int, *, dtype: Any) -> Any: + """Unit-pitch pixel-centre coordinates (CONVENTIONS 1.2) on this device. + + ``aocore.centered_coordinates`` builds them in float64 and casts to + ``dtype`` directly on the GPU, so no host array is copied over. The + CPU path still passes through :meth:`asarray`, which keeps an injected + namespace in the loop. + """ + if self.is_cpu: + return self.asarray(_aocore_centered_coordinates(n, dtype=dtype), dtype=dtype) + return _aocore_centered_coordinates(n, dtype=dtype, backend="gpu") + def zeros(self, shape: Any, *, dtype: Any) -> Any: """Allocate a zero-filled array on this backend.""" return self.xp.zeros(shape, dtype=dtype) diff --git a/src/makewfs/sampling.py b/src/makewfs/sampling.py index a24610e..5dd5c0f 100644 --- a/src/makewfs/sampling.py +++ b/src/makewfs/sampling.py @@ -9,7 +9,6 @@ import numpy as np from aocore import block_sum as _aocore_block_sum -from aocore import centered_coordinates from numpy.typing import NDArray from .backend import ArrayBackend, centered_fft_intensity, cpu_backend @@ -147,9 +146,11 @@ def block_sum( """Sum square pixel blocks while preserving total flux. A thin wrapper over ``aocore.block_sum``, which owns flux-conserving - binning (CONVENTIONS 9). It keeps this module's error messages and a - factor-two fast path. ``backend`` is accepted for compatibility; the - reduction runs on the array's own namespace. + binning (CONVENTIONS 9); it keeps this module's error messages. + ``backend`` is accepted for compatibility; the reduction runs on the + array's own namespace. Since aocore 0.1.3 its strided CPU adds and + single-kernel CuPy path are at least as fast as the factor-two shortcut + this module used to keep, on the spot stacks the Shack-Hartmann bins. """ if factor < 1: raise ValueError("factor must be positive") @@ -158,17 +159,6 @@ def block_sum( height, width = array.shape[-2:] if height % factor or width % factor: raise ValueError(f"shape {array.shape[-2:]} is not divisible by factor {factor}") - if factor == 2: - # Two-times oversampling is the common SH path. Direct strided sums - # avoid NumPy's disproportionately expensive multi-axis reduction over - # thousands of tiny spot images and work unchanged with CuPy arrays. - return cast( - NDArray[Any], - array[..., 0::2, 0::2] - + array[..., 0::2, 1::2] - + array[..., 1::2, 0::2] - + array[..., 1::2, 1::2], - ) return cast(NDArray[Any], _aocore_block_sum(array, factor)) @@ -306,8 +296,8 @@ def build( coordinate = backend.arange(nfft, dtype=np.float64) plan.half_sample = backend.exp(-1j * math.pi * coordinate / nfft) else: - detector_coordinate = backend.asarray( - centered_coordinates(high_resolution_pixels), dtype=np.float64 + detector_coordinate = backend.centered_coordinates( + high_resolution_pixels, dtype=np.float64 ) / (sampling * oversampling) pupil_coordinate = backend.arange(samples_per_lenslet, dtype=np.float64) kernel = backend.exp( @@ -319,9 +309,7 @@ def build( ) plan.dft_kernel = backend.astype(kernel, plan.field_dtype) if field_stop_radius_lambda_over_d is not None: - coordinates = backend.asarray( - centered_coordinates(high_resolution_pixels), dtype=np.float64 - ) + coordinates = backend.centered_coordinates(high_resolution_pixels, dtype=np.float64) y, x = backend.meshgrid(coordinates, coordinates, indexing="ij") radius_lambda_over_d = backend.hypot(x, y) / (oversampling * sampling) plan.field_stop_mask = radius_lambda_over_d <= field_stop_radius_lambda_over_d diff --git a/src/makewfs/wavefront.py b/src/makewfs/wavefront.py index 0b2bfbc..4d8664d 100644 --- a/src/makewfs/wavefront.py +++ b/src/makewfs/wavefront.py @@ -7,7 +7,7 @@ from typing import Any, cast import numpy as np -from aocore import centered_coordinates, phase_to_opd +from aocore import phase_to_opd from numpy.typing import ArrayLike, NDArray from .backend import ArrayBackend, cpu_backend @@ -23,14 +23,14 @@ def _coordinates( ) -> tuple[NDArray[np.float64], NDArray[np.float64]]: """Return centered physical ``(x, y)`` coordinates for an array shape. - Pixel centres follow ``aocore.centered_coordinates`` (CONVENTIONS 1.2). The - unit-pitch offsets are exact in either precision; scaling to metres happens - on the backend in ``dtype``. + Pixel centres follow ``aocore.centered_coordinates`` (CONVENTIONS 1.2), + built on the backend's device. The unit-pitch offsets are exact in either + precision; scaling to metres happens on the backend in ``dtype``. """ resolved = backend or cpu_backend() height, width = shape - x = resolved.asarray(centered_coordinates(width), dtype=dtype) * extent_m / width - y = resolved.asarray(centered_coordinates(height), dtype=dtype) * extent_m / height + x = resolved.centered_coordinates(width, dtype=dtype) * extent_m / width + y = resolved.centered_coordinates(height, dtype=dtype) * extent_m / height xx, yy = resolved.meshgrid(x, y) return xx, yy @@ -261,7 +261,9 @@ def grid_rms(opd: Any, *, backend: ArrayBackend | None = None) -> Any: This is ``rms_unweighted`` in the sense of aocore CONVENTIONS 4.1: every grid pixel counts equally, including pixels outside the pupil, and the mean - is not removed. + is not removed. It equals ``aocore.rms_unweighted(opd)``, which returns a + host float and so would synchronize on its own; this returns a device + scalar for the one batched metadata crossing. """ resolved = backend or cpu_backend() values = resolved.asarray(opd, dtype=np.float64) diff --git a/tests/test_input_rms.py b/tests/test_input_rms.py index fde893a..17ed1bd 100644 --- a/tests/test_input_rms.py +++ b/tests/test_input_rms.py @@ -9,6 +9,7 @@ import numpy as np import pytest from aocore import rms as aocore_rms +from aocore import rms_unweighted as aocore_rms_unweighted from aocore.conformance import check_rms from makewfs import WavefrontSensor, load_config @@ -80,7 +81,7 @@ def test_tilt_matches_the_analytic_pupil_weighted_rms(name: str) -> None: aocore_rms(opd, _pupil(sensor)), rel=1e-12 ) assert metadata["wfs_input_opd_rms_unweighted_m"] == pytest.approx( - math.sqrt(float(np.mean(opd**2))), rel=1e-12 + aocore_rms_unweighted(opd), rel=1e-12 ) @@ -111,7 +112,7 @@ def test_integrated_exposure_reports_the_rms_of_the_mean_opd(name: str) -> None: aocore_rms(mean_opd, _pupil(sensor)), rel=1e-12 ) assert metadata["wfs_input_opd_rms_unweighted_m"] == pytest.approx( - math.sqrt(float(np.mean(mean_opd**2))), rel=1e-12 + aocore_rms_unweighted(mean_opd), rel=1e-12 ) @@ -163,7 +164,7 @@ def test_phase_input_rms_is_reported_in_opd_metres(name: str) -> None: phase = 2.0 * np.pi * opd / wavelength weighted = aocore_rms(opd, _pupil(opd_sensor)) - unweighted = math.sqrt(float(np.mean(opd**2))) + unweighted = aocore_rms_unweighted(opd) for sensor, wavefront in ((opd_sensor, opd), (phase_sensor, phase)): metadata = sensor.expose(wavefront, seed=1).metadata assert metadata["wfs_input_opd_rms_m"] == pytest.approx(weighted, rel=1e-12) @@ -182,6 +183,7 @@ def test_empty_pupil_on_the_input_grid_is_rejected() -> None: def test_grid_rms_keeps_piston_and_counts_every_pixel() -> None: opd = np.array([[1.0, 1.0], [1.0, 5.0]]) assert float(grid_rms(opd)) == pytest.approx(math.sqrt(28.0 / 4.0)) + assert float(grid_rms(opd)) == pytest.approx(aocore_rms_unweighted(opd), rel=1e-15) @pytest.mark.parametrize(