makewfs 1.0.0__py3-none-any.whl

makewfs/__about__.py

@@ -0,0 +1,3 @@
+"""Package version metadata."""
+
+__version__ = "1.0.0"

makewfs/__init__.py

@@ -0,0 +1,15 @@
+"""Configuration-driven adaptive-optics wavefront-sensor image simulation."""
+
+from .__about__ import __version__
+from .api import WavefrontSensor, simulate
+from .config import Config, ConfigError, WFSConfig, load_config
+
+__all__ = [
+    "Config",
+    "ConfigError",
+    "WFSConfig",
+    "WavefrontSensor",
+    "__version__",
+    "load_config",
+    "simulate",
+]

makewfs/api.py

@@ -0,0 +1,210 @@
+"""Small public facade for configured wavefront-sensor simulations."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Iterator
+from pathlib import Path
+from time import perf_counter
+from typing import Any, cast
+
+import numpy as np
+from numpy.typing import ArrayLike, NDArray
+
+from .backend import ArrayBackend, cpu_backend, cupy_backend
+from .config import WFSConfig, load_config
+from .detector import DetectorAdapter
+from .provenance import metadata as build_metadata
+from .sensors.base import OpticalResult, SensorEngine
+from .sensors.pyramid import PyramidEngine
+from .sensors.shack_hartmann import ShackHartmannEngine
+from .wavefront import iter_phase_samples
+
+
+class WavefrontSensor:
+    """Configured wavefront-sensor facade.
+
+    Construct once and call :meth:`expose` for each closed-loop residual OPD.
+    ``photon_rate`` exposes the deterministic optical result for workflows that
+    use another detector or need an ideal reference image.
+    """
+
+    def __init__(self, config: WFSConfig, *, _backend: ArrayBackend | None = None) -> None:
+        self.config = config
+        self.backend = _backend or (
+            cupy_backend() if config.numerics.device == "gpu" else cpu_backend()
+        )
+        self.engine: SensorEngine
+        if config.sensor.kind == "shack_hartmann":
+            self.engine = ShackHartmannEngine(config, backend=self.backend)
+        elif config.sensor.kind == "pyramid":
+            self.engine = PyramidEngine(config, backend=self.backend)
+        else:
+            raise NotImplementedError(f"unsupported sensor kind {config.sensor.kind!r}")
+        self.detector = DetectorAdapter(
+            config.detector,
+            self.engine.output_shape,
+            device="cpu" if self.backend.is_cpu else "gpu",
+        )
+        self._metadata_base = build_metadata(
+            self.config,
+            sensor_kind=self.config.sensor.kind,
+            launched_rate=0.0,
+            captured_rate=0.0,
+            opd_rms_m=0.0,
+            seed=None,
+            source_states=self.engine.source_states,
+            file_digests=self.engine.file_digests,
+        )
+
+    @classmethod
+    def from_toml(cls, path: str | Path) -> WavefrontSensor:
+        """Build a sensor from a validated TOML file."""
+        return cls(load_config(path))
+
+    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 _frame_metadata(
+        self,
+        *,
+        launched_rate: float,
+        captured_rate: float,
+        opd_rms_m: float,
+        seed: int | None,
+    ) -> dict[str, Any]:
+        """Copy cached static provenance and fill the per-frame values."""
+        result = dict(self._metadata_base)
+        result.update(
+            {
+                "wfs_launched_photons_s": float(launched_rate),
+                "wfs_captured_photons_s": float(captured_rate),
+                "wfs_input_opd_rms_m": float(opd_rms_m),
+                "wfs_seed": seed if seed is not None else "internal",
+            }
+        )
+        return result
+
+    def photon_rate(self, wavefront: ArrayLike) -> Any:
+        """Return the ideal native-pixel photon rate before detector noise."""
+        return self._render(wavefront).photon_rate
+
+    def reference(self) -> Any:
+        """Return the ideal image for a zero dynamic OPD."""
+        dtype = np.dtype(self.config.numerics.dtype)
+        return self.photon_rate(self.backend.zeros(self.config.input.shape, dtype=dtype))
+
+    def expose(self, wavefront: ArrayLike, *, seed: int | None = None) -> Any:
+        """Render one wavefront and expose it through the configured detector."""
+        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)
+        )
+        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,
+            seed=seed,
+        )
+        detector_start = perf_counter()
+        frame = self.detector.expose(
+            result.photon_rate,
+            metadata=frame_metadata,
+            seed=seed,
+            spectral_photon_rate=(
+                None if result.spectral_photon_rate is None else result.spectral_photon_rate
+            ),
+            spectral_wavelengths_m=result.spectral_wavelengths_m,
+        )
+        detector_elapsed = perf_counter() - detector_start
+        frame.metadata["wfs_optical_render_s"] = optical_elapsed
+        frame.metadata["wfs_detector_expose_s"] = detector_elapsed
+        frame.metadata["wfs_total_expose_s"] = perf_counter() - total_start
+        return frame
+
+    def expose_many(
+        self, phases: Iterable[ArrayLike], seeds: Iterable[int | None] | None = None
+    ) -> Iterator[Any]:
+        """Yield one detector frame per phase sample without stacking the stream."""
+        seed_iter = iter(seeds) if seeds is not None else None
+        for phase in phases:
+            seed = next(seed_iter) if seed_iter is not None else None
+            yield self.expose(phase, seed=seed)
+
+    def expose_integrated(
+        self, phase_samples: ArrayLike | Iterable[ArrayLike], *, seed: int | None = None
+    ) -> Any:
+        """Expose one detector frame after uniformly averaging temporal OPD samples."""
+        total_start = perf_counter()
+        optical_start = total_start
+        rates: list[Any] = []
+        spectral_rates: list[Any] = []
+        spectral_wavelengths_m: tuple[float, ...] | None = None
+        opds: list[Any] = []
+        launched = 0.0
+        captured: Any = 0.0
+        for sample in iter_phase_samples(
+            phase_samples,
+            self.config.input.shape,
+            backend=self.backend,
+        ):
+            result = self._render(sample)
+            rates.append(result.photon_rate)
+            if result.spectral_photon_rate is not None:
+                spectral_rates.append(result.spectral_photon_rate)
+                if spectral_wavelengths_m is None:
+                    spectral_wavelengths_m = result.spectral_wavelengths_m
+                elif spectral_wavelengths_m != result.spectral_wavelengths_m:
+                    raise RuntimeError("spectral wavelength nodes changed within one exposure")
+            opds.append(result.opd_m)
+            launched = result.launched_rate_per_s
+            captured += result.captured_rate_per_s
+        if not rates:
+            raise ValueError("phase_samples must contain at least one sample")
+        average_rate = self.backend.mean(self.backend.stack(rates), axis=0)
+        average_spectral_rate = (
+            None
+            if not spectral_rates
+            else self.backend.mean(self.backend.stack(spectral_rates), axis=0)
+        )
+        average_opd = self.backend.mean(self.backend.stack(opds), axis=0)
+        captured_rate, opd_rms = self.backend.scalars(
+            captured / len(rates), self._opd_rms(average_opd)
+        )
+        frame_metadata = self._frame_metadata(
+            launched_rate=launched,
+            captured_rate=captured_rate,
+            opd_rms_m=opd_rms,
+            seed=seed,
+        )
+        frame_metadata["wfs_temporal_samples"] = len(rates)
+        optical_elapsed = perf_counter() - optical_start
+        detector_start = perf_counter()
+        frame = self.detector.expose(
+            average_rate,
+            metadata=frame_metadata,
+            seed=seed,
+            spectral_photon_rate=average_spectral_rate,
+            spectral_wavelengths_m=spectral_wavelengths_m,
+        )
+        frame.metadata["wfs_optical_render_s"] = optical_elapsed
+        frame.metadata["wfs_detector_expose_s"] = perf_counter() - detector_start
+        frame.metadata["wfs_total_expose_s"] = perf_counter() - total_start
+        return frame
+
+
+def simulate(
+    wavefront: ArrayLike, config: WFSConfig | str | Path, *, seed: int | None = None
+) -> Any:
+    """One-shot convenience wrapper around :class:`WavefrontSensor`."""
+    resolved = load_config(config) if isinstance(config, (str, Path)) else config
+    return WavefrontSensor(resolved).expose(wavefront, seed=seed)
+
+
+__all__ = ["WavefrontSensor", "simulate"]

makewfs/backend.py

@@ -0,0 +1,396 @@
+"""Array and FFT primitives used by the portable optical kernels.
+
+The CPU release uses :mod:`numpy` arrays and :mod:`scipy.fft`, but sensor
+mathematics calls this small backend object rather than allocating through
+NumPy directly.  That boundary is deliberately private today; it gives a
+future CuPy implementation one place to provide array creation, reductions,
+FFT, and interpolation semantics without changing the optical equations.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, cast
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+@dataclass(frozen=True)
+class ArrayBackend:
+    """Numerical namespace for one optical array backend.
+
+    ``xp`` is an Array API-compatible namespace.  The CPU instance is the only
+    supported instance in the public release; a future device backend may
+    provide the corresponding namespace and override the SciPy-only helpers.
+    Methods named ``scalar`` and ``to_host`` are explicit host-boundary points
+    for metadata and configuration diagnostics, rather than accidental scalar
+    extraction in a sensor kernel.
+    """
+
+    xp: Any
+    name: str = "cpu"
+
+    @property
+    def is_cpu(self) -> bool:
+        """Whether this backend uses host NumPy/SciPy arrays."""
+        return self.name == "cpu"
+
+    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 zeros(self, shape: Any, *, dtype: Any) -> Any:
+        """Allocate a zero-filled array on this backend."""
+        return self.xp.zeros(shape, dtype=dtype)
+
+    def zeros_like(self, value: Any) -> Any:
+        """Allocate an array matching ``value`` on this backend."""
+        return self.xp.zeros_like(value)
+
+    def astype(self, value: Any, dtype: Any) -> Any:
+        """Cast an array without changing its backend."""
+        return value.astype(dtype)
+
+    def full(self, shape: Any, value: Any, *, dtype: Any) -> Any:
+        """Allocate a constant-filled array on this backend."""
+        return self.xp.full(shape, value, dtype=dtype)
+
+    def empty(self, shape: Any, *, dtype: Any) -> Any:
+        """Allocate an uninitialized array on this backend."""
+        return self.xp.empty(shape, dtype=dtype)
+
+    def arange(self, *args: Any, **kwargs: Any) -> Any:
+        """Create a backend array of evenly spaced values."""
+        return self.xp.arange(*args, **kwargs)
+
+    def meshgrid(self, *args: Any, **kwargs: Any) -> Any:
+        """Create backend coordinate grids."""
+        return self.xp.meshgrid(*args, **kwargs)
+
+    def repeat(self, value: Any, repeats: Any, *, axis: int) -> Any:
+        """Repeat values along one backend axis."""
+        return self.xp.repeat(value, repeats, axis=axis)
+
+    def tile(self, value: Any, reps: Any) -> Any:
+        """Tile a backend array."""
+        return self.xp.tile(value, reps)
+
+    def stack(self, values: Any, *, axis: int = 0) -> Any:
+        """Stack backend arrays."""
+        return self.xp.stack(values, axis=axis)
+
+    def sum(self, value: Any, *, axis: Any = None) -> Any:
+        """Reduce a backend array by summation."""
+        return self.xp.sum(value, axis=axis)
+
+    def mean(self, value: Any, *, axis: Any = None) -> Any:
+        """Reduce a backend array by mean."""
+        return self.xp.mean(value, axis=axis)
+
+    def average(self, value: Any, *, weights: Any) -> Any:
+        """Compute a weighted backend average."""
+        return self.xp.average(value, weights=weights)
+
+    def any(self, value: Any) -> Any:
+        """Backend reduction testing whether any element is true."""
+        return self.xp.any(value)
+
+    def all(self, value: Any) -> Any:
+        """Backend reduction testing whether all elements are true."""
+        return self.xp.all(value)
+
+    def isfinite(self, value: Any) -> Any:
+        """Elementwise finite-value test."""
+        return self.xp.isfinite(value)
+
+    def abs(self, value: Any) -> Any:
+        """Elementwise absolute value."""
+        return self.xp.abs(value)
+
+    def exp(self, value: Any) -> Any:
+        """Elementwise exponential."""
+        return self.xp.exp(value)
+
+    def sqrt(self, value: Any) -> Any:
+        """Elementwise square root."""
+        return self.xp.sqrt(value)
+
+    def hypot(self, left: Any, right: Any) -> Any:
+        """Elementwise Euclidean norm."""
+        return self.xp.hypot(left, right)
+
+    def cos(self, value: Any) -> Any:
+        """Elementwise cosine."""
+        return self.xp.cos(value)
+
+    def sin(self, value: Any) -> Any:
+        """Elementwise sine."""
+        return self.xp.sin(value)
+
+    def arctan2(self, left: Any, right: Any) -> Any:
+        """Elementwise two-argument arctangent."""
+        return self.xp.arctan2(left, right)
+
+    def mod(self, value: Any, divisor: Any) -> Any:
+        """Elementwise remainder."""
+        return self.xp.mod(value, divisor)
+
+    def where(self, condition: Any, left: Any, right: Any) -> Any:
+        """Select values elementwise on the backend."""
+        return self.xp.where(condition, left, right)
+
+    def ptp(self, value: Any) -> Any:
+        """Backend peak-to-peak reduction."""
+        return self.xp.ptp(value)
+
+    def argmax(self, value: Any) -> Any:
+        """Return the flat index of a backend array maximum."""
+        return self.xp.argmax(value)
+
+    def fftfreq(self, value: int) -> Any:
+        """Return backend FFT frequency bins."""
+        return self.xp.fft.fftfreq(value)
+
+    def fftshift(self, value: Any, *, axes: Any = None) -> Any:
+        """Shift zero frequency to the center."""
+        return self.xp.fft.fftshift(value, axes=axes)
+
+    def ifftshift(self, value: Any, *, axes: Any = None) -> Any:
+        """Undo a centered FFT shift."""
+        return self.xp.fft.ifftshift(value, axes=axes)
+
+    def centered_fft2(self, array: Any, *, workers: int = 1) -> Any:
+        """Perform a centered, unitary two-dimensional FFT."""
+        axes = (-2, -1)
+        if self.is_cpu:
+            from scipy import fft
+
+            transformed = fft.fftshift(
+                fft.fft2(
+                    fft.ifftshift(array, axes=axes),
+                    axes=axes,
+                    workers=workers,
+                    norm="ortho",
+                    overwrite_x=True,
+                ),
+                axes=axes,
+            )
+        else:  # pragma: no cover - reserved for a future device backend
+            transformed = self.fftshift(
+                self.xp.fft.fft2(self.ifftshift(array, axes=axes), axes=axes, norm="ortho"),
+                axes=axes,
+            )
+        return transformed
+
+    def centered_fft_intensity(
+        self, array: Any, *, workers: int = 1, overwrite_input: bool = False
+    ) -> Any:
+        """Return centered unitary FFT intensity without an irrelevant input roll.
+
+        Translating an entrance field changes only Fourier phase, so an input
+        ``ifftshift`` cannot affect intensity. Shack-Hartmann propagation uses
+        this identity to avoid one detector-batch-sized array permutation.
+        """
+        axes = (-2, -1)
+        height, width = array.shape[-2:]
+        if height % 2 == 0 and width % 2 == 0:
+            working = array if overwrite_input else self.xp.array(array, copy=True)
+            working[..., ::2, 1::2] *= -1
+            working[..., 1::2, ::2] *= -1
+            if self.is_cpu:
+                from scipy import fft
+
+                transformed = fft.fft2(
+                    working,
+                    axes=axes,
+                    workers=workers,
+                    norm="ortho",
+                    overwrite_x=overwrite_input,
+                )
+            else:  # pragma: no cover - GPU optional
+                transformed = self.xp.fft.fft2(working, axes=axes, norm="ortho")
+            return self.abs(transformed) ** 2
+        if self.is_cpu:
+            from scipy import fft
+
+            transformed = fft.fftshift(
+                fft.fft2(array, axes=axes, workers=workers, norm="ortho"),
+                axes=axes,
+            )
+        else:  # pragma: no cover - GPU optional
+            transformed = self.fftshift(self.xp.fft.fft2(array, axes=axes, norm="ortho"), axes=axes)
+        return self.abs(transformed) ** 2
+
+    def centered_ifft2(self, array: Any, *, workers: int = 1) -> Any:
+        """Perform a centered, unitary two-dimensional inverse FFT."""
+        axes = (-2, -1)
+        if self.is_cpu:
+            from scipy import fft
+
+            transformed = fft.fftshift(
+                fft.ifft2(
+                    fft.ifftshift(array, axes=axes),
+                    axes=axes,
+                    workers=workers,
+                    norm="ortho",
+                    overwrite_x=True,
+                ),
+                axes=axes,
+            )
+        else:  # pragma: no cover - reserved for a future device backend
+            transformed = self.fftshift(
+                self.xp.fft.ifft2(self.ifftshift(array, axes=axes), axes=axes, norm="ortho"),
+                axes=axes,
+            )
+        return transformed
+
+    def map_coordinates(self, array: Any, coordinates: Any, *, order: int, mode: str) -> Any:
+        """Interpolate coordinates, using SciPy only for the CPU backend."""
+        if self.is_cpu:
+            from scipy.ndimage import map_coordinates
+
+            return map_coordinates(array, coordinates, order=order, mode=mode)
+        from cupyx.scipy.ndimage import map_coordinates  # pragma: no cover - GPU optional
+
+        device_coordinates = (
+            self.xp.stack(coordinates, axis=0)
+            if isinstance(coordinates, (list, tuple))
+            else coordinates
+        )
+        return map_coordinates(array, device_coordinates, order=order, mode=mode)
+
+    def convolve(self, array: Any, kernel: Any) -> Any:
+        """Convolve a batch of arrays with a backend-compatible kernel."""
+        if self.is_cpu:
+            from scipy.ndimage import convolve
+
+            return convolve(array, kernel, mode="constant", cval=0.0)
+        from cupyx.scipy.ndimage import convolve  # pragma: no cover - GPU optional
+
+        return convolve(array, kernel, mode="constant", cval=0.0)
+
+    def gaussian_filter(self, array: Any, sigma: Any) -> Any:
+        """Apply a Gaussian filter through the CPU numerical backend."""
+        if self.is_cpu:
+            from scipy.ndimage import gaussian_filter
+
+            return gaussian_filter(array, sigma=sigma, mode="constant")
+        from cupyx.scipy.ndimage import gaussian_filter  # pragma: no cover - GPU optional
+
+        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 int(next_fast_len(value))
+
+    def scalar(self, value: Any) -> float:
+        """Extract one host scalar at an explicit metadata/geometry boundary."""
+        item = value.item() if hasattr(value, "item") else value
+        return float(item)
+
+    def scalars(self, *values: Any) -> tuple[float, ...]:
+        """Extract several scalars with one device synchronization."""
+        if self.is_cpu:
+            return tuple(self.scalar(value) for value in values)
+        packed = self.xp.stack([self.xp.asarray(value) for value in values])
+        return tuple(float(value) for value in self.xp.asnumpy(packed))
+
+    def to_host(self, value: Any) -> NDArray[Any]:
+        """Copy an array to host NumPy storage at an explicit boundary."""
+        if self.is_cpu:
+            return cast(NDArray[Any], value)
+        return np.asarray(self.xp.asnumpy(value))
+
+
+_CPU_BACKEND = ArrayBackend(np, name="cpu")
+
+
+def cpu_backend() -> ArrayBackend:
+    """Return the shared CPU backend instance."""
+    return _CPU_BACKEND
+
+
+def cupy_backend() -> ArrayBackend:
+    """Return the private optional CuPy backend.
+
+    CuPy is intentionally imported lazily so the core package remains usable
+    without CUDA.  Callers should treat this as experimental and keep the
+    explicit host transfer before the ``getframes`` detector adapter.
+    """
+    try:
+        import cupy
+    except ImportError as exc:  # pragma: no cover - depends on optional install
+        raise ImportError("the private CuPy backend requires makewfs[gpu]") from exc
+    return ArrayBackend(cupy, name="cupy")
+
+
+def real_dtype(name: str) -> np.dtype[Any]:
+    """Return the configured real dtype."""
+    if name == "float32":
+        return np.dtype(np.float32)
+    if name == "float64":
+        return np.dtype(np.float64)
+    raise ValueError(f"unsupported dtype {name!r}")
+
+
+def complex_dtype(name: str) -> np.dtype[Any]:
+    """Return the complex dtype paired with a real dtype."""
+    if name == "float32":
+        return np.dtype(np.complex64)
+    if name == "float64":
+        return np.dtype(np.complex128)
+    raise ValueError(f"unsupported dtype {name!r}")
+
+
+def centered_fft2(
+    array: NDArray[Any], *, workers: int = 1, backend: ArrayBackend | None = None
+) -> NDArray[Any]:
+    """Centered two-dimensional FFT with unitary normalization."""
+    return cast(NDArray[Any], (backend or cpu_backend()).centered_fft2(array, workers=workers))
+
+
+def centered_ifft2(
+    array: NDArray[Any], *, workers: int = 1, backend: ArrayBackend | None = None
+) -> NDArray[Any]:
+    """Centered two-dimensional inverse FFT with unitary normalization."""
+    return cast(NDArray[Any], (backend or cpu_backend()).centered_ifft2(array, workers=workers))
+
+
+def centered_fft_intensity(
+    array: NDArray[Any],
+    *,
+    workers: int = 1,
+    backend: ArrayBackend | None = None,
+    overwrite_input: bool = False,
+) -> NDArray[Any]:
+    """Return centered unitary FFT intensity without shifting the input field."""
+    return cast(
+        NDArray[Any],
+        (backend or cpu_backend()).centered_fft_intensity(
+            array, workers=workers, overwrite_input=overwrite_input
+        ),
+    )
+
+
+def next_fast_length(value: int) -> int:
+    """Return a convenient CPU FFT length."""
+    return cpu_backend().next_fast_length(value)
+
+
+__all__ = [
+    "ArrayBackend",
+    "centered_fft2",
+    "centered_fft_intensity",
+    "centered_ifft2",
+    "complex_dtype",
+    "cpu_backend",
+    "cupy_backend",
+    "next_fast_length",
+    "real_dtype",
+]