diffpes.types¶

PyTree-compatible data structures and factories for physical inputs, forward outputs, certification records, transformation contracts, and provenance.

Define types and factory functions for diffpes.

Extended Summary¶

This package provides PyTree-compatible data structures and their factory functions for ARPES simulation data. The data includes crystal geometry, band structures, orbital projections, simulation parameters, and polarization configurations. JAX stores fields that participate in autodiff as array children. It stores shape values, such as SimulationParams.fidelity, as auxiliary data. It also stores code-path selectors, such as PolarizationConfig.polarization_type, as auxiliary data. These values remain concrete during tracing. Changes to these values trigger recompilation.

The package contains these submodules:

aliases
Define scalar type aliases for JAX-compatible numeric types.
bands
Define band-structure and orbital-projection data structures.
constants
Define numerical, physical, orbital, and VASP-format constants for diffpes.
dos
Define density-of-states data structures.
experiment
Define the geometry of an ARPES experiment.
geometry
Define crystal-geometry data structures for VASP crystal structures.
kpath
Define k-space path and grid data structures.
params
Define simulation-parameter data structures.
tables
Store small immutable numerical tables for simulation routines.
radial_params
Define radial-wavefunction parameter structures.
self_energy
Define self-energy configuration data structures.
tb_model
Define tight-binding model and diagonalized-band data structures.
volumetric
Define volumetric data structures for VASP CHGCAR files.
context
Define structured inputs for high-level VASP simulation workflows.
certification
Store JAX-native carriers for certified forward-model executions.
certification_constants
Define static identifiers and schema constants for forward certification.
contracts
Define static carriers for certified transformation contracts.
inspection
Store types-owned records from certificate inspection.
provenance
Store types-owned carriers for artifact provenance and information flow.

Routine Listings¶

ArtifactResolver: Resolve an artifact to normalized content and optional source bytes.
ArtifactRef: Store static identity and role for one source or derived artifact.
CertificationClaim: Store a named claim and its continuous numerical evidence.
CertificationContext: Store prepared selections and references for compiled certification.
CertifiedResult: Store a numerical result paired with its differentiable certificate.
CertificateDiff: Store categorized differences between two forward certificates.
CheckFunction: Callable signature for a pure JAX certification check.
CANONICAL_ARRAY_CHUNK_BYTES: Array chunk size used by canonical PyTree encoding in bytes.
CANONICAL_JSON_PREFIX: Domain prefix for canonical JSON consistency checksums.
CANONICAL_JSON_VERSION: Version of the canonical JSON representation.
CANONICAL_PYTREE_PREFIX: Domain prefix for canonical PyTree consistency checksums.
CANONICAL_PYTREE_VERSION: Version of the canonical PyTree representation.
CANONICAL_SUPPORTED_ARRAY_KINDS: NumPy dtype kinds accepted by canonical array encoding.
CERTIFICATE_ARRAY_KINDS: NumPy dtype kinds accepted in persisted certificates.
CERTIFICATE_ARRAY_PREVIEW_ITEMS: Maximum array elements shown by certificate inspection.
CERTIFICATE_DOCUMENT_KEYS: Required top-level keys in a certificate document.
CERTIFICATE_FORMAT: Stable identifier for the forward-certificate document format.
CERTIFICATE_H5_GROUP: Reserved HDF5 group containing attached certificates.
CERTIFICATE_SCHEMA_MAJOR: Supported major version of the certificate schema.
CERTIFICATE_SCHEMA_MINOR: Supported minor version of the certificate schema.
CERTIFICATE_SCHEMA_PATTERN: Pattern matching supported certificate schema versions.
CERTIFICATION_IDENTIFIER_PATTERN: Pattern matching permanent certification identifiers.
CERTIFICATION_LEVEL_IDS: Ordered cumulative scientific-certification level identifiers.
CERTIFICATION_LEVEL_PREFIXES: Evidence prefixes required by each certification level.
CERTIFICATION_POLICY_IDS: Stable identifiers of built-in cumulative policies.
CERTIFICATION_POLICY_LEVEL_COUNT: Number of required levels for each built-in policy.
CERTIFICATION_SEMVER_PATTERN: Pattern matching certification semantic versions.
CHECKSUM_ALGORITHM: Name of the non-security consistency-checksum algorithm.
CHECKSUM_FILE_CHUNK_BYTES: File chunk size used by streaming consistency checksums in bytes.
CHECKSUM_FORMAT_VERSION: Version of the consistency-checksum text format.
CHECKSUM_PATTERN: Pattern matching formatted consistency-checksum records.
CHECKSUM_RECORD_KIND_PATTERN: Pattern matching consistency-checksum record-kind identifiers.
TB_RADIAL_INPUT_COUNT: Number of positional inputs accepted by the radial ARPES model adapter.
TB_RADIAL_MODEL_ID: Permanent identifier of the radial ARPES forward model.
TB_RADIAL_MODEL_VERSION: Semantic version of the radial ARPES forward model.
TWO_ME_OVER_HBAR_SQ_INV_EV_ANG2: Store the inverse free-electron dispersion constant.
CompositionReport: Store a conservative transformation-composition result.
ConventionRef: Store a versioned semantic convention used by a scientific model.
DependencyMap: Store declared and JAXPR-observed dependency relations.
DerivativeEvidence: Store JVP, VJP, reference, and information-spectrum evidence.
DomainPredicate: Store a static declaration of one model-domain predicate.
DomainResult: Store the traced evaluation of one declared domain predicate.
EvidenceRef: Store numerical evidence with static method and source identity.
EvidenceReport: Store the offline consistency outcome for one evidence record.
ExecutionManifest: Store software and execution identity prepared at the I/O boundary.
ExperimentGeometry: Store the geometry and resolution of an ARPES experiment.
ForwardCertificate: Store the complete assurance record for one forward execution.
ForwardModelSpec: Store the identity of a differentiable forward model.
HandshakeReport: Store the validation outcome for one registration handshake.
InformationSpectrum: Store a matrix-free information spectrum in input coordinates.
InformationState: Store effective semantic state for one artifact or result node.
PolicyReport: Store a traced policy truth table for derived certification levels.
ProvenanceGraph: Store a validated lineage graph and its propagated semantics.
ProvenanceReport: Store a structural and semantic provenance-validation report.
RegisteredModel: Store a frozen binding between a model spec and its executor.
RegisteredTransformation: Store a frozen transformation and its consistency checksum.
RegistrationHandshake: Store declarative registration requirements for one plan owner.
RegistryReport: Store the structural validation result for one registry snapshot.
RegistrySnapshot: Store an immutable deterministic snapshot of registry entries.
ReproductionReport: Store a numerical comparison from deliberate forward re-execution.
SensitivityMap: Store scaled sensitivities from inputs to output projections.
TransformationContract: Store the static semantic contract for one registered transformation.
TransformationRecord: Store one transformation and its semantic information effects.
VerificationReport: Store an offline certificate-verification outcome.
WaiverRecord: Store a bounded policy-waiver declaration without changing claim status.
WaiverReport: Store the temporal validation outcome for one waiver.
make_artifact_ref(): Create a validated artifact reference.
make_certification_claim(): Create a claim retaining both continuous and discrete evidence.
make_certification_context(): Create a prepared certification context.
make_certified_result(): Pair any JAX-compatible result value with a forward certificate.
make_certificate_diff(): Construct a validated certificate-difference record.
make_composition_report(): Create a validated immutable transformation-composition report.
make_convention_ref(): Create a validated convention reference.
make_dependency_map(): Create a structural dependency map.
make_derivative_evidence(): Create validated derivative and local-information evidence.
make_domain_predicate(): Create a validated domain-predicate declaration.
make_domain_result(): Create one traced domain evaluation.
make_evidence_ref(): Create validated vector-valued numerical evidence.
make_evidence_report(): Create an offline evidence-verification report.
make_execution_manifest(): Create a validated execution manifest.
make_forward_certificate(): Create and cross-validate a complete forward certificate.
make_forward_model_spec(): Create a validated stable forward-model specification.
make_information_spectrum(): Create a validated local information spectrum.
make_information_state(): Create a validated semantic-information state for one graph node.
make_handshake_report(): Create a report for one registration handshake.
make_policy_report(): Create a validated policy truth table.
make_provenance_graph(): Create a validated immutable provenance graph carrier.
make_provenance_report(): Create a validated structural and semantic provenance report.
make_registered_model(): Create a validated model-registry binding.
make_registered_transformation(): Create a validated transformation-registry binding.
make_registry_report(): Create a validated structural registry report.
make_registry_snapshot(): Create an immutable registry snapshot.
make_registration_handshake(): Create declarative registration requirements for one plan owner.
make_reproduction_report(): Create a report comparing a result with its re-execution.
make_sensitivity_map(): Create a named, scaled local-sensitivity map.
make_transformation_contract(): Create a validated immutable transformation contract.
make_transformation_record(): Create a validated information-aware transformation record.
make_verification_report(): Create an offline certificate-verification report.
make_waiver_record(): Create a bounded policy-waiver declaration.
make_waiver_report(): Create a temporal waiver-validation report.
ArpesSpectrum: Store ARPES simulation output in a JAX PyTree.
ATTR_AUX: HDF5 attribute name storing auxiliary PyTree data as JSON.
ATTR_NONE: HDF5 attribute name listing PyTree fields stored as None.
ATTR_TYPE: HDF5 attribute name storing the PyTree type name.
BAND_LINE_MIN_VALUES: Minimum tokens on an EIGENVAL band line.
BAND_LINE_SPIN_VALUES: Tokens on a spin-polarized EIGENVAL band line.
BAND_NDIM: Expected dimensionality of band-energy arrays.
BandStructure: Store electronic band-structure data in a JAX PyTree.
BOHR_TO_ANGSTROM: Bohr radius in Angstrom.
COORDINATE_MODE_TOKENS: Recognized KPOINTS coordinate-mode tokens.
CROSS_SECTION_ENERGIES: Photon-energy tabulation grid for the cross-section tables in eV.
CROSS_SECTION_SIGMA_D: Yeh-Lindau d-subshell cross sections on the tabulation grid.
CROSS_SECTION_SIGMA_P: Yeh-Lindau p-subshell cross sections on the tabulation grid.
CROSS_SECTION_SIGMA_S: Yeh-Lindau s-subshell cross sections on the tabulation grid.
CrystalGeometry: Store VASP POSCAR crystal geometry in a JAX PyTree.
D_ORBITAL_SLICE: Slice selecting the five d orbitals.
DensityOfStates: Store density-of-states data in a JAX PyTree.
DiagonalizedBands: Store diagonalized electronic-structure data in a JAX PyTree.
DosType: Supported density-of-states containers.
EIG_DOWN_INDEX: Column index of spin-down eigenvalues in EIGENVAL.
EIG_UP_INDEX: Column index of spin-up eigenvalues in EIGENVAL.
ENERGY_AXIS_NDIM: Expected dimensionality of energy-axis arrays.
EKIN_FLOOR_EV: Set the physical kinetic-energy floor in eV.
EPS: Epsilon floor guarding divisions and norms.
FLOAT_TOKEN_RE: Compiled regex matching floating-point tokens.
FullDensityOfStates: Store spin-resolved total and projected DOS data in a JAX PyTree.
GAUNT_IMAG_TOL: Tolerance for discarding imaginary Gaunt residues.
HBAR_C_EV_A: Reduced Planck constant times c in eV Angstrom.
HBAR_EV_S: Reduced Planck constant in eV s.
HBAR_SQ_OVER_2ME_EV_ANG2: Store the free-electron dispersion constant in eV Angstrom squared.
INTENSITY_NDIM: Expected dimensionality of intensity arrays.
ISPIN2_BLOCKS: PROCAR block count for ISPIN=2 calculations.
ISPIN_SPIN_POLARIZED: ISPIN value marking spin-polarized VASP runs.
KB_EV_PER_K: Boltzmann constant in eV per kelvin.
KPATH_AUX_WITH_COMMENT_LEN: KPathInfo auxiliary-data length including a comment.
KPATH_AUX_WITH_COORD_MODE_LEN: KPathInfo auxiliary-data length including a coordinate mode.
KGrid: Store a fixed-shape raster in fractional k-space.
KPath: Store a generated path through fractional k-space.
KPathInfo: Store k-point path metadata in a JAX PyTree.
KPOINT_LINE_VALUES: Tokens on an EIGENVAL k-point line.
K_PREFACTOR_INV_ANG_SQRT_EV: Store the momentum prefactor in inverse Angstrom per square-root eV.
L_MAX: Maximum angular momentum supported by the precomputed table.
LATTICE_ROWS: Number of lattice-vector rows in POSCAR/CHGCAR headers.
M_D: Magnetic quantum numbers of the d orbitals.
M_P: Magnetic quantum numbers of the p orbitals.
make_1d_chain_model(): Create a 1D chain tight-binding model.
make_arpes_spectrum(): Create a validated ArpesSpectrum instance.
make_band_structure(): Create a validated BandStructure instance.
make_crystal_geometry(): Create a validated CrystalGeometry instance.
make_density_of_states(): Create a validated DensityOfStates instance.
make_diagonalized_bands(): Create a validated DiagonalizedBands instance.
make_expanded_simulation_params(): Build simulation parameters with auto-derived energy window.
make_full_density_of_states(): Create a validated FullDensityOfStates instance.
make_graphene_model(): Create a graphene pz tight-binding model.
make_experiment_geometry(): Create a validated geometry for an ARPES experiment.
make_kgrid(): Create a validated fixed-shape k-space raster.
make_kpath(): Create a validated path through fractional k-space.
make_kpath_info(): Create a validated KPathInfo instance.
make_orbital_basis(): Create a validated OrbitalBasis instance.
make_orbital_projection(): Create a validated OrbitalProjection instance.
make_polarization_config(): Create a validated PolarizationConfig instance.
make_self_energy_config(): Create a validated SelfEnergyConfig instance.
make_simulation_params(): Create a validated SimulationParams instance.
make_slater_params(): Create a validated SlaterParams instance.
make_soc_volumetric_data(): Create a validated SOCVolumetricData instance.
make_spin_band_structure(): Create a validated SpinBandStructure instance.
make_spin_orbital_projection(): Create a validated SpinOrbitalProjection instance.
make_tb_model(): Create a validated TBModel instance.
make_volumetric_data(): Create a validated VolumetricData instance.
make_workflow_context(): Create a workflow context from parsed VASP inputs.
ME_EV: Electron rest energy in eV.
MIN_SUM: Minimum-sum floor guarding normalizations.
N_ORBITALS: Number of orbitals in the VASP projection basis.
N_SOC_MAG_BLOCKS: Magnetization block count in SOC CHGCAR files.
N_SPIN_COMPONENTS: Spin-projection component count in PROCAR.
N_TAYLOR: Taylor-series order for the Faddeeva evaluation.
NON_S_ORBITAL_SLICE: Slice selecting all non-s orbitals.
NonJaxNumber: Union of int, float, and complex.
NONSPIN_COLS: DOSCAR column count without spin polarization.
NORM_EPS: Epsilon floor guarding eigenvector normalization.
ORBITAL_DIRS_NORMALIZED: Unit-normalized orbital directions in VASP orbital order.
ORBITAL_INDEX: Mapping from orbital name to VASP orbital index.
OrbitalBasis: Store orbital quantum-number metadata in a JAX PyTree.
OrbitalProjection: Store orbital-resolved band projections in a JAX PyTree.
P_ORBITAL_SLICE: Slice selecting the three p orbitals.
PHASE_LOSS_MESSAGE: Warning text for PROCAR magnitude-only eigenvectors.
PolarizationConfig: Store photon-polarization geometry in a JAX PyTree.
PRESET_NAMES: Recognized band-scatter plotting preset names.
ProjectionType: Supported orbital-projection containers.
S_IDX: Index of the s orbital.
SCALAR_LINE_COMPONENTS: Tokens on a scalar CHGCAR header line.
ScalarBool: Union of bool and Bool[Array, " "].
ScalarComplex: Union of complex and Complex[Array, " "].
ScalarFloat: Union of float and Float[Array, " "].
ScalarInteger: Union of int and Int[Array, " "].
ScalarNumeric: Union of int, float, complex, and Num[Array, " "].
SelfEnergyConfig: Store energy-dependent self-energy data in a JAX PyTree.
SimulationParams: Store ARPES simulation parameters in a JAX PyTree.
SlaterParams: Store Slater radial-wavefunction parameters in a JAX PyTree.
SMALL_ARGUMENT: Small-argument cutoff for spherical Bessel seeds.
SOC_BLOCKS: PROCAR block count for SOC calculations.
SOCVolumetricData: Store SOC CHGCAR volumetric-grid data in a JAX PyTree.
SPIN_COLS: DOSCAR column count with spin polarization.
SpinBandStructure: Store spin-resolved electronic band-structure data in a JAX PyTree.
SpinOrbitalProjection: Store orbital projections with spin data in a JAX PyTree.
TBModel: Store tight-binding parameters in a JAX PyTree.
VolumetricData: Store CHGCAR volumetric-grid data in a JAX PyTree.
WEIGHT_COMPONENT_COUNT: Tokens on a weighted k-point line.
WEIGHT_COMPONENT_INDEX: Index of the weight token on a k-point line.
WorkflowContext: Store parsed VASP inputs for high-level workflow helpers.
XYZ_COMPONENTS: Number of Cartesian vector components.

Notes

All structured carriers are immutable equinox.Module PyTrees. Array fields remain differentiable leaves, while shape and control-flow metadata use equinox.field(static=True).

class diffpes.types.ArpesSpectrum(intensity: Float[Array, 'K E'], energy_axis: Float[Array, 'E']) → None[source]¶

Bases: Module

Store ARPES simulation output in a JAX PyTree.

This type stores an ARPES simulation result. The result contains a two-dimensional photoemission intensity map I(k, E) and its energy axis. The k-point dimension indexes the momentum-resolved detector channels, and the energy dimension indexes the binding-energy grid for the spectral function.

This JAX-compatible PyTree lets jit-compiled loss functions compare simulated spectra with experimental data. grad can differentiate these functions during parameter fitting. Both fields contain dense JAX arrays and no static auxiliary data. JAX can differentiate the complete object.

See:

TestArpesSpectrum

Variables:

intensity (Float[Array, "K E"]) – Photoemission intensity for K k-points and E energies.
energy_axis (Float[Array, " E"]) – Energy axis values in eV.

Notes

Implemented as an immutable equinox.Module PyTree. Equinox derives the tree structure from the annotated fields; both fields are differentiable leaves.

intensity: Float[Array, 'K E']¶

energy_axis: Float[Array, 'E']¶

__init__(intensity: Float[Array, 'K E'], energy_axis: Float[Array, 'E']) → None¶

class diffpes.types.BandStructure(eigenvalues: Float[Array, 'K B'], kpoints: Float[Array, 'K 3'], kpoint_weights: Float[Array, 'K'], fermi_energy: Float[Array, '']) → None[source]¶

Bases: Module

Store electronic band-structure data in a JAX PyTree.

This type stores the core outputs of a DFT band-structure calculation. The outputs include E_n(k), the reciprocal-space k-point mesh, the k-point integration weights, and the Fermi energy. These fields describe the single-particle electronic structure for ARPES simulations.

This JAX-compatible PyTree passes through jax.jit, jax.vmap, and jax.grad without manual flattening. All four fields contain JAX-traced arrays and no static auxiliary data. JAX can differentiate the object with respect to each field.

See:

TestBandStructure

Variables:

eigenvalues (Float[Array, "K B"]) – Band energies in eV for K k-points and B bands.
kpoints (Float[Array, "K 3"]) – k-point coordinates in reciprocal space.
kpoint_weights (Float[Array, " K"]) – Integration weights for each k-point.
fermi_energy (Float[Array, " "]) – Fermi level energy in eV.

Notes

Implemented as an immutable equinox.Module PyTree. Equinox derives the tree structure from the annotated fields; all fields are differentiable leaves and no static metadata is present.

eigenvalues: Float[Array, 'K B']¶

kpoints: Float[Array, 'K 3']¶

kpoint_weights: Float[Array, 'K']¶

fermi_energy: Float[Array, '']¶

__init__(eigenvalues: Float[Array, 'K B'], kpoints: Float[Array, 'K 3'], kpoint_weights: Float[Array, 'K'], fermi_energy: Float[Array, '']) → None¶

class diffpes.types.OrbitalProjection(projections: Float[Array, 'K B A 9'], spin: Float[Array, 'K B A 6'] | None, oam: Float[Array, 'K B A 3'] | None) → None[source]¶

Bases: Module

Store orbital-resolved band projections in a JAX PyTree.

All array fields are differentiable PyTree leaves. Optional fields are empty subtrees when None; changing their presence changes the tree structure and may trigger recompilation.

See:

TestOrbitalProjection

Variables:

projections (Float[Array, "K B A 9"]) – Orbital projection weights.
spin (Optional[Float[Array, "K B A 6"]]) – Optional spin projections.
oam (Optional[Float[Array, "K B A 3"]]) – Optional orbital-angular-momentum projections.

projections: Float[Array, 'K B A 9']¶

spin: Float[Array, 'K B A 6'] | None¶

oam: Float[Array, 'K B A 3'] | None¶

__init__(projections: Float[Array, 'K B A 9'], spin: Float[Array, 'K B A 6'] | None, oam: Float[Array, 'K B A 3'] | None) → None¶

class diffpes.types.SpinBandStructure(eigenvalues_up: Float[Array, 'K B'], eigenvalues_down: Float[Array, 'K B'], kpoints: Float[Array, 'K 3'], kpoint_weights: Float[Array, 'K'], fermi_energy: Float[Array, '']) → None[source]¶

Bases: Module

Store spin-resolved electronic band-structure data in a JAX PyTree.

This type stores eigenvalues for both spin channels from an ISPIN=2 VASP calculation. The two spin channels share the same k-point mesh and weights. read_eigenval returns this type when return_mode="full" and the EIGENVAL file contains spin-polarized data.

This class is an immutable equinox.Module PyTree. JAX stores all five dense array fields as children and uses no auxiliary data. JAX can differentiate the complete object with respect to each field.

See:

TestSpinBandStructure

Variables:

eigenvalues_up (Float[Array, "K B"]) – Spin-up (majority) band energies in eV for K k-points and B bands. JAX-traced (differentiable).
eigenvalues_down (Float[Array, "K B"]) – Spin-down (minority) band energies in eV for K k-points and B bands. JAX-traced (differentiable).
kpoints (Float[Array, "K 3"]) – k-point coordinates in reciprocal (fractional) space, shared by both spin channels. JAX-traced (differentiable).
kpoint_weights (Float[Array, " K"]) – Integration weights for each k-point, used for Brillouin-zone averaging. Uniform weights (all ones) are the norm for band structure paths. JAX-traced (differentiable).
fermi_energy (Float[Array, " "]) – Fermi level energy in eV. A 0-D scalar array. JAX-traced (differentiable).

Notes

Implemented as an immutable equinox.Module PyTree. Equinox derives the tree structure from the annotated fields; all fields are differentiable leaves and no static metadata is present.

See also

BandStructure: Single-spin-channel variant.
make_spin_band_structure: Factory function with validation and float64 casting.

eigenvalues_up: Float[Array, 'K B']¶

eigenvalues_down: Float[Array, 'K B']¶

kpoints: Float[Array, 'K 3']¶

kpoint_weights: Float[Array, 'K']¶

fermi_energy: Float[Array, '']¶

__init__(eigenvalues_up: Float[Array, 'K B'], eigenvalues_down: Float[Array, 'K B'], kpoints: Float[Array, 'K 3'], kpoint_weights: Float[Array, 'K'], fermi_energy: Float[Array, '']) → None¶

class diffpes.types.SpinOrbitalProjection(projections: Float[Array, 'K B A 9'], spin: Float[Array, 'K B A 6'], oam: Float[Array, 'K B A 3'] | None) → None[source]¶

Bases: Module

Store orbital projections with spin data in a JAX PyTree.

All present arrays are differentiable PyTree leaves. oam=None is an empty subtree; changing its presence changes the tree structure.

See:

TestSpinOrbitalProjection

Variables:

projections (Float[Array, "K B A 9"]) – Orbital projection weights.
spin (Float[Array, "K B A 6"]) – Mandatory spin projections.
oam (Optional[Float[Array, "K B A 3"]]) – Optional orbital-angular-momentum projections.

projections: Float[Array, 'K B A 9']¶

spin: Float[Array, 'K B A 6']¶

oam: Float[Array, 'K B A 3'] | None¶

__init__(projections: Float[Array, 'K B A 9'], spin: Float[Array, 'K B A 6'], oam: Float[Array, 'K B A 3'] | None) → None¶

diffpes.types.make_arpes_spectrum(intensity: Float[Array, 'K Ei'], energy_axis: Float[Array, 'Ea']) → ArpesSpectrum[source]¶

Create a validated ArpesSpectrum instance.

The factory validates and normalizes simulated ARPES data before it constructs an ArpesSpectrum PyTree. The factory casts both input arrays to float64. Downstream loss functions therefore maintain double-precision accuracy.

@jaxtyped(typechecker=beartype) checks the energy dimension E at call time. This dimension must agree between intensity and energy_axis.

See:: TestMakeArpesSpectrum

Implementation Logic¶

Prepare the normalized values:
```
intensity_arr = jnp.asarray(intensity, dtype=jnp.float64)
```
This expression gives the later validation steps a stable shape and dtype.
Apply static validation:
```
intensity_arr.shape[1] != energy_arr.shape[0]
```
This predicate rejects invalid structure before JAX traces the numerical checks.
Apply traced validation:
```
~jnp.all(jnp.isfinite(intensity_arr))
```
This predicate remains active during eager and compiled execution.
Return the named instance:
```
return spectrum
```
The explicit name keeps the implementation and the Returns section synchronized.

param intensity:: Photoemission intensity map.
type intensity:: Float[Array, 'K Ei']
param energy_axis:: Energy axis values in eV.
type energy_axis:: Float[Array, 'Ea']
returns:: spectrum – Validated ARPES spectrum instance with all arrays in float64.
rtype:: ArpesSpectrum
raises ValueError:: If the intensity energy-axis length does not match energy_axis.
raises EquinoxRuntimeError:: If intensity is non-finite or the energy axis is not strictly increasing.

Notes

Static validation raises ValueError before traced construction when the two energy-axis lengths disagree. Traced validation uses eqx.error_if and raises EquinoxRuntimeError for non-finite intensity or an energy axis that is not strictly increasing.

diffpes.types.make_band_structure(eigenvalues: Float[Array, 'Ke B'], kpoints: Float[Array, 'Kk 3'], kpoint_weights: Float[Array, 'Kw'] | None = None, fermi_energy: int | float | complex | Num[Array, ''] = 0.0) → BandStructure[source]¶

Create a validated BandStructure instance.

The factory validates and normalizes raw band-structure data before it constructs a BandStructure PyTree. The factory casts all input arrays to float64 for numerical stability. Energy differences and Lorentzian broadening depend on precision. The factory replaces missing k-point weights with uniform weights. Callers therefore do not handle the common equal-weight case explicitly.

@jaxtyped(typechecker=beartype) checks input shapes and dtypes at call time. This check finds different dimensions before the simulation uses them.

See:: TestMakeBandStructure

Implementation Logic¶

Prepare the normalized values:
```
eigenvalues_arr = jnp.asarray(eigenvalues, dtype=jnp.float64)
```
This expression gives the later validation steps a stable shape and dtype.
Apply static validation:
```
kpoints_arr.shape[0] != nkpts
```
This predicate rejects invalid structure before JAX traces the numerical checks.
Apply traced validation:
```
~jnp.all(jnp.isfinite(eigenvalues_arr))
```
This predicate remains active during eager and compiled execution.
Return the named instance:
```
return bands
```
The explicit name keeps the implementation and the Returns section synchronized.

param eigenvalues:: Band energies in eV for K k-points and B bands.
type eigenvalues:: Float[Array, 'Ke B']
param kpoints:: k-point coordinates in reciprocal space.
type kpoints:: Float[Array, 'Kk 3']
param kpoint_weights:: Integration weights. Defaults to uniform weights.
type kpoint_weights:: Optional[Float[Array, 'Kw']], default: None
param fermi_energy:: Fermi level in eV. Default is 0.0.
type fermi_energy:: Union[int, float, complex, Num[Array, '']], default: 0.0
returns:: bands – Validated band structure instance with all arrays in float64.
rtype:: BandStructure
raises ValueError:: If eigenvalues, k-points, and weights disagree on their k-point count.
raises EquinoxRuntimeError:: If eigenvalues or k-points are non-finite, or weights are non-finite or negative.

Notes

Static validation raises ValueError before traced construction when k-point or weight counts disagree with the eigenvalues. Traced validation uses eqx.error_if and raises EquinoxRuntimeError for non-finite arrays or negative weights.

diffpes.types.make_orbital_projection(projections: Float[Array, 'Kp Bp Ap Op'], spin: Float[Array, 'Ks Bs As 6'] | None = None, oam: Float[Array, 'Ko Bo Ao 3'] | None = None) → OrbitalProjection[source]¶

Create a validated OrbitalProjection instance.

The factory validates and normalizes raw orbital projection data before constructing an OrbitalProjection PyTree. The factory casts the mandatory projections array to float64. It casts the optional spin and oam arrays only when they are present. Thus, None continues to identify calculations without spin-orbit coupling.

@jaxtyped(typechecker=beartype) checks the shape constraints at call time. The K, B, and A dimensions must agree across all arrays.

See:: TestMakeOrbitalProjection

Implementation Logic¶

