pyturb 0.2.0__py3-none-any.whl

pyturb/__init__.py

@@ -0,0 +1,97 @@
+"""pyturb — fast, GPU-optional atmospheric phase screens for adaptive optics.
+
+Three levels of API cover the common AO simulation needs:
+
+- :class:`Atmosphere` — a full layered atmosphere: many turbulent layers with
+  per-layer wind summed into pupil OPD, with frozen-flow time evolution,
+  off-axis directions, and standard site profiles. This is the high-level
+  entry point most users want.
+- :class:`PhaseScreen` — statistically independent Kolmogorov / von Kármán
+  screens via the FFT method with subharmonic low-frequency correction.
+- :class:`InfinitePhaseScreen` — an endless frozen-flow screen extruded row
+  by row (Assémat & Wilson 2006) for closed-loop temporal simulation.
+
+Pass ``device="gpu"`` to any of them to run on CUDA via CuPy. Phase screens
+are returned in radians; :class:`Atmosphere` returns OPD in metres (achromatic)
+unless a ``wavelength`` is given.
+"""
+
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _pkg_version
+
+from . import analysis
+from .analysis import zernike_basis, zernike_decompose
+from .atmosphere import Atmosphere, PeriodicWrapWarning
+from .backend import get_array_module, get_fft_workers, set_fft_workers, to_numpy
+from .benchmark import benchmark
+from .flow import FourierFlowScreen
+from .fourier import PhaseScreen
+from .infinite import InfinitePhaseScreen, phase_covariance
+from .io import load, save
+from .profiles import (
+    Layer,
+    bufton_wind,
+    coherence_time,
+    discretize_cn2,
+    effective_wind_speed,
+    get_profile,
+    greenwood_frequency,
+    hufnagel_valley,
+    isoplanatic_angle,
+    list_profiles,
+    mean_turbulence_height,
+)
+from .utils import (
+    air_refractivity,
+    opd_to_phase,
+    phase_to_opd,
+    r0_at_wavelength,
+    r0_from_seeing,
+    seeing_from_r0,
+    structure_function,
+    water_vapour_refractivity,
+)
+
+try:  # single source of truth is the installed package metadata (pyproject.toml)
+    __version__ = _pkg_version("pyturb")
+except PackageNotFoundError:  # not installed (e.g. running from a source tree)
+    __version__ = "0.0.0+unknown"
+
+__all__ = [
+    "Atmosphere",
+    "PeriodicWrapWarning",
+    "Layer",
+    "PhaseScreen",
+    "InfinitePhaseScreen",
+    "FourierFlowScreen",
+    "phase_covariance",
+    "structure_function",
+    "get_profile",
+    "list_profiles",
+    "hufnagel_valley",
+    "bufton_wind",
+    "discretize_cn2",
+    "isoplanatic_angle",
+    "coherence_time",
+    "greenwood_frequency",
+    "mean_turbulence_height",
+    "effective_wind_speed",
+    "r0_from_seeing",
+    "seeing_from_r0",
+    "r0_at_wavelength",
+    "opd_to_phase",
+    "phase_to_opd",
+    "air_refractivity",
+    "water_vapour_refractivity",
+    "save",
+    "load",
+    "analysis",
+    "zernike_basis",
+    "zernike_decompose",
+    "benchmark",
+    "to_numpy",
+    "get_array_module",
+    "set_fft_workers",
+    "get_fft_workers",
+    "__version__",
+]

pyturb/_accel.py

