diffpes.inout¶

VASP file parsers (POSCAR, EIGENVAL, KPOINTS, DOSCAR, PROCAR, CHGCAR), typed-result HDF5 persistence, portable JSON/HDF5 forward certificates, workflow helpers, and plotting utilities for ARPES simulation input.

Parse VASP files for ARPES simulation input.

Extended Summary¶

The subpackage parses VASP output files into PyTrees for ARPES simulations. It supports POSCAR, EIGENVAL, KPOINTS, DOSCAR, PROCAR, and CHGCAR files. It also provides HDF5 persistence, workflow helpers, and plotting utilities.

The following list describes the submodules:

chgcar
Parse a VASP CHGCAR file.
certificate
Persist forward-model certificates in portable formats.
doscar
Parse a VASP DOSCAR file.
eigenval
Parse a VASP EIGENVAL file.
hdf5
Serialize and deserialize diffpes PyTrees in HDF5.
helpers
Provide workflow helpers for simulation-ready parser arrays.
kpoints
Parse a VASP KPOINTS file.
plotting
Plot ARPES spectra with analysis utilities.
poscar
Parse a VASP POSCAR/CONTCAR file.
procar
Parse a VASP PROCAR file.

Routine Listings¶

aggregate_atoms(): Sum orbital projections over a set of atoms.
attach_certificate_h5(): Attach a certificate atomically to an HDF5 result file.
certificate_identity(): Compute the scientific identity of a canonical certificate.
apply_kpath_ticks(): Apply symmetry-point ticks/labels from KPathInfo to an axis.
check_consistency(): Check dimension agreement across parsed VASP files.
list_band_scatter_presets(): Return supported preset names for projected band scatter plots.
load_from_h5(): Load PyTrees from an HDF5 file.
finalize_certificate(): Replace the kernel placeholder with the canonical identity.
load_certificate_h5(): Load a certificate embedded in an HDF5 result file.
load_certificate_json(): Load a validated forward certificate from canonical JSON.
plot_arpes_spectrum(): Plot an ARPES intensity map from an ArpesSpectrum PyTree.
plot_arpes_with_kpath(): Plot ARPES spectrum and annotate k-axis using KPathInfo.
plot_band_scatter_preset(): Plot projected bands as marker-size-weighted scatter points.
plot_band_scatter_with_kpath(): Plot projected band scatter and annotate x-axis with k-path labels.
read_chgcar(): Parse a VASP CHGCAR file.
read_doscar(): Parse a VASP DOSCAR file.
read_eigenval(): Parse a VASP EIGENVAL file.
read_kpoints(): Parse a VASP KPOINTS file.
read_poscar(): Parse a VASP POSCAR/CONTCAR file.
read_procar(): Parse a VASP PROCAR file.
reduce_orbitals(): Reduce 9 orbital channels to s/p/d totals.
save_to_h5(): Save one or more named PyTrees to an HDF5 file.
save_certificate_json(): Save a forward certificate atomically as canonical JSON.
select_atoms(): Extract orbital projections for a subset of atoms.

Notes

All parsers use standard Python I/O because file parsing is sequential. Factory functions convert the parsed data to JAX arrays.

diffpes.inout.attach_certificate_h5(path: str | Path, name: str, certificate: ForwardCertificate) → None[source]¶

Attach a certificate atomically to an HDF5 result file.

The function updates the complete file through a same-directory temporary. It preserves existing numerical result groups.

See:: TestAttachCertificateH5
Return type:: None

Implementation Logic¶

Encode the certificate:

document = _certificate_document(certificate)
data = _json_bytes(document, newline=True)

The HDF5 record stores the same canonical bytes as JSON persistence.

Copy the current container:
```
shutil.copy2(destination, temporary)
```
An existing result file remains intact while the copy changes.

Write and replace the container:

_write_h5_record(temporary, name, data, certificate)
os.replace(temporary, destination)
temporary.unlink(missing_ok=True)

Replacement publishes the complete file. Failure removes the temporary.

param path:: Existing HDF5 result path, or a path for a new HDF5 container.
type path:: str | Path
param name:: Name of one result under the certificate index group.
type name:: str
param certificate:: Certificate associated with the named result.
type certificate:: ForwardCertificate
raises BaseException:: If copying, writing, or replacing the HDF5 file fails.

diffpes.inout.certificate_identity(certificate: ForwardCertificate) → str[source]¶

Compute the scientific identity of a canonical certificate.

The identity covers scientific and numerical fields. It excludes the self-reference, audit execution ID, and wall-clock timestamp. The CRC32 value detects accidental mismatches and does not authenticate the record.

See:: TestCertificateIdentity

Implementation Logic¶

Compute the canonical identity:
```
identity = _document_identity(document)
```
The canonical payload omits only the self-reference and audit fields.

param certificate:: Concrete certificate at the persistence boundary.
type certificate:: ForwardCertificate
returns:: identity – Stable non-security identity for the scientific execution record.
rtype:: str

diffpes.inout.finalize_certificate(certificate: ForwardCertificate) → ForwardCertificate[source]¶

Replace the kernel placeholder with the canonical identity.

Canonical encoding stays outside JAX tracing because it has no scientific derivative. The returned certificate retains every scientific leaf.

See:: TestFinalizeCertificate

Implementation Logic¶

Replace the placeholder:
```
identity = certificate_identity(certificate)
```
The factory copies every other certificate field without modification.

param certificate:: Concrete certificate produced by the JAX-native execution kernel.
type certificate:: ForwardCertificate
returns:: result – Equivalent certificate with its final canonical identity.
rtype:: ForwardCertificate

diffpes.inout.load_certificate_h5(path: str | Path, name: str) → ForwardCertificate[source]¶

Load a certificate embedded in an HDF5 result file.

The persistence operation retains the complete scientific-assurance record and its JAX array leaves. Consistency checks detect accidental storage corruption.

