diffpes.maths tests

Validation and regression tests for diffpes.maths.

Validate dipole matrix element assembly.

Extended Summary

The tests validate the complete matrix-element pipeline for the ARPES photoemission intensity. The pipeline includes the single-orbital function dipole_matrix_element_single, the per-orbital intensity wrapper dipole_intensity_orbital, and the multi-orbital batch function dipole_intensities_all_orbitals. The tests cover allowed transitions, finite values, JIT compatibility, and nonnegative intensities. They compare |M|^2 with the intensity function and verify the output shapes. They also verify derivatives for the polarization vector and Slater exponent.

tests.test_diffpes.test_maths.test_dipole.r_grid() → Float[Array, 'R']

Provide a dense radial grid for numerical integration.

The fixture returns 5000 points from 1e-6 to 50.0 Bohr. The nonzero lower bound avoids the 1/r singularity in Slater radial functions. The upper bound lets typical exponential tails decay below machine precision.

Return type:

Float[Array, 'R']

tests.test_diffpes.test_maths.test_dipole.z_polarized() → Complex[Array, '3']

Provide a z-polarized electric-field vector.

The fixture returns the complex polarization vector (0, 0, 1). This vector selects the q=0 component of the dipole operator. It represents linear polarization along the z-axis.

Return type:

Complex[Array, '3']

tests.test_diffpes.test_maths.test_dipole.x_polarized() → Complex[Array, '3']

Provide an x-polarized electric-field vector.

The fixture returns the complex polarization vector (1, 0, 0). This vector mixes the q=+1 and q=-1 spherical components. It represents linear polarization along the x-axis.

Return type:

Complex[Array, '3']

class tests.test_diffpes.test_maths.test_dipole.TestDipoleMatrixElementSingle

Bases: object

Validate dipole_matrix_element_single.

The tests validate the core single-orbital dipole matrix element. This element combines a radial integral with Gaunt angular coupling. The tests verify allowed transitions, finite values, and JIT compatibility. They also verify autodiff gradients for the electric-field vector.

See:

dipole_matrix_element_single()

test_s_orbital_nonzero(r_grid, z_polarized) → None

Verify the s -> p dipole transition is nonzero for z-polarization.

The test uses an n=1, zeta=1.0 Slater 1s orbital with z-polarized light and k along z. The dipole selection rule allows l=0 to l'=1. With q=0, the matrix element must be nonzero. The test asserts |M| > 1e-6 to confirm the angular (Gaunt) and radial integrals both contribute.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_finite_output(r_grid, z_polarized) → None

Verify the matrix element is numerically finite for a p-orbital.

The test uses an n=2, zeta=2.0 Slater 2p orbital. It uses the k-vector (1.0, 0.5, 0.3) and z-polarization. It asserts jnp.isfinite(M) to guard against NaN or Inf from the radial integration or angular coupling at arbitrary k-directions.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_jit_compatible(r_grid, z_polarized) → None

Verify dipole_matrix_element_single is JAX-JIT-compatible.

The test wraps the function in jax.jit with k as the traced argument. The closure captures the other inputs as constants. The test uses an s-orbital (l=0, m=0) with k along z. Asserts the JIT- compiled output is finite, confirming no Python-side control flow or non-JAX operations break tracing.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_gradient_wrt_efield(r_grid) → None

Verify autodiff gradient of intensity w.r.t. polarization vector.

The test defines loss = |M(E)|^2 for an s-orbital with k along z. It differentiates the complete dipole pipeline for the electric-field vector with jax.grad. It asserts all three gradient components are finite, confirming the Gaunt lookup, radial integral, and complex-valued inner product are all smoothly differentiable.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

class tests.test_diffpes.test_maths.test_dipole.TestDipoleIntensityOrbital

Bases: object

Validate dipole_intensity_orbital.

Validates the single-orbital intensity function, which returns the modulus squared of the dipole matrix element. Tests verify the physical constraint that intensity is non-negative and that the function is consistent with computing |M|^2 directly from dipole_matrix_element_single.

See:

dipole_intensity_orbital()

test_non_negative(r_grid, z_polarized) → None

Verify photoemission intensity is non-negative.