Prepare the normalized values:
```
proj_arr = jnp.asarray(projections, dtype=jnp.float64)
```
This expression gives the later validation steps a stable shape and dtype.
Apply static validation:
```
proj_arr.shape[3] != N_ORBITALS
```
This predicate rejects invalid structure before JAX traces the numerical checks.
Apply traced validation:
```
~jnp.all(jnp.isfinite(proj_arr))
```
This predicate remains active during eager and compiled execution.
Return the named instance:
```
return orb_proj
```
The explicit name keeps the implementation and the Returns section synchronized.

param projections:: Orbital projection weights.
type projections:: Float[Array, 'Kp Bp Ap Op']
param spin:: Spin projections. Default is None.
type spin:: Optional[Float[Array, 'Ks Bs As 6']], default: None
param oam:: Orbital angular momentum. Default is None.
type oam:: Optional[Float[Array, 'Ko Bo Ao 3']], default: None
returns:: orb_proj – Validated orbital projection instance with all non-None arrays in float64.
rtype:: OrbitalProjection
raises ValueError:: If optional channel axes disagree with the projection axes or the projection orbital axis does not contain 9 columns.
raises EquinoxRuntimeError:: If projections are non-finite or negative, or a present spin channel is non-finite.

Notes

Static validation raises ValueError before traced construction when the projection, spin, or OAM shapes violate their structural contract. Traced validation uses eqx.error_if and raises EquinoxRuntimeError for non-finite arrays or negative projections.

diffpes.types.make_spin_band_structure(eigenvalues_up: Float[Array, 'Ku Bu'], eigenvalues_down: Float[Array, 'Kd Bd'], kpoints: Float[Array, 'Kk 3'], kpoint_weights: Float[Array, 'Kw'] | None = None, fermi_energy: int | float | complex | Num[Array, ''] = 0.0) → SpinBandStructure[source]¶

Create a validated SpinBandStructure instance.

The factory validates and normalizes raw spin-resolved band structure data before constructing a SpinBandStructure PyTree. This is the spin-polarized (ISPIN=2) counterpart to make_band_structure(). The factory casts all input arrays to float64 for numerical stability. It replaces missing k-point weights with uniform weights. Callers therefore do not handle the common equal-weight case explicitly.

@jaxtyped(typechecker=beartype) checks the input shapes and dtypes at call time. This check finds different K or B dimensions before the simulation uses them.

See:: TestMakeSpinBandStructure

Implementation Logic¶

Prepare the normalized values:
```
up_arr = jnp.asarray(eigenvalues_up, dtype=jnp.float64)
```
This expression gives the later validation steps a stable shape and dtype.
Apply static validation:
```
up_arr.shape != down_arr.shape
```
This predicate rejects invalid structure before JAX traces the numerical checks.
Apply traced validation:
```
~jnp.all(jnp.isfinite(up_arr))
```
This predicate remains active during eager and compiled execution.
Return the named instance:
```
return bands
```
The explicit name keeps the implementation and the Returns section synchronized.

param eigenvalues_up:: Spin-up band energies in eV for K k-points and B bands.
type eigenvalues_up:: Float[Array, 'Ku Bu']
param eigenvalues_down:: Spin-down band energies in eV. Must share the same (K, B) shape as eigenvalues_up.
type eigenvalues_down:: Float[Array, 'Kd Bd']
param kpoints:: k-point coordinates in reciprocal (fractional) space.
type kpoints:: Float[Array, 'Kk 3']
param kpoint_weights:: Integration weights per k-point. Defaults to uniform weights jnp.ones(K).
type kpoint_weights:: Optional[Float[Array, 'Kw']], default: None
param fermi_energy:: Fermi level in eV. Default is 0.0.
type fermi_energy:: Union[int, float, complex, Num[Array, '']], default: 0.0
returns:: bands – Validated spin-resolved band structure with all arrays in float64.
rtype:: SpinBandStructure
raises ValueError:: If the spin channels disagree on their k-point or band counts, or the k-point and weight counts disagree with the eigenvalues.
raises EquinoxRuntimeError:: If eigenvalues or k-points are non-finite, or weights are non-finite or negative.

Notes

Static validation raises ValueError before traced construction when the spin, k-point, band, or weight dimensions disagree. Traced validation uses eqx.error_if and raises EquinoxRuntimeError for non-finite arrays or negative weights.

See also

make_band_structure: Factory for single-spin-channel data.
SpinBandStructure: The PyTree class constructed by this factory.

diffpes.types.make_spin_orbital_projection(projections: Float[Array, 'Kp Bp Ap Op'], spin: Float[Array, 'Ks Bs As Ss'], oam: Float[Array, 'Ko Bo Ao 3'] | None = None) → SpinOrbitalProjection[source]¶

Create a validated SpinOrbitalProjection instance.

The factory validates and normalizes orbital projection data with mandatory spin before constructing a SpinOrbitalProjection PyTree. This factory supports spin-orbit coupling. Unlike make_orbital_projection(), it requires the spin field. Therefore, downstream SOC simulation kernels receive complete spin data without runtime checks.

The factory casts all present arrays to float64 for numerical stability. @jaxtyped(typechecker=beartype) checks the shape constraints at call time. The K, B, and A dimensions must agree across all arrays.

See:: TestMakeSpinOrbitalProjection

Implementation Logic¶

Prepare the normalized values:
```
proj_arr = jnp.asarray(projections, dtype=jnp.float64)
```
This expression gives the later validation steps a stable shape and dtype.
Apply static validation:
```
proj_arr.shape[:3] != spin_arr.shape[:3]
```
This predicate rejects invalid structure before JAX traces the numerical checks.
Apply traced validation:
```
~jnp.all(jnp.isfinite(proj_arr))
```
This predicate remains active during eager and compiled execution.
Return the named instance:
```
return soc_proj
```
The explicit name keeps the implementation and the Returns section synchronized.

param projections:: Orbital projection weights |<psi|Y_{lm}>|^2 following VASP ordering. Must share the K, B, A dimensions with spin.
type projections:: Float[Array, 'Kp Bp Ap Op']
param spin:: Spin projections [Sx_up, Sx_dn, Sy_up, Sy_dn, Sz_up, Sz_dn]. Required (non-optional).
type spin:: Float[Array, 'Ks Bs As Ss']
param oam:: Orbital angular momentum [L_p, L_d, L_total]. Default is None.
type oam:: Optional[Float[Array, 'Ko Bo Ao 3']], default: None
returns:: soc_proj – Validated instance with all non-None arrays in float64.
rtype:: SpinOrbitalProjection
raises ValueError:: If the projection, spin, and optional OAM axes disagree. The function also rejects an orbital axis without 9 columns or a spin axis without 6 columns.
raises EquinoxRuntimeError:: If projection values are non-finite or negative, or spin values are non-finite.

Notes

See also

make_orbital_projection: Factory for the optional-spin variant.
SpinOrbitalProjection: The PyTree class constructed by this factory.

class diffpes.types.ArtifactRef(artifact_id: str, media_type: str, byte_checksum: str | None, content_checksum: str, semantic_checksum: str, locator: str | None, role: str) → None[source]¶

Bases: Module

Store static identity and role for one source or derived artifact.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestArtifactref

Variables:

artifact_id (str) – Artifact id (static – a compile-time constant; changing it triggers retracing).
media_type (str) – Media type (static – a compile-time constant; changing it triggers retracing).
byte_checksum (Optional[str]) – Byte checksum (static – a compile-time constant; changing it triggers retracing).
content_checksum (str) – Content checksum (static – a compile-time constant; changing it triggers retracing).
semantic_checksum (str) – Semantic checksum (static – a compile-time constant; changing it triggers retracing).
locator (Optional[str]) – Locator (static – a compile-time constant; changing it triggers retracing).
role (str) – Role (static – a compile-time constant; changing it triggers retracing).

artifact_id: str¶

media_type: str¶

byte_checksum: str | None¶

content_checksum: str¶

semantic_checksum: str¶

locator: str | None¶

role: str¶

__init__(artifact_id: str, media_type: str, byte_checksum: str | None, content_checksum: str, semantic_checksum: str, locator: str | None, role: str) → None¶

type diffpes.types.ArtifactResolver = Callable[[ArtifactRef], tuple[Any, bytes | None]]¶

class diffpes.types.CertificationClaim(claim_id: str, subject_id: str, predicate_id: str, evidence_ids: tuple[str, ...], measured: Float[Array, 'n_measure'], reference: Float[Array, 'n_measure'], residual: Float[Array, 'n_measure'], tolerance: Float[Array, 'n_measure'], passed: Bool[Array, ''], checked: Bool[Array, ''], in_domain: Bool[Array, ''], margin: Float[Array, ''], severity_code: Int[Array, '']) → None[source]¶

Bases: Module

Store a named claim and its continuous numerical evidence.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestCertificationclaim

Variables:

claim_id (str) – Claim id (static – a compile-time constant; changing it triggers retracing).
subject_id (str) – Subject id (static – a compile-time constant; changing it triggers retracing).
predicate_id (str) – Predicate id (static – a compile-time constant; changing it triggers retracing).
evidence_ids (tuple[str, ]) – Evidence ids (static – a compile-time constant; changing it triggers retracing).
measured (Float[Array, " n_measure"]) – Measured retained as a differentiable JAX leaf in the declared physical units.
reference (Float[Array, " n_measure"]) – Reference retained as a differentiable JAX leaf in the declared physical units.
residual (Float[Array, " n_measure"]) – Residual retained as a differentiable JAX leaf in the declared physical units.
tolerance (Float[Array, " n_measure"]) – Tolerance retained as a differentiable JAX leaf in the declared physical units.
passed (Bool[Array, ""]) – Passed retained as a differentiable JAX leaf in the declared physical units.
checked (Bool[Array, ""]) – Checked retained as a differentiable JAX leaf in the declared physical units.
in_domain (Bool[Array, ""]) – In domain retained as a differentiable JAX leaf in the declared physical units.
margin (Float[Array, ""]) – Margin retained as a differentiable JAX leaf in the declared physical units.
severity_code (Int[Array, ""]) – Severity code retained as a differentiable JAX leaf in the declared physical units.

claim_id: str¶

subject_id: str¶

predicate_id: str¶

evidence_ids: tuple[str, ...]¶

measured: Float[Array, 'n_measure']¶

reference: Float[Array, 'n_measure']¶

residual: Float[Array, 'n_measure']¶

tolerance: Float[Array, 'n_measure']¶

passed: Bool[Array, '']¶

checked: Bool[Array, '']¶

in_domain: Bool[Array, '']¶

margin: Float[Array, '']¶

severity_code: Int[Array, '']¶

__init__(claim_id: str, subject_id: str, predicate_id: str, evidence_ids: tuple[str, ...], measured: Float[Array, 'n_measure'], reference: Float[Array, 'n_measure'], residual: Float[Array, 'n_measure'], tolerance: Float[Array, 'n_measure'], passed: Bool[Array, ''], checked: Bool[Array, ''], in_domain: Bool[Array, ''], margin: Float[Array, ''], severity_code: Int[Array, '']) → None¶

class diffpes.types.CertificationContext(manifest: ExecutionManifest, model: ForwardModelSpec, artifacts: tuple[ArtifactRef, ...], transformations: tuple[TransformationRecord, ...], evidence: tuple[EvidenceRef, ...], policy_id: str, check_ids: tuple[str, ...], input_checksums: tuple[str, ...], waivers: tuple[WaiverRecord, ...]) → None[source]¶

Bases: Module

Store prepared selections and references for compiled certification.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestCertificationcontext

Variables:

manifest (ExecutionManifest) – Manifest retained as a differentiable JAX leaf in the declared physical units.
model (ForwardModelSpec) – Model retained as a differentiable JAX leaf in the declared physical units.
artifacts (tuple[ArtifactRef, ]) – Artifacts retained as a differentiable JAX leaf in the declared physical units.
transformations (tuple[TransformationRecord, ]) – Transformations retained as a differentiable JAX leaf in the declared physical units.
evidence (tuple[EvidenceRef, ]) – Evidence retained as a differentiable JAX leaf in the declared physical units.
policy_id (str) – Policy id (static – a compile-time constant; changing it triggers retracing).
check_ids (tuple[str, ]) – Check ids (static – a compile-time constant; changing it triggers retracing).
input_checksums (tuple[str, ]) – Input checksums (static – a compile-time constant; changing it triggers retracing).
waivers (tuple[WaiverRecord, ]) – Policy-waiver records (static; changing them causes retracing).

manifest: ExecutionManifest¶

model: ForwardModelSpec¶

artifacts: tuple[ArtifactRef, ...]¶

transformations: tuple[TransformationRecord, ...]¶

evidence: tuple[EvidenceRef, ...]¶

policy_id: str¶

check_ids: tuple[str, ...]¶

input_checksums: tuple[str, ...]¶

waivers: tuple[WaiverRecord, ...]¶

__init__(manifest: ExecutionManifest, model: ForwardModelSpec, artifacts: tuple[ArtifactRef, ...], transformations: tuple[TransformationRecord, ...], evidence: tuple[EvidenceRef, ...], policy_id: str, check_ids: tuple[str, ...], input_checksums: tuple[str, ...], waivers: tuple[WaiverRecord, ...]) → None¶

class diffpes.types.CertifiedResult(value: Any, certificate: ForwardCertificate) → None[source]¶

Bases: Module

Store a numerical result paired with its differentiable certificate.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestCertifiedresult

Variables:

value (Any) – Value retained as a differentiable JAX leaf in the declared physical units.
certificate (ForwardCertificate) – Certificate retained as a differentiable JAX leaf in the declared physical units.

value: Any¶

certificate: ForwardCertificate¶

__init__(value: Any, certificate: ForwardCertificate) → None¶

type diffpes.types.CheckFunction = Callable[[PyTree], DomainResult]¶

class diffpes.types.ConventionRef(convention_id: str, version: str, parameters_json: str) → None[source]¶

Bases: Module

Store a versioned semantic convention used by a scientific model.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestConventionref

Variables:

convention_id (str) – Convention id (static – a compile-time constant; changing it triggers retracing).
version (str) – Version (static – a compile-time constant; changing it triggers retracing).
parameters_json (str) – Parameters json (static – a compile-time constant; changing it triggers retracing).

convention_id: str¶

version: str¶

parameters_json: str¶

__init__(convention_id: str, version: str, parameters_json: str) → None¶

class diffpes.types.DependencyMap(model_id: str, input_paths: tuple[str, ...], output_paths: tuple[str, ...], structural: Bool[Array, 'n_output n_input'], traced: Bool[Array, 'n_output n_input']) → None[source]¶

Bases: Module

Store declared and JAXPR-observed dependency relations.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestDependencymap

Variables:

model_id (str) – Model id (static – a compile-time constant; changing it triggers retracing).
input_paths (tuple[str, ]) – Input paths (static – a compile-time constant; changing it triggers retracing).
output_paths (tuple[str, ]) – Output paths (static – a compile-time constant; changing it triggers retracing).
structural (Bool[Array, "n_output n_input"]) – Structural retained as a differentiable JAX leaf in the declared physical units.
traced (Bool[Array, "n_output n_input"]) – Traced retained as a differentiable JAX leaf in the declared physical units.

model_id: str¶

input_paths: tuple[str, ...]¶

output_paths: tuple[str, ...]¶

structural: Bool[Array, 'n_output n_input']¶

traced: Bool[Array, 'n_output n_input']¶

__init__(model_id: str, input_paths: tuple[str, ...], output_paths: tuple[str, ...], structural: Bool[Array, 'n_output n_input'], traced: Bool[Array, 'n_output n_input']) → None¶

class diffpes.types.DerivativeEvidence(input_paths: tuple[str, ...], output_projection_ids: tuple[str, ...], method: str, scales: Float[Array, 'n_input'], jvp_probes: Float[Array, 'n_probe n_output'], vjp_probes: Float[Array, 'n_probe n_input'], reference_derivatives: Float[Array, 'n_probe n_deriv'], derivative_residuals: Float[Array, 'n_probe n_deriv'], singular_values: Float[Array, 'n_sv'], effective_rank: Int[Array, ''], condition_estimate: Float[Array, ''], finite: Bool[Array, ''], fd_correct: Bool[Array, '']) → None[source]¶

Bases: Module

Store JVP, VJP, reference, and information-spectrum evidence.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestDerivativeevidence

Variables:

input_paths (tuple[str, ]) – Input paths (static – a compile-time constant; changing it triggers retracing).
output_projection_ids (tuple[str, ]) – Output projection ids (static – a compile-time constant; changing it triggers retracing).
method (str) – Method (static – a compile-time constant; changing it triggers retracing).
scales (Float[Array, " n_input"]) – Scales retained as a differentiable JAX leaf in the declared physical units.
jvp_probes (Float[Array, "n_probe n_output"]) – Jvp probes retained as a differentiable JAX leaf in the declared physical units.
vjp_probes (Float[Array, "n_probe n_input"]) – Vjp probes retained as a differentiable JAX leaf in the declared physical units.
reference_derivatives (Float[Array, "n_probe n_deriv"]) – Reference derivatives retained as a differentiable JAX leaf in the declared physical units.
derivative_residuals (Float[Array, "n_probe n_deriv"]) – Derivative residuals retained as a differentiable JAX leaf in the declared physical units.
singular_values (Float[Array, " n_sv"]) – Singular values retained as a differentiable JAX leaf in the declared physical units.
effective_rank (Int[Array, ""]) – Effective rank retained as a differentiable JAX leaf in the declared physical units.
condition_estimate (Float[Array, ""]) – Condition estimate retained as a differentiable JAX leaf in the declared physical units. Zero means that there is no active information direction.
finite (Bool[Array, ""]) – Finite retained as a differentiable JAX leaf in the declared physical units.
fd_correct (Bool[Array, ""]) – Fd correct retained as a differentiable JAX leaf in the declared physical units.

input_paths: tuple[str, ...]¶

output_projection_ids: tuple[str, ...]¶

method: str¶

scales: Float[Array, 'n_input']¶

jvp_probes: Float[Array, 'n_probe n_output']¶

vjp_probes: Float[Array, 'n_probe n_input']¶

reference_derivatives: Float[Array, 'n_probe n_deriv']¶

derivative_residuals: Float[Array, 'n_probe n_deriv']¶

singular_values: Float[Array, 'n_sv']¶

effective_rank: Int[Array, '']¶

condition_estimate: Float[Array, '']¶

finite: Bool[Array, '']¶

fd_correct: Bool[Array, '']¶

__init__(input_paths: tuple[str, ...], output_projection_ids: tuple[str, ...], method: str, scales: Float[Array, 'n_input'], jvp_probes: Float[Array, 'n_probe n_output'], vjp_probes: Float[Array, 'n_probe n_input'], reference_derivatives: Float[Array, 'n_probe n_deriv'], derivative_residuals: Float[Array, 'n_probe n_deriv'], singular_values: Float[Array, 'n_sv'], effective_rank: Int[Array, ''], condition_estimate: Float[Array, ''], finite: Bool[Array, ''], fd_correct: Bool[Array, '']) → None¶

class diffpes.types.DomainPredicate(predicate_id: str, expression_id: str, units: str | None, severity: str) → None[source]¶

Bases: Module

Store a static declaration of one model-domain predicate.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestDomainpredicate

Variables:

predicate_id (str) – Predicate id (static – a compile-time constant; changing it triggers retracing).
expression_id (str) – Expression id (static – a compile-time constant; changing it triggers retracing).
units (Optional[str]) – Units (static – a compile-time constant; changing it triggers retracing).
severity (str) – Severity (static – a compile-time constant; changing it triggers retracing).

predicate_id: str¶

expression_id: str¶

units: str | None¶

severity: str¶

__init__(predicate_id: str, expression_id: str, units: str | None, severity: str) → None¶

class diffpes.types.DomainResult(predicate_id: str, measured: Float[Array, ''], reference: Float[Array, ''], residual: Float[Array, ''], tolerance: Float[Array, ''], margin: Float[Array, ''], passed: Bool[Array, ''], checked: Bool[Array, ''], in_domain: Bool[Array, ''], severity_code: Int[Array, '']) → None[source]¶

Bases: Module

Store the traced evaluation of one declared domain predicate.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestDomainresult

Variables:

predicate_id (str) – Predicate id (static – a compile-time constant; changing it triggers retracing).
measured (Float[Array, ""]) – Measured retained as a differentiable JAX leaf in the declared physical units.
reference (Float[Array, ""]) – Reference retained as a differentiable JAX leaf in the declared physical units.
residual (Float[Array, ""]) – Residual retained as a differentiable JAX leaf in the declared physical units.
tolerance (Float[Array, ""]) – Tolerance retained as a differentiable JAX leaf in the declared physical units.
margin (Float[Array, ""]) – Margin retained as a differentiable JAX leaf in the declared physical units.
passed (Bool[Array, ""]) – Passed retained as a differentiable JAX leaf in the declared physical units.
checked (Bool[Array, ""]) – Checked retained as a differentiable JAX leaf in the declared physical units.
in_domain (Bool[Array, ""]) – In domain retained as a differentiable JAX leaf in the declared physical units.
severity_code (Int[Array, ""]) – Severity code retained as a differentiable JAX leaf in the declared physical units.

predicate_id: str¶

measured: Float[Array, '']¶

reference: Float[Array, '']¶

residual: Float[Array, '']¶

tolerance: Float[Array, '']¶

margin: Float[Array, '']¶

passed: Bool[Array, '']¶

checked: Bool[Array, '']¶

in_domain: Bool[Array, '']¶

severity_code: Int[Array, '']¶

__init__(predicate_id: str, measured: Float[Array, ''], reference: Float[Array, ''], residual: Float[Array, ''], tolerance: Float[Array, ''], margin: Float[Array, ''], passed: Bool[Array, ''], checked: Bool[Array, ''], in_domain: Bool[Array, ''], severity_code: Int[Array, '']) → None¶

class diffpes.types.EvidenceRef(evidence_id: str, method_id: str, artifact_refs: tuple[str, ...], source_type: str, independent: bool, measured: Float[Array, 'n_measure'], reference: Float[Array, 'n_measure'], residual: Float[Array, 'n_measure'], tolerance: Float[Array, 'n_measure']) → None[source]¶

Bases: Module

Store numerical evidence with static method and source identity.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestEvidenceref

Variables:

evidence_id (str) – Evidence id (static – a compile-time constant; changing it triggers retracing).
method_id (str) – Method id (static – a compile-time constant; changing it triggers retracing).
artifact_refs (tuple[str, ]) – Artifact refs (static – a compile-time constant; changing it triggers retracing).
source_type (str) – Source type (static – a compile-time constant; changing it triggers retracing).
independent (bool) – Independent (static – a compile-time constant; changing it triggers retracing).
measured (Float[Array, " n_measure"]) – Measured retained as a differentiable JAX leaf in the declared physical units.
reference (Float[Array, " n_measure"]) – Reference retained as a differentiable JAX leaf in the declared physical units.
residual (Float[Array, " n_measure"]) – Residual retained as a differentiable JAX leaf in the declared physical units.
tolerance (Float[Array, " n_measure"]) – Tolerance retained as a differentiable JAX leaf in the declared physical units.

evidence_id: str¶

method_id: str¶

artifact_refs: tuple[str, ...]¶

source_type: str¶

independent: bool¶

measured: Float[Array, 'n_measure']¶

reference: Float[Array, 'n_measure']¶

residual: Float[Array, 'n_measure']¶

tolerance: Float[Array, 'n_measure']¶

__init__(evidence_id: str, method_id: str, artifact_refs: tuple[str, ...], source_type: str, independent: bool, measured: Float[Array, 'n_measure'], reference: Float[Array, 'n_measure'], residual: Float[Array, 'n_measure'], tolerance: Float[Array, 'n_measure']) → None¶

class diffpes.types.EvidenceReport(evidence_id: str, resolved: Bool[Array, ''], compatible: Bool[Array, ''], passed: Bool[Array, ''], residual_norm: Float[Array, '']) → None[source]¶

Bases: Module

Store the offline consistency outcome for one evidence record.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestEvidencereport

Variables:

evidence_id (str) – Evidence id (static – a compile-time constant; changing it triggers retracing).
resolved (Bool[Array, ""]) – Resolved retained as a differentiable JAX leaf in the declared physical units.
compatible (Bool[Array, ""]) – Compatible retained as a differentiable JAX leaf in the declared physical units.
passed (Bool[Array, ""]) – Passed retained as a differentiable JAX leaf in the declared physical units.
residual_norm (Float[Array, ""]) – Residual norm retained as a differentiable JAX leaf in the declared physical units.

evidence_id: str¶

resolved: Bool[Array, '']¶

compatible: Bool[Array, '']¶

passed: Bool[Array, '']¶

residual_norm: Float[Array, '']¶

__init__(evidence_id: str, resolved: Bool[Array, ''], compatible: Bool[Array, ''], passed: Bool[Array, ''], residual_norm: Float[Array, '']) → None¶

class diffpes.types.ExecutionManifest(execution_id: str, model_ref: str, schema_version: str, package_version: str, source_checksum: str, environment_checksum: str, backend: str, precision_policy: str, deterministic: bool, started_at_utc: str) → None[source]¶

Bases: Module

Store software and execution identity prepared at the I/O boundary.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestExecutionmanifest

Variables:

execution_id (str) – Execution id (static – a compile-time constant; changing it triggers retracing).
model_ref (str) – Model ref (static – a compile-time constant; changing it triggers retracing).
schema_version (str) – Schema version (static – a compile-time constant; changing it triggers retracing).
package_version (str) – Package version (static – a compile-time constant; changing it triggers retracing).
source_checksum (str) – Source checksum (static – a compile-time constant; changing it triggers retracing).
environment_checksum (str) – Environment checksum (static – a compile-time constant; changing it triggers retracing).
backend (str) – Backend (static – a compile-time constant; changing it triggers retracing).
precision_policy (str) – Precision policy (static – a compile-time constant; changing it triggers retracing).
deterministic (bool) – Deterministic (static – a compile-time constant; changing it triggers retracing).
started_at_utc (str) – Started at utc (static – a compile-time constant; changing it triggers retracing).

execution_id: str¶

model_ref: str¶

schema_version: str¶

package_version: str¶

source_checksum: str¶

environment_checksum: str¶

backend: str¶

precision_policy: str¶

deterministic: bool¶

started_at_utc: str¶

__init__(execution_id: str, model_ref: str, schema_version: str, package_version: str, source_checksum: str, environment_checksum: str, backend: str, precision_policy: str, deterministic: bool, started_at_utc: str) → None¶

class diffpes.types.ForwardCertificate(manifest: ExecutionManifest, model: ForwardModelSpec, artifacts: tuple[ArtifactRef, ...], transformations: tuple[TransformationRecord, ...], evidence: tuple[EvidenceRef, ...], claims: tuple[CertificationClaim, ...], domains: tuple[DomainResult, ...], derivatives: DerivativeEvidence, dependencies: DependencyMap, sensitivities: SensitivityMap, information: InformationSpectrum, policy_report: PolicyReport, policy_id: str, certificate_checksum: str, extensions_json: str, waivers: tuple[WaiverRecord, ...]) → None[source]¶

Bases: Module

Store the complete assurance record for one forward execution.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestForwardcertificate

Variables:

manifest (ExecutionManifest) – Manifest retained as a differentiable JAX leaf in the declared physical units.
model (ForwardModelSpec) – Model retained as a differentiable JAX leaf in the declared physical units.
artifacts (tuple[ArtifactRef, ]) – Artifacts retained as a differentiable JAX leaf in the declared physical units.
transformations (tuple[TransformationRecord, ]) – Transformations retained as a differentiable JAX leaf in the declared physical units.
evidence (tuple[EvidenceRef, ]) – Evidence retained as a differentiable JAX leaf in the declared physical units.
claims (tuple[CertificationClaim, ]) – Claims retained as a differentiable JAX leaf in the declared physical units.
domains (tuple[DomainResult, ]) – Domains retained as a differentiable JAX leaf in the declared physical units.
derivatives (DerivativeEvidence) – Derivatives retained as a differentiable JAX leaf in the declared physical units.
dependencies (DependencyMap) – Dependencies retained as a differentiable JAX leaf in the declared physical units.
sensitivities (SensitivityMap) – Sensitivities retained as a differentiable JAX leaf in the declared physical units.
information (InformationSpectrum) – Information retained as a differentiable JAX leaf in the declared physical units.
policy_report (PolicyReport) – Policy report retained as a differentiable JAX leaf in the declared physical units.
policy_id (str) – Policy id (static – a compile-time constant; changing it triggers retracing).
certificate_checksum (str) – Certificate checksum (static – a compile-time constant; changing it triggers retracing).
extensions_json (str) – Extensions json (static – a compile-time constant; changing it triggers retracing).
waivers (tuple[WaiverRecord, ]) – Policy-waiver records (static; changing them causes retracing).

manifest: ExecutionManifest¶

model: ForwardModelSpec¶

artifacts: tuple[ArtifactRef, ...]¶

transformations: tuple[TransformationRecord, ...]¶

evidence: tuple[EvidenceRef, ...]¶

claims: tuple[CertificationClaim, ...]¶

domains: tuple[DomainResult, ...]¶

derivatives: DerivativeEvidence¶

dependencies: DependencyMap¶

sensitivities: SensitivityMap¶

information: InformationSpectrum¶

policy_report: PolicyReport¶

policy_id: str¶

certificate_checksum: str¶

extensions_json: str¶

waivers: tuple[WaiverRecord, ...]¶

__init__(manifest: ExecutionManifest, model: ForwardModelSpec, artifacts: tuple[ArtifactRef, ...], transformations: tuple[TransformationRecord, ...], evidence: tuple[EvidenceRef, ...], claims: tuple[CertificationClaim, ...], domains: tuple[DomainResult, ...], derivatives: DerivativeEvidence, dependencies: DependencyMap, sensitivities: SensitivityMap, information: InformationSpectrum, policy_report: PolicyReport, policy_id: str, certificate_checksum: str, extensions_json: str, waivers: tuple[WaiverRecord, ...]) → None¶

