"""Electromagnetic field metrics and normalization utilities.
This module provides functions for computing various electromagnetic field metrics
like energy density and Poynting flux, as well as field normalization operations.
All functions support JAX's automatic differentiation and work with the standard
FDTD field array shapes.
"""
import jax
import jax.numpy as jnp
from fdtdx.core.misc import expand_to_3x3
[docs]
def compute_energy(
E: jax.Array,
H: jax.Array,
inv_permittivity: jax.Array | float,
inv_permeability: jax.Array | float,
axis: int = 0,
) -> jax.Array:
"""Computes the total electromagnetic energy density of the field.
Args:
E (jax.Array): Electric field array with shape (3, nx, ny, nz)
H (jax.Array): Magnetic field array with shape (3, nx, ny, nz)
inv_permittivity (jax.Array | float): Inverse permittivity. Shape (3, nx, ny, nz) for anisotropic or scalar
inv_permeability (jax.Array | float): Inverse permeability. Shape (3, nx, ny, nz) for anisotropic or scalar
axis (int, optional): Axis index of the X,Y,Z component for the E and H field. Defaults to 0.
Returns:
jax.Array: Total energy density array with shape (nx, ny, nz)
"""
inv_eps_shape = getattr(inv_permittivity, "shape", ())
inv_mu_shape = getattr(inv_permeability, "shape", ())
if (inv_eps_shape and inv_eps_shape[0] == 9) or (inv_mu_shape and inv_mu_shape[0] == 9):
inv_permittivity = expand_to_3x3(inv_permittivity)
inv_permeability = expand_to_3x3(inv_permeability)
# Invert the 3x3 matrices to get eps and mu
perm = (2, 3, 4, 0, 1) # (3, 3, nx, ny, nz) -> (nx, ny, nz, 3, 3)
inv_perm = (3, 4, 0, 1, 2) # (nx, ny, nz, 3, 3) -> (3, 3, nx, ny, nz)
eps = jnp.linalg.inv(inv_permittivity.transpose(perm)).transpose(inv_perm)
mu = jnp.linalg.inv(inv_permeability.transpose(perm)).transpose(inv_perm)
# For fully anisotropic materials: 0.5 * sum_ij E_i* ε_ij E_j
# E/H has shape (3, nx, ny, nz), eps/mu has shape (3, 3, nx, ny, nz)
energy_E = 0.5 * jnp.real(jnp.einsum("ixyz,ijxyz,jxyz->xyz", jnp.conj(E), eps, E))
energy_H = 0.5 * jnp.real(jnp.einsum("ixyz,ijxyz,jxyz->xyz", jnp.conj(H), mu, H))
total_energy = energy_E + energy_H
return total_energy
else:
# For anisotropic materials: energy = 0.5 * sum_i(ε_i * |E_i|² + μ_i * |H_i|²)
# Component-wise calculation then sum
E_squared = jnp.square(jnp.abs(E)) # shape: (3, nx, ny, nz)
energy_E = 0.5 * (1 / inv_permittivity) * E_squared # component-wise multiplication
energy_E = jnp.sum(energy_E, axis=axis) # sum over components
H_squared = jnp.square(jnp.abs(H)) # shape: (3, nx, ny, nz)
energy_H = 0.5 * (1 / inv_permeability) * H_squared # component-wise multiplication
energy_H = jnp.sum(energy_H, axis=axis) # sum over components
total_energy = energy_E + energy_H
return total_energy
[docs]
def normalize_by_energy(
E: jax.Array,
H: jax.Array,
inv_permittivity: jax.Array | float,
inv_permeability: jax.Array | float,
) -> tuple[jax.Array, jax.Array]:
"""Normalizes electromagnetic fields by their total energy.
Args:
E (jax.Array): Electric field array with shape (3, nx, ny, nz)
H (jax.Array): Magnetic field array with shape (3, nx, ny, nz)
inv_permittivity (jax.Array | float): Inverse of the electric permittivity array
inv_permeability (jax.Array | float): Inverse of the magnetic permeability array
Returns:
tuple[jax.Array, jax.Array]: Tuple of (normalized E field, normalized H field)
"""
total_energy = compute_energy(
E=E,
H=H,
inv_permittivity=inv_permittivity,
inv_permeability=inv_permeability,
)
energy_root = jnp.sqrt(jnp.sum(total_energy))
norm_E = E / energy_root
norm_H = H / energy_root
return norm_E, norm_H
[docs]
def compute_poynting_flux(E: jax.Array, H: jax.Array, axis: int = 0) -> jax.Array:
"""Calculates the Poynting vector (energy flux) from E and H fields.
Args:
E (jax.Array): Electric field array with shape (3, nx, ny, nz)
H (jax.Array): Magnetic field array with shape (3, nx, ny, nz)
axis (int, optional): Axis for computing the poynting flux. Defaults to 0.
Returns:
jax.Array: Poynting vector array with shape (3, nx, ny, nz) representing
energy flux in each direction
"""
return jnp.cross(
E,
jnp.conj(H),
axisa=axis,
axisb=axis,
axisc=axis,
)
def net_poynting_flux_through_box(
poynting_vector: jax.Array,
active_axes: tuple[int, ...],
area_weights: tuple[jax.Array, ...],
) -> jax.Array:
"""Net outward Poynting flux through the faces of a rectangular box.
Integrates the outward-normal component of the Poynting vector over the six
(or fewer) faces of a box, accumulating **per cell** so that non-uniform
grids -- where the transverse cell areas differ from cell to cell within a
single face -- are handled exactly. For each active axis ``a`` the maximum
face (outward normal ``+a``) contributes ``+Sum(S_a * area)`` and the minimum
face (outward normal ``-a``) contributes ``-Sum(S_a * area)``.
A face pair on an axis of size one lands on the same cell and cancels to
zero, so thin or periodic axes may be omitted from ``active_axes`` without
changing the result.
Args:
poynting_vector (jax.Array): Real Poynting vector of shape
``(3, Nx, Ny, Nz)`` sampled on the box cells (already co-located onto
a common Yee point by the detector interpolation).
active_axes (tuple[int, ...]): Axes whose two faces contribute to the sum.
area_weights (tuple[jax.Array, ...]): Indexed by axis; ``area_weights[a]``
holds the per-cell transverse face areas for a face normal to ``a``,
broadcastable to ``poynting_vector[a]`` (size one along axis ``a``).
Per-cell weighting is what makes the integral exact on non-uniform
grids.
Returns:
jax.Array: Scalar net outward flux (positive means net power leaving the
box).
"""
net = jnp.zeros((), dtype=poynting_vector.dtype)
for a in active_axes:
weighted = poynting_vector[a] * area_weights[a]
net = net + jnp.take(weighted, -1, axis=a).sum() - jnp.take(weighted, 0, axis=a).sum()
return net
[docs]
def compute_integrated_power(
E: jax.Array,
H: jax.Array,
axis: int,
area_weights: jax.Array | None = None,
) -> jax.Array:
"""Computes the integrated power (Poynting flux) across a transverse plane.
Args:
E (jax.Array): Electric field array with component axis first.
H (jax.Array): Magnetic field array with component axis first.
axis (int): Physical propagation axis whose Poynting component is integrated.
area_weights (jax.Array | None, optional): Optional detector-plane area weights
broadcastable to ``E[axis]``. Defaults to None.
Returns:
jax.Array: The absolute integrated power.
"""
# Compute Poynting vector components
S_complex = jnp.cross(jnp.conj(E), H, axisa=0, axisb=0, axisc=0)
S_real = 0.5 * jnp.real(S_complex[axis]) # power flow in desired direction
if area_weights is not None:
# normalize area weights for numerical stability and consistency with the None case
relative_weights = area_weights / jnp.mean(area_weights)
S_real = S_real * relative_weights
# Integrate over transverse plane (axis orthogonal to `axis`)
power = jnp.abs(jnp.sum(S_real))
return power
[docs]
def normalize_by_poynting_flux(
E: jax.Array,
H: jax.Array,
axis: int,
area_weights: jax.Array | None = None,
) -> tuple[jax.Array, jax.Array]:
"""Normalize fields so the integrated Poynting flux along ``axis`` is one.
Args:
E (jax.Array): Electric field array with component axis first.
H (jax.Array): Magnetic field array with component axis first.
axis (int): Physical propagation axis whose Poynting component is integrated.
area_weights (jax.Array | None, optional): Optional detector-plane area weights
broadcastable to ``E[axis]``. Uniform-grid callers may omit this for the
historical raw-sum normalization; non-uniform callers should provide weights
so refinement alone does not change the normalization.
Returns:
tuple[jax.Array, jax.Array]: Tuple of (normalized E field, normalized H field)
"""
power = compute_integrated_power(E, H, axis, area_weights)
# Normalize
norm_factor = jnp.sqrt(power)
E_norm = E / norm_factor
H_norm = H / norm_factor
return E_norm, H_norm
def resample_to_uniform_2d(
field: jax.Array,
x_centers: jax.Array,
y_centers: jax.Array,
) -> tuple[jax.Array, jax.Array | float, jax.Array | float]:
"""Resample a ``(component, x, y)`` field onto a uniform transverse grid.
Interpolates each component from the given non-uniform physical center
coordinates onto a uniform grid with the same extent and number of points.
Returns the resampled field and the uniform grid spacings.
"""
nx, ny = field.shape[1], field.shape[2]
target_x = jnp.linspace(x_centers[0], x_centers[-1], nx)
target_y = jnp.linspace(y_centers[0], y_centers[-1], ny)
def interp_x(component):
return jax.vmap(lambda column: jnp.interp(target_x, x_centers, column), in_axes=1, out_axes=1)(component)
def interp_y(component):
return jax.vmap(lambda row: jnp.interp(target_y, y_centers, row))(component)
field = jax.vmap(interp_x)(field)
field = jax.vmap(interp_y)(field)
dx = (target_x[1] - target_x[0]) if nx > 1 else 1.0
dy = (target_y[1] - target_y[0]) if ny > 1 else 1.0
return field, dx, dy
