diffpes.simul¶

Free-electron kinematics and ARPES simulation functions at six complexity levels.

Provide ARPES simulation functions at six complexity levels.

Extended Summary¶

The subpackage provides a complete ARPES simulation pipeline from pseudo-Voigt convolution to full polarization-dependent dipole matrix element calculations. Also exports broadening functions, cross-section models, polarization utilities, and orbital angular momentum.

The following list describes the submodules:

broadening
Compute energy broadening functions for ARPES simulations.
crosssections
Compute photoionization cross-section weights for ARPES.
expanded
Run expanded-input workflows for ARPES simulation.
forward
Run an end-to-end differentiable ARPES forward model.
kinematics
Compute free-electron photoemission kinematics.
oam
Compute orbital angular momentum.
polarization
Compute photon polarization and detector-frame transformations.
resolution
Apply momentum resolution broadening to ARPES simulations.
self_energy
Evaluate energy-dependent self-energy for ARPES simulations.
spectrum
Simulate ARPES spectra at six complexity levels.
workflow
Run high-level workflows for VASP-to-ARPES simulation.

Routine Listings¶

apply_momentum_broadening(): Convolve I(k, E) with a Gaussian in k-space.
build_efield(): Compute electric field vector from polarization config.
build_polarization_vectors(): Construct s- and p-polarization basis vectors.
detector_angles_to_kpar(): Convert detector angles to parallel momentum.
detector_rotation(): Build the detector-frame rotation.
compute_oam(): Compute orbital angular momentum z-component.
dipole_matrix_elements(): Compute dipole matrix elements for all 9 orbitals.
evaluate_self_energy(): Evaluate the imaginary self-energy \(\Gamma(E)\).
emission_angles(): Convert Cartesian momentum to emission angles.
fermi_dirac(): Compute Fermi-Dirac distribution value.
final_state_k_inv_ang(): Convert kinetic energy to final-state momentum magnitude.
gaussian(): Compute normalized Gaussian broadening profile.
heuristic_weights(): Compute heuristic orbital weights based on photon energy.
kinetic_energy_ev(): Compute the floored photoelectron kinetic energy.
kpar_to_detector_angles(): Convert parallel momentum to detector angles.
kz_from_inner_potential(): Compute complex out-of-plane momentum from the inner potential.
load_vasp_context(): Load a simulation-ready context from VASP output files.
photon_wavevector(): Build the unit photon wavevector from incidence angles.
polarization_from_angles(): Construct polarization from incidence angles.
polarization_to_spherical(): Convert Cartesian polarization to spherical components.
prepare_projection(): Prepare orbital projections for simulation.
run_vasp_workflow(): Run an end-to-end VASP-to-ARPES workflow in one call.
rotate_frame_vectors(): Rotate a real vector across a detector-angle grid.
rotate_polarization_grid(): Rotate polarization across a detector-angle grid.
simulate_advanced(): Simulate ARPES with Gaussian broadening and polarization rules.
simulate_advanced_expanded(): Run advanced-level ARPES simulation from plain arrays.
simulate_basic(): Simulate ARPES spectrum with Gaussian broadening and heuristic weights.
simulate_basic_expanded(): Run basic-level ARPES simulation from plain arrays.
simulate_basicplus(): Simulate ARPES with Gaussian broadening and Yeh-Lindau cross-sections.
simulate_basicplus_expanded(): Run basicplus-level ARPES simulation from plain arrays.
simulate_context(): Run a level-dispatched simulation from a loaded workflow context.
simulate_expanded(): Dispatch an expanded-input simulation by complexity level.
simulate_expert(): Simulate ARPES with Voigt broadening and dipole matrix elements.
simulate_expert_expanded(): Run expert-level ARPES simulation from plain arrays.
simulate_novice(): Simulate ARPES spectrum with Voigt broadening and uniform weights.
simulate_novice_expanded(): Run novice-level ARPES simulation from plain arrays.
simulate_soc(): Simulate ARPES with spin-orbit coupling (spin-dependent intensity).
simulate_soc_expanded(): Run SOC (spin-orbit coupling) ARPES simulation from plain arrays.
simulate_tb_radial(): Run the end-to-end differentiable ARPES forward model.
voigt(): Compute a normalized Thompson-Cox-Hastings pseudo-Voigt profile.
yeh_lindau_weights(): Compute Yeh-Lindau cross-section weights per orbital.

Notes

All simulation functions are JAX-compatible and use jax.vmap for vectorized evaluation across k-points and bands.

diffpes.simul.fermi_dirac(energy: float | Float[Array, ''], fermi_energy: float | Float[Array, ''], temperature: float | Float[Array, '']) → Float[Array, ''][source]¶

Compute Fermi-Dirac distribution value.

Evaluates the Fermi-Dirac thermal occupation function at a given energy, Fermi level, and temperature:

f(E) = 1 / (1 + exp((E - Ef) / (kB * T)))

See:: TestFermiDirac

Implementation Logic¶

Compute thermal energy kT:
```
kt = kB * T
```
Multiplies the Boltzmann constant kB = 8.617333e-5 eV/K by the temperature in Kelvin to obtain the thermal energy scale. Both values are cast to float64 for numerical precision.
Guard against T = 0:
```
safe_kt = max(kt, 1e-10)
```
At zero temperature the distribution becomes a step function, but the exponential would diverge. Clamping kT to a small positive value (1e-10 eV) avoids division by zero while preserving the sharp step-function behavior numerically.
Evaluate Fermi-Dirac function:
```
exponent = (E - Ef) / safe_kt
occupation = sigmoid(-exponent)
```
Computes the occupation probability. For E << Ef the result approaches 1 (filled states); for E >> Ef it approaches 0 (empty states).

param energy:: Electron energy in eV.
type energy:: Union[float, Float[Array, '']]
param fermi_energy:: Fermi level energy in eV.
type fermi_energy:: Union[float, Float[Array, '']]
param temperature:: Temperature in Kelvin.
type temperature:: Union[float, Float[Array, '']]
returns:: occupation – Fermi-Dirac occupation (0 to 1).
rtype:: Float[Array, '']

Notes

Uses the Boltzmann constant kB = 8.617333e-5 eV/K, imported as KB_EV_PER_K. jax.nn.sigmoid is algebraically identical to the reciprocal-exponential expression but has an overflow-safe JVP. Values and derivatives therefore underflow to finite exact zeros far above the Fermi level instead of becoming NaN. The existing 1e-10 eV thermal-scale clamp keeps the public function total at nonpositive temperature; validated simulation parameters require a strictly positive temperature.

diffpes.simul.gaussian(energy_range: Float[Array, 'E'], center: float | Float[Array, ''], sigma: float | Float[Array, '']) → Float[Array, 'E'][source]¶

Compute normalized Gaussian broadening profile.

Evaluates a Gaussian lineshape centered at center with standard deviation sigma, normalized so that the integral over all energies equals unity.

See:: TestGaussian

Implementation Logic¶

The function evaluates the analytic Gaussian probability density:

G(E) = exp(-(E - E0)^2 / (2 * sigma^2))
       / (sqrt(2 * pi) * sigma)

Compute energy differences:
```
diff = energy_range - center
```
Shifts the energy axis so the peak is at the origin.
Compute normalization factor:
```
norm_factor = sqrt(2 * pi) * sigma
```
This prefactor ensures the profile integrates to unity over (-inf, +inf). Thus, the Gaussian has a unit area.
Evaluate Gaussian profile:
```
profile = exp(-diff^2 / (2 * sigma^2)) / norm_factor
```
Element-wise evaluation of the normalized Gaussian at each energy point.

param energy_range:: Energy axis values in eV.
type energy_range:: Float[Array, 'E']
param center:: Center energy of the peak in eV.
type center:: Union[float, Float[Array, '']]
param sigma:: Gaussian standard deviation in eV.
type sigma:: Union[float, Float[Array, '']]
returns:: profile – Normalized Gaussian profile values.
rtype:: Float[Array, 'E']

diffpes.simul.voigt(energy_range: Float[Array, 'E'], center: float | Float[Array, ''], sigma: float | Float[Array, ''], gamma: float | Float[Array, '']) → Float[Array, 'E'][source]¶

Compute a normalized Thompson-Cox-Hastings pseudo-Voigt profile.

The function uses the pseudo-Voigt method from Thompson, Cox, and Hastings (1987) [1]. This method expresses the Voigt profile as a linear combination of Gaussian and Lorentzian components:

V(E) = eta * L(E) + (1 - eta) * G(E)

See:: TestVoigt

where eta is an empirically determined mixing ratio. This approximation is accurate to better than 1% relative error.

Implementation Logic¶

The pseudo-Voigt approximation proceeds in four stages:

Compute component FWHMs:
```
f_G = 2 * sigma * sqrt(2 * ln2)   (Gaussian FWHM)
f_L = 2 * gamma                    (Lorentzian FWHM)
```
Converts the Gaussian standard deviation and Lorentzian half-width to their respective full-width at half-maximum values.
Compute Voigt FWHM via empirical formula:
```
f_V = (f_G^5 + 2.69269 * f_G^4 * f_L
       + 2.42843 * f_G^3 * f_L^2
       + 4.47163 * f_G^2 * f_L^3
       + 0.07842 * f_G * f_L^4
       + f_L^5)^(1/5)
```
The Thompson-Cox-Hastings empirical relation approximates the FWHM of the true Voigt convolution from the component FWHMs. The function sanitizes the width polynomial before the fractional power. Thus, an inactive branch cannot produce invalid gradients.
Compute mixing ratio eta:
```
ratio = f_L / f_V
eta = 1.36603 * ratio - 0.47719 * ratio^2
      + 0.11116 * ratio^3
```
The mixing ratio interpolates between pure Gaussian (eta = 0) and pure Lorentzian (eta = 1). For non-negative physical widths, the Thompson-Cox-Hastings polynomial keeps it in [0, 1].

Combine Gaussian and Lorentzian components:

sigma_V = f_V / (2 * sqrt(2 * ln2))
gamma_V = f_V / 2
G = gaussian(energy_range, center, sigma_V)
L = gamma_V / (pi * (diff^2 + gamma_V^2))
profile = eta * L + (1 - eta) * G

Both components use the Voigt FWHM (not the original widths) so that the combined profile has the correct total width.

param energy_range:: Energy axis values in eV.
type energy_range:: Float[Array, 'E']
param center:: Center energy of the peak in eV.
type center:: Union[float, Float[Array, '']]
param sigma:: Gaussian standard deviation in eV.
type sigma:: Union[float, Float[Array, '']]
param gamma:: Lorentzian half-width at half-maximum in eV.
type gamma:: Union[float, Float[Array, '']]
returns:: profile – Normalized pseudo-Voigt profile values.
rtype:: Float[Array, 'E']
raises EquinoxRuntimeError:: If sigma and gamma are both zero. The normalized profile and its directional derivative have no definition at this point.

Notes

Quotients use diffpes.maths.safe_divide(), so inactive zero-width branches cannot inject NaNs into reverse-mode gradients. The pure-Gaussian ray gamma = 0 and pure-Lorentzian ray sigma = 0 retain their finite boundary sensitivities. The function rejects only their singular intersection.

References

diffpes.simul.heuristic_weights(photon_energy: float | Float[Array, '']) → Float[Array, '9'][source]¶

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:: TestHeuristicWeights

Implementation Logic¶

