JAX backend#

The differentiable forward model, metrics and pulse parameterisation. These share the core conventions with the NumPy ForwardModel, croak.metrics and ArrayPulse, and back the JAX-native solvers — LBFGSAD, LM, LBFGSHand, COPRAJax, OptxLBFGS, OptxLM and CMAES.

This is the general gradient route: it supplies the derivatives for geometric smearing and the fitted extras that the closed-form adjoint does not cover, the Jacobians Levenberg–Marquardt and the Laplace covariance need, derivatives for interactions implemented in the JAX forward path, and on-device execution of a whole loop or population. It also serves as the independent oracle the analytic adjoints are tested against. See Gradients.

JAX forward model#

JAX/autodiff twin of croak.forward.

This module re-implements the nonlinear-process forward model of croak.forward.ForwardModel in JAX, so that gradients and Jacobians of the trace (and of any error metric built on it) can be obtained by automatic differentiation. This is the general gradient route: it reaches model features the hand-derived Wirtinger adjoints in croak.forward do not cover (geometric smearing, fitted thickness/delay-zero/smearing width) and supplies the Jacobians that Levenberg–Marquardt and croak.covariance need. It backs the LBFGSAD, LM, OptxLM and CMAES retrievers, and serves as an independent oracle in the tests that cross-check the analytic adjoints.

The core local/slab map follows croak.forward.ForwardModel.signal_single(). Additional paths implement chromatic focusing and coherent collection. The module

  • enables 64-bit JAX (jax_enable_x64) at import — single precision would not match the ~1e-10 agreement the tests assert;

  • works internally in DFT-bin order (ifftshift on input, fftshift on output), using the unnormalised forward DFT for spectrum→time and the 1/N inverse DFT for time→spectrum, exactly as croak.grid.raw_fft() / _ifft0;

  • treats the slab propagation constant \(\beta(\omega)\) as a static constant: it does not depend on the optimisation variables, so it is built once with the numpy croak.materials.beta() and captured (no need to evaluate the refractive-index models in JAX).

See croak.forward.ForwardModel for the physical model and notation.

croak.forward_jax.make_signal_fn(omega, delays, interaction, *, material=None, thickness=0.0, npoints=1, omega0=0.0, quadrature='gausslegendre')[source]#

Build the single-delay signal map \(\tilde E \mapsto \psi(\omega)\).

JAX twin of croak.forward.ForwardModel.signal_single(). The returned function signal(ew, tau) takes a centred complex spectrum ew and a delay tau and returns the centred frequency-domain signal psi. It is differentiable and vmap/jit-friendly; the static problem geometry (grid, delays, interaction, dispersion) is captured in the closure.

Parameters:

  • omega (array_like) – Centred angular-frequency grid (rad/s), relative to omega0.

  • delays (array_like) – Delay axis (s); only its length/values inform the closure, tau is a runtime argument (kept for signature parity with the numpy model).

  • interaction (str) – "shg", "sd" or "pg".

  • material (str | None) – Dispersive-slab parameters; see croak.forward.ForwardModel.

  • thickness (float) – Dispersive-slab parameters; see croak.forward.ForwardModel.

  • npoints (int) – Dispersive-slab parameters; see croak.forward.ForwardModel.

  • omega0 (float) – Dispersive-slab parameters; see croak.forward.ForwardModel.

  • quadrature (str) – Dispersive-slab parameters; see croak.forward.ForwardModel.

Returns:

callable – signal(ew, tau) -> psi (centred order).

Return type:

Callable[[Array, Array | float], Array]

croak.forward_jax.make_trace_fn(omega, delays, interaction, *, material=None, thickness=0.0, npoints=1, omega0=0.0, quadrature='gausslegendre', smearing=None, focal=None, normalize=False)[source]#

Build the full-trace map \(\tilde E \mapsto T(\omega,\tau)\).

JAX twin of croak.forward.ForwardModel.trace(). The returned trace(ew) maps a centred spectrum to an (Nomega, Ndelay) intensity trace \(|\psi|^2\), vectorising the single-delay signal over delays with jax.vmap(). It is differentiable and jit-friendly.

