"""Compute photoionization cross-section weights for ARPES.
Extended Summary
----------------
The module provides orbital-dependent photoionization cross-section
calculations based on Yeh-Lindau tabulated data and simple
heuristic models for different photon energies.
Routine Listings
----------------
:func:`heuristic_weights`
Compute heuristic orbital weights based on photon energy.
:func:`yeh_lindau_weights`
Compute Yeh-Lindau cross-section weights per orbital.
Notes
-----
Simplified tabulations from Yeh and Lindau provide the cross-section data.
Their 1985 article appears in Atomic Data and Nuclear Data Tables 32, 1-155.
Values at 20, 40, and 60 eV approximate the s, p, and d cross-sections.
"""
import jax.numpy as jnp
from beartype import beartype
from jaxtyping import Array, Float, jaxtyped
from diffpes.types import (
CROSS_SECTION_ENERGIES,
CROSS_SECTION_SIGMA_D,
CROSS_SECTION_SIGMA_P,
CROSS_SECTION_SIGMA_S,
ScalarFloat,
)
[docs]
@jaxtyped(typechecker=beartype)
def heuristic_weights(
photon_energy: ScalarFloat,
) -> Float[Array, " 9"]:
"""Compute heuristic orbital weights based on photon energy.
Provides a simple two-regime model for orbital-dependent
photoionization cross-sections. This is a simplified approximation
useful when tabulated cross-section data is unavailable.
:see: :class:`~.test_crosssections.TestHeuristicWeights`
Implementation Logic
--------------------
The function selects between two pre-defined weight vectors
based on a 50 eV energy threshold:
1. **Below 50 eV (low-energy regime)**::
weights = [1, 2, 2, 2, 1, 1, 1, 1, 1]
Give p-orbitals at indices 1-3 a weight of 2. This value
reflecting the stronger p-orbital cross-section at low
photon energies typical of He-I or laser ARPES.
2. **Above 50 eV (high-energy regime)**::
weights = [1, 1, 1, 1, 2, 2, 2, 2, 2]
Give d-orbitals at indices 4-8 a weight of 2. This value
reflecting the resonant enhancement of d-orbital
cross-sections at higher photon energies, for example He-II or
synchrotron).
The function uses ``jnp.where`` for JAX transformations without Python
branching.
Parameters
----------
photon_energy : ScalarFloat
Incident photon energy in eV.
Returns
-------
weights : Float[Array, " 9"]
Orbital weights for [s, py, pz, px, dxy, dyz, dz2,
dxz, dx2-y2].
Notes
-----
The 50 eV regime choice is a hard selector. Its photon-energy gradient is
zero away from the threshold and undefined at the threshold. Therefore,
this heuristic contributes no continuous photon-energy Jacobian column.
Plan 06 replaces this approximation with a differentiable matrix element
treatment.
"""
low_e: Float[Array, " 9"] = jnp.array(
[1.0, 2.0, 2.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.0],
dtype=jnp.float64,
)
high_e: Float[Array, " 9"] = jnp.array(
[1.0, 1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 2.0, 2.0],
dtype=jnp.float64,
)
_threshold_ev: float = 50.0
weights: Float[Array, " 9"] = jnp.where(
photon_energy < _threshold_ev, low_e, high_e
)
return weights
def _interp_cross_section(
photon_energy: Float[Array, " "],
sigma_vals: Float[Array, " 3"],
) -> Float[Array, " "]:
"""Interpolate a cross-section with linear extrapolation.
Performs piecewise-linear interpolation of tabulated cross-section
values at the given photon energy, with constant extrapolation
outside the tabulated range.
Implementation Logic
--------------------
1. **Interpolate via jnp.interp**:
sigma = jnp.interp(photon_energy, CROSS_SECTION_ENERGIES, sigma_vals)
- ``CROSS_SECTION_ENERGIES`` = [20, 40, 60] eV are the
tabulation points.
- ``sigma_vals`` are the corresponding cross-section values.
- ``jnp.interp`` performs piecewise-linear interpolation
between adjacent tabulation points.
- At the boundaries, energies below 20 eV return sigma_vals[0].
Energies above 60 eV return sigma_vals[-1]. ``jnp.interp`` uses this
constant extrapolation by default.
Parameters
----------
photon_energy : Float[Array, " "]
Photon energy in eV.
sigma_vals : Float[Array, " 3"]
Cross-section values at [20, 40, 60] eV.
Returns
-------
sigma : Float[Array, " "]
Interpolated cross-section value.
"""
sigma: Float[Array, " "] = jnp.interp(
photon_energy, CROSS_SECTION_ENERGIES, sigma_vals
)
return sigma
[docs]
@jaxtyped(typechecker=beartype)
def yeh_lindau_weights(
photon_energy: ScalarFloat,
) -> Float[Array, " 9"]:
"""Compute Yeh-Lindau cross-section weights per orbital.
Interpolates tabulated photoionization cross-sections from
Yeh & Lindau (1985) [2]_ to produce orbital-resolved weights
at the specified photon energy.
:see: :class:`~.test_crosssections.TestYehLindauWeights`
Implementation Logic
--------------------
1. **Cast photon energy to float64**::
pe = jnp.asarray(photon_energy, dtype=float64)
Ensures consistent precision for the interpolation.
2. **Interpolate s, p, d cross-sections independently**::
s_w = _interp_cross_section(pe, CROSS_SECTION_SIGMA_S)
p_w = _interp_cross_section(pe, CROSS_SECTION_SIGMA_P)
d_w = _interp_cross_section(pe, CROSS_SECTION_SIGMA_D)
Each call linearly interpolates the corresponding
tabulated values at 20, 40, and 60 eV. The tabulated
data (``CROSS_SECTION_SIGMA_S``, ``CROSS_SECTION_SIGMA_P``,
``CROSS_SECTION_SIGMA_D``) encodes simplified
Yeh-Lindau cross-sections for s, p, and d subshells.
3. **Broadcast to 9-orbital weight vector**::
weights = [s_w, p_w, p_w, p_w, d_w, d_w, d_w, d_w, d_w]
Maps the three subshell cross-sections onto the full
9-orbital basis: 1 s-orbital, 3 p-orbitals (each
receiving p_w), and 5 d-orbitals (each receiving d_w).
Parameters
----------
photon_energy : ScalarFloat
Incident photon energy in eV.
Returns
-------
weights : Float[Array, " 9"]
Cross-section weights for [s, py, pz, px, dxy, dyz,
dz2, dxz, dx2-y2].
References
----------
.. [2] Yeh & Lindau, "Atomic subshell photoionization cross
sections and asymmetry parameters: 1 <= Z <= 103", Atomic
Data and Nuclear Data Tables 32, 1-155 (1985).
"""
pe: Float[Array, " "] = jnp.asarray(photon_energy, dtype=jnp.float64)
s_w: Float[Array, " "] = _interp_cross_section(pe, CROSS_SECTION_SIGMA_S)
p_w: Float[Array, " "] = _interp_cross_section(pe, CROSS_SECTION_SIGMA_P)
d_w: Float[Array, " "] = _interp_cross_section(pe, CROSS_SECTION_SIGMA_D)
weights: Float[Array, " 9"] = jnp.array(
[s_w, p_w, p_w, p_w, d_w, d_w, d_w, d_w, d_w],
dtype=jnp.float64,
)
return weights
__all__: list[str] = [
"heuristic_weights",
"yeh_lindau_weights",
]