The function selects between two pre-defined weight vectors based on a 50 eV energy threshold:

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.
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.

param photon_energy:: Incident photon energy in eV.
type photon_energy:: Union[float, Float[Array, '']]
returns:: weights – Orbital weights for [s, py, pz, px, dxy, dyz, dz2, dxz, dx2-y2].
rtype:: Float[Array, '9']

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.

diffpes.simul.yeh_lindau_weights(photon_energy: float | Float[Array, '']) → Float[Array, '9'][source]¶

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:: TestYehLindauWeights

Implementation Logic¶

Cast photon energy to float64:
```
pe = jnp.asarray(photon_energy, dtype=float64)
```
Ensures consistent precision for the interpolation.
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.
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).

param photon_energy:: Incident photon energy in eV.
type photon_energy:: Union[float, Float[Array, '']]
returns:: weights – Cross-section weights for [s, py, pz, px, dxy, dyz, dz2, dxz, dx2-y2].
rtype:: Float[Array, '9']

References

diffpes.simul.simulate_advanced_expanded(eigenbands: Float[Array, 'K B'], surface_orb: Float[Array, 'K B A 9'], ef: float | Float[Array, ''], sigma: float | Float[Array, ''], fidelity: int, temperature: float | Float[Array, ''], photon_energy: float | Float[Array, ''], polarization: str = 'unpolarized', incident_theta: float | Float[Array, ''] = 45.0, incident_phi: float | Float[Array, ''] = 0.0, polarization_angle: float | Float[Array, ''] = 0.0) → ArpesSpectrum[source]¶

Run advanced-level ARPES simulation from plain arrays.

The function adds polarization-dependent selection rules to the basicplus level. The factor |e . d|^2 weights the photoemission intensity. Here, e is the electric field, and d is the orbital dipole matrix element. For unpolarized light, the function averages the s-polarization and p-polarization contributions.

See:: TestSimulateAdvancedExpanded

Implementation Logic¶

Build carrier inputs:

bands, orb_proj = _build_inputs(
    eigenbands=eigenbands,
    surface_orb=surface_orb,
    ef=ef,
)

This validates the raw arrays and binds the band and projection carriers.

Build simulation settings:

params = make_expanded_simulation_params(
    eigenbands=eigenbands,
    fidelity=fidelity,
    sigma=sigma,
    temperature=temperature,
    photon_energy=photon_energy,
)

This derives the Gaussian energy grid and traced physical settings.

Build polarization settings:

pol = _build_polarization(
    polarization=polarization,
    incident_theta=incident_theta,
    incident_phi=incident_phi,
    polarization_angle=polarization_angle,
)

This converts degree-valued incidence angles through the shared factory.

Evaluate the advanced model:
```
spectrum = simulate_advanced(bands, orb_proj, params, pol)
```
This delegates polarization-weighted intensity to the core model.

param eigenbands:: Band eigenvalues in eV, shape (n_kpoints, n_bands).
type eigenbands:: Float[Array, 'K B']
param surface_orb:: Orbital projection coefficients, shape (n_kpoints, n_bands, n_atoms, 9).
type surface_orb:: Float[Array, 'K B A 9']
param ef:: Fermi energy in eV.
type ef:: Union[float, Float[Array, '']]
param sigma:: Gaussian broadening width in eV.
type sigma:: Union[float, Float[Array, '']]
param fidelity:: Number of points in the energy axis.
type fidelity:: int
param temperature:: Electronic temperature in Kelvin for the Fermi-Dirac distribution.
type temperature:: Union[float, Float[Array, '']]
param photon_energy:: Incident photon energy in eV.
type photon_energy:: Union[float, Float[Array, '']]
param polarization:: Polarization type: "s", "p", "linear", or "unpolarized" (default).
type polarization:: str, default: "unpolarized"
param incident_theta:: Polar angle of the incident beam in degrees. Default 45.
type incident_theta:: Union[float, Float[Array, '']], default: 45.0
param incident_phi:: Azimuthal angle of the incident beam in degrees. Default 0.
type incident_phi:: Union[float, Float[Array, '']], default: 0.0
param polarization_angle:: Rotation angle for arbitrary linear polarization in radians. Default 0.
type polarization_angle:: Union[float, Float[Array, '']], default: 0.0
returns:: spectrum – Simulated ARPES spectrum containing the intensity map and energy axis.
rtype:: ArpesSpectrum

See also

simulate_advanced: Low-level implementation accepting PyTrees.

diffpes.simul.simulate_basic_expanded(eigenbands: Float[Array, 'K B'], surface_orb: Float[Array, 'K B A 9'], ef: float | Float[Array, ''], sigma: float | Float[Array, ''], fidelity: int, temperature: float | Float[Array, ''], photon_energy: float | Float[Array, '']) → ArpesSpectrum[source]¶

Run basic-level ARPES simulation from plain arrays.

Applies Gaussian broadening with energy-dependent heuristic orbital weights. The heuristic enhances p-orbital contributions below 50 eV photon energy and d-orbital contributions above, providing a rough approximation of photoionization cross-section effects without tabulated data.

See:: TestSimulateBasicExpanded

Implementation Logic¶

Build carrier inputs:

bands, orb_proj = _build_inputs(
    eigenbands=eigenbands,
    surface_orb=surface_orb,
    ef=ef,
)

This validates the raw arrays and binds the band and projection carriers.

Build simulation settings:

params = make_expanded_simulation_params(
    eigenbands=eigenbands,
    fidelity=fidelity,
    sigma=sigma,
    temperature=temperature,
    photon_energy=photon_energy,
)

This derives the energy window without introducing a Lorentzian width.

Evaluate the basic model:
```
spectrum = simulate_basic(bands, orb_proj, params)
```
This applies the core Gaussian model and its heuristic orbital weights.

param eigenbands:: Band eigenvalues in eV, shape (n_kpoints, n_bands).
type eigenbands:: Float[Array, 'K B']
param surface_orb:: Orbital projection coefficients, shape (n_kpoints, n_bands, n_atoms, 9).
type surface_orb:: Float[Array, 'K B A 9']
param ef:: Fermi energy in eV.
type ef:: Union[float, Float[Array, '']]
param sigma:: Gaussian broadening width in eV.
type sigma:: Union[float, Float[Array, '']]
param fidelity:: Number of points in the energy axis.
type fidelity:: int
param temperature:: Electronic temperature in Kelvin for the Fermi-Dirac distribution.
type temperature:: Union[float, Float[Array, '']]
param photon_energy:: Incident photon energy in eV. Determines the heuristic orbital weighting regime.
type photon_energy:: Union[float, Float[Array, '']]
returns:: spectrum – Simulated ARPES spectrum containing the intensity map and energy axis.
rtype:: ArpesSpectrum

See also

simulate_basic: Low-level implementation accepting PyTrees.

diffpes.simul.simulate_basicplus_expanded(eigenbands: Float[Array, 'K B'], surface_orb: Float[Array, 'K B A 9'], ef: float | Float[Array, ''], sigma: float | Float[Array, ''], fidelity: int, temperature: float | Float[Array, ''], photon_energy: float | Float[Array, '']) → ArpesSpectrum[source]¶

Run basicplus-level ARPES simulation from plain arrays.

The function applies Gaussian broadening with interpolated Yeh-Lindau photoionization cross-sections. Tabulated atomic data defines these cross-sections. They provide orbital-dependent intensity scaling at each photon energy.

See:: TestSimulateBasicplusExpanded

Implementation Logic¶

Build carrier inputs:

bands, orb_proj = _build_inputs(
    eigenbands=eigenbands,
    surface_orb=surface_orb,
    ef=ef,
)

This validates the raw arrays and binds the band and projection carriers.

Build simulation settings:

params = make_expanded_simulation_params(
    eigenbands=eigenbands,
    fidelity=fidelity,
    sigma=sigma,
    temperature=temperature,
    photon_energy=photon_energy,
)

This retains the photon energy used for Yeh-Lindau interpolation.

Evaluate the basicplus model:
```
spectrum = simulate_basicplus(bands, orb_proj, params)
```
This delegates cross-section weighting and Gaussian broadening to the core.

param eigenbands:: Band eigenvalues in eV, shape (n_kpoints, n_bands).
type eigenbands:: Float[Array, 'K B']
param surface_orb:: Orbital projection coefficients, shape (n_kpoints, n_bands, n_atoms, 9).
type surface_orb:: Float[Array, 'K B A 9']
param ef:: Fermi energy in eV.
type ef:: Union[float, Float[Array, '']]
param sigma:: Gaussian broadening width in eV.
type sigma:: Union[float, Float[Array, '']]
param fidelity:: Number of points in the energy axis.
type fidelity:: int
param temperature:: Electronic temperature in Kelvin for the Fermi-Dirac distribution.
type temperature:: Union[float, Float[Array, '']]
param photon_energy:: Incident photon energy in eV. Used to interpolate Yeh-Lindau cross-section tables.
type photon_energy:: Union[float, Float[Array, '']]
returns:: spectrum – Simulated ARPES spectrum containing the intensity map and energy axis.
rtype:: ArpesSpectrum

See also

simulate_basicplus: Low-level implementation accepting PyTrees.

diffpes.simul.simulate_expanded(level: str, eigenbands: Float[Array, 'K B'], surface_orb: Float[Array, 'K B A 9'], ef: float | Float[Array, ''] = 0.0, sigma: float | Float[Array, ''] = 0.04, gamma: float | Float[Array, ''] = 0.1, fidelity: int = 25000, temperature: float | Float[Array, ''] = 15.0, photon_energy: float | Float[Array, ''] = 11.0, polarization: str = 'unpolarized', incident_theta: float | Float[Array, ''] = 45.0, incident_phi: float | Float[Array, ''] = 0.0, polarization_angle: float | Float[Array, ''] = 0.0, surface_spin: Float[Array, 'K B A 6'] | None = None, ls_scale: float | Float[Array, ''] = 0.01) → ArpesSpectrum[source]¶

Dispatch an expanded-input simulation by complexity level.

This entry point selects one of six simulation functions from level. Only level, eigenbands, and surface_orb require explicit values. The function supplies defaults for all other parameters. It ignores parameters that the selected level does not use.

See:: TestSimulateExpanded

Implementation Logic¶

Normalize the level key:
```
level_key = level.lower()
```
This makes the static dispatcher case-insensitive.

Select the matching wrapper:

if level_key == "novice":
    spectrum = simulate_novice_expanded(
        eigenbands=eigenbands,
        surface_orb=surface_orb,
        ef=ef,
        sigma=sigma,
        gamma=gamma,
        fidelity=fidelity,
        temperature=temperature,
        photon_energy=photon_energy,
    )

The remaining elif branches use the same explicit forwarding pattern, so each level receives only its public parameters.

Reject unsupported levels:

msg: str = (
    "Unknown simulation level. "
    "Expected one of: novice, basic, basicplus, advanced, "
    "expert, soc."
)
raise ValueError(msg)

This fails before returning a partially defined spectrum.

Return the selected spectrum:
```
return spectrum
```
This preserves one annotated return variable across all valid branches.