class diffpes.types.ForwardModelSpec(model_id: str, model_version: str, observable_id: str, implementation_ref: str, assumptions: tuple[str, ...], conventions: tuple[ConventionRef, ...], domain: tuple[DomainPredicate, ...], differentiable_paths: tuple[str, ...], nondifferentiable_paths: tuple[str, ...]) → None[source]¶

Bases: Module

Store the identity of a differentiable forward model.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestForwardmodelspec

Variables:

model_id (str) – Model id (static – a compile-time constant; changing it triggers retracing).
model_version (str) – Model version (static – a compile-time constant; changing it triggers retracing).
observable_id (str) – Observable id (static – a compile-time constant; changing it triggers retracing).
implementation_ref (str) – Implementation ref (static – a compile-time constant; changing it triggers retracing).
assumptions (tuple[str, ]) – Assumptions (static – a compile-time constant; changing it triggers retracing).
conventions (tuple[ConventionRef, ]) – Conventions (static – a compile-time constant; changing it triggers retracing).
domain (tuple[DomainPredicate, ]) – Domain (static – a compile-time constant; changing it triggers retracing).
differentiable_paths (tuple[str, ]) – Differentiable paths (static – a compile-time constant; changing it triggers retracing).
nondifferentiable_paths (tuple[str, ]) – Nondifferentiable paths (static – a compile-time constant; changing it triggers retracing).

model_id: str¶

model_version: str¶

observable_id: str¶

implementation_ref: str¶

assumptions: tuple[str, ...]¶

conventions: tuple[ConventionRef, ...]¶

domain: tuple[DomainPredicate, ...]¶

differentiable_paths: tuple[str, ...]¶

nondifferentiable_paths: tuple[str, ...]¶

__init__(model_id: str, model_version: str, observable_id: str, implementation_ref: str, assumptions: tuple[str, ...], conventions: tuple[ConventionRef, ...], domain: tuple[DomainPredicate, ...], differentiable_paths: tuple[str, ...], nondifferentiable_paths: tuple[str, ...]) → None¶

class diffpes.types.HandshakeReport(owner_id: str, complete: Bool[Array, ''], missing_ids: tuple[str, ...]) → None[source]¶

Bases: Module

Store the validation outcome for one registration handshake.

The report keeps a JAX Boolean outcome and static missing identities.

See:

TestHandshakeReport

Variables:

owner_id (str) – Plan owner identity (static; changing it causes retracing).
complete (Bool[Array, ""]) – Whether every declared identity has a registry binding.
missing_ids (tuple[str, ]) – Missing declared identities (static; changes cause retracing).

owner_id: str¶

complete: Bool[Array, '']¶

missing_ids: tuple[str, ...]¶

__init__(owner_id: str, complete: Bool[Array, ''], missing_ids: tuple[str, ...]) → None¶

class diffpes.types.InformationSpectrum(input_paths: tuple[str, ...], singular_values: Float[Array, 'n_sv'], right_singular_vectors: Float[Array, 'n_sv n_input'], effective_rank: Int[Array, ''], condition_estimate: Float[Array, ''], threshold: Float[Array, '']) → None[source]¶

Bases: Module

Store a matrix-free information spectrum in input coordinates.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestInformationspectrum

Variables:

input_paths (tuple[str, ]) – Input paths (static – a compile-time constant; changing it triggers retracing).
singular_values (Float[Array, " n_sv"]) – Singular values retained as a differentiable JAX leaf in the declared physical units.
right_singular_vectors (Float[Array, "n_sv n_input"]) – Right singular vectors retained as a differentiable JAX leaf in the declared physical units.
effective_rank (Int[Array, ""]) – Effective rank retained as a differentiable JAX leaf in the declared physical units.
condition_estimate (Float[Array, ""]) – Condition estimate retained as a differentiable JAX leaf in the declared physical units. Zero means that there is no active information direction.
threshold (Float[Array, ""]) – Threshold retained as a differentiable JAX leaf in the declared physical units.

input_paths: tuple[str, ...]¶

singular_values: Float[Array, 'n_sv']¶

right_singular_vectors: Float[Array, 'n_sv n_input']¶

effective_rank: Int[Array, '']¶

condition_estimate: Float[Array, '']¶

threshold: Float[Array, '']¶

__init__(input_paths: tuple[str, ...], singular_values: Float[Array, 'n_sv'], right_singular_vectors: Float[Array, 'n_sv n_input'], effective_rank: Int[Array, ''], condition_estimate: Float[Array, ''], threshold: Float[Array, '']) → None¶

class diffpes.types.PolicyReport(policy_id: str, level_ids: tuple[str, ...], required_claim_ids: tuple[str, ...], claim_passed: Bool[Array, 'n_claim'], claim_checked: Bool[Array, 'n_claim'], claim_in_domain: Bool[Array, 'n_claim'], achieved: Bool[Array, 'n_level']) → None[source]¶

Bases: Module

Store a traced policy truth table for derived certification levels.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestPolicyreport

Variables:

policy_id (str) – Policy id (static – a compile-time constant; changing it triggers retracing).
level_ids (tuple[str, ]) – Level ids (static – a compile-time constant; changing it triggers retracing).
required_claim_ids (tuple[str, ]) – Required claim ids (static – a compile-time constant; changing it triggers retracing).
claim_passed (Bool[Array, " n_claim"]) – Claim passed retained as a differentiable JAX leaf in the declared physical units.
claim_checked (Bool[Array, " n_claim"]) – Claim checked retained as a differentiable JAX leaf in the declared physical units.
claim_in_domain (Bool[Array, " n_claim"]) – Claim in domain retained as a differentiable JAX leaf in the declared physical units.
achieved (Bool[Array, " n_level"]) – Achieved retained as a differentiable JAX leaf in the declared physical units.

policy_id: str¶

level_ids: tuple[str, ...]¶

required_claim_ids: tuple[str, ...]¶

claim_passed: Bool[Array, 'n_claim']¶

claim_checked: Bool[Array, 'n_claim']¶

claim_in_domain: Bool[Array, 'n_claim']¶

achieved: Bool[Array, 'n_level']¶

__init__(policy_id: str, level_ids: tuple[str, ...], required_claim_ids: tuple[str, ...], claim_passed: Bool[Array, 'n_claim'], claim_checked: Bool[Array, 'n_claim'], claim_in_domain: Bool[Array, 'n_claim'], achieved: Bool[Array, 'n_level']) → None¶

class diffpes.types.RegisteredModel(spec: ForwardModelSpec, executor: Callable[[...], Any], registration_checksum: str) → None[source]¶

Bases: Module

Store a frozen binding between a model spec and its executor.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestRegisteredmodel

Variables:

spec (ForwardModelSpec) – Spec retained as a differentiable JAX leaf in the declared physical units.
executor (Callable[..., Any]) – Executor (static – a compile-time constant; changing it triggers retracing).
registration_checksum (str) – Registration checksum (static – a compile-time constant; changing it triggers retracing).

spec: ForwardModelSpec¶

executor: Callable[[...], Any]¶

registration_checksum: str¶

__init__(spec: ForwardModelSpec, executor: Callable[[...], Any], registration_checksum: str) → None¶

class diffpes.types.RegisteredTransformation(contract: TransformationContract, registration_checksum: str) → None[source]¶

Bases: Module

Store a frozen transformation and its consistency checksum.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestRegisteredtransformation

Variables:

contract (TransformationContract) – Contract retained as a differentiable JAX leaf in the declared physical units.
registration_checksum (str) – Registration checksum (static – a compile-time constant; changing it triggers retracing).

contract: TransformationContract¶

registration_checksum: str¶

__init__(contract: TransformationContract, registration_checksum: str) → None¶

class diffpes.types.RegistrationHandshake(owner_id: str, model_refs: tuple[str, ...], transformation_refs: tuple[str, ...], convention_refs: tuple[str, ...], evidence_ids: tuple[str, ...]) → None[source]¶

Bases: Module

Store declarative registration requirements for one plan owner.

The record names required identities without importing an unfinished scientific kernel.

See:

TestRegistrationHandshake

Variables:

owner_id (str) – Plan owner identity (static; changing it causes retracing).
model_refs (tuple[str, ]) – Required model identities (static; changing them causes retracing).
transformation_refs (tuple[str, ]) – Required transformation identities (static; changes cause retracing).
convention_refs (tuple[str, ]) – Required convention identities (static; changes cause retracing).
evidence_ids (tuple[str, ]) – Required evidence identities (static; changing them causes retracing).

owner_id: str¶

model_refs: tuple[str, ...]¶

transformation_refs: tuple[str, ...]¶

convention_refs: tuple[str, ...]¶

evidence_ids: tuple[str, ...]¶

__init__(owner_id: str, model_refs: tuple[str, ...], transformation_refs: tuple[str, ...], convention_refs: tuple[str, ...], evidence_ids: tuple[str, ...]) → None¶

class diffpes.types.RegistryReport(valid: bool, errors: tuple[str, ...], model_count: int, transformation_count: int, checksum: str, frozen: bool) → None[source]¶

Bases: Module

Store the structural validation result for one registry snapshot.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestRegistryreport

Variables:

valid (bool) – Valid (static – a compile-time constant; changing it triggers retracing).
errors (tuple[str, ]) – Errors (static – a compile-time constant; changing it triggers retracing).
model_count (int) – Model count (static – a compile-time constant; changing it triggers retracing).
transformation_count (int) – Transformation count (static – a compile-time constant; changing it triggers retracing).
checksum (str) – Checksum (static – a compile-time constant; changing it triggers retracing).
frozen (bool) – Frozen (static – a compile-time constant; changing it triggers retracing).

valid: bool¶

errors: tuple[str, ...]¶

model_count: int¶

transformation_count: int¶

checksum: str¶

frozen: bool¶

__init__(valid: bool, errors: tuple[str, ...], model_count: int, transformation_count: int, checksum: str, frozen: bool) → None¶

class diffpes.types.RegistrySnapshot(models: tuple[RegisteredModel, ...], transformations: tuple[RegisteredTransformation, ...], checksum: str) → None[source]¶

Bases: Module

Store an immutable deterministic snapshot of registry entries.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestRegistrysnapshot

Variables:

models (tuple[RegisteredModel, ]) – Models retained as a differentiable JAX leaf in the declared physical units.
transformations (tuple[RegisteredTransformation, ]) – Transformations retained as a differentiable JAX leaf in the declared physical units.
checksum (str) – Checksum (static – a compile-time constant; changing it triggers retracing).

models: tuple[RegisteredModel, ...]¶

transformations: tuple[RegisteredTransformation, ...]¶

checksum: str¶

__init__(models: tuple[RegisteredModel, ...], transformations: tuple[RegisteredTransformation, ...], checksum: str) → None¶

class diffpes.types.ReproductionReport(execution_id: str, result_checksum: str, reproduced: Bool[Array, ''], max_abs_error: Float[Array, ''], max_rel_error: Float[Array, ''], tolerance: Float[Array, '']) → None[source]¶

Bases: Module

Store a numerical comparison from deliberate forward re-execution.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestReproductionreport

Variables:

execution_id (str) – Execution id (static – a compile-time constant; changing it triggers retracing).
result_checksum (str) – Result checksum (static – a compile-time constant; changing it triggers retracing).
reproduced (Bool[Array, ""]) – Reproduced retained as a differentiable JAX leaf in the declared physical units.
max_abs_error (Float[Array, ""]) – Max abs error retained as a differentiable JAX leaf in the declared physical units.
max_rel_error (Float[Array, ""]) – Max rel error retained as a differentiable JAX leaf in the declared physical units.
tolerance (Float[Array, ""]) – Tolerance retained as a differentiable JAX leaf in the declared physical units.

execution_id: str¶

result_checksum: str¶

reproduced: Bool[Array, '']¶

max_abs_error: Float[Array, '']¶

max_rel_error: Float[Array, '']¶

tolerance: Float[Array, '']¶

__init__(execution_id: str, result_checksum: str, reproduced: Bool[Array, ''], max_abs_error: Float[Array, ''], max_rel_error: Float[Array, ''], tolerance: Float[Array, '']) → None¶

class diffpes.types.SensitivityMap(input_paths: tuple[str, ...], output_projection_ids: tuple[str, ...], scales: Float[Array, 'n_input'], sensitivities: Float[Array, 'n_output n_input'], threshold: Float[Array, ''], active: Bool[Array, 'n_output n_input']) → None[source]¶

Bases: Module

Store scaled sensitivities from inputs to output projections.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestSensitivitymap

Variables:

input_paths (tuple[str, ]) – Input paths (static – a compile-time constant; changing it triggers retracing).
output_projection_ids (tuple[str, ]) – Output projection ids (static – a compile-time constant; changing it triggers retracing).
scales (Float[Array, " n_input"]) – Scales retained as a differentiable JAX leaf in the declared physical units.
sensitivities (Float[Array, "n_output n_input"]) – Sensitivities retained as a differentiable JAX leaf in the declared physical units.
threshold (Float[Array, ""]) – Threshold retained as a differentiable JAX leaf in the declared physical units.
active (Bool[Array, "n_output n_input"]) – Active retained as a differentiable JAX leaf in the declared physical units.

input_paths: tuple[str, ...]¶

output_projection_ids: tuple[str, ...]¶

scales: Float[Array, 'n_input']¶

sensitivities: Float[Array, 'n_output n_input']¶

threshold: Float[Array, '']¶

active: Bool[Array, 'n_output n_input']¶

__init__(input_paths: tuple[str, ...], output_projection_ids: tuple[str, ...], scales: Float[Array, 'n_input'], sensitivities: Float[Array, 'n_output n_input'], threshold: Float[Array, ''], active: Bool[Array, 'n_output n_input']) → None¶

class diffpes.types.TransformationRecord(transformation_id: str, transformation_version: str, parent_ids: tuple[str, ...], output_ids: tuple[str, ...], preserves: tuple[str, ...], introduces: tuple[str, ...], destroys: tuple[str, ...], invalidates_claims: tuple[str, ...], parameters_checksum: str) → None[source]¶

Bases: Module

Store one transformation and its semantic information effects.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestTransformationrecord

Variables:

transformation_id (str) – Transformation id (static – a compile-time constant; changing it triggers retracing).
transformation_version (str) – Transformation version (static – a compile-time constant; changing it triggers retracing).
parent_ids (tuple[str, ]) – Parent ids (static – a compile-time constant; changing it triggers retracing).
output_ids (tuple[str, ]) – Output ids (static – a compile-time constant; changing it triggers retracing).
preserves (tuple[str, ]) – Preserves (static – a compile-time constant; changing it triggers retracing).
introduces (tuple[str, ]) – Introduces (static – a compile-time constant; changing it triggers retracing).
destroys (tuple[str, ]) – Destroys (static – a compile-time constant; changing it triggers retracing).
invalidates_claims (tuple[str, ]) – Invalidates claims (static – a compile-time constant; changing it triggers retracing).
parameters_checksum (str) – Parameters checksum (static – a compile-time constant; changing it triggers retracing).

transformation_id: str¶

transformation_version: str¶

parent_ids: tuple[str, ...]¶

output_ids: tuple[str, ...]¶

preserves: tuple[str, ...]¶

introduces: tuple[str, ...]¶

destroys: tuple[str, ...]¶

invalidates_claims: tuple[str, ...]¶

parameters_checksum: str¶

__init__(transformation_id: str, transformation_version: str, parent_ids: tuple[str, ...], output_ids: tuple[str, ...], preserves: tuple[str, ...], introduces: tuple[str, ...], destroys: tuple[str, ...], invalidates_claims: tuple[str, ...], parameters_checksum: str) → None¶

class diffpes.types.VerificationReport(certificate_checksum: str, policy_id: str, structure_valid: Bool[Array, ''], evidence_valid: Bool[Array, ''], policy_report: PolicyReport) → None[source]¶

Bases: Module

Store an offline certificate-verification outcome.

Carry scientific vocabulary separately from traced leaves. The record remains stable under JIT, VMAP, JVP, and VJP transforms.

See:

TestVerificationreport

Variables:

certificate_checksum (str) – Certificate checksum (static – a compile-time constant; changing it triggers retracing).
policy_id (str) – Policy id (static – a compile-time constant; changing it triggers retracing).
structure_valid (Bool[Array, ""]) – Structure valid retained as a differentiable JAX leaf in the declared physical units.
evidence_valid (Bool[Array, ""]) – Evidence valid retained as a differentiable JAX leaf in the declared physical units.
policy_report (PolicyReport) – Policy report retained as a differentiable JAX leaf in the declared physical units.

certificate_checksum: str¶

policy_id: str¶

structure_valid: Bool[Array, '']¶

evidence_valid: Bool[Array, '']¶

policy_report: PolicyReport¶

__init__(certificate_checksum: str, policy_id: str, structure_valid: Bool[Array, ''], evidence_valid: Bool[Array, ''], policy_report: PolicyReport) → None¶

class diffpes.types.WaiverRecord(waiver_id: str, policy_id: str, claim_ids: tuple[str, ...], author: str, reason: str, issued_at_utc: str, expires_at_utc: str) → None[source]¶

Bases: Module

Store a bounded policy-waiver declaration without changing claim status.

A waiver records review context only. It never changes a failed claim to a passed claim.

See:

TestWaiverRecord

Variables:

waiver_id (str) – Permanent waiver identity (static; changing it causes retracing).
policy_id (str) – Applicable policy identity (static; changing it causes retracing).
claim_ids (tuple[str, ]) – Affected claim identities (static; changing them causes retracing).
author (str) – Responsible reviewer (static; changing it causes retracing).
reason (str) – Technical reason (static; changing it causes retracing).
issued_at_utc (str) – Absolute UTC issue time (static; changing it causes retracing).
expires_at_utc (str) – Absolute UTC expiry time (static; changing it causes retracing).

waiver_id: str¶

policy_id: str¶

claim_ids: tuple[str, ...]¶

author: str¶

reason: str¶

issued_at_utc: str¶

expires_at_utc: str¶

__init__(waiver_id: str, policy_id: str, claim_ids: tuple[str, ...], author: str, reason: str, issued_at_utc: str, expires_at_utc: str) → None¶

class diffpes.types.WaiverReport(waiver_id: str, valid: Bool[Array, ''], active: Bool[Array, ''], errors: tuple[str, ...]) → None[source]¶

Bases: Module

Store the temporal validation outcome for one waiver.

The report distinguishes valid structure from active temporal scope.

See:

TestWaiverReport

Variables:

waiver_id (str) – Permanent waiver identity (static; changing it causes retracing).
valid (Bool[Array, ""]) – Whether the record has valid absolute UTC fields.
active (Bool[Array, ""]) – Whether the waiver covers the selected UTC time.
errors (tuple[str, ]) – Validation errors (static; changing them causes retracing).

waiver_id: str¶

valid: Bool[Array, '']¶

active: Bool[Array, '']¶

errors: tuple[str, ...]¶

__init__(waiver_id: str, valid: Bool[Array, ''], active: Bool[Array, ''], errors: tuple[str, ...]) → None¶

diffpes.types.make_artifact_ref(artifact_id: str, media_type: str, byte_checksum: str | None, content_checksum: str, semantic_checksum: str, locator: str | None, role: str) → ArtifactRef[source]¶

Create a validated artifact reference.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakeArtifactRef

Parameters:

artifact_id (str) – Artifact id used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
media_type (str) – Media type used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
byte_checksum (Optional[str]) – Byte checksum used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
content_checksum (str) – Content checksum used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
semantic_checksum (str) – Semantic checksum used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
locator (Optional[str]) – Locator used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
role (str) – Role used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).

Returns:

result – Validated immutable carrier.

Return type:

ArtifactRef

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_certification_claim(claim_id: str, subject_id: str, predicate_id: str, evidence_ids: tuple[str, ...], measured: Any, reference: Any, residual: Any, tolerance: Any, passed: Any, checked: Any = True, in_domain: Any = True, margin: Any = 0.0, severity_code: Any = 0) → CertificationClaim[source]¶

Create a claim retaining both continuous and discrete evidence.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakeCertificationClaim

Parameters:

claim_id (str) – Claim id used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
subject_id (str) – Subject id used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
predicate_id (str) – Predicate id used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
evidence_ids (tuple[str, ...]) – Evidence ids used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
measured (Any) – Measured used to construct the validated carrier as a traced numerical value in the declared physical units.
reference (Any) – Reference used to construct the validated carrier as a traced numerical value in the declared physical units.
residual (Any) – Residual used to construct the validated carrier as a traced numerical value in the declared physical units.
tolerance (Any) – Tolerance used to construct the validated carrier as a traced numerical value in the declared physical units.
passed (Any) – Passed used to construct the validated carrier as a traced numerical value in the declared physical units.
checked (Any, default: True) – Checked used to construct the validated carrier as a traced numerical value in the declared physical units.
in_domain (Any, default: True) – In domain used to construct the validated carrier as a traced numerical value in the declared physical units.
margin (Any, default: 0.0) – Margin used to construct the validated carrier as a traced numerical value in the declared physical units.
severity_code (Any, default: 0) – Severity code used to construct the validated carrier as a traced numerical value in the declared physical units.

Returns:

result – Validated immutable carrier.

Return type:

CertificationClaim

Raises:

ValueError – If static structure or cross-record validation fails.

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_certification_context(manifest: ExecutionManifest, model: ForwardModelSpec, artifacts: tuple[ArtifactRef, ...] = (), transformations: tuple[TransformationRecord, ...] = (), evidence: tuple[EvidenceRef, ...] = (), policy_id: str = 'org.diffpes.policy.research.v1', check_ids: tuple[str, ...] = (), input_checksums: tuple[str, ...] = (), waivers: tuple[WaiverRecord, ...] = ()) → CertificationContext[source]¶

Create a prepared certification context.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakeCertificationContext

Parameters:

manifest (ExecutionManifest) – Manifest used to construct the validated carrier as a traced numerical value in the declared physical units.
model (ForwardModelSpec) – Model used to construct the validated carrier as a traced numerical value in the declared physical units.
artifacts (tuple[ArtifactRef, ...], default: ()) – Artifacts used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
transformations (tuple[TransformationRecord, ...], default: ()) – Transformations used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
evidence (tuple[EvidenceRef, ...], default: ()) – Evidence used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
policy_id (str, default: "org.diffpes.policy.research.v1") – Policy id used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
check_ids (tuple[str, ...], default: ()) – Check ids used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
input_checksums (tuple[str, ...], default: ()) – Input checksums used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
waivers (tuple[WaiverRecord, ...], default: ()) – Policy-waiver records. Default is an empty tuple.

Returns:

result – Validated immutable carrier.

Return type:

CertificationContext

Raises:

ValueError – If static structure or cross-record validation fails.

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_certified_result(value: Any, certificate: ForwardCertificate) → CertifiedResult[source]¶

Pair any JAX-compatible result value with a forward certificate.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakeCertifiedResult

Parameters:

value (Any) – Value used to construct the validated carrier as a traced numerical value in the declared physical units.
certificate (ForwardCertificate) – Certificate used to construct the validated carrier as a traced numerical value in the declared physical units.

Returns:

result – Validated immutable carrier.

Return type:

CertifiedResult

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_convention_ref(convention_id: str, version: str, parameters_json: str = '{}') → ConventionRef[source]¶

Create a validated convention reference.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakeConventionRef

Parameters:

convention_id (str) – Convention id used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
version (str) – Version used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
parameters_json (str, default: "{}") – Parameters json used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).

Returns:

result – Validated immutable carrier.

Return type:

ConventionRef

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_dependency_map(model_id: str, input_paths: tuple[str, ...], output_paths: tuple[str, ...], structural: Any, traced: Any) → DependencyMap[source]¶

Create a structural dependency map.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakeDependencyMap

Parameters:

model_id (str) – Model id used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
input_paths (tuple[str, ...]) – Input paths used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
output_paths (tuple[str, ...]) – Output paths used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
structural (Any) – Structural used to construct the validated carrier as a traced numerical value in the declared physical units.
traced (Any) – Traced used to construct the validated carrier as a traced numerical value in the declared physical units.

Returns:

result – Validated immutable carrier.

Return type:

DependencyMap

Raises:

ValueError – If static structure or cross-record validation fails.

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_derivative_evidence(input_paths: tuple[str, ...], output_projection_ids: tuple[str, ...], method: str, scales: Any, jvp_probes: Any, vjp_probes: Any, reference_derivatives: Any, derivative_residuals: Any, singular_values: Any, effective_rank: Any, condition_estimate: Any, finite: Any, fd_correct: Any) → DerivativeEvidence[source]¶

Create validated derivative and local-information evidence.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakeDerivativeEvidence

Parameters:

input_paths (tuple[str, ...]) – Input paths used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
output_projection_ids (tuple[str, ...]) – Output projection ids used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
method (str) – Method used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
scales (Any) – Scales used to construct the validated carrier as a traced numerical value in the declared physical units.
jvp_probes (Any) – Jvp probes used to construct the validated carrier as a traced numerical value in the declared physical units.
vjp_probes (Any) – Vjp probes used to construct the validated carrier as a traced numerical value in the declared physical units.
reference_derivatives (Any) – Reference derivatives used to construct the validated carrier as a traced numerical value in the declared physical units.
derivative_residuals (Any) – Derivative residuals used to construct the validated carrier as a traced numerical value in the declared physical units.
singular_values (Any) – Singular values used to construct the validated carrier as a traced numerical value in the declared physical units.
effective_rank (Any) – Effective rank used to construct the validated carrier as a traced numerical value in the declared physical units.
condition_estimate (Any) – Condition estimate used to construct the validated carrier as a traced numerical value in the declared physical units. Zero means that there is no active information direction.
finite (Any) – Finite used to construct the validated carrier as a traced numerical value in the declared physical units.
fd_correct (Any) – Fd correct used to construct the validated carrier as a traced numerical value in the declared physical units.

Returns:

result – Validated immutable carrier.

Return type:

DerivativeEvidence

Raises:

ValueError – If static structure or cross-record validation fails.

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_domain_predicate(predicate_id: str, expression_id: str, units: str | None = None, severity: str = 'error') → DomainPredicate[source]¶

Create a validated domain-predicate declaration.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakeDomainPredicate

Parameters:

predicate_id (str) – Predicate id used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
expression_id (str) – Expression id used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
units (Optional[str], default: None) – Units used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
severity (str, default: "error") – Severity used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).

Returns:

result – Validated immutable carrier.

Return type:

DomainPredicate

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_domain_result(predicate_id: str, measured: Any, reference: Any, residual: Any, tolerance: Any, margin: Any, passed: Any, checked: Any = True, in_domain: Any = True, severity_code: Any = 0) → DomainResult[source]¶

Create one traced domain evaluation.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakeDomainResult

Parameters:

predicate_id (str) – Predicate id used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
measured (Any) – Measured used to construct the validated carrier as a traced numerical value in the declared physical units.
reference (Any) – Reference used to construct the validated carrier as a traced numerical value in the declared physical units.
residual (Any) – Residual used to construct the validated carrier as a traced numerical value in the declared physical units.
tolerance (Any) – Tolerance used to construct the validated carrier as a traced numerical value in the declared physical units.
margin (Any) – Margin used to construct the validated carrier as a traced numerical value in the declared physical units.
passed (Any) – Passed used to construct the validated carrier as a traced numerical value in the declared physical units.
checked (Any, default: True) – Checked used to construct the validated carrier as a traced numerical value in the declared physical units.
in_domain (Any, default: True) – In domain used to construct the validated carrier as a traced numerical value in the declared physical units.
severity_code (Any, default: 0) – Severity code used to construct the validated carrier as a traced numerical value in the declared physical units.

Returns:

result – Validated immutable carrier.

Return type:

DomainResult

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_evidence_ref(evidence_id: str, method_id: str, artifact_refs: tuple[str, ...], source_type: str, independent: bool, measured: Any, reference: Any, residual: Any, tolerance: Any) → EvidenceRef[source]¶

Create validated vector-valued numerical evidence.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakeEvidenceRef

Parameters:

evidence_id (str) – Evidence id used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
method_id (str) – Method id used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
artifact_refs (tuple[str, ...]) – Artifact refs used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
source_type (str) – Source type used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
independent (bool) – Independent used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
measured (Any) – Measured used to construct the validated carrier as a traced numerical value in the declared physical units.
reference (Any) – Reference used to construct the validated carrier as a traced numerical value in the declared physical units.
residual (Any) – Residual used to construct the validated carrier as a traced numerical value in the declared physical units.
tolerance (Any) – Tolerance used to construct the validated carrier as a traced numerical value in the declared physical units.

Returns:

result – Validated immutable carrier.

Return type:

EvidenceRef

Raises:

ValueError – If static structure or cross-record validation fails.

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_evidence_report(evidence_id: str, resolved: Any, compatible: Any, passed: Any, residual_norm: Any) → EvidenceReport[source]¶

Create an offline evidence-verification report.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakeEvidenceReport

Parameters:

evidence_id (str) – Evidence id used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
resolved (Any) – Resolved used to construct the validated carrier as a traced numerical value in the declared physical units.
compatible (Any) – Compatible used to construct the validated carrier as a traced numerical value in the declared physical units.
passed (Any) – Passed used to construct the validated carrier as a traced numerical value in the declared physical units.
residual_norm (Any) – Residual norm used to construct the validated carrier as a traced numerical value in the declared physical units.

Returns:

result – Validated immutable carrier.

Return type:

EvidenceReport

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_execution_manifest(execution_id: str, model_ref: str, schema_version: str, package_version: str, source_checksum: str, environment_checksum: str, backend: str, precision_policy: str, deterministic: bool, started_at_utc: str) → ExecutionManifest[source]¶

Create a validated execution manifest.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakeExecutionManifest

Parameters:

execution_id (str) – Execution id used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
model_ref (str) – Model ref used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
schema_version (str) – Schema version used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
package_version (str) – Package version used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
source_checksum (str) – Source checksum used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
environment_checksum (str) – Environment checksum used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
backend (str) – Backend used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
precision_policy (str) – Precision policy used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
deterministic (bool) – Deterministic used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
started_at_utc (str) – Started at utc used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).

Returns:

result – Validated immutable carrier.

Return type:

ExecutionManifest

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_forward_certificate(manifest: ExecutionManifest, model: ForwardModelSpec, artifacts: tuple[ArtifactRef, ...], transformations: tuple[TransformationRecord, ...], evidence: tuple[EvidenceRef, ...], claims: tuple[CertificationClaim, ...], domains: tuple[DomainResult, ...], derivatives: DerivativeEvidence, dependencies: DependencyMap, sensitivities: SensitivityMap, information: InformationSpectrum, policy_report: PolicyReport, policy_id: str, certificate_checksum: str, extensions_json: str = '{}', waivers: tuple[WaiverRecord, ...] = ()) → ForwardCertificate[source]¶

Create and cross-validate a complete forward certificate.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakeForwardCertificate

Parameters:

manifest (ExecutionManifest) – Manifest used to construct the validated carrier as a traced numerical value in the declared physical units.
model (ForwardModelSpec) – Model used to construct the validated carrier as a traced numerical value in the declared physical units.
artifacts (tuple[ArtifactRef, ...]) – Artifacts used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
transformations (tuple[TransformationRecord, ...]) – Transformations used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
evidence (tuple[EvidenceRef, ...]) – Evidence used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
claims (tuple[CertificationClaim, ...]) – Claims used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
domains (tuple[DomainResult, ...]) – Domains used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
derivatives (DerivativeEvidence) – Derivatives used to construct the validated carrier as a traced numerical value in the declared physical units.
dependencies (DependencyMap) – Dependencies used to construct the validated carrier as a traced numerical value in the declared physical units.
sensitivities (SensitivityMap) – Sensitivities used to construct the validated carrier as a traced numerical value in the declared physical units.
information (InformationSpectrum) – Information used to construct the validated carrier as a traced numerical value in the declared physical units.
policy_report (PolicyReport) – Policy report used to construct the validated carrier as a traced numerical value in the declared physical units.
policy_id (str) – Policy id used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
certificate_checksum (str) – Certificate checksum used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
extensions_json (str, default: "{}") – Extensions json used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
waivers (tuple[WaiverRecord, ...], default: ()) – Policy-waiver records. Default is an empty tuple.

Returns:

result – Validated immutable carrier.

Return type:

ForwardCertificate

Raises:

ValueError – If static structure or cross-record validation fails.

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_forward_model_spec(model_id: str, model_version: str, observable_id: str, implementation_ref: str, assumptions: tuple[str, ...] = (), conventions: tuple[ConventionRef, ...] = (), domain: tuple[DomainPredicate, ...] = (), differentiable_paths: tuple[str, ...] = (), nondifferentiable_paths: tuple[str, ...] = ()) → ForwardModelSpec[source]¶

Create a validated stable forward-model specification.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakeForwardModelSpec

Parameters:

model_id (str) – Model id used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
model_version (str) – Model version used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
observable_id (str) – Observable id used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
implementation_ref (str) – Implementation ref used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
assumptions (tuple[str, ...], default: ()) – Assumptions used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
conventions (tuple[ConventionRef, ...], default: ()) – Conventions used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
domain (tuple[DomainPredicate, ...], default: ()) – Domain used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
differentiable_paths (tuple[str, ...], default: ()) – Differentiable paths used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
nondifferentiable_paths (tuple[str, ...], default: ()) – Nondifferentiable paths used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).

Returns:

result – Validated immutable carrier.

Return type:

ForwardModelSpec

Raises:

ValueError – If static structure or cross-record validation fails.

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_handshake_report(owner_id: str, complete: Any, missing_ids: tuple[str, ...] = ()) → HandshakeReport[source]¶

Create a report for one registration handshake.

The report keeps the completion outcome as a JAX Boolean leaf.

See:

TestMakeHandshakeReport

Parameters:

owner_id (str) – Plan owner identity (static; changing it causes retracing).
complete (Any) – Whether every declared identity has a registry binding.
missing_ids (tuple[str, ...], default: ()) – Missing declared identities. Default is an empty tuple.

Returns:

result – Validated completion outcome and missing identities.

Return type:

HandshakeReport

Notes

The factory converts the completion outcome to a scalar JAX Boolean.

diffpes.types.make_information_spectrum(input_paths: tuple[str, ...], singular_values: Any, right_singular_vectors: Any, effective_rank: Any, condition_estimate: Any, threshold: Any) → InformationSpectrum[source]¶

Create a validated local information spectrum.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakeInformationSpectrum

Parameters:

input_paths (tuple[str, ...]) – Input paths used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
singular_values (Any) – Singular values used to construct the validated carrier as a traced numerical value in the declared physical units.
right_singular_vectors (Any) – Right singular vectors used to construct the validated carrier as a traced numerical value in the declared physical units.
effective_rank (Any) – Effective rank used to construct the validated carrier as a traced numerical value in the declared physical units.
condition_estimate (Any) – Condition estimate used to construct the validated carrier as a traced numerical value in the declared physical units. Zero means that there is no active information direction.
threshold (Any) – Threshold used to construct the validated carrier as a traced numerical value in the declared physical units.

Returns:

result – Validated immutable carrier.

Return type:

InformationSpectrum

Raises:

ValueError – If static structure or cross-record validation fails.

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_policy_report(policy_id: str, level_ids: tuple[str, ...], required_claim_ids: tuple[str, ...], claim_passed: Any, claim_checked: Any, claim_in_domain: Any, achieved: Any) → PolicyReport[source]¶

Create a validated policy truth table.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakePolicyReport

Parameters:

policy_id (str) – Policy id used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
level_ids (tuple[str, ...]) – Level ids used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
required_claim_ids (tuple[str, ...]) – Required claim ids used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
claim_passed (Any) – Claim passed used to construct the validated carrier as a traced numerical value in the declared physical units.
claim_checked (Any) – Claim checked used to construct the validated carrier as a traced numerical value in the declared physical units.
claim_in_domain (Any) – Claim in domain used to construct the validated carrier as a traced numerical value in the declared physical units.
achieved (Any) – Achieved used to construct the validated carrier as a traced numerical value in the declared physical units.

Returns:

result – Validated immutable carrier.

Return type:

PolicyReport

Raises:

ValueError – If static structure or cross-record validation fails.

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_registered_model(spec: ForwardModelSpec, executor: Callable[[...], Any], registration_checksum: str) → RegisteredModel[source]¶

Create a validated model-registry binding.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakeRegisteredModel

Parameters:

spec (ForwardModelSpec) – Spec used to construct the validated carrier as a traced numerical value in the declared physical units.
executor (Callable[..., Any]) – Executor used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
registration_checksum (str) – Registration checksum used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).

Returns:

result – Validated immutable carrier.

Return type:

RegisteredModel

Raises:

TypeError – If static structure or cross-record validation fails.

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_registered_transformation(contract: TransformationContract, registration_checksum: str) → RegisteredTransformation[source]¶

Create a validated transformation-registry binding.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakeRegisteredTransformation

Parameters:

contract (TransformationContract) – Contract used to construct the validated carrier as a traced numerical value in the declared physical units.
registration_checksum (str) – Registration checksum used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).

Returns:

result – Validated immutable carrier.

Return type:

RegisteredTransformation

Raises:

ValueError – If static structure or cross-record validation fails.

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_registration_handshake(owner_id: str, model_refs: tuple[str, ...] = (), transformation_refs: tuple[str, ...] = (), convention_refs: tuple[str, ...] = (), evidence_ids: tuple[str, ...] = ()) → RegistrationHandshake[source]¶

Create declarative registration requirements for one plan owner.

The factory validates each static identity without importing a scientific executor.

See:

TestMakeRegistrationHandshake

Parameters:

owner_id (str) – Plan owner identity (static; changing it causes retracing).
model_refs (tuple[str, ...], default: ()) – Required model identities. Default is an empty tuple.
transformation_refs (tuple[str, ...], default: ()) – Required transformation identities. Default is an empty tuple.
convention_refs (tuple[str, ...], default: ()) – Required convention identities. Default is an empty tuple.
evidence_ids (tuple[str, ...], default: ()) – Required evidence identities. Default is an empty tuple.

Returns:

result – Validated declarative registration requirements.

Return type:

RegistrationHandshake

Notes

The factory validates each static identity before module construction.

diffpes.types.make_registry_report(valid: bool, errors: tuple[str, ...], model_count: int, transformation_count: int, checksum: str, frozen: bool) → RegistryReport[source]¶

Create a validated structural registry report.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakeRegistryReport

Parameters:

valid (bool) – Valid used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
errors (tuple[str, ...]) – Errors used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
model_count (int) – Model count used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
transformation_count (int) – Transformation count used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
checksum (str) – Checksum used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
frozen (bool) – Frozen used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).

Returns:

result – Validated immutable carrier.

Return type:

RegistryReport

Raises:

ValueError – If static structure or cross-record validation fails.

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_registry_snapshot(models: tuple[RegisteredModel, ...], transformations: tuple[RegisteredTransformation, ...], checksum: str) → RegistrySnapshot[source]¶

Create an immutable registry snapshot.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakeRegistrySnapshot

Parameters:

models (tuple[RegisteredModel, ...]) – Models used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
transformations (tuple[RegisteredTransformation, ...]) – Transformations used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
checksum (str) – Checksum used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).

Returns:

result – Validated immutable carrier.

Return type:

RegistrySnapshot

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_reproduction_report(execution_id: str, result_checksum: str, reproduced: Any, max_abs_error: Any, max_rel_error: Any, tolerance: Any) → ReproductionReport[source]¶

Create a report comparing a result with its re-execution.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakeReproductionReport

Parameters:

execution_id (str) – Execution id used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
result_checksum (str) – Result checksum used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
reproduced (Any) – Reproduced used to construct the validated carrier as a traced numerical value in the declared physical units.
max_abs_error (Any) – Max abs error used to construct the validated carrier as a traced numerical value in the declared physical units.
max_rel_error (Any) – Max rel error used to construct the validated carrier as a traced numerical value in the declared physical units.
tolerance (Any) – Tolerance used to construct the validated carrier as a traced numerical value in the declared physical units.

Returns:

result – Validated immutable carrier.

Return type:

ReproductionReport

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_sensitivity_map(input_paths: tuple[str, ...], output_projection_ids: tuple[str, ...], scales: Any, sensitivities: Any, threshold: Any, active: Any) → SensitivityMap[source]¶

Create a named, scaled local-sensitivity map.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakeSensitivityMap

Parameters:

input_paths (tuple[str, ...]) – Input paths used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
output_projection_ids (tuple[str, ...]) – Output projection ids used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
scales (Any) – Scales used to construct the validated carrier as a traced numerical value in the declared physical units.
sensitivities (Any) – Sensitivities used to construct the validated carrier as a traced numerical value in the declared physical units.
threshold (Any) – Threshold used to construct the validated carrier as a traced numerical value in the declared physical units.
active (Any) – Active used to construct the validated carrier as a traced numerical value in the declared physical units.

Returns:

result – Validated immutable carrier.

Return type:

SensitivityMap

Raises:

ValueError – If static structure or cross-record validation fails.

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_transformation_record(transformation_id: str, transformation_version: str, parent_ids: tuple[str, ...], output_ids: tuple[str, ...], preserves: tuple[str, ...] = (), introduces: tuple[str, ...] = (), destroys: tuple[str, ...] = (), invalidates_claims: tuple[str, ...] = (), parameters_checksum: str = 'none') → TransformationRecord[source]¶

Create a validated information-aware transformation record.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakeTransformationRecord

Parameters:

transformation_id (str) – Transformation id used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
transformation_version (str) – Transformation version used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
parent_ids (tuple[str, ...]) – Parent ids used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
output_ids (tuple[str, ...]) – Output ids used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
preserves (tuple[str, ...], default: ()) – Preserves used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
introduces (tuple[str, ...], default: ()) – Introduces used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
destroys (tuple[str, ...], default: ()) – Destroys used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
invalidates_claims (tuple[str, ...], default: ()) – Invalidates claims used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
parameters_checksum (str, default: "none") – Parameters checksum used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).

Returns:

result – Validated immutable carrier.

Return type:

TransformationRecord

Raises:

ValueError – If static structure or cross-record validation fails.

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_verification_report(certificate_checksum: str, policy_id: str, structure_valid: Any, evidence_valid: Any, policy_report: PolicyReport) → VerificationReport[source]¶

Create an offline certificate-verification report.

Carry static scientific vocabulary separately from traced numerical leaves while preserving the validation boundary defined by this factory.

See:

TestMakeVerificationReport

Parameters:

certificate_checksum (str) – Certificate checksum used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
policy_id (str) – Policy id used to construct the validated carrier (static – a compile-time constant; changing it triggers retracing).
structure_valid (Any) – Structure valid used to construct the validated carrier as a traced numerical value in the declared physical units.
evidence_valid (Any) – Evidence valid used to construct the validated carrier as a traced numerical value in the declared physical units.
policy_report (PolicyReport) – Policy report used to construct the validated carrier as a traced numerical value in the declared physical units.

Returns:

result – Validated immutable carrier.

Return type:

VerificationReport

Raises:

ValueError – If static structure or cross-record validation fails.

Notes

The factory checks static structure eagerly. JAX array operations validate numerical values and preserve differentiation behavior.

diffpes.types.make_waiver_record(waiver_id: str, policy_id: str, claim_ids: tuple[str, ...], author: str, reason: str, issued_at_utc: str, expires_at_utc: str) → WaiverRecord[source]¶

Create a bounded policy-waiver declaration.

The factory validates static vocabulary. The waiver validator checks the absolute UTC interval.

See:

TestMakeWaiverRecord

Parameters:

waiver_id (str) – Permanent waiver identity (static; changing it causes retracing).
policy_id (str) – Applicable policy identity (static; changing it causes retracing).
claim_ids (tuple[str, ...]) – Affected claim identities (static; changing them causes retracing).
author (str) – Responsible reviewer (static; changing it causes retracing).
reason (str) – Technical reason (static; changing it causes retracing).
issued_at_utc (str) – Absolute UTC issue time (static; changing it causes retracing).
expires_at_utc (str) – Absolute UTC expiry time (static; changing it causes retracing).

Returns:

result – Validated static waiver declaration.

Return type:

WaiverRecord

Raises:

ValueError – If a required text field or claim identity is empty.

Notes

The factory validates static text before it constructs the waiver record.

diffpes.types.make_waiver_report(waiver_id: str, valid: Any, active: Any, errors: tuple[str, ...] = ()) → WaiverReport[source]¶

Create a temporal waiver-validation report.

The report keeps validation outcomes as JAX Boolean leaves.

See:

TestMakeWaiverReport

Parameters:

waiver_id (str) – Permanent waiver identity (static; changing it causes retracing).
valid (Any) – Whether the record has valid absolute UTC fields.
active (Any) – Whether the waiver covers the selected UTC time.
errors (tuple[str, ...], default: ()) – Validation errors. Default is an empty tuple.

Returns:

result – Validated temporal outcome and errors.

Return type:

WaiverReport

Notes

The factory converts temporal outcomes to scalar JAX Boolean leaves.

class diffpes.types.WorkflowContext(bands: BandStructure, orb_proj: OrbitalProjection | SpinOrbitalProjection, kpath: KPathInfo | None, dos: DensityOfStates | FullDensityOfStates | None) → None[source]¶

Bases: Module

Store parsed VASP inputs for high-level workflow helpers.

Bundles the electronic bands and orbital projections required by an ARPES simulation with optional k-path and density-of-states metadata. Every non-None field is itself an Equinox PyTree, so the complete context can pass through JAX transformations as one immutable object.

See:

TestWorkflowContext

Variables:

bands (BandStructure) – Band eigenvalues and k-point coordinates from EIGENVAL.
orb_proj (ProjectionType) – Orbital projections from PROCAR, optionally including spin.
kpath (Optional[KPathInfo]) – Parsed KPOINTS metadata when available.
dos (Optional[DosType]) – Parsed DOSCAR data when available.

bands: BandStructure¶

orb_proj: OrbitalProjection | SpinOrbitalProjection¶

kpath: KPathInfo | None¶

dos: DensityOfStates | FullDensityOfStates | None¶

__init__(bands: BandStructure, orb_proj: OrbitalProjection | SpinOrbitalProjection, kpath: KPathInfo | None, dos: DensityOfStates | FullDensityOfStates | None) → None¶

diffpes.types.make_workflow_context(bands: BandStructure, orb_proj: OrbitalProjection | SpinOrbitalProjection, kpath: KPathInfo | None = None, dos: DensityOfStates | FullDensityOfStates | None = None) → WorkflowContext[source]¶

Create a workflow context from parsed VASP inputs.

This type collects parsed electronic bands and orbital projections with optional k-path and density-of-states data. The result is one immutable Equinox PyTree for high-level simulation workflows.

See:: TestMakeWorkflowContext

Implementation Logic¶

Check the shared axes:
```
bands.eigenvalues.shape[:2] != orb_proj.projections.shape[:2]
```
This static shape check raises ValueError before construction.
Construct the context:
```
context = WorkflowContext(...)
```
The Equinox module keeps all supplied carriers in one PyTree.

param bands:: Band eigenvalues and k-point coordinates.
type bands:: BandStructure
param orb_proj:: Orbital projections, optionally including spin channels.
type orb_proj:: Union[OrbitalProjection, SpinOrbitalProjection]
param kpath:: Parsed k-path metadata. Default is None.
type kpath:: Optional[KPathInfo], default: None
param dos:: Parsed density-of-states data. Default is None.
type dos:: Union[DensityOfStates, FullDensityOfStates, None], default: None
returns:: context – Immutable Equinox container holding the supplied inputs.
rtype:: WorkflowContext
raises ValueError:: If the k-point or band dimensions disagree between bands and orb_proj.

class diffpes.types.CompositionReport(valid: bool, errors: tuple[str, ...], available_semantics: tuple[str, ...], destroyed_information: tuple[str, ...], invalidated_claims: tuple[str, ...], transformation_refs: tuple[str, ...]) → None[source]¶

Bases: Module

Store a conservative transformation-composition result.

This carrier records whether a sequence of transformation contracts can compose without losing a required scientific semantic.

See:

TestCompositionReport

Variables:

valid (bool) – Whether composition succeeds (static – a compile-time constant; changing it triggers retracing).
errors (tuple[str, ]) – Composition failures (static – compile-time constants; changing them triggers retracing).
available_semantics (tuple[str, ]) – Semantics available after composition (static – compile-time constants; changing them triggers retracing).
destroyed_information (tuple[str, ]) – Information destroyed along the path (static – compile-time constants; changing them triggers retracing).
invalidated_claims (tuple[str, ]) – Claims invalidated along the path (static – compile-time constants; changing them triggers retracing).
transformation_refs (tuple[str, ]) – Ordered transformation references (static – compile-time constants; changing them triggers retracing).

Notes

Every field is static metadata. Code outside numerical kernels evaluates composition. Therefore, the report introduces no gradient path or reduction.

See also

make_composition_report: Create a validated immutable transformation-composition report.

valid: bool¶

errors: tuple[str, ...]¶

available_semantics: tuple[str, ...]¶

destroyed_information: tuple[str, ...]¶

invalidated_claims: tuple[str, ...]¶

transformation_refs: tuple[str, ...]¶

__init__(valid: bool, errors: tuple[str, ...], available_semantics: tuple[str, ...], destroyed_information: tuple[str, ...], invalidated_claims: tuple[str, ...], transformation_refs: tuple[str, ...]) → None¶

class diffpes.types.TransformationContract(transformation_id: str, transformation_version: str, requires: tuple[str, ...], produces: tuple[str, ...], preserves: tuple[str, ...], introduces: tuple[str, ...], destroys: tuple[str, ...], invalidates_claims: tuple[str, ...], jax_pure: bool) → None[source]¶

Bases: Module

Store the static semantic contract for one registered transformation.

This carrier declares how one transformation changes the named scientific semantics available to downstream forward-model claims.

See:

TestTransformationContract

Variables:

transformation_id (str) – Permanent transformation identifier (static – a compile-time constant; changing it triggers retracing).
transformation_version (str) – Semantic version (static – a compile-time constant; changing it triggers retracing).
requires (tuple[str, ]) – Semantics required on input (static – compile-time constants; changing them triggers retracing).
produces (tuple[str, ]) – Output semantics created by the transformation (static – compile-time constants; changing them triggers retracing).
preserves (tuple[str, ]) – Input semantics preserved exactly (static – compile-time constants; changing them triggers retracing).
introduces (tuple[str, ]) – New semantics introduced on output (static – compile-time constants; changing them triggers retracing).
destroys (tuple[str, ]) – Information made unavailable on output (static – compile-time constants; changing them triggers retracing).
invalidates_claims (tuple[str, ]) – Claim identifiers invalidated by the transformation (static – compile-time constants; changing them triggers retracing).
jax_pure (bool) – Whether the contract declares the transformation JAX-pure (static – a compile-time constant; changing it triggers retracing).

Notes

Every field is static metadata. The carrier has no differentiable leaves; it describes a JAX transformation but does not participate in its gradient.

See also

make_transformation_contract: Create a validated immutable transformation contract.

transformation_id: str¶

transformation_version: str¶

requires: tuple[str, ...]¶

produces: tuple[str, ...]¶

preserves: tuple[str, ...]¶

introduces: tuple[str, ...]¶

destroys: tuple[str, ...]¶

invalidates_claims: tuple[str, ...]¶

jax_pure: bool¶

__init__(transformation_id: str, transformation_version: str, requires: tuple[str, ...], produces: tuple[str, ...], preserves: tuple[str, ...], introduces: tuple[str, ...], destroys: tuple[str, ...], invalidates_claims: tuple[str, ...], jax_pure: bool) → None¶

diffpes.types.make_composition_report(valid: bool, errors: Sequence[str], available_semantics: Sequence[str], destroyed_information: Sequence[str], invalidated_claims: Sequence[str], transformation_refs: Sequence[str]) → CompositionReport[source]¶

Create a validated immutable transformation-composition report.

Freeze the cumulative semantic state and require the validity flag to agree exactly with whether the report contains errors.

See:: TestMakeCompositionReport

Implementation Logic¶

Normalize errors:
```
normalized_errors = _normalize_terms(errors, field_name="errors")
```
Freeze validation failures and reject blank or duplicate entries.
Check validity consistency:
```
if valid == bool(normalized_errors):
```
Require success exactly when no composition error is present.
Construct the report:
```
report = CompositionReport(...)
```
Normalize the cumulative semantic fields and bind the result.

param valid:: Whether composition succeeded (static – a compile-time constant; changing it triggers retracing).
type valid:: bool
param errors:: Composition errors (static – compile-time constants; changing them triggers retracing).
type errors:: Sequence[str]
param available_semantics:: Final available semantics (static – compile-time constants; changing them triggers retracing).
type available_semantics:: Sequence[str]
param destroyed_information:: Cumulative information losses (static – compile-time constants; changing them triggers retracing).
type destroyed_information:: Sequence[str]
param invalidated_claims:: Cumulative invalidated claims (static – compile-time constants; changing them triggers retracing).
type invalidated_claims:: Sequence[str]
param transformation_refs:: Ordered transformation references (static – compile-time constants; changing them triggers retracing).
type transformation_refs:: Sequence[str]
returns:: report – Validated immutable composition report.
rtype:: CompositionReport
raises ValueError:: If a sequence contains blank or duplicate entries, or valid does not agree with whether errors is empty.

Notes

Validation is entirely static because the report contains no numerical JAX leaves.

diffpes.types.make_transformation_contract(transformation_id: str, transformation_version: str, *, requires: Sequence[str] = (), produces: Sequence[str] = (), preserves: Sequence[str] = (), introduces: Sequence[str] = (), destroys: Sequence[str] = (), invalidates_claims: Sequence[str] = (), jax_pure: bool = True) → TransformationContract[source]¶

Create a validated immutable transformation contract.

Validate a permanent identity and freeze each semantic term sequence without hiding duplicates or contradictory information-loss declarations.

See:: TestMakeTransformationContract

Implementation Logic¶

Validate identity:
```
_validate_identity(transformation_id, transformation_version)
```
Require a permanent reverse-DNS-like identifier and semantic version.
Normalize semantic terms:
```
normalized_requires = _normalize_terms(requires, ...)
```
Freeze each semantic sequence while retaining duplicate detection.
Reject contradictory loss declarations:
```
contradictory = set(normalized_preserves) & set(normalized_destroys)
```
A term cannot remain available while also being declared destroyed.
Construct the carrier:
```
contract = TransformationContract(...)
```
Bind and return the immutable static contract.

param transformation_id:: Permanent reverse-DNS-like identifier (static – a compile-time constant; changing it triggers retracing).
type transformation_id:: str
param transformation_version:: Semantic version (static – a compile-time constant; changing it triggers retracing).
type transformation_version:: str
param requires:: Required input semantics (static – compile-time constants; changing them triggers retracing). Default is empty.
type requires:: Sequence[str], default: ()
param produces:: Produced output semantics (static – compile-time constants; changing them triggers retracing). Default is empty.
type produces:: Sequence[str], default: ()
param preserves:: Preserved input semantics (static – compile-time constants; changing them triggers retracing). Default is empty.
type preserves:: Sequence[str], default: ()
param introduces:: Introduced output semantics (static – compile-time constants; changing them triggers retracing). Default is empty.
type introduces:: Sequence[str], default: ()
param destroys:: Destroyed information labels (static – compile-time constants; changing them triggers retracing). Default is empty.
type destroys:: Sequence[str], default: ()
param invalidates_claims:: Invalidated claim identifiers (static – compile-time constants; changing them triggers retracing). Default is empty.
type invalidates_claims:: Sequence[str], default: ()
param jax_pure:: JAX-purity declaration (static – a compile-time constant; changing it triggers retracing). Default is True.
type jax_pure:: bool, default: True
returns:: contract – Validated immutable transformation contract.
rtype:: TransformationContract
raises ValueError:: If the identity or version is invalid; a semantic term is blank or duplicated; or a term is both preserved or introduced and destroyed.

Notes

The factory performs only static validation because all inputs become static fields. It does not convert a traced numerical value to Python control flow.

class diffpes.types.DensityOfStates(energy: Float[Array, 'E'], total_dos: Float[Array, 'E'], fermi_energy: Float[Array, '']) → None[source]¶

Bases: Module

Store density-of-states data in a JAX PyTree.

This type stores total density-of-states (DOS) data from VASP DOSCAR files. Two 1-D arrays share the same energy axis. A scalar value specifies the Fermi energy.

This JAX PyTree passes through jax.jit, jax.grad, jax.vmap, and other JAX transformations. All fields contain JAX arrays and no auxiliary data. Therefore, every field participates in autodiff tracing.

See:

TestDensityOfStates

Variables:

energy (Float[Array, " E"]) – Energy axis in eV.
total_dos (Float[Array, " E"]) – Total density of states.
fermi_energy (Float[Array, " "]) – Fermi level energy in eV.

Notes

JAX stores all three array fields as children and uses no auxiliary data. Therefore, JAX tracing and transformations can access every field.

See also

make_density_of_states: Factory function with validation and float64 casting.

energy: Float[Array, 'E']¶

total_dos: Float[Array, 'E']¶

fermi_energy: Float[Array, '']¶

__init__(energy: Float[Array, 'E'], total_dos: Float[Array, 'E'], fermi_energy: Float[Array, '']) → None¶

class diffpes.types.FullDensityOfStates(energy: Float[Array, 'E'], total_dos_up: Float[Array, 'E'], total_dos_down: Float[Array, 'E'] | None, integrated_dos_up: Float[Array, 'E'], integrated_dos_down: Float[Array, 'E'] | None, pdos: Float[Array, 'A E C'] | None, fermi_energy: Float[Array, ''], natoms: int) → None[source]¶

Bases: Module

Store spin-resolved total and projected DOS data in a JAX PyTree.

This type stores the full DOS data from a VASP DOSCAR file, including spin-resolved total DOS, integrated DOS, and per-atom site-projected DOS. read_doscar returns this type when return_mode="full". This type is the comprehensive counterpart to the simpler DensityOfStates, which only stores a single total-DOS channel.

JAX stores the array fields as PyTree children. JAX tracing can access these required and optional fields. JAX stores natoms as auxiliary data because this Python int is a structural constant.

See:

TestFullDensityOfStates

Variables:

energy (Float[Array, " E"]) – Energy axis in eV, shared by all DOS channels. JAX-traced (differentiable).
total_dos_up (Float[Array, " E"]) – Total DOS for spin-up channel (or the only channel if ISPIN=1). Units are states/eV. JAX-traced (differentiable).
total_dos_down (Optional[Float[Array, " E"]]) – Total DOS for spin-down channel, or None if the calculation is non-spin-polarized (ISPIN=1). Units are states/eV. JAX-traced when present.
integrated_dos_up (Float[Array, " E"]) – Integrated (cumulative) DOS for spin-up channel. The value at the Fermi energy gives the number of electrons for this spin channel. JAX-traced (differentiable).
integrated_dos_down (Optional[Float[Array, " E"]]) – Integrated DOS for spin-down channel, or None if ISPIN=1. JAX-traced when present.
pdos (Optional[Float[Array, "A E C"]]) – Per-atom site-projected DOS. A specifies the atom count. E specifies the energy-point count. C specifies the orbital-column count, which depends on the VASP LORBIT setting. None if no PDOS blocks are present in the DOSCAR file. JAX-traced when present.
fermi_energy (Float[Array, " "]) – Fermi level energy in eV. A 0-D scalar array. JAX-traced (differentiable).
natoms (int) – Number of atoms in the unit cell (static – a compile-time constant; changing it triggers retracing).

