diffpes 2026.3.1__py3-none-any.whl

diffpes/inout/poscar.py

@@ -0,0 +1,174 @@
+"""VASP POSCAR file parser.
+
+Extended Summary
+----------------
+Reads VASP POSCAR/CONTCAR crystal structure files and returns
+a :class:`~diffpes.types.CrystalGeometry` PyTree containing lattice
+vectors, atomic coordinates, element symbols, and atom counts.
+
+Routine Listings
+----------------
+:func:`read_poscar`
+    Parse a VASP POSCAR file into a CrystalGeometry.
+
+Notes
+-----
+Handles both direct (fractional) and Cartesian coordinate formats,
+optional selective dynamics, and automatic reciprocal lattice
+computation.
+"""
+
+from pathlib import Path
+
+import jax.numpy as jnp
+import numpy as np
+
+from diffpes.types import CrystalGeometry, make_crystal_geometry
+
+
+def read_poscar(
+    filename: str = "POSCAR",
+) -> CrystalGeometry:
+    """Parse a VASP POSCAR/CONTCAR file.
+
+    Reads a VASP POSCAR (or CONTCAR) file that defines the crystal
+    structure: lattice vectors, element symbols, atom counts, and
+    atomic coordinates in either direct (fractional) or Cartesian
+    form. The function returns a :class:`~diffpes.types.CrystalGeometry`
+    PyTree with coordinates always stored in fractional form and the
+    lattice pre-scaled by the universal scaling factor.
+
+    Extended Summary
+    ----------------
+    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. Additional columns (selective-dynamics T/F flags) are
+      silently ignored.
+
+    Implementation Logic
+    --------------------
+    1. **Read comment line** (line 1) -- consumed and discarded.
+
+    2. **Read scaling factor** (line 2) -- a single float that
+       multiplies all lattice vectors.
+
+    3. **Read lattice vectors** (lines 3-5) -- three rows of three
+       floats each, stored as a (3, 3) NumPy array and immediately
+       multiplied by the scaling factor.
+
+    4. **Detect optional species line** (line 6) -- if the line
+       contains no digits it is the VASP-5 style element-symbol line
+       (e.g. ``"Si O"``). The symbols are stored and the next line is
+       read. If the line contains digits, it is the atom-count line
+       (VASP-4 style) and ``symbols`` remains empty.
+
+    5. **Read atom counts** -- parse the current line as a list of
+       integers giving the number of atoms per species.
+
+    6. **Detect optional Selective Dynamics line** -- if the next line
+       starts with ``"s"`` or ``"S"``, selective-dynamics flags are
+       present; consume the line and read the coordinate-type line
+       that follows.
+
+    7. **Determine coordinate type** -- if the coordinate-type line
+       starts with ``"c"``, ``"C"``, ``"k"``, or ``"K"`` the
+       coordinates are Cartesian; otherwise they are direct
+       (fractional).
+
+    8. **Read coordinates** -- parse ``natoms`` lines, each with at
+       least three floats. Only the first three tokens are used
+       (selective-dynamics flags on positions 4-6 are ignored).
+
+    9. **Convert Cartesian to fractional** (if needed) -- scale the
+       Cartesian coordinates by the scaling factor, then solve
+       ``lattice^T @ frac^T = cart^T`` via ``np.linalg.solve`` to
+       obtain fractional coordinates.
+
+    10. **Construct PyTree** -- call ``make_crystal_geometry`` with the
+        scaled lattice, fractional coordinates, element symbols, and
+        atom counts. The reciprocal lattice is computed automatically
+        inside ``make_crystal_geometry`` as ``2 * pi * inv(lattice)^T``.
+
+    Parameters
+    ----------
+    filename : str, optional
+        Path to POSCAR file. Default is ``"POSCAR"``.
+
+    Returns
+    -------
+    geometry : CrystalGeometry
+        Crystal geometry with lattice, reciprocal lattice, fractional
+        coordinates, element symbols, and atom counts.
+
+    Raises
+    ------
+    IndexError
+        If the coordinate-type line is empty (malformed file).
+
+    Notes
+    -----
+    Coordinates are always returned in **fractional** (direct) form,
+    regardless of the input format. If the input is Cartesian, the
+    conversion ``frac = lattice^{-T} @ cart^T`` is performed via
+    ``np.linalg.solve`` for numerical stability. The optional
+    selective-dynamics flags (``T T F`` appended to coordinate lines)
+    are detected and skipped but not stored. VASP-4 style files that
+    lack the element-symbol line will produce an empty ``symbols``
+    tuple; the caller should supply symbols from an external source
+    (e.g. POTCAR) in that case.
+    """
+    path: Path = Path(filename)
+    with path.open("r") as fid:
+        _comment: str = fid.readline().strip()
+        scale: float = float(fid.readline().strip())
+        lattice: np.ndarray = np.zeros((3, 3), dtype=np.float64)
+        for i in range(3):
+            vals: list[float] = [float(x) for x in fid.readline().split()]
+            lattice[i, :] = vals
+        lattice = lattice * scale
+        line: str = fid.readline().strip()
+        symbols: tuple[str, ...] = ()
+        if not any(c.isdigit() for c in line):
+            symbols = tuple(line.split())
+            line = fid.readline().strip()
+        atom_counts: list[int] = [int(x) for x in line.split()]
+        natoms: int = sum(atom_counts)
+        line = fid.readline().strip()
+        selective: bool = False
+        if line[0].lower() == "s":
+            selective = True  # noqa: F841
+            line = fid.readline().strip()
+        cartesian: bool = line[0].lower() in ("c", "k")
+        coords: np.ndarray = np.zeros((natoms, 3), dtype=np.float64)
+        for i in range(natoms):
+            vals = [float(x) for x in fid.readline().split()[:3]]
+            coords[i, :] = vals
+        if cartesian:
+            coords = coords * scale
+            coords = np.linalg.solve(lattice.T, coords.T).T
+    geometry: CrystalGeometry = make_crystal_geometry(
+        lattice=jnp.asarray(lattice),
+        coords=jnp.asarray(coords),
+        symbols=symbols,
+        atom_counts=atom_counts,
+    )
+    return geometry
+
+
+__all__: list[str] = [
+    "read_poscar",
+]

