diffpes.utils

Mathematical utilities for the ARPES simulation pipeline: the Faddeeva function used by Voigt broadening and z-score normalization for spectra.

Provide utility functions for ARPES simulations.

Extended Summary

The subpackage provides mathematical utilities for ARPES simulations. The Faddeeva function uses a 64-term Taylor series. The Voigt broadening profile uses this function. Z-score normalization prepares spectra for comparisons with experiments. Complex packing functions provide the required real-valued optimizer boundary for complex physics parameters.

The following list describes the submodules:

• math

  Compute mathematical utilities for ARPES simulations.

Routine Listings

faddeeva()

Evaluate the Faddeeva function w(z) = exp(-z^2) erfc(-iz).

pack_complex()

Pack complex parameters as stacked real values.

unpack_complex()

Unpack stacked real parameters into complex values.

zscore_normalize()

Apply z-score normalization (zero-mean, unit-variance).

Notes

All functions support JAX transformations and automatic differentiation. The Faddeeva implementation uses jax.lax.scan for the coefficient recurrence.

diffpes.utils.faddeeva(z: Complex[Array, '...']) → Complex[Array, '...']

Evaluate the Faddeeva function w(z) = exp(-z^2) erfc(-iz).

The function computes the Faddeeva function for arbitrary complex arrays. It applies Horner’s method to a precomputed Taylor polynomial.

The following equation defines the Faddeeva function:

\[w(z) = e^{-z^2} \operatorname{erfc}(-iz) = e^{-z^2} \left(1 + \frac{2i}{\sqrt{\pi}} \int_0^z e^{t^2} dt \right)\]

The Voigt profile uses the real part of the Faddeeva function along the imaginary axis. ARPES simulations use it to convolve Lorentzian lifetime broadening with Gaussian instrument resolution.

Implementation:

1. Cast to complex128: Convert the input to jnp.complex128 with jnp.asarray. This type provides sufficient precision for the 64-term polynomial.

2. Apply Horner’s method: Call jnp.polyval with _W_POLY and the converted input. _W_POLY stores the coefficients in descending power order. The unroll=8 hint lets XLA pipeline the inner loop.

The ODE \(w'(z) = -2z w(z) + 2i/\sqrt{\pi}\) defines the Taylor coefficients. The _faddeeva_taylor_coeffs function computes them during module import. Its docstring gives the full recurrence derivation.

See:

TestFaddeeva

Parameters:

z (']) – Complex argument(s), arbitrary shape.

Returns:

w – Faddeeva function values, same shape as z.

Return type:

']

Notes

The result has approximately 15 significant digits for |z| < 6. The Taylor series converges slowly for |z| >= 6. Use an asymptotic expansion or a continued fraction outside the convergence domain. The function does not select these alternatives automatically.

The function is fully differentiable via JAX autodiff, which is important for gradient-based optimization of broadening parameters in inverse-fitting workflows.

diffpes.utils.pack_complex(z: Complex[Array, '...']) → Float[Array, '... 2']

Pack complex parameters as stacked real values.

Complex parameters cross the optimizer and Fisher-information boundary as stacked reals, while values remain complex inside the physics pipeline. This function is the sanctioned complex-to-real crossing point.

See:

TestPackComplex

Parameters:

z (']) – Complex-valued physics parameters of arbitrary shape.

Returns:

packed – Real-valued parameters with real and imaginary components in the final axis, in that order.

Return type:

2']

Notes

The function forms the final axis with jnp.stack([z.real, z.imag], axis=-1). This operation preserves the component dtype and exposes independent real optimizer coordinates.

See also

unpack_complex

Restore complex values inside the physics pipeline.

diffpes.utils.unpack_complex(p: Float[Array, '... 2']) → Complex[Array, '...']

Unpack stacked real parameters into complex values.

Complex parameters cross the optimizer and Fisher-information boundary as stacked reals, while values remain complex inside the physics pipeline. This function is the sanctioned real-to-complex crossing point.

See:

TestUnpackComplex

Parameters:

p ( 2']) – Real-valued optimizer parameters whose final axis stores real and imaginary components, in that order.

Returns:

unpacked – Complex-valued physics parameters with the packing axis removed.

Return type:

']

Notes

jax.lax.complex combines the final-axis components without changing their precision. For a real loss, JAX differentiates the two packed components as ordinary real optimizer coordinates.

See also

pack_complex

Expose complex parameters at the real optimizer boundary.

diffpes.utils.zscore_normalize(data: Float[Array, '...']) → Float[Array, '...']

Apply z-score normalization (zero-mean, unit-variance).

The function transforms a float array to zero mean and unit standard deviation. This transformation prepares simulated and experimental ARPES spectra for comparison. The z-score transformation is:

\[\hat{x}_i = \frac{x_i - \bar{x}}{\sigma}\]

where \(\bar{x}\) is the global mean and \(\sigma\) is the population standard deviation (\(\text{ddof}=0\)).

Implementation details:

1. Compute the statistics: Compute the global mean with jnp.mean. Compute the population standard deviation with jnp.std and ddof=0.

2. Guard against zero deviation: Pass the centered values and standard deviation to safe_divide(). The function selects zero for a constant input and defines a zero subgradient.

3. Normalize the values: Compute (data - mean) / safe_std for each element and return the result.

See:

TestZscoreNormalize

Parameters:

data (']) – Input data array of any shape.

Returns:

normalized – Normalized data with mean 0 and standard deviation 1 (or all zeros if the input is constant).

Return type:

']

Notes

The function computes one global mean and standard deviation over all elements. For each-axis normalization, reshape the data and call the function on each slice.

This function is differentiable via JAX autodiff with respect to the input data. The gradient propagates through both the mean-subtraction and the division by standard deviation.