The test uses an n=2, zeta=1.5 Slater p-orbital (l=1, m=0) with an off-axis k-vector and z-polarization. Intensity equals |M|^2 and must be nonnegative. The test asserts I >= 0.0 as a basic physical sanity check.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_equals_abs_squared(r_grid, z_polarized) → None

Verify intensity equals the modulus squared of the matrix element.

The test computes M and I independently for a 1s orbital. It uses k along z and z-polarization. It asserts I == |M|^2 within 1e-12. This comparison confirms that dipole_intensity_orbital is a faithful wrapper around dipole_matrix_element_single.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

class tests.test_diffpes.test_maths.test_dipole.TestDipoleIntensitiesAllOrbitals

Bases: object

Validate dipole_intensities_all_orbitals.

The batch function computes the dipole intensity for each orbital. The tests verify one intensity per orbital and nonnegative values. They also verify derivatives of the total intensity for the Slater exponent zeta.

See:

dipole_intensities_all_orbitals()

test_output_shape(r_grid, z_polarized) → None

Verify output has one intensity per orbital in the basis.

The test constructs a three-orbital basis (1s, 2pz, 2px) with corresponding Slater exponents and computes intensities at a single k-point. The test asserts the output shape is (3,), matching the number of orbitals in the basis.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_all_non_negative(r_grid, z_polarized) → None

Verify all orbital intensities are non-negative.

The test constructs a four-orbital basis (1s, 2py, 2pz, 2px) spanning all three p sub-orbitals plus s, with an off-axis k-vector and z-polarization. Asserts I >= 0 element-wise, confirming the |M|^2 definition holds for every orbital in the batch.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_gradient_wrt_zeta(r_grid, z_polarized) → None

Verify autodiff gradient of total intensity w.r.t. Slater exponent.

The test defines loss = sum(I) for one 1s orbital. It differentiates the loss for zeta with jax.grad. The test asserts that the gradient is finite. This result confirms that the complete forward pipeline supports reverse-mode autodiff. Inverse fits of orbital exponents require this property.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

Validate Gaunt coefficient table.

Extended Summary

The tests validate the precomputed Gaunt coefficient table for angular coupling in dipole matrix elements. Gaunt coefficients encode the integral of three spherical harmonics over the unit sphere. They enforce the dipole selection rules Delta l = +/-1 and Delta m = q. The tests verify the table shape and both selection rules. They also verify allowed transitions, reproducible construction, real values, and the positive fundamental s-p coupling.

class tests.test_diffpes.test_maths.test_gaunt.TestBuildGauntTable

Bases: object

Validate the precomputed Gaunt coefficient table.

Validates the module-level GAUNT_TABLE array and the gaunt_lookup accessor function. build_gaunt_table(l_max=4) builds the table at import time. The table stores real-valued Gaunt coefficients indexed by (l, m, q, l’, m’). Tests systematically check selection rules, allowed transitions, table shape, dtype, reproducibility, and a known analytical value.

See:

build_gaunt_table()

test_table_shape() → None

Verify the precomputed table has the expected 5-D shape.

The five axes represent l, m, q, l', and m'. Their respective sizes are 5, 9, 3, 6, and 11 for l_max=4. The test asserts the exact shape (5, 9, 3, 6, 11).

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_selection_rule_delta_l() → None

Verify the dipole selection rule Delta l = +/-1.

The test iterates over every valid table index combination. It asserts that whenever |l' - l| != 1 the Gaunt coefficient is zero (< 1e-12). This is the fundamental angular momentum selection rule for electric-dipole transitions.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_s_to_p_nonzero() → None

Verify the s -> p allowed transition has a nonzero Gaunt coefficient.

Looks up G(l=0, m=0, q=0, l’=1, m’=0), the prototypical electric-dipole transition from an s-orbital to pz. Asserts |G| > 1e-6, confirming the table correctly encodes the coupling.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_p_to_s_nonzero() → None

Verify the p -> s allowed transition has a nonzero Gaunt coefficient.

Looks up G(l=1, m=0, q=0, l’=0, m’=0), the reverse of the s->p transition. Asserts |G| > 1e-6, confirming reciprocity of the Gaunt integral.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_p_to_d_nonzero() → None

Verify the p -> d allowed transition has a nonzero Gaunt coefficient.