Parameters:

  • smearing (SmearingKernel or None, optional) – Geometric time-smearing kernel (see croak.smearing). Because the delay offset acts across the whole delay axis, the smeared model has no single-delay form; this simply delegates to make_param_trace_fn() with the slab thickness and delay zero held fixed. The same applies to focal.

  • focal (FocalMixture or None, optional) – Chromatic focal mixture, optionally carrying a collection aperture. As for make_param_trace_fn(), to which this delegates when either it or smearing is set.

  • normalize (bool, optional) – Scale the trace to unit maximum (default False; retrieval works with the unnormalised trace and absorbs the scale through mu).

  • parameters (other) – As make_signal_fn().

  • omega (ArrayLike)

  • delays (ArrayLike)

  • interaction (str)

  • material (str | None)

  • thickness (float)

  • npoints (int)

  • omega0 (float)

  • quadrature (str)

Returns:

callable – trace(ew) -> T of shape (Nomega, Ndelay).

Return type:

Callable[[Array], Array]

croak.forward_jax.make_param_trace_fn(omega, delays, interaction, *, depth_weight=None, depth_node_phase=None, material=None, thickness=0.0, npoints=1, omega0=0.0, quadrature='gausslegendre', smearing=None, focal=None, depth_transverse=False, fit_thickness=False, fit_tau0=False, fit_smearing=False, normalize=False)[source]#

Build a trace map that also takes the slab thickness and a delay-zero shift.

Like make_trace_fn(), but the returned trace(ew, thickness, tau0) is differentiable in the slab thickness L and a delay-zero offset tau0 as well as in the spectrum, so the AD retrievers (LBFGSAD, LM) and the covariance estimator can fit them alongside the pulse. See croak._jax_pulse.augment() for the parameter packing.

The delay-zero offset enters as a shift of the gate’s delay phase, \(e^{i\omega(\tau-\tau_0)}\) — it slides the simulated trace along the delay axis (for SHG, its symmetry centre), correcting a mis-set time zero.

When fit_thickness is set the depth quadrature uses normalised Gauss–Legendre nodes \(\xi\in[0,1]\) and weights \(\hat w\) (\(\sum\hat w=1\)) so that, at call time, the depth nodes are \(z_q=L\,\xi_q\), the weights are \(L\,\hat w_q\), and the propagation phases \(e^{i\beta L\xi_q}\), \(e^{i\beta L(1-\xi_q)}\) are differentiable in L. When it is not set, the nodes/weights are baked once from quadrature_nodes_weights() (so the thin-medium single-node case material=None is preserved for a tau0-only fit).

