Skip to content

Differentiable Solver API

bcm

Differentiable planar solver: an eigendecomposition-free cascade in JAX.

The NumPy engine (:mod:hyperbolic_optics.structure) diagonalises each layer's 4×4 Berreman matrix, whose eigenvectors are awkward to differentiate (arbitrary phase, divergent derivative at degeneracies, which every isotropic layer has). Here each layer's scattering matrix is built from matrix products and linear solves only, so jax.grad gives exact derivatives with respect to thicknesses, tensors, orientations, angle and wavelength.

For a thin slice, exp(iΔτ) − I is a short Taylor series; 2ᴺ slices are assembled by N doublings, each a Redheffer star product of the interval with itself (the precise integration method of Zhong & Williams 1994, as in Yang, Song & Cai, arXiv:2608.06185). Intervals are carried as wave amplitudes in a reference medium of permittivity 2kx² + 1, whose waves always propagate and are scaled to unit power, so no doubling meets a singular inverse.

Coefficients are Jones matrices [..., out, in] with 0 = p and 1 = s, on the NumPy engine's sign convention. The half-spaces are isotropic and lossless, except in :func:solve_stack_on_crystal. Inputs broadcast over leading batch dimensions. Reference accuracy needs jax.config.update("jax_enable_x64", True); without it a warning says so.

Classes:

  • PlanarResult –

    Complex Jones matrices and powers, each [..., out, in] with 0 = p, 1 = s.

  • PlanarSMatrix –

    The full bidirectional scattering matrix of a planar stack.

  • StackFields –

    Fields and power flow through a planar stack, for p and s light from the top.

  • CrystalResult –

    Response of a stack on a semi-infinite anisotropic (or lossy) substrate.

Functions:

  • rotation_matrix –

    Euler rotation Rz(z) · Ry(y) · Rx(x) in radians, [..., 3, 3].

  • rotate_tensor –

    Similarity transform R · tensor · Rᵀ, broadcasting over batch dimensions.

  • system_matrix –

    Berreman matrix Δ with dF/dz̃ = iΔF, F = [Ex, Ey, Hx, Hy].

  • solve_stack –

    Reflection and transmission of a planar stack of anisotropic layers.

  • stack_smatrix –

    The bidirectional, power-normalized scattering matrix of a planar stack.

  • cascade –

    Connect upper's bottom port to lower's top port.

  • stack_fields –

    Fields, interface fluxes and layer-resolved absorption of a planar stack.

  • solve_stack_on_crystal –

    A planar stack on a semi-infinite anisotropic, lossy or hyperbolic substrate.

PlanarResult

Bases: NamedTuple

Complex Jones matrices and powers, each [..., out, in] with 0 = p, 1 = s.

PlanarSMatrix

Bases: NamedTuple

The full bidirectional scattering matrix of a planar stack.

Attributes:

  • s (Array) –

    [..., 4, 4] indexed [out, in] over :data:PORTS. top is the incident-side half-space, bottom the exit side; p and s are defined in each half-space's own plane of incidence (x-z). Reference planes are the stack's two outer interfaces.

  • propagating (Array) –

    [..., 4], whether each port's wave propagates (it is evanescent in a half-space whose index is below kx).

Amplitudes are scaled by √kz of their half-space, so between propagating ports |s|² is the fraction of power transferred, and a lossless stack has a unitary s. Reciprocal media (symmetric ε and μ) give s(kx) = D · s(−kx)ᵀ · D with D = diag(−1, 1, −1, 1): reversing a wave reverses its in-plane wavevector, and the p waves' sign convention (they share the sign of Hy) flips on reversal. For an evanescent port the same scaling still applies (with complex kz), so matrices of stacks that share a half-space compose correctly through it -- frustrated total internal reflection included -- but |s|² is not a power there.

StackFields

Bases: NamedTuple

Fields and power flow through a planar stack, for p and s light from the top.