See:: TestLoadCertificateH5

Implementation Logic¶

Resolve the stored record:
```
root = file[CERTIFICATE_H5_GROUP]
group = root[name]
```
Missing groups or names raise KeyError before decoding.

Decode the canonical bytes:

data = stored.tobytes()
document = _read_document(data)
certificate = _certificate_from_document(document)

The decoder validates the persisted schema and consistency check.

Validate the convenience index:
```
msg: str = f"HDF5 certificate index mismatch for {key!r}"
```
Every HDF5 attribute must agree with the canonical JSON record.

param path:: HDF5 result path.
type path:: str | Path
param name:: Certificate name supplied to attach_certificate_h5().
type name:: str
returns:: certificate – Reconstructed and validated certificate.
rtype:: ForwardCertificate
raises KeyError:: If the certificate group or named record is absent.
raises ValueError:: If the exact JSON bytes or HDF5 convenience index are inconsistent.

diffpes.inout.load_certificate_json(path: str | Path) → ForwardCertificate[source]¶

Load a validated forward certificate from canonical JSON.

The persistence operation retains the complete scientific-assurance record and its JAX array leaves. Consistency checks detect accidental storage corruption.

See:: TestLoadCertificateJson

Implementation Logic¶

Read and validate the document:
```
data = Path(path).read_bytes()
document = _read_document(data)
```
The decoder checks the schema and consistency checksum before use.
Reconstruct the carrier:
```
certificate = _certificate_from_document(document)
```
The decoder restores persisted numerical leaves as JAX arrays.

param path:: Source JSON path.
type path:: str | Path
returns:: certificate – Reconstructed certificate with numerical leaves restored as JAX arrays.
rtype:: ForwardCertificate

diffpes.inout.save_certificate_json(certificate: ForwardCertificate, path: str | Path) → None[source]¶

Save a forward certificate atomically as canonical JSON.

The persistence operation retains the complete scientific-assurance record and its JAX array leaves. Consistency checks detect accidental storage corruption.

See:: TestSaveCertificateJson
Return type:: None

Implementation Logic¶

Build the certificate document:
```
document = _certificate_document(certificate)
data = _json_bytes(document, newline=True)
```
The document includes the schema and a non-security consistency check.
Replace the destination atomically:
```
_atomic_write(Path(path), data)
```
A same-directory temporary prevents a partial JSON record.

param certificate:: Validated scientific-assurance record to persist.
type certificate:: ForwardCertificate
param path:: Destination JSON path. Its parent directory must already exist.
type path:: str | Path

diffpes.inout.read_chgcar(filename: str = 'CHGCAR') → VolumetricData | SOCVolumetricData[source]¶

Parse a VASP CHGCAR file.

The parser supports three layouts:

ISPIN=1: 1 grid block (charge only).
ISPIN=2: 2 grid blocks (charge + scalar magnetization).
SOC (LSORBIT): 4 grid blocks (charge, mx, my, mz).

The CHGCAR file starts with a POSCAR-style structural header. One or more volumetric blocks follow the header on a real-space FFT grid. Each block starts with the three grid dimensions. VASP then writes the flattened values in Fortran column-major order.

See:: TestReadChgcar

Implementation Logic¶

Read the structural header and remaining records:

path: Path = Path(filename)
with path.open("r") as fid:
    lattice, coords, symbols, atom_counts = _read_poscar_header(fid)
    rest_lines: list[str] = [
        line.rstrip("\n") for line in fid
    ]

This separates the geometry from the subsequent volumetric blocks.

Normalize the charge grid by the cell volume:

charge_grid: Float[NDArray, "Nx Ny Nz"] = (
    charge_vals.reshape(grid_shape, order="F") / volume
)

This converts VASP charge-times-volume values to charge density.

Return the carrier for the detected layout:
```
return volumetric
```
The grid-count branch binds scalar or SOC data to one output name.

param filename:: Path to the CHGCAR file. Default is "CHGCAR".
type filename:: str, default: "CHGCAR"
returns:: volumetric – VolumetricData for ISPIN=1 or ISPIN=2 files. SOCVolumetricData for SOC files with four grid blocks.
rtype:: VolumetricData | SOCVolumetricData
raises ValueError:: If the lattice volume is zero, the grid dimensions are absent, or a data block is incomplete.

Notes

The returned densities use electrons per cubic Angstrom. The parser divides the raw VASP values by the cell volume. It stores coordinates in fractional form and preserves the VASP Fortran grid order.

diffpes.inout.read_doscar(filename: str = 'DOSCAR', return_mode: Literal['legacy', 'full'] = 'legacy') → DensityOfStates | FullDensityOfStates[source]¶

Parse a VASP DOSCAR file.

The function reads a VASP DOSCAR file that contains total and optional site-projected) density of states on a uniform energy grid.

The DOSCAR file format written by VASP consists of:

Header (6 lines):
- Line 1: system header with NATOMS as the first integer.
- Lines 2-5: additional metadata (unused here).
- Line 6: EMIN EMAX NEDOS EFERMI ... – energy window bounds, number of DOS grid points, and the Fermi energy.
Total DOS block (NEDOS lines): each line contains the energy value followed by density-of-states columns.
- ISPIN=1: 3 columns – energy, DOS_up, intDOS_up.
- ISPIN=2: 5 columns – energy, DOS_up, DOS_down, intDOS_up, intDOS_down.
Per-atom PDOS blocks: Each optional block contains NEDOS lines. The header has the same format as line 6 of the main header. Orbital-projected DOS values follow the header. LORBIT and spin polarization determine the column count.

See:: TestReadDoscar

Implementation Logic¶

Read the header and total-DOS dimensions:

path: Path = Path(filename)
with path.open("r") as fid:
    header: list[str] = fid.readline().split()
    natoms: int = int(header[0])

This establishes the atom count before the function allocates the total and projected data.

