"""Validate high-level workflow helpers in :mod:`diffpes.simul.workflow`.
Extended Summary
----------------
Validates VASP context loading, projection preparation, simulation dispatch,
and the combined file-to-spectrum workflow with controlled temporary inputs.
"""
from pathlib import Path
import chex
import jax
import jax.numpy as jnp
import pytest
from beartype.typing import Any, Callable
from jaxtyping import Array, Float
import diffpes
from diffpes.simul import (
load_vasp_context,
prepare_projection,
run_vasp_workflow,
simulate_context,
)
from diffpes.types import (
ArpesSpectrum,
SpinOrbitalProjection,
WorkflowContext,
make_band_structure,
make_orbital_projection,
make_spin_orbital_projection,
make_workflow_context,
)
from tests._gradients import gradient_gate
_FIXTURES_DIR: Path = (
Path(__file__).resolve().parents[1] / "test_inout" / "fixtures"
)
[docs]
class TestLoadVaspContextEdgeCases(chex.TestCase):
"""Validate additional paths in :func:`diffpes.simul.load_vasp_context`.
:see: :func:`~diffpes.simul.load_vasp_context`
"""
[docs]
def test_no_doscar_fermi_defaults_to_zero(self) -> None:
"""Verify Fermi energy is 0.0 when doscar_file=None and fermi_energy=None.
The test passes ``doscar_file=None`` and ``fermi_energy=None``, exercising
workflow.py line 142 (``resolved_fermi = 0.0``). Asserts the
returned band structure has fermi_energy == 0.0.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
context: diffpes.types.WorkflowContext
context = load_vasp_context(
directory=str(_FIXTURES_DIR),
eigenval_file="EIGENVAL_spin",
procar_file="PROCAR_spin",
doscar_file=None,
kpoints_file=None,
fermi_energy=None,
check_dimensions=True,
)
chex.assert_trees_all_close(
context.bands.fermi_energy, jnp.float64(0.0), atol=1e-12
)
assert context.dos is None
[docs]
def test_missing_doscar_raises(self) -> None:
"""Verify that a missing required DOSCAR raises FileNotFoundError.
The test passes a non-existent ``doscar_file`` with ``fermi_energy=None``,
exercising workflow.py lines 146-150 (FileNotFoundError path).
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
with pytest.raises(FileNotFoundError):
load_vasp_context(
directory=str(_FIXTURES_DIR),
eigenval_file="EIGENVAL_spin",
procar_file="PROCAR_spin",
doscar_file="DOES_NOT_EXIST",
kpoints_file=None,
fermi_energy=None,
check_dimensions=False,
)
[docs]
def test_explicit_fermi_reads_doscar_optionally(self) -> None:
"""Verify optional DOSCAR loading with an explicit Fermi energy.
The test passes ``fermi_energy=1.5`` and a valid ``doscar_file``, exercising
workflow.py lines 154-158 (optional DOSCAR read). The explicit
Fermi energy controls the bands, and the file supplies ``dos``.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
context: diffpes.types.WorkflowContext
context = load_vasp_context(
directory=str(_FIXTURES_DIR),
eigenval_file="EIGENVAL_spin",
procar_file="PROCAR_spin",
doscar_file="DOSCAR",
kpoints_file=None,
fermi_energy=1.5,
check_dimensions=True,
)
chex.assert_trees_all_close(
context.bands.fermi_energy, jnp.float64(1.5), atol=1e-12
)
assert context.dos is not None
[docs]
class TestLoadVaspContext(chex.TestCase):
"""Validate :func:`diffpes.simul.load_vasp_context`.
:see: :func:`~diffpes.simul.load_vasp_context`
"""
[docs]
def test_loads_context_with_optional_dos_and_kpath(self) -> None:
"""Verify context loading with inferred Fermi level and checks.
The test establishes the loads context with optional dos and kpath contract for
load vasp context with the concrete values and array shapes described below.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
context: diffpes.types.WorkflowContext
context = load_vasp_context(
directory=str(_FIXTURES_DIR),
eigenval_file="EIGENVAL_spin",
procar_file="PROCAR_spin",
doscar_file="DOSCAR",
kpoints_file="KPOINTS_line_fallback",
procar_mode="full",
check_dimensions=True,
)
chex.assert_shape(context.bands.eigenvalues, (2, 2))
chex.assert_shape(context.orb_proj.projections, (2, 2, 1, 9))
assert context.orb_proj.spin is not None
assert context.kpath is not None
assert context.dos is not None
chex.assert_trees_all_close(
context.bands.fermi_energy,
jnp.float64(0.5),
atol=1e-12,
)
[docs]
class TestPrepareProjection(chex.TestCase):
"""Validate :func:`diffpes.simul.prepare_projection`.
:see: :func:`~diffpes.simul.prepare_projection`
"""
[docs]
def test_spin_orbital_projection_attaches_oam(self) -> None:
"""Verify OAM attachment works for SpinOrbitalProjection input.
The test constructs a SpinOrbitalProjection and calls ``prepare_projection``
with ``attach_oam=True``. Asserts the returned object is still a
SpinOrbitalProjection with OAM attached, covering workflow.py
line 224 (make_spin_orbital_projection with oam).
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
projections: Array
spin: Array
orb: diffpes.types.SpinOrbitalProjection
prepared: (
diffpes.types.OrbitalProjection
| diffpes.types.SpinOrbitalProjection
)
projections = jnp.ones((2, 2, 2, 9), dtype=jnp.float64)
spin = jnp.ones((2, 2, 2, 6), dtype=jnp.float64)
orb = make_spin_orbital_projection(projections=projections, spin=spin)
prepared = prepare_projection(orb_proj=orb, attach_oam=True)
assert isinstance(prepared, SpinOrbitalProjection)
assert prepared.oam is not None
chex.assert_shape(prepared.oam, (2, 2, 2, 3))
[docs]
def test_selects_atoms_and_attaches_oam(self) -> None:
"""Verify atom sub-selection and OAM attachment in one call.
The test establishes the selects atoms and attaches oam contract for prepare
projection with the concrete values and array shapes described below.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
projections: Array
orb: diffpes.types.OrbitalProjection
prepared: (
diffpes.types.OrbitalProjection
| diffpes.types.SpinOrbitalProjection
)
projections = jnp.ones((2, 2, 3, 9), dtype=jnp.float64)
orb = make_orbital_projection(projections=projections)
prepared = prepare_projection(
orb_proj=orb,
atom_indices=[0, 2],
attach_oam=True,
)
chex.assert_shape(prepared.projections, (2, 2, 2, 9))
assert prepared.oam is not None
chex.assert_shape(prepared.oam, (2, 2, 2, 3))
[docs]
class TestSimulateContext(chex.TestCase):
"""Validate :func:`diffpes.simul.simulate_context`.
:see: :func:`~diffpes.simul.simulate_context`
"""
[docs]
def test_momentum_broadening_changes_output(self) -> None:
"""Verify nonzero dk changes simulated intensity.
The test establishes the momentum broadening changes output contract for
simulate context with the concrete values and array shapes described below.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
eigenbands: Array
kx: Array
kpoints: Array
projections: Array
bands: diffpes.types.BandStructure
orb: diffpes.types.OrbitalProjection
context: diffpes.types.WorkflowContext
base: diffpes.types.ArpesSpectrum
broadened: diffpes.types.ArpesSpectrum
nk: int = 12
nb: int = 4
na: int = 2
eigenbands = jnp.linspace(
-2.0, 0.6, nk * nb, dtype=jnp.float64
).reshape(nk, nb)
kx = jnp.linspace(0.0, 1.0, nk, dtype=jnp.float64)
kpoints = jnp.stack(
[kx, jnp.zeros_like(kx), jnp.zeros_like(kx)],
axis=1,
)
projections = jnp.ones((nk, nb, na, 9), dtype=jnp.float64) * 0.1
projections = projections.at[:, :, :, 4:9].set(0.3)
bands = make_band_structure(
eigenvalues=eigenbands,
kpoints=kpoints,
fermi_energy=0.0,
)
orb = make_orbital_projection(projections=projections)
context = make_workflow_context(
bands=bands,
orb_proj=orb,
kpath=None,
dos=None,
)
base = simulate_context(
context=context,
level="basic",
fidelity=320,
sigma=0.05,
temperature=20.0,
photon_energy=35.0,
normalize=False,
)
broadened = simulate_context(
context=context,
level="basic",
fidelity=320,
sigma=0.05,
temperature=20.0,
photon_energy=35.0,
dk=0.06,
normalize=False,
)
chex.assert_shape(base.intensity, (nk, 320))
chex.assert_shape(broadened.intensity, (nk, 320))
assert not jnp.allclose(base.intensity, broadened.intensity)
[docs]
def test_loaded_fermi_energy_gradient_matches_fd(self) -> None:
"""Propagate the loaded Fermi energy through the workflow gradient.
Extended Summary
----------------
The test traces an explicit scalar Fermi energy through
``load_vasp_context`` and ``simulate_context``. Its finite,
nonzero derivative matches central differences at ``rtol=1e-5``.
Notes
-----
The test loads two-k-point spin fixtures inside the traced loss.
It simulates the basic level at 300 K and checks the summed
intensity with the shared gradient harness.
"""
def workflow_loss(
fermi_energy: Float[Array, ""],
) -> Float[Array, ""]:
context: WorkflowContext = load_vasp_context(
directory=str(_FIXTURES_DIR),
eigenval_file="EIGENVAL_spin",
procar_file="PROCAR_spin",
doscar_file=None,
kpoints_file=None,
fermi_energy=fermi_energy,
procar_mode="legacy",
check_dimensions=True,
)
spectrum: ArpesSpectrum = simulate_context(
context=context,
level="basic",
fidelity=96,
sigma=0.05,
gamma=0.1,
temperature=300.0,
photon_energy=35.0,
)
loss: Float[Array, ""] = jnp.sum(spectrum.intensity)
return loss
fermi_energy: Float[Array, ""] = jnp.asarray(-0.4)
derivative: Float[Array, ""] = jax.grad(workflow_loss)(fermi_energy)
chex.assert_tree_all_finite(derivative)
assert float(jnp.abs(derivative)) > 1e-12
gradient_gate(
workflow_loss,
fermi_energy,
regime="stiff",
scale_floor=1.0,
)
[docs]
class TestRunVaspWorkflow(chex.TestCase):
"""Validate :func:`diffpes.simul.run_vasp_workflow`.
:see: :func:`~diffpes.simul.run_vasp_workflow`
"""
[docs]
def test_runs_end_to_end_with_normalization(self) -> None:
"""Verify one-call workflow runs and returns normalized spectrum.
The test establishes the runs end to end with normalization contract for run
vasp workflow with the concrete values and array shapes described below.
Notes
-----
The test builds the inputs in the test body and checks the stated property with the documented numerical or structural assertions."""
spectrum: diffpes.types.ArpesSpectrum
mean_val: Array
spectrum = run_vasp_workflow(
level="basic",
directory=str(_FIXTURES_DIR),
eigenval_file="EIGENVAL_spin",
procar_file="PROCAR",
doscar_file="DOSCAR",
kpoints_file="KPOINTS_line_fallback",
fidelity=180,
sigma=0.05,
temperature=15.0,
photon_energy=11.0,
normalize=True,
check_dimensions=True,
)
chex.assert_shape(spectrum.intensity, (2, 180))
chex.assert_shape(spectrum.energy_axis, (180,))
mean_val = jnp.mean(spectrum.intensity)
chex.assert_trees_all_close(mean_val, jnp.float64(0.0), atol=1e-8)