The last axis of every array is the incident polarization (0 = p, 1 = s), each of unit incident power.

Attributes:

  • R, (T) –

    [..., 2] total reflectance and transmittance (summed over the outgoing polarizations).

  • absorption (Array) –

    [..., L, 2] fraction of the incident power absorbed in each layer.

  • interface_flux (Array) –

    [..., L + 1, 2] normal energy flux S_z at each interface, from the top of the first layer down, over the incident flux: 1 − R at the top and T at the bottom.

  • z_um (Array) –

    [..., L·S] depths of the field samples below the top of the first layer (empty unless samples_per_layer was given).

  • fields (Array) –

    [..., L·S, 6, 2] complex (Ex, Ey, Ez, Hx, Hy, Hz) at those depths, H scaled by the free-space impedance.

CrystalResult

Bases: NamedTuple

Response of a stack on a semi-infinite anisotropic (or lossy) substrate.

Attributes:

  • r (Array) –

    [..., 2, 2] reflection Jones matrix [out, in] (0 = p, 1 = s), on the same convention as :func:solve_stack.

  • R (Array) –

    |r|².

  • T (Array) –

    [..., 2] total power carried into the substrate, per incident polarization, from the normal Poynting flux just inside it (1 − R − T is absorbed in the layers).

rotation_matrix

rotation_matrix(x, y, z)

Euler rotation Rz(z) · Ry(y) · Rx(x) in radians, [..., 3, 3].

The payload API's rotationX/rotationY/rotationZ convention, so a tensor rotated with :func:rotate_tensor matches the one a Crystal Layer would use for the same angles.

Source code in hyperbolic_optics/bcm.py
def rotation_matrix(x: jax.Array, y: jax.Array, z: jax.Array) -> jax.Array:
    """Euler rotation ``Rz(z) · Ry(y) · Rx(x)`` in radians, ``[..., 3, 3]``.

    The payload API's ``rotationX``/``rotationY``/``rotationZ`` convention, so a
    tensor rotated with :func:`rotate_tensor` matches the one a ``Crystal Layer``
    would use for the same angles.
    """
    x, y, z = jnp.broadcast_arrays(*(jnp.asarray(angle, dtype=float) for angle in (x, y, z)))
    zero, one = jnp.zeros_like(x), jnp.ones_like(x)

    def matrix(rows):
        return jnp.stack([jnp.stack(row, axis=-1) for row in rows], axis=-2)

    cx, sx, cy, sy, cz, sz = jnp.cos(x), jnp.sin(x), jnp.cos(y), jnp.sin(y), jnp.cos(z), jnp.sin(z)
    rot_x = matrix([[one, zero, zero], [zero, cx, -sx], [zero, sx, cx]])
    rot_y = matrix([[cy, zero, sy], [zero, one, zero], [-sy, zero, cy]])
    rot_z = matrix([[cz, -sz, zero], [sz, cz, zero], [zero, zero, one]])
    return rot_z @ rot_y @ rot_x

rotate_tensor

rotate_tensor(tensor, rotation)

Similarity transform R · tensor · Rᵀ, broadcasting over batch dimensions.

Source code in hyperbolic_optics/bcm.py
def rotate_tensor(tensor: jax.Array, rotation: jax.Array) -> jax.Array:
    """Similarity transform ``R · tensor · Rᵀ``, broadcasting over batch dimensions."""
    return rotation @ tensor @ jnp.swapaxes(rotation, -1, -2)

system_matrix

system_matrix(eps, mu, kx)

Berreman matrix Δ with dF/dz̃ = iΔF, F = [Ex, Ey, Hx, Hy].

Derived from Maxwell's equations with ∂x → i·kx and ∂y → 0 by eliminating the longitudinal components: Ez = a·F from (εE)_z = −kx·Hy and Hz = b·F from (μH)_z = kx·Ey.

Parameters:

  • eps (Array) –

    Permittivity tensor [..., 3, 3].

  • mu (Array) –

    Permeability tensor [..., 3, 3].

  • kx (Array) –

    Normalised in-plane wavevector [...].

