diffpes.utils tests

Validation and regression tests for diffpes.utils.

Validate the ARPES mathematical utilities.

Extended Summary

The tests exercise the Faddeeva function, z-score normalization, and the complex packing boundary. They cover JIT, vectorization, precision, round trips, and gradients.

class tests.test_diffpes.test_utils.test_math.TestFaddeeva(*args, **kwargs) → None

Bases: TestCase

Validate diffpes.utils.math.faddeeva().

The tests cover both coordinate axes and the analytic value at the origin. Each test uses compiled and uncompiled execution.

See:

faddeeva()

test_imaginary_axis__with_jit() → None

Verify faddeeva returns finite values on the imaginary axis.

The function preserves the input shape and produces finite real components for three imaginary inputs.

Notes

The test uses i*[0.5, 1.0, 2.0]. It checks shape (3,) and the finite real components under both JAX variants.

Return type:

None

test_imaginary_axis__without_jit() → None

Verify faddeeva returns finite values on the imaginary axis.

The function preserves the input shape and produces finite real components for three imaginary inputs.

Notes

The test uses i*[0.5, 1.0, 2.0]. It checks shape (3,) and the finite real components under both JAX variants.

Return type:

None

test_real_axis__with_jit() → None

Verify faddeeva returns finite values on the real axis.

The function preserves the input shape and produces finite real components for the specified interval.

Notes

The test uses 100 complex points across [-3, 3]. It checks shape (100,) and the finite real components under both JAX variants.

Return type:

None

test_real_axis__without_jit() → None

Verify faddeeva returns finite values on the real axis.

The function preserves the input shape and produces finite real components for the specified interval.

Notes

The test uses 100 complex points across [-3, 3]. It checks shape (100,) and the finite real components under both JAX variants.

Return type:

None

test_zero__with_jit() → None

Verify faddeeva(0) returns approximately 1.0.

The Faddeeva function satisfies w(0) = erfc(0) = 1.

Notes

The test supplies the complex scalar zero. It compares the real component with 1.0 at an absolute tolerance of 0.05 under both variants.

Return type:

None

test_zero__without_jit() → None

Verify faddeeva(0) returns approximately 1.0.

The Faddeeva function satisfies w(0) = erfc(0) = 1.

Notes

The test supplies the complex scalar zero. It compares the real component with 1.0 at an absolute tolerance of 0.05 under both variants.

Return type:

None

class tests.test_diffpes.test_utils.test_math.TestZscoreNormalize(*args, **kwargs) → None

Bases: TestCase

Validate diffpes.utils.math.zscore_normalize().

The tests cover standard data, constant data, and two-dimensional data. They verify the global normalization and the zero-variance guard.

See:

zscore_normalize()

test_2d_input__with_jit() → None

Verify zscore_normalize handles arrays with two dimensions.

The function preserves the two-dimensional shape and uses global statistics across both axes.

Notes

The test normalizes arange(12) with shape (3, 4). It checks the shape and compares the global mean with zero at a tolerance of 1e-10.

Return type:

None

test_2d_input__without_jit() → None

Verify zscore_normalize handles arrays with two dimensions.

The function preserves the two-dimensional shape and uses global statistics across both axes.

Notes

The test normalizes arange(12) with shape (3, 4). It checks the shape and compares the global mean with zero at a tolerance of 1e-10.

Return type:

None

test_constant_input__with_jit() → None

Verify zscore_normalize returns zeros for a constant array.

The zero-variance guard produces a finite zero array.

Notes

The test normalizes an array of ten ones under both JAX variants. It compares the result with ten zeros at an absolute tolerance of 1e-10.

Return type:

None

test_constant_input__without_jit() → None

Verify zscore_normalize returns zeros for a constant array.

The zero-variance guard produces a finite zero array.

Notes

The test normalizes an array of ten ones under both JAX variants. It compares the result with ten zeros at an absolute tolerance of 1e-10.

Return type:

None

test_normalized_stats__with_jit() → None

Verify zscore_normalize produces zero mean and unit deviation.

Standard z-score normalization produces a mean of zero and a standard deviation of one.

Notes

The test normalizes [1, 2, 3, 4, 5]. It compares both statistics with their analytic values at an absolute tolerance of 1e-10.

Return type:

None

test_normalized_stats__without_jit() → None

Verify zscore_normalize produces zero mean and unit deviation.

Standard z-score normalization produces a mean of zero and a standard deviation of one.

Notes

The test normalizes [1, 2, 3, 4, 5]. It compares both statistics with their analytic values at an absolute tolerance of 1e-10.

Return type:

None

class tests.test_diffpes.test_utils.test_math.TestPackComplex(*args, **kwargs) → None

Bases: TestCase

Validate pack_complex().

The tests cover exact round trips, dtype preservation, JIT, and vectorization. Asymmetric complex values reveal an incorrect coordinate order.

See:

pack_complex()

test_round_trip_and_dtype() → None

Preserve generic complex128 values exactly through a JIT round trip.

The input uses unequal real and imaginary components. These values reveal a component swap or a real-symmetric implementation.

Notes

The test packs and unpacks a (2, 2) array with jax.jit. It checks the packed shape, both dtypes, and exact value equality.

Return type:

None

test_vmap() → None

Vectorize packing independently over a leading parameter batch.

The result retains the batch and parameter axes. It appends only the two-component packing axis.

Notes

The test applies jax.vmap to three complex parameter vectors. It compares the result with direct packing and checks the exact round trip.

Return type:

None

class tests.test_diffpes.test_utils.test_math.TestUnpackComplex(*args, **kwargs) → None

Bases: TestCase

Validate unpack_complex().

The tests cover exact round trips and dtype preservation. They also compare complex magnitude gradients with gradients in real coordinates.

See:

unpack_complex()

test_round_trip_and_dtype() → None

Preserve generic stacked float64 values exactly through a JIT round trip.

The final input axis contains asymmetric real and imaginary coordinates. These values reveal an incorrect coordinate order.

Notes

The test unpacks and packs a (2, 3, 2) array with jax.jit. It checks the unpacked shape, both dtypes, and exact coordinate equality.

Return type:

None

test_gradient_equivalence() → None

Match complex-magnitude gradients to packed real coordinates exactly.

For p = stack([x, y]), the real gradient equals stack([2x, 2y]).

Notes

The test differentiates a compiled loss on generic float64 coordinates. It compares the gradient with twice the input by exact equality.

Return type:

None

class tests.test_diffpes.test_utils.test_math.TestComplexAutodiffConvention(*args, **kwargs) → None

Bases: TestCase

Pin JAX’s complex-gradient convention at the packing boundary.

The tests fix the conjugated Wirtinger convention at the boundary between complex physics values and real optimizer coordinates.

See:

pack_complex()

See:

unpack_complex()

test_wirtinger_convention() → None

Pin grad(abs(z)**2) at 1+1j to exactly 2-2j.

The convention determines the relation between complex gradients and gradients of stacked real coordinates.

Notes

The test applies reverse-mode autodiff at 1+1j. It checks exact equality with 2-2j to detect a change in the JAX convention.

Return type:

None