diffpes 2026.3.1__py3-none-any.whl

diffpes/types/geometry.py

@@ -0,0 +1,306 @@
+"""Crystal geometry data structure for VASP crystal structures.
+
+Extended Summary
+----------------
+Defines the :class:`CrystalGeometry` PyTree for representing
+crystal structures parsed from VASP POSCAR files. Includes real-space
+lattice vectors, reciprocal lattice, atomic coordinates, element
+symbols, and atom counts per species.
+
+Routine Listings
+----------------
+:class:`CrystalGeometry`
+    PyTree NamedTuple for crystal geometry data.
+:func:`make_crystal_geometry`
+    Factory function with validation and reciprocal lattice computation.
+
+Notes
+-----
+The ``symbols`` field is stored as auxiliary data in the PyTree
+since JAX cannot trace Python strings. All numeric fields are
+stored as JAX arrays for compatibility with JAX transformations.
+"""
+
+import jax.numpy as jnp
+from beartype import beartype
+from beartype.typing import NamedTuple, Tuple, Union
+from jax.tree_util import register_pytree_node_class
+from jaxtyping import Array, Float, Int, jaxtyped
+
+from .aliases import ScalarNumeric
+
+
+@register_pytree_node_class
+class CrystalGeometry(NamedTuple):
+    """PyTree for crystal geometry from VASP POSCAR.
+
+    Encapsulates the full crystal structure information parsed from a
+    VASP POSCAR file: real-space lattice vectors, the corresponding
+    reciprocal lattice, fractional atomic coordinates, per-species
+    element symbols, and per-species atom counts. Together these fields
+    fully describe the periodic crystal needed for ARPES simulation.
+
+    This class is registered as a JAX PyTree via
+    ``@register_pytree_node_class``. Numeric fields (lattice,
+    reciprocal_lattice, coords, atom_counts) are stored as children
+    visible to JAX tracing, while the ``symbols`` tuple of Python
+    strings is stored as auxiliary data because JAX cannot trace
+    Python strings.
+
+    Attributes
+    ----------
+    lattice : Float[Array, "3 3"]
+        Real-space lattice vectors as rows (angstroms).
+    reciprocal_lattice : Float[Array, "3 3"]
+        Reciprocal lattice vectors as rows (1/angstroms).
+    coords : Float[Array, "N 3"]
+        Fractional atomic coordinates.
+    symbols : tuple[str, ...]
+        Element symbols for each species.
+    atom_counts : Int[Array, " S"]
+        Number of atoms per species.
+
+    Notes
+    -----
+    Registered as a JAX PyTree with ``@register_pytree_node_class``.
+    The ``symbols`` field is a tuple of Python strings and therefore
+    must be stored as PyTree auxiliary data (not children). JAX treats
+    auxiliary data as compile-time constants: changing ``symbols``
+    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_lattice: Float[Array, "3 3"]
+    coords: Float[Array, "N 3"]
+    symbols: tuple[str, ...]
+    atom_counts: Int[Array, " S"]
+
+    def tree_flatten(
+        self,
+    ) -> Tuple[
+        Tuple[
+            Float[Array, "3 3"],
+            Float[Array, "3 3"],
+            Float[Array, "N 3"],
+            Int[Array, " S"],
+        ],
+        tuple[str, ...],
+    ]:
+        """Flatten into JAX children and auxiliary data.
+
+        Separates the PyTree into children (JAX-traced arrays) and
+        auxiliary data (static Python values). For CrystalGeometry,
+        the four numeric fields are children and the ``symbols`` tuple
+        of strings is auxiliary data.
+
+        Implementation Logic
+        --------------------
+        1. **Children** (JAX arrays, participate in autodiff):
+           ``(lattice, reciprocal_lattice, coords, atom_counts)``
+        2. **Auxiliary data** (static, not traced by JAX):
+           ``symbols`` -- a tuple of Python strings. JAX treats this
+           as a compile-time constant; any change triggers JIT
+           recompilation.
+
+        Returns
+        -------
+        children : tuple of Array
+            Tuple of ``(lattice, reciprocal_lattice, coords,
+            atom_counts)`` JAX arrays.
+        aux_data : tuple[str, ...]
+            Element symbol strings, stored outside JAX tracing.
+        """
+        return (
+            (
+                self.lattice,
+                self.reciprocal_lattice,
+                self.coords,
+                self.atom_counts,
+            ),
+            self.symbols,
+        )
+
+    @classmethod
+    def tree_unflatten(
+        cls,
+        aux_data: tuple[str, ...],
+        children: Tuple[
+            Float[Array, "3 3"],
+            Float[Array, "3 3"],
+            Float[Array, "N 3"],
+            Int[Array, " S"],
+        ],
+    ) -> "CrystalGeometry":
+        """Reconstruct a CrystalGeometry from flattened components.
+
+        Inverse of :meth:`tree_flatten`. JAX calls this method
+        automatically when unflattening a PyTree after a
+        transformation (e.g., inside ``jax.jit`` or ``jax.grad``).
+
+        Implementation Logic
+        --------------------
+        1. Unpack ``children`` into four JAX arrays:
+           ``(lattice, reciprocal_lattice, coords, atom_counts)``.
+        2. Receive ``aux_data`` as the ``symbols`` tuple of strings.
+        3. Pass all five fields to the constructor, re-interleaving
+           ``symbols`` from ``aux_data`` into its correct position
+           between ``coords`` and ``atom_counts``.
+
+        Parameters
+        ----------
+        aux_data : tuple[str, ...]
+            Element symbol strings recovered from auxiliary data.
+        children : tuple of Array
+            Tuple of ``(lattice, reciprocal_lattice, coords,
+            atom_counts)`` JAX arrays.
+
+        Returns
+        -------
+        geometry : CrystalGeometry
+            Reconstructed instance with identical data.
+        """
+        lattice: Float[Array, "3 3"]
+        reciprocal_lattice: Float[Array, "3 3"]
+        coords: Float[Array, "N 3"]
+        atom_counts: Int[Array, " S"]
+        lattice, reciprocal_lattice, coords, atom_counts = children
+        geometry: CrystalGeometry = cls(
+            lattice=lattice,
+            reciprocal_lattice=reciprocal_lattice,
+            coords=coords,
+            symbols=aux_data,
+            atom_counts=atom_counts,
+        )
+        return geometry
+
+
+def _compute_reciprocal_lattice(
+    lattice: Float[Array, "3 3"],
+) -> Float[Array, "3 3"]:
+    r"""Compute reciprocal lattice vectors from real-space lattice.
+
+    Derives the reciprocal lattice from the real-space lattice using
+    the standard cross-product formula. The reciprocal lattice vectors
+    ``b_i`` satisfy ``a_i . b_j = 2 pi delta_{ij}``.
+
+    Implementation Logic
+    --------------------
+    Given real-space lattice vectors ``a1, a2, a3`` (rows of
+    ``lattice``):
+
+    1. **Extract rows**: ``a1 = lattice[0]``, ``a2 = lattice[1]``,
+       ``a3 = lattice[2]``.
+    2. **Compute unit-cell volume**:
+       ``V = a1 . (a2 x a3)`` (scalar triple product).
+    3. **Compute reciprocal vectors** via the cross-product formula:
+
+       .. math::
+
+           b_1 = 2\\pi \\, (a_2 \\times a_3) / V
+
+           b_2 = 2\\pi \\, (a_3 \\times a_1) / V
+
+           b_3 = 2\\pi \\, (a_1 \\times a_2) / V
+
+    4. **Stack** ``[b1, b2, b3]`` into a (3, 3) array with
+       reciprocal vectors as rows.
+
+    Parameters
+    ----------
+    lattice : Float[Array, "3 3"]
+        Real-space lattice vectors as rows.
+
+    Returns
+    -------
+    reciprocal : Float[Array, "3 3"]
+        Reciprocal lattice vectors as rows (units of 1/angstrom,
+        with the 2 pi prefactor included).
+    """
+    a1: Float[Array, " 3"] = lattice[0]
+    a2: Float[Array, " 3"] = lattice[1]
+    a3: Float[Array, " 3"] = lattice[2]
+    volume: Float[Array, " "] = jnp.dot(a1, jnp.cross(a2, a3))
+    b1: Float[Array, " 3"] = 2.0 * jnp.pi * jnp.cross(a2, a3) / volume
+    b2: Float[Array, " 3"] = 2.0 * jnp.pi * jnp.cross(a3, a1) / volume
+    b3: Float[Array, " 3"] = 2.0 * jnp.pi * jnp.cross(a1, a2) / volume
+    reciprocal: Float[Array, "3 3"] = jnp.stack([b1, b2, b3])
+    return reciprocal
+
+
+@jaxtyped(typechecker=beartype)
+def make_crystal_geometry(
+    lattice: Union[Float[Array, "3 3"], "list[list[ScalarNumeric]]"],
+    coords: Float[Array, "N 3"],
+    symbols: tuple[str, ...],
+    atom_counts: Union[Int[Array, " S"], "list[int]"],
+) -> CrystalGeometry:
+    """Create a validated CrystalGeometry instance.
+
+    Validates and normalises the inputs, then automatically computes
+    the reciprocal lattice from the real-space lattice so the caller
+    does not need to supply it. Numeric arrays are cast to the
+    appropriate JAX dtypes (float64 for real-valued fields, int32
+    for atom counts).
+
+    Implementation Logic
+    --------------------
+    1. **Cast lattice** to ``jnp.float64`` via ``jnp.asarray``.
+       Accepts both JAX arrays and nested Python lists.
+    2. **Cast coords** to ``jnp.float64`` via ``jnp.asarray``.
+    3. **Cast atom_counts** to ``jnp.int32`` via ``jnp.asarray``.
+       Accepts both JAX integer arrays and Python ``list[int]``.
+    4. **Auto-compute reciprocal lattice** by calling
+       ``_compute_reciprocal_lattice(lattice_arr)``. This derives
+       ``b_i = 2 pi (a_j x a_k) / V`` from the validated
+       real-space lattice.
+    5. **Construct** the ``CrystalGeometry`` NamedTuple from all
+       five fields (including the computed reciprocal lattice) and
+       return it.
+
+    Parameters
+    ----------
+    lattice : Union[Float[Array, "3 3"], list]
+        Real-space lattice vectors as rows (angstroms).
+    coords : Float[Array, "N 3"]
+        Fractional atomic coordinates.
+    symbols : tuple[str, ...]
+        Element symbols for each species.
+    atom_counts : Union[Int[Array, " S"], list[int]]
+        Number of atoms per species.
+
+    Returns
+    -------
+    geometry : CrystalGeometry
+        Validated crystal geometry instance with the reciprocal
+        lattice pre-computed.
+
+    See Also
+    --------
+    CrystalGeometry : The PyTree class constructed by this factory.
+    _compute_reciprocal_lattice : Cross-product formula used to
+        derive the reciprocal lattice.
+    """
+    lattice_arr: Float[Array, "3 3"] = jnp.asarray(lattice, dtype=jnp.float64)
+    coords_arr: Float[Array, "N 3"] = jnp.asarray(coords, dtype=jnp.float64)
+    counts_arr: Int[Array, " S"] = jnp.asarray(atom_counts, dtype=jnp.int32)
+    reciprocal: Float[Array, "3 3"] = _compute_reciprocal_lattice(lattice_arr)
+    geometry: CrystalGeometry = CrystalGeometry(
+        lattice=lattice_arr,
+        reciprocal_lattice=reciprocal,
+        coords=coords_arr,
+        symbols=symbols,
+        atom_counts=counts_arr,
+    )
+    return geometry
+
+
+__all__: list[str] = [
+    "CrystalGeometry",
+    "make_crystal_geometry",
+]

diffpes/types/kpath.py

@@ -0,0 +1,365 @@
+"""K-point path information data structure.
+
+Extended Summary
+----------------
+Defines the :class:`KPathInfo` PyTree for storing Brillouin-zone
+metadata parsed from VASP KPOINTS files, including both plotting
+labels and mode-specific metadata (automatic grids, explicit weights,
+line-mode segments/endpoints).
+
+Routine Listings
+----------------
+:class:`KPathInfo`
+    PyTree for k-point path metadata.
+:func:`make_kpath_info`
+    Create a validated KPathInfo instance.
+"""
+
+import jax.numpy as jnp
+from beartype import beartype
+from beartype.typing import NamedTuple, Optional, Tuple, Union
+from jax.tree_util import register_pytree_node_class
+from jaxtyping import Array, Float, Int, jaxtyped
+
+
+@register_pytree_node_class
+class KPathInfo(NamedTuple):
+    """PyTree for k-point path metadata.
+
+    Extended Summary
+    ----------------
+    Stores Brillouin-zone path information parsed from VASP KPOINTS
+    files. Includes plotting fields (labels + label indices) and
+    mode-specific metadata needed for full parser completeness:
+    automatic-mode grid/shift, explicit-mode k-points/weights, and
+    line-mode segment endpoints.
+
+    This class is registered as a JAX PyTree via
+    ``@register_pytree_node_class``. Array-valued metadata is stored
+    as children. String metadata is stored as auxiliary data because
+    JAX cannot trace Python strings.
+
+    Attributes
+    ----------
+    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).
+    labels : tuple[str, ...]
+        Symmetry point labels (e.g., Gamma, M, K).
+    comment : str
+        Raw comment from KPOINTS line 1.
+    coordinate_mode : str
+        Coordinate/scheme line metadata:
+        line/explicit -> Reciprocal/Cartesian line,
+        automatic -> scheme line (e.g., Monkhorst-Pack).
+
+    Notes
+    -----
+    Registered as a JAX PyTree with ``@register_pytree_node_class``.
+    The ``mode`` string and ``labels`` tuple of strings are stored as
+    auxiliary data (not children). JAX treats auxiliary data as
+    compile-time constants: changing either 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: Optional[Float[Array, "K 3"]]
+    weights: Optional[Float[Array, " K"]]
+    grid: Optional[Int[Array, " 3"]]
+    shift: Optional[Float[Array, " 3"]]
+    mode: str
+    labels: tuple[str, ...]
+    comment: str
+    coordinate_mode: str
+
+    def tree_flatten(
+        self,
+    ) -> Tuple[
+        Tuple[
+            Int[Array, " "],
+            Int[Array, " L"],
+            Int[Array, " "],
+            Int[Array, " "],
+            Optional[Float[Array, "K 3"]],
+            Optional[Float[Array, " K"]],
+            Optional[Int[Array, " 3"]],
+            Optional[Float[Array, " 3"]],
+        ],
+        Tuple[str, tuple[str, ...], str, str],
+    ]:
+        """Flatten into JAX children and auxiliary data.
+
+        Separates the PyTree into children (JAX-traced arrays) and
+        auxiliary data (static Python values). For KPathInfo, all
+        array-valued metadata is stored in children and string metadata
+        in aux_data.
+
+        Implementation Logic
+        --------------------
+        1. **Children** (JAX arrays, participate in tracing):
+           ``(num_kpoints, label_indices, points_per_segment,
+           segments, kpoints, weights, grid, shift)``.
+        2. **Auxiliary data** (static, not traced by JAX):
+           ``(mode, labels, comment, coordinate_mode)``.
+           JAX treats these as compile-time constants; any change
+           triggers JIT recompilation.
+
+        Returns
+        -------
+        children : tuple of Array
+            Tuple of array-valued metadata fields.
+        aux_data : tuple[str, tuple[str, ...], str, str]
+            String metadata stored outside JAX tracing.
+        """
+        return (
+            (
+                self.num_kpoints,
+                self.label_indices,
+                self.points_per_segment,
+                self.segments,
+                self.kpoints,
+                self.weights,
+                self.grid,
+                self.shift,
+            ),
+            (
+                self.mode,
+                self.labels,
+                self.comment,
+                self.coordinate_mode,
+            ),
+        )
+
+    @classmethod
+    def tree_unflatten(
+        cls,
+        aux_data: Tuple[str, tuple[str, ...], str, str],
+        children: Tuple[
+            Int[Array, " "],
+            Int[Array, " L"],
+            Int[Array, " "],
+            Int[Array, " "],
+            Optional[Float[Array, "K 3"]],
+            Optional[Float[Array, " K"]],
+            Optional[Int[Array, " 3"]],
+            Optional[Float[Array, " 3"]],
+        ],
+    ) -> "KPathInfo":
+        """Reconstruct a KPathInfo from flattened components.
+
+        Inverse of :meth:`tree_flatten`. JAX calls this method
+        automatically when unflattening a PyTree after a
+        transformation (e.g., inside ``jax.jit`` or ``jax.grad``).
+
+        Implementation Logic
+        --------------------
+        1. Unpack ``children`` into the eight array-valued fields.
+        2. Unpack ``aux_data`` into the four string fields.
+        3. Pass all fields to the constructor.
+
+        Parameters
+        ----------
+        aux_data : tuple[str, tuple[str, ...], str, str]
+            String metadata recovered from auxiliary data.
+        children : tuple of Array
+            Array-valued metadata tuple.
+
+        Returns
+        -------
+        kpath : KPathInfo
+            Reconstructed instance with identical data.
+        """
+        num_kpoints: Int[Array, " "]
+        label_indices: Int[Array, " L"]
+        points_per_segment: Int[Array, " "]
+        segments: Int[Array, " "]
+        kpoints: Optional[Float[Array, "K 3"]]
+        weights: Optional[Float[Array, " K"]]
+        grid: Optional[Int[Array, " 3"]]
+        shift: Optional[Float[Array, " 3"]]
+        (
+            num_kpoints,
+            label_indices,
+            points_per_segment,
+            segments,
+            kpoints,
+            weights,
+            grid,
+            shift,
+        ) = children
+        mode: str
+        labels: tuple[str, ...]
+        comment: str
+        coordinate_mode: str
+        mode, labels, comment, coordinate_mode = aux_data
+        kpath: KPathInfo = cls(
+            num_kpoints=num_kpoints,
+            label_indices=label_indices,
+            points_per_segment=points_per_segment,
+            segments=segments,
+            kpoints=kpoints,
+            weights=weights,
+            grid=grid,
+            shift=shift,
+            mode=mode,
+            labels=labels,
+            comment=comment,
+            coordinate_mode=coordinate_mode,
+        )
+        return kpath
+
+
+@jaxtyped(typechecker=beartype)
+def make_kpath_info(  # noqa: PLR0913
+    num_kpoints: Union[int, Int[Array, " "]],
+    label_indices: Union[Int[Array, " L"], "list[int]"],
+    points_per_segment: Union[int, Int[Array, " "]] = 0,
+    segments: Union[int, Int[Array, " "]] = 0,
+    kpoints: Optional[Float[Array, "K 3"]] = None,
+    weights: Optional[Float[Array, " K"]] = None,
+    grid: Optional[Union[Int[Array, " 3"], "list[int]"]] = None,
+    shift: Optional[Float[Array, " 3"]] = None,
+    mode: str = "Line-mode",
+    labels: tuple[str, ...] = (),
+    comment: str = "",
+    coordinate_mode: str = "",
+) -> KPathInfo:
+    """Create a validated KPathInfo instance.
+
+    Extended Summary
+    ----------------
+    Factory function that validates and normalises raw k-path
+    metadata before constructing a ``KPathInfo`` PyTree. Integer
+    scalars and arrays are cast to ``int32``; float arrays are cast
+    to ``float64``. Optional fields (``kpoints``, ``weights``,
+    ``grid``, ``shift``) are cast only when present, preserving
+    ``None`` for modes that do not use them.
+
+    The function is decorated with ``@jaxtyped(typechecker=beartype)``
+    so that shape constraints are checked at call time. String
+    metadata (``mode``, ``labels``, ``comment``, ``coordinate_mode``)
+    is passed through unchanged and stored 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.
+
+    Implementation Logic
+    --------------------
+    1. **Cast integer fields** (``num_kpoints``, ``label_indices``,
+       ``points_per_segment``, ``segments``) to ``jnp.int32`` via
+       ``jnp.asarray``. Accepts both Python ints and JAX arrays.
+    2. **Cast optional float fields** (``kpoints``, ``weights``,
+       ``shift``) to ``jnp.float64`` when not ``None``.
+    3. **Cast optional int field** (``grid``) to ``jnp.int32`` when
+       not ``None``.
+    4. **Pass string metadata** (``mode``, ``labels``, ``comment``,
+       ``coordinate_mode``) through unchanged -- stored as auxiliary
+       data in the PyTree.
+    5. **Construct** the ``KPathInfo`` NamedTuple from all validated
+       fields and return it.
+
+    Parameters
+    ----------
+    num_kpoints : Union[int, Int[Array, " "]]
+        Total number of k-points along the path.
+    label_indices : Union[Int[Array, " L"], list[int]]
+        Indices of symmetry points along the path.
+    points_per_segment : Union[int, Int[Array, " "]], optional
+        Raw value from line 2 of KPOINTS. Default is 0.
+    segments : Union[int, Int[Array, " "]], optional
+        Number of path segments in line mode. Default is 0.
+    kpoints : Optional[Float[Array, "K 3"]], optional
+        Mode-specific k-point coordinates. Default is None.
+    weights : Optional[Float[Array, " K"]], optional
+        Explicit-mode weights. Default is None.
+    grid : Optional[Union[Int[Array, " 3"], list[int]]], optional
+        Automatic-mode MP/Gamma grid. Default is None.
+    shift : Optional[Float[Array, " 3"]], optional
+        Automatic-mode grid shift. Default is None.
+    mode : str, optional
+        KPOINTS file mode. Default is ``"Line-mode"``.
+    labels : tuple[str, ...], optional
+        Symmetry point labels. Default is empty tuple.
+    comment : str, optional
+        KPOINTS comment line. Default is empty string.
+    coordinate_mode : str, optional
+        Coordinate/scheme line metadata. Default is empty string.
+
+    Returns
+    -------
+    kpath : KPathInfo
+        Validated k-path info instance.
+
+    See Also
+    --------
+    KPathInfo : The PyTree class constructed by this factory.
+    """
+    nkpts_arr: Int[Array, " "] = jnp.asarray(num_kpoints, dtype=jnp.int32)
+    indices_arr: Int[Array, " L"] = jnp.asarray(label_indices, dtype=jnp.int32)
+    pps_arr: Int[Array, " "] = jnp.asarray(points_per_segment, dtype=jnp.int32)
+    segments_arr: Int[Array, " "] = jnp.asarray(segments, dtype=jnp.int32)
+    kpoints_arr: Optional[Float[Array, "K 3"]] = None
+    if kpoints is not None:
+        kpoints_arr = jnp.asarray(kpoints, dtype=jnp.float64)
+    weights_arr: Optional[Float[Array, " K"]] = None
+    if weights is not None:
+        weights_arr = jnp.asarray(weights, dtype=jnp.float64)
+    grid_arr: Optional[Int[Array, " 3"]] = None
+    if grid is not None:
+        grid_arr = jnp.asarray(grid, dtype=jnp.int32)
+    shift_arr: Optional[Float[Array, " 3"]] = None
+    if shift is not None:
+        shift_arr = jnp.asarray(shift, dtype=jnp.float64)
+    kpath: KPathInfo = KPathInfo(
+        num_kpoints=nkpts_arr,
+        label_indices=indices_arr,
+        points_per_segment=pps_arr,
+        segments=segments_arr,
+        kpoints=kpoints_arr,
+        weights=weights_arr,
+        grid=grid_arr,
+        shift=shift_arr,
+        mode=mode,
+        labels=labels,
+        comment=comment,
+        coordinate_mode=coordinate_mode,
+    )
+    return kpath
+
+
+__all__: list[str] = [
+    "KPathInfo",
+    "make_kpath_info",
+]