# SPDX-FileCopyrightText: 2026 John C. Travers and contributors
#
# SPDX-License-Identifier: GPL-3.0-or-later
"""Read fields, spectra and spectral densities out of a saved run.
The field, spectrum and density accessors, in solprop's conventions. Each one takes an
:class:`~solprop.output.schema.Output` -- a memory handler or an HDF5 handler, through
the same protocol -- selects a saved plane and a set of modes, and returns an abscissa
and an array whose trivial mode and polarisation axes have been squeezed away.
Three rules hold throughout.
**Convention.** Arrays come back in the optics convention by default and in the FFT
convention on request. For a spectrum the difference is the conjugation *and*
the phase origin, both of which the output layer applies; for a time-domain array it is
the conjugation alone. A real field, an energy and a density are convention-free.
**Normalisation.** :func:`get_spectrum` carries the spectral normalisation, so that
``abs(spectrum) ** 2`` is the spectral energy density in J/(rad/s) --
:func:`solprop.fourier.spectral_density` of the same plane, to round-off. The stored
array itself, the raw forward DFT, is ``output.fields()``.
**Band.** The default is the grid's active band, on an ascending
frequency axis. Pass ``band="all"`` for every stored bin, remembering that a field grid
then includes the zero-frequency bin, where a wavelength is undefined.
No window is ever applied implicitly. ``bandpass`` is a keyword, it is a wavelength
range, and nothing happens without it.
"""
from typing import Literal
import numpy as np
from jaxtyping import Complex, Float, Inexact
from solprop import fourier
from solprop.exceptions import ProcessingError
from solprop.grid import TimeGrid
from solprop.optics import TransferFunction, stored_factor
from solprop.output.schema import Convention, Output
from solprop.processing import spectra
from solprop.processing.axes import (
Band,
ModeSelection,
SpectralAxis,
_apply_modes,
_check_choice,
_keep_index,
_nearest_plane,
_plane_field,
_squeeze,
grid_of,
spectral_axis,
)
__all__ = [
"FIELD_KINDS",
"FieldKind",
"ModeSelection",
"get_field",
"get_phase",
"get_power",
"get_spectral_density",
"get_spectrum",
"nearest_plane",
]
FieldKind = Literal["analytic", "real", "envelope"]
"""Which time-domain representation :func:`get_field` returns."""
FIELD_KINDS: tuple[FieldKind, ...] = ("analytic", "real", "envelope")
"""The analytic signal, the real field, and the analytic signal without its carrier."""
[docs]
def nearest_plane(output: Output, z: float) -> int:
"""Return the index of the saved plane nearest a position.
Parameters
----------
output : solprop.output.schema.Output
An opened output handler holding at least one plane.
z : float
Position along the propagation, m.
Returns
-------
int
Index into the saved planes.
Raises
------
solprop.exceptions.OutputError
If no plane has been saved.
Examples
--------
>>> from solprop.processing._examples import example_output
>>> from solprop.processing.fields import nearest_plane
>>> nearest_plane(example_output(), 0.049)
2
"""
return _nearest_plane(output, z)
[docs]
def get_spectrum(
output: Output,
*,
plane: int | None = None,
z: float | None = None,
modes: ModeSelection = None,
band: Band = "active",
convention: Convention = "optics",
) -> tuple[Float[np.ndarray, " n_x"], Complex[np.ndarray, "*shape n_x"]]:
"""Return the complex spectrum of one plane, or of every plane.
The stored raw forward DFT is scaled by the spectral normalisation --
``2 dt / sqrt(2 pi)`` on a field grid and ``dt / sqrt(2 pi)`` on an envelope grid --
so that the modulus squared is the spectral energy density in J/(rad/s)
(``conventions.md`` section 4). The axis is ascending and, by default, the active
band.
Parameters
----------
output : solprop.output.schema.Output
An opened output handler.
plane : int or None
Index of the plane wanted, negative from the end. With ``z``, at most one of the
two; with neither, every plane is returned and the result carries a leading
``n_z`` axis.
z : float or None
Position of the plane wanted, m; the nearest saved plane is used.
modes : int, sequence of int, 'sum' or None
Which modes to return. ``'sum'`` is refused here: modal amplitudes are complex
and adding them is meaningless away from a point in space.
band : {'active', 'all'}
Which part of the stored axis to return.
convention : {'optics', 'fft'}
Sign convention of the returned spectrum.
Returns
-------
omega : numpy.ndarray
Angular frequency, shape ``(n_x,)``, rad/s, ascending.
spectrum : numpy.ndarray
The spectrum, in the square root of J/(rad/s), with the leading ``n_z`` axis
when no plane was selected and trivial mode and polarisation axes squeezed.
Raises
------
solprop.exceptions.ModeError
If ``modes='sum'``.
Examples
--------
>>> import numpy as np
>>> from solprop.processing._examples import example_output
>>> from solprop.processing.fields import get_spectrum
>>> omega, spectrum = get_spectrum(example_output(), plane=-1)
>>> spectrum.shape == omega.shape
True
"""
grid = grid_of(output)
values = _plane_field(output, plane, z, convention=convention)
values = values * _fft_norm(grid)
values, kept = _apply_modes(values, modes, allow_sum=False)
keep = _keep_index(grid, band)
return grid.omega[keep], _squeeze(values[..., keep], mode_axis=kept)
[docs]
def get_spectral_density( # noqa: PLR0913 - the abscissa and the smoothing on top
output: Output,
axis: SpectralAxis = "wavelength",
*,
plane: int | None = None,
z: float | None = None,
modes: ModeSelection = None,
band: Band = "active",
spec_range: tuple[float, float] | None = None,
resolution: float | None = None,
) -> tuple[Float[np.ndarray, " n_x"], Float[np.ndarray, "*shape n_x"]]:
"""Return the spectral energy density of one plane, or of every plane.
The density is per unit of the abscissa asked for: J/(rad/s),
J/Hz or J/m, with the Jacobian of :func:`solprop.processing.axes.spectral_axis`
applied. The wavelength abscissa is returned ascending, so the density is reversed
with it. Polarisation components are summed, because they are two components of one
field and their densities add.
Parameters
----------
output : solprop.output.schema.Output
An opened output handler.
axis : {'wavelength', 'f', 'omega'}
The abscissa, in m, Hz or rad/s.
plane : int or None
Index of the plane wanted, negative from the end; see :func:`get_spectrum`.
z : float or None
Position of the plane wanted, m.
modes : int, sequence of int, 'sum' or None
Which modes to return; ``'sum'`` adds their densities.
band : {'active', 'all'}
Which part of the stored axis to use.
spec_range : tuple of float or None
Limits on the abscissa, in its own units, applied after the change of variable.
resolution : float or None
Instrument resolution, in the units of the abscissa, as a full width at half
maximum. When given, the density is resampled onto a uniform axis and smoothed
by :func:`solprop.processing.spectra.specres`.
Returns
-------
x : numpy.ndarray
The abscissa, shape ``(n_x,)``, ascending.
density : numpy.ndarray
The spectral energy density, J per unit of ``x``.
Examples
--------
>>> import numpy as np
>>> from solprop.processing._examples import example_output
>>> from solprop.processing.fields import get_spectral_density
>>> wavelength, density = get_spectral_density(example_output(), plane=0)
>>> bool(np.all(np.diff(wavelength) > 0))
True
"""
grid = grid_of(output)
values = _plane_field(output, plane, z, convention="optics")
density = np.abs(values * _fft_norm(grid)) ** 2
density = np.sum(density, axis=-2, keepdims=True)
density, kept = _apply_modes(density, modes, allow_sum=True)
keep = _keep_index(grid, band)
omega = grid.omega[keep]
x, jacobian = spectral_axis(omega, axis)
density = _squeeze(density[..., keep] * jacobian, mode_axis=kept)
if axis == "wavelength":
# The wavelength abscissa decreases where omega increases; both are reversed so
# that every accessor hands back an ascending axis.
x, density = x[::-1], density[..., ::-1]
if resolution is not None:
# Integrating in omega with a window evaluated on the target axis is the same
# sum as the target-axis integral with the measure d omega / |d omega / dx|.
weights = np.abs(grid.d_omega / jacobian)
if axis == "wavelength":
weights = weights[::-1]
return spectra.specres(
x, density, resolution, weights=weights, spec_range=spec_range
)
if spec_range is not None:
inside = (x >= min(spec_range)) & (x <= max(spec_range))
return x[inside], density[..., inside]
return x, density
[docs]
def get_field( # noqa: PLR0913 - a plane, a mode set, a representation and a window
output: Output,
*,
kind: FieldKind = "analytic",
plane: int | None = None,
z: float | None = None,
modes: ModeSelection = None,
oversampling: int = 1,
bandpass: spectra.Bandpass | None = None,
convention: Convention = "optics",
transfer: TransferFunction | None = None,
transform_limit: bool = False,
) -> tuple[Float[np.ndarray, " n_t"], Inexact[np.ndarray, "*shape n_t"]]:
"""Return the time-domain field of one plane, or of every plane.
On either grid and in three representations: the analytic signal
``E_a(t)``, whose modulus squared is the instantaneous power in watts; the real
field ``Re E_a``; and the envelope ``A(t) = E_a(t) exp(-i omega0 t)``. On an
envelope grid the stored state already is the envelope, and on a field grid the
analytic signal is formed by :func:`solprop.fourier.analytic_signal`; one definition
serves both.
Oversampling zero-pads the spectrum, so the returned time axis is the grid's own
fine axis: centred on zero, of length ``oversampling * n_t``.
Parameters
----------
output : solprop.output.schema.Output
An opened output handler.
kind : {'analytic', 'real', 'envelope'}
Which representation to return.
plane : int or None
Index of the plane wanted, negative from the end; see :func:`get_spectrum`.
z : float or None
Position of the plane wanted, m.
modes : int, sequence of int, 'sum' or None
Which modes to return. ``'sum'`` is refused: modal amplitudes are complex.
oversampling : int
Factor by which to resample the time axis, at least 1.
bandpass : tuple of float, numpy.ndarray or None
Wavelength limits, m, of an explicit Planck-tapered window applied to the
spectrum before the transform, or a ready-made window on the stored axis.
convention : {'optics', 'fft'}
Sign convention of the complex representations; ``'real'`` is convention-free.
transfer : TransferFunction or None
A linear optical element of :mod:`solprop.optics` applied to the spectrum before
the transform -- a compressor, a window, a stretch of air -- in the optics
convention.
transform_limit : bool
Whether to discard the spectral phase first, giving the shortest pulse the
spectrum supports, centred on ``t = 0``. Applied after
``bandpass`` and before ``transfer``.
Returns
-------
t : numpy.ndarray
Time axis, shape ``(n_t,)`` or ``(oversampling n_t,)``, s, centred on zero.
field : numpy.ndarray
The field, in the square root of watts, with trivial axes squeezed; real for
``kind='real'`` and complex otherwise.
Raises
------
solprop.exceptions.ProcessingError
If ``kind`` is unknown or ``oversampling`` is below one.
solprop.exceptions.ModeError
If ``modes='sum'``.
Examples
--------
>>> import numpy as np
>>> from solprop.processing._examples import example_output
>>> from solprop.processing.fields import get_field
>>> t, field = get_field(example_output(), plane=0, kind="real")
>>> bool(np.isrealobj(field))
True
"""
_check_choice(kind, FIELD_KINDS, "kind")
grid, t, amplitude = _plane_amplitude(
output,
plane=plane,
z=z,
oversampling=oversampling,
bandpass=bandpass,
transfer=transfer,
transform_limit=transform_limit,
)
amplitude, kept = _apply_modes(amplitude, modes, allow_sum=False)
# The stored envelope has its carrier removed and the stored real field does not, so
# exactly one of the two representations needs the phasor, whichever grid this is.
if (grid.kind == "envelope") != (kind == "envelope"):
sign = 1.0 if grid.kind == "envelope" else -1.0
amplitude = amplitude * np.exp(sign * 1j * grid.omega0 * t)
if kind == "real":
return t, _squeeze(np.real(amplitude), mode_axis=kept)
if convention == "optics":
amplitude = fourier.to_optics(amplitude)
return t, _squeeze(amplitude, mode_axis=kept)
[docs]
def get_power( # noqa: PLR0913 - a plane, a mode set, a window and an optical element
output: Output,
*,
plane: int | None = None,
z: float | None = None,
modes: ModeSelection = None,
oversampling: int = 1,
bandpass: spectra.Bandpass | None = None,
transfer: TransferFunction | None = None,
transform_limit: bool = False,
) -> tuple[Float[np.ndarray, " n_t"], Float[np.ndarray, "*shape n_t"]]:
"""Return the instantaneous power of one plane, or of every plane.
``P(t) = sum_pol abs(E_a(t)) ** 2`` in watts, the quantity every temporal metric and
every temporal plot is built on. Polarisation components are summed because they are
two components of one field whose powers add at each instant, which is the reduction
:func:`solprop.stats.peak_power` uses.
Parameters
----------
output : solprop.output.schema.Output
An opened output handler.
plane : int or None
Index of the plane wanted, negative from the end; see :func:`get_spectrum`.
z : float or None
Position of the plane wanted, m.
modes : int, sequence of int, 'sum' or None
Which modes to return; ``'sum'`` adds their powers.
oversampling : int
Factor by which to resample the time axis, at least 1.
bandpass : tuple of float, numpy.ndarray or None
Wavelength limits, m, of an explicit window applied before the transform, or a
ready-made window on the stored axis.
transfer : TransferFunction or None
A linear optical element of :mod:`solprop.optics` applied to the spectrum before
the transform -- a compressor, a window, a stretch of air -- in the optics
convention.
transform_limit : bool
Whether to discard the spectral phase first, giving the shortest pulse the
spectrum supports, centred on ``t = 0``. Applied after
``bandpass`` and before ``transfer``.
Returns
-------
t : numpy.ndarray
Time axis, shape ``(n_t,)`` or ``(oversampling n_t,)``, s.
power : numpy.ndarray
Instantaneous power, W, with trivial axes squeezed.
Examples
--------
>>> import numpy as np
>>> from solprop.processing._examples import example_output
>>> from solprop.processing.fields import get_power
>>> t, power = get_power(example_output(), plane=0)
>>> bool(np.all(power >= 0.0))
True
"""
_, t, amplitude = _plane_amplitude(
output,
plane=plane,
z=z,
oversampling=oversampling,
bandpass=bandpass,
transfer=transfer,
transform_limit=transform_limit,
)
power = np.sum(np.abs(amplitude) ** 2, axis=-2, keepdims=True)
power, kept = _apply_modes(power, modes, allow_sum=True)
return t, _squeeze(power, mode_axis=kept)
[docs]
def get_phase(
output: Output,
*,
plane: int | None = None,
z: float | None = None,
modes: ModeSelection = None,
band: Band = "active",
) -> tuple[Float[np.ndarray, " n_x"], Float[np.ndarray, "*shape n_x"]]:
"""Return the unwrapped spectral phase of one plane, or of every plane.
The phase of the optics-convention spectrum, unwrapped along the frequency axis.
The optics convention has already moved the phase origin from the first time sample
to the centre of the window (``conventions.md`` section 2), so no linear ramp
``omega T / 2`` is subtracted and a flat phase means a transform-limited pulse.
Parameters
----------
output : solprop.output.schema.Output
An opened output handler.
plane : int or None
Index of the plane wanted, negative from the end; see :func:`get_spectrum`.
z : float or None
Position of the plane wanted, m.
modes : int, sequence of int, 'sum' or None
Which modes to return. ``'sum'`` is refused: a phase is not additive.
band : {'active', 'all'}
Which part of the stored axis to use.
Returns
-------
omega : numpy.ndarray
Angular frequency, shape ``(n_x,)``, rad/s, ascending.
phase : numpy.ndarray
Unwrapped spectral phase, rad, with trivial axes squeezed.
Examples
--------
>>> import numpy as np
>>> from solprop.processing._examples import example_output
>>> from solprop.processing.fields import get_phase
>>> omega, phase = get_phase(example_output(), plane=0)
>>> bool(np.all(np.isfinite(phase)))
True
"""
omega, spectrum = get_spectrum(
output, plane=plane, z=z, modes=modes, band=band, convention="optics"
)
return omega, np.unwrap(np.angle(spectrum), axis=-1)
def _fft_norm(grid: TimeGrid) -> float:
"""Return the factor that turns a stored spectrum into a spectral amplitude.
The square root of the prefactor of :func:`solprop.fourier.spectral_density`, so
that the modulus squared of the scaled spectrum is the energy density in
J/(rad/s).
Parameters
----------
grid : solprop.grid.TimeGrid
The grid the spectrum lives on.
Returns
-------
float
``dt sqrt(2 / pi)`` on a field grid, ``dt / sqrt(2 pi)`` on an envelope grid.
"""
if grid.kind == "field":
return grid.dt * np.sqrt(2.0 / np.pi)
return grid.dt / np.sqrt(2.0 * np.pi)
def _to_time(
grid: TimeGrid,
stored: Complex[np.ndarray, "*rest n_omega"],
*,
oversampling: int,
bandpass: spectra.Bandpass | None,
transfer: TransferFunction | None = None,
transform_limit: bool = False,
) -> tuple[Float[np.ndarray, " n_t"], Complex[np.ndarray, "*rest n_t"]]:
"""Return the time axis and the amplitude whose modulus squared is the power.
That amplitude is the analytic signal on a field grid and the carrier-free envelope
on an envelope grid -- the same pair :func:`solprop.stats` works with, so that a
power computed here is bit-identical to a power computed there. Restoring the
carrier is the caller's business, because doing it unconditionally would cost a
complex multiply and the last bit of every power.
Parameters
----------
grid : solprop.grid.TimeGrid
The grid the spectrum lives on.
stored : numpy.ndarray
Stored spectrum, internal convention, shape ``(..., n_omega)``.
oversampling : int
Factor by which to resample the time axis, at least 1.
bandpass : tuple of float, numpy.ndarray or None
Wavelength limits, m, of an explicit window applied before the transform, or a
ready-made window on the stored axis.
transfer : TransferFunction or None
An optics-convention element applied to the spectrum, after the band pass and
the transform limit.
transform_limit : bool
Whether to discard the spectral phase first.
Returns
-------
t : numpy.ndarray
Time axis, s, centred on zero.
amplitude : numpy.ndarray
The complex amplitude, internal convention, in the square root of watts.
Raises
------
solprop.exceptions.ProcessingError
If ``oversampling`` is below one.
"""
if oversampling < 1:
msg = f"oversampling must be at least 1, got {oversampling}"
raise ProcessingError(msg)
if bandpass is not None:
stored = stored * spectra._as_window(grid, bandpass)
if transform_limit:
# Zero spectral phase about the window centre, which in a stored spectrum is
# the alternating sign of the half-window origin shift (conventions section 2).
stored = fourier.spectrum_to_internal(
np.abs(stored).astype(np.complex128), kind=grid.kind
)
if transfer is not None:
stored = stored * stored_factor(grid, transfer)
n_t = oversampling * grid.n_t
if oversampling > 1:
stored = fourier.oversample(stored, n_t, kind=grid.kind)
samples = fourier.inverse(stored, kind=grid.kind)
amplitude = fourier.analytic_signal(samples) if grid.kind == "field" else samples
t = (np.arange(n_t, dtype=float) - n_t / 2) * (grid.dt / oversampling)
return t, amplitude
def _plane_amplitude( # noqa: PLR0913 - the plane and everything done to its spectrum
output: Output,
*,
plane: int | None,
z: float | None,
oversampling: int,
bandpass: spectra.Bandpass | None,
transfer: TransferFunction | None = None,
transform_limit: bool = False,
) -> tuple[TimeGrid, Float[np.ndarray, " n_t"], Complex[np.ndarray, "*planes n_t"]]:
"""Return the grid, the time axis and the amplitude of the selected planes.
The one path from a saved run to the time domain, so that a power, a field and an
autocorrelation all come from the same transform with the same normalisation.
Parameters
----------
output : solprop.output.schema.Output
An opened output handler.
plane : int or None
Plane index, negative from the end, or ``None`` for every plane.
z : float or None
Plane position, m, or ``None``.
oversampling : int
Factor by which to resample the time axis, at least 1.
bandpass : tuple of float, numpy.ndarray or None
Wavelength limits, m, of an explicit window applied before the transform, or a
ready-made window on the stored axis.
transfer : TransferFunction or None
An optics-convention element applied to the spectrum.
transform_limit : bool
Whether to discard the spectral phase first.
Returns
-------
grid : solprop.grid.TimeGrid
The rebuilt grid.
t : numpy.ndarray
Time axis, s, centred on zero.
amplitude : numpy.ndarray
Complex amplitude with every axis kept, shape ``(..., n_modes, n_pol, n_t)``, in
the internal convention and the square root of watts.
"""
grid = grid_of(output)
stored = _plane_field(output, plane, z, convention="fft")
t, amplitude = _to_time(
grid,
stored,
oversampling=oversampling,
bandpass=bandpass,
transfer=transfer,
transform_limit=transform_limit,
)
return grid, t, amplitude