Returns:

  • Array –

    [..., 4, 4].

Source code in hyperbolic_optics/bcm.py
def system_matrix(eps: jax.Array, mu: jax.Array, kx: jax.Array) -> jax.Array:
    """Berreman matrix ``Δ`` with ``dF/dz̃ = iΔF``, ``F = [Ex, Ey, Hx, Hy]``.

    Derived from Maxwell's equations with ``∂x → i·kx`` and ``∂y → 0`` by
    eliminating the longitudinal components: ``Ez = a·F`` from
    ``(εE)_z = −kx·Hy`` and ``Hz = b·F`` from ``(μH)_z = kx·Ey``.

    Args:
        eps: Permittivity tensor ``[..., 3, 3]``.
        mu: Permeability tensor ``[..., 3, 3]``.
        kx: Normalised in-plane wavevector ``[...]``.

    Returns:
        ``[..., 4, 4]``.
    """
    batch = jnp.broadcast_shapes(eps.shape[:-2], mu.shape[:-2], jnp.shape(kx))
    eps = jnp.broadcast_to(eps, batch + (3, 3))
    mu = jnp.broadcast_to(mu, batch + (3, 3))
    a, b = _longitudinal(eps, mu, kx)
    kx = jnp.broadcast_to(kx, batch).astype(eps.dtype)[..., None]

    def e(i, j):
        return eps[..., i, j][..., None]

    def m(i, j):
        return mu[..., i, j][..., None]

    unit = jnp.eye(4, dtype=eps.dtype)
    ex, ey, hx, hy = unit[0], unit[1], unit[2], unit[3]

    row_ex = kx * a + m(1, 0) * hx + m(1, 1) * hy + m(1, 2) * b
    row_ey = -(m(0, 0) * hx + m(0, 1) * hy + m(0, 2) * b)
    row_hx = kx * b - (e(1, 0) * ex + e(1, 1) * ey + e(1, 2) * a)
    row_hy = e(0, 0) * ex + e(0, 1) * ey + e(0, 2) * a
    return jnp.stack([row_ex, row_ey, row_hx, row_hy], axis=-2)

solve_stack

solve_stack(
    eps,
    thickness_um,
    wavelength_um,
    incidence_rad,
    eps_incident,
    eps_exit,
    mu=None,
    doublings=DEFAULT_DOUBLINGS,
)

Reflection and transmission of a planar stack of anisotropic layers.

Parameters:

  • eps (Array) –

    Layer permittivity tensors in the lab frame, [..., L, 3, 3], ordered from the incident side. L may be zero.

  • thickness_um (Array) –

    Layer thicknesses in µm, [..., L].

  • wavelength_um (Array) –

    Vacuum wavelength in µm, [...]. (From a frequency in cm⁻¹: 1e4 / frequency.)

  • incidence_rad (Array) –

    Angle of incidence in the incident medium, radians.

  • eps_incident (Array) –

    Real, positive permittivity of the incident half-space.

  • eps_exit (Array) –

    Real, positive permittivity of the exit half-space.

  • mu (Array | None, default: None ) –

    Layer permeability tensors [..., L, 3, 3]; identity if omitted.

  • doublings (int, default: DEFAULT_DOUBLINGS ) –

    Doublings per layer (static). See :data:DEFAULT_DOUBLINGS.

Returns:

  • PlanarResult –

    class:PlanarResult with r, t, R, T, each

  • PlanarResult –

    [..., 2, 2] indexed [out, in], 0 = p and 1 = s. Wrap in

  • PlanarResult –

    jax.jit(..., static_argnames="doublings") for repeated calls.