Allocate and populate the total-DOS table:
```
data: Float[NDArray, "E C"] = np.zeros(
    (nedos, ncols), dtype=np.float64
)
```
This preserves each column until the function knows the return mode.
Return the selected DOS carrier:
```
return dos
```
Both branches bind their validated result to dos.

param filename:: Path to DOSCAR file. Default is "DOSCAR".
type filename:: str, default: "DOSCAR"
param return_mode:: "legacy" (default) returns a DensityOfStates with only spin-up total DOS (backward-compatible). "full" returns a FullDensityOfStates with both spin channels, integrated DOS, and PDOS blocks when present.
type return_mode:: Literal['legacy', 'full'], default: "legacy"
returns:: dos – Density of states data.
rtype:: Union[DensityOfStates, FullDensityOfStates]

Notes

In "full" mode, the parser also reads each PDOS block after the total DOS section. Each PDOS block has NEDOS energy points. LORBIT determines the VASP orbital order. For example, LORBIT=11 starts with s, p_y, p_z, p_x, d_{xy}. The parser reads the Fermi energy from column 4 of line 6.

diffpes.inout.read_eigenval(filename: str = 'EIGENVAL', fermi_energy: float | Float[Array, ''] = 0.0, return_mode: Literal['legacy', 'full'] = 'legacy') → BandStructure | SpinBandStructure[source]¶

Parse a VASP EIGENVAL file.

The function reads the electronic eigenvalues at each k-point from a VASP EIGENVAL file. It supports nonpolarized ISPIN=1 and spin-polarized ISPIN=2 computations.

The EIGENVAL file format written by VASP consists of:

Header (6 lines):
- Line 1: four integers including ISPIN as the 4th value. ISPIN=2 indicates spin-polarized eigenvalues.
- Lines 2-5: additional metadata (unused here).
- Line 6: NELECT NKPOINTS NBANDS – number of electrons, k-points, and bands.
K-point / band blocks (NKPOINTS blocks): each block starts with a k-point line of 4 floats (kx ky kz weight), followed by NBANDS eigenvalue lines.
- ISPIN=1: each band line has 2 values – band_index energy.
- ISPIN=2: each band line has 3 values – band_index energy_up energy_down.

See:: TestReadEigenval

Implementation Logic¶

Read the spin and dimension metadata:

path: Path = Path(filename)
with path.open("r") as fid:
    first_line: list[str] = fid.readline().split()
    ispin: int = int(first_line[3])

This step fixes the static file layout before the function allocates the band arrays.

Sort each spin channel by band energy:
```
eigen_up = np.take_along_axis(eigen_up, order_up, axis=1)
```
This gives every k-point the ascending band convention used downstream.
Return the requested band carrier:
```
return bands
```
The return-mode branch binds scalar or spin-resolved data to one output.

param filename:: Path to EIGENVAL file. Default is "EIGENVAL".
type filename:: str, default: "EIGENVAL"
param fermi_energy:: Fermi level in eV used to reference the eigenvalues. The function accepts Python scalars and traced scalar arrays. The default is 0.0.
type fermi_energy:: Union[float, Float[Array, '']], default: 0.0
param return_mode:: "legacy" (default) returns a BandStructure with only spin-up eigenvalues (backward-compatible). "full" returns a SpinBandStructure with both spin channels when ISPIN=2, or a BandStructure when ISPIN=1.
type return_mode:: Literal['legacy', 'full'], default: "legacy"
returns:: bands – Band structure with eigenvalues and k-points. The type depends on return_mode and the spin polarization.
rtype:: Union[BandStructure, SpinBandStructure]
raises ValueError:: If a k-point line has fewer than four values. If a band line has fewer values than the selected ISPIN mode requires. If the parser reaches EOF before it reads all blocks.

Notes

For ISPIN=2 computations, each eigenvalue line contains the index and two energy columns. In "legacy" mode, the parser reads only the spin-up energy. In "full" mode, the parser retains both channels in a SpinBandStructure. The EIGENVAL file does not contain the Fermi energy. Read it separately from a DOSCAR or OUTCAR file. The parser retains the eigenvalues without a fermi_energy shift. It stores the Fermi energy in the returned PyTree for downstream shifts.

diffpes.inout.load_from_h5(path: str | Path, name: str | None = None) → Any[source]¶

Load PyTrees from an HDF5 file.

The function deserializes HDF5 groups into diffpes PyTrees. It reads the datasets as JAX arrays and reconstructs each Equinox module with keyword arguments.

See:: TestLoadFromH5

Implementation Logic¶

Open the requested HDF5 path:
```
file_path: Path = Path(path)
with h5py.File(file_path, "r") as f:
```
This gives named and all-group loads the same read-only file boundary.
Reconstruct registered groups recursively:
```
loaded: Any = _load_group(f[name])
```
The nested loader restores child arrays, optional fields, and metadata. It then calls the registered Equinox carrier class.
Return the selected load result:
```
return loaded
```
A named request returns one carrier. Other requests return a mapping.

param path:: File path to the HDF5 file to read.
type path:: Union[str, Path]
param name:: Name of a specific group to load. If None, the function loads all groups and returns a dictionary.
type name:: Optional[str], default: None
returns:: loaded – One PyTree when name identifies a group. Otherwise, a dictionary that maps group names to PyTree instances.
rtype:: Any
raises KeyError:: If name identifies no group in the file.
raises TypeError:: If a group’s _pytree_type is not in the registry.

diffpes.inout.save_to_h5(path: str | Path, /, *, compression: str | None = None, compression_opts: Any = None, shuffle: bool = False, fletcher32: bool = False, chunks: bool | tuple[int, ...] | None = None, **pytrees: Any) → None[source]¶

Save one or more named PyTrees to an HDF5 file.

The function serializes each keyword PyTree into a named HDF5 group. JAX array fields become datasets with their Equinox field names. The codec stores static metadata in a JSON group attribute.

