diffpes.maths

Differentiable Cartesian rotations and angular matrix elements for dipole photoemission.

Compute angular matrix elements for dipole photoemission.

Extended Summary

The subpackage provides Gaunt coefficients and real spherical harmonics. It also assembles the full dipole matrix element for the differentiable ARPES forward model. The dipole matrix element for orbital \((n, l, m)\) combines radial integrals, Gaunt coefficients, and spherical harmonics of the photoelectron direction:

\[M(\mathbf{k}, n, l, m) = \sum_{l', m'} B_{n,l}^{l'}(|\mathbf{k}|) \cdot G(l, m, l', m') \cdot Y_{l'}^{m'}(\hat{k}) \cdot \hat{e}_{q(m'-m)}\]

The following list describes the submodules:

  • dipole

    Assemble full dipole matrix elements.

  • gaunt

    Build the Gaunt coefficient table for dipole transitions.

  • rotations

    Construct differentiable three-dimensional rotations.

  • safe

    Provide named gradient-safe elementary operations.

  • spherical_harmonics

    Compute real spherical harmonics in JAX.

Routine Listings

GAUNT_TABLE

Module-level precomputed Gaunt coefficient table for l_max=4.

build_gaunt_table()

Build the dipole Gaunt coefficient lookup table.

dipole_intensities_all_orbitals()

Compute |M|^2 for all orbitals in the basis.

dipole_intensity_orbital()

Compute |M|^2 for one orbital.

dipole_matrix_element_single()

Compute dipole matrix element for a single orbital (n, l, m).

gaunt_lookup()

Look up a single Gaunt coefficient from the precomputed table.

real_spherical_harmonic()

Evaluate a single real spherical harmonic.

real_spherical_harmonics_all()

Evaluate all real spherical harmonics up to l_max.

rodrigues_rotation()

Construct a rotation matrix with Rodrigues’ formula.

safe_arccos()

Evaluate arccos with saturated values and zero boundary gradients.

safe_arctan2()

Evaluate arctan2 with a zero value and gradient at the origin.

safe_divide()

Divide with a fallback and zero quotient gradients at zero denominators.

safe_log()

Evaluate log with a finite floor and zero gradients below it.

safe_norm()

Compute a Euclidean norm with a zero gradient at zero vectors.

safe_power()

Raise positive inputs to a power and return zero otherwise.

safe_sqrt()

Evaluate sqrt on positive inputs and return zero otherwise.

Notes

All functions support JAX transformations and automatic differentiation. Pure Python computes the Gaunt table once during import. The module stores the table as a JAX array for constant-time lookup during traced computation.

diffpes.maths.dipole_intensities_all_orbitals(k_vec: Float[Array, '3'], r_grid: Float[Array, 'R'], slater_params: SlaterParams, efield: Complex[Array, '3']) → Float[Array, 'O'][source]

Compute |M|^2 for all orbitals in the basis.

The function scans every orbital in the Slater basis. It computes each radial wavefunction from slater_params and then computes the squared dipole matrix element. jax.lax.switch specializes one branch for each static (n, l, m) tuple. jax.lax.scan carries the differentiable arrays along the orbital axis.

For each orbital o, the function:

  1. Extracts quantum numbers \((n_o, l_o, m_o)\) and Slater exponent \(\zeta_o\) from slater_params.

  2. Evaluates the normalized Slater radial function \(R(r) = N r^{n-1} e^{-\zeta r}\) on the supplied grid.

  3. Weights by the multi-zeta coefficient slater_params.coefficients[o, 0] (first column for single-zeta bases).

  4. Calls dipole_intensity_orbital to compute \(|M|^2\).

The function stacks the results into a one-dimensional array. Its length equals the number of orbitals in the basis.

See:

TestDipoleIntensitiesAllOrbitals

Parameters:

  • k_vec (Float[Array, '3']) – Photoelectron wavevector.

  • r_grid (Float[Array, 'R']) – Radial grid.

  • slater_params (SlaterParams) – Slater exponents and orbital basis.

  • efield (Complex[Array, '3']) – Polarization vector.

Returns:

intensities – |M|^2 per orbital.

Return type:

Float[Array, 'O']

Notes

The number and quantum numbers of orbitals remain static PyTree metadata, so changing the basis structure retraces the function. Slater exponents and coefficients remain traced leaves with gradients through the scan.

diffpes.maths.dipole_intensity_orbital(k_vec: Float[Array, '3'], r_grid: Float[Array, 'R'], radial_values: Float[Array, 'R'], l: int, m: int, efield: Complex[Array, '3']) → Float[Array, ''][source]

