# SPDX-FileCopyrightText: 2026 John C. Travers and contributors
#
# SPDX-License-Identifier: GPL-3.0-or-later
"""Freeze spatial geometry and separable coefficients for the kernel contract.
These records contain arrays and scalars only. Geometry is independent of the
temporal grid. A separable operator keeps frequency-only material coefficients and
one transverse-frequency square per column, avoiding a full coefficient volume.
Kernel schema 15 adds the hybrid planar kind ``"slab"`` (slab modes in x, an FFT in
y) and the one per-channel table it needs, ``kappa_sq``.
"""
import math
from dataclasses import dataclass, replace
from typing import Literal
import numpy as np
from jaxtyping import Bool, Complex, Float
from solprop.adiabatic import AdiabaticAdmission
from solprop.constants import C
from solprop.exceptions import OutputError, TranslationError
from solprop.grid import FreeGrid, RadialGrid, TimeGrid
from solprop.spatial import (
ABSORBED_CHANNELS,
FAR_FIELD_STATISTICS,
ApertureOutline,
ApertureTables,
aperture_outline,
aperture_scale,
attenuation_extremes,
)
from solprop.stats import ApertureSet, StatsSpec
from solprop.translate.chebyshev import PiecewiseChebyshev
from solprop.translate.masks import Grids
__all__ = [
"FluxPlan",
"RateBounds",
"SeparableOperator",
"SpatialGeometry",
"aperture_tables",
"axial_admittance",
"axial_factors",
"longitudinal_root",
"rate_bounds",
"separable_values",
]
[docs]
@dataclass(frozen=True, eq=False)
class SpatialGeometry:
"""Carry spatial axes, quadrature, reconstruction and absorption across the seam.
Attributes
----------
kind : {'radial', 'free', 'slab'}
Radial QDHT, Cartesian FFT, or the hybrid planar transform: slab modes in x
(the transform's own synthesis and projection on the slab rule) times an FFT
in y. Modal transforms carry no such record.
shape : tuple of int
``(n_r,)`` or ``(n_x, n_y)``: the **physical** node grid. Cartesian column
``ix + n_x*iy`` has x fastest. Every entry is at least 2, except ``n_y = 1``
for a line uniform in y, whose one y sample is ``y = k_y = 0`` with unit
measure (per metre of y). For ``"slab"``, ``n_x`` counts the slab rule's x
nodes and the state's columns are ``(p, k_y)``, ``p + n_p*iy`` with the mode
index fastest, so ``n_modes = n_p*n_y`` while ``n_nodes = n_x*n_y``.
axes : tuple of numpy.ndarray
``(r,)`` or ``(x, y)``, float64 in m, strictly increasing. A slab's x axis is
its quadrature nodes inside the gap, not uniformly spaced.
spectral_axes : tuple of numpy.ndarray
``(k_r,)`` or ``(k_x, k_y)``, in rad/m; for ``"slab"`` ``(k_y,)`` alone, the
x direction's spectrum being the mode index. Cartesian axes are in FFT order.
spatial_weights : numpy.ndarray
Physical area weights in m^2, shape ``(n_nodes,)`` in node order.
spectral_weights : numpy.ndarray
Parseval weights for the stored transverse spectrum, shape ``(n_modes,)``.
Units are m^-2 for radial Hankel spectra and m^2 for Cartesian raw DFTs. A
slab column is a modal amplitude in sqrt(W/m) and its weight ``2 dy / (eps0
c n_y)`` is the raw DFT's ``dy/n_y`` divided by ``eps0 c / 2``, in m ohm, so
that the common ``eps0 c / 2`` energy measure gives joules on every kind.
coordinates : numpy.ndarray
``(r, 0)`` or ``(x, y)`` per physical node, in m, shape ``(n_nodes, 2)``.
window : numpy.ndarray
Spatial collar window in [0, 1], shape ``(n_nodes,)``. For ``"slab"`` it
is the y collar alone, the same at every x node of a row.
absorber_rate : numpy.ndarray
In-stage spatial power absorption rate, in 1/m, shape ``(n_nodes,)``; for
``"slab"`` a function of y alone, like `window`.
on_axis : numpy.ndarray
Linear reconstruction weights from transverse spectral columns to the axis,
complex128, shape ``(n_modes,)``. Temporal oversampling is not folded in.
forward_scale : float
Common transverse forward scaling; m^2 radially, dimensionless for the FFT.
inverse_scale : float
Common transverse inverse scaling; m^-2 radially, dimensionless for the FFT
(``1/n_y`` for the slab's y transform).
Raises
------
TranslationError
If dimensions, axes, measures, collar values or scalings disagree.
"""
kind: Literal["radial", "free", "slab"]
shape: tuple[int, ...]
axes: tuple[Float[np.ndarray, " ..."], ...]
spectral_axes: tuple[Float[np.ndarray, " ..."], ...]
spatial_weights: Float[np.ndarray, " n_area"]
spectral_weights: Float[np.ndarray, " n_spectral_area"]
coordinates: Float[np.ndarray, "n_points 2"]
window: Float[np.ndarray, " n_window"]
absorber_rate: Float[np.ndarray, " n_rate"]
on_axis: Complex[np.ndarray, " n_axis_columns"]
forward_scale: float
inverse_scale: float
def __post_init__(self) -> None:
"""Own immutable arrays and check their common spatial shape."""
dimension = 1 if self.kind == "radial" else 2
if self.kind not in ("radial", "free", "slab") or len(self.shape) != dimension:
msg = (
"spatial kind and shape must be radial (n_r,), free (n_x, n_y) or "
"slab (n_x, n_y)"
)
raise TranslationError(msg)
# A Cartesian line uniform in y (solprop.grid.LineGrid) is the one
# exception to two samples per axis: n_y = 1, carrying only k_y = 0.
if any(
isinstance(size, bool)
or not isinstance(size, int)
or size < (1 if axis == 1 and self.kind == "free" else 2)
for axis, size in enumerate(self.shape)
):
msg = (
"spatial shape entries must be integers of at least 2, "
"except a Cartesian line's n_y = 1"
)
raise TranslationError(msg)
for name in ("axes", "spectral_axes"):
axes = tuple(_frozen_array(one, np.float64) for one in getattr(self, name))
# A slab's x spectrum is its mode index, carried by the transform's
# matrices, so its only spectral axis is k_y.
sizes = (
self.shape[1:]
if self.kind == "slab" and name == "spectral_axes"
else self.shape
)
if len(axes) != len(sizes) or any(
one.shape != (size,) for one, size in zip(axes, sizes, strict=True)
):
msg = f"spatial {name} must match shape {self.shape}"
raise TranslationError(msg)
object.__setattr__(self, name, axes)
for axis in self.axes:
if np.any(np.diff(axis) <= 0):
msg = "physical spatial axes must be strictly increasing"
raise TranslationError(msg)
self._adopt_vectors()
if self.kind == "slab":
self._check_slab_columns()
self._check_coordinates()
if np.any(self.spatial_weights <= 0) or np.any(self.spectral_weights <= 0):
msg = "spatial and spectral integration weights must be positive"
raise TranslationError(msg)
if np.any((self.window < 0) | (self.window > 1)) or np.any(
self.absorber_rate < 0
):
msg = "spatial window must be in [0, 1] and absorber_rate nonnegative"
raise TranslationError(msg)
if any(
not math.isfinite(value) or value <= 0
for value in (self.forward_scale, self.inverse_scale)
):
msg = "spatial transform scalings must be finite and positive"
raise TranslationError(msg)
def _adopt_vectors(self) -> None:
"""Own the node and column vectors after checking their lengths."""
size = math.prod(self.shape)
for name in (
"spatial_weights",
"spectral_weights",
"coordinates",
"window",
"absorber_rate",
"on_axis",
):
values = _frozen_array(
getattr(self, name), np.complex128 if name == "on_axis" else np.float64
)
expected = (size, 2) if name == "coordinates" else (size,)
if self.kind == "slab" and name in ("spectral_weights", "on_axis"):
# The state's columns are (p, k_y), not the nodes: n_p*n_y of them.
expected = (values.size,)
if values.ndim != 1 or values.size % self.shape[1] != 0:
msg = (
f"slab {name} must be a vector of n_p*n_y column values, "
f"a multiple of n_y = {self.shape[1]}, got {values.shape}"
)
raise TranslationError(msg)
if values.shape != expected:
msg = f"spatial {name} must have shape {expected}, got {values.shape}"
raise TranslationError(msg)
object.__setattr__(self, name, values)
[docs]
@property
def n_columns(self) -> int:
"""Return the number of state columns the geometry describes.
Returns
-------
int
``n_r`` or ``n_x*n_y`` for the radial and Cartesian kinds, where columns
and nodes coincide, and ``n_p*n_y`` for a slab, where they do not.
"""
return int(self.spectral_weights.size)
def _check_slab_columns(self) -> None:
"""Check a slab's per-column and per-row arrays agree with its two shapes."""
if self.on_axis.size != self.spectral_weights.size:
msg = "slab on_axis and spectral_weights must both have n_p*n_y entries"
raise TranslationError(msg)
rows = (self.shape[1], self.shape[0])
# The plates carry no collar: the window and the rate are the y collar's,
# the same at every x node of a row, so a backend may apply them per (p, y).
for name in ("window", "absorber_rate"):
values = getattr(self, name).reshape(rows)
if np.any(values != values[:, :1]):
msg = f"slab {name} must depend on y alone (constant along each x row)"
raise TranslationError(msg)
def _check_coordinates(self) -> None:
"""Validate the stated ordering and the transverse spectral meaning."""
epsilon = np.finfo(float).eps
if self.kind == "radial":
if (
np.any(self.axes[0] <= 0)
or np.any(self.spectral_axes[0] <= 0)
or np.any(np.diff(self.spectral_axes[0]) <= 0)
):
msg = "radial spatial and spectral axes must be positive and increasing"
raise TranslationError(msg)
expected = np.column_stack((self.axes[0], np.zeros(self.shape[0])))
else:
pairs = (
zip(self.axes[1:], self.spectral_axes, strict=True)
if self.kind == "slab"
else zip(self.axes, self.spectral_axes, strict=True)
)
for axis, spectral in pairs:
if axis.size == 1:
# The line's y: one sample at the origin carrying only k_y = 0.
if axis[0] != 0 or spectral[0] != 0:
msg = "a Cartesian line's single y sample must be y = k_y = 0"
raise TranslationError(msg)
continue
spacing = (axis[-1] - axis[0]) / (axis.size - 1)
error = 32 * epsilon * max(float(np.max(np.abs(axis))), spacing)
if np.any(np.abs(np.diff(axis) - spacing) > error):
msg = "Cartesian physical axes must be uniformly spaced"
raise TranslationError(msg)
expected_frequency = 2 * np.pi * np.fft.fftfreq(axis.size, spacing)
spectral_error = (
64 * epsilon * max(float(np.max(np.abs(expected_frequency))), 1.0)
)
if np.any(np.abs(spectral - expected_frequency) > spectral_error):
msg = (
"Cartesian spectral axes must be in FFT order "
"with reciprocal spacing"
)
raise TranslationError(msg)
expected = np.column_stack(
(
np.tile(self.axes[0], self.shape[1]),
np.repeat(self.axes[1], self.shape[0]),
)
)
if not np.array_equal(self.coordinates, expected):
msg = "spatial coordinates must match the axes with x fastest (ix + n_x*iy)"
raise TranslationError(msg)
count = (
1
if self.kind == "radial"
else self.shape[1]
if self.kind == "slab"
else math.prod(self.shape)
)
if abs(count * self.forward_scale * self.inverse_scale - 1) > 32 * epsilon:
msg = (
"spatial scalings must normalize the QDHT or the unnormalised FFT pair"
)
raise TranslationError(msg)
[docs]
@dataclass(frozen=True, eq=False)
class SeparableOperator:
"""Store a constant spatial operator and prefactor without coefficient volumes.
Attributes
----------
k_squared : numpy.ndarray
Shared complex material wave-number square in m^-2, shape ``(n_omega,)``.
kperp_sq : numpy.ndarray
Transverse wave-number square in m^-2, shape ``(n_modes,)``, nonnegative.
k_squared_low : numpy.ndarray
Complex residual of the material square in m^-2, shape ``(n_omega,)``.
kperp_sq_low : numpy.ndarray
Real residual of the transverse square in m^-2, shape ``(n_modes,)``.
frame : numpy.ndarray
Shared real frame subtraction in m^-1, shape ``(n_omega,)``.
spectral_rate : numpy.ndarray
Shared spectral power absorption in m^-1, shape ``(n_omega,)``, nonnegative.
prefactor_numerator : numpy.ndarray
Shared purely imaginary numerator of the nonlinear prefactor, with the
temporal crop scale included, shape ``(n_omega,)``.
prefactor_clamped : numpy.ndarray
Shared prefactor when the real longitudinal square is nonpositive, with the
temporal crop scale included, shape ``(n_omega,)``.
active : numpy.ndarray
Physical frequency mask, bool, shape ``(n_omega,)``.
evanescent_rate : float
Largest evanescent field attenuation rate in m^-1, positive and finite.
n_pol : int
State polarisation components, 1 or 2; every coefficient is shared by them.
z_span : tuple of float
Valid propagation interval in m, finite and increasing.
kappa_sq : numpy.ndarray or None
Kernel schema 15, the hybrid planar kind only: each slab channel's complex
transverse wave-number square ``kappa_p(omega)**2 = k_squared -
beta_p(omega)**2`` in m^-2, shape ``(n_channels, n_omega)``, where
``beta_p = omega n_eff,p / c`` is the mode's own complex propagation
constant (the wall term under the mode's loss policy). Column ``c`` belongs
to channel ``c % n_channels`` (the channel index fastest). Zero at inactive
frequencies. ``None`` (the default) for every other kind, which leaves the
operator exactly as before.
kperp_rate : numpy.ndarray or None
Kernel schema 31: the transverse-wavevector power absorption rate
``alpha_perp`` of a run's safety ramp (SEP 40), in
m^-1, one per column, shape ``(n_modes,)``, finite and nonnegative. It enters
the operator on the active band only, beside the spectral rate,
``D -= alpha_perp / 2`` (SEP 22 section 5.3), and never the unitary phase.
``None`` (the default) is stored as zeros, whose operator is today's to the
last bit. On the hybrid planar kind it acts per ``(p, k_y)`` column (SEP 40).
user_rate : numpy.ndarray or None
The user's share of `spectral_rate`, in m^-1, shape ``(n_omega,)``; only
``absorbed_power`` reads it, since the operator takes the total. ``None`` (the
default) is stored as zeros.
Raises
------
TranslationError
If arrays, scales, frequency counts or the valid interval disagree, the
transverse rate is negative or misshaped, or the user's share is not a share of
the spectral rate.
"""
k_squared: Complex[np.ndarray, " n_k"]
kperp_sq: Float[np.ndarray, " n_transverse"]
k_squared_low: Complex[np.ndarray, " n_k_low"]
kperp_sq_low: Float[np.ndarray, " n_transverse_low"]
frame: Float[np.ndarray, " n_frame"]
spectral_rate: Float[np.ndarray, " n_spectral_rate"]
prefactor_numerator: Complex[np.ndarray, " n_numerator"]
prefactor_clamped: Complex[np.ndarray, " n_clamped"]
active: Bool[np.ndarray, " n_active"]
evanescent_rate: float
n_pol: int
z_span: tuple[float, float]
kappa_sq: Complex[np.ndarray, "n_channels n_kappa"] | None = None
kperp_rate: Float[np.ndarray, " n_kperp_rate"] | None = None
user_rate: Float[np.ndarray, " n_user_rate"] | None = None
def __post_init__(self) -> None:
"""Own immutable coefficients and validate the separable representation."""
for name in (
"k_squared",
"kperp_sq",
"k_squared_low",
"kperp_sq_low",
"frame",
"spectral_rate",
"prefactor_numerator",
"prefactor_clamped",
"active",
):
dtype = (
np.bool_
if name == "active"
else np.complex128
if name
in (
"k_squared",
"k_squared_low",
"prefactor_numerator",
"prefactor_clamped",
)
else np.float64
)
values = _frozen_array(getattr(self, name), dtype)
if values.ndim != 1 or values.size == 0:
msg = f"separable {name} must be a nonempty vector"
raise TranslationError(msg)
if name == "kperp_sq_low" and values.size != self.kperp_sq.size:
msg = "separable kperp_sq_low must match the transverse count"
raise TranslationError(msg)
if (
name not in ("kperp_sq", "kperp_sq_low")
and values.size != self.k_squared.size
):
msg = f"separable {name} must match k_squared's frequency count"
raise TranslationError(msg)
object.__setattr__(self, name, values)
if np.any(self.kperp_sq < 0) or np.any(self.spectral_rate < 0):
msg = "separable kperp_sq and spectral_rate must be nonnegative"
raise TranslationError(msg)
if np.any(np.real(self.prefactor_numerator) != 0) or np.any(
np.real(self.prefactor_clamped) != 0
):
msg = "separable prefactors must be purely imaginary"
raise TranslationError(msg)
if self.n_pol not in (1, 2) or isinstance(self.n_pol, bool):
msg = "separable n_pol must be 1 or 2"
raise TranslationError(msg)
if not math.isfinite(self.evanescent_rate) or self.evanescent_rate <= 0:
msg = "separable evanescent_rate must be finite and positive"
raise TranslationError(msg)
if (
len(self.z_span) != 2
or not all(math.isfinite(value) for value in self.z_span)
or not self.z_span[0] < self.z_span[1]
):
msg = "separable z_span must be finite and increasing"
raise TranslationError(msg)
if self.kappa_sq is not None:
kappa = _frozen_array(self.kappa_sq, np.complex128)
if (
kappa.ndim != 2
or kappa.shape[0] == 0
or kappa.shape[1] != self.k_squared.size
or self.kperp_sq.size % kappa.shape[0] != 0
):
msg = (
"separable kappa_sq must be (n_channels, n_omega), with n_omega "
"matching k_squared and n_channels dividing the column count"
)
raise TranslationError(msg)
if np.any(kappa[:, ~self.active] != 0):
msg = "separable kappa_sq must vanish at inactive frequencies"
raise TranslationError(msg)
object.__setattr__(self, "kappa_sq", kappa)
self._adopt_kperp_rate()
self._adopt_user_rate()
def _adopt_user_rate(self) -> None:
"""Own the user's share of the spectral rate, zeros when absent, and check it.
Raises
------
TranslationError
If it is not a finite share of the spectral rate on every bin.
"""
rate = (
np.zeros(self.spectral_rate.size)
if self.user_rate is None
else _frozen_array(self.user_rate, np.float64)
)
rate.setflags(write=False)
if (
rate.shape != self.spectral_rate.shape
or np.any(rate < 0)
or np.any(rate > self.spectral_rate)
):
msg = (
"separable user_rate must be the user's share of spectral_rate, "
f"0 <= user_rate <= spectral_rate, shape {self.spectral_rate.shape}"
)
raise TranslationError(msg)
object.__setattr__(self, "user_rate", rate)
def _adopt_kperp_rate(self) -> None:
"""Own the transverse rate, zeros when absent, and validate it.
Raises
------
TranslationError
If it is not one finite nonnegative rate per column.
"""
rate = (
np.zeros(self.kperp_sq.size)
if self.kperp_rate is None
else _frozen_array(self.kperp_rate, np.float64)
)
rate.setflags(write=False)
if rate.shape != self.kperp_sq.shape or np.any(rate < 0):
msg = (
"separable kperp_rate must be one nonnegative rate per transverse "
f"column, shape {self.kperp_sq.shape}"
)
raise TranslationError(msg)
# On the hybrid planar kind the rate acts on k_y alone, per (p, k_y) column,
# which is diagonal there (SEP 40 section 5.4).
object.__setattr__(self, "kperp_rate", rate)
[docs]
@property
def n_channels(self) -> int:
"""Return the number of slab channels the columns cycle through.
Returns
-------
int
The rows of `kappa_sq`, or 1 when it is absent.
"""
return 1 if self.kappa_sq is None else int(self.kappa_sq.shape[0])
[docs]
@property
def value_shape(self) -> tuple[int, int, int, int]:
"""Return the full state shape represented by these shared coefficients.
Returns
-------
tuple of int
``(1, n_modes, n_pol, n_omega)`` without allocating that volume.
"""
return (1, self.kperp_sq.size, self.n_pol, self.k_squared.size)
[docs]
@dataclass(frozen=True)
class RateBounds:
"""Bounds on a separable operator's real part, for the exponent budget.
SEP 22 section 4.5: the decaying factor ``exp(lambda_minus h)`` of the inverse
propagator and the growing one ``exp(lambda_plus h)`` of a gain medium overflow
once their exponent passes ``log(DBL_MAX)``. Each bound is the largest over the
frequencies of a sufficient per-frequency bound, and the terms are those of the
frequency where the loss bound is largest.
Attributes
----------
loss : float
``lambda_minus``, a bound on ``max(-Re D)`` in m^-1, nonnegative.
gain : float
``lambda_plus``, a bound on ``max(Re D)`` in m^-1, nonnegative.
propagating : float
The largest propagating attenuation at the binding frequency, m^-1.
evanescent : float
The capped evanescent attenuation there, ``min(|Im k_z|, evanescent_rate)``.
spectral : float
Half the spectral rate there.
transverse : float
Half the transverse collar's largest rate, applied on active frequencies.
"""
loss: float
gain: float
propagating: float
evanescent: float
spectral: float
transverse: float
[docs]
def rate_bounds(coefficients: SeparableOperator) -> RateBounds:
"""Bound the real part of a separable operator without forming it.
The search of :func:`solprop.spatial.attenuation_extremes` per slab channel (one
on every kind but the hybrid planar), so the cost is
``O(n_omega log n_modes)``. On an active frequency the loss bound is the larger
of the propagating attenuation's positive part and the capped evanescent one,
plus half the largest transverse rate, plus half the spectral rate; on an
inactive one it is half the spectral rate alone (SEP 22, Q6). The gain bound is
the negated smallest propagating attenuation where it is negative.
Parameters
----------
coefficients : SeparableOperator
The operator.
Returns
-------
RateBounds
The bounds and the binding frequency's terms.
Examples
--------
>>> import numpy as np
>>> from solprop.translate.spatial import SeparableOperator, rate_bounds
>>> operator = SeparableOperator(
... k_squared=np.array([0j, 4 + 0j]), kperp_sq=np.array([0.0, 9.0]),
... k_squared_low=np.zeros(2, complex), kperp_sq_low=np.zeros(2),
... frame=np.zeros(2), spectral_rate=np.array([6.0, 0.0]),
... prefactor_numerator=np.zeros(2, complex),
... prefactor_clamped=np.zeros(2, complex),
... active=np.array([False, True]), evanescent_rate=2.0, n_pol=1,
... z_span=(0.0, 1.0), kperp_rate=np.array([0.0, 4.0]),
... )
>>> bounds = rate_bounds(operator)
>>> bounds.loss, bounds.evanescent, bounds.transverse
(4.0, 2.0, 2.0)
"""
data = coefficients
active = data.active
n_omega = data.k_squared.size
largest = np.full(n_omega, -np.inf)
smallest = np.full(n_omega, np.inf)
evanescent = np.zeros(n_omega)
columns = np.arange(data.kperp_sq.size)
for channel in range(data.n_channels):
mine = columns[columns % data.n_channels == channel]
square = (
data.k_squared
if data.kappa_sq is None
else data.k_squared - data.kappa_sq[channel]
)
high, low, ev = attenuation_extremes(
square,
data.kperp_sq[mine],
square_low=data.k_squared_low,
kperp_sq_low=data.kperp_sq_low[mine],
)
largest = np.maximum(largest, high)
smallest = np.minimum(smallest, low)
evanescent = np.maximum(evanescent, ev)
propagating = np.where(active, np.maximum(largest, 0.0), 0.0)
capped = np.where(active, np.minimum(evanescent, data.evanescent_rate), 0.0)
assert data.kperp_rate is not None # adopted as zeros
transverse = 0.5 * float(np.max(data.kperp_rate, initial=0.0))
loss = (
np.maximum(propagating, capped)
+ np.where(active, transverse, 0.0)
+ 0.5 * data.spectral_rate
)
binding = int(np.argmax(loss))
gain = np.where(active & np.isfinite(smallest), np.maximum(-smallest, 0.0), 0.0)
return RateBounds(
loss=float(loss[binding]),
gain=float(np.max(gain, initial=0.0)),
propagating=float(propagating[binding]),
evanescent=float(capped[binding]),
spectral=float(0.5 * data.spectral_rate[binding]),
transverse=transverse if bool(active[binding]) else 0.0,
)
def _frozen_array(values: np.ndarray, dtype: type) -> np.ndarray:
"""Own a finite, contiguous array without freezing a caller's writable buffer.
Parameters
----------
values : numpy.ndarray
Numerical coefficients or geometry values to copy.
dtype : type
Required seam dtype: float64, complex128 or bool.
Returns
-------
numpy.ndarray
Finite, contiguous, read-only copy with the input shape.
"""
result = np.array(values, dtype=dtype, order="C", copy=True)
if not np.all(np.isfinite(result)):
msg = "spatial contract arrays must be finite"
raise TranslationError(msg)
result.setflags(write=False)
return result
[docs]
def separable_values(
coefficients: SeparableOperator,
columns: slice | None = None,
) -> tuple[
Complex[np.ndarray, "1 n_modes n_pol n_omega"],
Complex[np.ndarray, "n_modes n_pol n_omega"],
]:
"""Evaluate a separable operator and prefactor for the host reference.
The reference may select bounded transverse columns. Compiled backends
evaluate the same formula per element and retain only shared vectors.
Parameters
----------
coefficients : SeparableOperator
Constant coefficients, already validated and in the internal convention.
columns : slice or None, optional
Transverse columns to evaluate; None selects all columns.
Returns
-------
tuple of numpy.ndarray
Linear rate in m^-1, shape ``(1, n_modes, n_pol, n_omega)``, and nonlinear
prefactor, shape ``(n_modes, n_pol, n_omega)``.
"""
root, propagating, attenuation = longitudinal_root(coefficients, columns)
transverse_rate = _kperp_rate(coefficients, columns)[:, None]
# The transverse rate joins the field attenuation inside the active factor (SEP
# 22, Q6), so an inactive bin's operator is today's whatever the collar; a zero
# rate leaves every element as it was, signed zeros included: -a - 0 = -a.
operator = (
(-attenuation - 0.5 * transverse_rate)
- 1j * (np.real(root) - coefficients.frame)
) * coefficients.active - 0.5 * coefficients.spectral_rate
prefactor = np.broadcast_to(coefficients.prefactor_clamped, root.shape).copy()
np.divide(
coefficients.prefactor_numerator,
np.real(root),
out=prefactor,
where=propagating,
)
prefactor *= coefficients.active
shape = (root.shape[0], coefficients.n_pol, len(coefficients.k_squared))
return (
np.broadcast_to(operator[:, None, :], shape)[None],
np.broadcast_to(prefactor[:, None, :], shape),
)
def _kperp_rate(
coefficients: SeparableOperator, columns: slice | None
) -> Float[np.ndarray, " n_columns"]:
"""Return the transverse rate of the selected columns.
Parameters
----------
coefficients : SeparableOperator
The operator.
columns : slice or None
Transverse columns; None selects all.
Returns
-------
numpy.ndarray
The power rate in m^-1, one per selected column.
"""
rate = coefficients.kperp_rate
assert rate is not None # adopted as zeros in __post_init__
return rate if columns is None else rate[columns]
[docs]
def longitudinal_root(
coefficients: SeparableOperator,
columns: slice | None = None,
) -> tuple[
Complex[np.ndarray, "n_modes n_omega"],
Bool[np.ndarray, "n_modes n_omega"],
Float[np.ndarray, "n_modes n_omega"],
]:
"""Return the longitudinal wave number, its propagating test and its field rate.
The one evaluation of the root the stages use: the principal square root ``k_z``
of the compensated square ``q = (k^2 - k_perp^2) + (k_low^2 - k_perp,low^2)``
(with a slab channel's wall term subtracted last), the test ``Re q > 0`` that
separates propagating from evanescent columns, and the field attenuation,
``Im k_z`` where propagating and ``min(|Im k_z|, evanescent_rate)`` where not.
:func:`separable_values` builds the operator from it, and the far-field
statistics read it (:func:`axial_factors`), so both see the same root
(SEP 23 section 6).
Parameters
----------
coefficients : SeparableOperator
Constant coefficients in the internal convention.
columns : slice or None, optional
Transverse columns to evaluate; None selects all columns.
Returns
-------
tuple of numpy.ndarray
``k_z`` in m^-1 (complex), the propagating mask (bool) and the field
attenuation in m^-1, each of shape ``(n_columns, n_omega)``. Inactive bins
are not masked here: every consumer applies ``coefficients.active``.
"""
transverse = (
coefficients.kperp_sq if columns is None else coefficients.kperp_sq[columns]
)
transverse_low = (
coefficients.kperp_sq_low
if columns is None
else coefficients.kperp_sq_low[columns]
)
square = (coefficients.k_squared[None, :] - transverse[:, None]) + (
coefficients.k_squared_low[None, :] - transverse_low[:, None]
)
if coefficients.kappa_sq is not None:
# q = k^2 eps_co - kappa_p^2 - k_y^2, the wall term subtracted last so a
# free-space column's arithmetic is untouched (kernel schema 15).
indices = np.arange(coefficients.kperp_sq.size)
channels = (indices if columns is None else indices[columns]) % (
coefficients.n_channels
)
square = square - coefficients.kappa_sq[channels]
root = np.sqrt(square)
propagating = np.real(square) > 0
attenuation = np.where(
propagating,
np.imag(root),
np.minimum(np.abs(np.imag(root)), coefficients.evanescent_rate),
)
return root, propagating, attenuation
[docs]
def axial_admittance(
coefficients: SeparableOperator,
omega: Float[np.ndarray, " n_omega"],
columns: slice,
) -> Complex[np.ndarray, "n_chunk n_omega"]:
"""Return the dimensionless forward admittance ``c k_z / omega`` of a chunk.
The physical axial admittance is ``eps0 c`` times this value (SEP 34). It is the
root of :func:`longitudinal_root` on active propagating bins at positive
frequency, in the optics convention (positive imaginary part for loss), and zero
elsewhere: a nonpropagating or inactive column is not a forward power port.
Parameters
----------
coefficients : SeparableOperator
The plane's or the run's constant coefficients.
omega : numpy.ndarray
Physical angular frequencies in rad/s, shape ``(n_omega,)``.
columns : slice
The transverse columns of the chunk.
Returns
-------
numpy.ndarray
Complex admittance, shape ``(n_chunk, n_omega)``.
Examples
--------
>>> import numpy as np
>>> from solprop.constants import C
>>> from solprop.translate.spatial import SeparableOperator, axial_admittance
>>> omega = np.array([0.0, C, 2 * C])
>>> vacuum = SeparableOperator(
... k_squared=(omega / C) ** 2 + 0j, kperp_sq=np.array([0.0, 2.25]),
... k_squared_low=np.zeros(3, complex), kperp_sq_low=np.zeros(2),
... frame=np.zeros(3), spectral_rate=np.zeros(3),
... prefactor_numerator=np.zeros(3, complex),
... prefactor_clamped=np.zeros(3, complex),
... active=np.array([False, True, True]), evanescent_rate=200.0, n_pol=1,
... z_span=(0.0, 1.0),
... )
>>> np.round(axial_admittance(vacuum, omega, slice(None)).real, 12).tolist()
[[0.0, 1.0, 1.0], [0.0, 0.0, 0.661437827766]]
"""
root, propagating, _ = longitudinal_root(coefficients, columns)
live = propagating & coefficients.active[None, :] & (omega > 0)[None, :]
admittance = np.zeros(root.shape, dtype=np.complex128)
np.divide(C * root, omega[None, :], out=admittance, where=live)
return admittance
[docs]
def axial_factors(
coefficients: SeparableOperator,
omega: Float[np.ndarray, " n_omega"],
columns: slice,
) -> tuple[
Complex[np.ndarray, "n_chunk n_omega"], Float[np.ndarray, "5 n_chunk n_omega"]
]:
"""Return a chunk's admittance and absorption rates for the far-field statistics.
The weights of the far-field statistics (SEP 23 sections 4.8 and 4.9, measured
by SEP 34's axial flux), from :func:`longitudinal_root`. The admittance is
:func:`axial_admittance`. The power rates, in the order of
:data:`solprop.spatial.ABSORBED_CHANNELS`, are the operator's spectral rate on
every bin, as the stages apply it, less the user's share, its own channel, which
comes last; the transverse collar's rate on active bins
(zero without one); ``2 min(|Im k_z|, cap)`` on active evanescent bins; and
``2 Im k_z`` on active propagating bins, the material loss (negative for gain).
Together they are twice the negated real part of the operator; weighted by the
real admittance, their power is the physical flux the linear flow removes. The
evanescent channel, whose admittance is zero, removes no physical flux. The user's
share of the spectral rate is the fifth channel; the share and the rest sum to the
rate the operator applies, so the channels' sum is unmoved by the split.
Parameters
----------
coefficients : SeparableOperator
The run's constant coefficients.
omega : numpy.ndarray
Physical angular frequencies in rad/s, shape ``(n_omega,)``.
columns : slice
The transverse columns of the chunk.
Returns
-------
tuple of numpy.ndarray
The dimensionless complex admittance, shape ``(n_chunk, n_omega)``, and the
rates in m^-1, shape ``(5, n_chunk, n_omega)``.
"""
root, propagating, attenuation = longitudinal_root(coefficients, columns)
active = coefficients.active[None, :]
rates = np.zeros((len(ABSORBED_CHANNELS), *root.shape))
assert coefficients.user_rate is not None # adopted as zeros
rates[0] = (coefficients.spectral_rate - coefficients.user_rate)[None, :]
rates[4] = coefficients.user_rate[None, :]
# The transverse collar's rate acts on the active band, as the operator applies
# it (SEP 22, Q6).
rates[1] = np.where(active, _kperp_rate(coefficients, columns)[:, None], 0.0)
rates[2] = np.where(active & ~propagating, 2.0 * attenuation, 0.0)
rates[3] = np.where(active & propagating, 2.0 * attenuation, 0.0)
return axial_admittance(coefficients, omega, columns), rates
[docs]
@dataclass(frozen=True, eq=False)
class FluxPlan:
"""The plane-measure coefficients of a spatial run at any position (SEP 34).
Kernel schema 35. The run's statistics read the axial admittance of the medium
at the plane they measure through :meth:`at`: a uniform medium's is the run's
own separable coefficients, and a radial gas path's replaces the material square
by its z table, so the root is :func:`longitudinal_root`'s at that position.
Attributes
----------
base : SeparableOperator
Transverse squares, wall terms and active band, with the material square at
the entrance.
material : solprop.translate.chebyshev.PiecewiseChebyshev or None
The complex material square ``(omega n(omega, z) / c)**2`` in m^-2 along z,
value shape ``(n_omega,)``, for a radial medium that varies; None for a
uniform medium.
admission : solprop.adiabatic.AdiabaticAdmission or None
The varying path's adiabatic certificate (SEP 34 section 4.5); None for a
uniform medium, which needs none.
"""
base: SeparableOperator
material: PiecewiseChebyshev | None = None
admission: AdiabaticAdmission | None = None
[docs]
def at(self, z: float) -> SeparableOperator:
"""Return the coefficients of the plane at ``z``.
Parameters
----------
z : float
Position in m inside the translated span.
Returns
-------
SeparableOperator
The coefficients whose root is the plane's ``k_z``.
"""
if self.material is None:
return self.base
square = np.asarray(self.material(float(z)), np.complex128)
return replace(self.base, k_squared=np.where(self.base.active, square, 0))
def _far_field_columns(
geometry: "SpatialGeometry | FreeGrid | RadialGrid | None", # noqa: UP037
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""Return the Cartesian columns a far-field statistic reads, or refuse the geometry.
Parameters
----------
geometry : SpatialGeometry, solprop.grid.FreeGrid, solprop.grid.RadialGrid or None
The run's geometry record, or the grid of a saved plane.
Returns
-------
tuple of numpy.ndarray
The stored axes ``k_x`` and ``k_y`` in rad/m, the spectral weights and the
on-axis weights.
Raises
------
OutputError
For a radial or planar-hybrid geometry, or none at all.
"""
if isinstance(geometry, FreeGrid):
return geometry.k_x, geometry.k_y, geometry.spectral_weights, geometry.on_axis
if isinstance(geometry, SpatialGeometry) and geometry.kind == "free":
k_x, k_y = geometry.spectral_axes
return k_x, k_y, geometry.spectral_weights, geometry.on_axis
kind = (
"a radial grid"
if isinstance(geometry, RadialGrid)
or (isinstance(geometry, SpatialGeometry) and geometry.kind == "radial")
else "the hybrid planar kind, whose columns are slab modes,"
if isinstance(geometry, SpatialGeometry)
else "a problem without spatial geometry"
)
msg = (
f"the far-field statistics {FAR_FIELD_STATISTICS} need a Cartesian grid or a "
f"line, and {kind} has no Cartesian transverse wavevector to map onto a mask "
"plane"
)
raise OutputError(msg)
[docs]
def aperture_tables(
spec: StatsSpec,
grids: Grids | TimeGrid,
geometry: "SpatialGeometry | FreeGrid | RadialGrid | None", # noqa: UP037
) -> ApertureTables | None:
"""Form the far-field statistics' tables once, for every backend (SEP 23 §6).
The per-aperture constants -- each row's outline of
:func:`solprop.spatial.aperture_outline`, whose half-width under quadrature
weighting is the corrected ``D_q/2`` of a circle or an ellipse (unchanged on a
line, whose edges are straight, and on the straight edges of a rectangle or a
half-plane) -- and the scale table of :func:`solprop.spatial.aperture_scale`, with
the geometry's column axes, weights and on-axis vector and the spec's regions. On
a line every column has ``kappa_y = 0`` and the light is uniform in y, so a
rectangle has no y edges to meet: it is a slit of its x width, the row a circle of
that diameter forms. Every backend calls this at construction and
:func:`solprop.stats.collect` calls it for a saved plane, so all of them use the
same scale to the bit.
Parameters
----------
spec : solprop.stats.StatsSpec
The request and its far-field tables.
grids : solprop.translate.masks.Grids or solprop.grid.TimeGrid
Anything carrying the physical frequency axis ``omega`` in rad/s.
geometry : SpatialGeometry, solprop.grid.FreeGrid, solprop.grid.RadialGrid or None
The run's geometry record or a saved plane's grid.
Returns
-------
solprop.spatial.ApertureTables or None
The tables, or ``None`` when the spec names no far-field statistic.
Raises
------
OutputError
For a radial or planar-hybrid geometry or none at all; and on a line, for an
aperture whose centre is off ``y = 0``, a half-plane bounded in y or an ellipse,
which have no meaning uniform in y, or a region whose ``kappa_y`` interval
excludes zero, which would select nothing.
"""
if not set(spec.names) & set(FAR_FIELD_STATISTICS):
return None
k_x, k_y, weights, on_axis = _far_field_columns(geometry)
apertures = spec.apertures
regions = spec.spectral_regions
line = k_y.size == 1
if line:
if apertures is not None and np.any(np.asarray(apertures.centres)[:, 1] != 0):
msg = (
"on a line uniform in y an aperture is a slit along x, so its y centre "
"must be zero"
)
raise OutputError(msg)
if apertures is not None:
_refuse_off_line(apertures)
excluded = np.flatnonzero(~((regions[:, 2] <= 0) & (regions[:, 3] > 0)))
if excluded.size:
msg = (
"on a line every column has kappa_y = 0, so regions "
f"{excluded.tolist()}, whose kappa_y interval excludes zero, would "
"select nothing"
)
raise OutputError(msg)
if apertures is None:
empty = np.empty(0)
return ApertureTables(
k_x=k_x,
k_y=k_y,
weights=weights,
on_axis=on_axis,
regions=regions,
centres=np.empty((0, 2)),
half_widths=empty,
edges=empty,
distances=empty,
scale=np.empty((0, grids.omega.size)),
)
outlines = _outlines(apertures, line=line)
distances = np.asarray(apertures.distances)
return ApertureTables(
k_x=k_x,
k_y=k_y,
weights=weights,
on_axis=on_axis,
regions=regions,
centres=np.asarray(apertures.centres),
half_widths=np.array([one.half_width for one in outlines]),
edges=np.array([one.edge for one in outlines]),
distances=distances,
scale=aperture_scale(distances, grids.omega),
profiles=tuple(one.kind for one in outlines),
exponent=1 if apertures.weighting == "quadrature" else 2,
shapes=tuple(one.shape for one in outlines),
half_heights=np.array([one.half_height for one in outlines]),
stretches=np.array([one.stretch for one in outlines]),
)
def _refuse_off_line(apertures: ApertureSet) -> None:
"""Refuse the outlines a line uniform in y gives no meaning.
Parameters
----------
apertures : solprop.stats.ApertureSet
The set.
Raises
------
OutputError
For a half-plane bounded in y, which holds every column or none, or an
ellipse, which is an edge in y as well as in x.
"""
for row, shape in enumerate(apertures.row_shapes):
if shape == "ellipse" or shape.startswith("half_plane_y"):
msg = (
f"on a line every column has kappa_y = 0 and the light is uniform in "
f"y, so aperture {row}, a {shape}, has no meaning there; a line takes "
"a circle or a rectangle, each a slit along x, or a half-plane bounded "
"in x"
)
raise OutputError(msg)
def _outlines(apertures: ApertureSet, *, line: bool) -> list[ApertureOutline]:
"""Return each row's outline on a plane or a line.
Parameters
----------
apertures : solprop.stats.ApertureSet
The set.
line : bool
Whether the columns are a line's.
Returns
-------
list of solprop.spatial.ApertureOutline
One per aperture.
"""
count = apertures.n_apertures
diameters = np.asarray(apertures.diameters)
edges = np.zeros(count) if apertures.edges is None else np.asarray(apertures.edges)
extents = (
np.zeros((count, 2))
if apertures.extents is None
else np.asarray(apertures.extents)
)
corrected = apertures.weighting == "quadrature" and not line
outlines = []
for row in range(count):
shape = apertures.row_shapes[row]
diameter, extent = float(diameters[row]), extents[row]
if line and shape == "rectangle":
# sqrt(x*x) is |x| in correctly rounded arithmetic, so the slit decides a
# hard edge as a circle of this diameter does (SEP 23 section 4.10).
shape, diameter = "circle", float(extent[0])
outlines.append(
aperture_outline(
shape,
diameter,
(float(extent[0]), float(extent[1])),
float(edges[row]),
apertures.row_profiles[row],
corrected=corrected,
)
)
return outlines