See:: TestSaveToH5
Return type:: None

Implementation Logic¶

Reject an empty save request:

if not pytrees:
    msg: str = "At least one PyTree must be provided."
    raise ValueError(msg)

This prevents creation of a file with no registered carrier groups.

Write each carrier through the registry codec:

file_path: Path = Path(path)
with h5py.File(file_path, "w") as f:
    for group_name, pytree in pytrees.items():
        grp: h5py.Group = f.create_group(group_name)
        _write_module(grp, pytree)

The recursive writer preserves child arrays and static metadata. It applies the storage flags under one group name.

param path:: File path for the HDF5 file to create.
type path:: Union[str, Path]
param compression:: HDF5 compression filter name, for example "gzip" or "lzf". Applied to non-scalar datasets only.
type compression:: Optional[str], default: None
param compression_opts:: Compression options for h5py, for example the gzip level. Must be None when compression is None.
type compression_opts:: Any, default: None
param shuffle:: If True, enable HDF5 shuffle filter on non-scalar datasets.
type shuffle:: bool, default: False
param fletcher32:: If True, enable HDF5 Fletcher32 checksum on non-scalar datasets.
type fletcher32:: bool, default: False
param chunks:: Chunking policy for non-scalar datasets. True enables auto-chunking, or provide an explicit chunk-shape tuple.
type chunks:: Union[bool, tuple[int, ...], None], default: None
param **pytrees:: Named PyTree instances. Each keyword argument name becomes an HDF5 group name.
type **pytrees:: Any
raises ValueError:: If the caller provides no PyTrees.
raises ValueError:: If the caller provides compression_opts without compression.
raises TypeError:: If a PyTree’s class is not in the registry.

Notes

Scalar datasets (shape ()) are always written without HDF5 filter/chunk flags because those options are invalid for scalar dataspace in HDF5.

diffpes.inout.aggregate_atoms(orb: OrbitalProjection, atom_indices: list[int] | None = None) → Float[Array, 'K B 9'][source]¶

Sum orbital projections over a set of atoms.

The function sums the atom axis and produces a (K, B, 9) array. Therefore, each k-point and band pair contains the total orbital weight. The ARPES intensity computation uses this reduction because it needs the aggregate orbital character instead of individual atom contributions.

See:: TestAggregateAtoms

Implementation Logic¶

Select the requested atom data:
```
proj = orb.projections[:, :, idx, :]
proj = orb.projections
```
The optional index limits the reduction to the requested atoms.
Sum the atom axis:
```
result = jnp.sum(proj, axis=2)
```
The reduction keeps the k-point, band, and orbital axes explicit.

param orb:: Full orbital projections with shape (K, B, A, 9).
type orb:: OrbitalProjection
param atom_indices:: 0-based indices of atoms to sum over. If None, sums over all atoms.
type atom_indices:: Optional[list[int]], default: None
returns:: result – Atom-summed orbital projections.
rtype:: Float[Array, 'K B 9']

Notes

This function operates only on the projections field and ignores spin and oam fields. For spin-resolved aggregation, use select_atoms() first and then perform the reduction manually.

diffpes.inout.check_consistency(bands: BandStructure, orb: OrbitalProjection | SpinOrbitalProjection, kpath: KPathInfo | None = None) → None[source]¶

Check dimension agreement across parsed VASP files.

Validates that the k-point and band dimensions are consistent between the EIGENVAL-derived band structure, the PROCAR-derived orbital projections, and (optionally) the KPOINTS-derived path metadata. This is a defensive check intended to catch mismatches caused by mixing output files from different VASP runs.

See:: TestCheckConsistency
Return type:: None

Implementation Logic¶

Read the shared dimensions:

nk_bands = int(bands.eigenvalues.shape[0])
nb_bands = int(bands.eigenvalues.shape[1])
nk_procar = int(orb.projections.shape[0])
nb_procar = int(orb.projections.shape[1])

These static dimensions identify incompatible parser outputs early.

Compare the band and k-point axes:
```
if nk_bands != nk_procar:
if nb_bands != nb_procar:
```
Each mismatch raises ValueError with both observed sizes.
Check optional line-mode metadata:
```
if nk_kpath > 0 and nk_bands != nk_kpath:
```
A positive KPOINTS count must agree with the EIGENVAL count.

param bands:: Parsed EIGENVAL data.
type bands:: BandStructure
param orb:: Parsed PROCAR data.
type orb:: Union[OrbitalProjection, SpinOrbitalProjection]
param kpath:: Parsed KPOINTS data.
type kpath:: Optional[KPathInfo], default: None
raises ValueError:: If k-point or band counts disagree between files.

Notes

This function does not verify atom counts (PROCAR vs POSCAR) because BandStructure does not carry atom information. For atom-count checks, compare orb.projections.shape[2] against geometry.positions.shape[0] manually.

diffpes.inout.reduce_orbitals(projections: Float[Array, 'K B A 9']) → Float[Array, 'K B A 3'][source]¶

Reduce 9 orbital channels to s/p/d totals.

Collapses the 9-channel VASP orbital decomposition (s, py, pz, px, dxy, dyz, dz2, dxz, dx2-y2) into three angular-momentum shell totals. This is useful for coarse-grained orbital-character analysis, for example fat-band colors from s/p/d weight).

See:: TestReduceOrbitals

Implementation Logic¶

Compute shell totals:

s_total = projections[..., S_IDX]
p_total = jnp.sum(projections[..., P_ORBITAL_SLICE], axis=-1)
d_total = jnp.sum(projections[..., D_ORBITAL_SLICE], axis=-1)

The fixed slices apply the public VASP orbital ordering.

Stack the shell axis:
```
reduced = jnp.stack([s_total, p_total, d_total], axis=-1)
```
The new trailing axis stores the s, p, and d totals in that order.