Notes

The natoms field is a Python int declared with eqx.field(static=True) rather than as a traced leaf. Changing natoms triggers recompilation of any jit-compiled function that receives this PyTree.

Optional fields (total_dos_down, integrated_dos_down, pdos) may be None for non-spin-polarized calculations or DOSCAR files without projected DOS. JAX handles None leaves transparently.

See also

DensityOfStates: Simplified single-channel variant.
make_full_density_of_states: Factory function with validation and float64 casting.

energy: Float[Array, 'E']¶

total_dos_up: Float[Array, 'E']¶

total_dos_down: Float[Array, 'E'] | None¶

integrated_dos_up: Float[Array, 'E']¶

integrated_dos_down: Float[Array, 'E'] | None¶

pdos: Float[Array, 'A E C'] | None¶

fermi_energy: Float[Array, '']¶

natoms: int¶

__init__(energy: Float[Array, 'E'], total_dos_up: Float[Array, 'E'], total_dos_down: Float[Array, 'E'] | None, integrated_dos_up: Float[Array, 'E'], integrated_dos_down: Float[Array, 'E'] | None, pdos: Float[Array, 'A E C'] | None, fermi_energy: Float[Array, ''], natoms: int) → None¶

diffpes.types.make_density_of_states(energy: Float[Array, 'Ee'], total_dos: Float[Array, 'Ed'], fermi_energy: int | float | complex | Num[Array, ''] = 0.0) → DensityOfStates[source]¶

Create a validated DensityOfStates instance.

The factory validates and normalizes the inputs before it constructs the DensityOfStates PyTree. It casts all numerical inputs to float64 JAX arrays. Downstream JAX transformations therefore use double precision without silent dtype promotion.

See:: TestMakeDensityOfStates

Implementation Logic¶

Prepare the normalized values:
```
energy_arr = jnp.asarray(energy, dtype=jnp.float64)
```
This expression gives the later validation steps a stable shape and dtype.
Apply static validation:
```
energy_arr.shape[0] != dos_arr.shape[0]
```
This predicate rejects invalid structure before JAX traces the numerical checks.
Apply traced validation:
```
~jnp.all(jnp.diff(energy_arr) > 0.0)
```
This predicate remains active during eager and compiled execution.
Return the named instance:
```
return dos
```
The explicit name keeps the implementation and the Returns section synchronized.

param energy:: Energy axis in eV.
type energy:: Float[Array, 'Ee']
param total_dos:: Total density of states.
type total_dos:: Float[Array, 'Ed']
param fermi_energy:: Fermi level in eV. Default is 0.0.
type fermi_energy:: Union[int, float, complex, Num[Array, '']], default: 0.0
returns:: dos – Validated density of states instance.
rtype:: DensityOfStates
raises ValueError:: If the energy and DOS channel lengths disagree.
raises EquinoxRuntimeError:: If the energy axis is not strictly increasing or DOS values are non-finite.

Notes

Static validation raises ValueError before traced construction when the energy and DOS lengths differ. Traced validation uses eqx.error_if and raises EquinoxRuntimeError when the energy axis is not increasing or the DOS contains a non-finite value.

See also

DensityOfStates: The PyTree class constructed by this factory.

diffpes.types.make_full_density_of_states(energy: Float[Array, 'Ee'], total_dos_up: Float[Array, 'Eu'], integrated_dos_up: Float[Array, 'Eiu'], fermi_energy: int | float | complex | Num[Array, ''] = 0.0, total_dos_down: Float[Array, 'Ed'] | None = None, integrated_dos_down: Float[Array, 'Eid'] | None = None, pdos: Float[Array, 'A Ep C'] | None = None, natoms: int = 0) → FullDensityOfStates[source]¶

Create a validated FullDensityOfStates instance.

The factory validates and normalizes full density-of-states states data before constructing a FullDensityOfStates PyTree. The factory casts all present numerical arrays to float64 for numerical stability. It casts optional fields only when they are present. Thus, None continues to identify non-spin-polarized or non-projected calculations.

@jaxtyped(typechecker=beartype) checks the energy dimension E across all provided arrays at call time.

Use this factory when you need the complete DOSCAR output including spin-resolved channels and site-projected DOS. For the simpler case of a single total-DOS channel, prefer make_density_of_states() instead.

See:: TestMakeFullDensityOfStates

Implementation Logic¶

Prepare the normalized values:
```
energy_arr = jnp.asarray(energy, dtype=jnp.float64)
```
This expression gives the later validation steps a stable shape and dtype.
Apply static validation:
```
up_arr.shape[0] != nenergy
```
This predicate rejects invalid structure before JAX traces the numerical checks.
Apply traced validation:
```
~jnp.all(jnp.diff(energy_arr) > 0.0)
```
This predicate remains active during eager and compiled execution.
Return the named instance:
```
return dos
```
The explicit name keeps the implementation and the Returns section synchronized.

param energy:: Energy axis in eV.
type energy:: Float[Array, 'Ee']
param total_dos_up:: Spin-up total DOS (states/eV).
type total_dos_up:: Float[Array, 'Eu']
param integrated_dos_up:: Spin-up integrated (cumulative) DOS.
type integrated_dos_up:: Float[Array, 'Eiu']
param fermi_energy:: Fermi level in eV. Default is 0.0.
type fermi_energy:: Union[int, float, complex, Num[Array, '']], default: 0.0
param total_dos_down:: Spin-down total DOS. Default is None (ISPIN=1).
type total_dos_down:: Optional[Float[Array, 'Ed']], default: None
param integrated_dos_down:: Spin-down integrated DOS. Default is None.
type integrated_dos_down:: Optional[Float[Array, 'Eid']], default: None
param pdos:: Per-atom site-projected DOS with C orbital columns. Default is None.
type pdos:: Optional[Float[Array, 'A Ep C']], default: None
param natoms:: Number of atoms in the unit cell (static – a compile-time constant; changing it triggers retracing). Default is 0.
type natoms:: int, default: 0
returns:: dos – Validated full density of states with all non-None arrays in float64.
rtype:: FullDensityOfStates
raises ValueError:: If a total, integrated, or projected DOS channel disagrees with the energy-axis length.
raises EquinoxRuntimeError:: If the energy axis is not strictly increasing or any DOS channel contains non-finite values.

Notes

Static validation raises ValueError before traced construction when any present DOS channel disagrees with the energy-axis length. Traced validation uses eqx.error_if and raises EquinoxRuntimeError when the energy axis is not increasing or a present DOS channel is non-finite.

See also

make_density_of_states: Factory for the simplified variant.
FullDensityOfStates: The PyTree class constructed by this factory.

class diffpes.types.ExperimentGeometry(photon_energy_ev: Float[Array, ''], polarization: Complex[Array, '3'], incidence_theta: Float[Array, ''], incidence_phi: Float[Array, ''], sample_azimuth: Float[Array, ''], work_function_ev: Float[Array, ''], inner_potential_ev: Float[Array, ''], temperature_k: Float[Array, ''], energy_resolution_ev: Float[Array, ''], momentum_resolution_inv_ang: Float[Array, ''], mean_free_path_ang: Float[Array, ''], slit: str) → None[source]¶

Bases: Module

Store the geometry and resolution of an ARPES experiment.

This PyTree groups the traced quantities that define one ARPES measurement. JAX differentiates every numerical field. The static slit field selects the detector frame and causes retracing when it changes.

See:

TestExperimentGeometry

Variables:

photon_energy_ev (Float[Array, ""]) – Photon energy in eV.
polarization (Complex[Array, "3"]) – Unit polarization vector in the laboratory frame.
incidence_theta (Float[Array, ""]) – Incidence angle from the surface normal in radians.
incidence_phi (Float[Array, ""]) – Azimuthal incidence angle in radians.
sample_azimuth (Float[Array, ""]) – Sample rotation about the surface normal in radians.
work_function_ev (Float[Array, ""]) – Work function in eV.
inner_potential_ev (Float[Array, ""]) – Inner potential in eV.
temperature_k (Float[Array, ""]) – Sample temperature in kelvin.
energy_resolution_ev (Float[Array, ""]) – Full width at half maximum of the energy resolution in eV.
momentum_resolution_inv_ang (Float[Array, ""]) – Full width at half maximum of the momentum resolution in 1/Angstrom.
mean_free_path_ang (Float[Array, ""]) – Mean free path of the photoelectron in Angstrom.
slit (str) – Detector slit orientation. This field is static. A change causes JAX to retrace the receiving function.

Notes

The traced fields support calibration and geometry inversion. The normalized polarization removes its intensity-scale gauge.

See also

make_experiment_geometry: Create a validated geometry for an ARPES experiment.

photon_energy_ev: Float[Array, '']¶

polarization: Complex[Array, '3']¶

incidence_theta: Float[Array, '']¶

incidence_phi: Float[Array, '']¶

sample_azimuth: Float[Array, '']¶

work_function_ev: Float[Array, '']¶

inner_potential_ev: Float[Array, '']¶

temperature_k: Float[Array, '']¶

energy_resolution_ev: Float[Array, '']¶

momentum_resolution_inv_ang: Float[Array, '']¶

mean_free_path_ang: Float[Array, '']¶

slit: str¶

__init__(photon_energy_ev: Float[Array, ''], polarization: Complex[Array, '3'], incidence_theta: Float[Array, ''], incidence_phi: Float[Array, ''], sample_azimuth: Float[Array, ''], work_function_ev: Float[Array, ''], inner_potential_ev: Float[Array, ''], temperature_k: Float[Array, ''], energy_resolution_ev: Float[Array, ''], momentum_resolution_inv_ang: Float[Array, ''], mean_free_path_ang: Float[Array, ''], slit: str) → None¶

diffpes.types.make_experiment_geometry(photon_energy_ev: float | Float[Array, ''], polarization: Complex[Array, '3'], incidence_theta: float | Float[Array, ''] = 0.0, incidence_phi: float | Float[Array, ''] = 0.0, sample_azimuth: float | Float[Array, ''] = 0.0, work_function_ev: float | Float[Array, ''] = 4.0, inner_potential_ev: float | Float[Array, ''] = 10.0, temperature_k: float | Float[Array, ''] = 10.0, energy_resolution_ev: float | Float[Array, ''] = 0.02, momentum_resolution_inv_ang: float | Float[Array, ''] = 0.01, mean_free_path_ang: float | Float[Array, ''] = 10.0, slit: str = 'H') → ExperimentGeometry[source]¶

Create a validated geometry for an ARPES experiment.

The factory converts all numerical inputs to JAX arrays. It then fixes the intensity-scale gauge by normalizing the complex polarization vector.

See:: TestMakeExperimentGeometry

Implementation Logic¶

Convert the inputs:
```
photon_energy = jnp.asarray(photon_energy_ev, dtype=jnp.float64)
```
The conversion gives each traced field a stable JAX dtype.
Validate the traced values:
```
photon_energy = eqx.error_if(photon_energy, invalid, message)
```
The runtime checks remain active during compiled execution.
Normalize the polarization:
```
normalized_polarization = checked_polarization / safe_norm
```
The safe denominator keeps the invalid zero-norm branch finite.
Return the named carrier:
```
return geometry
```
The carrier preserves each valid numerical input as a traced leaf.

param photon_energy_ev:: Photon energy in eV.
type photon_energy_ev:: Union[float, Float[Array, '']]
param polarization:: Nonzero complex polarization vector in the laboratory frame.
type polarization:: Complex[Array, '3']
param incidence_theta:: Incidence angle from the surface normal in radians. Default is 0.0.
type incidence_theta:: Union[float, Float[Array, '']], default: 0.0
param incidence_phi:: Azimuthal incidence angle in radians. Default is 0.0.
type incidence_phi:: Union[float, Float[Array, '']], default: 0.0
param sample_azimuth:: Sample rotation about the surface normal in radians. Default is 0.0.
type sample_azimuth:: Union[float, Float[Array, '']], default: 0.0
param work_function_ev:: Work function in eV. Default is 4.0.
type work_function_ev:: Union[float, Float[Array, '']], default: 4.0
param inner_potential_ev:: Inner potential in eV. Default is 10.0.
type inner_potential_ev:: Union[float, Float[Array, '']], default: 10.0
param temperature_k:: Sample temperature in kelvin. Default is 10.0.
type temperature_k:: Union[float, Float[Array, '']], default: 10.0
param energy_resolution_ev:: Energy-resolution width in eV. Default is 0.02.
type energy_resolution_ev:: Union[float, Float[Array, '']], default: 0.02
param momentum_resolution_inv_ang:: Momentum-resolution width in 1/Angstrom. Default is 0.01.
type momentum_resolution_inv_ang:: Union[float, Float[Array, '']], default: 0.01
param mean_free_path_ang:: Mean free path in Angstrom. Default is 10.0.
type mean_free_path_ang:: Union[float, Float[Array, '']], default: 10.0
param slit:: Detector slit orientation. This value is static. A change causes retracing. Default is "H".
type slit:: str, default: "H"
returns:: geometry – Validated experiment geometry with a unit polarization vector.
rtype:: ExperimentGeometry
raises ValueError:: If slit is not "H" or "V".
raises EquinoxRuntimeError:: If an input is non-finite or violates its physical range. The factory also rejects a zero polarization vector.

Notes

The normalization is differentiable for every accepted polarization. The factory assigns no derivative at the rejected zero vector.

class diffpes.types.CrystalGeometry(lattice: Float[Array, '3 3'], reciprocal: Float[Array, '3 3'], positions: Float[Array, 'N 3'], species: tuple[str, ...]) → None[source]¶

Bases: Module

Store VASP POSCAR crystal geometry in a JAX PyTree.

This type contains the complete crystal structure from a VASP POSCAR file. It contains real-space lattice vectors, the reciprocal lattice, fractional positions, and per-atom species. Together, these fields describe the periodic crystal for an ARPES simulation.

JAX stores the numerical fields as PyTree children. JAX tracing can access these fields. JAX stores the species tuple as auxiliary data because JAX cannot trace Python strings.

See:

TestCrystalGeometry

Variables:

lattice (Float[Array, "3 3"]) – Real-space lattice vectors as rows (angstroms).
reciprocal (Float[Array, "3 3"]) – Reciprocal lattice vectors as rows (1/angstroms).
positions (Float[Array, "N 3"]) – Fractional atomic positions.
species (tuple[str, ]) – Per-atom species symbols (static – compile-time constants; changing them triggers retracing). An empty tuple records unknown species.

Notes

The species field is a tuple of Python strings declared with eqx.field(static=True) rather than as a traced leaf. Changing it triggers recompilation of any jit-compiled function that receives this PyTree.

See also

make_crystal_geometry: Factory function with validation, float64 casting, and automatic reciprocal lattice computation.

lattice: Float[Array, '3 3']¶

reciprocal: Float[Array, '3 3']¶

positions: Float[Array, 'N 3']¶

species: tuple[str, ...]¶

__init__(lattice: Float[Array, '3 3'], reciprocal: Float[Array, '3 3'], positions: Float[Array, 'N 3'], species: tuple[str, ...]) → None¶

diffpes.types.make_crystal_geometry(lattice: Float[Array, '3 3'] | list[list[int | float | complex | Num[Array, '']]], positions: Float[Array, 'N 3'], species: tuple[str, ...]) → CrystalGeometry[source]¶

Create a validated CrystalGeometry instance.

The factory validates the real-space lattice and fractional positions. It computes the reciprocal lattice and casts numerical arrays to float64.

See:: TestMakeCrystalGeometry

Implementation Logic¶

Prepare the normalized values:
```
lattice_arr = jnp.asarray(lattice, dtype=jnp.float64)
```
This expression gives the later validation steps a stable shape and dtype.
Apply static validation:
```
species and len(species) != positions_arr.shape[0]
```
This predicate rejects invalid structure before JAX traces the numerical checks.
Apply traced validation:
```
~jnp.all(jnp.isfinite(positions_arr))
```
This predicate remains active during eager and compiled execution.
Return the named instance:
```
return geometry
```
The explicit name keeps the implementation and the Returns section synchronized.

param lattice:: Real-space lattice vectors as rows (angstroms).
type lattice:: Union[Float[Array, '3 3'], list[list[Union[int, float, complex, Num[Array, '']]]]]
param positions:: Fractional atomic positions.
type positions:: Float[Array, 'N 3']
param species:: Per-atom species symbols (static – compile-time constants; changing them triggers retracing). Use an empty tuple when the source does not identify species.
type species:: tuple[str, ...]
returns:: geometry – Validated crystal geometry instance with the reciprocal lattice pre-computed.
rtype:: CrystalGeometry
raises ValueError:: If a non-empty species tuple differs from the number of positions. Empty species remain valid for VASP 4 files without a species line.
raises EquinoxRuntimeError:: If positions or lattice entries are non-finite. The function also rejects a left-handed lattice or a lattice that violates the numerical stability limits.

Notes

Static validation raises ValueError before traced construction when species and position counts disagree. Traced validation uses eqx.error_if and raises EquinoxRuntimeError for non-finite data or a lattice that fails the orientation and conditioning limits.

See also

CrystalGeometry: The PyTree class constructed by this factory.
_compute_reciprocal_lattice: Cross-product formula used to derive the reciprocal lattice.

class diffpes.types.CertificateDiff(scientific: tuple[str, ...], numerical: tuple[str, ...], environment: tuple[str, ...], audit: tuple[str, ...]) → None[source]¶

Bases: Module

Store categorized differences between two forward certificates.

The record groups differences by scientific meaning, numerical evidence, execution environment, and audit metadata. It does not run the model again.

See:

TestCertificateDiff

Variables:

scientific (tuple[str, ]) – Differing scientific fields (static – compile-time constants; changing them triggers retracing).
numerical (tuple[str, ]) – Differing numerical-evidence fields (static – compile-time constants; changing them triggers retracing).
environment (tuple[str, ]) – Differing execution-environment fields (static – compile-time constants; changing them triggers retracing).
audit (tuple[str, ]) – Differing audit fields (static – compile-time constants; changing them triggers retracing).

Notes

Inspection compares persisted metadata only. This carrier has no numerical leaves and does not reevaluate or differentiate a forward model.

See also

make_certificate_diff: Construct a validated certificate-difference record.

scientific: tuple[str, ...]¶

numerical: tuple[str, ...]¶

environment: tuple[str, ...]¶

audit: tuple[str, ...]¶

property identical: bool¶

Return whether the record contains no categorized difference.

See:: TestCertificateDiff
Returns:: identical – Whether every difference category is empty.
Return type:: bool

property summary: str¶

Return a one-line categorized comparison summary.

See:: TestCertificateDiff
Returns:: summary – Human-readable list of nonempty difference categories.
Return type:: str

__init__(scientific: tuple[str, ...], numerical: tuple[str, ...], environment: tuple[str, ...], audit: tuple[str, ...]) → None¶

diffpes.types.make_certificate_diff(*, scientific: tuple[str, ...] = (), numerical: tuple[str, ...] = (), environment: tuple[str, ...] = (), audit: tuple[str, ...] = ()) → CertificateDiff[source]¶

Construct a validated certificate-difference record.

Validate and freeze field names in each comparison category.

See:: TestMakeCertificateDiff

Implementation Logic¶

Validate category names:
```
scientific=_difference_names(scientific, "scientific")
```
Require immutable tuples of nonempty field names in every category.
Construct the difference:
```
difference = CertificateDiff(...)
```
Bind and return the categorized comparison carrier.

param scientific:: Differing scientific fields (static – compile-time constants; changing them triggers retracing). Default is empty.
type scientific:: tuple[str, ...], default: ()
param numerical:: Differing numerical fields (static – compile-time constants; changing them triggers retracing). Default is empty.
type numerical:: tuple[str, ...], default: ()
param environment:: Differing environment fields (static – compile-time constants; changing them triggers retracing). Default is empty.
type environment:: tuple[str, ...], default: ()
param audit:: Differing audit fields (static – compile-time constants; changing them triggers retracing). Default is empty.
type audit:: tuple[str, ...], default: ()
returns:: difference – Validated immutable certificate difference.
rtype:: CertificateDiff
raises ValueError:: If a category is not a tuple of nonempty field names.

Notes

Validation is static and does not introduce a gradient path.

class diffpes.types.KGrid(kpoints: Float[Array, 'n_k 3'], kz: Float[Array, ''] | None, photon_energy_axis_ev: Float[Array, 'n_rows'] | None, mesh_shape: tuple[int, int]) → None[source]¶

Bases: Module

Store a fixed-shape raster in fractional k-space.

The flattened coordinates remain traced for geometry inversion. The static mesh shape lets compiled consumers restore the raster without a data-dependent shape operation.

See:

TestKGrid

Variables:

kpoints (Float[Array, "n_k 3"]) – Flattened fractional k-points.
kz (Optional[Float[Array, ""]]) – Fixed Cartesian out-of-plane momentum in 1/Angstrom. None marks a grid with varying or unspecified out-of-plane momentum.
photon_energy_axis_ev (Optional[Float[Array, " n_rows"]]) – Photon energy for each raster row in eV. None marks a grid without a photon-energy axis.
mesh_shape (tuple[int, int]) – Raster shape as (n_rows, n_cols). This field is static. A change causes JAX to retrace the receiving function.

See also

make_kgrid: Create a validated fixed-shape k-space raster.

kpoints: Float[Array, 'n_k 3']¶

kz: Float[Array, ''] | None¶

photon_energy_axis_ev: Float[Array, 'n_rows'] | None¶

mesh_shape: tuple[int, int]¶

__init__(kpoints: Float[Array, 'n_k 3'], kz: Float[Array, ''] | None, photon_energy_axis_ev: Float[Array, 'n_rows'] | None, mesh_shape: tuple[int, int]) → None¶

class diffpes.types.KPath(kpoints: Float[Array, 'n_k 3'], kz: Float[Array, ''] | None, labels: tuple[str, ...], label_indices: tuple[int, ...], n_per_segment: int) → None[source]¶

Bases: Module

Store a generated path through fractional k-space.

This PyTree carries the dense path from a k-space builder. The numerical coordinates remain traced. Labels, their indices, and the segment size remain static because they describe the compiled path shape.

See:

TestKPath

Variables:

kpoints (Float[Array, "n_k 3"]) – Fractional k-points.
kz (Optional[Float[Array, ""]]) – Fixed Cartesian out-of-plane momentum in 1/Angstrom. None means that the path has no separate fixed value.
labels (tuple[str, ]) – Labels for the path anchors. This field is static. A change causes JAX to retrace the receiving function.
label_indices (tuple[int, ]) – Indices for the labels. This field is static. A change causes JAX to retrace the receiving function.
n_per_segment (int) – Number of points in each segment. This field is static. A change causes JAX to retrace the receiving function.

See also

KPathInfo: Store k-point path metadata in a JAX PyTree.
make_kpath: Create a validated path through fractional k-space.

kpoints: Float[Array, 'n_k 3']¶

kz: Float[Array, ''] | None¶

labels: tuple[str, ...]¶

label_indices: tuple[int, ...]¶

n_per_segment: int¶

__init__(kpoints: Float[Array, 'n_k 3'], kz: Float[Array, ''] | None, labels: tuple[str, ...], label_indices: tuple[int, ...], n_per_segment: int) → None¶

class diffpes.types.KPathInfo(num_kpoints: Int[Array, ''], label_indices: Int[Array, 'L'], points_per_segment: Int[Array, ''], segments: Int[Array, ''], kpoints: Float[Array, 'K 3'] | None, weights: Float[Array, 'K'] | None, grid: Int[Array, '3'] | None, shift: Float[Array, '3'] | None, mode: str, labels: tuple[str, ...], comment: str, coordinate_mode: str) → None[source]¶

Bases: Module

Store k-point path metadata in a JAX PyTree.

This type stores Brillouin-zone path information from VASP KPOINTS files. It includes plotting labels and label indices. Mode-specific metadata includes an automatic grid and shift, explicit k-points and weights, and line-mode segment endpoints.

JAX stores numerical fields as PyTree children. It stores string metadata as auxiliary data because JAX cannot trace Python strings.

See:

TestKPathInfo

Variables:

num_kpoints (Int[Array, " "]) – Total number of k-points in the path (line mode) or header count (explicit mode).
label_indices (Int[Array, " L"]) – Indices of symmetry points along the path.
points_per_segment (Int[Array, " "]) – Raw integer from line 2 of KPOINTS (line mode: points per segment).
segments (Int[Array, " "]) – Number of line segments in line mode.
kpoints (Optional[Float[Array, "K 3"]]) – Mode-specific k-points: line mode -> segment endpoints (segments + 1), explicit mode -> listed k-points, automatic mode -> None.
weights (Optional[Float[Array, " K"]]) – Explicit-mode per-k-point weights (None otherwise).
grid (Optional[Int[Array, " 3"]]) – Automatic-mode Monkhorst-Pack/Gamma grid (None otherwise).
shift (Optional[Float[Array, " 3"]]) – Automatic-mode grid shift (None otherwise).
mode (str) – KPOINTS file mode (Automatic, Line-mode, Explicit; static – a compile-time constant; changing it triggers retracing).
labels (tuple[str, ]) – Symmetry point labels (e.g., Gamma, M, K; static – compile-time constants; changing them triggers retracing).
comment (str) – Raw comment from KPOINTS line 1 (static – a compile-time constant; changing it triggers retracing).
coordinate_mode (str) – Coordinate/scheme line metadata: line/explicit -> Reciprocal/Cartesian line, automatic -> scheme line (e.g., Monkhorst-Pack; static – a compile-time constant; changing it triggers retracing).

Notes

eqx.field(static=True) declares string metadata as auxiliary data instead of traced leaves. Changing any static value triggers recompilation of any jit-compiled function that receives this PyTree.

See also

make_kpath_info: Factory function with validation and int32 casting.

num_kpoints: Int[Array, '']¶

label_indices: Int[Array, 'L']¶

points_per_segment: Int[Array, '']¶

segments: Int[Array, '']¶

kpoints: Float[Array, 'K 3'] | None¶

weights: Float[Array, 'K'] | None¶

grid: Int[Array, '3'] | None¶

shift: Float[Array, '3'] | None¶

mode: str¶

labels: tuple[str, ...]¶

comment: str¶

coordinate_mode: str¶

__init__(num_kpoints: Int[Array, ''], label_indices: Int[Array, 'L'], points_per_segment: Int[Array, ''], segments: Int[Array, ''], kpoints: Float[Array, 'K 3'] | None, weights: Float[Array, 'K'] | None, grid: Int[Array, '3'] | None, shift: Float[Array, '3'] | None, mode: str, labels: tuple[str, ...], comment: str, coordinate_mode: str) → None¶

diffpes.types.make_kgrid(kpoints: Float[Array, 'n_k 3'], mesh_shape: tuple[int, int], kz: float | Float[Array, ''] | None = None, photon_energy_axis_ev: Float[Array, 'n_rows'] | None = None) → KGrid[source]¶

Create a validated fixed-shape k-space raster.

The factory validates the static raster shape and traced numerical data. It preserves an optional photon-energy axis for a photon-energy map.

See:: TestMakeKGrid

Implementation Logic¶

Validate the raster shape:
```
mesh_shape[0] * mesh_shape[1] == kpoints.shape[0]
```
This check keeps the flattened point count consistent with the raster.
Validate the traced values:
```
checked_kpoints = eqx.error_if(kpoints_array, invalid, message)
```
The runtime checks remain active during compiled execution.
Return the named grid:
```
return kgrid
```
The result keeps only the raster shape static.

param kpoints:: Flattened fractional k-points.
type kpoints:: Float[Array, 'n_k 3']
param mesh_shape:: Raster shape as (n_rows, n_cols). This value is static. A change causes retracing.
type mesh_shape:: tuple[int, int]
param kz:: Fixed Cartesian out-of-plane momentum in 1/Angstrom. Default is None.
type kz:: Union[float, Float[Array, ''], None], default: None
param photon_energy_axis_ev:: Photon energy for each raster row in eV. Default is None.
type photon_energy_axis_ev:: Optional[Float[Array, 'n_rows']], default: None
returns:: kgrid – Validated grid with traced fractional coordinates.
rtype:: KGrid
raises ValueError:: If the mesh dimensions are not positive or their product is wrong. The factory also rejects a photon-energy axis with the wrong length.
raises EquinoxRuntimeError:: If traced data is non-finite or a photon energy is not positive.

Notes

The mesh shape controls array reshaping and remains static. All numerical fields carry gradients when they are present.

diffpes.types.make_kpath(kpoints: Float[Array, 'n_k 3'], labels: tuple[str, ...] = (), label_indices: tuple[int, ...] = (), n_per_segment: int = 1, kz: float | Float[Array, ''] | None = None) → KPath[source]¶

Create a validated path through fractional k-space.

The factory validates the static plotting metadata and the traced path coordinates. It preserves an optional fixed Cartesian out-of-plane value.

See:: TestMakeKPath

Implementation Logic¶

Validate the path structure:
```
label_indices[-1] < kpoints.shape[0]
```
The static checks keep labels within the fixed path shape.
Validate the traced values:
```
checked_kpoints = eqx.error_if(kpoints_array, invalid, message)
```
The runtime checks remain active during compiled execution.
Return the named path:
```
return kpath
```
The result keeps only the plotting metadata static.