diffpes/inout/procar.py

@@ -0,0 +1,331 @@
+"""VASP PROCAR file parser.
+
+Extended Summary
+----------------
+Reads VASP PROCAR files containing orbital-resolved band
+projections and returns an
+:class:`~diffpes.types.OrbitalProjection` PyTree. Supports
+non-spin, spin-polarized (ISPIN=2), and SOC layouts.
+
+Routine Listings
+----------------
+:func:`read_procar`
+    Parse a VASP PROCAR file into an OrbitalProjection.
+
+Notes
+-----
+Orbital ordering follows VASP convention:
+``[s, py, pz, px, dxy, dyz, dz2, dxz, dx2-y2]``.
+"""
+
+import re
+from pathlib import Path
+
+import jax.numpy as jnp
+import numpy as np
+from beartype.typing import Literal, Optional, Union
+from jaxtyping import Array, Float
+
+from diffpes.types import (
+    OrbitalProjection,
+    SpinOrbitalProjection,
+    make_orbital_projection,
+    make_spin_orbital_projection,
+)
+
+_NORBS: int = 9
+_NSPIN_COMPONENTS: int = 6
+_ISPIN2_BLOCKS: int = 2
+_SOC_BLOCKS: int = 4
+
+
+def read_procar(
+    filename: str = "PROCAR",
+    return_mode: Literal["legacy", "full"] = "legacy",
+) -> Union[OrbitalProjection, SpinOrbitalProjection]:
+    r"""Parse a VASP PROCAR file.
+
+    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).
+
+    Extended Summary
+    ----------------
+    The PROCAR file written by VASP (when ``LORBIT=11`` or ``12``)
+    contains site- and orbital-resolved projections of each Kohn-Sham
+    eigenstate. The file is organised into one or more *blocks*, each
+    containing:
+
+    * A header line with ``# of k-points``, ``# of bands``,
+      ``# of ions``.
+    * For each k-point: a coordinate line, then for each band: a
+      band-energy line, an orbital-name header, one projection line
+      per atom (columns: ion index, s, py, pz, px, dxy, dyz, dz2,
+      dxz, dx2-y2, tot), a ``tot`` summation line, and a blank line.
+
+    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.
+
+    Implementation Logic
+    --------------------
+    1. **Read entire file** into a single string and delegate to
+       :func:`_parse_procar_blocks` to extract structured block data.
+
+    2. **Determine spin layout** from the number of blocks.
+
+    3. **Legacy mode or non-spin**: return an ``OrbitalProjection``
+       wrapping only the first block's projection array (shape
+       ``(K, B, A, 9)``).
+
+    4. **Spin-polarized (ISPIN=2), full mode**:
+
+       a. Average spin-up and spin-down projections to get a
+          spin-averaged orbital weight per atom.
+       b. Construct a ``(K, B, A, 6)`` spin texture array where
+          channels 4-5 encode ``Sz+`` and ``Sz-`` (the positive and
+          negative parts of the spin-up minus spin-down difference
+          summed over orbitals).
+       c. Return a ``SpinOrbitalProjection``.
+
+    5. **SOC (4 blocks), full mode**:
+
+       a. Use the total-projection block (block 0) as the orbital
+          weights.
+       b. Construct a ``(K, B, A, 6)`` spin texture array where
+          channels 0-1 encode ``Sx+`` / ``Sx-``, 2-3 encode
+          ``Sy+`` / ``Sy-``, and 4-5 encode ``Sz+`` / ``Sz-``.
+          Each component is the positive/negative part of the
+          orbital-summed projection from the corresponding
+          spin block.
+       c. Return a ``SpinOrbitalProjection``.
+
+    Parameters
+    ----------
+    filename : str, optional
+        Path to PROCAR file. Default is ``"PROCAR"``.
+    return_mode : {"legacy", "full"}, optional
+        ``"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.
+
+    Returns
+    -------
+    orb_proj : OrbitalProjection or SpinOrbitalProjection
+        ``OrbitalProjection`` for legacy mode or non-spin data.
+        ``SpinOrbitalProjection`` for full mode with spin data.
+
+    Raises
+    ------
+    ValueError
+        If no valid PROCAR blocks are found in the file.
+
+    Notes
+    -----
+    The 9 orbital channels follow the VASP convention:
+    ``[s, py, pz, px, dxy, dyz, dz2, dxz, dx2-y2]``.
+    The ``tot`` column printed by VASP is **not** stored; only the
+    individual orbital columns are kept. For ISPIN=2 in full mode, the
+    spin-averaged projection ``(up + down) / 2`` is used as the
+    orbital weight because many downstream consumers expect a single
+    projection array rather than separate spin channels. The spin
+    texture is encoded as 6 non-negative channels
+    ``[Sx+, Sx-, Sy+, Sy-, Sz+, Sz-]`` following the convention used
+    by the ARPES simulation pipeline.
+    """
+    path: Path = Path(filename)
+    with path.open("r") as fid:
+        content: str = fid.read()
+
+    blocks: list[dict] = _parse_procar_blocks(content)
+
+    if not blocks:
+        msg = "No valid PROCAR blocks found."
+        raise ValueError(msg)
+
+    nblocks: int = len(blocks)
+    nkpts: int = blocks[0]["nkpts"]
+    nbands: int = blocks[0]["nbands"]
+    natoms: int = blocks[0]["natoms"]
+
+    is_spin_polarized: bool = nblocks == _ISPIN2_BLOCKS
+    is_soc: bool = nblocks == _SOC_BLOCKS
+
+    if return_mode == "legacy" or (not is_spin_polarized and not is_soc):
+        proj_arr: Float[Array, " K B A 9"] = jnp.asarray(
+            blocks[0]["projections"], dtype=jnp.float64
+        )
+        return make_orbital_projection(projections=proj_arr)
+
+    if is_spin_polarized:
+        proj_up: np.ndarray = blocks[0]["projections"]
+        proj_down: np.ndarray = blocks[1]["projections"]
+        avg: np.ndarray = (proj_up + proj_down) / 2.0
+        proj_arr = jnp.asarray(avg, dtype=jnp.float64)
+        spin_data: np.ndarray = np.zeros(
+            (nkpts, nbands, natoms, _NSPIN_COMPONENTS), dtype=np.float64
+        )
+        sz_diff: np.ndarray = np.sum(proj_up - proj_down, axis=-1)
+        spin_data[:, :, :, 4] = np.maximum(sz_diff, 0.0)
+        spin_data[:, :, :, 5] = np.maximum(-sz_diff, 0.0)
+        spin_arr: Float[Array, " K B A 6"] = jnp.asarray(
+            spin_data, dtype=jnp.float64
+        )
+        return make_spin_orbital_projection(
+            projections=proj_arr, spin=spin_arr
+        )
+
+    # SOC: 4 blocks = total, Sx, Sy, Sz
+    proj_total: np.ndarray = blocks[0]["projections"]
+    proj_sx: np.ndarray = blocks[1]["projections"]
+    proj_sy: np.ndarray = blocks[2]["projections"]
+    proj_sz: np.ndarray = blocks[3]["projections"]
+    proj_arr = jnp.asarray(proj_total, dtype=jnp.float64)
+
+    spin_data = np.zeros(
+        (nkpts, nbands, natoms, _NSPIN_COMPONENTS), dtype=np.float64
+    )
+    sx_sum: np.ndarray = np.sum(proj_sx, axis=-1)
+    sy_sum: np.ndarray = np.sum(proj_sy, axis=-1)
+    sz_sum: np.ndarray = np.sum(proj_sz, axis=-1)
+    spin_data[:, :, :, 0] = np.maximum(sx_sum, 0.0)
+    spin_data[:, :, :, 1] = np.maximum(-sx_sum, 0.0)
+    spin_data[:, :, :, 2] = np.maximum(sy_sum, 0.0)
+    spin_data[:, :, :, 3] = np.maximum(-sy_sum, 0.0)
+    spin_data[:, :, :, 4] = np.maximum(sz_sum, 0.0)
+    spin_data[:, :, :, 5] = np.maximum(-sz_sum, 0.0)
+    spin_arr = jnp.asarray(spin_data, dtype=jnp.float64)
+    return make_spin_orbital_projection(projections=proj_arr, spin=spin_arr)
+
+
+def _parse_procar_blocks(
+    content: str,
+) -> list[dict]:
+    """Parse all PROCAR blocks from the full file content string.
+
+    Extended Summary
+    ----------------
+    A PROCAR file may contain 1, 2, or 4 consecutive blocks depending
+    on the spin configuration. Each block begins with a header line
+    matching ``"# of k-points: K  # of bands: B  # of ions: A"`` and
+    is followed by nested k-point / band / atom projection data.
+
+    Implementation Logic
+    --------------------
+    1. Split the content into lines and scan for lines containing the
+       substring ``"k-points"`` (the block header).
+    2. Extract ``(nkpts, nbands, natoms)`` from the header using a
+       regex that captures all integers on the line.
+    3. Allocate a ``(nkpts, nbands, natoms, 9)`` NumPy array for the
+       orbital projections.
+    4. For each k-point within the block:
+
+       a. Search forward for a line matching the pattern
+          ``k-point <index> : kx ky kz`` using a regex.
+       b. For each band within the k-point:
+
+          i.   Skip forward until a line containing ``"band"`` is
+               found (the band energy header).
+          ii.  Skip the orbital-name header line (``ion  s  py ...``).
+          iii. Read ``natoms`` lines, parsing columns 1 through 9
+               (skipping the ion index in column 0) as the orbital
+               projections.
+          iv.  Skip the ``tot`` summation line and the trailing blank
+               line.
+
+    5. Append a dict with keys ``'nkpts'``, ``'nbands'``,
+       ``'natoms'``, and ``'projections'`` for each block.
+    6. Return the list of block dicts.
+
+    Parameters
+    ----------
+    content : str
+        The entire PROCAR file content as a single string.
+
+    Returns
+    -------
+    blocks : list[dict]
+        List of parsed blocks. Each dict contains:
+
+        * ``'nkpts'`` (int): number of k-points.
+        * ``'nbands'`` (int): number of bands.
+        * ``'natoms'`` (int): number of atoms (ions).
+        * ``'projections'`` (np.ndarray): orbital projections with
+          shape ``(nkpts, nbands, natoms, 9)`` and dtype ``float64``.
+
+    Notes
+    -----
+    The parser uses 1-based k-point indices from the file to place
+    data into the 0-based NumPy array (``k_idx = parsed_index - 1``).
+    Band and atom lines are read positionally (sequential order) rather
+    than by parsed index. The ``tot`` column (column 10 in the PROCAR
+    line) is not stored.
+    """
+    blocks: list[dict] = []
+    lines: list[str] = content.splitlines()
+    i: int = 0
+
+    while i < len(lines):
+        if "k-points" not in lines[i]:
+            i += 1
+            continue
+        header: str = lines[i]
+        params: list[int] = [int(x) for x in re.findall(r"\d+", header)]
+        nkpts: int = params[0]
+        nbands: int = params[1]
+        natoms: int = params[2]
+        projections: np.ndarray = np.zeros(
+            (nkpts, nbands, natoms, _NORBS), dtype=np.float64
+        )
+        i += 1
+
+        k_re: str = (
+            r"k-point\s+(\d+)\s*:\s*" r"([-\d.]+)\s+([-\d.]+)\s+([-\d.]+)"
+        )
+        kpts_found: int = 0
+        while i < len(lines) and kpts_found < nkpts:
+            k_match: Optional[re.Match[str]] = re.search(k_re, lines[i])
+            if k_match is None:
+                i += 1
+                continue
+            k_idx: int = int(k_match.group(1)) - 1
+            i += 1
+            for b in range(nbands):
+                while i < len(lines) and "band" not in lines[i]:
+                    i += 1
+                i += 1  # skip band header
+                i += 1  # skip orbital-name header
+                for a in range(natoms):
+                    vals: list[float] = [float(x) for x in lines[i].split()]
+                    projections[k_idx, b, a, :] = vals[1 : _NORBS + 1]
+                    i += 1
+                i += 1  # skip tot line
+                i += 1  # skip blank line
+            kpts_found += 1
+
+        blocks.append(
+            {
+                "nkpts": nkpts,
+                "nbands": nbands,
+                "natoms": natoms,
+                "projections": projections,
+            }
+        )
+
+    return blocks
+
+
+__all__: list[str] = [
+    "read_procar",
+]