Source code in hyperbolic_optics/bcm.py
def solve_stack(
    eps: jax.Array,
    thickness_um: jax.Array,
    wavelength_um: jax.Array,
    incidence_rad: jax.Array,
    eps_incident: jax.Array,
    eps_exit: jax.Array,
    mu: jax.Array | None = None,
    doublings: int = DEFAULT_DOUBLINGS,
) -> PlanarResult:
    """Reflection and transmission of a planar stack of anisotropic layers.

    Args:
        eps: Layer permittivity tensors in the lab frame, ``[..., L, 3, 3]``,
            ordered from the incident side. ``L`` may be zero.
        thickness_um: Layer thicknesses in µm, ``[..., L]``.
        wavelength_um: Vacuum wavelength in µm, ``[...]``. (From a frequency in
            cm⁻¹: ``1e4 / frequency``.)
        incidence_rad: Angle of incidence in the incident medium, radians.
        eps_incident: Real, positive permittivity of the incident half-space.
        eps_exit: Real, positive permittivity of the exit half-space.
        mu: Layer permeability tensors ``[..., L, 3, 3]``; identity if omitted.
        doublings: Doublings per layer (static). See :data:`DEFAULT_DOUBLINGS`.

    Returns:
        :class:`PlanarResult` with ``r``, ``t``, ``R``, ``T``, each
        ``[..., 2, 2]`` indexed ``[out, in]``, 0 = p and 1 = s. Wrap in
        ``jax.jit(..., static_argnames="doublings")`` for repeated calls.
    """
    _check_x64()
    kx = jnp.sqrt(jnp.asarray(eps_incident)) * jnp.sin(jnp.asarray(incidence_rad))
    t, _, r, _ = _scattering(
        eps, thickness_um, wavelength_um, kx, eps_incident, eps_exit, mu, doublings
    )
    flux = jnp.real(_kz(eps_exit, kx)) / jnp.real(_kz(eps_incident, kx))
    return PlanarResult(r=r, t=t, R=jnp.abs(r) ** 2, T=jnp.abs(t) ** 2 * flux[..., None, None])

stack_smatrix

stack_smatrix(
    eps,
    thickness_um,
    wavelength_um,
    kx,
    eps_top,
    eps_bottom,
    mu=None,
    doublings=DEFAULT_DOUBLINGS,
)

The bidirectional, power-normalized scattering matrix of a planar stack.

Where :func:solve_stack gives the response to light from the incident side, this gives all four ports: light from either side, in either polarization, as needed to connect stacks to one another (for instance in a SAX circuit, :mod:hyperbolic_optics.integrations.sax).

Parameters:

  • eps (Array) –

    Layer permittivity tensors [..., L, 3, 3], from the top down.

  • thickness_um (Array) –

    Layer thicknesses in µm, [..., L].

  • wavelength_um (Array) –

    Vacuum wavelength in µm.

  • kx (Array) –

    In-plane wavevector over k₀ (n·sin θ, conserved through the stack), rather than an angle, so stacks sharing a half-space share it too.

  • eps_top, eps_bottom –

    Real, positive permittivities of the half-spaces.

  • mu (Array | None, default: None ) –

    Layer permeability tensors; identity if omitted.

  • doublings (int, default: DEFAULT_DOUBLINGS ) –

    Doublings per layer (static).

Returns:

Source code in hyperbolic_optics/bcm.py
def stack_smatrix(
    eps: jax.Array,
    thickness_um: jax.Array,
    wavelength_um: jax.Array,
    kx: jax.Array,
    eps_top: jax.Array,
    eps_bottom: jax.Array,
    mu: jax.Array | None = None,
    doublings: int = DEFAULT_DOUBLINGS,
) -> PlanarSMatrix:
    """The bidirectional, power-normalized scattering matrix of a planar stack.

    Where :func:`solve_stack` gives the response to light from the incident
    side, this gives all four ports: light from either side, in either
    polarization, as needed to connect stacks to one another (for instance in a
    SAX circuit, :mod:`hyperbolic_optics.integrations.sax`).

    Args:
        eps: Layer permittivity tensors ``[..., L, 3, 3]``, from the top down.
        thickness_um: Layer thicknesses in µm, ``[..., L]``.
        wavelength_um: Vacuum wavelength in µm.
        kx: In-plane wavevector over k₀ (``n·sin θ``, conserved through the
            stack), rather than an angle, so stacks sharing a half-space share
            it too.
        eps_top, eps_bottom: Real, positive permittivities of the half-spaces.
        mu: Layer permeability tensors; identity if omitted.
        doublings: Doublings per layer (static).

    Returns:
        :class:`PlanarSMatrix` over :data:`PORTS`.
    """
    _check_x64()
    kx = jnp.asarray(kx)
    s00, s01, s10, s11 = _scattering(
        eps, thickness_um, wavelength_um, kx, eps_top, eps_bottom, mu, doublings
    )
    return _normalize_smatrix((s00, s01, s10, s11), kx, eps_top, eps_bottom)

cascade

cascade(upper, lower)

Connect upper's bottom port to lower's top port.

The two must share that half-space (same permittivity, same kx): the result is the stack of both, with the shared medium of zero thickness -- what a SAX circuit computes, without SAX. Evanescent channels in the shared medium are carried, so a gap below total internal reflection tunnels.

Source code in hyperbolic_optics/bcm.py
def cascade(upper: PlanarSMatrix, lower: PlanarSMatrix) -> PlanarSMatrix:
    """Connect ``upper``'s bottom port to ``lower``'s top port.

    The two must share that half-space (same permittivity, same ``kx``): the
    result is the stack of both, with the shared medium of zero thickness --
    what a SAX circuit computes, without SAX. Evanescent channels in the shared
    medium are carried, so a gap below total internal reflection tunnels.
    """

    def blocks(m):
        s = m.s
        return s[..., :2, :2], s[..., 2:, :2], s[..., :2, 2:], s[..., 2:, 2:]

    a_tt, a_bt, a_tb, a_bb = blocks(upper)  # [out side][in side]: t = top, b = bottom
    b_tt, b_bt, b_tb, b_bb = blocks(lower)
    eye = jnp.eye(2, dtype=a_tt.dtype)
    down = _solve(eye - b_tt @ a_bb, b_tt @ a_bt)  # lower's reflection seen from above
    up = _solve(eye - a_bb @ b_tt, a_bb)
    top_top = a_tt + a_tb @ down
    bottom_top = b_bt @ _solve(eye - a_bb @ b_tt, a_bt)
    top_bottom = a_tb @ _solve(eye - b_tt @ a_bb, b_tb)
    bottom_bottom = b_bb + b_bt @ up @ b_tb
    s = jnp.concatenate(
        [
            jnp.concatenate([top_top, top_bottom], axis=-1),
            jnp.concatenate([bottom_top, bottom_bottom], axis=-1),
        ],
        axis=-2,
    )
    propagating = jnp.concatenate([upper.propagating[..., :2], lower.propagating[..., 2:]], axis=-1)
    return PlanarSMatrix(s=s, propagating=propagating)

stack_fields

stack_fields(
    eps,
    thickness_um,
    wavelength_um,
    incidence_rad,
    eps_incident,
    eps_exit,
    mu=None,
    samples_per_layer=0,
    doublings=DEFAULT_DOUBLINGS,
)

Fields, interface fluxes and layer-resolved absorption of a planar stack.

The fields at each plane come from the scattering matrices of the stack above and below it, never from transfer matrices, so thick, evanescent and strongly absorbing layers stay stable. Everything is differentiable, including with respect to thicknesses: the samples inside a layer sit at fixed fractions of its thickness.

Parameters:

  • samples_per_layer (int, default: 0 ) –

    Field samples per layer (static), evenly spaced from its top to its bottom, both included. 0 skips the depth profile.

The other arguments are as in :func:solve_stack.

Returns:

Source code in hyperbolic_optics/bcm.py
def stack_fields(
    eps: jax.Array,
    thickness_um: jax.Array,
    wavelength_um: jax.Array,
    incidence_rad: jax.Array,
    eps_incident: jax.Array,
    eps_exit: jax.Array,
    mu: jax.Array | None = None,
    samples_per_layer: int = 0,
    doublings: int = DEFAULT_DOUBLINGS,
) -> StackFields:
    """Fields, interface fluxes and layer-resolved absorption of a planar stack.

    The fields at each plane come from the scattering matrices of the stack
    above and below it, never from transfer matrices, so thick, evanescent and
    strongly absorbing layers stay stable. Everything is differentiable,
    including with respect to thicknesses: the samples inside a layer sit at
    fixed fractions of its thickness.

    Args:
        samples_per_layer: Field samples per layer (static), evenly spaced from
            its top to its bottom, both included. ``0`` skips the depth profile.

    The other arguments are as in :func:`solve_stack`.

    Returns:
        :class:`StackFields`.
    """
    _check_x64()
    eps, mu = _materials(eps, mu)
    thickness_um = jnp.asarray(thickness_um)
    eps_incident, eps_exit = jnp.asarray(eps_incident), jnp.asarray(eps_exit)
    n_layers = eps.shape[-3]
    kx = jnp.sqrt(eps_incident) * jnp.sin(jnp.asarray(incidence_rad))
    k0 = 2 * jnp.pi / jnp.asarray(wavelength_um)
    w_ref = _reference_modes(kx)

    deltas = _deltas(eps, mu, kx, w_ref)
    layers = [_layer(deltas[i], k0 * thickness_um[..., i], doublings) for i in range(n_layers)]
    above = [_interface(_isotropic_modes(eps_incident, kx), w_ref)]
    for layer in layers:
        above.append(_star(above[-1], layer))
    below = [_interface(w_ref, _isotropic_modes(eps_exit, kx))]
    for layer in reversed(layers):
        below.insert(0, _star(layer, below[0]))

    incident_flux = 0.5 * jnp.real(_kz(eps_incident, kx))[..., None]
    planes = [_plane_fields(above[j], below[j], w_ref) for j in range(n_layers + 1)]
    flux = jnp.stack([_poynting_z(f) for f in planes], axis=-2) / incident_flux[..., None, :]
    absorption = flux[..., :-1, :] - flux[..., 1:, :]

    batch = flux.shape[:-2]
    z_um = jnp.zeros(batch + (0,))
    fields = jnp.zeros(batch + (0, 6, 2), dtype=eps.dtype)
    if samples_per_layer and n_layers:
        fractions = jnp.linspace(0.0, 1.0, samples_per_layer)
        tops = jnp.cumsum(thickness_um, axis=-1) - thickness_um
        depths, samples = [], []
        for i in range(n_layers):
            k0d = (k0 * thickness_um[..., i])[..., None]
            delta = deltas[i][..., None, :, :]
            upper = _layer(delta, k0d * fractions, doublings)
            lower = _layer(delta, k0d * (1 - fractions), doublings)

            def expand(blocks):
                return tuple(block[..., None, :, :] for block in blocks)

            tangential = _plane_fields(
                _star(expand(above[i]), upper), _star(lower, expand(below[i + 1])),
                w_ref[..., None, :, :],
            )  # fmt: skip
            a, b = _longitudinal(eps[..., i, :, :], mu[..., i, :, :], kx)
            ez = jnp.einsum("...j,...sjn->...sn", a, tangential)
            hz = jnp.einsum("...j,...sjn->...sn", b, tangential)
            full = jnp.stack(
                [tangential[..., 0, :], tangential[..., 1, :], ez,
                 tangential[..., 2, :], tangential[..., 3, :], hz],
                axis=-2,
            )  # fmt: skip
            samples.append(full)
            depths.append(tops[..., i, None] + thickness_um[..., i, None] * fractions)
        fields = jnp.concatenate(jnp.broadcast_arrays(*samples), axis=-3)
        z_um = jnp.concatenate(jnp.broadcast_arrays(*depths), axis=-1)

    return StackFields(
        R=1.0 - flux[..., 0, :],
        T=flux[..., -1, :],
        absorption=absorption,
        interface_flux=flux,
        z_um=z_um,
        fields=fields,
    )