@@ -0,0 +1,175 @@
+"""Optional Numba-accelerated CPU kernels (transparent, with NumPy fallback).
+
+The CPU hot paths — the spectral engine's per-frame layer sum and the
+extruder's per-frame bicubic readout — are memory-bandwidth bound in plain
+NumPy because each fuses several broadcast multiplies into large ``(L, n, n)``
+temporaries. A single fused, ``prange``-parallel pass over the data does the
+same arithmetic reading the inputs once and writing the output once, across all
+cores. These mirror the GPU's fused kernels so both backends do the same work
+per frame.
+
+Numba is an *optional* dependency: if it is not importable, :data:`HAVE_NUMBA`
+is ``False`` and callers fall back to the NumPy expressions. Nothing here
+changes results beyond float round-off (the fused reductions accumulate in the
+input dtype for the spectral sum and in double for the readout, matching each
+NumPy path).
+"""
+
+from __future__ import annotations
+
+from importlib.util import find_spec
+
+HAVE_NUMBA = find_spec("numba") is not None
+
+
+if HAVE_NUMBA:
+    import numpy as np
+    from numba import njit, prange
+
+    @njit(fastmath=True, cache=True)
+    def spectral_layer_sum(spectra, px, py, out):  # pragma: no cover - jit
+        """``out[i,j] = sum_l spectra[l,i,j] * px[l,i] * py[l,j]``.
+
+        The separable frozen-flow shift and layer sum in one fused pass: the
+        ``(L, n, n)`` spectrum stack is read once and the ``(n, n)`` shifted sum
+        written once, instead of materialising two ``(L, n, n)`` complex
+        products and reducing them.
+
+        Kept single-threaded on purpose: reducing a power-of-two-sized complex
+        stack over the leading (layer) axis strides the inner loop by a whole
+        2-D plane, so a ``prange`` over pixels hits cache-set aliasing exactly
+        at the power-of-two grids AO uses (256/512/1024). The serial fused pass
+        is a robust, monotonic win over NumPy at every size; the far larger
+        parallel win goes to the extruder readout, whose gather is
+        data-dependent and does not alias.
+        """
+        n_layers, n_rows, n_cols = spectra.shape
+        for i in range(n_rows):
+            for j in range(n_cols):
+                acc = spectra[0, i, j] * px[0, i] * py[0, j]
+                for lyr in range(1, n_layers):
+                    acc += spectra[lyr, i, j] * px[lyr, i] * py[lyr, j]
+                out[i, j] = acc
+
+    @njit(parallel=True, fastmath=True, cache=True)
+    def extrude_cubic_readout(buf, along, perp, sa, sp, fillm1, out):  # pragma: no cover
+        """Fused multi-layer Catmull-Rom pupil readout, summed over layers.
+
+        ``buf`` is the ``(L, cap, W)`` ring-buffer stack; ``along``/``perp`` are
+        the ``(L, n, n)`` rotated pupil grids; ``sa``/``sp`` the per-layer
+        along/perp readout shifts; ``fillm1`` the per-layer last valid row.
+        Accumulates in double, matching the tap-broadcast gather it replaces.
+        """
+        n_layers = buf.shape[0]
+        width = buf.shape[2]
+        n = out.shape[0]
+        for i in prange(n):
+            for j in range(n):
+                acc = 0.0
+                for lyr in range(n_layers):
+                    row = along[lyr, i, j] + sa[lyr]
+                    col = perp[lyr, i, j] + sp[lyr]
+                    r0 = int(np.floor(row))
+                    c0 = int(np.floor(col))
+                    fr = row - r0
+                    fc = col - c0
+                    fr2 = fr * fr
+                    fr3 = fr2 * fr
+                    fc2 = fc * fc
+                    fc3 = fc2 * fc
+                    wr = (
+                        0.5 * (-fr + 2.0 * fr2 - fr3),
+                        0.5 * (2.0 - 5.0 * fr2 + 3.0 * fr3),
+                        0.5 * (fr + 4.0 * fr2 - 3.0 * fr3),
+                        0.5 * (-fr2 + fr3),
+                    )
+                    wc = (
+                        0.5 * (-fc + 2.0 * fc2 - fc3),
+                        0.5 * (2.0 - 5.0 * fc2 + 3.0 * fc3),
+                        0.5 * (fc + 4.0 * fc2 - 3.0 * fc3),
+                        0.5 * (-fc2 + fc3),
+                    )
+                    fmax = fillm1[lyr]
+                    val = 0.0
+                    for a in range(4):
+                        rr = r0 + (a - 1)
+                        if rr < 0:
+                            rr = 0
+                        elif rr > fmax:
+                            rr = fmax
+                        rs = 0.0
+                        for b in range(4):
+                            cc = c0 + (b - 1)
+                            if cc < 0:
+                                cc = 0
+                            elif cc > width - 1:
+                                cc = width - 1
+                            rs += wc[b] * buf[lyr, rr, cc]
+                        val += wr[a] * rs
+                    acc += val
+                out[i, j] = acc
+
+    @njit(fastmath=True, inline="always")
+    def _sincpi(x):  # pragma: no cover - jit helper
+        if x == 0.0:
+            return 1.0
+        px = 3.141592653589793 * x
+        return np.sin(px) / px
+
+    @njit(parallel=True, fastmath=True, cache=True)
+    def extrude_lanczos_readout(buf, along, perp, sa, sp, fillm1, out):  # noqa: E501  pragma: no cover
+        """Fused multi-layer Lanczos-3 (6-tap) pupil readout, summed over layers.
+
+        The higher-fidelity counterpart of :func:`extrude_cubic_readout`: a
+        flatter sub-Nyquist windowed-sinc kernel (6 taps per axis) that halves
+        the extruder's finest-scale structure-function deficit. Accumulates in
+        double, matching the tap-broadcast Lanczos gather it replaces.
+        """
+        n_layers = buf.shape[0]
+        width = buf.shape[2]
+        n = out.shape[0]
+        for i in prange(n):
+            for j in range(n):
+                acc = 0.0
+                for lyr in range(n_layers):
+                    row = along[lyr, i, j] + sa[lyr]
+                    col = perp[lyr, i, j] + sp[lyr]
+                    r0 = int(np.floor(row))
+                    c0 = int(np.floor(col))
+                    tr = row - r0
+                    tc = col - c0
+                    wr0 = _sincpi(tr + 2.0) * _sincpi((tr + 2.0) / 3.0)
+                    wr1 = _sincpi(tr + 1.0) * _sincpi((tr + 1.0) / 3.0)
+                    wr2 = _sincpi(tr) * _sincpi(tr / 3.0)
+                    wr3 = _sincpi(tr - 1.0) * _sincpi((tr - 1.0) / 3.0)
+                    wr4 = _sincpi(tr - 2.0) * _sincpi((tr - 2.0) / 3.0)
+                    wr5 = _sincpi(tr - 3.0) * _sincpi((tr - 3.0) / 3.0)
+                    sr = wr0 + wr1 + wr2 + wr3 + wr4 + wr5
+                    wr = (wr0 / sr, wr1 / sr, wr2 / sr, wr3 / sr, wr4 / sr, wr5 / sr)
+                    wc0 = _sincpi(tc + 2.0) * _sincpi((tc + 2.0) / 3.0)
+                    wc1 = _sincpi(tc + 1.0) * _sincpi((tc + 1.0) / 3.0)
+                    wc2 = _sincpi(tc) * _sincpi(tc / 3.0)
+                    wc3 = _sincpi(tc - 1.0) * _sincpi((tc - 1.0) / 3.0)
+                    wc4 = _sincpi(tc - 2.0) * _sincpi((tc - 2.0) / 3.0)
+                    wc5 = _sincpi(tc - 3.0) * _sincpi((tc - 3.0) / 3.0)
+                    sc = wc0 + wc1 + wc2 + wc3 + wc4 + wc5
+                    wc = (wc0 / sc, wc1 / sc, wc2 / sc, wc3 / sc, wc4 / sc, wc5 / sc)
+                    fmax = fillm1[lyr]
+                    val = 0.0
+                    for a in range(6):
+                        rr = r0 + (a - 2)
+                        if rr < 0:
+                            rr = 0
+                        elif rr > fmax:
+                            rr = fmax
+                        rs = 0.0
+                        for b in range(6):
+                            cc = c0 + (b - 2)
+                            if cc < 0:
+                                cc = 0
+                            elif cc > width - 1:
+                                cc = width - 1
+                            rs += wc[b] * buf[lyr, rr, cc]
+                        val += wr[a] * rs
+                    acc += val
+                out[i, j] = acc

