"""Differentiable resonance extraction from short time-domain ringdowns.
The routines in this module deliberately operate only on JAX arrays. A scalar
objective constructed from :func:`fit_resonance` can therefore be differentiated
through an FDTD detector trace and back to material or geometry parameters.
The current estimator targets one isolated, already-localized resonance. It
coherently demodulates equal translated windows and fits the complex ratio
between adjacent windows. Translating a window over a single damped mode
``E(t) = Re[a exp((-decay_rate + 2j*pi*frequency) t)]``
multiplies its demodulated amplitude by one complex number. Fitting that number
is substantially better conditioned in float32 than fitting the nearly-unit
per-timestep pole of an optical-frequency FDTD signal.
"""
from __future__ import annotations
import math
from typing import NamedTuple
import jax
import jax.numpy as jnp
from fdtdx.constants import c as C_0
class ResonanceEstimate(NamedTuple):
"""JAX-array outputs of a differentiable single-resonance fit.
Attributes:
frequency: Resonant frequency in Hz.
decay_rate: Field-amplitude decay rate in rad/s (equivalently 1/s).
quality_factor: ``pi * frequency / decay_rate``.
pole: Complex translation factor between adjacent coherent windows.
residual: Relative RMS residual of the complex exponential regression.
window_consistency: Relative RMS spread of adjacent-window decay rates.
amplitudes: Complex coherently demodulated amplitude in every window.
"""
frequency: jax.Array
decay_rate: jax.Array
quality_factor: jax.Array
pole: jax.Array
residual: jax.Array
window_consistency: jax.Array
amplitudes: jax.Array
class EnergyDecayEstimate(NamedTuple):
"""JAX-array outputs of a differentiable stored-energy decay fit.
Unlike :class:`ResonanceEstimate`, this estimator does not infer a
frequency. It fits the logarithm of equal-window mean electromagnetic
energies to one exponential and converts the resulting *energy* decay
rate to ``Q = 2*pi*frequency / decay_rate``.
"""
decay_rate: jax.Array
quality_factor: jax.Array
residual: jax.Array
window_consistency: jax.Array
window_energies: jax.Array
class RingdownAcquisition(NamedTuple):
"""Differentiable, fit-free measures of a time-domain ringdown.
These measures are intended for the acquisition phase of an inverse-design
run, before a sufficiently isolated pole exists for :func:`fit_resonance`.
The constant component is removed before measuring the signal; without
that operation a small Yee-grid DC offset can dominate a weak optical
ringdown and provide a misleading optimization gradient.
"""
early_rms: jax.Array
late_rms: jax.Array
retention: jax.Array
score: jax.Array
class CoherentRingdownAcquisition(NamedTuple):
"""Target-frequency measures for acquiring an atom-coupled resonance.
Unlike broadband RMS, these amplitudes are coherently demodulated at one
specified frequency. They therefore reject late DC, grid transients, and
modes outside the atomic linewidth while remaining differentiable.
"""
early_amplitude: jax.Array
late_amplitude: jax.Array
retention: jax.Array
late_coherence: jax.Array
score: jax.Array
class CavityQEDMetrics(NamedTuple):
"""Differentiable mode-volume and Purcell estimates in stable log form."""
mode_volume: jax.Array
purcell_factor: jax.Array
log_mode_volume: jax.Array
log_purcell_factor: jax.Array
def predict_continuation_frequency(
previous_frequency: jax.Array | float,
current_frequency: jax.Array | float,
*,
branch_spacing: jax.Array | float,
step_ratio: jax.Array | float = 1.0,
maximum_prediction_fraction: float = 0.25,
) -> jax.Array:
"""Predict the next point on one continuous resonance branch.
An isolated Maxwell eigenvalue is locally analytic under a smooth material
continuation (the usual Kato/Rellich eigenvalue-continuation result). A
secant is therefore a first-order predictor for the next pole. Clipping
that prediction inside a fixed fraction of the translated-window branch
spacing prevents a noisy previous step from jumping to a neighboring
``1 / window_duration`` alias.
This helper only operates on already measured scalar frequencies. It adds
no FDTD solve, and its result is stop-gradient because selecting a pole
branch is a discrete controller decision that must remain frozen during a
forward/adjoint pair and all of its trust-region replays.
"""
if not 0.0 < maximum_prediction_fraction < 0.5:
raise ValueError("maximum_prediction_fraction must lie strictly inside (0, 0.5)")
previous = jnp.asarray(previous_frequency)
current = jnp.asarray(current_frequency)
spacing = jnp.asarray(branch_spacing, dtype=current.dtype)
ratio = jnp.asarray(step_ratio, dtype=current.dtype)
limit = jnp.asarray(maximum_prediction_fraction, dtype=current.dtype) * spacing
trend = ratio * (current - previous)
predicted = current + jnp.clip(trend, -limit, limit)
valid = (
jnp.isfinite(previous)
& jnp.isfinite(current)
& jnp.isfinite(spacing)
& jnp.isfinite(ratio)
& (previous > 0.0)
& (current > 0.0)
& (spacing > 0.0)
& (ratio >= 0.0)
)
predicted = jnp.where(valid, predicted, current)
return jax.lax.stop_gradient(predicted)
def dominant_frequency_branch(
signal: jax.Array,
*,
time_step: float | jax.Array,
center_frequency: float | jax.Array,
samples_per_window: int,
start_index: int = 0,
num_windows: int = 3,
maximum_branch_offset: int = 4,
) -> jax.Array:
"""Select the dominant translated-window frequency branch.
A pole inferred from the phase between windows is ambiguous modulo
``1 / window_duration``. A coarse periodogram over all fit windows has
finer spacing by ``num_windows`` and therefore selects the branch before
:func:`fit_resonance` performs its precise complex-pole regression. The
discrete branch is intentionally stop-gradient; derivatives flow through
the subsequent fixed-branch resonance fit.
"""
signal = jnp.asarray(signal)
if signal.ndim != 1:
raise ValueError(f"signal must be one-dimensional, got shape {signal.shape}")
if samples_per_window < 2:
raise ValueError("samples_per_window must be at least 2")
if num_windows < 3:
raise ValueError("at least three windows are required")
if maximum_branch_offset < 0:
raise ValueError("maximum_branch_offset must be nonnegative")
sample_count = samples_per_window * num_windows
if start_index < 0 or start_index + sample_count > signal.shape[0]:
raise ValueError("requested branch-search windows do not fit in signal")
real_dtype = jnp.result_type(signal.dtype, jnp.float32)
samples = signal[start_index : start_index + sample_count].astype(real_dtype)
window = jnp.hanning(sample_count + 2).astype(real_dtype)[1:-1]
spectrum = jnp.abs(jnp.fft.rfft(samples * window))
dt = jnp.asarray(time_step, dtype=real_dtype)
frequencies = jnp.arange(spectrum.shape[0], dtype=real_dtype) / (dt * jnp.asarray(sample_count, dtype=real_dtype))
branch_spacing = 1.0 / (dt * jnp.asarray(samples_per_window, dtype=real_dtype))
half_width = (maximum_branch_offset + 0.5) * branch_spacing
center = jnp.asarray(center_frequency, dtype=real_dtype)
in_band = jnp.abs(frequencies - center) <= half_width
scores = jnp.where(in_band, spectrum, -jnp.inf)
return jax.lax.stop_gradient(frequencies[jnp.argmax(scores)])
def late_time_frequency_branch(
signal: jax.Array,
*,
time_step: float | jax.Array,
frequency_window: tuple[float, float],
samples_per_window: int,
start_index: int = 0,
num_windows: int | None = None,
) -> jax.Array:
"""Select the observable long-lived branch from the final fit window.
A broadband pulse can leave a large, rapidly decaying resonance near the
beginning of a ringdown and a much smaller high-Q resonance at late time.
Selecting a frequency from the complete trace then biases a single-pole
regression toward the prompt mode. This helper removes local DC, applies
a Hann window to the *last* complete fit window, and returns its strongest
positive-frequency Fourier bin inside ``frequency_window``. The existing
:func:`fit_resonance` regression subsequently refines the bin to sub-bin
frequency and decay-rate accuracy.
The returned branch is explicitly stop-gradient. Pole discovery is a
discrete choice: run this helper on an accepted checkpoint trace, store
the resulting scalar frequency, and keep it fixed throughout the next
optimizer epoch or line search. Derivatives through :func:`fit_resonance`
remain ordinary exact JAX derivatives on that fixed branch. Re-selecting
the branch for every trial can still make the *forward* objective
discontinuous even though no derivative passes through the selection.
``frequency_window`` should be a fixed physical/source bandwidth, not a
per-structure hand-tuned interval. A peak returned from noise is rejected
by the residual, window-consistency, time-prefix, and energy checks applied
to the subsequent resonance fit.
Args:
signal: One-dimensional real detector trace containing the ringdown.
time_step: Finite positive scalar time between adjacent samples in
seconds. It is a fixed acquisition setting rather than a traced
branch-selection input.
frequency_window: Inclusive positive-frequency search interval in Hz,
ordered ``(minimum, maximum)``.
samples_per_window: Samples in each translated fit window. Using the
same value in this selector and :func:`fit_resonance` guarantees
that the selected Fourier bin lies within the fit's unambiguous
phase-refinement interval.
start_index: First ringdown sample, after the source has cleared.
num_windows: Number of complete windows in the intended pole fit. By
default, use every complete window after ``start_index``.
Returns:
A scalar JAX array containing the selected frequency in Hz. If the
interval contains no finite bin or the final window has no numerically
observable oscillation, the result is ``nan`` so downstream trust
checks cannot accidentally certify it.
Raises:
ValueError: If the signal, frequency interval, or static window
configuration is invalid.
"""
signal = jnp.asarray(signal)
if signal.ndim != 1:
raise ValueError(f"signal must be one-dimensional, got shape {signal.shape}")
if jnp.iscomplexobj(signal):
raise ValueError("signal must be real-valued")
if samples_per_window < 2:
raise ValueError("samples_per_window must be at least 2")
if start_index < 0 or start_index >= signal.shape[0]:
raise ValueError(f"start_index must be in [0, {signal.shape[0]}), got {start_index}")
time_step_array = jnp.asarray(time_step)
if time_step_array.ndim != 0:
raise ValueError(f"time_step must be a scalar, got shape {time_step_array.shape}")
try:
time_step_value = float(time_step)
except (TypeError, ValueError, jax.errors.ConcretizationTypeError) as exc:
raise ValueError("time_step must be a concrete finite positive scalar") from exc
if not math.isfinite(time_step_value) or time_step_value <= 0.0:
raise ValueError("time_step must be finite and strictly positive")
if len(frequency_window) != 2:
raise ValueError("frequency_window must contain exactly two values")
minimum_frequency, maximum_frequency = (float(value) for value in frequency_window)
if (
not math.isfinite(minimum_frequency)
or not math.isfinite(maximum_frequency)
or minimum_frequency <= 0.0
or maximum_frequency <= minimum_frequency
):
raise ValueError("frequency_window must be finite, positive, and strictly increasing")
available_windows = (signal.shape[0] - start_index) // samples_per_window
if num_windows is None:
num_windows = available_windows
if num_windows < 3:
raise ValueError("at least three complete coherent windows are required")
if num_windows > available_windows:
raise ValueError(f"requested {num_windows} windows but only {available_windows} fit in the signal")
last_start = start_index + (num_windows - 1) * samples_per_window
real_dtype = jnp.result_type(signal.dtype, jnp.float32)
raw_samples = signal[last_start : last_start + samples_per_window].astype(real_dtype)
samples_are_finite = jnp.all(jnp.isfinite(raw_samples))
raw_scale = jnp.max(jnp.abs(raw_samples))
samples = raw_samples - jnp.mean(raw_samples)
oscillating_scale = jnp.max(jnp.abs(samples))
window = jnp.hanning(samples_per_window + 2).astype(real_dtype)[1:-1]
windowed_samples = samples * window
spectrum = jnp.abs(jnp.fft.rfft(windowed_samples))
dt = jnp.asarray(time_step_value, dtype=real_dtype)
frequencies = jnp.arange(spectrum.shape[0], dtype=real_dtype) / (
dt * jnp.asarray(samples_per_window, dtype=real_dtype)
)
in_band = (frequencies >= minimum_frequency) & (frequencies <= maximum_frequency)
valid_bins = in_band & jnp.isfinite(spectrum)
scores = jnp.where(valid_bins, spectrum, -jnp.inf)
peak = jnp.max(scores)
selected = frequencies[jnp.argmax(scores)]
tiny = jnp.asarray(jnp.finfo(real_dtype).tiny, dtype=real_dtype)
eps = jnp.asarray(jnp.finfo(real_dtype).eps, dtype=real_dtype)
time_domain_floor = jnp.maximum(tiny, eps * raw_scale)
spectral_floor = jnp.maximum(tiny, eps * jnp.sum(jnp.abs(windowed_samples)))
observable = (
samples_are_finite
& jnp.any(valid_bins)
& jnp.isfinite(peak)
& (oscillating_scale > time_domain_floor)
& (peak > spectral_floor)
)
selected = jnp.where(observable, selected, jnp.asarray(jnp.nan, dtype=real_dtype))
return jax.lax.stop_gradient(selected)
def windowed_mode_field_rms(
demodulated_amplitude: jax.Array,
*,
decay_rate: float | jax.Array,
time_step: float | jax.Array,
samples_per_window: int,
) -> jax.Array:
"""Convert a real-signal Hann demodulate to same-window modal RMS.
Demodulating ``A cos(omega t + phi)`` returns ``A / 2``, not the peak
amplitude. The additional decay-envelope factor maps the Hann-weighted
coherent amplitude onto the rectangular mean-square window used by
:func:`fit_energy_decay`, so stored energy and local field refer to the
same interval.
"""
if samples_per_window < 2:
raise ValueError("samples_per_window must be at least 2")
amplitude = jnp.asarray(demodulated_amplitude)
real_dtype = jnp.result_type(jnp.real(amplitude).dtype, jnp.float32)
gamma = jnp.asarray(decay_rate, dtype=real_dtype)
dt = jnp.asarray(time_step, dtype=real_dtype)
index = jnp.arange(samples_per_window, dtype=real_dtype)
envelope = jnp.exp(-gamma * dt * index)
window = jnp.hanning(samples_per_window + 2).astype(real_dtype)[1:-1]
coherent_envelope = jnp.sum(window * envelope) / jnp.sum(window)
rms_envelope = jnp.sqrt(jnp.mean(jnp.square(envelope)))
tiny = jnp.asarray(jnp.finfo(real_dtype).tiny, dtype=real_dtype)
return (
jnp.sqrt(jnp.asarray(2.0, dtype=real_dtype))
* jnp.abs(amplitude)
* rms_envelope
/ jnp.maximum(coherent_envelope, tiny)
)
def cavity_qed_metrics(
*,
quality_factor: jax.Array,
frequency: jax.Array,
local_field_rms: jax.Array,
modal_energy: jax.Array,
symmetry_factor: float | jax.Array = 1.0,
length_unit: float | jax.Array = 1e-6,
) -> CavityQEDMetrics:
"""Estimate effective mode volume and Purcell factor without unstable ratios.
For a late-time single-mode ringdown, total time-averaged electromagnetic
energy divided by the squared local-field RMS is the electric mode volume.
The computation stays in log space until its returned physical metrics are
formed. This matters in float32 FDTD: differentiating the algebraically
equivalent nested ratio ``U / E**2`` can overflow even when its value and
true derivative are finite.
``length_unit`` is metres per desired mode-volume length unit; its default
returns mode volume in cubic micrometres. ``symmetry_factor`` restores the
multiplicity of an energy integral taken on a reduced simulation domain.
"""
dtype = jnp.result_type(
quality_factor,
frequency,
local_field_rms,
modal_energy,
jnp.float32,
)
tiny = jnp.asarray(jnp.finfo(dtype).tiny, dtype=dtype)
quality_factor = jnp.maximum(jnp.asarray(quality_factor, dtype=dtype), tiny)
frequency = jnp.maximum(jnp.asarray(frequency, dtype=dtype), tiny)
local_field_rms = jnp.maximum(jnp.asarray(local_field_rms, dtype=dtype), jnp.sqrt(tiny))
modal_energy = jnp.maximum(jnp.asarray(modal_energy, dtype=dtype), tiny)
symmetry_factor = jnp.maximum(jnp.asarray(symmetry_factor, dtype=dtype), tiny)
length_unit = jnp.maximum(jnp.asarray(length_unit, dtype=dtype), tiny)
log_mode_volume = (
jnp.log(symmetry_factor) + jnp.log(modal_energy) - 2.0 * jnp.log(local_field_rms) - 3.0 * jnp.log(length_unit)
)
wavelength_units = C_0 / frequency / length_unit
log_purcell_factor = (
jnp.log(jnp.asarray(3.0 / (4.0 * jnp.pi**2), dtype=dtype))
+ 3.0 * jnp.log(wavelength_units)
+ jnp.log(quality_factor)
- log_mode_volume
)
return CavityQEDMetrics(
mode_volume=jnp.exp(log_mode_volume),
purcell_factor=jnp.exp(log_purcell_factor),
log_mode_volume=log_mode_volume,
log_purcell_factor=log_purcell_factor,
)
[docs]
def fixed_frequency_cavity_ldos(
*,
peak_resonant_ldos: jax.Array,
quality_factor: jax.Array,
pole_frequency: jax.Array,
emitter_frequency: jax.Array,
background_ldos: float | jax.Array = 0.0,
) -> jax.Array:
"""Evaluate a fitted cavity pole at one immutable emitter frequency.
``peak_resonant_ldos`` is the resonant contribution at the pole, excluding
``background_ldos``. The returned value applies the single-pole
Lorentzian detuning factor
``1 / (1 + (2 Q (f_emitter - f_pole) / f_pole)**2)``.
Unlike optimizing peak ``Q/V`` while following a moving pole, this metric
gives no benefit to a resonance that walks away from the emitter. Every
operation remains differentiable with respect to the fitted pole, Q, mode
volume, and ultimately the FDTD geometry.
"""
dtype = jnp.result_type(
peak_resonant_ldos,
quality_factor,
pole_frequency,
emitter_frequency,
background_ldos,
jnp.float32,
)
tiny = jnp.asarray(jnp.finfo(dtype).tiny, dtype=dtype)
peak = jnp.maximum(jnp.asarray(peak_resonant_ldos, dtype=dtype), 0.0)
quality = jnp.maximum(jnp.asarray(quality_factor, dtype=dtype), tiny)
pole = jnp.maximum(jnp.asarray(pole_frequency, dtype=dtype), tiny)
emitter = jnp.maximum(jnp.asarray(emitter_frequency, dtype=dtype), tiny)
background = jnp.maximum(jnp.asarray(background_ldos, dtype=dtype), 0.0)
linewidth_offset = 2.0 * quality * (emitter - pole) / pole
return background + peak / (1.0 + linewidth_offset**2)
def remove_ringdown_dc(signal: jax.Array) -> jax.Array:
"""Remove the constant component of a one-dimensional detector trace."""
signal = jnp.asarray(signal)
if signal.ndim != 1:
raise ValueError(f"signal must be one-dimensional, got shape {signal.shape}")
return signal - jnp.mean(signal)
def ringdown_acquisition(
signal: jax.Array,
*,
samples_per_window: int,
num_windows: int | None = None,
) -> RingdownAcquisition:
"""Return a fit-free objective for acquiring an atom-coupled resonance.
The score is the logarithm of the RMS field in the final complete window.
Maximizing it rewards both excitation at the detector and persistence
after the source pulse, without assuming that a single exponential pole is
already identifiable. Once a reliable pole appears, callers can switch to
:func:`fit_resonance` for direct Q optimization.
"""
signal = remove_ringdown_dc(signal)
if samples_per_window < 2:
raise ValueError("samples_per_window must be at least 2")
available_windows = signal.shape[0] // samples_per_window
if num_windows is None:
num_windows = available_windows
if num_windows < 2:
raise ValueError("at least two complete windows are required")
if num_windows > available_windows:
raise ValueError(f"requested {num_windows} windows but only {available_windows} fit in the signal")
windows = signal[: num_windows * samples_per_window].reshape((num_windows, samples_per_window))
tiny = jnp.asarray(jnp.finfo(windows.dtype).tiny, dtype=windows.dtype)
rms = jnp.sqrt(jnp.mean(windows**2, axis=1) + tiny)
early_rms = rms[0]
late_rms = rms[-1]
retention = late_rms / jnp.maximum(early_rms, tiny)
return RingdownAcquisition(
early_rms=early_rms,
late_rms=late_rms,
retention=retention,
score=jnp.log(jnp.maximum(late_rms, tiny)),
)
def coherent_ringdown_acquisition(
signal: jax.Array,
*,
time_step: float | jax.Array,
center_frequency: float | jax.Array,
samples_per_window: int,
num_windows: int | None = None,
retention_weight: float | jax.Array = 0.25,
) -> CoherentRingdownAcquisition:
"""Acquire a persistent ringdown specifically at ``center_frequency``.
Each complete window is Hann-weighted and coherently demodulated. The
score primarily maximizes the final-window target-frequency amplitude and
adds a smaller log-retention term. The latter favors persistence without
allowing an optimizer to win merely by suppressing initial excitation:
scaling the whole trace still changes the score with unit log slope.
This is intended for the topology-discovery interval before a clean pole
can be fitted. Once a reliable single pole exists, use
:func:`fit_resonance` and a direct Q/V objective.
"""
signal = jnp.asarray(signal)
if signal.ndim != 1:
raise ValueError(f"signal must be one-dimensional, got shape {signal.shape}")
if samples_per_window < 2:
raise ValueError("samples_per_window must be at least 2")
available_windows = signal.shape[0] // samples_per_window
if num_windows is None:
num_windows = available_windows
if num_windows < 2:
raise ValueError("at least two complete windows are required")
if num_windows > available_windows:
raise ValueError(f"requested {num_windows} windows but only {available_windows} fit in the signal")
count = num_windows * samples_per_window
windows = signal[:count].reshape((num_windows, samples_per_window))
dtype = windows.dtype
local_index = jnp.arange(samples_per_window, dtype=dtype)
hann = 0.5 - 0.5 * jnp.cos(2.0 * jnp.pi * local_index / jnp.asarray(samples_per_window - 1, dtype=dtype))
absolute_index = jnp.arange(count, dtype=dtype).reshape((num_windows, samples_per_window))
phase = (
2.0
* jnp.pi
* jnp.asarray(center_frequency, dtype=dtype)
* (jnp.asarray(time_step, dtype=dtype) * absolute_index)
)
# Fit cosine, sine, and a constant simultaneously. Orthogonalizing the
# target-frequency basis against DC per window avoids leakage when a window
# contains a non-integer number of optical cycles. It also makes the
# result independent of unused tail samples when an optimizer shortens its
# FDTD tape to the active acquisition window.
basis = jnp.stack((jnp.cos(phase), jnp.sin(phase), jnp.ones_like(phase)), axis=-1)
weighted_basis = basis * hann[jnp.newaxis, :, jnp.newaxis]
gram = jnp.einsum("wsi,wsj->wij", weighted_basis, basis)
rhs = jnp.einsum("wsi,ws->wi", weighted_basis, windows)
regularization = 10.0 * jnp.finfo(dtype).eps * jnp.trace(gram, axis1=-2, axis2=-1) / 3.0
gram = gram + regularization[:, jnp.newaxis, jnp.newaxis] * jnp.eye(3, dtype=dtype)
coefficients = jnp.linalg.solve(gram, rhs[..., jnp.newaxis])[..., 0]
tiny = jnp.asarray(jnp.finfo(dtype).tiny, dtype=dtype)
amplitudes = jnp.sqrt(coefficients[:, 0] ** 2 + coefficients[:, 1] ** 2 + tiny)
early_amplitude = amplitudes[0]
late_amplitude = amplitudes[-1]
retention = late_amplitude / jnp.maximum(early_amplitude, tiny)
late_centered = windows[-1] - coefficients[-1, 2]
late_rms = jnp.sqrt(jnp.mean(late_centered**2) + tiny)
late_coherence = late_amplitude / jnp.maximum(jnp.sqrt(2.0) * late_rms, tiny)
weight = jnp.asarray(retention_weight, dtype=dtype)
score = jnp.log(jnp.maximum(late_amplitude, tiny)) + weight * jnp.log(jnp.maximum(retention, tiny))
return CoherentRingdownAcquisition(
early_amplitude=early_amplitude,
late_amplitude=late_amplitude,
retention=retention,
late_coherence=late_coherence,
score=score,
)
def fit_resonance(
signal: jax.Array,
*,
time_step: float | jax.Array,
center_frequency: float | jax.Array,
samples_per_window: int,
start_index: int = 0,
num_windows: int | None = None,
) -> ResonanceEstimate:
"""Fit one complex resonance to a real-valued FDTD ringdown.
The fit is GPU-native and differentiable with respect to ``signal``,
``time_step``, and ``center_frequency``. ``samples_per_window``,
``start_index``, and ``num_windows`` define array shapes and are therefore
static Python values. The center frequency should already identify the
desired mode closely enough that its phase advances by less than ``pi``
between coherent windows.
Args:
signal: One-dimensional real detector trace after the excitation pulse.
time_step: Time between adjacent samples in seconds.
center_frequency: Tracking/demodulation frequency in Hz.
samples_per_window: Samples in each translated Hann window. A window
spanning tens of optical periods suppresses the counter-rotating
component of a real trace.
start_index: First sample included in the fit.
num_windows: Number of complete windows. By default, use every complete
window after ``start_index``.
Returns:
A :class:`ResonanceEstimate` containing only JAX arrays.
Raises:
ValueError: If the signal or static window configuration is invalid.
Notes:
Mode selection and early-stop decisions are intentionally outside this
function. Those operations are discrete and should be held fixed while
differentiating an optimization step.
"""
signal = jnp.asarray(signal)
if signal.ndim != 1:
raise ValueError(f"signal must be one-dimensional, got shape {signal.shape}")
if samples_per_window < 2:
raise ValueError("samples_per_window must be at least 2")
if start_index < 0 or start_index >= signal.shape[0]:
raise ValueError(f"start_index must be in [0, {signal.shape[0]}), got {start_index}")
available_windows = (signal.shape[0] - start_index) // samples_per_window
if num_windows is None:
num_windows = available_windows
if num_windows < 3:
raise ValueError("at least three complete coherent windows are required")
if num_windows > available_windows:
raise ValueError(f"requested {num_windows} windows but only {available_windows} fit in the signal")
sample_count = num_windows * samples_per_window
samples = signal[start_index : start_index + sample_count]
real_dtype = jnp.result_type(samples.dtype, jnp.float32)
samples = samples.astype(real_dtype)
dt = jnp.asarray(time_step, dtype=real_dtype)
reference = jnp.asarray(center_frequency, dtype=real_dtype)
# Avoid zero-weight endpoints while retaining the sidelobe suppression of a
# Hann window. Every block uses an identical translated window, which makes
# the single-exponential ratio exact in the absence of the counter-rotating
# component and other modes.
window = jnp.hanning(samples_per_window + 2).astype(real_dtype)[1:-1]
indices = start_index + jnp.arange(sample_count, dtype=real_dtype)
carrier = jnp.exp(-2j * jnp.pi * reference * dt * indices)
blocks = (samples * carrier).reshape((num_windows, samples_per_window))
amplitudes = jnp.sum(blocks * window[None, :], axis=1) / jnp.sum(window)
previous = amplitudes[:-1]
following = amplitudes[1:]
tiny = jnp.asarray(jnp.finfo(real_dtype).tiny, dtype=real_dtype)
denominator = jnp.real(jnp.vdot(previous, previous)) + tiny
pole = jnp.vdot(previous, following) / denominator
window_duration = dt * samples_per_window
pole_magnitude = jnp.abs(pole)
decay_rate = -jnp.log(pole_magnitude) / window_duration
frequency = reference + jnp.angle(pole) / (2 * jnp.pi * window_duration)
quality_factor = jnp.pi * frequency / decay_rate
prediction_error = following - pole * previous
residual = jnp.sqrt(
jnp.real(jnp.vdot(prediction_error, prediction_error))
/ (jnp.real(jnp.vdot(following, following)) + tiny)
+ tiny
)
adjacent_magnitudes = jnp.maximum(jnp.abs(following / previous), tiny)
adjacent_decay = -jnp.log(adjacent_magnitudes) / window_duration
decay_mean = jnp.mean(adjacent_decay)
window_consistency = jnp.sqrt(
jnp.mean((adjacent_decay - decay_mean) ** 2) + tiny
) / (jnp.abs(decay_mean) + tiny)
return ResonanceEstimate(
frequency=frequency,
decay_rate=decay_rate,
quality_factor=quality_factor,
pole=pole,
residual=residual,
window_consistency=window_consistency,
amplitudes=amplitudes,
)
def fit_energy_decay(
energy: jax.Array,
*,
time_step: float | jax.Array,
frequency: float | jax.Array,
samples_per_window: int,
start_index: int = 0,
num_windows: int | None = None,
) -> EnergyDecayEstimate:
"""Fit a differentiable exponential lifetime to total modal energy.
This is a complementary check on an atom-point field fit. A local field
trace can overestimate lifetime when a weak, slowly decaying tail is not
the dominant energy-carrying cavity mode. Fitting total electromagnetic
energy makes that failure mode visible while retaining gradients through
the FDTD tape.
``energy`` must be a non-negative one-dimensional detector trace. Equal
windows are averaged before fitting in log space, which suppresses the
electric/magnetic exchange visible in instantaneous Yee-grid samples.
At least three windows are required so fit residual and adjacent-window
consistency remain independently measurable.
"""
energy = jnp.asarray(energy)
if energy.ndim != 1:
raise ValueError(f"energy must be one-dimensional, got shape {energy.shape}")
if samples_per_window < 2:
raise ValueError("samples_per_window must be at least 2")
if start_index < 0 or start_index >= energy.shape[0]:
raise ValueError(f"start_index must be in [0, {energy.shape[0]}), got {start_index}")
available_windows = (energy.shape[0] - start_index) // samples_per_window
if num_windows is None:
num_windows = available_windows
if num_windows < 3:
raise ValueError("at least three complete energy windows are required")
if num_windows > available_windows:
raise ValueError(f"requested {num_windows} windows but only {available_windows} fit in the energy trace")
count = num_windows * samples_per_window
real_dtype = jnp.result_type(energy.dtype, jnp.float32)
samples = energy[start_index : start_index + count].astype(real_dtype)
windows = samples.reshape((num_windows, samples_per_window))
tiny = jnp.asarray(jnp.finfo(real_dtype).tiny, dtype=real_dtype)
window_energies = jnp.maximum(jnp.mean(windows, axis=1), tiny)
log_energy = jnp.log(window_energies)
# Regress against a centered unitless window index. This avoids the poor
# conditioning of optical times expressed as ~1e-12 seconds in float32.
index = jnp.arange(num_windows, dtype=real_dtype)
centered_index = index - jnp.mean(index)
centered_log_energy = log_energy - jnp.mean(log_energy)
slope_per_window = jnp.vdot(centered_index, centered_log_energy) / jnp.vdot(centered_index, centered_index)
window_duration = jnp.asarray(time_step, dtype=real_dtype) * jnp.asarray(samples_per_window, dtype=real_dtype)
decay_rate = -slope_per_window / window_duration
frequency = jnp.asarray(frequency, dtype=real_dtype)
quality_factor = 2.0 * jnp.pi * frequency / decay_rate
prediction = jnp.mean(log_energy) + slope_per_window * centered_index
residual = jnp.sqrt(jnp.mean((log_energy - prediction) ** 2) + tiny)
adjacent_decay = -jnp.diff(log_energy) / window_duration
decay_mean = jnp.mean(adjacent_decay)
window_consistency = jnp.sqrt(
jnp.mean((adjacent_decay - decay_mean) ** 2) + tiny
) / (jnp.abs(decay_mean) + tiny)
return EnergyDecayEstimate(
decay_rate=decay_rate,
quality_factor=quality_factor,
residual=residual,
window_consistency=window_consistency,
window_energies=window_energies,
)