Parameters:

  • fit_thickness (bool, optional) – Make the trace differentiable in thickness. Requires a dispersive slab (material set and thickness > 0); raises otherwise.

  • fit_tau0 (bool, optional) – Make the trace differentiable in the delay-zero offset tau0. Marker only — tau0 is always a runtime argument; the flag documents intent and parallels fit_thickness.

  • smearing (SmearingKernel or None, optional) – Geometric time-smearing kernel (see croak.smearing). When set, the returned trace also accepts a fourth argument smear_scale (default 1.0) that multiplies the kernel widths, so the smearing strength — equivalently the mask ratio d/D — can be fitted. A scalar scales both channels together; a length-2 vector (scale_p, scale_delta) scales the gate-shape (p) and delay (delta) channels independently (see _smear_scale_pair). None (default) leaves the model untouched.

  • focal (FocalMixture or None, optional) –

    Chromatic focal mixture (croak.focal) in place of smearing: the focal average is carried explicitly, each node applying its own \(A(r,\omega)\) and its own definite (p, theta) rather than a Gaussian distribution over them. Mutually exclusive with smearing.

    If the mixture carries a collection (croak.collection.CollectionAperture) the per-node fields are transformed to transverse wavevector and integrated over the aperture instead of summed incoherently — the difference between an instrument that collects the whole signal beam and one that collects it through a hole. PG (TG) only; see croak.collection.

    If the mixture carries an evolve spec (fc-z, croak.focal.evolved_arm_filters()), each of the three arms at each depth node applies its own linearly evolved complex profile at the arm’s walked radius, in place of the single shared entrance-face filter. The spec’s slab parameters must match this function’s material/thickness/npoints/quadrature (verified; the filters must sit on the same depth nodes as the propagation phases). PG only; fitting thickness with evolved profiles is unsupported because the table is built for fixed depth nodes. Cost is close to the entrance mixture’s: the depth axis is already batched, only the build-time table is new.

  • fit_smearing (bool, optional) – Make the trace differentiable in the smearing scale. Marker only, for the same reason as fit_tau0; it does require smearing to be set.

  • depth_weight (array_like or None, optional) – Complex generation-envelope weights, one per quadrature node, multiplying the depth-quadrature weights: models the transverse geometry a 1D depth integral cannot see (focal-spot evolution across the slab, beamlet walk-through, Gouy phase), which enters as a smooth per-depth factor \(g(z_q)\) on the locally generated signal. The overall scale is irrelevant (absorbed by \(\mu\)); only the shape matters. Requires a dispersive slab and length npoints. With fit_thickness the envelope is pinned to the normalised nodes (it rescales with the fitted slab rather than staying fixed in absolute \(z\)). None (default) is the uniform-generation model. Negligible for slabs much thinner than the beams’ Rayleigh range.

  • depth_node_phase (array_like or None, optional) – Complex multiplicative factors of shape (focal.nodes, npoints) applied to each focal/depth contribution. Requires focal. Use for explicit model diagnostics; these factors are not fitted automatically.

  • depth_transverse (bool, optional) –

    Model the transverse decoherence of the depth integral. croak’s depth quadrature normally adds the signal generated at every depth with the same transverse phase — fully coherently in transverse wavevector. Real propagation dephases the angular spectrum: signal born at depth \(z_q\) accumulates, over the remaining slab, the extra phase

    \[\phi_{jq}(\omega) = \bigl[k_z(\omega, k_{\perp j}) - k_z(\omega, 0)\bigr]\, (L - z_q),\]

    applied per aperture node \(j\) inside the collection transform, before the coherent depth sum (see _depth_phase_rate for \(k_{\perp j}\) and the exact \(k_z\) used). Parameter-free and opt-in; the default False omits this transverse exit phase. Requires focal with a collection and a dispersive slab. For a full-collection variant use a collection window much wider than the signal footprint (a large chromatic hole, or a wide chromatic=False k window) — there is deliberately no separate path. Costs one aperture-transform contraction per depth node instead of one total; with depth_weight and depth_node_phase it composes (both keep their existing meaning). See Modelling the collection aperture for when the missing physics matters (it grows with slab thickness).

  • parameters (other) – As make_trace_fn().

  • omega (ArrayLike)

  • delays (ArrayLike)

  • interaction (str)

  • material (str | None)

  • thickness (float)

  • npoints (int)

  • omega0 (float)

  • quadrature (str)

  • normalize (bool)

Returns:

callable – trace(ew, thickness, tau0, smear_scale=1.0) -> T of shape (Nomega, Ndelay). thickness/tau0/smear_scale may be Python floats (held fixed) or JAX scalars (differentiated); smear_scale may also be a length-2 vector for split channel scaling.

Return type:

Callable[[…], Array]

croak.forward_jax.make_forward_adjoint_fns(omega, delays, interaction, *, material=None, thickness=0.0, npoints=1, omega0=0.0, quadrature='gausslegendre')[source]#

Build the matched JAX forward (with tape) and hand-derived adjoint.

JAX twin of the forward/adjoint pair on croak.forward.ForwardModel (signal_single() + adjoint_single()). Both share one set of static constants and are vmap/jit-friendly, so a caller can backpropagate a trace cotangent without automatic differentiation — the basis of LBFGSHand and COPRAJax.