param kpoints:: Fractional k-points.
type kpoints:: Float[Array, 'n_k 3']
param labels:: Labels for the path anchors. This value is static. A change causes retracing. Default is an empty tuple.
type labels:: tuple[str, ...], default: ()
param label_indices:: Indices for the labels. This value is static. A change causes retracing. Default is an empty tuple.
type label_indices:: tuple[int, ...], default: ()
param n_per_segment:: Number of points in each segment. This value is static. A change causes retracing. Default is 1.
type n_per_segment:: int, default: 1
param kz:: Fixed Cartesian out-of-plane momentum in 1/Angstrom. Default is None.
type kz:: Union[float, Float[Array, ''], None], default: None
returns:: kpath – Validated path with traced fractional coordinates.
rtype:: KPath
raises ValueError:: If the static metadata has invalid lengths, indices, or segment size.
raises EquinoxRuntimeError:: If a traced k-point or fixed out-of-plane value is non-finite.

Notes

The k-points and optional kz value carry gradients. The labels, indices, and segment size do not carry gradients.

diffpes.types.make_kpath_info(num_kpoints: int | Int[Array, ''], label_indices: Int[Array, 'L'] | list[int], points_per_segment: int | Int[Array, ''] = 0, segments: int | Int[Array, ''] = 0, kpoints: Float[Array, 'K 3'] | None = None, weights: Float[Array, 'K'] | None = None, grid: Int[Array, '3'] | list[int] | None = None, shift: Float[Array, '3'] | None = None, mode: str = 'Line-mode', labels: tuple[str, ...] = (), comment: str = '', coordinate_mode: str = '') → KPathInfo[source]¶

Create a validated KPathInfo instance.

The factory validates and normalizes raw k-path metadata before constructing a KPathInfo PyTree. Integer The factory casts integer scalars and arrays to int32. It casts float arrays to float64. The factory casts optional fields only when they are present. Thus, None continues to identify modes that omit them.

@jaxtyped(typechecker=beartype) checks shape constraints at call time. The factory passes string metadata unchanged and stores it as PyTree auxiliary data.

Use this factory when constructing KPathInfo from parsed KPOINTS data or when building synthetic k-paths for testing. The factory ensures consistent dtypes and handles the three KPOINTS modes (Line-mode, Explicit, Automatic) through the optional fields.

See:: TestMakeKPathInfo

Implementation Logic¶

Prepare the normalized values:
```
nkpts_arr = jnp.asarray(num_kpoints, dtype=jnp.int32)
```
This expression gives the later validation steps a stable shape and dtype.
Apply static validation:
```
mode not in _KPATH_MODES
```
This predicate rejects invalid structure before JAX traces the numerical checks.
Apply traced validation:
```
~(nkpts_arr >= 0)
```
This predicate remains active during eager and compiled execution.
Return the named instance:
```
return kpath
```
The explicit name keeps the implementation and the Returns section synchronized.

param num_kpoints:: Total number of k-points along the path.
type num_kpoints:: Union[int, Int[Array, '']]
param label_indices:: Indices of symmetry points along the path.
type label_indices:: Union[Int[Array, 'L'], list[int]]
param points_per_segment:: Raw value from line 2 of KPOINTS. Default is 0.
type points_per_segment:: Union[int, Int[Array, '']], default: 0
param segments:: Number of path segments in line mode. Default is 0.
type segments:: Union[int, Int[Array, '']], default: 0
param kpoints:: Mode-specific k-point coordinates. Default is None.
type kpoints:: Optional[Float[Array, 'K 3']], default: None
param weights:: Explicit-mode weights. Default is None.
type weights:: Optional[Float[Array, 'K']], default: None
param grid:: Automatic-mode MP/Gamma grid. Default is None.
type grid:: Union[Int[Array, '3'], list[int], None], default: None
param shift:: Automatic-mode grid shift. Default is None.
type shift:: Optional[Float[Array, '3']], default: None
param mode:: KPOINTS file mode (static – a compile-time constant; changing it triggers retracing). Default is "Line-mode".
type mode:: str, default: "Line-mode"
param labels:: Symmetry point labels (static – compile-time constants; changing them triggers retracing). Default is empty tuple.
type labels:: tuple[str, ...], default: ()
param comment:: KPOINTS comment line (static – a compile-time constant; changing it triggers retracing). Default is empty string.
type comment:: str, default: ""
param coordinate_mode:: Coordinate/scheme line metadata (static – a compile-time constant; changing it triggers retracing). Default is empty string.
type coordinate_mode:: str, default: ""
returns:: kpath – Validated k-path info instance.
rtype:: KPathInfo
raises ValueError:: If mode has an unsupported value. The function also rejects incomplete or inconsistent line-mode metadata.
raises EquinoxRuntimeError:: If num_kpoints is negative, a traced line-mode segment count is less than one, or supplied kpoints, weights, or shift values are non-finite.

Notes

Static validation raises ValueError before traced construction for an unsupported mode or inconsistent line-mode metadata. Traced validation uses eqx.error_if and raises EquinoxRuntimeError for negative k-point counts, invalid traced segment counts, and non-finite optional floating-point arrays.

See also

KPathInfo: The PyTree class constructed by this factory.

class diffpes.types.PolarizationConfig(theta: Float[Array, ''], phi: Float[Array, ''], polarization_angle: Float[Array, ''], polarization_type: str) → None[source]¶

Bases: Module

Store photon-polarization geometry in a JAX PyTree.

This type describes the photon-polarization state and incidence geometry for an ARPES experiment. Three angular fields define the light direction and polarization orientation. The string polarization_type selects an optics convention: s-pol, p-pol, circular, arbitrary linear, or unpolarized.

This class is an immutable equinox.Module PyTree. JAX stores the three angular fields as array children that support autodiff. It stores polarization_type as auxiliary data because JAX cannot trace strings.

See:

TestPolarizationConfig

Variables:

theta (Float[Array, " "]) – Incident angle from surface normal in radians.
phi (Float[Array, " "]) – In-plane azimuthal angle in radians.
polarization_angle (Float[Array, " "]) – Arbitrary linear polarization angle in radians.
polarization_type (str) – One of LVP, LHP, RCP, LCP, LAP, unpolarized (static – a compile-time constant; changing it triggers retracing).

Notes

Equinox derives the PyTree structure from the annotated fields. The polarization_type field is a Python string and is therefore stored as auxiliary data rather than as a child. JAX treats auxiliary data as a compile-time constant: changing polarization_type triggers recompilation of any jit-compiled function that receives this PyTree. This is intentional because the polarization type selects different code branches in the matrix-element calculation.

See also

make_polarization_config: Factory function with validation and float64 casting.

theta: Float[Array, '']¶

phi: Float[Array, '']¶

polarization_angle: Float[Array, '']¶

polarization_type: str¶

__init__(theta: Float[Array, ''], phi: Float[Array, ''], polarization_angle: Float[Array, ''], polarization_type: str) → None¶

class diffpes.types.SimulationParams(energy_min: Float[Array, ''], energy_max: Float[Array, ''], sigma: Float[Array, ''], gamma: Float[Array, ''], temperature: Float[Array, ''], photon_energy: Float[Array, ''], fidelity: int) → None[source]¶

Bases: Module

Store ARPES simulation parameters in a JAX PyTree.

This type collects the scalar physical parameters that control an ARPES simulation. These parameters define the energy window, discretization, broadening widths, sample temperature, and photon energy. The spectral function and Fermi-Dirac weighting routines use these parameters.

This class is an immutable equinox.Module PyTree. All JAX stores the float-valued fields as array children that support autodiff. It stores the Python int field fidelity as auxiliary data. This field controls array shapes and must remain concrete during compilation.

See:

TestSimulationParams

Variables:

energy_min (Float[Array, " "]) – Lower bound of energy window in eV.
energy_max (Float[Array, " "]) – Upper bound of energy window in eV.
sigma (Float[Array, " "]) – Gaussian instrumental broadening width in eV.
gamma (Float[Array, " "]) – Lorentzian lifetime broadening half-width in eV.
temperature (Float[Array, " "]) – Sample temperature in Kelvin.
photon_energy (Float[Array, " "]) – Incident photon energy in eV.
fidelity (int) – Number of points along the energy axis (static – a compile-time constant; changing it triggers retracing).

Notes

Equinox derives the PyTree structure from the annotated fields. The fidelity field is a Python int (not a JAX array) and is therefore stored as auxiliary data rather than as a child. JAX treats auxiliary data as a compile-time constant: changing fidelity triggers recompilation of any jit-compiled function that receives this PyTree. This is intentional because fidelity determines the length of the energy axis array and JAX requires static shapes.

See also

make_simulation_params: Factory function with validation and float64 casting.

energy_min: Float[Array, '']¶

energy_max: Float[Array, '']¶

sigma: Float[Array, '']¶

gamma: Float[Array, '']¶

temperature: Float[Array, '']¶

photon_energy: Float[Array, '']¶

fidelity: int¶

__init__(energy_min: Float[Array, ''], energy_max: Float[Array, ''], sigma: Float[Array, ''], gamma: Float[Array, ''], temperature: Float[Array, ''], photon_energy: Float[Array, ''], fidelity: int) → None¶

diffpes.types.make_expanded_simulation_params(eigenbands: Float[Array, 'K B'], fidelity: int = 25000, sigma: float | Float[Array, ''] = 0.04, gamma: float | Float[Array, ''] = 0.1, temperature: float | Float[Array, ''] = 15.0, photon_energy: float | Float[Array, ''] = 11.0, energy_padding: float | Float[Array, ''] = 1.0) → SimulationParams[source]¶

Build simulation parameters with auto-derived energy window.

The factory constructs a SimulationParams PyTree. It derives the energy window from the actual band energies instead of fixed defaults. The window spans [min(eigenbands) - energy_padding, max(eigenbands) + energy_padding], ensuring every band falls within the simulated range.

See:: TestMakeExpandedSimulationParams

Implementation Logic¶

Prepare the normalized values:
```
bands_arr = jnp.asarray(eigenbands, dtype=jnp.float64)
```
This expression gives the later validation steps a stable shape and dtype.
Apply static validation:
```
bands_arr.size == 0
```
This predicate rejects invalid structure before JAX traces the numerical checks.
Apply traced validation:
```
~jnp.all(jnp.isfinite(bands_arr))
```
This predicate remains active during eager and compiled execution.
Return the named instance:
```
return params
```
The explicit name keeps the implementation and the Returns section synchronized.

param eigenbands:: Band eigenvalues in eV, shape (n_kpoints, n_bands). Used only to derive energy_min and energy_max.
type eigenbands:: Float[Array, 'K B']
param fidelity:: Number of points in the energy axis (static – a compile-time constant; changing it triggers retracing). Default is 25000.
type fidelity:: int, default: 25000
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. Default is 0.1.
type gamma:: Union[float, Float[Array, '']], default: 0.1
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 energy_padding:: Symmetric padding around band extrema in eV. Default is 1.
type energy_padding:: Union[float, Float[Array, '']], default: 1.0
returns:: params – Simulation parameters with data-derived energy window.
rtype:: SimulationParams
raises ValueError:: If eigenbands is empty, resolved statically from its shape.
raises EquinoxRuntimeError:: If eigenbands or energy_padding is non-finite, or if energy_padding is negative under eager or compiled execution.

Notes

Static validation raises ValueError before tracing when the band array is empty. Traced validation uses eqx.error_if and raises EquinoxRuntimeError for non-finite bands or for padding that is non-finite or negative. The delegated factory applies its own documented static and traced validation contract.

See also

make_simulation_params: Lower-level factory with explicit energy bounds.

diffpes.types.make_polarization_config(theta: float | Float[Array, ''] = 0.7854, phi: float | Float[Array, ''] = 0.0, polarization_angle: float | Float[Array, ''] = 0.0, polarization_type: str = 'unpolarized') → PolarizationConfig[source]¶

Create a validated PolarizationConfig instance.

The factory validates and normalizes the inputs before it constructs the PolarizationConfig PyTree. It casts angular parameters to 0-D float64 JAX arrays. It keeps polarization_type as a Python string in the PyTree auxiliary data.

The default incidence angle theta = 0.7854 rad equals pi/4, or approximately 45 degrees. This common geometry provides sensitivity to in-plane and out-of-plane orbital components.

See:: TestMakePolarizationConfig

Implementation Logic¶

Prepare the normalized values:
```
theta_arr = jnp.asarray(theta, dtype=jnp.float64)
```
This expression gives the later validation steps a stable shape and dtype.
Apply static validation:
```
polarization_type not in _POLARIZATION_TYPES
```
This predicate rejects invalid structure before JAX traces the numerical checks.
Apply traced validation:
```
~jnp.isfinite(theta_arr)
```
This predicate remains active during eager and compiled execution.
Return the named instance:
```
return config
```
The explicit name keeps the implementation and the Returns section synchronized.

param theta:: Incident angle in radians. Default is pi/4 ~ 0.7854 rad (45 degrees).
type theta:: Union[float, Float[Array, '']], default: 0.7854
param phi:: Azimuthal angle in radians. Default is 0.
type phi:: Union[float, Float[Array, '']], default: 0.0
param polarization_angle:: Linear polarization angle in radians. Default is 0.
type polarization_angle:: Union[float, Float[Array, '']], default: 0.0
param polarization_type:: Polarization type (static – a compile-time constant; changing it triggers retracing). Default is "unpolarized".
type polarization_type:: str, default: "unpolarized"
returns:: config – Validated polarization configuration.
rtype:: PolarizationConfig
raises ValueError:: If polarization_type is not one of LVP, LHP, RCP, LCP, LAP, or unpolarized.
raises EquinoxRuntimeError:: If any angular parameter is non-finite.

Notes

Static validation raises ValueError before traced construction for an unsupported polarization_type. Traced validation uses eqx.error_if and raises EquinoxRuntimeError when an angle is non-finite.

See also

PolarizationConfig: The PyTree class constructed by this factory.

Create a validated SimulationParams instance.

The factory validates and normalizes the inputs before it constructs the SimulationParams PyTree. It casts float-valued parameters to 0-D float64 JAX arrays. It keeps fidelity as a Python int in the PyTree auxiliary data.

The defaults represent a typical low-temperature ARPES experiment. The energy window defaults to [-3.0, 1.0] eV. Higher-level simulation drivers can override this window from the supplied eigenband energy range.

See:: TestMakeSimulationParams

Implementation Logic¶

Prepare the normalized values:
```
emin_arr = jnp.asarray(energy_min, dtype=jnp.float64)
```
This expression gives the later validation steps a stable shape and dtype.
Apply static validation:
```
fidelity < _MIN_FIDELITY
```
This predicate rejects invalid structure before JAX traces the numerical checks.
Apply traced validation:
```
~jnp.isfinite(emin_arr)
```
This predicate remains active during eager and compiled execution.
Return the named instance:
```
return params
```
The explicit name keeps the implementation and the Returns section synchronized.

param energy_min:: Lower energy bound in eV. Default is -3.0.
type energy_min:: Union[int, float, complex, Num[Array, '']], default: -3.0
param energy_max:: Upper energy bound in eV. Default is 1.0.
type energy_max:: Union[int, float, complex, Num[Array, '']], default: 1.0
param fidelity:: Energy axis resolution (static – a compile-time constant; changing it triggers retracing). Default is 25000.
type fidelity:: int, default: 25000
param sigma:: Gaussian broadening in eV. Default is 0.04.
type sigma:: Union[float, Float[Array, '']], default: 0.04
param gamma:: Lorentzian broadening in eV. Default is 0.1.
type gamma:: Union[float, Float[Array, '']], default: 0.1
param temperature:: Temperature in Kelvin. Default is 15.0.
type temperature:: Union[float, Float[Array, '']], default: 15.0
param photon_energy:: Photon energy in eV. Default is 11.0.
type photon_energy:: Union[float, Float[Array, '']], default: 11.0
returns:: params – Validated simulation parameters.
rtype:: SimulationParams
raises ValueError:: If fidelity is less than two.
raises EquinoxRuntimeError:: If an energy bound or physical scalar is non-finite, if energy_min >= energy_max, or if any broadening width, temperature, or photon energy is not positive.

Notes

Static validation raises ValueError before traced construction when fidelity is less than two. Traced validation uses eqx.error_if and raises EquinoxRuntimeError for non-finite values, an unordered energy window, or non-positive physical scalars.

See also

SimulationParams: The PyTree class constructed by this factory.

class diffpes.types.InformationState(node_id: str, active_semantics: tuple[str, ...], destroyed_information: tuple[str, ...], invalidated_claims: tuple[str, ...]) → None[source]¶

Bases: Module

Store effective semantic state for one artifact or result node.

This carrier separates semantics that remain available from information and claims invalidated along the node’s provenance path.

See:

TestInformationState

Variables:

node_id (str) – Artifact or result identifier (static – a compile-time constant; changing it triggers retracing).
active_semantics (tuple[str, ]) – Available scientific semantics (static – compile-time constants; changing them triggers retracing).
destroyed_information (tuple[str, ]) – Information lost before this node (static – compile-time constants; changing them triggers retracing).
invalidated_claims (tuple[str, ]) – Claims invalidated before this node (static – compile-time constants; changing them triggers retracing).

Notes

The state is declarative static metadata. It records information flow but does not alter or differentiate through the associated physical arrays.

See also

make_information_state: Create a validated semantic-information state for one graph node.

node_id: str¶

active_semantics: tuple[str, ...]¶

destroyed_information: tuple[str, ...]¶

invalidated_claims: tuple[str, ...]¶

__init__(node_id: str, active_semantics: tuple[str, ...], destroyed_information: tuple[str, ...], invalidated_claims: tuple[str, ...]) → None¶

class diffpes.types.ProvenanceGraph(records: tuple[TransformationRecord, ...], external_inputs: tuple[str, ...], initial_semantics: tuple[tuple[str, tuple[str, ...]], ...], topological_order: tuple[str, ...], information: tuple[InformationState, ...], validation_errors: tuple[str, ...], graph_checksum: str) → None[source]¶

Bases: Module

Store a validated lineage graph and its propagated semantics.

The graph retains every transformation edge, external root, and effective semantic state needed to inspect information flow without reevaluation.

See:

TestProvenanceGraph

Variables:

records (tuple[TransformationRecord, ]) – Transformation records in graph order.
external_inputs (tuple[str, ]) – External root identifiers (static – compile-time constants; changing them triggers retracing).
initial_semantics (tuple[tuple[str, tuple[str, ]], ]) – Initial semantics per external root (static – compile-time constants; changing them triggers retracing).
topological_order (tuple[str, ]) – Validated node order (static – compile-time constants; changing them triggers retracing).
information (tuple[InformationState, ]) – Propagated semantic state for graph nodes.
validation_errors (tuple[str, ]) – Structural or semantic validation errors (static – compile-time constants; changing them triggers retracing).
graph_checksum (str) – Deterministic consistency checksum (static – a compile-time constant; changing it triggers retracing).

Notes

The graph is an immutable audit carrier evaluated outside forward kernels. Its records and information states contain no differentiable array leaves.

See also

make_provenance_graph: Create a validated immutable provenance graph carrier.

records: tuple[TransformationRecord, ...]¶

external_inputs: tuple[str, ...]¶

initial_semantics: tuple[tuple[str, tuple[str, ...]], ...]¶

topological_order: tuple[str, ...]¶

information: tuple[InformationState, ...]¶

validation_errors: tuple[str, ...]¶

graph_checksum: str¶

__init__(records: tuple[TransformationRecord, ...], external_inputs: tuple[str, ...], initial_semantics: tuple[tuple[str, tuple[str, ...]], ...], topological_order: tuple[str, ...], information: tuple[InformationState, ...], validation_errors: tuple[str, ...], graph_checksum: str) → None¶

class diffpes.types.ProvenanceReport(valid: bool, errors: tuple[str, ...], roots: tuple[str, ...], terminal_outputs: tuple[str, ...], orphaned_inputs: tuple[str, ...], topological_order: tuple[str, ...], graph_checksum: str) → None[source]¶

Bases: Module

Store a structural and semantic provenance-validation report.

This carrier summarizes graph validity and exposes its roots, terminal outputs, orphaned inputs, and deterministic traversal identity.

See:

TestProvenanceReport

Variables:

valid (bool) – Whether validation succeeded (static – a compile-time constant; changing it triggers retracing).
errors (tuple[str, ]) – Validation failures (static – compile-time constants; changing them triggers retracing).
roots (tuple[str, ]) – Root node identifiers (static – compile-time constants; changing them triggers retracing).
terminal_outputs (tuple[str, ]) – Terminal output identifiers (static – compile-time constants; changing them triggers retracing).
orphaned_inputs (tuple[str, ]) – Unconsumed external input identifiers (static – compile-time constants; changing them triggers retracing).
topological_order (tuple[str, ]) – Validated node order (static – compile-time constants; changing them triggers retracing).
graph_checksum (str) – Deterministic consistency checksum (static – a compile-time constant; changing it triggers retracing).

Notes

The report contains only static graph metadata and therefore contributes no gradient path to a certified forward execution.

See also

make_provenance_report: Create a validated structural and semantic provenance report.

valid: bool¶

errors: tuple[str, ...]¶

roots: tuple[str, ...]¶

terminal_outputs: tuple[str, ...]¶

orphaned_inputs: tuple[str, ...]¶

topological_order: tuple[str, ...]¶

graph_checksum: str¶

__init__(valid: bool, errors: tuple[str, ...], roots: tuple[str, ...], terminal_outputs: tuple[str, ...], orphaned_inputs: tuple[str, ...], topological_order: tuple[str, ...], graph_checksum: str) → None¶

diffpes.types.make_information_state(node_id: str, active_semantics: tuple[str, ...] = (), destroyed_information: tuple[str, ...] = (), invalidated_claims: tuple[str, ...] = ()) → InformationState[source]¶

Create a validated semantic-information state for one graph node.

Freeze named semantic sets while rejecting blank or duplicate entries.

See:: TestMakeInformationState

Implementation Logic¶

Validate node identity:
```
node_id=_require_text(node_id, "node_id")
```
Reject a blank node identifier.
Freeze semantic sets:
```
active_semantics=_text_tuple(active_semantics, "active_semantics")
```
Convert each sequence to an immutable, duplicate-free tuple.
Construct the state:
```
state = InformationState(...)
```
Bind and return the semantic state carrier.

param node_id:: Artifact or result identifier (static – a compile-time constant; changing it triggers retracing).
type node_id:: str
param active_semantics:: Available semantics (static – compile-time constants; changing them triggers retracing). Default is empty.
type active_semantics:: tuple[str, ...], default: ()
param destroyed_information:: Lost information labels (static – compile-time constants; changing them triggers retracing). Default is empty.
type destroyed_information:: tuple[str, ...], default: ()
param invalidated_claims:: Invalidated claim identifiers (static – compile-time constants; changing them triggers retracing). Default is empty.
type invalidated_claims:: tuple[str, ...], default: ()
returns:: state – Validated immutable semantic-information state.
rtype:: InformationState
raises ValueError:: If an identifier is blank or a semantic sequence contains duplicates.

Notes

Validation is static; the resulting carrier contains no numerical leaves.

diffpes.types.make_provenance_graph(records: tuple[TransformationRecord, ...], external_inputs: tuple[str, ...], initial_semantics: tuple[tuple[str, tuple[str, ...]], ...], topological_order: tuple[str, ...], information: tuple[InformationState, ...], validation_errors: tuple[str, ...], graph_checksum: str) → ProvenanceGraph[source]¶

Create a validated immutable provenance graph carrier.

Validate carrier types, root coverage, and unique semantic-state IDs, then freeze the graph in deterministic topological order.

See:: TestMakeProvenanceGraph

Implementation Logic¶

Validate carrier sequences:
```
frozen_records = tuple(records)
```
Require transformation records and information states to use their canonical types.
Validate root semantics:
```
if set(semantic_nodes) != set(inputs):
```
Require one unique initial-semantic entry for every external input.
Validate state identities:
```
if len(state_ids) != len(set(state_ids)):
```
Prevent ambiguous propagated semantic states.
Construct the graph:
```
graph = ProvenanceGraph(...)
```
Freeze the validated graph and bind the result.

param records:: Transformation records in graph order.
type records:: tuple[TransformationRecord, ...]
param external_inputs:: External root identifiers (static – compile-time constants; changing them triggers retracing).
type external_inputs:: tuple[str, ...]
param initial_semantics:: Initial semantics for every root (static – compile-time constants; changing them triggers retracing).
type initial_semantics:: tuple[tuple[str, tuple[str, ...]], ...]
param topological_order:: Validated node order (static – compile-time constants; changing them triggers retracing).
type topological_order:: tuple[str, ...]
param information:: Propagated semantic states for graph nodes.
type information:: tuple[InformationState, ...]
param validation_errors:: Graph validation errors (static – compile-time constants; changing them triggers retracing).
type validation_errors:: tuple[str, ...]
param graph_checksum:: Deterministic consistency checksum (static – a compile-time constant; changing it triggers retracing).
type graph_checksum:: str
returns:: graph – Validated immutable provenance graph.
rtype:: ProvenanceGraph
raises TypeError:: If records or information contains the wrong carrier type.
raises ValueError:: If text values are blank or duplicated. The function also rejects incomplete initial semantics or duplicate information-state node IDs.

Notes

Graph validation uses only static identities and carrier structure; it does not inspect or reduce physical model arrays.

diffpes.types.make_provenance_report(valid: bool, errors: tuple[str, ...], roots: tuple[str, ...], terminal_outputs: tuple[str, ...], orphaned_inputs: tuple[str, ...], topological_order: tuple[str, ...], graph_checksum: str) → ProvenanceReport[source]¶

Create a validated structural and semantic provenance report.

Freeze the graph summary and require the validity flag to agree exactly with whether validation errors are present.

See:: TestMakeProvenanceReport

Implementation Logic¶

Normalize errors:
```
normalized_errors = _text_tuple(errors, "errors", unique=False)
```
Freeze validation failures without hiding repeated diagnostics.
Check validity consistency:
```
if valid == bool(normalized_errors):
```
Require success exactly when no graph error is present.
Construct the report:
```
report = ProvenanceReport(...)
```
Freeze graph endpoints and bind the result.

param valid:: Whether graph validation succeeded (static – a compile-time constant; changing it triggers retracing).
type valid:: bool
param errors:: Validation failures (static – compile-time constants; changing them triggers retracing).
type errors:: tuple[str, ...]
param roots:: Root identifiers (static – compile-time constants; changing them triggers retracing).
type roots:: tuple[str, ...]
param terminal_outputs:: Terminal output identifiers (static – compile-time constants; changing them triggers retracing).
type terminal_outputs:: tuple[str, ...]
param orphaned_inputs:: Unconsumed input identifiers (static – compile-time constants; changing them triggers retracing).
type orphaned_inputs:: tuple[str, ...]
param topological_order:: Validated node order (static – compile-time constants; changing them triggers retracing).
type topological_order:: tuple[str, ...]
param graph_checksum:: Deterministic consistency checksum (static – a compile-time constant; changing it triggers retracing).
type graph_checksum:: str
returns:: report – Validated immutable provenance report.
rtype:: ProvenanceReport
raises ValueError:: If a text value is blank, a unique sequence contains duplicates, or valid does not agree with whether errors is empty.

Notes

Report construction is static and introduces no differentiable leaves.

class diffpes.types.OrbitalBasis(n_values: tuple[int, ...], l_values: tuple[int, ...], m_values: tuple[int, ...], labels: tuple[str, ...]) → None[source]¶

Bases: Module

Store orbital quantum-number metadata in a JAX PyTree.

This type describes the orbital basis for dipole matrix-element calculations for the differentiable Chinook pipeline. The quantum numbers (n, l, m) parameterize the radial wavefunctions (via Slater-type orbitals) and angular parts (spherical harmonics) that enter the photoemission matrix element.

All fields contain static auxiliary data because quantum numbers control code paths. They determine recurrence depths in spherical Bessel functions and associated Legendre polynomials. They also index the Gaunt coefficient table. A quantum-number change alters the computational graph. JAX must therefore recompile after this change.

See:

TestOrbitalBasis

Variables:

n_values (tuple[int, ]) – Principal quantum numbers, one per orbital. Each value determines the radial node count and the power of r in the Slater-type radial function R_nl(r) ~ r^{n-1} exp(-zeta*r). These values are static.
l_values (tuple[int, ]) – Angular momentum quantum numbers, one per orbital (0=s, 1=p, 2=d, 3=f). Determines the spherical harmonic Y_l^m used in the matrix element integral (static – compile-time constants; changing them triggers retracing).
m_values (tuple[int, ]) – Magnetic quantum numbers, one per orbital. Ranges from -l to +l for each orbital. Selects the specific spherical harmonic component (static – compile-time constants; changing them triggers retracing).
labels (tuple[str, ]) – Human-readable orbital labels (e.g. ("2s", "2px", ...)). Used for plotting and debugging (static – compile-time constants; changing them triggers retracing).

Notes

Implemented as an immutable equinox.Module PyTree. All fields are auxiliary data (no JAX array children) because changing any quantum number changes the computational graph and requires JIT recompilation. The children tuple is always empty.

See also

SlaterParams: Wraps differentiable Slater exponents alongside this static basis metadata.
make_orbital_basis: Factory function with length validation and default label generation.

n_values: tuple[int, ...]¶

l_values: tuple[int, ...]¶

m_values: tuple[int, ...]¶

labels: tuple[str, ...]¶

__init__(n_values: tuple[int, ...], l_values: tuple[int, ...], m_values: tuple[int, ...], labels: tuple[str, ...]) → None¶

class diffpes.types.SlaterParams(zeta: Float[Array, 'O'], coefficients: Float[Array, 'O C'], orbital_basis: OrbitalBasis) → None[source]¶

Bases: Module

Store Slater radial-wavefunction parameters in a JAX PyTree.

This type contains differentiable Slater exponents and linear-combination coefficients alongside the static orbital basis metadata. Inverse fitting workflows optimize the Slater exponents (zeta). jax.grad gives the gradient of the simulated ARPES intensity with respect to the radial wavefunction shape.