solve_stack_on_crystal

solve_stack_on_crystal(
    eps,
    thickness_um,
    wavelength_um,
    incidence_rad,
    eps_incident,
    eps_substrate,
    mu=None,
    mu_substrate=None,
    doublings=DEFAULT_DOUBLINGS,
)

A planar stack on a semi-infinite anisotropic, lossy or hyperbolic substrate.

The NumPy engine's Semi Infinite Anisotropic Layer, in the differentiable solver: bulk crystals (calcite, sapphire, hBN, SiC, α-MoO₃, a LiNbO₃ wafer) as substrates, in any orientation. The waves the substrate carries away are found without eigendecomposition, by a matrix sign function (see :func:_substrate_modes), so reflection is exact and differentiable even where the substrate's modes are degenerate.

The other arguments are as in :func:solve_stack (L may be zero: a bare crystal).

Parameters:

  • eps_substrate (Array) –

    The substrate's tensor [..., 3, 3] in the lab frame (rotate it with :func:rotate_tensor); lossy is fine.

  • mu_substrate (Array | None, default: None ) –

    Its permeability; identity if omitted.

Returns:

Source code in hyperbolic_optics/bcm.py
def solve_stack_on_crystal(
    eps: jax.Array,
    thickness_um: jax.Array,
    wavelength_um: jax.Array,
    incidence_rad: jax.Array,
    eps_incident: jax.Array,
    eps_substrate: jax.Array,
    mu: jax.Array | None = None,
    mu_substrate: jax.Array | None = None,
    doublings: int = DEFAULT_DOUBLINGS,
) -> CrystalResult:
    """A planar stack on a semi-infinite anisotropic, lossy or hyperbolic substrate.

    The NumPy engine's ``Semi Infinite Anisotropic Layer``, in the
    differentiable solver: bulk crystals (calcite, sapphire, hBN, SiC,
    α-MoO₃, a LiNbO₃ wafer) as substrates, in any orientation. The waves the
    substrate carries away are found without eigendecomposition, by a matrix
    sign function (see :func:`_substrate_modes`), so reflection is exact and
    differentiable even where the substrate's modes are degenerate.

    The other arguments are as in :func:`solve_stack` (``L`` may be zero: a
    bare crystal).

    Args:
        eps_substrate: The substrate's tensor ``[..., 3, 3]`` in the lab frame
            (rotate it with :func:`rotate_tensor`); lossy is fine.
        mu_substrate: Its permeability; identity if omitted.

    Returns:
        :class:`CrystalResult`: ``r``, ``R`` and the power transmitted into
        the substrate ``T``.
    """
    _check_x64()
    eps_substrate, mu_substrate = _materials(eps_substrate, mu_substrate)
    eps_incident = jnp.asarray(eps_incident)
    kx = jnp.sqrt(eps_incident) * jnp.sin(jnp.asarray(incidence_rad))
    # a lossless substrate gets a relative loss of 1e-8 to decide which waves carry energy
    # in; the subspaces are then refined to the lossless ones, so it costs no accuracy
    w_sub = _substrate_modes(eps_substrate, mu_substrate, kx, 1e-8, 60)
    w_in = _isotropic_modes(eps_incident, kx)
    t, _, r, _ = _cascade(eps, mu, thickness_um, wavelength_um, kx, w_in, w_sub, doublings)
    flux = _poynting_z(w_sub[..., :, :2] @ t)  # into the substrate, per input
    incident = 0.5 * jnp.real(_kz(eps_incident, kx))[..., None]
    return CrystalResult(r=r, R=jnp.abs(r) ** 2, T=flux / incident)