Compute |M|^2 for one orbital.

The function computes the photoemission intensity for one initial-state orbital. The quantum numbers (l, m) and the sampled radial wavefunction define this orbital. The intensity is the squared modulus of the complex dipole matrix element:

\[I(\mathbf{k}) = |M(\mathbf{k}, l, m)|^2\]

This wrapper calls dipole_matrix_element_single and returns \(|M|^2 = M \cdot M^*\). The construction gives a real, nonnegative result. JAX differentiates the result with respect to all continuous inputs.

See:

TestDipoleIntensityOrbital

Parameters:

  • k_vec (Float[Array, '3']) – Photoelectron wavevector.

  • r_grid (Float[Array, 'R']) – Radial grid.

  • radial_values (Float[Array, 'R']) – Radial wavefunction on grid.

  • l (int) – Angular momentum.

  • m (int) – Magnetic quantum number.

  • efield (Complex[Array, '3']) – Polarization vector.

Returns:

intensity – Squared modulus of the matrix element.

Return type:

Float[Array, '']

Notes

The function computes the complex matrix element first. It then applies the modulus squared once, after the coherent channel sum is complete. Gradients flow through both operations for every continuous input.

diffpes.maths.dipole_matrix_element_single(k_vec: Float[Array, '3'], r_grid: Float[Array, 'R'], radial_values: Float[Array, 'R'], l: int, m: int, efield: Complex[Array, '3']) → Complex[Array, ''][source]

Compute dipole matrix element for a single orbital (n, l, m).

\[M = \sum_{q} \hat{e}_q \sum_{l'} B^{l'}(|k|) \cdot G(l, m, l', m+q) \cdot Y_{l'}^{m+q}(\hat{k})\]

where the sum is over dipole components q in {-1, 0, +1} and final-state angular momenta l’ in {l-1, l+1}.

The function assembles the full photoemission dipole matrix element from four quantities:

  1. Radial integral \(B^{l'}(|k|)\) – the overlap between the initial radial wavefunction and the final-state spherical Bessel function \(j_{l'}(kr)\). The radial integral applies the \(r^3\) weight and trapezoidal quadrature.

  2. Gaunt coefficient \(G(l, m, l', m')\) – the angular integral that couples the initial and final states through the dipole operator. The function reads this coefficient from GAUNT_TABLE.

  3. Real spherical harmonic \(Y_{l'}^{m'}(\hat{k})\) – the angular part of the final-state plane wave expansion. The function computes it at the direction of the photoelectron wavevector.

  4. Polarization component \(\hat{e}_q\) – the q-th spherical component of the polarization vector. The _cartesian_to_spherical_dipole function maps the Cartesian components to the real harmonic basis.

The dipole selection rule \(l' = l \pm 1\) restricts the final-state sum to at most two terms per q value. The magnetic selection rule \(m' = m + q\) with \(|q| \le 1\) means at most three q values contribute.

Numerical stability routes the wavevector norm, normalization, polar angle, and azimuthal angle through the named safe-math primitives. Their guard conventions select zero subgradients at the zero vector and at the angular coordinate singularities without perturbing non-singular values.

See:

TestDipoleMatrixElementSingle

Parameters:

  • k_vec (Float[Array, '3']) – Photoelectron wavevector in Cartesian coordinates.

  • r_grid (Float[Array, 'R']) – Radial grid for integration.

  • radial_values (Float[Array, 'R']) – R(r) sampled on r_grid.

  • l (int) – Angular momentum quantum number of the initial orbital.

  • m (int) – Magnetic quantum number of the initial orbital.

  • efield (Complex[Array, '3']) – Polarization vector in Cartesian coordinates.

Returns:

M – Complex dipole matrix element.

Return type:

Complex[Array, '']

Notes

Python unrolls the loop over q and l’ during tracing. Static quantum numbers determine the iteration bounds. This process produces one fixed computation graph for each (l, m) pair. Different orbitals therefore produce distinct XLA programs.

diffpes.maths.build_gaunt_table(l_max: int = 4) → Float[Array, 'L_src M_src 3 L_dst M_dst'][source]

Build the dipole Gaunt coefficient lookup table.

The function computes each nonzero real Gaunt coefficient through l_max. It stores the coefficients in a dense five-dimensional NumPy array. It then converts the table to a JAX array for constant-time lookup.