Returns:

  • signal_tape (callable) – signal_tape(ew, tau) -> (psi, test, gate) (centred psi; test / gate are the time-domain fields the adjoint reuses).

  • adjoint (callable) – adjoint(psi_bar, test, gate, tau) -> ew_bar (centred), the Wirtinger adjoint mapping \(\partial L/\partial\psi^*\) to \(\partial L/\partial\tilde E^*\).

Parameters:

  • omega (ArrayLike)

  • delays (ArrayLike)

  • interaction (str)

  • material (str | None)

  • thickness (float)

  • npoints (int)

  • omega0 (float)

  • quadrature (str)

Return type:

tuple[Callable[[Array, Array | float], tuple[Array, Array, Array]], Callable[[Array, Array, Array, Array | float], Array]]

croak.forward_jax.maketrace_jax(omega, delays, ew, interaction, *, material=None, thickness=0.0, npoints=1, omega0=0.0, quadrature='gausslegendre', smearing=None, normalize=True)[source]#

Generate a (normalised) trace with the JAX model, returned as numpy.

Convenience wrapper around make_trace_fn() paralleling croak.forward.maketrace(), for tests and examples.

Parameters:

  • omega (ArrayLike)

  • delays (ArrayLike)

  • ew (NDArray[complex128])

  • interaction (str)

  • material (str | None)

  • thickness (float)

  • npoints (int)

  • omega0 (float)

  • quadrature (str)

  • smearing (SmearingKernel | None)

  • normalize (bool)

Return type:

NDArray[float64]

JAX metrics#

JAX/autodiff twin of croak.metrics plus the solver objective pieces.

These are JAX clones of the trace-error metrics in croak.metrics, written so that the FROG error (and the least-squares residual vector) is a differentiable function of the pulse parameters. They back LBFGSAD (which minimises the scalar frog_error()) and LM (which drives residual_vector() to zero), and are checked against the numpy metrics in the tests.

Conventions match croak.metrics and croak.lbfgs:

  • t_meas / t_sim are (Nomega, Ndelay) intensity traces;

  • weights is the per-frequency weight vector of length Nomega (broadcast along the delay axis), as passed by the solvers;

  • the scalar error is \(R=\sqrt{\sum\mathrm{resid}^2/\mathrm{denom}}\) with \(\mathrm{denom}=MN\,\max(T_\mathrm{meas}w)^2\) — identical to croak.metrics.compute_R() and to the value minimised in croak.lbfgs.LBFGS.

The second-difference smoothness penalties mirror croak.pulses._second_diff_penalty so regularised retrieval matches the analytic-gradient solver.

croak.metrics_jax.mu_global(t_meas, t_sim, weights)[source]#

Compute the optimal global trace scale factor with JAX.

Mirrors croak.metrics.compute_mu(). weights is a per-frequency vector broadcast along the delay axis.

Parameters:

  • t_meas (Array)

  • t_sim (Array)

  • weights (Array)

Return type:

Array

croak.metrics_jax.mu_per_freq(t_meas, t_sim, weights)[source]#

Compute the optimal scale factor for each frequency row with JAX.

Mirrors croak.metrics.compute_mu_per_freq() with freq_axis=0 (sum over delays). As there, the weights cancel between numerator and denominator — the factors are weight-independent, and a zero-weight row falls back to 1.0. Returns one factor per frequency row.

Parameters:

  • t_meas (Array)

  • t_sim (Array)

  • weights (Array)

Return type:

Array

croak.metrics_jax.frog_error(t_meas, t_sim, weights, *, R_omega=False)[source]#

Normalised FROG error \(R\) minimised by the gradient solvers.

Identical to the value computed in croak.lbfgs.LBFGS (and to croak.metrics.frog_error()): the optimal scale mu is applied, the weighted residual is formed, and \(R=\sqrt{\sum\mathrm{resid}^2/ \mathrm{denom}}\) is returned.

Parameters:

  • R_omega (bool, optional) – Use per-frequency adaptive scaling factors instead of a single global mu.

  • t_meas (Array)

  • t_sim (Array)

  • weights (Array)