param projections:: Full 9-channel orbital projections.
type projections:: Float[Array, 'K B A 9']
returns:: reduced – Reduced projections: [s_total, p_total, d_total].
rtype:: Float[Array, 'K B A 3']

Notes

The VASP orbital ordering assumed here is: [s, py, pz, px, dxy, dyz, dz2, dxz, dx2-y2] (indices 0-8). This matches the standard PROCAR output when LORBIT=11 or LORBIT=12.

diffpes.inout.select_atoms(orb: OrbitalProjection | SpinOrbitalProjection, atom_indices: list[int]) → OrbitalProjection | SpinOrbitalProjection[source]¶

Extract orbital projections for a subset of atoms.

The function creates a projection object that contains only the requested atoms. Use it to isolate contributions from specified sites. Examples include surface atoms and atoms of one element.

See:: TestSelectAtoms

Implementation Logic¶

Build the atom index:
```
idx = jnp.asarray(atom_indices, dtype=jnp.int32)
```
The JAX index keeps the selection compatible with traced array access.

Select each present array leaf:

proj_sub = orb.projections[:, :, idx, :]
spin_sub = orb.spin[:, :, idx, :]
oam_sub = orb.oam[:, :, idx, :]

Each selection uses the same atom axis and preserves the other axes.

Preserve the carrier type:
```
result = SpinOrbitalProjection(...)
result = OrbitalProjection(...)
```
The branch returns the same projection carrier variant as the input.

param orb:: Full orbital projections with shape (K, B, A, 9).
type orb:: Union[OrbitalProjection, SpinOrbitalProjection]
param atom_indices:: 0-based indices of atoms to select.
type atom_indices:: list[int]
returns:: result – Projections restricted to the specified atoms. Shape (K, B, len(atom_indices), 9). Preserves the input type.
rtype:: Union[OrbitalProjection, SpinOrbitalProjection]

Notes

The returned object shares no memory with the original because JAX advanced indexing always produces a copy. The pure function works inside code that jax.jit compiles.

diffpes.inout.read_kpoints(filename: str = 'KPOINTS') → KPathInfo[source]¶

Parse a VASP KPOINTS file.

The function reads a VASP KPOINTS file that specifies Brillouin-zone sampling. It supports three standard modes: Line-mode (path segments), Automatic (Monkhorst-Pack/Gamma grids), and Explicit (listed k-points with optional weights).

See:: TestReadKpoints

Implementation Logic¶

Read the mode selector:
```
mode_token: str = lines[2].strip()
mode_lower: str = mode_token.lower()
```
This determines the static parser branch from the third KPOINTS record.
Dispatch the selected layout:
```
if "line" in mode_lower:
    mode = "Line-mode"
```
The other branches parse automatic grids or explicit weighted points.
Return normalized path metadata:
```
return kpath
```
Each layout binds its fields through one factory.

param filename:: Path to KPOINTS file. Default is "KPOINTS".
type filename:: str, default: "KPOINTS"
returns:: kpath – K-point metadata including labels/indices and mode-specific parsed fields.
rtype:: KPathInfo

Notes

In Line-mode, line 2 is points-per-segment. num_kpoints in the returned object is the total count segments * points_per_segment, preserving existing plotting behavior.

diffpes.inout.apply_kpath_ticks(ax: Axes, kpath: KPathInfo, draw_symmetry_lines: bool = True, line_color: str = 'white', line_width: float = 0.5, line_alpha: float = 0.35) → Axes[source]¶

Apply symmetry-point ticks/labels from KPathInfo to an axis.

The function adds k-path symmetry labels, such as G, M, and K, to an axis using KPathInfo. Optionally draws vertical guide lines at interior symmetry points to visually separate path segments.

See:: TestApplyKpathTicks

Implementation Logic¶

Normalize label positions and text:

indices: list[int] = np.asarray(
    kpath.label_indices, dtype=np.int32
).tolist()
labels: list[str] = list(kpath.labels)

Host-side lists match the Matplotlib tick API and preserve label order.

Apply the shared label count:

n_labels: int = min(len(indices), len(labels))
ticks: list[float] = [float(idx) for idx in indices[:n_labels]]

Truncation tolerates incomplete legacy metadata without shifting labels.

Return the annotated axis:
```
return ax
```
The caller receives the same axis after optional guide-line drawing.

param ax:: Target axis.
type ax:: Axes
param kpath:: K-path metadata containing symmetry labels and their indices.
type kpath:: KPathInfo
param draw_symmetry_lines:: If True, draw vertical guide lines at interior symmetry points.
type draw_symmetry_lines:: bool, default: True
param line_color:: Color of symmetry guide lines.
type line_color:: str, default: "white"
param line_width:: Width of symmetry guide lines.
type line_width:: float, default: 0.5
param line_alpha:: Alpha of symmetry guide lines.
type line_alpha:: float, default: 0.35
returns:: ax – The same axis, modified in place.
rtype:: Axes

Notes

This function mutates ax and returns it for convenient chaining.

diffpes.inout.list_band_scatter_presets() → tuple[str, ...][source]¶

Return supported preset names for projected band scatter plots.

Returns the canonical immutable names accepted by the projected-band plotting functions.

See:: TestListBandScatterPresets
Returns:: presets – Available preset names accepted by plot_band_scatter_preset().
Return type:: tuple[str, ...]

Notes

The function returns the tuple without allocating or reordering it.

diffpes.inout.plot_arpes_spectrum(spectrum: ArpesSpectrum, ax: Axes | None = None, cmap: str = 'gray', colorbar: bool = True, clim: tuple[float, float] | None = None, interpolation: str = 'nearest', aspect: Literal['equal', 'auto'] = 'auto', xlabel: str = 'k-point index', ylabel: str = 'Energy (eV)', title: str = 'Simulated ARPES Spectrum') → tuple[Figure | SubFigure, Axes, AxesImage][source]¶