The five axes correspond to:

  • axis 0 (l): initial angular momentum, size l_max + 1.

  • axis 1 (m + l_max): offset initial magnetic quantum number, size 2 * l_max + 1, so that m = -l_max maps to index 0.

  • axis 2 (q + 1): dipole component index, size 3, for q in {-1, 0, +1}.

  • axis 3 (lp): final angular momentum, size l_max + 2 (because the dipole operator can promote l to l + 1).

  • axis 4 (mp + l_max): offset final magnetic quantum number, size 2 * (l_max + 1) + 1.

The construction visits each valid combination of quantum numbers. It applies the dipole selection rules and calls _real_gaunt_dipole. The initialized zero values represent forbidden transitions.

Use the following expression to index the table: GAUNT_TABLE[l, m + l_max, q + 1, lp, mp + l_max] where q in {-1, 0, +1} indexes the three dipole components.

See:

TestBuildGauntTable

Implementation Logic

  1. Allocate the dense coefficient table:

    table: Float[NDArray, "L1 M1 Q L2 M2"] = np.zeros(
    (l_src_dim, m_src_dim, q_dim, l_dst_dim, m_dst_dim),
    dtype=np.float64,
)

    The dense layout gives constant-time lookup for valid quantum numbers.

  2. Convert the completed table to JAX:

    gaunt_table: Float[
    Array, "L_src M_src 3 L_dst M_dst"
] = jnp.asarray(table, dtype=jnp.float64)

    The JAX constant can participate in differentiable forward models.

param l_max:

Maximum angular momentum. Default 4 (s through g).

type l_max:

int, default: 4

returns:

gaunt_table – Dense array of Gaunt coefficients. Shape: (l_max+1, 2*l_max+1, 3, l_max+2, 2*(l_max+1)+1).

rtype:

Float[Array, 'L_src M_src 3 L_dst M_dst']

Notes

The function uses pure Python and NumPy because it runs once during module import. The module stores the result as a JAX array constant. Therefore, gradient computations do not include it as a trainable parameter.

diffpes.maths.gaunt_lookup(l: int, m: int, q: int, lp: int, mp: int) → float[source]

Look up a single Gaunt coefficient from the precomputed table.

The function provides an accessor for the module-level GAUNT_TABLE. It converts the physical quantum numbers into offsets for the dense storage layout. The index mapping is:

  • l indexes axis 0 directly.

  • m + L_MAX offsets the magnetic quantum number to a non-negative index on axis 1.

  • q + 1 maps q in {-1, 0, +1} to indices {0, 1, 2} on axis 2.

  • lp indexes axis 3 directly.

  • mp + L_MAX offsets m’ on axis 4.

The function converts the returned value to a Python float for non-JAX contexts. In JAX-traced code, index GAUNT_TABLE directly to avoid Python overhead.

See:

TestGauntLookup

Implementation Logic

  1. Index and convert the coefficient:

    coefficient: float = float(
    GAUNT_TABLE[l, m + L_MAX, q + 1, lp, mp + L_MAX]
)

    The offsets map signed quantum numbers to dense array indices.

param l:

Initial state angular momentum.

type l:

int

param m:

Initial state magnetic quantum number.

type m:

int

param q:

Dipole component (-1, 0, or +1).

type q:

int

param lp:

Final state angular momentum.

type lp:

int

param mp:

Final state magnetic quantum number.

type mp:

int

returns:

coefficient – The Gaunt coefficient.

rtype:

float

diffpes.maths.rodrigues_rotation(axis: Float[Array, '3'], angle: float | Float[Array, '']) → Float[Array, '3 3'][source]

Construct a rotation matrix with Rodrigues’ formula.

The function safely normalizes the rotation axis. It then constructs an active rotation matrix for column vectors.

See:

TestRodriguesRotation

Parameters:

  • axis (Float[Array, '3']) – Rotation axis in Cartesian coordinates.

  • angle (Union[float, Float[Array, '']]) – Active rotation angle in radians.

Returns:

rotation – Active Cartesian rotation matrix.

Return type:

Float[Array, '3 3']

Notes

The function uses \(R = I + \sin(\alpha)[\hat n]_\times + (1 - \cos(\alpha))[\hat n]_\times^2\). The safe normalization gives the identity matrix for a zero axis.

diffpes.maths.safe_arccos(x: Float[Array, '...']) → Float[Array, '...'][source]

Evaluate arccos with saturated values and zero boundary gradients.

Computes real angles on the closed cosine range. Values outside the range use the nearest physical endpoint.

See:

TestSafeArccos

