thunder4d 0.9.0__py3-none-any.whl

src/__init__.py

@@ -0,0 +1,16 @@
+"""THUNDER4D - Transverse-resolved Herriott-cell UPPE solver for Nonlinear Dynamics close to Experimental Reality, 4D.
+
+4D (x, y, z, t) UPPE simulation of multipass cells.
+"""
+__version__ = "0.9.0"
+from .backend import gpu_available, to_host
+from .grid import Grid, suggest_grid, C0
+from .materials import Material
+from .pulse import (Pulse, Beam, gaussian, super_gaussian, top_hat, knife_edge, measured_profile,
+                    hermite_gauss, laguerre_gauss, mode_superposition,
+                    gaussian_spectrum, measured_spectrum, spectrum_from_temporal_field,
+                    add_spectral_phase, zernike_wavefront)
+from .cell import MPC
+from .solver import UPPESolver
+from .simulation import MPCSimulation
+from . import diagnostics

src/backend.py

@@ -0,0 +1,64 @@
+"""Array backend: NumPy on the CPU, CuPy on an NVIDIA GPU.
+
+Only the heavy arrays live on the device: the field A(x, y, Om), the propagators, the absorbers and
+the mirror transfer function. Axes (grid.x, grid.lam, ...) and everything returned by
+`thunder4d.diagnostics` are always NumPy arrays; `to_host` converts a device array.
+"""
+import numpy as np
+
+_GPU_OK = None
+
+
+def gpu_available():
+    """True if CuPy is installed and a working CUDA device can run a kernel and an FFT."""
+    global _GPU_OK
+    if _GPU_OK is None:
+        try:
+            import cupy as cp
+            import cupyx.scipy.fft as cfft
+            if cp.cuda.runtime.getDeviceCount() < 1:
+                raise RuntimeError("no CUDA device")
+            a = cp.ones((8, 8), cp.complex64)
+            cfft.fft2(a, axes=(0, 1))
+            float(cp.sum(cp.abs(a) ** 2))       # compiles a kernel: fails on unsupported GPUs
+            _GPU_OK = True
+        except Exception:
+            _GPU_OK = False
+    return _GPU_OK
+
+
+def resolve_device(device="auto"):
+    """'auto' -> 'gpu' if CuPy works else 'cpu'; 'cpu' / 'gpu' are forced."""
+    device = "cpu" if device is None else str(device).lower()
+    if device == "auto":
+        return "gpu" if gpu_available() else "cpu"
+    if device == "cpu":
+        return "cpu"
+    if device in ("gpu", "cuda"):
+        if not gpu_available():
+            raise RuntimeError("device='gpu' requested but CuPy with a working CUDA device was not "
+                               "found (pip install cupy-cuda12x); use device='auto' or 'cpu'")
+        return "gpu"
+    raise ValueError("device must be 'auto', 'cpu' or 'gpu'")
+
+
+def get_xp(device):
+    """Array module for a resolved device name."""
+    if device == "gpu":
+        import cupy
+        return cupy
+    return np
+
+
+def is_device_array(a):
+    return type(a).__module__.split(".")[0] == "cupy"
+
+
+def xp_of(a):
+    """Array module that owns `a` (numpy for anything that is not a CuPy array)."""
+    return get_xp("gpu") if is_device_array(a) else np
+
+
+def to_host(a):
+    """NumPy view/copy of an array that may live on the GPU (no copy for NumPy input)."""
+    return a.get() if is_device_array(a) else np.asarray(a)

src/cell.py