pyturb/analysis.py

@@ -0,0 +1,246 @@
+"""Analysis and validation utilities: Zernikes, temporal PSDs, decorrelation.
+
+These turn a pile of phase screens into the diagnostics AO people actually
+check turbulence against:
+
+- :func:`zernike_basis` / :func:`zernike_decompose` — a Noll-ordered Zernike
+  basis on a circular pupil and a least-squares projection onto it.
+- :func:`noll_variance` / :func:`noll_residual_variance` — the Kolmogorov
+  Zernike-mode variances and post-correction residuals of Noll (1976), the
+  textbook thing to validate a decomposition against.
+- :func:`temporal_psd` / :func:`fit_power_law` — a one-sided temporal power
+  spectrum and a log-log slope fit (frozen-flow gives power-law regimes).
+- :func:`differential_variance` — angular (anisoplanatism) decorrelation.
+
+Everything works on NumPy or CuPy input (device arrays are brought to the host
+for the reductions). Reference: Noll, R. J. (1976), JOSA 66, 207.
+"""
+
+from __future__ import annotations
+
+from math import factorial
+from typing import Optional, Tuple
+
+import numpy as np
+from numpy.typing import ArrayLike
+
+from .backend import to_numpy
+
+__all__ = [
+    "noll_to_zernike",
+    "zernike_basis",
+    "zernike_decompose",
+    "noll_variance",
+    "noll_residual_variance",
+    "temporal_psd",
+    "fit_power_law",
+    "differential_variance",
+]
+
+# Noll (1976) Table IV: residual wavefront variance after perfectly correcting
+# the first J Zernike modes, in units of (D/r0)^{5/3} rad^2. Delta[1] is the
+# residual after removing piston (i.e. the total minus piston).
+_NOLL_RESIDUAL = {
+    1: 1.0299, 2: 0.582, 3: 0.134, 4: 0.111, 5: 0.0880, 6: 0.0648,
+    7: 0.0587, 8: 0.0525, 9: 0.0463, 10: 0.0401, 11: 0.0377, 12: 0.0352,
+    13: 0.0328, 14: 0.0304, 15: 0.0279, 16: 0.0267, 17: 0.0255, 18: 0.0243,
+    19: 0.0232, 20: 0.0220, 21: 0.0208,
+}
+
+
+def _noll_residual_coeff(j):
+    """Delta_j in (D/r0)^{5/3} units, tabulated then asymptotic (Noll 1976)."""
+    if j < 1:
+        raise ValueError("Noll index j must be >= 1")
+    if j in _NOLL_RESIDUAL:
+        return _NOLL_RESIDUAL[j]
+    # Large-J asymptote Delta_J ~ 0.2944 J^{-sqrt(3)/2}.
+    return 0.2944 * j ** (-np.sqrt(3.0) / 2.0)
+
+
+def noll_to_zernike(j: int) -> Tuple[int, int]:
+    """Radial/azimuthal orders ``(n, m)`` for Noll single index ``j`` (>= 1)."""
+    if j < 1:
+        raise ValueError("Noll index j must be >= 1")
+    n = 0
+    j1 = j - 1
+    while j1 > n:
+        n += 1
+        j1 -= n
+    m = (-1) ** j * ((n % 2) + 2 * ((j1 + ((n + 1) % 2)) // 2))
+    return n, m
+
+
+def _radial(n, m, rho):
+    m = abs(m)
+    out = np.zeros_like(rho)
+    for k in range((n - m) // 2 + 1):
+        c = ((-1) ** k * factorial(n - k)) / (
+            factorial(k)
+            * factorial((n + m) // 2 - k)
+            * factorial((n - m) // 2 - k)
+        )
+        out += c * rho ** (n - 2 * k)
+    return out
+
+
+def zernike_basis(
+    n_modes: int, n_pixels: int, diameter_pixels: Optional[float] = None
+) -> np.ndarray:
+    """Noll-ordered Zernike basis over a circular pupil.
+
+    Parameters
+    ----------
+    n_modes : int
+        Number of modes, starting at Noll ``j = 1`` (piston).
+    n_pixels : int
+        Grid size; the returned array is ``(n_modes, n_pixels, n_pixels)``.
+    diameter_pixels : float, optional
+        Pupil diameter in pixels (default ``n_pixels``). The pupil is the
+        inscribed disc; values are 0 outside it.
+
+    Returns
+    -------
+    basis : ndarray
+        ``(n_modes, n_pixels, n_pixels)``. Each mode is orthonormalised so that
+        its variance over the pupil is 1 (Noll normalisation), and the pupil
+        mask is ``basis[0] != 0`` (piston).
+    """
+    if n_modes < 1 or n_pixels < 2:
+        raise ValueError("n_modes >= 1 and n_pixels >= 2 required")
+    radius = (n_pixels if diameter_pixels is None else diameter_pixels) / 2.0
+    grid = (np.arange(n_pixels) - (n_pixels - 1) / 2.0) / radius
+    xx, yy = np.meshgrid(grid, grid, indexing="ij")
+    rho = np.hypot(xx, yy)
+    theta = np.arctan2(yy, xx)
+    mask = rho <= 1.0
+
+    basis = np.zeros((n_modes, n_pixels, n_pixels))
+    for idx in range(n_modes):
+        j = idx + 1
+        n, m = noll_to_zernike(j)
+        norm = np.sqrt(n + 1.0) * (1.0 if m == 0 else np.sqrt(2.0))
+        ang = np.cos(m * theta) if m >= 0 else np.sin(-m * theta)
+        z = norm * _radial(n, m, rho) * ang
+        basis[idx] = np.where(mask, z, 0.0)
+    return basis
+
+
+def zernike_decompose(
+    phase: ArrayLike, n_modes: int, basis: Optional[np.ndarray] = None
+) -> np.ndarray:
+    """Least-squares Zernike coefficients of ``phase`` over the pupil.
+
+    Parameters
+    ----------
+    phase : ndarray
+        ``(n, n)`` screen or ``(count, n, n)`` stack (NumPy or CuPy).
+    n_modes : int
+        Number of Noll modes to fit.
+    basis : ndarray, optional
+        A precomputed :func:`zernike_basis` (reused across many calls to avoid
+        rebuilding it). Must match ``n_modes`` and the screen size.
+
+    Returns
+    -------
+    coeffs : ndarray
+        ``(n_modes,)`` or ``(count, n_modes)`` coefficients in the same units
+        as ``phase``.
+    """
+    phase = to_numpy(phase).astype(np.float64)
+    single = phase.ndim == 2
+    if single:
+        phase = phase[None]
+    n_pixels = phase.shape[-1]
+    if basis is None:
+        basis = zernike_basis(n_modes, n_pixels)
+    mask = basis[0] != 0
+    design = basis[:, mask].T  # (n_pupil, n_modes)
+    data = phase[:, mask].T  # (n_pupil, count)
+    coeffs, *_ = np.linalg.lstsq(design, data, rcond=None)
+    coeffs = coeffs.T  # (count, n_modes)
+    return coeffs[0] if single else coeffs
+
+
+def noll_variance(j: int, diameter: float, r0: float) -> float:
+    """Kolmogorov variance of Zernike mode ``j`` [rad^2] (Noll 1976).
+
+    The per-mode variance is ``Delta_{j-1} - Delta_j`` in ``(D/r0)^{5/3}``
+    units. ``j = 1`` (piston) has no finite variance, so ``j >= 2``.
+    """
+    if j < 2:
+        raise ValueError("piston (j=1) has no finite variance; use j >= 2")
+    coeff = _noll_residual_coeff(j - 1) - _noll_residual_coeff(j)
+    return coeff * (diameter / r0) ** (5.0 / 3.0)
+
+
+def noll_residual_variance(j: int, diameter: float, r0: float) -> float:
+    """Residual wavefront variance [rad^2] after correcting the first ``j``
+    Zernike modes (Noll 1976)."""
+    return _noll_residual_coeff(j) * (diameter / r0) ** (5.0 / 3.0)
+
+
+def temporal_psd(series: ArrayLike, dt: float) -> Tuple[np.ndarray, np.ndarray]:
+    """One-sided temporal power spectral density of a time series.
+
+    Parameters
+    ----------
+    series : array_like
+        Samples along the **last** axis (e.g. a pupil pixel or a Zernike
+        coefficient over frames); leading axes are treated as independent
+        series and averaged.
+    dt : float
+        Sample spacing [s].
+
+    Returns
+    -------
+    freq, psd : ndarray
+        Positive frequencies [Hz] (excluding DC) and the averaged PSD, scaled
+        so that ``sum(psd) * df`` approximates the series variance.
+    """
+    series = to_numpy(series).astype(np.float64)
+    series = series - series.mean(axis=-1, keepdims=True)
+    n = series.shape[-1]
+    spectrum = np.fft.rfft(series, axis=-1)
+    psd = (np.abs(spectrum) ** 2) * (2.0 * dt / n)
+    psd = psd.reshape(-1, psd.shape[-1]).mean(axis=0)
+    freq = np.fft.rfftfreq(n, d=dt)
+    return freq[1:], psd[1:]
+
+
+def fit_power_law(
+    freq: ArrayLike,
+    psd: ArrayLike,
+    fmin: Optional[float] = None,
+    fmax: Optional[float] = None,
+) -> Tuple[float, float]:
+    """Fit ``psd ~ freq**slope`` over ``[fmin, fmax]`` (log-log least squares).
+
+    Returns
+    -------
+    slope, amplitude : float
+        Such that ``psd ≈ amplitude * freq**slope`` in the band.
+    """
+    freq = np.asarray(freq, dtype=np.float64)
+    psd = np.asarray(psd, dtype=np.float64)
+    band = np.ones(freq.shape, dtype=bool)
+    if fmin is not None:
+        band &= freq >= fmin
+    if fmax is not None:
+        band &= freq <= fmax
+    band &= psd > 0
+    if band.sum() < 2:
+        raise ValueError("need at least two positive points in the band")
+    slope, intercept = np.polyfit(np.log(freq[band]), np.log(psd[band]), 1)
+    return float(slope), float(np.exp(intercept))
+
+
+def differential_variance(reference: ArrayLike, other: ArrayLike) -> float:
+    """Variance of ``other - reference`` [same units squared].
+
+    With ``reference`` the on-axis OPD/phase and ``other`` an off-axis one, this
+    is the angular (anisoplanatism) error; it grows as ``(theta/theta0)^{5/3}``
+    and reaches ~1 rad^2 at the isoplanatic angle.
+    """
+    diff = to_numpy(other).astype(np.float64) - to_numpy(reference).astype(np.float64)
+    return float(np.var(diff))