Looks up G(l=1, m=0, q=0, l’=2, m’=0), a higher-l allowed dipole transition. Asserts |G| > 1e-6, confirming that the table covers transitions beyond the lowest-order s<->p coupling.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_forbidden_delta_m() → None

Verify the magnetic selection rule forbids |Delta m| > 1.

Looks up G(l=2, m=0, q=0, l’=1, m’=2) where m’ - m = 2 != q = 0. The selection rule requires m' = m + q and forbids m'=2 here. The test asserts the coefficient is zero (< 1e-12).

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_rebuild_matches_precomputed() → None

Verify that rebuilding the table reproduces the precomputed values.

The test calls build_gaunt_table(l_max=4) at test time and compares to the module-level GAUNT_TABLE with jnp.allclose. This guards against silent corruption of the cached table and confirms deterministic construction.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_real_valued() → None

Verify the table dtype is float64 (no imaginary residuals).

The Gaunt coefficients for real spherical harmonics are purely real. Asserts the table dtype is jnp.float64, confirming the construction did not accidentally introduce complex values.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_known_value_y00_dipole() → None

Verify the s-to-p Gaunt coefficient is a known positive value.

The integral G(0, 0, 0, 1, 0) = integral Y_0^0 * Y_1^0 * Y_1^0 dOmega is analytically positive. While the exact numerical value depends on normalization conventions, the sign is unambiguous. Asserts val > 0.0 as a consistency check against sign errors in the Condon-Shortley phase convention.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

class tests.test_diffpes.test_maths.test_gaunt.TestGauntLookup

Bases: object

Validate gaunt_lookup().

Covers indexed retrieval from the canonical dipole Gaunt table for one allowed transition and one forbidden magnetic transition.

See:

gaunt_lookup()

test_matches_canonical_table_entries() → None

Match lookup results to the canonical table at allowed and zero entries.

The accessor must preserve both the positive s-to-p coefficient and an exactly forbidden magnetic channel under the package indexing convention.

Notes

The test evaluates two scalar lookups and compares them with their directly indexed GAUNT_TABLE entries at zero absolute and relative tolerance.

Return type:

None

class tests.test_diffpes.test_maths.test_gaunt.TestWigner3jSelectionRules

Bases: object

Validate internal Wigner 3-j and complex Gaunt zero paths.

The tests exercise early returns in _wigner3j and _complex_gaunt. These branches return 0.0 for violations of selection rules.

See:

build_gaunt_table()

test_abs_m_exceeds_j_returns_zero() → None

Verify |m1| > j1 causes _wigner3j to return 0.0.

The test constructs a call where |m1| = 2 > j1 = 1, violating the |mi| <= ji constraint, and asserts the result is 0.0 (line 111).

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_triangle_inequality_violated_returns_zero() → None

Verify triangle inequality violation causes _wigner3j to return 0.0.

The test uses j1=2, j2=1, j3=0 where j3 < |j1 - j2| = 1, violating the triangle inequality, and asserts the result is 0.0 (line 113).

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_complex_gaunt_zero_w3j_000_returns_zero() → None

Verify that zero w3j_000 (parity violation) causes _complex_gaunt to return 0.0.

The three-j symbol (2,1,0 | 0,0,0) is zero for this input. The parity rule gives zero because l1+l2+l3 = 3 is odd. The test asserts that the complex Gaunt integral returns 0.0.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

Validate differentiable Cartesian rotations.

Extended Summary

The tests compare Rodrigues matrices with SciPy rotations. They check orthogonality, orientation, fixed axes, JAX transforms, and gradients.

class tests.test_diffpes.test_maths.test_rotations.TestRodriguesRotation(*args, **kwargs) → None

Bases: TestCase

Validate rodrigues_rotation().

The tests cover external values, rotation invariants, the zero-axis guard, JIT, vmap, and gradients in the axis and angle.

See:

rodrigues_rotation()

test_matches_scipy_rotation() → None

Match SciPy rotation matrices to relative tolerance 1e-12.

The test normalizes three external axes. SciPy then constructs each active matrix from the corresponding rotation vector.

Notes

Compare each float64 matrix with Rotation.from_rotvec. Use rtol=1e-12 and atol=1e-12.

Return type:

None

test_preserves_rotation_invariants() → None

Verify orthogonality, orientation, and the fixed-axis property.

A proper rotation has R @ R.T = I and determinant one. It also leaves its normalized axis unchanged.

Notes

Use a generic axis and a 0.73 radian angle. Check each invariant at absolute tolerance 1e-14.

Return type:

None

test_zero_axis_returns_identity_with_finite_gradient() → None

Verify the identity value and finite gradient at a zero axis.

The safe axis normalization selects a zero subgradient at the origin. Thus, the matrix remains finite for this boundary input.

Notes

JIT the matrix calculation. Differentiate its sum with respect to the three axis components.

Return type:

None

test_jit_and_vmap_match_eager_results() → None

Verify consistent matrices under JIT and vmap.

The test uses one generic axis and five traced angles. It compares the batched matrices with an eager stack.

Notes

JIT a vmapped Rodrigues call over the angle axis. Compare all entries at rtol=1e-14.

Return type:

None

test_axis_and_angle_gradients_match_finite_differences() → None

Match axis and angle gradients with central finite differences.

A generic matrix contraction depends on both traced inputs. The shared harness checks forward mode, reverse mode, and central differences.

Notes

Contract the matrix with asymmetric weights. Require nonzero gradient norms for the axis leaf and the angle leaf.

Return type:

None

Validate named gradient-safe elementary operations.

Extended Summary

The tests compare all seven safe-math helpers with NumPy away from their guards. They compare both autodiff modes with central finite differences. They also verify each documented guard value and subgradient. Property tests cover broadcasting and sign symmetry on generated finite inputs.

class tests.test_diffpes.test_maths.test_safe.TestSafeDivide

Bases: object

Validate safe_divide().

Covers ordinary broadcast division, both autodiff modes, and the finite fallback with zero operand subgradients at a zero denominator.

See:

safe_divide()

test_values_gradients_and_guard() → None

Match division and pin zero-denominator guard behavior.

Return type:

None

Extended Summary

Ordinary values must match NumPy to rtol=1e-15 and derivatives must pass the smooth program-wide gradient harness. At a zero divisor, the function returns the configured fallback. Both operand gradients equal zero and contain no NaNs.

Notes

The test uses a length-three numerator and a nonzero divisor. It differentiates a scalar guarded call for both operands with JIT-backed autodiff.

test_broadcasting_and_sign_symmetry() → None

Preserve broadcasting and odd symmetry in the numerator.

Return type:

None

Extended Summary

For finite vectors and a positive scalar divisor, safe_divide must broadcast the scalar. It must also preserve odd symmetry within floating-point tolerance.

Notes

Hypothesis generates twenty three-component vectors and divisors away from zero; Chex checks shape and sign symmetry at rtol=1e-15.

class tests.test_diffpes.test_maths.test_safe.TestSafeSqrt

Bases: object

Validate safe_sqrt().

Covers NumPy agreement, both autodiff modes, and the selected zero value and subgradient on non-positive inputs.

See:

safe_sqrt()

test_values_gradients_and_guard() → None

Match square roots and pin the non-positive guard.

Return type:

None

Extended Summary

Positive values must match NumPy to rtol=1e-15 and pass the smooth gradient harness. Zero and negative inputs must both return zero with exact zero gradients and finite traced values.

Notes

The test evaluates three positive values away from the branch boundary and a JIT-compiled two-value guard probe before differentiating each guard.

class tests.test_diffpes.test_maths.test_safe.TestSafeNorm

Bases: object

Validate safe_norm().

Covers Euclidean values, axis handling, sign symmetry, both derivative modes, and the zero-vector subgradient convention.

See:

safe_norm()

test_values_gradients_and_guard() → None

Match Euclidean norms and pin the zero-vector gradient.

Return type:

None

Extended Summary

Batched vector norms must match NumPy to rtol=1e-15 with retained dimensions and pass the gradient harness. An all-zero vector must return zero and have an exactly zero, finite gradient.

Notes

The test uses two generic three-vectors for the ordinary path and an independent JIT and reverse-mode probe at the origin.

test_sign_symmetry() → None

Preserve norm symmetry under vector sign reversal.

Return type:

None

Extended Summary

