meanas/meanas/fdtd/pml.py
2026-04-21 21:13:34 -07:00

279 lines
9 KiB
Python

"""
Convolutional perfectly matched layer (CPML) support for FDTD updates.
The helpers in this module construct per-face CPML parameters and then wrap the
standard Yee updates with the additional auxiliary `psi` fields needed by the
CPML recurrence.
The intended call pattern is:
1. Build a `cpml_params[axis][polarity_index]` table with `cpml_params(...)`.
2. Pass that table into `updates_with_cpml(...)` together with `dt`, `dxes`, and
`epsilon`.
3. Advance the returned `update_E` / `update_H` closures in the simulation loop.
Each face can be enabled or disabled independently by replacing one table entry
with `None`.
"""
# TODO retest pmls!
from typing import Any
from collections.abc import Callable, Sequence
from copy import deepcopy
import numpy
from numpy.typing import NDArray, DTypeLike
from ..fdmath import fdfield, dx_lists_t
from ..fdmath.functional import deriv_forward, deriv_back
__author__ = 'Jan Petykiewicz'
def cpml_params(
axis: int,
polarity: int,
dt: float,
thickness: int = 8,
ln_R_per_layer: float = -1.6,
epsilon_eff: float = 1,
mu_eff: float = 1,
m: float = 3.5,
ma: float = 1,
cfs_alpha: float = 0,
) -> dict[str, Any]:
"""
Construct the parameter block for one CPML face.
Args:
axis: Which Cartesian axis the CPML is normal to (`0`, `1`, or `2`).
polarity: Which face along that axis (`-1` for the low-index face,
`+1` for the high-index face).
dt: Timestep used by the Yee update.
thickness: Number of Yee cells occupied by the CPML region.
ln_R_per_layer: Logarithmic attenuation target per layer.
epsilon_eff: Effective permittivity used when choosing the CPML scaling.
mu_eff: Effective permeability used when choosing the CPML scaling.
m: Polynomial grading exponent for `sigma` and `kappa`.
ma: Polynomial grading exponent for the complex-frequency shift `alpha`.
cfs_alpha: Maximum complex-frequency shift parameter.
Returns:
Dictionary with:
- `param_e`: `(p0, p1, p2)` arrays for the E update
- `param_h`: `(p0, p1, p2)` arrays for the H update
- `region`: slice tuple selecting the CPML cells on that face
"""
if axis not in range(3):
raise ValueError(f'Invalid axis: {axis}')
if polarity not in (-1, 1):
raise ValueError(f'Invalid polarity: {polarity}')
if thickness <= 2:
raise ValueError('It would be wise to have a pml with 4+ cells of thickness')
if epsilon_eff <= 0:
raise ValueError('epsilon_eff must be positive')
sigma_max = -ln_R_per_layer / 2 * (m + 1)
kappa_max = numpy.sqrt(epsilon_eff * mu_eff)
alpha_max = cfs_alpha
xe = numpy.arange(1, thickness + 1, dtype=float) # TODO: pass in dtype?
xh = numpy.arange(1, thickness + 1, dtype=float)
if polarity > 0:
xe -= 0.5
elif polarity < 0:
xh -= 0.5
xe = xe[::-1]
xh = xh[::-1]
expand_slice_l: list[Any] = [None, None, None]
expand_slice_l[axis] = slice(None)
expand_slice = tuple(expand_slice_l)
def par(x: NDArray[numpy.float64]) -> tuple[NDArray[numpy.float64], NDArray[numpy.float64], NDArray[numpy.float64]]:
scaling = (x / thickness) ** m
sigma = scaling * sigma_max
kappa = 1 + scaling * (kappa_max - 1)
alpha = ((1 - x / thickness) ** ma) * alpha_max
p0 = numpy.exp(-(sigma / kappa + alpha) * dt)
p1 = sigma / (sigma + kappa * alpha) * (p0 - 1)
p2 = 1 / kappa
return p0[expand_slice], p1[expand_slice], p2[expand_slice]
p0e, p1e, p2e = par(xe)
p0h, p1h, p2h = par(xh)
region_list = [slice(None), slice(None), slice(None)]
if polarity < 0:
region_list[axis] = slice(None, thickness)
elif polarity > 0:
region_list[axis] = slice(-thickness, None)
region = tuple(region_list)
return {
'param_e': (p0e, p1e, p2e),
'param_h': (p0h, p1h, p2h),
'region': region,
}
def updates_with_cpml(
cpml_params: Sequence[Sequence[dict[str, Any] | None]],
dt: float,
dxes: dx_lists_t,
epsilon: fdfield,
*,
dtype: DTypeLike = numpy.float32,
) -> tuple[Callable[..., None], Callable[..., None]]:
"""
Build Yee-step update closures augmented with CPML terms.
Args:
cpml_params: Three-by-two sequence indexed as `[axis][polarity_index]`.
Entries are the dictionaries returned by `cpml_params(...)`; use
`None` to disable CPML on one face.
dt: Timestep.
dxes: Yee-grid spacing lists `[dx_e, dx_h]`.
epsilon: Electric material distribution used by the E update.
dtype: Storage dtype for the auxiliary CPML state arrays.
Returns:
`(update_E, update_H)` closures with the same call shape as the basic
Yee updates:
- `update_E(e, h, epsilon)`
- `update_H(e, h, mu)`
The closures retain the CPML auxiliary state internally.
"""
Dfx, Dfy, Dfz = deriv_forward(dxes[1])
Dbx, Dby, Dbz = deriv_back(dxes[1])
psi_E: list[list[tuple[Any, Any]]] = [[(None, None) for _ in range(2)] for _ in range(3)]
psi_H: list[list[tuple[Any, Any]]] = deepcopy(psi_E)
params_E: list[list[tuple[Any, Any, Any, Any]]] = [[(None, None, None, None) for _ in range(2)] for _ in range(3)]
params_H: list[list[tuple[Any, Any, Any, Any]]] = deepcopy(params_E)
for axis in range(3):
for pp, _polarity in enumerate((-1, 1)):
cpml_param = cpml_params[axis][pp]
if cpml_param is None:
psi_E[axis][pp] = (None, None)
psi_H[axis][pp] = (None, None)
continue
region = cpml_param['region']
region_shape = epsilon[0][region].shape
psi_E[axis][pp] = (
numpy.zeros(region_shape, dtype=dtype),
numpy.zeros(region_shape, dtype=dtype),
)
psi_H[axis][pp] = (
numpy.zeros(region_shape, dtype=dtype),
numpy.zeros(region_shape, dtype=dtype),
)
params_E[axis][pp] = cpml_param['param_e'] + (region,)
params_H[axis][pp] = cpml_param['param_h'] + (region,)
pE = numpy.empty_like(epsilon, dtype=dtype)
pH = numpy.empty_like(epsilon, dtype=dtype)
def update_E(
e: fdfield,
h: fdfield,
epsilon: fdfield,
) -> None:
dyHx = Dby(h[0])
dzHx = Dbz(h[0])
dxHy = Dbx(h[1])
dzHy = Dbz(h[1])
dxHz = Dbx(h[2])
dyHz = Dby(h[2])
dH = ((dxHy, dxHz),
(dyHx, dyHz),
(dzHx, dzHy))
pE.fill(0)
for axis in range(3):
se = (-1, 1, -1)[axis]
transverse = numpy.delete(range(3), axis)
u, v = transverse
dHu, dHv = dH[axis]
for pp in range(2):
psi_Eu, psi_Ev = psi_E[axis][pp]
if psi_Eu is None:
# No pml in this direction
continue
p0e, p1e, p2e, region = params_E[axis][pp]
dHu[region] *= p2e
dHv[region] *= p2e
psi_Eu *= p0e
psi_Ev *= p0e
psi_Eu += p1e * dHv[region] # note reversed u,v mapping
psi_Ev += p1e * dHu[region]
pE[u][region] += +se * psi_Eu
pE[v][region] += -se * psi_Ev
e[0] += dt / epsilon[0] * (dyHz - dzHy + pE[0])
e[1] += dt / epsilon[1] * (dzHx - dxHz + pE[1])
e[2] += dt / epsilon[2] * (dxHy - dyHx + pE[2])
def update_H(
e: fdfield,
h: fdfield,
mu: fdfield | tuple[int, int, int] = (1, 1, 1),
) -> None:
dyEx = Dfy(e[0])
dzEx = Dfz(e[0])
dxEy = Dfx(e[1])
dzEy = Dfz(e[1])
dxEz = Dfx(e[2])
dyEz = Dfy(e[2])
dE = ((dxEy, dxEz),
(dyEx, dyEz),
(dzEx, dzEy))
pH.fill(0)
for axis in range(3):
se = (-1, 1, -1)[axis]
transverse = numpy.delete(range(3), axis)
u, v = transverse
dEu, dEv = dE[axis]
for pp in range(2):
psi_Hu, psi_Hv = psi_H[axis][pp]
if psi_Hu is None:
# No pml here
continue
p0h, p1h, p2h, region = params_H[axis][pp]
dEu[region] *= p2h # modifies d_E_
dEv[region] *= p2h
psi_Hu *= p0h
psi_Hv *= p0h
psi_Hu += p1h * dEv[region] # note reversed u,v mapping
psi_Hv += p1h * dEu[region]
pH[u][region] += +se * psi_Hu
pH[v][region] += -se * psi_Hv
h[0] -= dt / mu[0] * (dyEz - dzEy + pH[0])
h[1] -= dt / mu[1] * (dzEx - dxEz + pH[1])
h[2] -= dt / mu[2] * (dxEy - dyEx + pH[2])
return update_E, update_H