Parameters:

x (']) – Real cosine values.

Returns:

angles – Angles in radians, saturated to pi below -1 and zero above 1.

Return type:

']

Notes

Inputs strictly inside (-1, 1) use the ordinary arccos operation. Constants supply values at or beyond either endpoint. This selection gives zero subgradients and avoids the infinite endpoint derivative.

diffpes.maths.safe_arctan2(y: Float[Array, '...'], x: Float[Array, '...']) → Float[Array, '...'][source]

Evaluate arctan2 with a zero value and gradient at the origin.

Computes four-quadrant angles for broadcast-compatible Cartesian coordinates. The origin uses an explicit boundary convention.

See:

TestSafeArctan2

Parameters:

  • y (']) – Vertical coordinates.

  • x (']) – Horizontal coordinates broadcast-compatible with y.

Returns:

angles – Four-quadrant angles in radians, with zero at (0, 0).

Return type:

']

Notes

At the indeterminate origin, sanitized coordinates (0, 1) keep the inactive branch finite. The selected value and both coordinate subgradients at the origin are zero.

diffpes.maths.safe_divide(numerator: Float[Array, '...'], denominator: Float[Array, '...'], fallback: float | Float[Array, ''] = 0.0) → Float[Array, '...'][source]

Divide with a fallback and zero quotient gradients at zero denominators.

Applies elementwise guarded division to broadcast-compatible real arrays. The fallback defines the value on the zero-denominator boundary.

See:

TestSafeDivide

Parameters:

  • numerator (']) – Dividend array.

  • denominator (']) – Divisor array broadcast-compatible with numerator.

  • fallback (Union[float, Float[Array, '']], default: 0.0) – Value returned wherever denominator is zero.

Returns:

quotient – Broadcast quotient with fallback at zero denominators.

Return type:

']

Notes

The function replaces a zero denominator with one before it evaluates the inactive division branch. At this boundary, both quotient operands have zero selected subgradients. A traced fallback retains its usual selected-value gradient.

diffpes.maths.safe_log(x: Float[Array, '...'], floor: float | Float[Array, ''] = 1e-300) → Float[Array, '...'][source]

Evaluate log with a finite floor and zero gradients below it.

Computes the natural logarithm on a positive guarded domain. The floor keeps the returned values finite for small or nonpositive inputs.

See:

TestSafeLog

Parameters:

  • x (']) – Real input array.

  • floor (Union[float, Float[Array, '']], default: 1e-300) – Positive lower bound used before taking the logarithm.

Returns:

logarithms – Natural logarithms of maximum(x, floor).

Return type:

']

Notes

The function replaces inputs at or below the positive floor before it computes the logarithm. Their selected subgradient with respect to x is zero.

diffpes.maths.safe_norm(x: Float[Array, '... n'], axis: int = -1, keepdims: bool = False) → Float[Array, '...'][source]

Compute a Euclidean norm with a zero gradient at zero vectors.

Reduces the selected vector axis with a guarded square root. The operation supports batched vectors and an optional retained axis.

See:

TestSafeNorm

Parameters:

  • x ( n']) – Real vectors.

  • axis (int, default: -1) – (static — a compile-time constant; changing it triggers retracing) Axis containing vector components.

  • keepdims (bool, default: False) – (static — a compile-time constant; changing it triggers retracing) Whether the reduced axis remains with length one.

Returns:

norms – Euclidean norms reduced along axis.

Return type:

']

Notes

The function passes the squared norm to safe_sqrt(). A zero vector has value zero and a zero selected gradient.

diffpes.maths.safe_power(x: Float[Array, '...'], exponent: float | Float[Array, '']) → Float[Array, '...'][source]

Raise positive inputs to a power and return zero otherwise.

Computes real powers on positive bases for arbitrary real exponents. A guarded branch keeps fractional powers outside that domain finite.

See:

TestSafePower

Parameters:

  • x (']) – Real bases.

  • exponent (Union[float, Float[Array, '']]) – Real exponent, including non-integer values.

Returns:

powers – x**exponent for positive x and zero otherwise.

Return type:

']

Notes

The function replaces nonpositive bases with one before exponentiation. This replacement prevents complex or invalid results from fractional powers. Both inputs have zero selected subgradients on the guarded set.

diffpes.maths.safe_sqrt(x: Float[Array, '...']) → Float[Array, '...'][source]

Evaluate sqrt on positive inputs and return zero otherwise.

Applies the principal real square root only on its positive domain. The guarded branch defines a finite value outside that domain.

See:

TestSafeSqrt

Parameters:

x (']) – Real input array.

Returns:

roots – Principal square roots, with zero for x <= 0.

Return type:

']

Notes

The function replaces nonpositive inputs with one before it computes the square root. The selected value and subgradient are zero for x <= 0.

diffpes.maths.real_spherical_harmonic(l: int, m: int, theta: Float[Array, '...'], phi: Float[Array, '...']) → Float[Array, '...'][source]

Evaluate a single real spherical harmonic.

The function computes \(Y_l^m(\theta, \varphi)\).

The function computes the real-valued spherical harmonic with the Condon-Shortley phase convention:

\[ \begin{align}\begin{aligned}Y_l^m(\theta, \varphi) = \sqrt{2} \, N_l^{|m|} \, P_l^{|m|}(\cos\theta) \, \cos(m\varphi) \quad (m > 0)\\Y_l^0(\theta, \varphi) = N_l^0 \, P_l^0(\cos\theta)\\Y_l^m(\theta, \varphi) = (-1)^{|m|} \sqrt{2} \, N_l^{|m|} \, P_l^{|m|}(\cos\theta) \, \sin(|m|\varphi) \quad (m < 0)\end{aligned}\end{align} \]

For negative m, the \((-1)^{|m|}\) factor cancels the Condon-Shortley phase. The _associated_legendre_plm function includes this phase in \(P_l^{|m|}\). The result follows the real-to-complex convention of the Gaunt table. Consequently, the complex-basis Gaunt coefficients match direct integrals of these real harmonics.

The _normalization function computes \(N_l^{|m|}\). It uses exact integer arithmetic for the factorial ratio.

JAX can transform this function and differentiate it with respect to theta and phi.

See:

TestRealSphericalHarmonic

Implementation Logic

  1. Compute the associated Legendre values:

    cos_theta: Float[Array, " ..."] = jnp.cos(theta)
am: int = abs(m)
plm: Float[Array, " ..."] = _associated_legendre_plm(
    l, am, cos_theta
)

    These values provide the angular basis for each real branch.

  2. Select the real harmonic branch:

    if m > 0:
    ylm = jnp.sqrt(2.0) * norm * plm * jnp.cos(m * phi)
    return ylm

    The sign of m selects the cosine, constant, or sine basis.

param l:

Degree (0 <= l).

type l:

int

param m:

Order (-l <= m <= l).

type m:

int

param theta:

Polar angle from z-axis in radians.

type theta:

']

param phi:

Azimuthal angle in radians.

type phi:

']

returns:

ylm – Real spherical harmonic values.

rtype:

']

raises ValueError:

If l is negative or abs(m) exceeds l.

diffpes.maths.real_spherical_harmonics_all(l_max: int, theta: Float[Array, '...'], phi: Float[Array, '...']) → Float[Array, '... N'][source]

Evaluate all real spherical harmonics up to l_max.

The function computes every real spherical harmonic \(Y_l^m(\theta, \varphi)\) for \(0 \le l \le l_{\max}\) and \(-l \le m \le l\), returning them stacked along a new trailing axis.

The ordering is the standard “degree-then-order” layout:

\[(0,0),\;(1,-1),(1,0),(1,1),\;(2,-2),\ldots,(2,2),\;\ldots\]

so that the index into the last axis for a given (l, m) pair is \(l^2 + l + m\). The total number of harmonics is \((l_{\max}+1)^2\).

Python unrolls the loop over (l, m) during tracing. The real_spherical_harmonic function computes each harmonic independently. jnp.stack joins the results. Each l value therefore starts a separate Legendre recurrence chain. The implementation does not share recurrence values across l values.

See:

TestRealSphericalHarmonicsAll

Implementation Logic

  1. Collect harmonics in canonical order:

    for l in range(l_max + 1):
    for m in range(-l, l + 1):
        results.append(
            real_spherical_harmonic(l, m, theta, phi)
        )

    The nested ranges produce the documented degree-then-order sequence.

  2. Stack the harmonic fields:

    ylm_all: Float[Array, " ... N"] = jnp.stack(
    results, axis=-1
)

    The new trailing axis indexes each (l, m) pair.

param l_max:

Maximum angular momentum.

type l_max:

int

param theta:

Polar angle from z-axis.

type theta:

']

param phi:

Azimuthal angle.

type phi:

']

returns:

ylm_all – All spherical harmonics stacked along the last axis, where N = (l_max + 1)**2.

rtype:

N']