The safe guard must not change the Euclidean invariant norm(-x) = norm(x) for generated finite three-vectors.

Notes

Hypothesis generates twenty vectors across both signs and Chex checks the two scalar norms at rtol=1e-15.

class tests.test_diffpes.test_maths.test_safe.TestSafeArccos

Bases: object

Validate safe_arccos().

Covers NumPy agreement, both derivative modes, saturation outside the cosine domain, and zero subgradients at both endpoints.

See:

safe_arccos()

test_values_gradients_and_guard() → None

Match arccos and pin saturated endpoint gradients.

Return type:

None

Extended Summary

Interior values must match NumPy to rtol=1e-15 and pass the smooth gradient harness. Inputs at and beyond -1 and 1 must saturate to pi and zero with exact zero, finite gradients.

Notes

The test uses three generic interior cosines and a four-value JIT guard probe spanning both endpoints and both out-of-domain sides.

class tests.test_diffpes.test_maths.test_safe.TestSafeArctan2

Bases: object

Validate safe_arctan2().

Covers NumPy agreement in all quadrants, broadcasting, both derivative modes, and the zero value and coordinate gradients at the origin.

See:

safe_arctan2()

test_values_gradients_and_guard() → None

Match arctan2 and pin the origin convention.

Return type:

None

Extended Summary

Away from the origin, values must match NumPy to rtol=1e-15 and derivatives must pass the gradient harness. At (0, 0), the value and both coordinate gradients must be exactly zero and finite.

Notes

The test uses points in three quadrants for ordinary behavior, then runs the scalar origin through JIT and differentiates both arguments.

class tests.test_diffpes.test_maths.test_safe.TestSafeLog

Bases: object

Validate safe_log().

Covers NumPy agreement, both derivative modes, and the finite floor value with a zero input subgradient at and below the floor.

See:

safe_log()

test_values_gradients_and_guard() → None

Match logarithms and pin floor-clamped gradients.

Return type:

None

Extended Summary

Positive values above the floor must match NumPy to rtol=1e-15 and pass the gradient harness. Inputs at and below a 1e-3 floor must return log(1e-3) with exact zero, finite gradients.

Notes

The test uses three ordinary positive values and a three-value JIT probe containing the floor, zero, and a negative input.

class tests.test_diffpes.test_maths.test_safe.TestSafePower

Bases: object

Validate safe_power().

Covers fractional powers, both derivative modes, and the zero value and base/exponent subgradients on non-positive bases.

See:

safe_power()

test_values_gradients_and_guard() → None

Match fractional powers and pin non-positive gradients.

Return type:

None

Extended Summary

Positive fractional powers must match NumPy to rtol=1e-15 and pass the gradient harness. Zero and negative bases must return zero with exact zero gradients for both the base and exponent.

Notes

The test uses exponent 1.7 on three positive bases, then evaluates and differentiates a two-base guarded sum under JIT-compatible primitives.

Validate real spherical harmonics.

Extended Summary

The tests validate the real spherical harmonic functions real_spherical_harmonic and real_spherical_harmonics_all. They cover known analytical values for low-order harmonics with the Condon-Shortley phase convention. They verify numerical orthonormality, JIT compatibility, and autodiff gradients for theta and phi. They also verify input validation, batch shape, and consistency between the single and batch interfaces.

class tests.test_diffpes.test_maths.test_spherical_harmonics.TestRealSphericalHarmonic

Bases: object

Validate real_spherical_harmonic.

The tests validate the scalar real spherical harmonic with the Condon-Shortley phase convention. They compare l=0 and l=1 harmonics with closed-form expressions. They verify orthonormality with numerical quadrature. They also check JIT, autodiff, and input validation.

See:

real_spherical_harmonic()

test_y00_constant() → None

Verify Y_0^0 = 1/(2*sqrt(pi)) is constant over the sphere.

The test evaluates Y_0^0 at four different (theta, phi) pairs using jax.vmap. The l=0 harmonic is the unique isotropic solution with the well-known value 1/(2*sqrt(pi)). Asserts all four outputs match this constant to within 1e-12.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_y10_cosine() → None

Verify Y_1^0 = sqrt(3/(4*pi)) * cos(theta) along the meridian.