param level:: One of "novice", "basic", "basicplus", "advanced", "expert", or "soc" (case-insensitive).
type level:: str
param eigenbands:: Band eigenvalues in eV, shape (n_kpoints, n_bands).
type eigenbands:: Float[Array, 'K B']
param surface_orb:: Orbital projection coefficients, shape (n_kpoints, n_bands, n_atoms, 9).
type surface_orb:: Float[Array, 'K B A 9']
param ef:: Fermi energy in eV. Default is 0.
type ef:: Union[float, Float[Array, '']], default: 0.0
param sigma:: Gaussian broadening width in eV. Default is 0.04.
type sigma:: Union[float, Float[Array, '']], default: 0.04
param gamma:: Lorentzian broadening width in eV. Only used by novice and expert levels. Default is 0.1.
type gamma:: Union[float, Float[Array, '']], default: 0.1
param fidelity:: Number of points in the energy axis. Default is 25000.
type fidelity:: int, default: 25000
param temperature:: Electronic temperature in Kelvin. Default is 15.
type temperature:: Union[float, Float[Array, '']], default: 15.0
param photon_energy:: Incident photon energy in eV. Default is 11.
type photon_energy:: Union[float, Float[Array, '']], default: 11.0
param polarization:: Polarization type: "s", "p", "linear", or "unpolarized". Only used by advanced and expert.
type polarization:: str, default: "unpolarized"
param incident_theta:: Polar angle of the incident beam in degrees. Only used by advanced and expert. Default is 45.
type incident_theta:: Union[float, Float[Array, '']], default: 45.0
param incident_phi:: Azimuthal angle of the incident beam in degrees. Only used by advanced and expert. Default is 0.
type incident_phi:: Union[float, Float[Array, '']], default: 0.0
param polarization_angle:: Rotation angle for arbitrary linear polarization in radians. Only used by advanced and expert. Default is 0.
type polarization_angle:: Union[float, Float[Array, '']], default: 0.0
param surface_spin:: Spin projections; required when level='soc'. Default None.
type surface_spin:: Optional[Float[Array, 'K B A 6']], default: None
param ls_scale:: Spin-orbit coupling strength when level='soc'. Default 0.01.
type ls_scale:: Union[float, Float[Array, '']], default: 0.01
returns:: spectrum – Simulated ARPES spectrum containing the intensity map and energy axis.
rtype:: ArpesSpectrum
raises ValueError:: If level is not one of the six recognized levels, or if level='soc' and surface_spin is None.

See also

simulate_novice_expanded: Voigt broadening, uniform weights.
simulate_basic_expanded: Gaussian, heuristic weights.
simulate_basicplus_expanded: Gaussian, Yeh-Lindau weights.
simulate_advanced_expanded: Gaussian, Yeh-Lindau, polarization.
simulate_expert_expanded: Voigt, Yeh-Lindau, polarization, dipole matrix elements.
simulate_soc_expanded: Expert plus spin-orbit (S·k_photon) correction.

diffpes.simul.simulate_expert_expanded(eigenbands: Float[Array, 'K B'], surface_orb: Float[Array, 'K B A 9'], ef: float | Float[Array, ''], sigma: float | Float[Array, ''], gamma: float | Float[Array, ''], fidelity: int, temperature: float | Float[Array, ''], photon_energy: float | Float[Array, ''], polarization: str = 'unpolarized', incident_theta: float | Float[Array, ''] = 45.0, incident_phi: float | Float[Array, ''] = 0.0, polarization_angle: float | Float[Array, ''] = 0.0) → ArpesSpectrum[source]¶

Run expert-level ARPES simulation from plain arrays.

Most physically complete model. Combines Voigt broadening (Gaussian + Lorentzian), Yeh-Lindau photoionization cross-sections, polarization selection rules, and full dipole matrix element weighting. For unpolarized light, the function averages the s-polarization and p-polarization contributions.

See:: TestSimulateExpertExpanded

Implementation Logic¶

Build carrier inputs:

bands, orb_proj = _build_inputs(
    eigenbands=eigenbands,
    surface_orb=surface_orb,
    ef=ef,
)

This validates the raw arrays and binds the band and projection carriers.

Build simulation settings:

params = make_expanded_simulation_params(
    eigenbands=eigenbands,
    fidelity=fidelity,
    sigma=sigma,
    gamma=gamma,
    temperature=temperature,
    photon_energy=photon_energy,
)

This derives the energy window and preserves both Voigt widths.

Build polarization settings:

pol = _build_polarization(
    polarization=polarization,
    incident_theta=incident_theta,
    incident_phi=incident_phi,
    polarization_angle=polarization_angle,
)

This converts degree-valued incidence angles through the shared factory.

Evaluate the expert model:
```
spectrum = simulate_expert(bands, orb_proj, params, pol)
```
This delegates dipole weighting and Voigt broadening to the core model.

param eigenbands:: Band eigenvalues in eV, shape (n_kpoints, n_bands).
type eigenbands:: Float[Array, 'K B']
param surface_orb:: Orbital projection coefficients, shape (n_kpoints, n_bands, n_atoms, 9).
type surface_orb:: Float[Array, 'K B A 9']
param ef:: Fermi energy in eV.
type ef:: Union[float, Float[Array, '']]
param sigma:: Gaussian broadening width in eV.
type sigma:: Union[float, Float[Array, '']]
param gamma:: Lorentzian broadening width in eV.
type gamma:: Union[float, Float[Array, '']]
param fidelity:: Number of points in the energy axis.
type fidelity:: int
param temperature:: Electronic temperature in Kelvin for the Fermi-Dirac distribution.
type temperature:: Union[float, Float[Array, '']]
param photon_energy:: Incident photon energy in eV.
type photon_energy:: Union[float, Float[Array, '']]
param polarization:: Polarization type: "s", "p", "linear", or "unpolarized" (default).
type polarization:: str, default: "unpolarized"
param incident_theta:: Polar angle of the incident beam in degrees. Default 45.
type incident_theta:: Union[float, Float[Array, '']], default: 45.0
param incident_phi:: Azimuthal angle of the incident beam in degrees. Default 0.
type incident_phi:: Union[float, Float[Array, '']], default: 0.0
param polarization_angle:: Rotation angle for arbitrary linear polarization in radians. Default 0.
type polarization_angle:: Union[float, Float[Array, '']], default: 0.0
returns:: spectrum – Simulated ARPES spectrum containing the intensity map and energy axis.
rtype:: ArpesSpectrum

See also

simulate_expert: Low-level implementation accepting PyTrees.

diffpes.simul.simulate_novice_expanded(eigenbands: Float[Array, 'K B'], surface_orb: Float[Array, 'K B A 9'], ef: float | Float[Array, ''], sigma: float | Float[Array, ''], gamma: float | Float[Array, ''], fidelity: int, temperature: float | Float[Array, ''], photon_energy: float | Float[Array, '']) → ArpesSpectrum[source]¶

Run novice-level ARPES simulation from plain arrays.

Simplest physical model: applies Voigt broadening (combined Gaussian + Lorentzian) with uniform orbital weights. All non-s orbitals contribute equally to the photoemission intensity.

See:: TestSimulateNoviceExpanded

Implementation Logic¶

Build carrier inputs:

bands, orb_proj = _build_inputs(
    eigenbands=eigenbands,
    surface_orb=surface_orb,
    ef=ef,
)

This validates the raw arrays and binds the band and projection carriers.

Build simulation settings:

params = make_expanded_simulation_params(
    eigenbands=eigenbands,
    fidelity=fidelity,
    sigma=sigma,
    gamma=gamma,
    temperature=temperature,
    photon_energy=photon_energy,
)

This derives the energy window and preserves both Voigt widths.

Evaluate the novice model:
```
spectrum = simulate_novice(bands, orb_proj, params)
```
This delegates the differentiable numerical work to the carrier API.

param eigenbands:: Band eigenvalues in eV, shape (n_kpoints, n_bands).
type eigenbands:: Float[Array, 'K B']
param surface_orb:: Orbital projection coefficients, shape (n_kpoints, n_bands, n_atoms, 9).
type surface_orb:: Float[Array, 'K B A 9']
param ef:: Fermi energy in eV.
type ef:: Union[float, Float[Array, '']]
param sigma:: Gaussian broadening width in eV.
type sigma:: Union[float, Float[Array, '']]
param gamma:: Lorentzian broadening width in eV.
type gamma:: Union[float, Float[Array, '']]
param fidelity:: Number of points in the energy axis.
type fidelity:: int
param temperature:: Electronic temperature in Kelvin for the Fermi-Dirac distribution.
type temperature:: Union[float, Float[Array, '']]
param photon_energy:: Incident photon energy in eV.
type photon_energy:: Union[float, Float[Array, '']]
returns:: spectrum – Simulated ARPES spectrum containing the intensity map and energy axis.
rtype:: ArpesSpectrum

See also

simulate_novice: Low-level implementation accepting PyTrees.

diffpes.simul.simulate_soc_expanded(eigenbands: Float[Array, 'K B'], surface_orb: Float[Array, 'K B A 9'], surface_spin: Float[Array, 'K B A 6'], ef: float | Float[Array, ''], sigma: float | Float[Array, ''], gamma: float | Float[Array, ''], fidelity: int, temperature: float | Float[Array, ''], photon_energy: float | Float[Array, ''], polarization: str = 'unpolarized', incident_theta: float | Float[Array, ''] = 45.0, incident_phi: float | Float[Array, ''] = 0.0, polarization_angle: float | Float[Array, ''] = 0.0, ls_scale: float | Float[Array, ''] = 0.01) → ArpesSpectrum[source]¶

Run SOC (spin-orbit coupling) ARPES simulation from plain arrays.

Expert model plus spin-dependent intensity correction (S·k_photon). Requires spin projections of shape (n_kpoints, n_bands, n_atoms, 6) (up/down for x, y, z). The function interprets the incident angles in degrees.

See:: TestSimulateSocExpanded

Implementation Logic¶

Build the band carrier:

bands, _projection = _build_inputs(
    eigenbands=eigenbands,
    surface_orb=surface_orb,
    ef=ef,
)

This validates the common raw arrays and preserves the Fermi-energy shift.

Attach spin-resolved projections:

soc_proj = make_spin_orbital_projection(
    projections=jnp.asarray(surface_orb, dtype=jnp.float64),
    spin=surface_spin,
)

This keeps the six spin channels explicit for the SOC correction.

Build simulation and polarization settings:

params = make_expanded_simulation_params(
    eigenbands=eigenbands,
    fidelity=fidelity,
    sigma=sigma,
    gamma=gamma,
    temperature=temperature,
    photon_energy=photon_energy,
)
pol = _build_polarization(
    polarization=polarization,
    incident_theta=incident_theta,
    incident_phi=incident_phi,
    polarization_angle=polarization_angle,
)

These carriers retain the traced physical parameters and the static mode choice.

Evaluate the SOC model:
```
spectrum = simulate_soc(
    bands, soc_proj, params, pol, ls_scale=ls_scale
)
```
This applies the spin-dependent correction only after carrier validation.