Return type:

Array

croak.metrics_jax.residual_vector(t_meas, t_sim, weights, *, R_omega=False)[source]#

Flattened least-squares residual for LM.

Returns \((T_\mathrm{meas}-\mu T_\mathrm{sim})\,w/\sqrt{\mathrm{denom}}\) flattened to 1-D, so that its Euclidean norm equals the FROG error frog_error(). The least-squares solvers minimise its squared norm, with any penalty rows appended separately.

Parameters:

  • t_meas (Array)

  • t_sim (Array)

  • weights (Array)

  • R_omega (bool)

Return type:

Array

croak.metrics_jax.amplitude_penalty(amp_params)[source]#

Amplitude smoothness penalty (max-normalised), as in ArrayPulse.

Parameters:

amp_params (Array)

Return type:

Array

croak.metrics_jax.phase_penalty(phase_params)[source]#

Phase smoothness penalty (mean squared curvature), as in ArrayPulse.

Parameters:

phase_params (Array)

Return type:

Array

croak.metrics_jax.spectral_penalty(amp_params, target)[source]#

Relative spectral-amplitude mismatch ‖a - target‖² / ‖target‖².

JAX twin of croak.pulses.ArrayPulse.spectral_penalty(). target is the target spectral amplitude over the optimised support (same length as amp_params).

Parameters:

  • amp_params (Array)

  • target (Array)

Return type:

Array

JAX pulse parameterisation#

The differentiable u → spectrum map shared by the JAX solvers, the reduced cubic-B-spline phase basis, and the smoothness / spectral / temporal penalties.

Shared parameterisation helper for the JAX retrievers.

Both LBFGSAD and LM optimise the same real parameter vector as croak.lbfgs.LBFGS: the spectral amplitude and phase over an optional support window (u = [A; phi]), or just the phase when phase_only holds the amplitude fixed. This module provides the single differentiable map from those parameters to a complex spectrum, mirroring croak.pulses.ArrayPulse.spectrum()/set_params so the AD solvers share the analytic solver’s exact parameterisation (and are directly comparable to it).

croak._jax_pulse.make_spectrum_fn(amplitude, phase, *, support=None, phase_only=False)[source]#

Build the differentiable map u -> spectrum and its initial parameters.

JAX twin of croak.pulses.ArrayPulse. Parameters outside the support window keep the seed amplitude/phase and are not optimised.

Parameters:

  • amplitude (array_like) – Seed spectral amplitude \(|\tilde E|\) and phase (length N).

  • phase (array_like) – Seed spectral amplitude \(|\tilde E|\) and phase (length N).

  • support (tuple of int, optional) – Half-open (start, stop) index range of the optimised window; the full grid by default.

  • phase_only (bool, optional) – Optimise only the phase, holding amplitude fixed.

Returns:

  • spectrum_fn (callable) – spectrum_fn(u) -> ew returning the centred complex spectrum \(A\,e^{i\phi}\).

  • u0 (numpy.ndarray) – Initial parameter vector ([A[support]; phi[support]] or, for phase_only, phi[support]).

  • nidcs (int) – Number of optimised frequency indices — the split point between the amplitude and phase blocks of u (for the regularisation penalties).

Return type:

tuple[Callable[[Array], Array], NDArray[float64], int]

croak._jax_pulse.make_param_grad_fns(amplitude, phase, *, support=None, phase_only=False)[source]#

Build the hand-derived parameter-space adjoints for u.

Same u as make_spectrum_fn(). JAX twin of the adjoint side of croak.pulses.ArrayPulse (ew_vjp() and the *_penalty_grad methods). The parameterisation (support window, phase_only) matches make_spectrum_fn() exactly so the two are used together.

Parameters:

  • amplitude (ArrayLike)

  • phase (ArrayLike)

  • support (tuple[int, int] | None)

  • phase_only (bool)

Return type:

ParamGradFns