Plot an ARPES intensity map from an ArpesSpectrum PyTree.

The function renders a two-dimensional ARPES map with matplotlib.axes.Axes.imshow. The vertical axis contains energy, and the horizontal axis contains the k-point index. The function accepts an existing axis. It creates a figure and axis when the caller supplies none.

See:: TestPlotArpesSpectrum

Implementation Logic¶

Normalize plotting arrays:
```
intensity, energy_axis = _prepare_plot_arrays(spectrum)
```
This validates the image rank and energy-axis length before plotting.

Draw the transposed intensity image:

image: AxesImage = ax.imshow(
    intensity.T,
    origin="lower",
    aspect=aspect,
    extent=extent,
    cmap=cmap,
)

The transpose places energy vertically and k-point index horizontally.

Return the Matplotlib objects:
```
return plot_result
```
This binding keeps axis reuse and colorbar state explicit.

param spectrum:

Spectrum containing intensity of shape (K, E) and energy_axis of shape (E,).

type spectrum:

ArpesSpectrum

param ax:

Existing axis for the image. If None, the function creates a figure and axis.

type ax:

Optional[Axes], default: None

param cmap:

Matplotlib colormap name. Default is "gray".

type cmap:

str, default: "gray"

param colorbar:

If True, add a colorbar labeled "Intensity (a.u.)".

type colorbar:

bool, default: True

param clim:

Optional (vmin, vmax) color limits.

type clim:

Optional[tuple[float, float]], default: None

param interpolation:

Image interpolation mode. Default is "nearest".

type interpolation:

str, default: "nearest"

param aspect:

Image aspect ratio passed to imshow. Default is "auto".

type aspect:

Literal['equal', 'auto'], default: "auto"

param xlabel:

x-axis label text. Default is "k-point index".

type xlabel:

str, default: "k-point index"

param ylabel:

y-axis label text. Default is "Energy (eV)".

type ylabel:

str, default: "Energy (eV)"

param title:

Axis title text. Default is "Simulated ARPES Spectrum".

type title:

str, default: "Simulated ARPES Spectrum"

rtype:

tuple[Union[Figure, SubFigure], Axes, AxesImage]

returns:

fig (Figure) – Matplotlib figure object.
ax (Axes) – Axis used for plotting.
image (AxesImage) – Image artist created by imshow.

diffpes.inout.plot_arpes_with_kpath(spectrum: ArpesSpectrum, kpath: KPathInfo, ax: Axes | None = None, cmap: str = 'gray', colorbar: bool = True, clim: tuple[float, float] | None = None, interpolation: str = 'nearest', aspect: Literal['equal', 'auto'] = 'auto', xlabel: str = 'Momentum (k)', ylabel: str = 'Energy (eV)', title: str = 'Simulated ARPES Spectrum', draw_symmetry_lines: bool = True) → tuple[Figure | SubFigure, Axes, AxesImage][source]¶

Plot ARPES spectrum and annotate k-axis using KPathInfo.

This wrapper combines plot_arpes_spectrum() and apply_kpath_ticks(). Use it to show symmetry labels on a line-mode ARPES image.

See:: TestPlotArpesWithKpath

Implementation Logic¶

Render the base spectrum:

fig, ax, image = plot_arpes_spectrum(
    spectrum=spectrum,
    ax=ax,
    cmap=cmap,
    aspect=aspect,
    clim=clim,
    colorbar=colorbar,
    xlabel=xlabel,
    ylabel=ylabel,
    title=title,
)

This centralizes image validation and styling in the base plotter.

Apply symmetry labels:

apply_kpath_ticks(
    ax=ax,
    kpath=kpath,
    draw_symmetry_lines=draw_symmetry_lines,
)

This adds line-mode metadata without recreating the image artist.

Return the composed plot:
```
return plot_result
```
The result contains the figure, axis, and image from the base call.

param spectrum:

Spectrum to plot.

type spectrum:

ArpesSpectrum

param kpath:

Symmetry-point metadata used for x-axis annotation.

type kpath:

KPathInfo

param ax:

Existing axis to draw on. If None, create a new one.

type ax:

Optional[Axes], default: None

param cmap:

Matplotlib colormap name. Default is "gray".

type cmap:

str, default: "gray"

param colorbar:

If True, add a colorbar.

type colorbar:

bool, default: True

param clim:

Optional (vmin, vmax) color limits.

type clim:

Optional[tuple[float, float]], default: None

param interpolation:

Image interpolation mode for imshow.

type interpolation:

str, default: "nearest"

param aspect:

Image aspect ratio for imshow.

type aspect:

Literal['equal', 'auto'], default: "auto"

param xlabel:

x-axis label. Default is "Momentum (k)".

type xlabel:

str, default: "Momentum (k)"

param ylabel:

y-axis label. Default is "Energy (eV)".

type ylabel:

str, default: "Energy (eV)"

param title:

Plot title.

type title:

str, default: "Simulated ARPES Spectrum"

param draw_symmetry_lines:

If True, draw vertical guide lines at interior symmetry points.

type draw_symmetry_lines:

bool, default: True

rtype:

tuple[Union[Figure, SubFigure], Axes, AxesImage]

returns:

fig (Figure) – Matplotlib figure object.
ax (Axes) – Axis used for plotting.
image (AxesImage) – Image artist created by imshow.

Implementation Logic¶

Resolve the preset weights:

weights, signed = _weights_from_preset(
    orb_proj, preset, atom_indices
)

This reduces the requested orbital data to one value per band.

Scale the marker areas:

marker_sizes: Float[NDArray, "K B"] = np.maximum(
    np.abs(weights) * size_scale,
    min_size,
)

The magnitude controls area while the lower bound keeps points visible.

Return the plot objects:
```
return plot_result
```
The caller receives the figure, axis, and scatter artist.

param bands:

Band structure with eigenvalues shape (K, B).

type bands:

BandStructure

param orb_proj:

Projection object containing orbital weights and optional spin/OAM.

type orb_proj:

Union[OrbitalProjection, SpinOrbitalProjection]

param preset:

Preset key from list_band_scatter_presets().

type preset:

str, default: "p"

param atom_indices:

Optional 0-based atom indices used before reduction.

type atom_indices:

Optional[list[int]], default: None

param ax:

Existing axis for the plot. If None, the function creates a figure and axis.

type ax:

Optional[Axes], default: None

param shift_fermi:

If True, plot E - E_F on y-axis.

type shift_fermi:

bool, default: True

param size_scale:

Linear scale factor for marker sizes.

type size_scale:

float, default: 250.0

param min_size:

Minimum marker size in points^2.

type min_size:

float, default: 0.5

param alpha:

Marker alpha.

type alpha:

float, default: 0.75

param color:

Solid marker color for non-signed presets.

type color:

str, default: "tab:blue"

param cmap:

Colormap name used for signed presets.

type cmap:

str, default: "coolwarm"

param colorbar:

If True and the preset has signs, draw a colorbar.

type colorbar:

bool, default: False

param xlabel:

X-axis label.

type xlabel:

str, default: "Momentum (k)"

param ylabel:

Y-axis label.

type ylabel:

str, default: "Energy (eV)"

param title:

Plot title.

type title:

str, default: "Projected Band Scatter"

rtype:

tuple[Union[Figure, SubFigure], Axes, PathCollection]

returns:

fig (Figure or SubFigure) – Parent figure.
ax (Axes) – Axis used for plotting.
scatter (PathCollection) – Scatter artist returned by Matplotlib.

raises ValueError:

If the preset weights and band eigenvalues have different shapes.

diffpes.inout.plot_band_scatter_with_kpath(bands: BandStructure, orb_proj: OrbitalProjection | SpinOrbitalProjection, kpath: KPathInfo, preset: str = 'p', atom_indices: list[int] | None = None, ax: Axes | None = None, shift_fermi: bool = True, size_scale: float = 250.0, min_size: float = 0.5, alpha: float = 0.75, color: str = 'tab:blue', cmap: str = 'coolwarm', colorbar: bool = False, xlabel: str = 'Momentum (k)', ylabel: str = 'Energy (eV)', title: str = 'Projected Band Scatter', draw_symmetry_lines: bool = True) → tuple[Figure | SubFigure, Axes, PathCollection][source]¶

Plot projected band scatter and annotate x-axis with k-path labels.

Combines a projected-band scatter plot with the symmetry labels from one KPathInfo carrier.

See:: TestPlotBandScatterWithKpath

Implementation Logic¶

Create the projected-band scatter plot:

fig, ax, scatter = plot_band_scatter_preset(
    bands=bands,
    orb_proj=orb_proj,
    preset=preset,
    atom_indices=atom_indices,
    ax=ax,
    shift_fermi=shift_fermi,
    size_scale=size_scale,
    min_size=min_size,
    alpha=alpha,
    color=color,
    cmap=cmap,
    colorbar=colorbar,
    xlabel=xlabel,
    ylabel=ylabel,
    title=title,
)

The base plotter owns the marker construction and axis styling.

Apply the k-path labels:

apply_kpath_ticks(
    ax=ax,
    kpath=kpath,
    draw_symmetry_lines=draw_symmetry_lines,
    line_color="black",
    line_alpha=0.35,
)

This adds symmetry metadata to the existing axis.

Return the plot objects:
```
return plot_result
```
The caller receives the figure, axis, and scatter artist.

param bands:

Band structure with eigenvalues shape (K, B).

type bands:

BandStructure

param orb_proj:

Projection object containing orbital weights and optional spin/OAM.

type orb_proj:

Union[OrbitalProjection, SpinOrbitalProjection]

param kpath:

Symmetry labels and indices for the plotted k-point path.

type kpath:

KPathInfo

param preset:

Preset key from list_band_scatter_presets().

type preset:

str, default: "p"

param atom_indices:

Optional zero-based atom indices used before reduction.

type atom_indices:

Optional[list[int]], default: None

param ax:

Existing axis to draw on, or None to create one.

type ax:

Optional[Axes], default: None

param shift_fermi:

Whether to plot energies relative to the Fermi level.

type shift_fermi:

bool, default: True

param size_scale:

Linear scale factor for marker sizes.

type size_scale:

float, default: 250.0

param min_size:

Minimum marker size in points squared.

type min_size:

float, default: 0.5

param alpha:

Marker opacity.

type alpha:

float, default: 0.75

param color:

Solid marker color for unsigned presets.

type color:

str, default: "tab:blue"

param cmap:

Colormap used for signed presets.

type cmap:

str, default: "coolwarm"

param colorbar:

Whether to draw a colorbar for signed presets.

type colorbar:

bool, default: False

param xlabel:

Horizontal axis label.

type xlabel:

str, default: "Momentum (k)"

param ylabel:

Vertical energy-axis label in eV.

type ylabel:

str, default: "Energy (eV)"

param title:

Plot title.

type title:

str, default: "Projected Band Scatter"

param draw_symmetry_lines:

Whether to draw vertical lines at symmetry points.

type draw_symmetry_lines:

bool, default: True

rtype:

tuple[Union[Figure, SubFigure], Axes, PathCollection]

returns:

fig (Figure or SubFigure) – Parent figure.
ax (Axes) – Axis used for plotting.
scatter (PathCollection) – Scatter artist returned by Matplotlib.

diffpes.inout.read_poscar(filename: str = 'POSCAR') → CrystalGeometry[source]¶

Parse a VASP POSCAR/CONTCAR file.