In a multi-zeta basis (C > 1), the coefficients matrix combines C Slater-type functions for each orbital. These functions have different exponents. In a single-zeta basis (C=1), each orbital has one exponent and a coefficient of 1.0.

See:

TestSlaterParams

Variables:

zeta (Float[Array, " O"]) – Slater exponents (inverse Bohr), one per orbital. Controls the radial decay rate of each Slater-type orbital R(r) ~ r^{n-1} exp(-zeta * r). JAX-traced (differentiable) – this is the primary optimization target for inverse fitting.
coefficients (Float[Array, "O C"]) – Linear combination coefficients for multi-zeta expansion. Shape is (O, C) where O is the number of orbitals and C is the contraction length. For single-zeta basis sets, C=1 and all coefficients are 1.0. JAX-traced (differentiable).
orbital_basis (OrbitalBasis) – Quantum numbers (n, l, m) and labels for each orbital. Static – a compile-time constant; changing it triggers retracing.

Notes

Implemented as an immutable equinox.Module PyTree. zeta and coefficients are JAX array children (on the gradient tape), while orbital_basis is auxiliary data (static at trace time). This separation means that jax.grad(loss)(slater_params) differentiates through the exponents and coefficients while keeping the quantum number structure fixed.

See also

OrbitalBasis: The static quantum number metadata nested inside.
make_slater_params: Factory function with validation and default coefficient generation.

zeta: Float[Array, 'O']¶

coefficients: Float[Array, 'O C']¶

orbital_basis: OrbitalBasis¶

__init__(zeta: Float[Array, 'O'], coefficients: Float[Array, 'O C'], orbital_basis: OrbitalBasis) → None¶

diffpes.types.make_orbital_basis(n_values: tuple[int, ...], l_values: tuple[int, ...], m_values: tuple[int, ...], labels: tuple[str, ...] | None = None) → OrbitalBasis[source]¶

Create a validated OrbitalBasis instance.

The factory validates quantum number tuples and constructs an OrbitalBasis PyTree. The three quantum number tuples must all have the same length (one entry per orbital). If labels is absent, the factory generates generic labels such as "orb_0" and "orb_1".

Use this factory instead of the raw OrbitalBasis constructor to get automatic length validation and default label generation.

See:: TestMakeOrbitalBasis

Implementation Logic¶

Prepare the normalized values:
```
n_orbitals = len(n_values)
```
This expression gives the later validation steps a stable shape and dtype.
Apply static validation:
```
len(l_values) != n_orbitals or len(m_values) != n_orbitals
```
This predicate rejects invalid structure before JAX traces the numerical checks.
Return the named instance:
```
return basis
```
The explicit name keeps the implementation and the Returns section synchronized.

param n_values:: Principal quantum numbers (static – compile-time constants; changing them triggers retracing), one per orbital.
type n_values:: tuple[int, ...]
param l_values:: Angular momentum quantum numbers (static – compile-time constants; changing them triggers retracing), one per orbital.
type l_values:: tuple[int, ...]
param m_values:: Magnetic quantum numbers (static – compile-time constants; changing them triggers retracing), one per orbital.
type m_values:: tuple[int, ...]
param labels:: Human-readable orbital labels (static – compile-time constants; changing them triggers retracing). Defaults to ("orb_0", "orb_1", ...).
type labels:: Optional[tuple[str, ...]], default: None
returns:: basis – Validated orbital basis with consistent lengths.
rtype:: OrbitalBasis
raises ValueError:: If the quantum-number tuples have different lengths or labels has a different length. The function also rejects invalid values of n, l, or m.

Notes

Every OrbitalBasis field uses eqx.field(static=True), so the factory performs static validation. Invalid tuple lengths or quantum numbers raise ValueError before tracing. No eqx.error_if checks apply.

See also

OrbitalBasis: The PyTree class constructed by this factory.

diffpes.types.make_slater_params(zeta: Float[Array, 'Oz'], orbital_basis: OrbitalBasis, coefficients: Float[Array, 'Oc C'] | None = None) → SlaterParams[source]¶

Create a validated SlaterParams instance.

The factory validates Slater radial wavefunction parameters and constructs a SlaterParams PyTree. The zeta array length must match the orbital count in orbital_basis. If coefficients is absent, the factory uses a single-zeta basis. It creates a column of ones with shape (O, 1).

Use this factory for automatic size validation and default coefficient generation. The factory also guarantees float64 dtypes.

See:: TestMakeSlaterParams

Implementation Logic¶

Prepare the normalized values:
```
zeta_arr = jnp.asarray(zeta, dtype=jnp.float64)
```
This expression gives the later validation steps a stable shape and dtype.
Apply static validation:
```
len(orbital_basis.n_values) != n_orbitals
```
This predicate rejects invalid structure before JAX traces the numerical checks.
Apply traced validation:
```
~jnp.all(jnp.isfinite(zeta_arr))
```
This predicate remains active during eager and compiled execution.
Return the named instance:
```
return params
```
The explicit name keeps the implementation and the Returns section synchronized.

param zeta:: Slater exponents (inverse Bohr), one per orbital.
type zeta:: Float[Array, 'Oz']
param orbital_basis:: Orbital quantum number metadata (static – a compile-time constant; changing it triggers retracing). Must have the same number of orbitals as the length of zeta.
type orbital_basis:: OrbitalBasis
param coefficients:: Multi-zeta linear combination coefficients. Shape (O, C) where C is the contraction length. Defaults to ones with C=1 (single-zeta basis).
type coefficients:: Optional[Float[Array, 'Oc C']], default: None
returns:: params – Validated Slater parameters with float64 arrays.
rtype:: SlaterParams
raises ValueError:: If zeta and orbital_basis contain different orbital counts. The function also rejects a different orbital count in coefficients.
raises EquinoxRuntimeError:: If zeta is non-finite or non-positive, or if coefficients contains a non-finite value.

Notes

Static validation raises ValueError before traced construction when the orbital count differs from either leading array dimension. Traced validation uses eqx.error_if and raises EquinoxRuntimeError when an exponent is non-finite or non-positive, or a coefficient is non-finite.

See also

SlaterParams: The PyTree class constructed by this factory.
make_orbital_basis: Factory for the OrbitalBasis argument.

class diffpes.types.SelfEnergyConfig(coefficients: Float[Array, 'P'], energy_nodes: Float[Array, 'P'] | None, mode: str) → None[source]¶

Bases: Module

Store energy-dependent self-energy data in a JAX PyTree.

This type models the imaginary part of the electronic self-energy Im[Sigma(E)] as a function of binding energy. In the forward ARPES simulator this replaces the scalar Lorentzian half-width gamma with an energy-dependent broadening that captures quasiparticle lifetime effects more realistically.

The mode string selects one of three modes:

constant: a single scalar broadening gamma applied uniformly at all energies. Equivalent to the standard SimulationParams.gamma.
polynomial: Im[Sigma(E)] = a0 + a1*E + a2*E^2 + …, with coefficients stored in coefficients.
tabulated: Discrete energy nodes specify Im[Sigma]. The function interpolates between them and requires energy_nodes.

The coefficients array is the primary differentiable quantity. jax.grad gives the sensitivity of the simulated ARPES spectrum to the self-energy shape. This sensitivity supports inverse fitting of lifetime broadening from experimental data.

See:

TestSelfEnergyConfig

Variables:

coefficients (Float[Array, " P"]) –
Parameters for the self-energy model. JAX-traced (differentiable).
- mode=”constant”: P=1, [gamma] in eV.
- mode=”polynomial”: P=degree+1, [a0, a1, ...] where Im[Sigma(E)] = sum_i a_i * E^i.
- mode=”tabulated”: P=N, [gamma_1, ..., gamma_N] in eV at the corresponding energy_nodes.
energy_nodes (Optional[Float[Array, " P"]]) – Energy grid (eV) for tabulated mode. Must have the same length P as coefficients in tabulated mode. None for constant and polynomial modes. JAX-traced when present.
mode (str) – One of "constant", "polynomial", "tabulated". Static – a compile-time constant; changing it triggers retracing.

Notes

Implemented as an immutable equinox.Module PyTree. JAX stores coefficients and energy_nodes as children on the gradient tape. It stores mode as compile-time auxiliary data. Changing mode triggers JIT recompilation because it alters the computation graph.

See also

make_self_energy_config: Factory function with mode validation and default coefficient generation.

coefficients: Float[Array, 'P']¶

energy_nodes: Float[Array, 'P'] | None¶

mode: str¶

__init__(coefficients: Float[Array, 'P'], energy_nodes: Float[Array, 'P'] | None, mode: str) → None¶

diffpes.types.make_self_energy_config(gamma: float | Float[Array, ''] = 0.1, mode: str = 'constant', coefficients: Float[Array, 'Pc'] | None = None, energy_nodes: Float[Array, 'Pn'] | None = None) → SelfEnergyConfig[source]¶

Create a validated SelfEnergyConfig instance.

The factory validates inputs and constructs a SelfEnergyConfig PyTree. It accepts only the "constant", "polynomial", and "tabulated" modes. It requires energy_nodes for the tabulated mode.

The convenience parameter gamma provides a shortcut for the common constant-broadening case. When coefficients is None, the factory creates the single-element array [gamma].

See:: TestMakeSelfEnergyConfig

Implementation Logic¶

Prepare the normalized values:
```
nodes_arr = None
```
This expression gives the later validation steps a stable shape and dtype.
Apply static validation:
```
mode not in ('constant', 'polynomial', 'tabulated')
```
This predicate rejects invalid structure before JAX traces the numerical checks.
Apply traced validation:
```
~jnp.all(jnp.isfinite(coeff_arr))
```
This predicate remains active during eager and compiled execution.
Return the named instance:
```
return config
```
The explicit name keeps the implementation and the Returns section synchronized.

param gamma:: Constant broadening in eV. Used as the sole coefficient when coefficients is None and mode is "constant". Default is 0.1.
type gamma:: Union[float, Float[Array, '']], default: 0.1
param mode:: Self-energy model. One of "constant", "polynomial", "tabulated" (static – a compile-time constant; changing it triggers retracing). Default is "constant".
type mode:: str, default: "constant"
param coefficients:: Explicit model coefficients. If None, defaults to [gamma]. For polynomial mode, these are the polynomial coefficients [a0, a1, ...]. For tabulated mode, these are the broadening values at each energy node.
type coefficients:: Optional[Float[Array, 'Pc']], default: None
param energy_nodes:: Energy grid (eV) for tabulated mode. Must have the same length as coefficients in tabulated mode. Ignored for other modes. Default is None.
type energy_nodes:: Optional[Float[Array, 'Pn']], default: None
returns:: config – Validated self-energy configuration with float64 arrays.
rtype:: SelfEnergyConfig
raises ValueError:: If mode has an unsupported value. The function also rejects an invalid presence or length of energy_nodes.
raises EquinoxRuntimeError:: If coefficients are non-finite or tabulated energy nodes are non-finite or not strictly increasing.

Notes

Static validation raises ValueError before traced construction for an unsupported mode, an invalid node-mode combination, or unequal tabulated lengths. Traced validation uses eqx.error_if and raises EquinoxRuntimeError for non-finite coefficients or for tabulated nodes that are non-finite or not strictly increasing.

See also

SelfEnergyConfig: The PyTree class constructed by this factory.

class diffpes.types.DiagonalizedBands(eigenvalues: Float[Array, 'K B'], eigenvectors: Complex[Array, 'K B O'], kpoints: Float[Array, 'K 3'], fermi_energy: Float[Array, '']) → None[source]¶

Bases: Module

Store diagonalized electronic-structure data in a JAX PyTree.

This type provides the common interface between VASP-derived and TB-derived inputs for the forward simulator simulate_tb_radial. The native diffpes.tightb.diagonalize_tb producer constructs this PyTree from a TBModel; the VASP adapter vasp_to_diagonalized constructs it from VASP eigenvectors.

This single PyTree type lets the differentiable forward simulator accept either tight-binding or first-principles inputs without code branches. The eigenvectors carry the orbital-decomposition information needed to compute dipole matrix elements in the Chinook pipeline.

All four fields are dense JAX arrays stored as children (no auxiliary data), making the entire object fully differentiable. This enables end-to-end gradient computation from TB hopping parameters through diagonalization to simulated ARPES intensity.

See:

TestDiagonalizedBands

Variables:

eigenvalues (Float[Array, "K B"]) – Band energies E_n(k) in eV for K k-points and B bands. JAX-traced (differentiable).
eigenvectors (Complex[Array, "K B O"]) – Complex orbital coefficients c_{k,b,orb} from the Hamiltonian diagonalization, where O is the number of orbitals in the basis. These encode the orbital character of each Bloch state and enter the dipole matrix element calculation. JAX-traced (differentiable, complex128).
kpoints (Float[Array, "K 3"]) – k-point coordinates in reciprocal (fractional) space. JAX-traced (differentiable).
fermi_energy (Float[Array, " "]) – Fermi level in eV. A 0-D scalar array. JAX-traced (differentiable).

Notes

Implemented as an immutable equinox.Module PyTree. All four fields are differentiable leaves and no static metadata is present.

See also

TBModel: The tight-binding model whose diagonalization produces a DiagonalizedBands.
make_diagonalized_bands: Factory function with validation and dtype casting.

eigenvalues: Float[Array, 'K B']¶

eigenvectors: Complex[Array, 'K B O']¶

kpoints: Float[Array, 'K 3']¶

fermi_energy: Float[Array, '']¶

__init__(eigenvalues: Float[Array, 'K B'], eigenvectors: Complex[Array, 'K B O'], kpoints: Float[Array, 'K 3'], fermi_energy: Float[Array, '']) → None¶

class diffpes.types.TBModel(hopping_params: Float[Array, 'H'], lattice_vectors: Float[Array, '3 3'], hopping_indices: tuple, n_orbitals: int, orbital_basis: OrbitalBasis) → None[source]¶

Bases: Module

Store tight-binding parameters in a JAX PyTree.

hopping_params and lattice_vectors are differentiable leaves. The connectivity, orbital count, and basis are compile-time metadata; changing any static field changes the PyTree definition and retraces JIT compiled consumers.

See:

TestTBModel

Variables:

hopping_params (Float[Array, " H"]) – Hopping amplitudes in eV.
lattice_vectors (Float[Array, "3 3"]) – Real-space lattice vectors in Angstrom.
hopping_indices (tuple) – Orbital connectivity and lattice translations (static – a compile-time constant; changing it triggers retracing).
n_orbitals (int) – Number of orbitals in the unit cell (static – a compile-time constant; changing it triggers retracing).
orbital_basis (OrbitalBasis) – Orbital quantum-number metadata (static – a compile-time constant; changing it triggers retracing).

hopping_params: Float[Array, 'H']¶

lattice_vectors: Float[Array, '3 3']¶

hopping_indices: tuple¶

n_orbitals: int¶

orbital_basis: OrbitalBasis¶

__init__(hopping_params: Float[Array, 'H'], lattice_vectors: Float[Array, '3 3'], hopping_indices: tuple, n_orbitals: int, orbital_basis: OrbitalBasis) → None¶

diffpes.types.make_1d_chain_model(t: float | Float[Array, ''] = -1.0) → TBModel[source]¶

Create a 1D chain tight-binding model.

The model has one orbital in each unit cell and nearest-neighbor hopping t.

The 1D chain has one s-orbital in each unit cell. Hopping connects only the two nearest neighbors at lattice vectors +a1 and -a1. The function uses an identity lattice matrix. Therefore, fractional and Cartesian coordinates coincide for a lattice constant of 1 in arbitrary units.

The resulting band dispersion is the standard cosine band:

\[E(k) = 2t \cos(2 \pi k)\]

with bandwidth |4t|. This model provides a minimal test for the Hamiltonian builder, diagonalizer, and gradient machinery.

The hopping list contains (0, 0, (+1,0,0)) and (0, 0, (-1,0,0)). These entries connect the single orbital to itself in adjacent unit cells. Hermitianization in build_hamiltonian_k makes these entries redundant because each is its own conjugate. The diagonal entry therefore receives 2t cos(2 pi k).

See:: TestMake1dChainModel

Implementation Logic¶

Create the orbital basis:
```
basis = make_orbital_basis(...)
```
The basis contains one s orbital per unit cell.
Define nearest-neighbor hops:
```
hopping_indices = ((0, 0, (1, 0, 0)), (0, 0, (-1, 0, 0)))
```
The two entries connect the orbital to both adjacent unit cells.
Construct the model:
```
model = make_tb_model(...)
```
The validated factory preserves gradients through the hopping value.

param t:: Hopping amplitude. Default -1.0.
type t:: Union[float, Float[Array, '']], default: -1.0
returns:: model – 1D chain model.
rtype:: TBModel

diffpes.types.make_diagonalized_bands(eigenvalues: Float[Array, 'Ke Be'], eigenvectors: Complex[Array, 'Kv Bv O'], kpoints: Float[Array, 'Kk 3'], fermi_energy: int | float | complex | Num[Array, ''] = 0.0) → DiagonalizedBands[source]¶

Create a validated DiagonalizedBands instance.

The factory validates and normalizes diagonalized electronic structure data before constructing a DiagonalizedBands PyTree. It casts real-valued arrays to float64 and the complex eigenvector array to complex128 to maintain full double-precision accuracy in the orbital decomposition.

@jaxtyped(typechecker=beartype) checks the shape constraints at call time. K must agree across all arrays, and B and O must be consistent.

Use this factory when constructing DiagonalizedBands from either TB diagonalization output or VASP eigenvector data.

See:: TestMakeDiagonalizedBands

Implementation Logic¶

Prepare the normalized values:
```
eig_arr = jnp.asarray(eigenvalues, dtype=jnp.float64)
```
This expression gives the later validation steps a stable shape and dtype.
Apply static validation:
```
eig_arr.shape != vec_arr.shape[:2]
```
This predicate rejects invalid structure before JAX traces the numerical checks.
Apply traced validation:
```
~jnp.all(jnp.isfinite(eig_arr))
```
This predicate remains active during eager and compiled execution.
Return the named instance:
```
return bands
```
The explicit name keeps the implementation and the Returns section synchronized.

param eigenvalues:: Band energies E_n(k) in eV for K k-points and B bands.
type eigenvalues:: Float[Array, 'Ke Be']
param eigenvectors:: Complex orbital coefficients c_{k,b,orb} from Hamiltonian diagonalization, where O is the number of orbitals.
type eigenvectors:: Complex[Array, 'Kv Bv O']
param kpoints:: k-point coordinates in reciprocal (fractional) space.
type kpoints:: Float[Array, 'Kk 3']
param fermi_energy:: Fermi level in eV. Default is 0.0.
type fermi_energy:: Union[int, float, complex, Num[Array, '']], default: 0.0
returns:: bands – Validated instance with float64 eigenvalues/kpoints and complex128 eigenvectors.
rtype:: DiagonalizedBands
raises ValueError:: If the k-point or band dimensions disagree across the arrays.
raises EquinoxRuntimeError:: If eigenvalues, eigenvectors, k-points, or the Fermi energy contain a non-finite value.

Notes

Static validation raises ValueError before traced construction when the k-point or band dimensions disagree. Traced validation uses eqx.error_if and raises EquinoxRuntimeError when any numerical field contains a non-finite value.

See also

DiagonalizedBands: The PyTree class constructed by this factory.

diffpes.types.make_graphene_model(t: float | Float[Array, ''] = -2.7) → TBModel[source]¶

Create a graphene pz tight-binding model.

The model has two orbitals (A/B sublattices) on a honeycomb lattice. It uses nearest-neighbor hopping t.

Graphene’s honeycomb lattice has two atoms (sublattices A and B) per primitive cell. The lattice vectors used here are:

a1 = (a, 0, 0)
a2 = (a/2, a*sqrt(3)/2, 0)
a3 = (0, 0, 10) (vacuum slab for 2-D periodicity)

with a = 2.46 Angstrom (the experimental graphene lattice constant). The function labels the two orbitals A_pz and B_pz. Their quantum numbers are (n=2, l=1, m=0). These orbitals represent carbon p_z orbitals on each sublattice.

Each A-site atom has three nearest-neighbor B-site atoms. In fractional coordinates the three A -> B hoppings connect to cells (0,0,0), (-1,0,0), and (0,-1,0). The reverse B -> A hoppings use (0,0,0), (+1,0,0), and (0,+1,0). The function lists these reverse hoppings explicitly. The raw Hamiltonian is therefore nearly Hermitian before build_hamiltonian_k applies Hermitianization.

The resulting 2x2 Hamiltonian produces the classic Dirac-cone band structure with linear dispersion near the K and K’ points and a bandwidth of |6t|.

See:: TestMakeGrapheneModel

Implementation Logic¶

Create the orbital basis:

basis: OrbitalBasis = make_orbital_basis(
    n_values=(2, 2),
    l_values=(1, 1),
    m_values=(0, 0),
    labels=("A_pz", "B_pz"),
)

These static quantum numbers identify the two carbon p-z orbitals.

Construct the primitive lattice:

a: float = 2.46
a1 = jnp.array([a, 0.0, 0.0], dtype=jnp.float64)
a2 = jnp.array(
    [a / 2.0, a * jnp.sqrt(3.0) / 2.0, 0.0],
    dtype=jnp.float64,
)
a3 = jnp.array([0.0, 0.0, 10.0], dtype=jnp.float64)
lattice = jnp.stack([a1, a2, a3])

The first two vectors define the honeycomb plane. The third vector adds vacuum normal to that plane.

Enumerate the nearest-neighbor hoppings:

hopping_indices = (
    (0, 1, (0, 0, 0)),
    (0, 1, (-1, 0, 0)),
    (0, 1, (0, -1, 0)),
    (1, 0, (0, 0, 0)),
    (1, 0, (1, 0, 0)),
    (1, 0, (0, 1, 0)),
)

The forward and reverse entries preserve the Hermitian lattice model.

Construct the validated model:

t_val = jnp.asarray(t, dtype=jnp.float64)
hopping_params = jnp.array(
    [t_val, t_val, t_val, t_val, t_val, t_val],
    dtype=jnp.float64,
)
model = make_tb_model(
    hopping_params=hopping_params,
    lattice_vectors=lattice,
    hopping_indices=hopping_indices,
    n_orbitals=2,
    orbital_basis=basis,
)

The shared JAX scalar keeps one differentiable hopping value for all six directed bonds. The factory applies the model validation contract.

Return the named instance:
```
return model
```
The explicit name keeps the implementation and the Returns section synchronized.

param t:: Nearest-neighbor hopping. Default -2.7 eV.
type t:: Union[float, Float[Array, '']], default: -2.7
returns:: model – Graphene model.
rtype:: TBModel

Notes

The default hopping value of -2.7 eV reproduces the standard nearest-neighbor graphene band structure. Castro Neto et al. use this value in Rev. Mod. Phys. 81, 109. The negative sign follows the convention that bonding states are lower in energy.

The conversion with jnp.asarray keeps t on the JAX gradient tape. Gradients of downstream spectra therefore propagate to the shared hopping.

diffpes.types.make_tb_model(hopping_params: Float[Array, 'H'], lattice_vectors: Float[Array, '3 3'], hopping_indices: tuple, n_orbitals: int, orbital_basis: OrbitalBasis) → TBModel[source]¶

Create a validated TBModel instance.

The factory validates and normalizes tight-binding model parameters before constructing a TBModel PyTree. The factory casts the two differentiable arrays to float64. It passes the three static fields unchanged as auxiliary data.

@jaxtyped(typechecker=beartype) checks the shape constraints on the differentiable arrays at call time.

Use this factory to build a TBModel from raw hopping data. Then pass the model to the Hamiltonian construction and diagonalization routines.

See:: TestMakeTBModel

Implementation Logic¶

Prepare the normalized values:
```
hop_arr = jnp.asarray(hopping_params, dtype=jnp.float64)
```
This expression gives the later validation steps a stable shape and dtype.
Apply static validation:
```
hop_arr.shape[0] != len(hopping_indices)
```
This predicate rejects invalid structure before JAX traces the numerical checks.
Apply traced validation:
```
~jnp.all(jnp.isfinite(hop_arr))
```
This predicate remains active during eager and compiled execution.
Return the named instance:
```
return model
```
The explicit name keeps the implementation and the Returns section synchronized.

param hopping_params:: Hopping amplitudes t_{ij,R} in eV, one per connection.
type hopping_params:: Float[Array, 'H']
param lattice_vectors:: Real-space lattice vectors as rows, in Angstroms.
type lattice_vectors:: Float[Array, '3 3']
param hopping_indices:: Connectivity: (orb_i, orb_j, (R_x, R_y, R_z)) per hopping, where R_x, R_y, R_z are integer lattice translation indices (static – a compile-time constant; changing it triggers retracing).
type hopping_indices:: tuple
param n_orbitals:: Number of orbitals in the unit cell (static – a compile-time constant; changing it triggers retracing).
type n_orbitals:: int
param orbital_basis:: Quantum number metadata for each orbital (static – a compile-time constant; changing it triggers retracing).
type orbital_basis:: OrbitalBasis
returns:: model – Validated tight-binding model with float64 arrays and static structural metadata.
rtype:: TBModel
raises ValueError:: If the hopping parameter count does not match the connectivity, an orbital index lies outside [0, n_orbitals), or the orbital basis size differs from n_orbitals.
raises EquinoxRuntimeError:: If hoppings or lattice vectors are non-finite, or if the lattice is degenerate.

Notes

Static validation raises ValueError before traced construction when hopping metadata, orbital counts, or orbital indices are inconsistent. Traced validation uses eqx.error_if and raises EquinoxRuntimeError for non-finite arrays or a degenerate lattice.

See also

TBModel: The PyTree class constructed by this factory.
make_orbital_basis: Factory for the OrbitalBasis argument.

class diffpes.types.SOCVolumetricData(lattice: Float[Array, '3 3'], coords: Float[Array, 'N 3'], charge: Float[Array, 'Nx Ny Nz'], magnetization: Float[Array, 'Nx Ny Nz'], magnetization_vector: Float[Array, 'Nx Ny Nz 3'], atom_counts: Int[Array, 'S'], grid_shape: tuple[int, int, int], symbols: tuple[str, ...]) → None[source]¶

Bases: Module

Store SOC CHGCAR volumetric-grid data in a JAX PyTree.

This type is a VolumetricData variant for spin-orbit coupling calculations. VASP writes four CHGCAR grid blocks: total charge, mx, my, and mz magnetization components. The magnetization contains the mz component for compatibility with ISPIN=2 consumers. magnetization_vector contains the complete 3-component vector (mx, my, mz) at each grid point.

This class is an immutable equinox.Module PyTree. JAX stores numerical fields as children that support tracing. It stores grid_shape and symbols as auxiliary data.

See:

TestSOCVolumetricData

Variables:

lattice (Float[Array, "3 3"]) – Real-space lattice vectors as rows, in Angstroms. JAX-traced (differentiable).
coords (Float[Array, "N 3"]) – Fractional atomic coordinates for all N atoms. JAX-traced (differentiable).
charge (Float[Array, "Nx Ny Nz"]) – Total charge density on the 3-D grid (electrons per unit cell volume). JAX-traced (differentiable).
magnetization (Float[Array, "Nx Ny Nz"]) – Scalar magnetization density, specifically the mz component. Provided for backward compatibility with code that expects ISPIN=2-style scalar magnetization. JAX-traced (differentiable).
magnetization_vector (Float[Array, "Nx Ny Nz 3"]) – Full vector magnetization (mx, my, mz) at each grid point. The last axis indexes the three Cartesian components. JAX-traced (differentiable).
atom_counts (Int[Array, " S"]) – Number of atoms per species. JAX-traced (int32).
grid_shape (tuple[int, int, int]) – Grid dimensions (Nx, Ny, Nz) (static – a compile-time constant; changing it triggers retracing).
symbols (tuple[str, ]) – Element symbols per species (static – compile-time constants; changing them triggers retracing).

Notes

Equinox derives the PyTree structure from the annotated fields. JAX stores six numerical fields as children. It stores grid_shape and symbols as auxiliary data. Unlike VolumetricData, both magnetization and magnetization_vector are mandatory (non-optional) because SOC calculations always produce all four grid blocks.

See also

VolumetricData: Non-SOC variant with optional scalar magnetization.
make_soc_volumetric_data: Factory function with validation and float64 casting.

lattice: Float[Array, '3 3']¶

coords: Float[Array, 'N 3']¶

charge: Float[Array, 'Nx Ny Nz']¶

magnetization: Float[Array, 'Nx Ny Nz']¶

magnetization_vector: Float[Array, 'Nx Ny Nz 3']¶

atom_counts: Int[Array, 'S']¶

grid_shape: tuple[int, int, int]¶

symbols: tuple[str, ...]¶

__init__(lattice: Float[Array, '3 3'], coords: Float[Array, 'N 3'], charge: Float[Array, 'Nx Ny Nz'], magnetization: Float[Array, 'Nx Ny Nz'], magnetization_vector: Float[Array, 'Nx Ny Nz 3'], atom_counts: Int[Array, 'S'], grid_shape: tuple[int, int, int], symbols: tuple[str, ...]) → None¶

class diffpes.types.VolumetricData(lattice: Float[Array, '3 3'], coords: Float[Array, 'N 3'], charge: Float[Array, 'Nx Ny Nz'], magnetization: Float[Array, 'Nx Ny Nz'] | None, atom_counts: Int[Array, 'S'], grid_shape: tuple[int, int, int], symbols: tuple[str, ...]) → None[source]¶

Bases: Module

Store CHGCAR volumetric-grid data in a JAX PyTree.

This type stores the charge density on a 3-D real-space grid together with the crystal lattice needed to interpret the grid coordinates. It includes an optional magnetization density grid for a spin-polarized CHGCAR calculation (ISPIN=2). The magnetization equals the difference between spin-up and spin-down charge densities.