The test evaluates Y_1^0 at 50 theta values from 0 to pi with phi=0 using jax.vmap. Compares against the analytical expression sqrt(3/(4*pi)) * cos(theta). Asserts agreement to within 1e-10, validating the associated Legendre polynomial P_1^0 = cos(theta) and the normalization constant.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_y11_sin_cos() → None

Verify Y_1^1 = -sqrt(3/(4*pi)) * sin(theta) * cos(phi).

The test evaluates Y_1^{+1} at theta=pi/4, phi=0 and compares against the Condon-Shortley convention expression. The negative sign comes from the (-1)^m phase factor for m > 0. Asserts agreement to within 1e-10.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_y1m1_sin_sin() → None

Verify Y_1^{-1} = +sqrt(3/(4*pi)) * sin(theta) * sin(phi).

The test evaluates Y_1^{-1} at theta=pi/3, phi=pi/4 and compares against the analytical expression. For m < 0, the real spherical harmonic uses sin(|m|*phi) and the Condon-Shortley phase cancels, yielding a positive prefactor. Asserts agreement to within 1e-10.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_orthonormality_low_l() → None

Verify orthonormality of Y_0^0 and Y_1^0 via numerical quadrature.

The test constructs a (100 x 200) grid in (theta, phi) using midpoint quadrature in cos(theta) and uniform spacing in phi. Computes the cross-overlap integral Y_0^0 * Y_1^0 * sin(theta) and the self-overlap integral |Y_0^0|^2 * sin(theta). Asserts the cross-overlap is < 0.01 (orthogonality) and the self-overlap is within 0.01 of 1.0 (normalization). The moderate grid density gives ~1% numerical accuracy, sufficient for a smoke test of the spherical harmonic implementation.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_jit_compatible() → None

Verify real_spherical_harmonic is JAX-JIT-compatible.

The test wraps a call to Y_2^1(theta, phi) in jax.jit and evaluates at theta=1.0, phi=0.5. Asserts the output is finite, confirming no Python-side operations (e.g., if-branches on traced values) break JAX tracing for the l=2 case.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_gradient_theta() → None

Verify autodiff gradient of Y_1^0 w.r.t. theta is finite.

The test differentiates Y_1^0(theta, phi=0) at theta=pi/4 using jax.grad. The analytical derivative is -sqrt(3/(4*pi)) * sin(theta), which is nonzero at pi/4. Asserts finiteness, confirming the associated Legendre polynomial implementation supports reverse-mode AD through theta.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_gradient_phi() → None

Verify autodiff gradient of Y_1^1 w.r.t. phi is finite.

The test differentiates Y_1^1(theta=pi/4, phi) at phi=0.5 using jax.grad. The m=1 harmonic has a cos(phi) dependence, so the derivative involves -sin(phi) and is nonzero at phi=0.5. The test asserts finiteness, confirming the azimuthal (phi) branch of the implementation supports AD.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_invalid_l_raises() → None

Verify that negative l raises ValueError.

The test calls real_spherical_harmonic(-1, 0, 0, 0) and expects a ValueError with a matching message. This tests the guard for the boundary condition l >= 0.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_invalid_m_raises() → None

Verify that |m| > l raises ValueError.

The test calls real_spherical_harmonic(1, 2, 0, 0) with m > l. It expects a ValueError with a matching message. This tests the constraint |m| <= l.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

class tests.test_diffpes.test_maths.test_spherical_harmonics.TestRealSphericalHarmonicsAll

Bases: object

Validate real_spherical_harmonics_all.

The batch function returns all real spherical harmonics through l_max at one (theta, phi) point. The tests verify the output shape and consistency with individual real_spherical_harmonic calls.

See:

real_spherical_harmonics_all()

test_output_shape() → None

Verify output has (l_max+1)^2 entries.

The test calls real_spherical_harmonics_all at one point. It expects shape (9,) because nine harmonics span l=0, 1, 2.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None

test_matches_individual() → None

Verify batch results match individual real_spherical_harmonic calls.

The test evaluates both interfaces at theta=1.2 and phi=0.7. It compares each batch element with the corresponding scalar call. It iterates through l and m in the documented order. The values agree within 1e-12 and confirm the batch ordering.

Notes

The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions.

Return type:

None