param eigenbands:: Band eigenvalues in eV, shape (n_kpoints, n_bands).
type eigenbands:: Float[Array, 'K B']
param surface_orb:: Orbital projection coefficients, shape (n_kpoints, n_bands, n_atoms, 9).
type surface_orb:: Float[Array, 'K B A 9']
param surface_spin:: Spin projections for SOC, with shape (n_kpoints, n_bands, n_atoms, 6). The channels are x up, x down, y up, y down, z up, and z down.
type surface_spin:: Float[Array, 'K B A 6']
param ef:: Fermi energy in eV.
type ef:: Union[float, Float[Array, '']]
param sigma:: Gaussian broadening width in eV.
type sigma:: Union[float, Float[Array, '']]
param gamma:: Lorentzian broadening width in eV.
type gamma:: Union[float, Float[Array, '']]
param fidelity:: Number of points on the energy axis.
type fidelity:: int
param temperature:: Electronic temperature in Kelvin.
type temperature:: Union[float, Float[Array, '']]
param photon_energy:: Incident photon energy in eV.
type photon_energy:: Union[float, Float[Array, '']]
param polarization:: Polarization type, for example "unpolarized" or "LHP". Default is "unpolarized".
type polarization:: str, default: "unpolarized"
param incident_theta:: Polar angle of the incident beam in degrees. Default 45.
type incident_theta:: Union[float, Float[Array, '']], default: 45.0
param incident_phi:: Azimuthal angle of the incident beam in degrees. Default 0.
type incident_phi:: Union[float, Float[Array, '']], default: 0.0
param polarization_angle:: Rotation angle for arbitrary linear polarization in radians. Default 0.
type polarization_angle:: Union[float, Float[Array, '']], default: 0.0
param ls_scale:: Spin-orbit coupling strength for the S·k_photon correction. Default 0.01.
type ls_scale:: Union[float, Float[Array, '']], default: 0.01
returns:: spectrum – Simulated ARPES spectrum with spin-orbit correction; intensity shape (K, E), energy_axis shape (E,).
rtype:: ArpesSpectrum

See also

simulate_soc: Low-level implementation accepting PyTrees.
simulate_expert_expanded: Same pipeline without spin data.

diffpes.simul.simulate_tb_radial(diag_bands: DiagonalizedBands, slater_params: SlaterParams, params: SimulationParams, pol_config: PolarizationConfig, work_function: float | Float[Array, ''] = 4.5, self_energy: SelfEnergyConfig | None = None, r_grid: Float[Array, 'R'] | None = None, dk: float | Float[Array, ''] | None = None) → ArpesSpectrum[source]¶

Run the end-to-end differentiable ARPES forward model.

Computes dipole matrix elements from first principles using radial integrals, Gaunt coefficients, and real spherical harmonics, then produces a simulated ARPES spectrum. The entire pipeline is JAX-traceable and supports jax.grad with respect to all continuous parameters.

This function provides the Chinook-style tight-binding ARPES forward model. The six-level functions in spectrum.py use VASP orbital projections as weights. In contrast, this function computes dipole matrix elements from Slater-type radial wavefunctions and tight-binding eigenvectors. JAX can therefore differentiate the complete simulation with respect to:

slater_params.zeta (Slater exponents controlling radial extent)
diag_bands.eigenvalues and diag_bands.eigenvectors
params.sigma, params.gamma, params.temperature
pol_config.theta, pol_config.phi
work_function
self_energy.coefficients (if provided)

See:: TestSimulateTBRadial

Implementation Logic¶

The simulation proceeds in five stages:

Create the numerical grids:
```
energy_axis: Float[Array, " E"] = jnp.linspace(
    params.energy_min, params.energy_max, params.fidelity
)
```
The energy axis fixes the output sampling. The optional radial grid controls the integration sampling in Bohr.
Evaluate the radial wavefunctions:
```
radial_scan: tuple[None, Float[Array, "O R"]] = jax.lax.scan(
    _scan_radial,
    None,
    orbital_indices,
)
```
The scan traverses the differentiable Slater parameters. jax.lax.switch selects each static principal quantum number.
Compute the coherent band intensities:
```
band_intensity: Float[Array, "K B"] = (i_s + i_p) / 2.0
```
Nested jax.vmap calls evaluate all k-points and bands. Each evaluation forms the coherent orbital sum before the squared modulus.

The coherent matrix element is:

\[M_{k,b} = \sum_o c_{k,b,o} \cdot M_o(k, \hat{\epsilon})\]

Here, \(c_{k,b,o}\) is an eigenvector coefficient and \(M_o\) is one orbital matrix element. Unpolarized light averages the s-polarized and p-polarized intensities.
Broaden the band contributions:
```
intensity: Float[Array, "K E"] = jax.vmap(_single_kpoint)(
    diag_bands.eigenvalues,
    band_intensity,
)
```
Each band contributes a Fermi-weighted Voigt profile. The optional self-energy makes the Lorentzian width depend on energy.
Apply the momentum response:
```
intensity = apply_momentum_broadening(
    intensity, k_distances, dk
)
```
The Gaussian response represents finite angular acceptance. The operation remains differentiable with respect to the spectrum and the response width.

param diag_bands:: Diagonalized electronic structure with eigenvalues of shape (K, B) and eigenvectors of shape (K, B, O). O is the orbital count. The carrier also contains k-points and the Fermi energy.
type diag_bands:: DiagonalizedBands
param slater_params:: Slater radial wavefunction parameters including Slater exponents zeta of shape (O,), expansion coefficients, and the orbital_basis specifying (n, l, m) quantum numbers for each orbital.
type slater_params:: SlaterParams
param params:: Simulation parameters including energy_min, energy_max, fidelity, sigma (Gaussian width), gamma (Lorentzian width), temperature, and photon_energy.
type params:: SimulationParams
param pol_config:: Photon polarization configuration specifying polarization_type, incidence angles theta and phi, and polarization_angle.
type pol_config:: PolarizationConfig
param work_function:: Material work function in eV. Default is 4.5.
type work_function:: Union[float, Float[Array, '']], default: 4.5
param self_energy:: Energy-dependent self-energy model for the Lorentzian broadening. If None, the function uses params.gamma for all energies.
type self_energy:: Optional[SelfEnergyConfig], default: None
param r_grid:: Radial integration grid in Bohr. Default is 10000 points linearly spaced on [1e-6, 50.0].
type r_grid:: Optional[Float[Array, 'R']], default: None
param dk:: Momentum broadening width in inverse Angstroms. If None, the function does not apply k-space convolution.
type dk:: Union[float, Float[Array, ''], None], default: None
returns:: spectrum – Simulated ARPES spectrum with intensity of shape (K, E) and energy_axis of shape (E,).
rtype:: ArpesSpectrum

Notes

The inner function _single_k_band uses safe_norm() and safe_divide() to select a zero direction and zero gradient at the Gamma point. It preserves the fractional crystal-momentum direction when it forms the photoelectron momentum. Both radial construction and coherent orbital summation use jax.lax.scan. Static quantum numbers select specialized branches with jax.lax.switch.

See also

simulate_expert: Projection-based simulation (uses VASP weights rather than ab-initio dipole matrix elements).
dipole_matrix_element_single: Core dipole integral computation.
slater_radial: Slater-type radial wavefunction evaluation.
evaluate_self_energy: Energy-dependent broadening model.
apply_momentum_broadening: k-space Gaussian convolution.

diffpes.simul.detector_angles_to_kpar(tx: Float[Array, '...'], ty: Float[Array, '...'], kinetic_energy_ev: Float[Array, '...'], slit: str) → Float[Array, '... 2'][source]¶

Convert detector angles to parallel momentum.

The function rotates the positive z direction with the Plan 03 detector convention. It broadcasts all traced inputs over their leading axes.

See:

TestDetectorAnglesToKpar

Parameters:

tx (']) – First detector angles in radians.
ty (']) – Second detector angles in radians.
kinetic_energy_ev (']) – Photoelectron kinetic energies in eV.
slit (str) – Static slit orientation, "H" or "V". A change causes retracing.

Returns:

k_parallel – Cartesian parallel momenta (kx, ky) in 1/Angstrom.

Return type:

2']

Raises:

ValueError – If slit is not "H" or "V".

Notes

The horizontal slit uses Rx(ty) @ Ry(tx). The vertical slit uses Rx(tx) @ Ry(ty). These active rotations act on the positive z vector.

diffpes.simul.emission_angles(k_cart_inv_ang: Float[Array, '... 3']) → tuple[Float[Array, '...'], Float[Array, '...']][source]¶

Convert Cartesian momentum to emission angles.

The function returns the polar angle from positive z and the azimuth from positive x. It selects zero azimuth at normal emission.

See:

TestEmissionAngles

Parameters:

k_cart_inv_ang ( 3']) – Cartesian momentum vectors in 1/Angstrom.

Return type:

tuple['], ']]

Returns:

theta (Float[Array, " ..."]) – Polar emission angles in radians.
phi (Float[Array, " ..."]) – Azimuthal emission angles in radians.

Notes

The polar angle uses arctan2(norm([kx, ky]), kz). The azimuth uses arctan2(ky, kx). Safe primitives give zero coordinate gradients at their undefined origins.

diffpes.simul.final_state_k_inv_ang(kinetic_energy_ev: Float[Array, '...']) → Float[Array, '...'][source]¶

Convert kinetic energy to final-state momentum magnitude.

The function applies the free-electron dispersion. A second floor guard keeps direct calls finite outside the physical domain.

See:: TestFinalStateKInvAng
Parameters:: kinetic_energy_ev (']) – Photoelectron kinetic energies in eV.
Returns:: momentum_magnitudes – Final-state momentum magnitudes in 1/Angstrom.
Return type:: ']

Notes

The function computes K_PREFACTOR_INV_ANG_SQRT_EV * sqrt(E_kin). Inputs at or below EKIN_FLOOR_EV have a zero selected gradient.

diffpes.simul.kinetic_energy_ev(photon_energy_ev: float | Float[Array, ''], work_function_ev: float | Float[Array, ''], binding_energy_ev: Float[Array, '...']) → Float[Array, '...'][source]¶

Compute the floored photoelectron kinetic energy.

The function applies energy conservation in the three-step photoemission model [3]. A physical floor defines the low-energy validity boundary.

See:

TestKineticEnergyEv

Parameters:

photon_energy_ev (Union[float, Float[Array, '']]) – Photon energy in eV.
work_function_ev (Union[float, Float[Array, '']]) – Work function in eV.
binding_energy_ev (']) – Binding energies in eV. The function accepts either sign convention.

Returns:

kinetic_energies – Kinetic energies in eV with a lower bound of EKIN_FLOOR_EV.

Return type:

']

Notes

The function computes \(h\nu-W-|E_b|\). Values at or below the floor have a zero selected gradient. Values above the floor keep exact gradients.

References

diffpes.simul.kpar_to_detector_angles(k_par_inv_ang: Float[Array, '... 2'], kinetic_energy_ev: Float[Array, '...'], slit: str) → tuple[Float[Array, '...'], Float[Array, '...']][source]¶

Convert parallel momentum to detector angles.

The function gives the exact inverse detector map on the physical domain. This domain requires norm(k_parallel) < k_f.

See:

TestKparToDetectorAngles

Parameters:

k_par_inv_ang ( 2']) – Cartesian parallel momenta (kx, ky) in 1/Angstrom.
kinetic_energy_ev (']) – Photoelectron kinetic energies in eV.
slit (str) – Static slit orientation, "H" or "V". A change causes retracing.

Return type:

tuple['], ']]

Returns:

tx (Float[Array, " ..."]) – First detector angles in radians.
ty (Float[Array, " ..."]) – Second detector angles in radians.

Raises:

ValueError – If slit is not "H" or "V".

Notes

The inverse uses the positive detector-normal branch. Safe square roots select finite boundary values outside the open physical domain.

diffpes.simul.kz_from_inner_potential(photon_energy_ev: float | Float[Array, ''], work_function_ev: float | Float[Array, ''], inner_potential_ev: float | Float[Array, ''], k_par_inv_ang: Float[Array, '...']) → tuple[Complex[Array, '...'], Bool[Array, '...']][source]¶

Compute complex out-of-plane momentum from the inner potential.

The function implements the free-electron final-state approximation [4]. Its principal complex root retains evanescent channels.

See:

TestKzFromInnerPotential

Parameters:

photon_energy_ev (Union[float, Float[Array, '']]) – Photon energy in eV.
work_function_ev (Union[float, Float[Array, '']]) – Work function in eV.
inner_potential_ev (Union[float, Float[Array, '']]) – Inner potential in eV.
k_par_inv_ang (']) – Parallel momentum magnitudes in 1/Angstrom.

Return type:

tuple['], ']]

Returns:

kz_values (Complex[Array, " ..."]) – Principal out-of-plane momenta in 1/Angstrom.
propagating (Bool[Array, " ..."]) – Mask that identifies positive real radicands.

Notes

The radicand equals \((2m_e/\hbar^2)(h\nu-W+V_0)-k_\parallel^2\) above the energy floor. Negative radicands give positive imaginary roots. The branch point has no assigned derivative.

For a propagating channel, \(\partial k_z/\partial V_0=(2\,\hbar^2/2m_e)^{-1}/k_z\).

References

diffpes.simul.compute_oam(projections: Float[Array, 'K B A 9']) → Float[Array, 'K B A 3'][source]¶

Compute orbital angular momentum z-component.

Evaluates the expectation value of the z-component of orbital angular momentum from orbital-resolved projections:

OAM_z = sum_m  m * |c_m|^2

Here, m is the magnetic quantum number, and c_m is the orbital projection coefficient. The function computes the p-orbital and d-orbital contributions separately. It then adds them.

See:: TestComputeOam

Implementation Logic¶

Extract the p-orbital projections:
```
p_proj: Float[Array, "K B A 3"] = projections[
    ..., P_ORBITAL_SLICE
]
```
Selects the three p-orbital coefficients [py, pz, px] corresponding to magnetic quantum numbers m = {+1, 0, -1}.
Compute the p-orbital OAM:
```
p_oam: Float[Array, "K B A"] = jnp.sum(
    M_P * p_proj**2, axis=-1
)
```
Weights each squared projection by its magnetic quantum number m_p = [+1, 0, -1] and sums over the p-orbital subspace.
Extract the d-orbital projections:
```
d_proj: Float[Array, "K B A 5"] = projections[
    ..., D_ORBITAL_SLICE
]
```
Selects the five d-orbital coefficients [dxy, dyz, dz2, dxz, dx2-y2] corresponding to m = {-2, -1, 0, +1, +2}.
Compute the d-orbital OAM:
```
d_oam: Float[Array, "K B A"] = jnp.sum(
    M_D * d_proj**2, axis=-1
)
```
Weights each squared projection by its magnetic quantum number m_d = [-2, -1, 0, +1, +2] and sums over the d-orbital subspace.

Stack the p, d, and total results:

total_oam: Float[Array, "K B A"] = p_oam + d_oam
oam: Float[Array, "K B A 3"] = jnp.stack(
    [p_oam, d_oam, total_oam], axis=-1
)

Returns all three components so that downstream analysis can inspect orbital-resolved or total OAM.

param projections:: Orbital projections with 9 orbitals per atom.
type projections:: Float[Array, 'K B A 9']
returns:: oam – OAM_z for [p-contribution, d-contribution, total].
rtype:: Float[Array, 'K B A 3']

Notes

Orbital indices follow VASP ordering: [s(0), py(1), pz(2), px(3), dxy(4), dyz(5), dz2(6), dxz(7), dx2-y2(8)]. The s-orbital (index 0) has m = 0 and does not contribute to the OAM.

diffpes.simul.build_efield(config: PolarizationConfig) → Complex[Array, '3'][source]¶

Compute electric field vector from polarization config.

Constructs the complex electric field polarization vector for the specified photon geometry and polarization type.

See:: TestBuildEfield

Implementation Logic¶

Build s- and p-polarization basis:
```
e_s, e_p = build_polarization_vectors(
    config.theta, config.phi
)
```
This computes the real orthonormal basis from the incidence angles.
Select the static polarization branch:
```
index: int = pol_index_map.get(pol_type, 5)
```
The string configuration selects one branch before JAX execution.
Evaluate the selected branch with JAX:
```
efield: Complex[Array, " 3"] = jax.lax.switch(
    index, branches, operand
)
```
The JAX switch preserves differentiation through the selected field.

param config:: Polarization geometry specification.
type config:: PolarizationConfig
returns:: efield – Complex electric field polarization vector.
rtype:: Complex[Array, '3']

diffpes.simul.build_polarization_vectors(theta: float | Float[Array, ''], phi: float | Float[Array, '']) → tuple[Float[Array, '3'], Float[Array, '3']][source]¶

Construct s- and p-polarization basis vectors.

The function constructs an orthonormal pair of polarization vectors from the photon incidence angles. The s-polarization is perpendicular to the incidence plane. The p-polarization is in the incidence plane and perpendicular to the wavevector.

See:: TestBuildPolarizationVectors

Implementation Logic¶

Construct the s-polarization vector:
```
e_s = [sin(phi), -cos(phi), 0]
```
This closed form is perpendicular to the incidence plane. It is the normalized k cross z convention continued to normal incidence.
Construct the p-polarization vector:
```
e_p = [-cos(theta) cos(phi),
       -cos(theta) sin(phi),
        sin(theta)]
```
This vector equals e_s cross k. It is perpendicular to the photon direction and completes the orthonormal transverse basis.

param theta:

Incident angle from surface normal in radians.

type theta:

Union[float, Float[Array, '']]

param phi:

In-plane azimuthal angle in radians.

type phi:

Union[float, Float[Array, '']]

rtype:

tuple[Float[Array, '3'], Float[Array, '3']]

returns:

e_s (Float[Array, " 3"]) – s-polarization unit vector (perpendicular to incidence plane).
e_p (Float[Array, " 3"]) – p-polarization unit vector (in incidence plane, perpendicular to photon wavevector).

Notes

The direct trigonometric form has no artificial collinearity threshold. At normal incidence, phi fixes the otherwise free transverse-frame gauge. The basis is smooth in both input angles for that gauge choice.

diffpes.simul.detector_rotation(tx: float | Float[Array, ''], ty: float | Float[Array, ''], slit: str) → Float[Array, '3 3'][source]¶

Build the detector-frame rotation.

The function composes the two analyzer-angle rotations in the order set by the static slit orientation.

See:: TestDetectorRotation

Implementation Logic¶

Build the Cartesian axis rotations:
```
rotation_x = rodrigues_rotation(x_axis, ty)
```
Rodrigues matrices retain derivatives with respect to both angles.
Compose the slit convention:
```
horizontal = rotation_x_ty @ rotation_y_tx
vertical = rotation_x_tx @ rotation_y_ty
```
Horizontal and vertical slits use the registered composition orders.

param tx:: First detector angle in radians.
type tx:: Union[float, Float[Array, '']]
param ty:: Second detector angle in radians.
type ty:: Union[float, Float[Array, '']]
param slit:: Slit orientation (static). Use "H" or "V".
type slit:: str
returns:: rotation – Proper rotation from the reference detector frame.
rtype:: Float[Array, '3 3']
raises ValueError:: If slit is not "H" or "V".

Notes

"H" uses R_x(ty) R_y(tx). "V" uses R_x(tx) R_y(ty). The same matrix rotates the emitted direction, polarization, and spin axis.

Chinook uses opposite raw signs for its horizontal Ty momentum and polarization coordinates. The declared source-coordinate mappings give one active DiffPES frame.

diffpes.simul.dipole_matrix_elements(efield: Complex[Array, '3']) → Float[Array, '9'][source]¶

Compute dipole matrix elements for all 9 orbitals.

Evaluates the squared modulus of the dipole transition matrix element for each orbital:

M_i = |e . d_i|^2

where e is the electric field polarization vector and d_i is the normalized direction vector of orbital i.

See:: TestDipoleMatrixElements

Implementation Logic¶

Dot product of E-field with each orbital direction:
```
dots = ORBITAL_DIRS_NORMALIZED @ efield
```
- Computes the inner product of the complex electric field vector with each of the 9 normalized orbital direction vectors (shape [9, 3] @ [3] -> [9]). The result is a complex-valued array of length 9.
- The s-orbital direction vector is [0, 0, 0] (zero vector), so its dot product is always zero regardless of the E-field, reflecting the isotropic (zero directionality) character of the s-orbital.
Square modulus |e . d|^2:
```
matrix_elements = |dots|^2
```
Takes the absolute value squared of each complex dot product. For real-valued E-fields this reduces to the squared real dot product. For circular polarization (complex E-field) this correctly accounts for the phase.

param efield:: Electric field polarization vector.
type efield:: Complex[Array, '3']
returns:: matrix_elements – |e dot d_orbital|^2 for each orbital.
rtype:: Float[Array, '9']

Notes

The s-orbital has a zero direction vector and therefore always produces a zero dipole matrix element with any polarization.

diffpes.simul.photon_wavevector(theta: float | Float[Array, ''], phi: float | Float[Array, '']) → Float[Array, '3'][source]¶

Build the unit photon wavevector from incidence angles.

The function constructs the unit photon propagation vector from spherical coordinates. Theta starts at the surface normal, and phi is the azimuthal angle. Spin-orbit ARPES simulations use the vector in the S·k_photon correction for circular dichroism.

See:: TestPhotonWavevector

Notes

Form the Cartesian spherical-coordinate vector, normalize it with its Euclidean norm, bind the unit result to k_hat, and return it.

Parameters:

theta (Union[float, Float[Array, '']]) – Incident angle from surface normal in radians.
phi (Union[float, Float[Array, '']]) – In-plane azimuthal angle in radians.

Returns:

k_photon – Unit wavevector in Cartesian coordinates.

Return type:

Float[Array, '3']

See also

build_polarization_vectors: Build the same k for the s-polarization and p-polarization basis. Use this function only for the propagation direction.

diffpes.simul.polarization_from_angles(incidence_theta: float | Float[Array, ''], incidence_phi: float | Float[Array, ''], kind: str, polarization_angle: float | Float[Array, ''] = 0.0) → Complex[Array, '3'][source]¶

Construct polarization from incidence angles.

The function returns an explicit complex Cartesian vector for a standard polarization state. The incidence angles use the laboratory frame.

See:: TestPolarizationFromAngles

Implementation Logic¶

Construct the transverse basis:
```
e_s, e_p = build_polarization_vectors(theta, phi)
```
The basis is orthonormal and perpendicular to the photon direction.
Select the requested state:
```
polarization = coefficients[0] * e_s + coefficients[1] * e_p
```
The static selector chooses s, p, circular, or linear coefficients.
Return the complex vector:
```
return polarization
```
The result retains phase information for later coherent contraction.

param incidence_theta:: Photon angle from the surface normal in radians.
type incidence_theta:: Union[float, Float[Array, '']]
param incidence_phi:: Photon azimuth in radians.
type incidence_phi:: Union[float, Float[Array, '']]
param kind:: Polarization kind (static). Use "s", "p", "c+", "c-", or "linear".
type kind:: str
param polarization_angle:: Linear-basis angle in radians. Default is 0.0.
type polarization_angle:: Union[float, Float[Array, '']], default: 0.0
returns:: polarization – Unit polarization vector in the laboratory frame.
rtype:: Complex[Array, '3']
raises ValueError:: If kind is not a supported polarization kind.

Notes

The kind value is static and selects a Python branch before tracing. JAX differentiates the result with respect to all angle arguments.

See also

build_polarization_vectors: Construct the transverse real basis.
polarization_to_spherical: Convert the result to spherical components.

diffpes.simul.polarization_to_spherical(polarization: Complex[Array, '3']) → Complex[Array, '3'][source]¶

Convert Cartesian polarization to spherical components.

The result uses component order (q=-1, q=0, q=+1) and the Condon-Shortley phase convention.

See:: TestPolarizationToSpherical

Implementation Logic¶

Read Cartesian components:
```
epsilon_x, epsilon_y, epsilon_z = polarization
```
The input remains complex so the operation preserves optical phase.
Apply the spherical-basis transform:
```
epsilon_minus = (epsilon_x - 1j * epsilon_y) / sqrt(2)
```
The transform follows the registered Condon-Shortley convention.
Stack the ordered result:
```
spherical = stack((epsilon_minus, epsilon_z, epsilon_plus))
```
This order matches the transitions q = (-1, 0, +1).

param polarization:: Cartesian polarization vector.
type polarization:: Complex[Array, '3']
returns:: spherical – Spherical components in (q=-1, q=0, q=+1) order.
rtype:: Complex[Array, '3']

Notes

The transform is complex-linear. It preserves the squared vector norm and supports JVP, VJP, and complex-step checks without a conjugation.

diffpes.simul.rotate_frame_vectors(vector: Float[Array, '3'], tx: Float[Array, 'n_tx'], ty: Float[Array, 'n_ty'], slit: str) → Float[Array, 'n_tx n_ty 3'][source]¶

Rotate a real vector across a detector-angle grid.

The function applies each detector-frame rotation to one real laboratory vector. It preserves both detector axes in the output.

See:: TestRotateFrameVectors

Implementation Logic¶

Map over both angle axes:
```
rotations = vmap(vmap(detector_rotation))(tx, ty)
```
Nested mapping builds one rotation for every detector coordinate.
Apply each rotation:
```
rotated = rotations @ vector
```
Matrix multiplication preserves the vector norm.

param vector:: Real vector in the reference laboratory frame.
type vector:: Float[Array, '3']
param tx:: First detector-angle axis in radians.
type tx:: Float[Array, 'n_tx']
param ty:: Second detector-angle axis in radians.
type ty:: Float[Array, 'n_ty']
param slit:: Slit orientation (static). Use "H" or "V".
type slit:: str
returns:: rotated – Rotated vector at each detector coordinate.
rtype:: Float[Array, 'n_tx n_ty 3']

Notes

The output has fixed shape for fixed angle-axis lengths. JAX can compile and differentiate the two mapped angle axes without Python data loops.

diffpes.simul.rotate_polarization_grid(polarization: Complex[Array, '3'], tx: Float[Array, 'n_tx'], ty: Float[Array, 'n_ty'], slit: str) → Complex[Array, 'n_tx n_ty 3'][source]¶

Rotate polarization across a detector-angle grid.

The function applies the shared detector frame to a complex polarization vector without reducing its phase or amplitude.

See:: TestRotatePolarizationGrid

Implementation Logic¶

Map over both angle axes:
```
rotated = vmap(vmap(rotate_one))(tx, ty)
```
Nested mapping applies the same frame convention at every coordinate.
Return the complex vectors:
```
return rotated
```
The result retains coherent complex components for later models.

param polarization:: Complex polarization in the reference laboratory frame.
type polarization:: Complex[Array, '3']
param tx:: First detector-angle axis in radians.
type tx:: Float[Array, 'n_tx']
param ty:: Second detector-angle axis in radians.
type ty:: Float[Array, 'n_ty']
param slit:: Slit orientation (static). Use "H" or "V".
type slit:: str
returns:: rotated – Rotated polarization at each detector coordinate.
rtype:: Complex[Array, 'n_tx n_ty 3']

Notes

Real rotation matrices act on the complex vector components. Therefore, the map is complex-linear in polarization and differentiable in both detector-angle axes.

diffpes.simul.apply_momentum_broadening(intensity: Float[Array, 'K E'], k_distances: Float[Array, 'K'], dk: float | Float[Array, '']) → Float[Array, 'K E'][source]¶

Convolve I(k, E) with a Gaussian in k-space.

Simulates the finite angular (momentum) resolution of an ARPES analyser by applying a Gaussian convolution along the k-axis. This smears sharp spectral features in momentum, mimicking the experimental point-spread function in the angular direction.

In an ARPES experiment, the electron analyzer has a finite angular acceptance. The photon beam also has a finite spot size. These properties produce a momentum-space resolution function. A Gaussian with width dk approximates this function. The implementation multiplies the intensity by a normalized Gaussian kernel. JAX can differentiate this operation.

See:: TestApplyMomentumBroadening

Implementation Logic¶

Guard the Gaussian width:
```
safe_dk: Float[Array, ""] = jnp.maximum(dk_arr, EPS)
```
This guard prevents division by zero. A zero width produces an effectively diagonal kernel.
Build the Gaussian kernel:
```
kernel: Float[Array, " K K"] = jnp.exp(
    -0.5 * scaled_distances**2
)
```
Each matrix element measures the Gaussian overlap between two k-points. The cumulative distances support nonuniform spacing.
Normalize each row:
```
kernel = safe_divide(kernel, row_sum)
```
The normalization preserves the spectral weight at each output k-point. The safe division also protects an empty numerical row.
Apply the kernel:
```
broadened: Float[Array, "K E"] = kernel @ intensity
```
The matrix product applies the momentum response independently to every energy column.

param intensity:: ARPES intensity map with shape (n_kpoints, n_energies).
type intensity:: Float[Array, 'K E']
param k_distances:: Cumulative k-path distances of shape (n_kpoints,), in inverse Angstroms. Must be monotonically increasing.
type k_distances:: Float[Array, 'K']
param dk:: Gaussian broadening standard deviation in inverse Angstroms. Typical experimental values range from 0.01 to 0.05.
type dk:: Union[float, Float[Array, '']]
returns:: broadened – Momentum-broadened intensity map, same shape as intensity.
rtype:: Float[Array, 'K E']

Notes

The kernel matrix is dense (K, K) and fully traced by JAX, so this function supports jax.grad with respect to dk and intensity. For very large numbers of k-points, memory usage scales as O(K^2).

diffpes.simul.evaluate_self_energy(energy: Float[Array, '...'], config: SelfEnergyConfig) → Float[Array, '...'][source]¶

Evaluate the imaginary self-energy \(\Gamma(E)\).

The function computes an energy-dependent Lorentzian broadening width from the specified self-energy model. The result replaces params.gamma in the Voigt profile when the lifetime changes with energy. This behavior can occur near the Fermi level or in correlated materials.

The imaginary part of the electron self-energy determines the quasiparticle lifetime and thus the Lorentzian component of the spectral linewidth. The function supports three parametric models:

constant: A single scalar broadening applied uniformly at all energies. Equivalent to using params.gamma directly.
polynomial: \(\Gamma(E) = \sum_n c_n E^n\). Store coefficients in descending degree order for jnp.polyval. This model can represent the quadratic Fermi-liquid self-energy \(\Gamma \propto (E - E_F)^2\).
tabulated: Piecewise-linear interpolation of user-supplied \((\varepsilon_i, \Gamma_i)\) pairs via jnp.interp. Suitable for self-energies obtained from many-body calculations such as GW or DMFT.

See:: TestEvaluateSelfEnergy

Implementation Logic¶

Read the static mode:
```
mode = config.mode
```
The Python string selects one computation while JAX traces its arrays.

Compute the selected self-energy:

result = jnp.broadcast_to(config.coefficients[0], energy.shape)
result = jnp.polyval(config.coefficients, energy)
result = jnp.interp(
    energy, config.energy_nodes, config.coefficients
)

Each branch returns a width array with the same shape as energy.

Reject an unknown mode:
```
raise ValueError(msg)
```
This static error prevents an unsupported model from entering a trace.

param energy:

Energy values in eV at which to evaluate the self-energy. Any shape. The output has the same shape.

type energy:

']

param config:

Self-energy model specification containing:

mode : str – "constant", "polynomial", or "tabulated".
coefficients : array – Model parameters (scalar for constant, polynomial coefficients for polynomial, gamma values for tabulated).
energy_nodes : array (only for "tabulated") – Energy grid for the tabulated gamma values.

type config:

SelfEnergyConfig

returns:

result – Energy-dependent Lorentzian broadening in eV, same shape as energy.

rtype:

']

raises ValueError:

If config.mode is not one of "constant", "polynomial", or "tabulated".

Notes

All three modes are fully JAX-differentiable with respect to config.coefficients. The "constant" and "polynomial" modes are also differentiable with respect to energy. The "tabulated" mode uses jnp.interp which has limited gradient support (piecewise-constant gradients).

diffpes.simul.simulate_advanced(bands: BandStructure, orb_proj: OrbitalProjection, params: SimulationParams, pol_config: PolarizationConfig) → ArpesSpectrum[source]¶

Simulate ARPES with Gaussian broadening and polarization rules.

The function adds light-polarization dependence to simulate_basicplus through dipole matrix elements. The factor |E . d_orbital|^2 weights each orbital channel. E is the electric field, and d_orbital is the dipole selection vector. The function supports polarized and unpolarized light.

See:: TestSimulateAdvanced

Implementation Logic¶

Build the energy axis:

energy_axis: Float[Array, " E"] = jnp.linspace(
    params.energy_min,
    params.energy_max,
    params.fidelity,
)

This defines the common energy grid for every k-point.

Vectorize the polarized band accumulation:

intensity: Float[Array, "K E"] = jax.vmap(_single_kpoint)(
    bands.eigenvalues, band_intensity
)

JAX maps the differentiable polarization model across all k-points.

Build the spectrum carrier:

spectrum: ArpesSpectrum = make_arpes_spectrum(
    intensity=intensity,
    energy_axis=energy_axis,
)

The factory validates and stores the computed arrays.

param bands:: Electronic band structure containing eigenvalues of shape (K, B) and the Fermi energy.
type bands:: BandStructure
param orb_proj:: Orbital projections with shape (K, B, A, 9). A is the atom count, and 9 is the orbital channel count. The channels follow the VASP order.
type orb_proj:: OrbitalProjection
param params:: Simulation parameters including sigma, photon_energy, temperature, energy_min, energy_max, and fidelity.
type params:: SimulationParams
param pol_config:: Light polarization configuration specifying polarization_type ("unpolarized", "linear", etc.), incidence angles theta and phi, and any additional polarization parameters.
type pol_config:: PolarizationConfig
returns:: spectrum – Simulated ARPES intensity map with intensity of shape (K, E) and energy_axis of shape (E,).
rtype:: ArpesSpectrum

Notes

Uses Gaussian broadening with Yeh-Lindau cross-sections and polarization-dependent orbital selection via |E . d_orbital|^2 weighting. For unpolarized light, the function averages the s-polarization and p-polarization intensities. This average is exact for an incoherent superposition of orthogonal polarization states.

diffpes.simul.simulate_basic(bands: BandStructure, orb_proj: OrbitalProjection, params: SimulationParams) → ArpesSpectrum[source]¶

Simulate ARPES spectrum with Gaussian broadening and heuristic weights.

Intermediate simulation that replaces the Voigt profile with a pure Gaussian and introduces energy-dependent heuristic orbital weights. Below about 50 eV, the heuristic weights enhance p-orbital contributions. Above this energy, they enhance d-orbital contributions. This method approximates photoionization cross-sections without tabulated atomic data.

See:: TestSimulateBasic

Implementation Logic¶

Build the energy axis:

energy_axis: Float[Array, " E"] = jnp.linspace(
    params.energy_min,
    params.energy_max,
    params.fidelity,
)

This defines the common energy grid for every k-point.

Vectorize the band accumulation:

intensity: Float[Array, "K E"] = jax.vmap(_single_kpoint)(
    bands.eigenvalues, weights
)

JAX maps the differentiable Gaussian calculation across all k-points.

Build the spectrum carrier:

spectrum: ArpesSpectrum = make_arpes_spectrum(
    intensity=intensity,
    energy_axis=energy_axis,
)

The factory validates and stores the computed arrays.

param bands:: Electronic band structure containing eigenvalues of shape (K, B) and the Fermi energy.
type bands:: BandStructure
param orb_proj:: Orbital projections with shape (K, B, A, 9). A is the atom count, and 9 is the orbital channel count. The channels follow the VASP order.
type orb_proj:: OrbitalProjection
param params:: Simulation parameters including sigma, photon_energy, temperature, energy_min, energy_max, and fidelity.
type params:: SimulationParams
returns:: spectrum – Simulated ARPES intensity map with intensity of shape (K, E) and energy_axis of shape (E,).
rtype:: ArpesSpectrum

Notes

Uses Gaussian broadening with energy-dependent heuristic orbital weights (p-enhanced below 50 eV, d-enhanced above). This level is useful without Yeh-Lindau cross-section tables. Use it when the simulation needs some orbital selectivity.

diffpes.simul.simulate_basicplus(bands: BandStructure, orb_proj: OrbitalProjection, params: SimulationParams) → ArpesSpectrum[source]¶

Simulate ARPES with Gaussian broadening and Yeh-Lindau cross-sections.

The function replaces the heuristic weights with interpolated Yeh-Lindau photoionization cross-sections. It applies the applicable cross-section to each orbital projection, including s. It then sums the non-s channels. This order preserves the cross-section magnitudes before orbital selection.

See:: TestSimulateBasicplus

Implementation Logic¶

Build the energy axis:

energy_axis: Float[Array, " E"] = jnp.linspace(
    params.energy_min,
    params.energy_max,
    params.fidelity,
)

This defines the common energy grid for every k-point.

Vectorize the band accumulation:

intensity: Float[Array, "K E"] = jax.vmap(_single_kpoint)(
    bands.eigenvalues, weights
)

JAX maps the weighted Gaussian calculation across all k-points.

Build the spectrum carrier:

spectrum: ArpesSpectrum = make_arpes_spectrum(
    intensity=intensity,
    energy_axis=energy_axis,
)

The factory validates and stores the computed arrays.

param bands:: Electronic band structure containing eigenvalues of shape (K, B) and the Fermi energy.
type bands:: BandStructure
param orb_proj:: Orbital projections with shape (K, B, A, 9). A is the atom count, and 9 is the orbital channel count. The channels follow the VASP order.
type orb_proj:: OrbitalProjection
param params:: Simulation parameters including sigma, photon_energy, temperature, energy_min, energy_max, and fidelity.
type params:: SimulationParams
returns:: spectrum – Simulated ARPES intensity map with intensity of shape (K, E) and energy_axis of shape (E,).
rtype:: ArpesSpectrum

Notes

Uses Gaussian broadening with interpolated Yeh-Lindau photoionization cross-section weights per orbital type. The the function computes the cross-sections from tabulated atomic data. It interpolates them to the specified photon energy.

diffpes.simul.simulate_expert(bands: BandStructure, orb_proj: OrbitalProjection, params: SimulationParams, pol_config: PolarizationConfig) → ArpesSpectrum[source]¶

Simulate ARPES with Voigt broadening and dipole matrix elements.

The most physically complete simulation model. Combines Voigt broadening (capturing both instrumental Gaussian and lifetime Lorentzian contributions via sigma and gamma), Yeh-Lindau photoionization cross-sections, and full polarization-dependent dipole matrix element weighting. Use this level for quantitative comparisons with experimental spectra.

See:: TestSimulateExpert

Implementation Logic¶

Build the energy axis:

energy_axis: Float[Array, " E"] = jnp.linspace(
    params.energy_min,
    params.energy_max,
    params.fidelity,
)

This defines the common energy grid for every k-point.

Vectorize the polarized band accumulation:

intensity: Float[Array, "K E"] = jax.vmap(_single_kpoint)(
    bands.eigenvalues, band_intensity
)

JAX maps the differentiable Voigt model across all k-points.

Build the spectrum carrier:

spectrum: ArpesSpectrum = make_arpes_spectrum(
    intensity=intensity,
    energy_axis=energy_axis,
)

The factory validates and stores the computed arrays.

param bands:: Electronic band structure containing eigenvalues of shape (K, B) and the Fermi energy.
type bands:: BandStructure
param orb_proj:: Orbital projections with shape (K, B, A, 9). A is the atom count, and 9 is the orbital channel count. The channels follow the VASP order.
type orb_proj:: OrbitalProjection
param params:: Simulation parameters including sigma, gamma, photon_energy, temperature, energy_min, energy_max, and fidelity.
type params:: SimulationParams
param pol_config:: Light polarization configuration specifying polarization_type ("unpolarized", "linear", etc.), incidence angles theta and phi, and any additional polarization parameters.
type pol_config:: PolarizationConfig
returns:: spectrum – Simulated ARPES intensity map with intensity of shape (K, E) and energy_axis of shape (E,).
rtype:: ArpesSpectrum

Notes

Uses Voigt broadening with Yeh-Lindau cross-sections, polarization selection rules, and dipole matrix element weighting. This is the most physically complete model. For unpolarized light, averages s and p contributions. The Voigt profile supports accurate line shape fits when the analysis must separate instrument and intrinsic broadening.

diffpes.simul.simulate_novice(bands: BandStructure, orb_proj: OrbitalProjection, params: SimulationParams) → ArpesSpectrum[source]¶

Simulate ARPES spectrum with Voigt broadening and uniform weights.

This entry-level simulation convolves each band with a Voigt profile and applies the Fermi-Dirac occupation. The function adds all non-s orbital projections with equal weights. This simple model includes lifetime and instrument broadening.

See:: TestSimulateNovice

Implementation Logic¶

Build the energy axis:

energy_axis: Float[Array, " E"] = jnp.linspace(
    params.energy_min,
    params.energy_max,
    params.fidelity,
)

This defines the common energy grid for every k-point.

Vectorize the band accumulation:

intensity: Float[Array, "K E"] = jax.vmap(_single_kpoint)(
    bands.eigenvalues, weights
)

JAX maps the differentiable band calculation across all k-points.

Build the spectrum carrier:

spectrum: ArpesSpectrum = make_arpes_spectrum(
    intensity=intensity,
    energy_axis=energy_axis,
)

The factory validates and stores the computed arrays.

param bands:: Electronic band structure containing eigenvalues of shape (K, B) and the Fermi energy.
type bands:: BandStructure
param orb_proj:: Orbital projections with shape (K, B, A, 9). A is the atom count, and 9 is the orbital channel count. The channels follow the VASP order.
type orb_proj:: OrbitalProjection
param params:: Simulation parameters including sigma, gamma, temperature, energy_min, energy_max, and fidelity.
type params:: SimulationParams
returns:: spectrum – Simulated ARPES intensity map with intensity of shape (K, E) and energy_axis of shape (E,).
rtype:: ArpesSpectrum

Notes

Uses Voigt profile (combined Gaussian-Lorentzian) and sums all non-s orbital contributions with equal weight. This is appropriate when orbital cross-section data is unavailable or when a quick qualitative comparison with experiment is sufficient.

See also

simulate_novice_expanded: Expanded variant that returns per-band contributions before summation.

diffpes.simul.simulate_soc(bands: BandStructure, orb_proj: SpinOrbitalProjection, params: SimulationParams, pol_config: PolarizationConfig, ls_scale: float | Float[Array, ''] = 0.01) → ArpesSpectrum[source]¶

Simulate ARPES with spin-orbit coupling (spin-dependent intensity).

The function extends the expert model with a spin-orbit correction. The spin projection along the photon wavevector modulates the orbital band intensity. This model supports spin-ARPES and circular dichroism. It requires orb_proj.spin with shape (K, B, A, 6) (spin up/down for x, y, z). Uses Voigt broadening and the same Yeh-Lindau and polarization logic as simulate_expert.

See:: TestSimulateSoc

Implementation Logic¶

Compute the spin-orbit correction:

band_intensity_soc: Float[Array, "K B"] = band_intensity * (
    1.0 + ls_arr * spin_dot_k
)

This modulates the orbital intensity by the photon-aligned spin.

Vectorize the corrected band accumulation:

intensity: Float[Array, "K E"] = jax.vmap(_single_kpoint)(
    bands.eigenvalues, band_intensity_soc
)

JAX maps the differentiable spin-orbit model across all k-points.

Build the spectrum carrier:

spectrum: ArpesSpectrum = make_arpes_spectrum(
    intensity=intensity,
    energy_axis=energy_axis,
)

The factory validates and stores the computed arrays.

param bands:: Band structure with eigenvalues and Fermi energy.
type bands:: BandStructure
param orb_proj:: Orbital projections with mandatory spin data (shape (K, B, A, 6)).
type orb_proj:: SpinOrbitalProjection
param params:: Simulation parameters (sigma, gamma, fidelity, etc.).
type params:: SimulationParams
param pol_config:: Light polarization and incidence angles.
type pol_config:: PolarizationConfig
param ls_scale:: Spin-orbit coupling strength for the S·k_photon correction. Default is 0.01.
type ls_scale:: Union[float, Float[Array, '']], default: 0.01
returns:: spectrum – Simulated ARPES intensity and energy axis.
rtype:: ArpesSpectrum

Notes

The spin array has six channels for the two directions on each axis. The function adds the two components for each axis. It then sums over atoms to obtain one spin vector for each band. The modulation uses 1 + ls_scale * S·k_photon. With ls_scale=0, the result equals the output from simulate_expert.

See also

simulate_expert: Same physics without spin correction.
photon_wavevector: Builds k_photon from incidence angles.

diffpes.simul.load_vasp_context(directory: str = '.', eigenval_file: str = 'EIGENVAL', procar_file: str = 'PROCAR', doscar_file: str | None = 'DOSCAR', kpoints_file: str | None = 'KPOINTS', fermi_energy: float | Float[Array, ''] | None = None, procar_mode: Literal['legacy', 'full'] = 'full', doscar_mode: Literal['legacy', 'full'] = 'legacy', check_dimensions: bool = True) → WorkflowContext[source]¶

Load a simulation-ready context from VASP output files.

Parses the required band and projection files. It also loads optional density-of-states and k-path data when those files are available.

See:: TestLoadVaspContext

Implementation Logic¶

Resolve the input directory:
```
root = Path(directory)
```
All requested VASP filenames are relative to this path.
Resolve the Fermi energy:
```
resolved_fermi = dos.fermi_energy
resolved_fermi = fermi_energy
```
An explicit value takes priority. Otherwise, DOSCAR supplies the value.
Parse the required carriers:
```
bands = read_eigenval(...)
orb_proj = read_procar(...)
```
EIGENVAL and PROCAR provide the arrays required by every workflow.
Load optional path metadata:
```
kpath = read_kpoints(str(kpoints_path))
```
The parser runs only when the optional KPOINTS file exists.
Validate and construct the context:
```
check_consistency(bands, orb_proj, kpath)
context = make_workflow_context(...)
```
The check rejects incompatible files before the factory builds a PyTree.

param directory:: Directory containing VASP files. Default is current directory.
type directory:: str, default: "."
param eigenval_file:: EIGENVAL filename. Default is "EIGENVAL".
type eigenval_file:: str, default: "EIGENVAL"
param procar_file:: PROCAR filename. Default is "PROCAR".
type procar_file:: str, default: "PROCAR"
param doscar_file:: DOSCAR filename used to infer Fermi energy when fermi_energy is not provided. Use None to skip DOSCAR.
type doscar_file:: Optional[str], default: "DOSCAR"
param kpoints_file:: KPOINTS filename for optional path metadata. Use None to skip KPOINTS parsing.
type kpoints_file:: Optional[str], default: "KPOINTS"
param fermi_energy:: Manual Fermi energy in eV. If None, the function reads DOSCAR when available. Otherwise, it uses 0.0.
type fermi_energy:: Union[float, Float[Array, ''], None], default: None
param procar_mode:: PROCAR return mode. "full" preserves spin data when present.
type procar_mode:: Literal['legacy', 'full'], default: "full"
param doscar_mode:: DOSCAR return mode when DOSCAR is read.
type doscar_mode:: Literal['legacy', 'full'], default: "legacy"
param check_dimensions:: If True, run cross-file consistency checks.
type check_dimensions:: bool, default: True
returns:: context – Loaded VASP data bundled for downstream workflow calls.
rtype:: WorkflowContext
raises FileNotFoundError:: If the function needs DOSCAR to infer the Fermi energy but cannot find the file.

diffpes.simul.prepare_projection(orb_proj: OrbitalProjection | SpinOrbitalProjection, atom_indices: list[int] | None = None, attach_oam: bool = False) → OrbitalProjection | SpinOrbitalProjection[source]¶

Prepare orbital projections for simulation.

Applies common pre-processing steps used in MATLAB-like workflows: selecting atom subsets and attaching OAM channels derived from orbital projections.

See:: TestPrepareProjection

Implementation Logic¶

Select atoms:
```
prepared = select_atoms(prepared, atom_indices)
```
The optional selection keeps only the requested zero-based atom axes.
Attach OAM channels:
```
oam = compute_oam(prepared.projections)
```
The function computes missing channels and rebuilds the same projection carrier type with the new data.

param orb_proj:: Input projection object.
type orb_proj:: Union[OrbitalProjection, SpinOrbitalProjection]
param atom_indices:: Optional 0-based atom indices to keep.
type atom_indices:: Optional[list[int]], default: None
param attach_oam:: If True and OAM is absent, compute OAM and attach it.
type attach_oam:: bool, default: False
returns:: prepared – Prepared projection object, preserving spin-aware type.
rtype:: Union[OrbitalProjection, SpinOrbitalProjection]

diffpes.simul.run_vasp_workflow(level: str = 'advanced', directory: str = '.', eigenval_file: str = 'EIGENVAL', procar_file: str = 'PROCAR', doscar_file: str | None = 'DOSCAR', kpoints_file: str | None = 'KPOINTS', fermi_energy: float | Float[Array, ''] | None = None, atom_indices: list[int] | None = None, attach_oam: bool = False, normalize: bool = False, dk: float | Float[Array, ''] | None = None, procar_mode: Literal['legacy', 'full'] = 'full', doscar_mode: Literal['legacy', 'full'] = 'legacy', check_dimensions: bool = True, sigma: float | Float[Array, ''] = 0.04, gamma: float | Float[Array, ''] = 0.1, fidelity: int = 25000, temperature: float | Float[Array, ''] = 15.0, photon_energy: float | Float[Array, ''] = 11.0, polarization: str = 'unpolarized', incident_theta: float | Float[Array, ''] = 45.0, incident_phi: float | Float[Array, ''] = 0.0, polarization_angle: float | Float[Array, ''] = 0.0, ls_scale: float | Float[Array, ''] = 0.01) → ArpesSpectrum[source]¶

Run an end-to-end VASP-to-ARPES workflow in one call.

This helper loads VASP files into a WorkflowContext and immediately delegates to simulate_context().

See:: TestRunVaspWorkflow

Implementation Logic¶

Load the workflow context:
```
context = load_vasp_context(...)
```
The loader parses the files and checks their shared dimensions.
Run the configured simulation:
```
spectrum = simulate_context(...)
```
The dispatcher applies the requested physics and post-processing steps.

param level:: Simulation complexity level (static; changing it retraces).
type level:: str, default: "advanced"
param directory:: Directory containing the VASP files.
type directory:: str, default: "."
param eigenval_file:: EIGENVAL filename relative to directory.
type eigenval_file:: str, default: "EIGENVAL"
param procar_file:: PROCAR filename relative to directory.
type procar_file:: str, default: "PROCAR"
param doscar_file:: Optional DOSCAR filename used to infer the Fermi energy.
type doscar_file:: Optional[str], default: "DOSCAR"
param kpoints_file:: Optional KPOINTS filename used for path metadata.
type kpoints_file:: Optional[str], default: "KPOINTS"
param fermi_energy:: Explicit Fermi energy in eV, or None to infer it.
type fermi_energy:: Union[float, Float[Array, ''], None], default: None
param atom_indices:: Optional zero-based atom subset.
type atom_indices:: Optional[list[int]], default: None
param attach_oam:: Whether to derive OAM channels before simulation.
type attach_oam:: bool, default: False
param normalize:: Whether to z-score normalize the final intensity.
type normalize:: bool, default: False
param dk:: Momentum broadening width in 1/Angstrom, or None.
type dk:: Union[float, Float[Array, ''], None], default: None
param procar_mode:: PROCAR parsing mode.
type procar_mode:: Literal['legacy', 'full'], default: "full"
param doscar_mode:: DOSCAR parsing mode.
type doscar_mode:: Literal['legacy', 'full'], default: "legacy"
param check_dimensions:: Whether to validate dimensions across input files.
type check_dimensions:: bool, default: True
param sigma:: Gaussian energy broadening width in eV.
type sigma:: Union[float, Float[Array, '']], default: 0.04
param gamma:: Lorentzian energy broadening width in eV.
type gamma:: Union[float, Float[Array, '']], default: 0.1
param fidelity:: Number of energy samples (static; changing it retraces).
type fidelity:: int, default: 25000
param temperature:: Electronic temperature in Kelvin.
type temperature:: Union[float, Float[Array, '']], default: 15.0
param photon_energy:: Incident photon energy in eV.
type photon_energy:: Union[float, Float[Array, '']], default: 11.0
param polarization:: Polarization mode (static; changing it retraces).
type polarization:: str, default: "unpolarized"
param incident_theta:: Incident polar angle in degrees.
type incident_theta:: Union[float, Float[Array, '']], default: 45.0
param incident_phi:: Incident azimuth angle in degrees.
type incident_phi:: Union[float, Float[Array, '']], default: 0.0
param polarization_angle:: Linear polarization angle in radians.
type polarization_angle:: Union[float, Float[Array, '']], default: 0.0
param ls_scale:: Spin-orbit intensity scale.
type ls_scale:: Union[float, Float[Array, '']], default: 0.01
returns:: spectrum – Final simulated spectrum.
rtype:: ArpesSpectrum

diffpes.simul.simulate_context(context: WorkflowContext, level: str = 'advanced', atom_indices: list[int] | None = None, attach_oam: bool = False, normalize: bool = False, dk: float | Float[Array, ''] | None = None, sigma: float | Float[Array, ''] = 0.04, gamma: float | Float[Array, ''] = 0.1, fidelity: int = 25000, temperature: float | Float[Array, ''] = 15.0, photon_energy: float | Float[Array, ''] = 11.0, polarization: str = 'unpolarized', incident_theta: float | Float[Array, ''] = 45.0, incident_phi: float | Float[Array, ''] = 0.0, polarization_angle: float | Float[Array, ''] = 0.0, ls_scale: float | Float[Array, ''] = 0.01) → ArpesSpectrum[source]¶

Run a level-dispatched simulation from a loaded workflow context.

Uses a parsed workflow context as the single input carrier. Optional post-processing preserves its energy axis and returns a new spectrum.

See:: TestSimulateContext

Implementation Logic¶

Prepare the projection carrier:
```
prepared = prepare_projection(...)
```
This step applies the requested atom selection and OAM construction.
Run the selected simulation:
```
spectrum = simulate_expanded(...)
```
The level selects the forward model while JAX traces continuous inputs.
Apply optional post-processing:
```
intensity = apply_momentum_broadening(intensity, k_dist, dk)
intensity = zscore_normalize(intensity)
```
Each enabled operation transforms the intensity without changing the energy axis.
Construct the result:
```
result = make_arpes_spectrum(...)
```
The factory validates the final spectrum carrier.

param context:: Parsed VASP context from load_vasp_context().
type context:: WorkflowContext
param level:: Simulation level for diffpes.simul.simulate_expanded().
type level:: str, default: "advanced"
param atom_indices:: Optional atom subset used before simulation.
type atom_indices:: Optional[list[int]], default: None
param attach_oam:: If True and OAM is absent, compute OAM before simulation.
type attach_oam:: bool, default: False
param normalize:: If True, apply global z-score normalization to output intensity.
type normalize:: bool, default: False
param dk:: If provided, apply momentum broadening along the k-axis.
type dk:: Union[float, Float[Array, ''], None], default: None
param sigma:: Gaussian broadening width in eV.
type sigma:: Union[float, Float[Array, '']], default: 0.04
param gamma:: Lorentzian broadening width in eV.
type gamma:: Union[float, Float[Array, '']], default: 0.1
param fidelity:: Number of points on the energy axis.
type fidelity:: int, default: 25000
param temperature:: Electronic temperature in Kelvin.
type temperature:: Union[float, Float[Array, '']], default: 15.0
param photon_energy:: Incident photon energy in eV.
type photon_energy:: Union[float, Float[Array, '']], default: 11.0
param polarization:: Polarization type, for example "unpolarized" or "LHP".
type polarization:: str, default: "unpolarized"
param incident_theta:: Incident polar angle in degrees.
type incident_theta:: Union[float, Float[Array, '']], default: 45.0
param incident_phi:: Incident azimuth angle in degrees.
type incident_phi:: Union[float, Float[Array, '']], default: 0.0
param polarization_angle:: Linear polarization angle in radians.
type polarization_angle:: Union[float, Float[Array, '']], default: 0.0
param ls_scale:: SOC scale used only when level="soc".
type ls_scale:: Union[float, Float[Array, '']], default: 0.01
returns:: result – Simulated ARPES spectrum after optional post-processing.
rtype:: ArpesSpectrum