This class is an immutable equinox.Module PyTree. JAX stores the numerical fields as children that support tracing. It stores grid_shape and symbols as auxiliary data because JAX cannot trace their Python tuples.

See:

TestVolumetricData

Variables:

lattice (Float[Array, "3 3"]) – Real-space lattice vectors as rows, in Angstroms. Defines the unit cell geometry for interpreting grid coordinates. JAX-traced (differentiable).
coords (Float[Array, "N 3"]) – Fractional atomic coordinates for all N atoms in the cell. JAX-traced (differentiable).
charge (Float[Array, "Nx Ny Nz"]) – Charge density on the 3-D real-space grid, in units of electrons per unit cell volume (VASP convention). Nx, Ny, Nz are the grid dimensions along the three lattice directions. JAX-traced (differentiable).
magnetization (Optional[Float[Array, "Nx Ny Nz"]]) – Scalar magnetization density (spin-up minus spin-down), or None for non-spin-polarized calculations (ISPIN=1). Same units and grid as charge. JAX-traced when present.
atom_counts (Int[Array, " S"]) – Number of atoms per species, with S = number of species. JAX-traced (differentiable, int32).
grid_shape (tuple[int, int, int]) – Grid dimensions (Nx, Ny, Nz) (static – a compile-time constant; changing it triggers retracing).
symbols (tuple[str, ]) – Element symbols for each species (e.g. ("Bi", "Se")). static – compile-time constants; changing them triggers retracing.

Notes

Equinox derives the PyTree structure from the annotated fields. JAX stores the Python tuples grid_shape and symbols as PyTree auxiliary data. JAX treats auxiliary data as compile-time constants: changing them triggers recompilation of any jit-compiled function.

See also

SOCVolumetricData: Variant for spin-orbit coupling with vector magnetization.
make_volumetric_data: Factory function with validation and float64 casting.

lattice: Float[Array, '3 3']¶

coords: Float[Array, 'N 3']¶

charge: Float[Array, 'Nx Ny Nz']¶

magnetization: Float[Array, 'Nx Ny Nz'] | None¶

atom_counts: Int[Array, 'S']¶

grid_shape: tuple[int, int, int]¶

symbols: tuple[str, ...]¶

__init__(lattice: Float[Array, '3 3'], coords: Float[Array, 'N 3'], charge: Float[Array, 'Nx Ny Nz'], magnetization: Float[Array, 'Nx Ny Nz'] | None, atom_counts: Int[Array, 'S'], grid_shape: tuple[int, int, int], symbols: tuple[str, ...]) → None¶

diffpes.types.make_soc_volumetric_data(lattice: Float[Array, '3 3'], coords: Float[Array, 'N 3'], charge: Float[Array, 'Cx Cy Cz'], magnetization: Float[Array, 'Mx My Mz'], magnetization_vector: Float[Array, 'Vx Vy Vz 3'], grid_shape: tuple[int, int, int] = (1, 1, 1), symbols: tuple[str, ...] = (), atom_counts: Int[Array, 'S'] | None = None) → SOCVolumetricData[source]¶

Create a validated SOCVolumetricData instance.

The factory validates and normalizes SOC CHGCAR volumetric data before constructing a SOCVolumetricData PyTree. It casts numerical arrays to float64 and atom counts to int32. Unlike make_volumetric_data(), both magnetization and magnetization_vector are mandatory because SOC calculations always produce all four grid blocks.

@jaxtyped(typechecker=beartype) checks shape constraints at call time. The grid dimensions must agree across charge, magnetization, and magnetization_vector.

See:: TestMakeSOCVolumetricData

Implementation Logic¶

Prepare the normalized values:
```
soc_lattice_arr = jnp.asarray(lattice, dtype=jnp.float64)
```
This expression gives the later validation steps a stable shape and dtype.
Apply static validation:
```
soc_charge_arr.shape != grid_shape
```
This predicate rejects invalid structure before JAX traces the numerical checks.
Apply traced validation:
```
~jnp.all(jnp.isfinite(soc_lattice_arr))
```
This predicate remains active during eager and compiled execution.
Return the named instance:
```
return vol
```
The explicit name keeps the implementation and the Returns section synchronized.

param lattice:: Real-space lattice vectors as rows, in Angstroms.
type lattice:: Float[Array, '3 3']
param coords:: Fractional atomic coordinates.
type coords:: Float[Array, 'N 3']
param charge:: Total charge density on 3-D grid (electrons per unit cell volume).
type charge:: Float[Array, 'Cx Cy Cz']
param magnetization:: Scalar magnetization density (mz component), for backward compatibility with ISPIN=2 consumers.
type magnetization:: Float[Array, 'Mx My Mz']
param magnetization_vector:: Full vector magnetization (mx, my, mz) at each grid point.
type magnetization_vector:: Float[Array, 'Vx Vy Vz 3']
param grid_shape:: Grid dimensions (Nx, Ny, Nz) (static – a compile-time constant; changing it triggers retracing). Default is (1, 1, 1).
type grid_shape:: tuple[int, int, int], default: (1, 1, 1)
param symbols:: Element symbols per species (static – compile-time constants; changing them triggers retracing). Default is empty tuple.
type symbols:: tuple[str, ...], default: ()
param atom_counts:: Number of atoms per species. Default is None (replaced by an empty int32 array).
type atom_counts:: Optional[Int[Array, 'S']], default: None
returns:: vol – Validated SOC volumetric data with float64/int32 arrays.
rtype:: SOCVolumetricData
raises ValueError:: If grid_shape does not match the charge, scalar magnetization, or vector-magnetization grids.
raises EquinoxRuntimeError:: If the lattice or any volumetric grid contains a non-finite value.

Notes

Static validation raises ValueError before traced construction when grid_shape differs from a scalar or vector grid shape. Traced validation uses eqx.error_if and raises EquinoxRuntimeError when the lattice or any charge or magnetization grid is non-finite.

See also

SOCVolumetricData: The PyTree class constructed by this factory.
make_volumetric_data: Factory for the non-SOC variant.

diffpes.types.make_volumetric_data(lattice: Float[Array, '3 3'], coords: Float[Array, 'N 3'], charge: Float[Array, 'Cx Cy Cz'], magnetization: Float[Array, 'Mx My Mz'] | None = None, grid_shape: tuple[int, int, int] = (1, 1, 1), symbols: tuple[str, ...] = (), atom_counts: Int[Array, 'S'] | None = None) → VolumetricData[source]¶

Create a validated VolumetricData instance.

The factory validates and normalizes CHGCAR volumetric data before it constructs a VolumetricData PyTree. It casts numerical arrays to float64 and atom counts to int32. The factory casts magnetization only when it is present. Thus, None continues to identify non-spin-polarized calculations. The factory substitutes an empty int32 array when atom_counts is None.

@jaxtyped(typechecker=beartype) checks the shape and dtype constraints at call time.

See:: TestMakeVolumetricData

Implementation Logic¶

Prepare the normalized values:
```
lattice_arr = jnp.asarray(lattice, dtype=jnp.float64)
```
This expression gives the later validation steps a stable shape and dtype.
Apply static validation:
```
charge_arr.shape != grid_shape
```
This predicate rejects invalid structure before JAX traces the numerical checks.
Apply traced validation:
```
~jnp.all(jnp.isfinite(lattice_arr))
```
This predicate remains active during eager and compiled execution.
Return the named instance:
```
return vol
```
The explicit name keeps the implementation and the Returns section synchronized.

param lattice:: Real-space lattice vectors as rows, in Angstroms.
type lattice:: Float[Array, '3 3']
param coords:: Fractional atomic coordinates.
type coords:: Float[Array, 'N 3']
param charge:: Charge density on 3-D grid (electrons per unit cell volume).
type charge:: Float[Array, 'Cx Cy Cz']
param magnetization:: Magnetization density (spin-up minus spin-down). Default is None (non-spin-polarized).
type magnetization:: Optional[Float[Array, 'Mx My Mz']], default: None
param grid_shape:: Grid dimensions (Nx, Ny, Nz) (static – a compile-time constant; changing it triggers retracing). Default is (1, 1, 1).
type grid_shape:: tuple[int, int, int], default: (1, 1, 1)
param symbols:: Element symbols per species (static – compile-time constants; changing them triggers retracing). Default is empty tuple.
type symbols:: tuple[str, ...], default: ()
param atom_counts:: Number of atoms per species. Default is None (replaced by an empty int32 array).
type atom_counts:: Optional[Int[Array, 'S']], default: None
returns:: vol – Validated volumetric data with float64/int32 arrays.
rtype:: VolumetricData
raises ValueError:: If grid_shape does not match charge or the optional magnetization grid.
raises EquinoxRuntimeError:: If the lattice or any volumetric grid contains a non-finite value.

Notes

Static validation raises ValueError before traced construction when grid_shape differs from the charge or magnetization shape. Traced validation uses eqx.error_if and raises EquinoxRuntimeError when the lattice, charge, or present magnetization contains a non-finite value.

See also

VolumetricData: The PyTree class constructed by this factory.
make_soc_volumetric_data: Factory for the SOC variant with vector magnetization.

Experiment geometry¶

Define the geometry of an ARPES experiment.

Extended Summary¶

This module defines a JAX-compatible carrier for the beamline, sample, and instrument geometry. The numerical fields remain traced for forward and inverse calculations. The slit selector remains static because it controls the detector convention.

Routine Listings¶

ExperimentGeometry: Store the geometry and resolution of an ARPES experiment.
make_experiment_geometry(): Create a validated geometry for an ARPES experiment.

Notes

The factory removes the norm of the polarization vector. This operation fixes the intensity-scale gauge at the experiment boundary.

class diffpes.types.experiment.ExperimentGeometry(photon_energy_ev: Float[Array, ''], polarization: Complex[Array, '3'], incidence_theta: Float[Array, ''], incidence_phi: Float[Array, ''], sample_azimuth: Float[Array, ''], work_function_ev: Float[Array, ''], inner_potential_ev: Float[Array, ''], temperature_k: Float[Array, ''], energy_resolution_ev: Float[Array, ''], momentum_resolution_inv_ang: Float[Array, ''], mean_free_path_ang: Float[Array, ''], slit: str) → None[source]¶

Bases: Module

Store the geometry and resolution of an ARPES experiment.

See:

TestExperimentGeometry

Variables:

photon_energy_ev (Float[Array, ""]) – Photon energy in eV.
polarization (Complex[Array, "3"]) – Unit polarization vector in the laboratory frame.
incidence_theta (Float[Array, ""]) – Incidence angle from the surface normal in radians.
incidence_phi (Float[Array, ""]) – Azimuthal incidence angle in radians.
sample_azimuth (Float[Array, ""]) – Sample rotation about the surface normal in radians.
work_function_ev (Float[Array, ""]) – Work function in eV.
inner_potential_ev (Float[Array, ""]) – Inner potential in eV.
temperature_k (Float[Array, ""]) – Sample temperature in kelvin.
energy_resolution_ev (Float[Array, ""]) – Full width at half maximum of the energy resolution in eV.
momentum_resolution_inv_ang (Float[Array, ""]) – Full width at half maximum of the momentum resolution in 1/Angstrom.
mean_free_path_ang (Float[Array, ""]) – Mean free path of the photoelectron in Angstrom.
slit (str) – Detector slit orientation. This field is static. A change causes JAX to retrace the receiving function.

Notes

The traced fields support calibration and geometry inversion. The normalized polarization removes its intensity-scale gauge.

See also

make_experiment_geometry: Create a validated geometry for an ARPES experiment.

photon_energy_ev: Float[Array, '']¶

polarization: Complex[Array, '3']¶

incidence_theta: Float[Array, '']¶

incidence_phi: Float[Array, '']¶

sample_azimuth: Float[Array, '']¶

work_function_ev: Float[Array, '']¶

inner_potential_ev: Float[Array, '']¶

temperature_k: Float[Array, '']¶

energy_resolution_ev: Float[Array, '']¶

momentum_resolution_inv_ang: Float[Array, '']¶

mean_free_path_ang: Float[Array, '']¶

slit: str¶

__init__(photon_energy_ev: Float[Array, ''], polarization: Complex[Array, '3'], incidence_theta: Float[Array, ''], incidence_phi: Float[Array, ''], sample_azimuth: Float[Array, ''], work_function_ev: Float[Array, ''], inner_potential_ev: Float[Array, ''], temperature_k: Float[Array, ''], energy_resolution_ev: Float[Array, ''], momentum_resolution_inv_ang: Float[Array, ''], mean_free_path_ang: Float[Array, ''], slit: str) → None¶

diffpes.types.experiment.make_experiment_geometry(photon_energy_ev: float | Float[Array, ''], polarization: Complex[Array, '3'], incidence_theta: float | Float[Array, ''] = 0.0, incidence_phi: float | Float[Array, ''] = 0.0, sample_azimuth: float | Float[Array, ''] = 0.0, work_function_ev: float | Float[Array, ''] = 4.0, inner_potential_ev: float | Float[Array, ''] = 10.0, temperature_k: float | Float[Array, ''] = 10.0, energy_resolution_ev: float | Float[Array, ''] = 0.02, momentum_resolution_inv_ang: float | Float[Array, ''] = 0.01, mean_free_path_ang: float | Float[Array, ''] = 10.0, slit: str = 'H') → ExperimentGeometry[source]¶

Create a validated geometry for an ARPES experiment.

The factory converts all numerical inputs to JAX arrays. It then fixes the intensity-scale gauge by normalizing the complex polarization vector.

See:: TestMakeExperimentGeometry

Implementation Logic¶

Convert the inputs:
```
photon_energy = jnp.asarray(photon_energy_ev, dtype=jnp.float64)
```
The conversion gives each traced field a stable JAX dtype.
Validate the traced values:
```
photon_energy = eqx.error_if(photon_energy, invalid, message)
```
The runtime checks remain active during compiled execution.
Normalize the polarization:
```
normalized_polarization = checked_polarization / safe_norm
```
The safe denominator keeps the invalid zero-norm branch finite.
Return the named carrier:
```
return geometry
```
The carrier preserves each valid numerical input as a traced leaf.

param photon_energy_ev:: Photon energy in eV.
type photon_energy_ev:: Union[float, Float[Array, '']]
param polarization:: Nonzero complex polarization vector in the laboratory frame.
type polarization:: Complex[Array, '3']
param incidence_theta:: Incidence angle from the surface normal in radians. Default is 0.0.
type incidence_theta:: Union[float, Float[Array, '']], default: 0.0
param incidence_phi:: Azimuthal incidence angle in radians. Default is 0.0.
type incidence_phi:: Union[float, Float[Array, '']], default: 0.0
param sample_azimuth:: Sample rotation about the surface normal in radians. Default is 0.0.
type sample_azimuth:: Union[float, Float[Array, '']], default: 0.0
param work_function_ev:: Work function in eV. Default is 4.0.
type work_function_ev:: Union[float, Float[Array, '']], default: 4.0
param inner_potential_ev:: Inner potential in eV. Default is 10.0.
type inner_potential_ev:: Union[float, Float[Array, '']], default: 10.0
param temperature_k:: Sample temperature in kelvin. Default is 10.0.
type temperature_k:: Union[float, Float[Array, '']], default: 10.0
param energy_resolution_ev:: Energy-resolution width in eV. Default is 0.02.
type energy_resolution_ev:: Union[float, Float[Array, '']], default: 0.02
param momentum_resolution_inv_ang:: Momentum-resolution width in 1/Angstrom. Default is 0.01.
type momentum_resolution_inv_ang:: Union[float, Float[Array, '']], default: 0.01
param mean_free_path_ang:: Mean free path in Angstrom. Default is 10.0.
type mean_free_path_ang:: Union[float, Float[Array, '']], default: 10.0
param slit:: Detector slit orientation. This value is static. A change causes retracing. Default is "H".
type slit:: str, default: "H"
returns:: geometry – Validated experiment geometry with a unit polarization vector.
rtype:: ExperimentGeometry
raises ValueError:: If slit is not "H" or "V".
raises EquinoxRuntimeError:: If an input is non-finite or violates its physical range. The factory also rejects a zero polarization vector.

Notes

The normalization is differentiable for every accepted polarization. The factory assigns no derivative at the rejected zero vector.

K-space carriers¶

Define k-space path and grid data structures.

Extended Summary¶

This module defines carriers for generated k-space paths and fixed-shape grids. It also keeps KPathInfo for metadata from VASP KPOINTS files.

Routine Listings¶

KPathInfo: Store k-point path metadata in a JAX PyTree.
KPath: Store a generated path through fractional k-space.
KGrid: Store a fixed-shape raster in fractional k-space.
make_kpath_info(): Create a validated KPathInfo instance.
make_kpath(): Create a validated path through fractional k-space.
make_kgrid(): Create a validated fixed-shape k-space raster.

class diffpes.types.kpath.KPathInfo(num_kpoints: Int[Array, ''], label_indices: Int[Array, 'L'], points_per_segment: Int[Array, ''], segments: Int[Array, ''], kpoints: Float[Array, 'K 3'] | None, weights: Float[Array, 'K'] | None, grid: Int[Array, '3'] | None, shift: Float[Array, '3'] | None, mode: str, labels: tuple[str, ...], comment: str, coordinate_mode: str) → None[source]¶

Bases: Module

Store k-point path metadata in a JAX PyTree.

JAX stores numerical fields as PyTree children. It stores string metadata as auxiliary data because JAX cannot trace Python strings.

See:

TestKPathInfo

Variables:

num_kpoints (Int[Array, " "]) – Total number of k-points in the path (line mode) or header count (explicit mode).
label_indices (Int[Array, " L"]) – Indices of symmetry points along the path.
points_per_segment (Int[Array, " "]) – Raw integer from line 2 of KPOINTS (line mode: points per segment).
segments (Int[Array, " "]) – Number of line segments in line mode.
kpoints (Optional[Float[Array, "K 3"]]) – Mode-specific k-points: line mode -> segment endpoints (segments + 1), explicit mode -> listed k-points, automatic mode -> None.
weights (Optional[Float[Array, " K"]]) – Explicit-mode per-k-point weights (None otherwise).
grid (Optional[Int[Array, " 3"]]) – Automatic-mode Monkhorst-Pack/Gamma grid (None otherwise).
shift (Optional[Float[Array, " 3"]]) – Automatic-mode grid shift (None otherwise).
mode (str) – KPOINTS file mode (Automatic, Line-mode, Explicit; static – a compile-time constant; changing it triggers retracing).
labels (tuple[str, ]) – Symmetry point labels (e.g., Gamma, M, K; static – compile-time constants; changing them triggers retracing).
comment (str) – Raw comment from KPOINTS line 1 (static – a compile-time constant; changing it triggers retracing).
coordinate_mode (str) – Coordinate/scheme line metadata: line/explicit -> Reciprocal/Cartesian line, automatic -> scheme line (e.g., Monkhorst-Pack; static – a compile-time constant; changing it triggers retracing).

Notes

See also

make_kpath_info: Factory function with validation and int32 casting.

num_kpoints: Int[Array, '']¶

label_indices: Int[Array, 'L']¶

points_per_segment: Int[Array, '']¶

segments: Int[Array, '']¶

kpoints: Float[Array, 'K 3'] | None¶

weights: Float[Array, 'K'] | None¶

grid: Int[Array, '3'] | None¶

shift: Float[Array, '3'] | None¶

mode: str¶

labels: tuple[str, ...]¶

comment: str¶

coordinate_mode: str¶

__init__(num_kpoints: Int[Array, ''], label_indices: Int[Array, 'L'], points_per_segment: Int[Array, ''], segments: Int[Array, ''], kpoints: Float[Array, 'K 3'] | None, weights: Float[Array, 'K'] | None, grid: Int[Array, '3'] | None, shift: Float[Array, '3'] | None, mode: str, labels: tuple[str, ...], comment: str, coordinate_mode: str) → None¶

class diffpes.types.kpath.KPath(kpoints: Float[Array, 'n_k 3'], kz: Float[Array, ''] | None, labels: tuple[str, ...], label_indices: tuple[int, ...], n_per_segment: int) → None[source]¶

Bases: Module

Store a generated path through fractional k-space.

See:

TestKPath

Variables:

kpoints (Float[Array, "n_k 3"]) – Fractional k-points.
kz (Optional[Float[Array, ""]]) – Fixed Cartesian out-of-plane momentum in 1/Angstrom. None means that the path has no separate fixed value.
labels (tuple[str, ]) – Labels for the path anchors. This field is static. A change causes JAX to retrace the receiving function.
label_indices (tuple[int, ]) – Indices for the labels. This field is static. A change causes JAX to retrace the receiving function.
n_per_segment (int) – Number of points in each segment. This field is static. A change causes JAX to retrace the receiving function.

See also

KPathInfo: Store k-point path metadata in a JAX PyTree.
make_kpath: Create a validated path through fractional k-space.

kpoints: Float[Array, 'n_k 3']¶

kz: Float[Array, ''] | None¶

labels: tuple[str, ...]¶

label_indices: tuple[int, ...]¶

n_per_segment: int¶

__init__(kpoints: Float[Array, 'n_k 3'], kz: Float[Array, ''] | None, labels: tuple[str, ...], label_indices: tuple[int, ...], n_per_segment: int) → None¶

class diffpes.types.kpath.KGrid(kpoints: Float[Array, 'n_k 3'], kz: Float[Array, ''] | None, photon_energy_axis_ev: Float[Array, 'n_rows'] | None, mesh_shape: tuple[int, int]) → None[source]¶

Bases: Module

Store a fixed-shape raster in fractional k-space.

The flattened coordinates remain traced for geometry inversion. The static mesh shape lets compiled consumers restore the raster without a data-dependent shape operation.

See:

TestKGrid

Variables:

kpoints (Float[Array, "n_k 3"]) – Flattened fractional k-points.
kz (Optional[Float[Array, ""]]) – Fixed Cartesian out-of-plane momentum in 1/Angstrom. None marks a grid with varying or unspecified out-of-plane momentum.
photon_energy_axis_ev (Optional[Float[Array, " n_rows"]]) – Photon energy for each raster row in eV. None marks a grid without a photon-energy axis.
mesh_shape (tuple[int, int]) – Raster shape as (n_rows, n_cols). This field is static. A change causes JAX to retrace the receiving function.

See also

make_kgrid: Create a validated fixed-shape k-space raster.

kpoints: Float[Array, 'n_k 3']¶

kz: Float[Array, ''] | None¶

photon_energy_axis_ev: Float[Array, 'n_rows'] | None¶

mesh_shape: tuple[int, int]¶

__init__(kpoints: Float[Array, 'n_k 3'], kz: Float[Array, ''] | None, photon_energy_axis_ev: Float[Array, 'n_rows'] | None, mesh_shape: tuple[int, int]) → None¶

diffpes.types.kpath.make_kpath_info(num_kpoints: int | Int[Array, ''], label_indices: Int[Array, 'L'] | list[int], points_per_segment: int | Int[Array, ''] = 0, segments: int | Int[Array, ''] = 0, kpoints: Float[Array, 'K 3'] | None = None, weights: Float[Array, 'K'] | None = None, grid: Int[Array, '3'] | list[int] | None = None, shift: Float[Array, '3'] | None = None, mode: str = 'Line-mode', labels: tuple[str, ...] = (), comment: str = '', coordinate_mode: str = '') → KPathInfo[source]¶

Create a validated KPathInfo instance.

@jaxtyped(typechecker=beartype) checks shape constraints at call time. The factory passes string metadata unchanged and stores it as PyTree auxiliary data.

See:: TestMakeKPathInfo

Implementation Logic¶

Prepare the normalized values:
```
nkpts_arr = jnp.asarray(num_kpoints, dtype=jnp.int32)
```
This expression gives the later validation steps a stable shape and dtype.
Apply static validation:
```
mode not in _KPATH_MODES
```
This predicate rejects invalid structure before JAX traces the numerical checks.
Apply traced validation:
```
~(nkpts_arr >= 0)
```
This predicate remains active during eager and compiled execution.
Return the named instance:
```
return kpath
```
The explicit name keeps the implementation and the Returns section synchronized.

param num_kpoints:: Total number of k-points along the path.
type num_kpoints:: Union[int, Int[Array, '']]
param label_indices:: Indices of symmetry points along the path.
type label_indices:: Union[Int[Array, 'L'], list[int]]
param points_per_segment:: Raw value from line 2 of KPOINTS. Default is 0.
type points_per_segment:: Union[int, Int[Array, '']], default: 0
param segments:: Number of path segments in line mode. Default is 0.
type segments:: Union[int, Int[Array, '']], default: 0
param kpoints:: Mode-specific k-point coordinates. Default is None.
type kpoints:: Optional[Float[Array, 'K 3']], default: None
param weights:: Explicit-mode weights. Default is None.
type weights:: Optional[Float[Array, 'K']], default: None
param grid:: Automatic-mode MP/Gamma grid. Default is None.
type grid:: Union[Int[Array, '3'], list[int], None], default: None
param shift:: Automatic-mode grid shift. Default is None.
type shift:: Optional[Float[Array, '3']], default: None
param mode:: KPOINTS file mode (static – a compile-time constant; changing it triggers retracing). Default is "Line-mode".
type mode:: str, default: "Line-mode"
param labels:: Symmetry point labels (static – compile-time constants; changing them triggers retracing). Default is empty tuple.
type labels:: tuple[str, ...], default: ()
param comment:: KPOINTS comment line (static – a compile-time constant; changing it triggers retracing). Default is empty string.
type comment:: str, default: ""
param coordinate_mode:: Coordinate/scheme line metadata (static – a compile-time constant; changing it triggers retracing). Default is empty string.
type coordinate_mode:: str, default: ""
returns:: kpath – Validated k-path info instance.
rtype:: KPathInfo
raises ValueError:: If mode has an unsupported value. The function also rejects incomplete or inconsistent line-mode metadata.
raises EquinoxRuntimeError:: If num_kpoints is negative, a traced line-mode segment count is less than one, or supplied kpoints, weights, or shift values are non-finite.

Notes

See also

KPathInfo: The PyTree class constructed by this factory.

diffpes.types.kpath.make_kpath(kpoints: Float[Array, 'n_k 3'], labels: tuple[str, ...] = (), label_indices: tuple[int, ...] = (), n_per_segment: int = 1, kz: float | Float[Array, ''] | None = None) → KPath[source]¶

Create a validated path through fractional k-space.

The factory validates the static plotting metadata and the traced path coordinates. It preserves an optional fixed Cartesian out-of-plane value.

See:: TestMakeKPath

Implementation Logic¶

Validate the path structure:
```
label_indices[-1] < kpoints.shape[0]
```
The static checks keep labels within the fixed path shape.
Validate the traced values:
```
checked_kpoints = eqx.error_if(kpoints_array, invalid, message)
```
The runtime checks remain active during compiled execution.
Return the named path:
```
return kpath
```
The result keeps only the plotting metadata static.

param kpoints:: Fractional k-points.
type kpoints:: Float[Array, 'n_k 3']
param labels:: Labels for the path anchors. This value is static. A change causes retracing. Default is an empty tuple.
type labels:: tuple[str, ...], default: ()
param label_indices:: Indices for the labels. This value is static. A change causes retracing. Default is an empty tuple.
type label_indices:: tuple[int, ...], default: ()
param n_per_segment:: Number of points in each segment. This value is static. A change causes retracing. Default is 1.
type n_per_segment:: int, default: 1
param kz:: Fixed Cartesian out-of-plane momentum in 1/Angstrom. Default is None.
type kz:: Union[float, Float[Array, ''], None], default: None
returns:: kpath – Validated path with traced fractional coordinates.
rtype:: KPath
raises ValueError:: If the static metadata has invalid lengths, indices, or segment size.
raises EquinoxRuntimeError:: If a traced k-point or fixed out-of-plane value is non-finite.

Notes

The k-points and optional kz value carry gradients. The labels, indices, and segment size do not carry gradients.

diffpes.types.kpath.make_kgrid(kpoints: Float[Array, 'n_k 3'], mesh_shape: tuple[int, int], kz: float | Float[Array, ''] | None = None, photon_energy_axis_ev: Float[Array, 'n_rows'] | None = None) → KGrid[source]¶

Create a validated fixed-shape k-space raster.

The factory validates the static raster shape and traced numerical data. It preserves an optional photon-energy axis for a photon-energy map.

See:: TestMakeKGrid

Implementation Logic¶

Validate the raster shape:
```
mesh_shape[0] * mesh_shape[1] == kpoints.shape[0]
```
This check keeps the flattened point count consistent with the raster.
Validate the traced values:
```
checked_kpoints = eqx.error_if(kpoints_array, invalid, message)
```
The runtime checks remain active during compiled execution.
Return the named grid:
```
return kgrid
```
The result keeps only the raster shape static.

param kpoints:: Flattened fractional k-points.
type kpoints:: Float[Array, 'n_k 3']
param mesh_shape:: Raster shape as (n_rows, n_cols). This value is static. A change causes retracing.
type mesh_shape:: tuple[int, int]
param kz:: Fixed Cartesian out-of-plane momentum in 1/Angstrom. Default is None.
type kz:: Union[float, Float[Array, ''], None], default: None
param photon_energy_axis_ev:: Photon energy for each raster row in eV. Default is None.
type photon_energy_axis_ev:: Optional[Float[Array, 'n_rows']], default: None
returns:: kgrid – Validated grid with traced fractional coordinates.
rtype:: KGrid
raises ValueError:: If the mesh dimensions are not positive or their product is wrong. The factory also rejects a photon-energy axis with the wrong length.
raises EquinoxRuntimeError:: If traced data is non-finite or a photon energy is not positive.

Notes

The mesh shape controls array reshaping and remains static. All numerical fields carry gradients when they are present.