@@ -0,0 +1,142 @@
+"""Module 2: the multipass cell.
+
+Unfolded picture used by the solver: one pass = cell centre -> d/2 -> curved mirror -> d/2 -> cell
+centre. The field is recorded at the cell centre (waist plane of the eigenmode, flat phase) after
+every pass, which is where the modal decomposition on the cell eigenmodes is exact.
+
+Nonlinear mode matching (aberration-free Kerr lens, uniform medium, power P constant along a pass):
+the envelope equation of a gaussian becomes the linear one with lambda -> lambda sqrt(1 - P/Pc),
+Pc = lambda^2/(2 pi n0 n2). Hence the matched waist and mirror spot are the linear ones scaled by
+(1 - P/Pc)^(1/4), and the B-integral per pass (on axis, peak) is
+    B = 8 pi n2 P / (lambda lambda_eff) * atan(d / 2 zR')      (-> 2 theta P/Pc for P << Pc).
+A pulse can only be matched at one power: the wings of the pulse see the linear mode. This
+time-dependent mismatch is precisely one of the couplings the 4D code resolves.
+"""
+import numpy as np
+from .grid import C0
+
+FS = 1e-15
+
+
+class MPC:
+    """Symmetric two-mirror cell.
+
+    R            : mirror radius of curvature [m]
+    d            : mirror separation [m]
+    medium       : Material
+    n_passes     : number of passes (centre -> mirror -> centre = 1 pass = 1 bounce)
+    mirror_gdd_fs2, mirror_tod_fs3 : dispersion per bounce
+    mirror_reflectivity : scalar, or callable R(lambda[m]) (intensity)
+    mirror_phase : optional callable phi(lambda[m]) [rad] per bounce (measured GDD curves, ...)
+    mirror_diameter : clear aperture [m] (soft super-gaussian edge), None = infinite
+    """
+
+    def __init__(self, R, d, medium, n_passes, *, mirror_gdd_fs2=0.0, mirror_tod_fs3=0.0,
+                 mirror_reflectivity=1.0, mirror_phase=None, mirror_diameter=None):
+        if not 0 < d < 2 * R:
+            raise ValueError("unstable cell: need 0 < d < 2R")
+        self.R, self.d, self.medium, self.n_passes = float(R), float(d), medium, int(n_passes)
+        self.gdd, self.tod = mirror_gdd_fs2, mirror_tod_fs3
+        self.refl, self.mphase = mirror_reflectivity, mirror_phase
+        self.diameter = mirror_diameter
+        self.g = 1 - d / R
+        self.theta = np.arccos(self.g)  # Gouy phase per pass
+        self._bounce_cache = None
+
+    @classmethod
+    def herriott(cls, R, N, k, medium, branch="concentric", n_passes=None, **kw):
+        """Re-entrant Herriott cell: N passes (twice the spots per mirror), d = R (1 +- cos(2*pi k/N)).
+        N must be even for the concentric branch. Default n_passes = N (one full re-entrant cycle);
+        pass n_passes to simulate any other number of passes in the same cell."""
+        value = np.cos(2 * np.pi * k / N)
+        if branch == "concentric":
+            if N % 2:
+                raise ValueError(
+                    f"N = {N} is odd: the concentric branch is only re-entrant for even N. N is the "
+                    f"number of passes of one re-entrant cycle (e.g. N = {2 * N}, k = {2 * k}); "
+                    f"to simulate a different number of passes, set n_passes separately.")
+            d = R * (1+ value)
+        elif branch == "planar":
+            d = R * (1 - value)
+        else:
+            raise ValueError(f"unknown branch {branch!r}: use 'concentric' or 'planar'")
+        return cls(R, d, medium, n_passes or N, **kw)
+
+    # ------------------------------------------------------------ modes
+    def linear_mode(self, lam):
+        n0 = self.medium.n(2 * np.pi * C0 / lam)
+        lm = lam / n0
+        w0 = np.sqrt(lm / (2 * np.pi) * np.sqrt(self.d * (2 * self.R - self.d)))
+        wm = np.sqrt(lm * self.R / np.pi * np.sqrt(self.d / (2 * self.R - self.d)))
+        return {"w0": w0, "w_mirror": wm, "zR": np.pi * w0 ** 2 / lm, "gouy_per_pass": self.theta}
+
+    def nonlinear_mode(self, P, lam=None, verbose=False):
+        """Kerr-lens matched eigenmode for power P [W]."""
+        lam = lam or self._lam0
+        Pc = self.medium.critical_power(lam)
+        r = P / Pc
+        if r >= 1:
+            raise ValueError(f"P/Pc = {r:.2f} >= 1: no matched mode (self-focusing)")
+        lin = self.linear_mode(lam)
+        s = (1 - r) ** 0.25
+        w0, wm = lin["w0"] * s, lin["w_mirror"] * s
+        lam_eff = lam * np.sqrt(1 - r)
+        zRp = np.pi * w0 ** 2 / lam_eff
+        B = 8 * np.pi * self.medium.n2 * P / (lam * lam_eff) * np.arctan(self.d / (2 * zRp))
+        out = {"P": P, "Pc": Pc, "P_over_Pc": r, "w0": w0, "w_mirror": wm,
+               "w0_linear": lin["w0"], "w_mirror_linear": lin["w_mirror"], "B_per_pass": B,
+               "B_total": B * self.n_passes,
+               "I_peak_center_W_cm2": 2 * P / (np.pi * w0 ** 2) * 1e-4}
+        if verbose:
+            print(self.report(out))
+        return out
+
+    def match(self, pulse, power="peak", verbose=True):
+        """Nonlinear mode matching from a Pulse (or a power in W). power: 'peak' | 'weighted'."""
+        self._lam0 = pulse.grid.lambda0 if hasattr(pulse, "grid") else self._lam0
+        if hasattr(pulse, "grid"):
+            P = pulse.peak_power() if power == "peak" else pulse.weighted_power()
+        else:
+            P = float(pulse)
+        return self.nonlinear_mode(P, self._lam0, verbose=verbose)
+
+    _lam0 = 1030e-9
+
+    def report(self, m=None):
+        s = (f"MPC: R = {self.R:.3f} m, d = {self.d:.4f} m, {self.n_passes} passes, "
+             f"Gouy/pass = {np.degrees(self.theta):.2f} deg, {self.medium}")
+        if m:
+            s += (f"\n  P = {m['P']/1e9:.3g} GW, Pc = {m['Pc']/1e9:.3g} GW, P/Pc = {m['P_over_Pc']:.3f}"
+                  f"\n  waist  {m['w0_linear']*1e6:.0f} um (linear) -> {m['w0']*1e6:.0f} um (matched)"
+                  f"\n  mirror {m['w_mirror_linear']*1e6:.0f} um (linear) -> {m['w_mirror']*1e6:.0f} um (matched)"
+                  f"\n  B/pass ~ {m['B_per_pass']:.3f} rad, total ~ {m['B_total']:.2f} rad, "
+                  f"I_peak(centre) ~ {m['I_peak_center_W_cm2']:.3g} W/cm^2")
+        return s
+
+    # ------------------------------------------------------------ mirror
+    def bounce_factor(self, grid):
+        """Complex transfer of one mirror in (x, y, Om) (cached)."""
+        if self._bounce_cache is not None and self._bounce_cache[0] is grid:
+            return self._bounce_cache[1]
+        W = grid.Omega * FS
+        spec = np.exp(1j * (self.gdd / 2 * W ** 2 + self.tod / 6 * W ** 3)).astype(complex)
+        lam = grid.lam
+        if callable(self.refl):
+            spec *= np.sqrt(np.clip(self.refl(np.where(grid.valid, lam, 1.0)), 0, 1))
+        else:
+            spec *= np.sqrt(self.refl)
+        if self.mphase is not None:
+            spec *= np.exp(1j * self.mphase(np.where(grid.valid, lam, 1.0)))
+        spec *= grid.valid
+        F = np.exp(-1j * (grid.r2[:, :, None] * (grid.omega / C0 / self.R)[None, None, :]))
+        F *= spec[None, None, :]
+        if self.diameter is not None:
+            ap = np.exp(-(np.sqrt(grid.r2) / (self.diameter / 2)) ** 30)
+            F *= ap[:, :, None]
+        F = grid.asarray(F.astype(grid.dtype))     # built on the host, stored on the grid's device
+        self._bounce_cache = (grid, F)
+        return F
+
+    def bounce(self, Aw, grid):
+        Aw *= self.bounce_factor(grid)
+        return Aw