class croak._jax_pulse.ParamGradFns(ew_vjp, amplitude_penalty_grad, phase_penalty_grad, spectral_penalty_grad)[source]#

Bases: NamedTuple

Hand-derived parameter-space adjoints, the JAX twin of ArrayPulse.

These mirror croak.pulses.ArrayPulse.ew_vjp() and the regularisation gradients, expressed in jnp (no autodiff). They let the non-AD retriever map a spectral cotangent ew_bar and the parameters u to dL/du.

Parameters:

  • ew_vjp (Callable[[Array, Array], Array])

  • amplitude_penalty_grad (Callable[[Array], Array])

  • phase_penalty_grad (Callable[[Array], Array])

  • spectral_penalty_grad (Callable[[Array, Array], Array])

ew_vjp: Callable[[Array, Array], Array]#

spectral cotangent → parameter gradient.

Type:

ew_vjp(u, ew_bar) -> du

amplitude_penalty_grad: Callable[[Array], Array]#

amplitude_penalty_grad(u) -> du (full-length; phase block zero).

phase_penalty_grad: Callable[[Array], Array]#

phase_penalty_grad(u) -> du (full-length; amplitude block zero).

spectral_penalty_grad: Callable[[Array, Array], Array]#

spectral_penalty_grad(u, target_support) -> du (amplitude block; zero in phase-only mode). target_support is the target amplitude over the optimised support (length nidcs).

croak._jax_pulse.make_spline_phase_basis(amplitude, phase, omega, *, n_nodes, degree=3, support=None, support_threshold=0.001)[source]#

Phase-only basis whose phase is a cubic B-spline over the spectrum support.

The spectral phase on the support window is a fixed linear map \(\phi = B\,c\) of n_nodes control points c (the optimised parameters), with the amplitude held at the seed amplitude. Because B (the scipy.interpolate.BSpline.design_matrix()) is a constant, the map is exactly \(C^2\) (cubic) and trivially AD-differentiable. Knots are clamped over the non-zero-spectrum window only, so nodes don’t waste degrees of freedom where there is no signal.

Returns (spectrum_fn, c0, phase_penalty, support) where c0 fits the seed phase and phase_penalty(c) is the P-spline roughness penalty \(\mathrm{mean}((D_2 c)^2)\) (a 2nd-difference on the coefficients).

Parameters:

  • amplitude (ArrayLike)

  • phase (ArrayLike)

  • omega (ArrayLike)

  • n_nodes (int)

  • degree (int)

  • support (tuple[int, int] | None)

  • support_threshold (float)

Return type:

tuple[Callable[[Array], Array], NDArray[float64], Callable, tuple[int, int]]

croak._jax_pulse.build_parameterisation(ew0, omega, *, phase_only, phase_basis='pointwise', n_nodes=20, support_threshold=0.001, spectral_target=None, time_mask=None)[source]#

Build the u -> spectrum map shared by the AD retrievers.

phase_basis="pointwise" is the existing per-frequency parameterisation (amplitude+phase, or phase-only); "bspline" is the reduced cubic B-spline phase basis (implies phase-only). The returned Penalties already handle slicing/mode so each solver just weights them by its reg_*.

time_mask (length n, on the centred grid.t axis) enables the temporal penalty: the energy fraction of E(t) = ifft(spectrum) falling where the mask is non-zero — a self-consistent, mode-agnostic alternative to the post-retrieval temporal filter. See time_window_mask().

Parameters:

  • ew0 (NDArray[complex128])

  • omega (ArrayLike)

  • phase_only (bool)

  • phase_basis (str)

  • n_nodes (int)

  • support_threshold (float)

  • spectral_target (NDArray[float64] | None)

  • time_mask (NDArray[float64] | None)

Return type:

Parameterisation

croak._jax_pulse.ls_parameterisation(ew0, omega, *, phase_only, phase_basis='pointwise', n_nodes=20, support_threshold=0.001)[source]#

Build the least-squares (LM) solver parameterisation.

