diffpes.radial

Differentiable radial primitives for ARPES matrix elements: spherical Bessel functions, atomic radial wavefunctions, and fixed-grid radial integrals.

Provide differentiable radial primitives for ARPES matrix elements.

Extended Summary

The subpackage provides JAX-compatible spherical Bessel functions and atomic radial wavefunctions. It also computes radial integrals on a fixed grid. The dipole matrix element pipeline uses these functions. The central quantity is the radial integral

\[B^{l'}(k) = (i)^{l'} \int_0^\infty R(r)\, r^3\, j_{l'}(kr)\, dr\]

The functions compute this integral with composite trapezoidal quadrature.

The following list describes the submodules:

  • bessel

    Compute spherical Bessel functions in JAX.

  • integrate

    Evaluate radial integrals on fixed grids.

  • wavefunctions

    Evaluate atomic radial wavefunction models in JAX.

Routine Listings

hydrogenic_radial()

Evaluate normalized hydrogenic radial function.

radial_integral()

Evaluate dipole radial integral on a fixed radial grid.

slater_radial()

Evaluate normalized Slater-type radial function.

spherical_bessel_jl()

Evaluate spherical Bessel function \(j_l(x)\).

Notes

All functions support JAX transformations and automatic differentiation. The Bessel and Laguerre recurrences use jax.lax.fori_loop.

diffpes.radial.spherical_bessel_jl(order: int, x: Float[Array, '...']) → Float[Array, '...'][source]

Evaluate spherical Bessel function \(j_l(x)\).

The function computes the spherical Bessel function of the first kind. It uses closed-form seeds for l=0 and l=1. It uses an upward recurrence for l >= 2.

Seed values:

\[ \begin{align}\begin{aligned}j_0(x) = \frac{\sin x}{x}\\j_1(x) = \frac{\sin x}{x^2} - \frac{\cos x}{x}\end{aligned}\end{align} \]

Upward recurrence (for l >= 2, starting from l=1):

\[j_{l+1}(x) = \frac{2l + 1}{x} \, j_l(x) - j_{l-1}(x)\]

The function implements the recurrence with jax.lax.fori_loop. The loop runs from index 1 to order - 1. It carries the pair \((j_{l-1}, j_l)\) and produces \(j_{l+1}\) at each step.

Small-argument limit:

For \(|x| < 10^{-8}\) (the module constant SMALL_ARGUMENT), the function uses the leading-order Taylor expansion:

\[j_l(x) \approx \frac{x^l}{(2l+1)!!}\]

This expansion avoids division by zero in the seed formulas. The small_mask value selects the recurrence or the small-argument limit. The x_safe input replaces the masked entries with 1.0. This replacement prevents invalid recurrence values from affecting the gradients.

Numerical stability notes:

  • The upward recurrence remains stable when \(j_l(x)\) dominates the recurrence at moderate x.

  • A downward Miller recurrence offers more stability when l is very large relative to x.

  • ARPES simulations do not use this regime. They use l <= 5 and usually use kr values from O(1) to O(10).

  • The jnp.where mask defines the gradients everywhere, including at x = 0.

See:

TestSphericalBesselJl

Parameters:

  • order (int) – Non-negative angular momentum order.

  • x (']) – Real argument array.

Returns:

values – j_l(x) evaluated element-wise with the same shape as x.

Return type:

']

Raises:

ValueError – If order is negative.

Notes

The order parameter is a Python integer, not a JAX tracer. A change to this parameter causes JAX to trace the function again. The loop bounds and the double-factorial constant depend on order statically. The series branch preserves the nonzero analytic limiting gradient near zero.

diffpes.radial.radial_integral(k: Float[Array, '...'], r: Float[Array, 'R'], radial_values: Float[Array, 'R'], l_prime: int) → Complex[Array, '...'][source]

Evaluate dipole radial integral on a fixed radial grid.

The function computes the dipole radial integral in the partial-wave expansion of the photoemission matrix element:

\[B^{l'}(k) = i^{l'} \int_0^\infty R(r) \, r^3 \, j_{l'}(k r) \, dr\]

The computation proceeds in five steps:

  1. Form the outer product: Form the 2-D array \(kr\) by broadcasting k (shape ...) against r (shape R) via jnp.expand_dims. This produces shape (... , R).

  2. Evaluate the Bessel function: Evaluate \(j_{l'}(kr)\) element-wise using spherical_bessel_jl. This function handles the small-argument limit internally.

  3. Assemble the integrand: Multiply the Bessel values by \(R(r) \cdot r^3\). The \(r^3\) factor combines the volume element and the dipole factor.

  4. Apply trapezoidal quadrature: Integrate over the radial grid using jnp.trapezoid along the last axis. The grid need not be uniform; jnp.trapezoid handles non-uniform spacing via the x argument.

  5. Apply the phase factor: Multiply the real integral by the complex phase \(i^{l'}\) to obtain the final complex-valued result.

The function supports batched evaluation. If k has shape (...), the output retains its leading shape. The integration contracts the last axis.

See:

TestRadialIntegral

Parameters:

  • k (']) – Momentum magnitudes at which the function computes the integral.

  • r (Float[Array, 'R']) – Monotonic radial grid.

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

  • l_prime (int) – Final-state angular momentum order.

Returns:

values – Complex radial integral values with the same leading shape as k.

Return type:

']

Raises:

ValueError – If l_prime is negative.

Notes

The result accuracy depends on the density and extent of the radial grid. Slater-type orbitals decay as \(e^{-\zeta r}\). A grid with 200–500 points usually suffices when it extends to \(r \sim 10/\zeta\).

The function uses the trapezoidal rule for its simplicity and JAX compatibility. Higher-order quadrature can improve accuracy for smooth integrands on coarse grids. Such quadrature requires a custom JAX-traceable implementation.

diffpes.radial.hydrogenic_radial(r: Float[Array, '...'], n: int, angular_momentum: int, z_eff: float | Float[Array, '']) → Float[Array, '...'][source]

Evaluate normalized hydrogenic radial function.

The function computes the exact radial wavefunction for a hydrogenic atom. The atom has one electron and an effective nuclear charge \(Z_{\text{eff}}\):

\[R_{n,l}(r) = N_{n,l} \, e^{-\rho/2} \, \rho^l \, L_{n-l-1}^{2l+1}(\rho)\]

Here, \(\rho = 2 Z_{\text{eff}} r / n\) is the scaled radial coordinate. The _associated_laguerre function computes the generalized Laguerre polynomial \(L_{n-l-1}^{2l+1}\).

Normalization:

The normalization constant is:

\[N_{n,l} = \left(\frac{2 Z_{\text{eff}}}{n}\right)^{3/2} \sqrt{\frac{(n - l - 1)!}{2n \cdot (n + l)!}}\]

This ensures \(\int_0^\infty |R_{n,l}(r)|^2 r^2 dr = 1\). The function computes the factorial ratio with Python’s math.factorial. This operation uses exact integer arithmetic. The function then converts the ratio to a JAX scalar with jnp.sqrt.

Hydrogenic vs. Slater model:

Unlike Slater-type orbitals (which are node-free exponentials), hydrogenic radial functions have \(n - l - 1\) radial nodes encoded by the zeros of the Laguerre polynomial. This makes them exact solutions for hydrogen-like atoms but less commonly used as basis functions in multi-electron calculations.

Laguerre polynomial recurrence:

The _associated_laguerre function computes \(L_{n-l-1}^{2l+1}(\rho)\) with an upward three-term recurrence. The recurrence starts at order 0 and ends at \(n - l - 1\). It remains stable in the upward direction. The implementation uses jax.lax.fori_loop for JAX transformations.

See:

TestHydrogenicRadial

Parameters:

  • r (']) – Radial coordinate in atomic units.

  • n (int) – Principal quantum number.

  • angular_momentum (int) – Angular momentum quantum number (0 <= angular_momentum < n).

  • z_eff (Union[float, Float[Array, '']]) – Effective nuclear charge.

Returns:

values – R_{n,l}(r) for hydrogenic orbitals.

Return type:

']

Raises:

ValueError – If n is less than one or angular_momentum lies outside [0, n).

Notes

The z_eff parameter is a JAX array that supports automatic differentiation. The quantum numbers n and angular_momentum are Python integers. They control the Laguerre polynomial order and remain static in the traced computation graph.

diffpes.radial.slater_radial(r: Float[Array, '...'], n: int, zeta: float | Float[Array, '']) → Float[Array, '...'][source]

Evaluate normalized Slater-type radial function.

The function computes the Slater-type orbital (STO) radial function:

\[R(r) = N \, r^{n-1} \, e^{-\zeta r}\]

The normalization constant \(N\) satisfies \(\int_0^\infty |R(r)|^2 r^2 dr = 1\):

\[N = \frac{(2\zeta)^{n + 1/2}}{\sqrt{(2n)!}}\]

Slater vs. hydrogenic models:

Slater-type orbitals are simpler than hydrogenic radial functions because they lack the associated Laguerre polynomial factor. They have the correct exponential decay and cusp behavior at the nucleus, making them popular as basis functions in quantum chemistry. However, they do not possess radial nodes (except at r = 0 and r = infinity), unlike the exact hydrogenic solutions.

The Slater exponent \(\zeta\) represents the effective nuclear charge and screening. A fit to Hartree-Fock atomic orbitals usually determines this exponent. Variational optimization provides another method.

Normalization derivation:

The radial normalization integral is:

\[\int_0^\infty r^{2(n-1)} e^{-2\zeta r} r^2 dr = \int_0^\infty r^{2n} e^{-2\zeta r} dr = \frac{(2n)!}{(2\zeta)^{2n+1}}\]

Setting \(N^2 \cdot (2n)! / (2\zeta)^{2n+1} = 1\) gives the formula above.

See:

TestSlaterRadial

Parameters:

  • r (']) – Radial coordinate in atomic units.

  • n (int) – Principal quantum number (n >= 1).

  • zeta (Union[float, Float[Array, '']]) – Slater exponent.

Returns:

values – Normalized radial function R(r) = N r^(n-1) exp(-zeta * r).

Return type:

']

Raises:

ValueError – If n is less than one.

Notes

The zeta parameter is a JAX array, not a Python float. Therefore, automatic differentiation can include this parameter. Inverse workflows can use its gradient to optimize Slater exponents.