src/diagnostics.py

@@ -0,0 +1,362 @@
+"""Diagnostics on a field A(x, y, Om) (FFT order). Outputs are returned SORTED (ascending axes).
+
+far field   = the plane where the field is (the cell centre = focus / waist plane when recorded
+              by MPCSimulation)
+near field  = the collimated beam: angular spectrum on a common angle grid theta (what you see after
+              collimating the output with a lens of focal f: x_out = f * theta). Computed with an
+              exact per-colour matrix Fourier transform, so a given theta means the same angle for
+              all colours.
+"""
+import numpy as np
+from .grid import C0
+from .backend import to_host
+
+FS = 1e-15
+
+
+# ------------------------------------------------------------------ basic quantities
+def spectral_support(grid, Aw, rel=1e-4):
+    Aw = to_host(Aw)
+    S = np.sum(np.abs(Aw) ** 2, axis=(0, 1))
+    return np.nonzero((S > rel * S.max()) & grid.valid)[0]
+
+
+def total_spectrum(grid, Aw, per_lambda=True):
+    """Spatially integrated spectrum, sorted by wavelength. Returns (lam [m], S)."""
+    Aw = to_host(Aw)
+    S = np.sum(np.abs(Aw) ** 2, axis=(0, 1)) * grid.dA * grid.dt / grid.Nt
+    return _to_lambda(grid, S, per_lambda)
+
+
+def _to_lambda(grid, S_w, per_lambda=True, axis=-1):
+    """Om-ordered array -> ascending-wavelength array (Jacobian 2 pi c / lam^2 if per_lambda)."""
+    idx = np.nonzero(grid.valid)[0]
+    idx = idx[np.argsort(grid.lam[idx])]
+    lam = grid.lam[idx]
+    S = np.take(S_w, idx, axis=axis)
+    if per_lambda:
+        shape = [1] * S.ndim
+        shape[axis] = -1
+        S = S * (2 * np.pi * C0 / lam ** 2).reshape(shape)
+    return lam, S
+
+
+def fluence(grid, Aw):
+    """F(x, y) [J/m^2], sorted."""
+    Aw = to_host(Aw)
+    F = np.sum(np.abs(Aw) ** 2, axis=2) * grid.dt / grid.Nt
+    return grid.shift_xy(F)
+
+
+def temporal_power(grid, Aw):
+    """P(t) [W], sorted time axis."""
+    A = grid.hfft.ift_t(to_host(Aw))
+    return grid.ts, grid.shift_w(np.sum(np.abs(A) ** 2, axis=(0, 1)) * grid.dA)
+
+
+# ------------------------------------------------------------------ near field (collimated beam)
+def near_field(grid, Aw, theta_max=None, n_theta=None, idx=None, theta_y=None):
+    """Near field (collimated beam): angular spectrum E(theta_x, theta_y, Om) on a common angle grid (sorted, odd size so that
+    theta = 0 is sampled). theta_max default: half the largest kx at lambda0.
+    theta_y: optional explicit theta_y values (e.g. [0.] for a single line).
+    Returns (theta_x, Ef) with Ef of shape (n_theta, len(theta_y), Nt), FFT order along Om."""
+    Aw = to_host(Aw)
+    n_theta = n_theta or (grid.Nx // 2) * 2 + 1
+    if theta_max is None:
+        theta_max = 0.5 * np.abs(grid.kx).max() / (grid.omega0 / C0)
+    th = np.linspace(-theta_max, theta_max, n_theta)
+    thy = th if theta_y is None else np.atleast_1d(np.asarray(theta_y, float))
+    idx = spectral_support(grid, Aw) if idx is None else idx
+    Ef = np.zeros((len(th), len(thy), grid.Nt), complex)
+    for j in idx:
+        k = grid.omega[j] / C0
+        Fx = np.exp(-1j * k * np.outer(th, grid.x)) * grid.dx * k / (2 * np.pi)
+        Fy = np.exp(-1j * k * np.outer(thy, grid.y)) * grid.dy * k / (2 * np.pi)
+        Ef[:, :, j] = Fx @ Aw[:, :, j] @ Fy.T
+    return th, Ef
+
+
+# ------------------------------------------------------------------ spectrograms
+def spectrogram_x_lambda(grid, Aw, y0=0.0, slit=None, axis="x", per_lambda=True):
+    """Spatio-spectral trace S(x, lambda) along a line (imaging spectrometer).
+    slit : width [m] integrated around y0 (None = single row). axis 'x' or 'y'.
+    Returns (x_sorted [m], lam [m], S[x, lam])."""
+    Aw = to_host(Aw)
+    if axis == "y":
+        Aw = np.swapaxes(Aw, 0, 1)
+        coord, other, d = grid.y, grid.x, grid.dx
+    else:
+        coord, other, d = grid.x, grid.y, grid.dy
+    if slit is None:
+        sel = [int(np.argmin(np.abs(other - y0)))]
+    else:
+        sel = np.nonzero(np.abs(other - y0) <= slit / 2)[0]
+    S = np.sum(np.abs(Aw[:, sel, :]).astype(np.float64) ** 2, axis=1)
+    lam, S = _to_lambda(grid, S, per_lambda, axis=1)
+    return np.fft.fftshift(coord), lam, np.fft.fftshift(S, axes=0)
+
+
+def spectrogram_theta_lambda(grid, Aw, theta_max=None, n_theta=None, per_lambda=True):
+    """Angle-resolved spectrum S(theta_x, lambda) at theta_y = 0 (output after collimation:
+    x_out = f * theta). Returns (theta [rad], lam [m], S[theta, lam])."""
+    th, Ef = near_field(grid, Aw, theta_max, n_theta, theta_y=[0.0])
+    S = np.abs(Ef[:, 0, :]).astype(np.float64) ** 2
+    lam, S = _to_lambda(grid, S, per_lambda, axis=1)
+    return th, lam, S
+
+
+# ------------------------------------------------------------------ homogeneity
+def homogeneity(grid, Aw=None, S_xyw=None, rel=1e-3):
+    """Spectral homogeneity V(x, y) = (sum sqrt(S S0))^2 / (sum S sum S0), S0 = total spectrum,
+    and its fluence-weighted mean. Pass either a field or an intensity cube S_xyw."""
+    S = (np.abs(to_host(Aw)) ** 2 if S_xyw is None else S_xyw).astype(np.float64)
+    S = S / S.max()
+    S0 = S.sum(axis=(0, 1))
+    num = np.sum(np.sqrt(S * S0[None, None, :]), axis=2) ** 2
+    den = S.sum(axis=2) * S0.sum()
+    F = S.sum(axis=2)
+    V = np.where(F > rel * F.max(), num / np.maximum(den, 1e-300), np.nan)
+    Vm = np.nansum(V * F) / np.sum(F[~np.isnan(V)])
+    return np.fft.fftshift(V, axes=(0, 1)) if S_xyw is None else V, float(Vm)
+
+
+# ------------------------------------------------------------------ modal content
+def _hermite_functions(x, w, M):
+    """Normalised 1D HG functions u_m(x) for m = 0..M, shape (M+1, len(x))."""
+    xi = np.sqrt(2) * x / w
+    psi = np.zeros((M + 1,) + xi.shape)
+    psi[0] = np.pi ** -0.25 * np.exp(-xi ** 2 / 2)
+    if M >= 1:
+        psi[1] = np.sqrt(2) * xi * psi[0]
+    for m in range(1, M):
+        psi[m + 1] = np.sqrt(2 / (m + 1)) * xi * psi[m] - np.sqrt(m / (m + 1)) * psi[m - 1]
+    return psi * np.sqrt(np.sqrt(2) / w)
+
+
+def mode_content(grid, Aw, w0_of_lambda, max_order=6, idx=None):
+    """Decomposition on Hermite-Gauss modes of waist w0(lambda) with flat phase (cell centre).
+    Returns dict: lam, eta00(lam) (fundamental fraction per colour), order_frac[N] (energy fraction
+    in modes of order m+n = N, spectrally integrated, last bin = everything above max_order)."""
+    Aw = to_host(Aw)
+    idx = spectral_support(grid, Aw, 1e-3) if idx is None else idx
+    lam = grid.lam[idx]
+    M = max_order
+    P = np.sum(np.abs(Aw[:, :, idx]) ** 2, axis=(0, 1)) * grid.dA
+    order_E = np.zeros(M + 2)
+    eta = np.zeros(len(idx))
+    for n, j in enumerate(idx):
+        w = w0_of_lambda(lam[n])
+        ux = _hermite_functions(grid.x, w, M)
+        uy = _hermite_functions(grid.y, w, M)
+        c = ux @ Aw[:, :, j] @ uy.T * grid.dA          # (M+1, M+1)
+        p = np.abs(c) ** 2
+        eta[n] = p[0, 0] / P[n]
+        mm, nn = np.meshgrid(np.arange(M + 1), np.arange(M + 1), indexing="ij")
+        for N in range(M + 1):
+            order_E[N] += p[mm + nn == N].sum()
+        order_E[M + 1] += max(P[n] - p[mm + nn <= M].sum(), 0.0)
+    o = np.argsort(lam)
+    return {"lam": lam[o], "eta00": eta[o], "order_frac": order_E / P.sum(),
+            "eta00_total": float(np.sum(eta * P) / P.sum())}
+
+
+def lg_mode_content(grid, Aw, w0_of_lambda, p_max=2, l_max=2, idx=None):
+    """Decomposition on Laguerre-Gauss modes LG_pl (exp(i l phi) convention, waist w0(lambda), flat
+    phase, cell centre), spectrally integrated.
+    Returns dict: p = 0..p_max, l = -l_max..l_max, frac[p, l] (fraction of the total energy in each
+    mode, shape (p_max+1, 2*l_max+1)) and residual (energy outside the listed modes)."""
+    from .pulse import laguerre_gauss
+    Aw = to_host(Aw)
+    idx = spectral_support(grid, Aw, 1e-3) if idx is None else idx
+    lam = grid.lam[idx]
+    ps, ls = np.arange(p_max + 1), np.arange(-l_max, l_max + 1)
+    P = np.sum(np.abs(Aw[:, :, idx]) ** 2, axis=(0, 1)) * grid.dA
+    E = np.zeros((len(ps), len(ls)))
+    for n, j in enumerate(idx):
+        w = w0_of_lambda(lam[n])
+        for a, p in enumerate(ps):
+            for b, l in enumerate(ls):
+                u = laguerre_gauss(w, int(p), int(l))(grid.X, grid.Y)
+                E[a, b] += np.abs(np.sum(np.conj(u) * Aw[:, :, j]) * grid.dA) ** 2
+    frac = E / P.sum()
+    return {"p": ps, "l": ls, "frac": frac, "residual": max(1.0 - float(frac.sum()), 0.0)}
+
+
+# ------------------------------------------------------------------ beam size / M2
+def beam_moments(grid, Aw, idx=None):
+    """Per colour: second-moment diameters/2 (wx, wy) [m] and M2x, M2y. Sorted by wavelength."""
+    Aw = to_host(Aw)
+    idx = spectral_support(grid, Aw, 1e-3) if idx is None else idx
+    E = Aw[:, :, idx].astype(np.complex128)
+    E /= np.abs(E).max()
+    I = np.abs(E) ** 2
+    P = I.sum(axis=(0, 1))
+    out = {}
+    for ax, coord, kc in ((0, grid.x, grid.kx), (1, grid.y, grid.ky)):
+        sh = [1, 1, 1]
+        sh[ax] = -1
+        c = coord.reshape(sh)
+        kk = kc.reshape(sh)
+        xm = (I * c).sum(axis=(0, 1)) / P
+        x2 = (I * c ** 2).sum(axis=(0, 1)) / P - xm ** 2
+        Ek = np.fft.fft(E, axis=ax)
+        Ik = np.abs(Ek) ** 2
+        Pk = Ik.sum(axis=(0, 1))
+        km = (Ik * kk).sum(axis=(0, 1)) / Pk
+        k2 = (Ik * kk ** 2).sum(axis=(0, 1)) / Pk - km ** 2
+        dE = np.fft.ifft(1j * kk * Ek, axis=ax)
+        xk = (c * np.imag(np.conj(E) * dE)).sum(axis=(0, 1)) / P - xm * km
+        name = "x" if ax == 0 else "y"
+        out["w" + name] = 2 * np.sqrt(np.maximum(x2, 0))
+        out["M2" + name] = 2 * np.sqrt(np.maximum(x2 * k2 - xk ** 2, 0))
+    lam = grid.lam[idx]
+    o = np.argsort(lam)
+    out = {k: v[o] for k, v in out.items()}
+    out["lam"] = lam[o]
+    wts = P[o] / P.sum()
+    out["M2x_mean"] = float(np.sum(out["M2x"] * wts))
+    out["M2y_mean"] = float(np.sum(out["M2y"] * wts))
+    return out
+
+
+# ------------------------------------------------------------------ compression
+def _fwhm(t, y):
+    y = y / y.max()
+    above = np.nonzero(y >= 0.5)[0]
+    if len(above) < 2:
+        return np.nan
+    i0, i1 = above[0], above[-1]
+    tl = np.interp(0.5, [y[i0 - 1], y[i0]], [t[i0 - 1], t[i0]]) if i0 > 0 else t[i0]
+    tr = np.interp(0.5, [y[i1 + 1], y[i1]], [t[i1 + 1], t[i1]]) if i1 < len(t) - 1 else t[i1]
+    return tr - tl
+
+
+def compress(grid, Aw, gdd_bounds_fs2=(-30000, 30000), tod_fs3=0.0, n_scan=61, rel=1e-4):
+    """Best uniform GDD (e.g. chirped mirrors) maximising the peak of the spatially integrated
+    power: coarse scan over gdd_bounds_fs2 then golden-section refinement.
+    Returns dict with gdd_fs2, t_fs, P(t) [W] (compressed), FWHM, transform-limited FWHM,
+    peak power and peak power of the (spatially resolved) transform limit."""
+    Aw = to_host(Aw)
+    F = np.sum(np.abs(Aw) ** 2, axis=2)
+    mask = F > rel * F.max()
+    E = Aw[mask]                                   # (npix, Nt)
+    W = grid.Omega * FS
+
+    def P_of(gdd):
+        ph = np.exp(1j * (gdd / 2 * W ** 2 + tod_fs3 / 6 * W ** 3)).astype(E.dtype)
+        A = grid.hfft.ift_t(E * ph[None, :])
+        return np.sum(np.abs(A).astype(np.float64) ** 2, axis=0) * grid.dA
+
+    scan = np.linspace(*gdd_bounds_fs2, n_scan)
+    pk = [P_of(g).max() for g in scan]
+    i = int(np.argmax(pk))
+    a, b = scan[max(i - 1, 0)], scan[min(i + 1, n_scan - 1)]
+    gr = (np.sqrt(5) - 1) / 2
+    c, d = b - gr * (b - a), a + gr * (b - a)
+    fc, fd = P_of(c).max(), P_of(d).max()
+    for _ in range(30):
+        if fc > fd:
+            b, d, fd = d, c, fc
+            c = b - gr * (b - a)
+            fc = P_of(c).max()
+        else:
+            a, c, fc = c, d, fd
+            d = a + gr * (b - a)
+            fd = P_of(d).max()
+    g = (a + b) / 2
+    P = np.fft.fftshift(P_of(g))
+    t = grid.ts / FS
+    Etl = np.abs(E).astype(np.float64)             # flat phase in every pixel
+    Ptl = np.fft.fftshift(np.sum(np.abs(grid.hfft.ift_t(Etl)) ** 2, axis=0) * grid.dA)
+    return {"gdd_fs2": g, "t_fs": t, "P": P, "fwhm_fs": _fwhm(t, P), "tl_fwhm_fs": _fwhm(t, Ptl),
+            "peak_power": float(P.max()), "peak_power_TL": float(Ptl.max()),
+            "gdd_scan_fs2": scan, "peak_scan": np.array(pk)}
+
+
+def local_gdd_map(grid, Aw, rel_fluence=1e-2, rel_spectrum=1e-2):
+    """Pixel-wise GDD [fs^2] from a weighted quadratic fit of the spectral phase
+    (radial chirp map). Returns sorted GDD(x, y) with NaN outside the beam."""
+    Aw = to_host(Aw)
+    F = np.sum(np.abs(Aw) ** 2, axis=2)
+    S0 = np.sum(np.abs(Aw) ** 2, axis=(0, 1))
+    idx = np.nonzero((S0 > rel_spectrum * S0.max()) & grid.valid)[0]
+    idx = idx[np.argsort(grid.Omega[idx])]
+    W = grid.Omega[idx] * FS
+    mask = F > rel_fluence * F.max()
+    E = Aw[mask][:, idx].astype(np.complex128)
+    E /= np.abs(E).max()
+    ph = np.unwrap(np.angle(E), axis=1)
+    wt = np.abs(E) ** 2
+    V = np.stack([np.ones_like(W), W, W ** 2], axis=1)            # (n, 3)
+    G = np.einsum("pn,ni,nj->pij", wt, V, V)
+    r = np.einsum("pn,ni,pn->pi", wt, V, ph)
+    coef = np.linalg.solve(G, r[..., None])[..., 0]
+    out = np.full(F.shape, np.nan)
+    out[mask] = 2 * coef[:, 2]
+    return grid.shift_xy(out)
+
+
+def radial_profile(grid, M2d, nbins=None):
+    """Azimuthal average of a sorted 2D map, returns (r, value)."""
+    X, Y = np.meshgrid(grid.xs, grid.ys, indexing="ij")
+    r = np.sqrt(X ** 2 + Y ** 2)
+    nb = nbins or grid.Nx // 2
+    edges = np.linspace(0, r.max() / np.sqrt(2), nb + 1)
+    ok = ~np.isnan(M2d)
+    s, _ = np.histogram(r[ok], edges, weights=M2d[ok])
+    n, _ = np.histogram(r[ok], edges)
+    with np.errstate(invalid="ignore"):
+        return 0.5 * (edges[1:] + edges[:-1]), s / n
+
+
+# ------------------------------------------------------------------ wavelength-resolved profiles
+def band_indices(grid, lam_c, bandwidth=0.0):
+    """Om indices inside [lam_c - bw/2, lam_c + bw/2]; bandwidth 0 -> nearest bin only
+    (the grid spectral resolution is lambda^2 / (c T))."""
+    if bandwidth > 0:
+        idx = np.nonzero(grid.valid & (np.abs(grid.lam - lam_c) <= bandwidth / 2))[0]
+        if len(idx):
+            return idx
+    return np.array([int(np.argmin(np.where(grid.valid, np.abs(grid.lam - lam_c), np.inf)))])
+
+
+def profiles_1d(grid, Aw, wavelengths, bandwidth=0.0, axis="x", mode="cut"):
+    """1D beam profiles for selected wavelengths (bandpass of width `bandwidth`) + the
+    spectrally integrated ("global") profile, in the plane where the field is (the far field
+    when it is recorded at the cell centre).
+
+    mode 'cut'        : line through y = 0 (x = 0 if axis='y')
+    mode 'integrated' : projection on the axis (like a 1D camera / slit-less lineout)
+    Returns (coord_sorted [m], profiles[n_lambda + 1, N] (last row = global, not normalised),
+             band_energy[n_lambda + 1] (J, fraction of the pulse in each band; last = total)).
+    """
+    Aw = to_host(Aw)
+    A = Aw if axis == "x" else np.swapaxes(Aw, 0, 1)
+    coord = grid.x if axis == "x" else grid.y
+    d_other = grid.dy if axis == "x" else grid.dx
+
+    def lineout(idx):
+        sub = np.abs(A[:, :, idx]).astype(np.float64) ** 2
+        s = sub.sum(axis=2)
+        return s[:, 0] if mode == "cut" else s.sum(axis=1) * d_other
+
+    rows, en = [], []
+    norm = grid.dA * grid.dt / grid.Nt
+    for lc in wavelengths:
+        idx = band_indices(grid, lc, bandwidth)
+        rows.append(lineout(idx))
+        en.append(float(np.sum(np.abs(A[:, :, idx]).astype(np.float64) ** 2)) * norm)
+    rows.append(lineout(np.nonzero(grid.valid)[0]))
+    en.append(grid.energy(Aw))
+    return np.fft.fftshift(coord), np.fft.fftshift(np.array(rows), axes=1), np.array(en)
+
+
+def near_profiles_1d(grid, Aw, wavelengths, bandwidth=0.0, theta_max=None, n_theta=None):
+    """Same as profiles_1d (cut through theta_y = 0) in the near field (collimated beam: angle,
+    common for all colours). Returns (theta [rad], profiles[n_lambda + 1, n_theta])."""
+    idx_all = spectral_support(grid, Aw)
+    th, Ef = near_field(grid, Aw, theta_max, n_theta, idx=idx_all, theta_y=[0.0])
+    S = np.abs(Ef[:, 0, :]) ** 2
+    rows = [S[:, band_indices(grid, lc, bandwidth)].sum(axis=1) for lc in wavelengths]
+    rows.append(S.sum(axis=1))
+    return th, np.array(rows)