The LM solvers add penalties as residual rows rather than scalar terms. Returns (spectrum_fn, u0, nidcs, phase_only_eff). For the B-spline basis nidcs=0 and phase_only_eff=True (no amplitude block) — and the LM solvers’ 2nd-difference phase rows on u then are the P-spline penalty on the spline coefficients.

Parameters:

  • ew0 (NDArray[complex128])

  • omega (ArrayLike)

  • phase_only (bool)

  • phase_basis (str)

  • n_nodes (int)

  • support_threshold (float)

Return type:

tuple[Callable[[Array], Array], NDArray[float64], int, bool]

class croak._jax_pulse.Parameterisation(spectrum_fn, u0, penalties)[source]#

Bases: NamedTuple

A differentiable u -> spectrum map plus its initial parameters/penalties.

Parameters:

  • spectrum_fn (Callable[[Array], Array])

  • u0 (NDArray[float64])

  • penalties (Penalties)

spectrum_fn: Callable[[Array], Array]#

Alias for field number 0

u0: NDArray[float64]#

Alias for field number 1

penalties: Penalties#

Alias for field number 2

class croak._jax_pulse.Penalties(phase, amplitude, spectral, temporal)[source]#

Bases: NamedTuple

Regularisation terms for a parameterisation (each maps u -> scalar).

Parameters:

  • phase (Callable[[Array], Array])

  • amplitude (Callable[[Array], Array])

  • spectral (Callable[[Array], Array])

  • temporal (Callable[[Array], Array])

phase: Callable[[Array], Array]#

Alias for field number 0

amplitude: Callable[[Array], Array]#

Alias for field number 1

spectral: Callable[[Array], Array]#

Alias for field number 2

temporal: Callable[[Array], Array]#

Alias for field number 3

croak._jax_pulse.time_window_mask(t, win_lo, win_hi, dt)[source]#

Weight (in [0, 1]) that penalises field energy outside a window.

The complement of a Planck taper over [win_lo, win_hi] (roll-off 20*dt each side, mirroring croak.processing.post_filter()): 0 inside the window, rising smoothly to 1 outside it. Multiplied against |E(t)|^2 by the temporal regularisation penalty.

Parameters:

  • t (ArrayLike)

  • win_lo (float)

  • win_hi (float)

  • dt (float)

Return type:

NDArray[float64]

class croak._jax_pulse.ExtraParamSpec(fit_thickness=False, thickness0=0.0, scale_L=0.0, fit_tau0=False, tau0_0=0.0, scale_tau=0.0, fit_smearing=False, smear_scale0=1.0, scale_smear=1.0, fit_smearing_split=False, smear_delta0=1.0)[source]#

Bases: object

Specification of the optional scalar parameters fitted with the pulse.

The slab thickness L and a delay-zero offset \(\tau_0\) are optimised as dimensionless, O(1)-scaled offsets appended to the pulse parameter vector, so the augmented problem stays well-conditioned for NLopt / SciPy trf despite the tiny physical scales (L ~ 1e-5 m, \(\tau_0 \sim 10^{-15}\) s):

\[L = L_0 + \delta_L\,s_L, \qquad \tau_0 = \tau_{0,0} + \delta_\tau\,s_\tau\]

with the optimised variables \(\delta_L, \delta_\tau\) initialised at zero. The returned lower bound \(-L_0/s_L\) keeps L non-negative for the bounded solvers (the Optimistix solvers ignore it). The centres \(L_0\), \(\tau_{0,0}\) are the priors: the nominal values for a cold-started retrieval, the previous fit’s values when warm-starting (see croak.pipeline.retrieve_from_tracedata()), and the fitted values when re-linearising for covariance. An extra whose fit_* flag is off is held at its centre.