diffpes/maths/__init__.py

@@ -0,0 +1,68 @@
+r"""Angular matrix elements for dipole photoemission.
+
+Extended Summary
+----------------
+Provides Gaunt coefficients, real spherical harmonics, and full
+dipole matrix element assembly for the differentiable ARPES
+forward model.  The dipole matrix element for orbital
+:math:`(n, l, m)` combines radial integrals, Gaunt coefficients,
+and spherical harmonics of the photoelectron direction:
+
+.. math::
+
+    M(\mathbf{k}, n, l, m) = \sum_{l', m'} B_{n,l}^{l'}(|\mathbf{k}|)
+        \cdot G(l, m, l', m') \cdot Y_{l'}^{m'}(\hat{k})
+        \cdot \hat{e}_{q(m'-m)}
+
+Routine Listings
+----------------
+:data:`GAUNT_TABLE`
+    Precomputed dense array of real-valued Gaunt coefficients for
+    dipole transitions up to l_max = 4 (s through g orbitals).
+:data:`L_MAX`
+    Maximum angular momentum supported by the precomputed Gaunt table.
+:func:`build_gaunt_table`
+    Build a dipole Gaunt coefficient lookup table for a given l_max.
+:func:`gaunt_lookup`
+    Look up a single Gaunt coefficient from the precomputed table.
+:func:`real_spherical_harmonic`
+    Evaluate a single real spherical harmonic :math:`Y_l^m(\theta, \varphi)`.
+:func:`real_spherical_harmonics_all`
+    Evaluate all real spherical harmonics up to a given l_max.
+:func:`dipole_matrix_element_single`
+    Compute the complex dipole matrix element for one orbital (n, l, m).
+:func:`dipole_intensity_orbital`
+    Compute :math:`|M|^2` for a single orbital.
+:func:`dipole_intensities_all_orbitals`
+    Compute :math:`|M|^2` for every orbital in a Slater basis set.
+
+Notes
+-----
+All functions are JAX-compatible and support JIT compilation and
+automatic differentiation.  The Gaunt table is computed once at
+import time using pure Python and stored as a JAX array for O(1)
+lookup during traced computation.
+"""
+
+from .dipole import (
+    dipole_intensities_all_orbitals,
+    dipole_intensity_orbital,
+    dipole_matrix_element_single,
+)
+from .gaunt import GAUNT_TABLE, L_MAX, build_gaunt_table, gaunt_lookup
+from .spherical_harmonics import (
+    real_spherical_harmonic,
+    real_spherical_harmonics_all,
+)
+
+__all__: list[str] = [
+    "GAUNT_TABLE",
+    "L_MAX",
+    "build_gaunt_table",
+    "dipole_intensities_all_orbitals",
+    "dipole_intensity_orbital",
+    "dipole_matrix_element_single",
+    "gaunt_lookup",
+    "real_spherical_harmonic",
+    "real_spherical_harmonics_all",
+]