The function reads a VASP POSCAR or CONTCAR file. The file defines lattice vectors, element symbols, atom counts, and atomic coordinates. The coordinates can use direct or Cartesian form. The function returns a CrystalGeometry PyTree with fractional coordinates. It applies the universal scaling factor to the lattice. A negative scalar specifies the target cell volume, as defined by the VASP format.

The POSCAR file format used by VASP has the following structure:

Line 1: Comment / system name (discarded).
Line 2: Universal scaling factor applied to all lattice vectors (and to Cartesian coordinates, if applicable).
Lines 3-5: Three rows of three floats defining the lattice vectors a, b, c in Angstroms (before scaling).
Line 6: Either element symbols (VASP >= 5) or atom counts (VASP 4). If non-numeric, it is the species line.
Line 6 or 7: Atom counts per species (list of integers).
Next line: Optional "Selective dynamics" flag. If present, the following line is the coordinate-type specifier.
Coordinate-type line: "Direct" / "Fractional" for fractional coordinates, or "Cartesian" / "Cart" for Cartesian.
Coordinate lines: natoms lines, each with at least three floats. The parser ignores additional selective-dynamics columns.

See:: TestReadPoscar

Implementation Logic¶

Read and scale the lattice:
```
lattice = lattice * scaling_factor
```
A positive value is the direct scale. For a negative value, derive the positive scale that gives a cell volume of abs(raw_scale).
Convert Cartesian coordinates when required:
```
coords = np.linalg.solve(lattice.T, coords.T).T
```
Solving against the scaled lattice produces fractional coordinates.
Return validated crystal geometry:
```
return geometry
```
The result includes the reciprocal lattice and per-atom species.

param filename:: Path to POSCAR file. Default is "POSCAR".
type filename:: str, default: "POSCAR"
returns:: geometry – Crystal geometry with lattice, reciprocal lattice, fractional positions, and per-atom species.
rtype:: CrystalGeometry
raises ValueError:: If a negative scale accompanies a zero-volume raw lattice.

Notes

The function always returns fractional coordinates. For Cartesian input, np.linalg.solve computes frac = lattice^{-T} @ cart^T with numerical stability. The parser detects and ignores optional selective-dynamics flags. VASP-4 files can omit the element-symbol line. In this case, the function returns an empty species tuple. Supply the species from an external source such as POTCAR.

diffpes.inout.read_procar(filename: str = 'PROCAR', return_mode: Literal['legacy', 'full'] = 'legacy') → OrbitalProjection | SpinOrbitalProjection[source]¶

Parse a VASP PROCAR file.

The function reads a VASP PROCAR file that contains the orbital-resolved projections of Kohn-Sham wave functions onto site-centred spherical harmonics. Supports three layouts:

Non-spin (ISPIN=1, no SOC): single block of k-points.
Spin-polarized (ISPIN=2): two consecutive blocks of k-points (one per spin channel).
SOC (LSORBIT=.TRUE.): four consecutive blocks per k-point (total, Sx, Sy, Sz projections).

The PROCAR file written by VASP (when LORBIT=11 or 12) contains site- and orbital-resolved projections of each Kohn-Sham eigenstate. The file contains one or more blocks with the following data:

A header line with # of k-points, # of bands, # of ions.
For each k-point: one coordinate line and one record for each band. Each band record contains the energy and an orbital header. It also contains one projection line for each atom. The projection columns are the ion index, nine orbitals, and tot. A tot line and a blank line follow each record.

The number of blocks determines the spin layout:

1 block: non-spin-polarized (ISPIN=1).
2 blocks: spin-polarized (ISPIN=2), block 0 = spin-up, block 1 = spin-down.
4 blocks: spin-orbit coupling (SOC), blocks = total, Sx, Sy, Sz.

See:: TestReadProcar

Implementation Logic¶

Parse the file blocks:
```
content = fid.read()
blocks = _parse_procar_blocks(content)
```
Each block carries one orbital projection table and its dimensions.
Identify the spin layout:
```
is_spin_polarized = nblocks == ISPIN2_BLOCKS
is_soc = nblocks == SOC_BLOCKS
```
The number of blocks distinguishes non-spin, ISPIN=2, and SOC data.

Build the projection and spin arrays:

avg = (proj_up + proj_down) / 2.0
sx_sum = np.sum(proj_sx, axis=-1)
sy_sum = np.sum(proj_sy, axis=-1)
sz_sum = np.sum(proj_sz, axis=-1)

Full mode stores signed spin components as separate non-negative pairs.

Construct the matching carrier:

projection_result = make_orbital_projection(projections=proj_arr)
projection_result = make_spin_orbital_projection(
    projections=proj_arr, spin=spin_arr
)

Legacy and non-spin data use the orbital carrier. Spin data uses the mandatory-spin carrier.

param filename:: Path to PROCAR file. Default is "PROCAR".
type filename:: str, default: "PROCAR"
param return_mode:: "legacy" (default) returns an OrbitalProjection from the first spin block only (backward-compatible). "full" returns a SpinOrbitalProjection (with mandatory spin field) for ISPIN=2 and SOC data, or an OrbitalProjection for non-spin data.
type return_mode:: Literal['legacy', 'full'], default: "legacy"
returns:: projection_result – OrbitalProjection for legacy mode or non-spin data. SpinOrbitalProjection for full mode with spin data.
rtype:: Union[OrbitalProjection, SpinOrbitalProjection]
raises ValueError:: If the parser finds no valid PROCAR blocks in the file.

Notes

The 9 orbital channels follow the VASP convention: [s, py, pz, px, dxy, dyz, dz2, dxz, dx2-y2]. The parser does not store the VASP tot column. It retains only the individual orbital columns. In the ISPIN=2 full mode, the parser uses (up + down) / 2 as the orbital weight. Downstream consumers expect one projection array instead of separate spin channels. The parser encodes the spin texture as six nonnegative channels. These channels follow the ARPES simulation convention [Sx+, Sx-, Sy+, Sy-, Sz+, Sz-].