Parameters:

  • fit_thickness (bool, optional) – Fit the slab thickness (requires a dispersive slab; thickness0 > 0).

  • thickness0 (float, optional) – Centre thickness \(L_0\) (m) — the prior/held value.

  • scale_L (float, optional) – Normalisation \(s_L\) (m) for \(\delta_L\); defaults to thickness0 (so \(\delta_L\) is the fractional change).

  • fit_tau0 (bool, optional) – Fit the delay-zero offset.

  • tau0_0 (float, optional) – Centre delay-zero offset \(\tau_{0,0}\) (s); 0 for retrieval.

  • scale_tau (float, optional) – Normalisation \(s_\tau\) (s) for \(\delta_\tau\); typically one delay step. Required when fit_tau0 is set.

  • fit_smearing (bool, optional) – Fit the geometric-smearing strength as a dimensionless multiplier on both kernel widths — equivalently, the mask ratio d/D, the only geometric lever on the effect in the reduced model. A fitted value can also absorb omitted physics. Requires a smearing kernel.

  • smear_scale0 (float, optional) – Centre multiplier (1.0 = the kernel as supplied). For a split fit this is the p-channel centre.

  • scale_smear (float, optional) – Normalisation for the multiplier offset. It is already O(1), so the default 1.0 is almost always right.

  • fit_smearing_split (bool, optional) – Fit the gate-shape (p) and delay (delta) kernel widths as two independent multipliers instead of one joint one — the diagnostic for which channel deviates from the geometric prediction. Requires fit_smearing.

  • smear_delta0 (float, optional) – Centre multiplier of the delay channel for a split fit (1.0 = the kernel as supplied).

property any: bool#

Whether any extra parameter is fitted.

class croak._jax_pulse.Augmentation(n_pulse, u0, unpack, physical, lower_bounds, idx_thickness, idx_tau0, idx_smear, idx_smear_delta, scales)[source]#

Bases: NamedTuple

The augmented parameter vector plus its (un)packing maps.

See augment(). unpack is differentiable (used inside the JAX objective/Jacobian); physical returns plain floats for reporting; the scales convert a covariance σ on \(\delta\) back to physical units.

Parameters:

  • n_pulse (int)

  • u0 (NDArray[float64])

  • unpack (Callable[[Array], tuple[Array, Array | float, Array | float, Array | float]])

  • physical (Callable[[ArrayLike], tuple[float, float, float | tuple[float, float]]])

  • lower_bounds (NDArray[float64])

  • idx_thickness (int | None)

  • idx_tau0 (int | None)

  • idx_smear (int | None)

  • idx_smear_delta (int | None)

  • scales (tuple[float, float, float])

n_pulse: int#

Alias for field number 0

u0: NDArray[float64]#

Augmented vector [u_pulse; delta_L?; delta_tau?; delta_smear?; delta_smear_d?].

unpack: Callable[[Array], tuple[Array, Array | float, Array | float, Array | float]]#

unpack(u_ext) -> (u_pulse, thickness, tau0, smear_scale) (JAX-friendly). smear_scale is a scalar for a joint fit and a length-2 vector (scale_p, scale_delta) when the channels are fitted separately.

physical: Callable[[ArrayLike], tuple[float, float, float | tuple[float, float]]]#

physical(u_ext) -> (thickness, tau0, smear_scale) as floats; smear_scale is a (scale_p, scale_delta) tuple for a split fit.

lower_bounds: NDArray[float64]#

Per-component lower bounds (-inf for the pulse block and tau0).

idx_thickness: int | None#

Index of the thickness offset in u_ext (None when not fitted).

idx_tau0: int | None#

Index of the delay-zero offset in u_ext (None when not fitted).

idx_smear: int | None#

Index of the smearing-scale offset in u_ext (None when not fitted). For a split fit this is the p-channel scale.

idx_smear_delta: int | None#

Index of the delay-channel scale offset (None unless split-fitted).

scales: tuple[float, float, float]#

Resolved scales (scale_L, scale_tau, scale_smear) (0 where not fitted). The split delay-channel offset shares scale_smear.

croak._jax_pulse.augment(u0_pulse, spec)[source]#

Append the fitted extra scalars to a pulse parameter vector.

Parameters:
Returns:

Augmentation – The augmented vector and its packing maps (see the class docstring).

Raises:

ValueError – If a fitted extra has a non-positive scale.

Return type:

Augmentation