kups.core.cell ¶

@dataclass
class Cell[P: tuple[bool, bool, bool]](Sliceable):
    """A [Frame][kups.core.cell.Frame] plus per-axis boundary semantics.

    Generic over the periodicity literal ``P`` (a length-3 tuple of
    booleans). [PeriodicCell][kups.core.cell.PeriodicCell] and
    [VacuumCell][kups.core.cell.VacuumCell] are subclasses that pin ``P``
    to a literal-typed default; for slab and wire geometries, construct
    ``Cell(frame, periodic=mask)`` directly — the literal tuple narrows
    ``P`` to the corresponding [Slab2D][kups.core.cell.Slab2D] or
    [Wire1D][kups.core.cell.Wire1D] alias.

    The cell delegates geometry queries (``volume``, ``vectors``, etc.) to
    its frame. Periodic-mask-aware operations
    ([`wrap`][kups.core.cell.Cell.wrap], [`fold`][kups.core.cell.Cell.fold],
    [`minimum_image_shifts`][kups.core.cell.Cell.minimum_image_shifts])
    live on the cell so callers don't need to access ``cell.periodic``
    directly.
    """

    frame: Frame
    periodic: P = field(static=True)

    @property
    def vectors(self) -> Array:
        return self.frame.vectors

    @property
    def inverse_vectors(self) -> Array:
        return self.frame.inverse_vectors

    @property
    def volume(self) -> Array:
        return self.frame.volume

    @property
    def perpendicular_lengths(self) -> Array:
        return self.frame.perpendicular_lengths

    def wrap(
        self,
        r: Array,
        *,
        input_space: CoordinateSpace = CoordinateSpace.REAL,
        output_space: CoordinateSpace = CoordinateSpace.REAL,
    ) -> Array:
        return _wrap(self.frame, self.periodic, r, input_space, output_space)

    def fold(self, r_frac: Array) -> tuple[Array, Array]:
        """Fold fractional coords into ``[0, 1)`` on periodic axes; non-periodic
        axes pass through unchanged.

        Returns ``(folded, in_cell)`` where ``in_cell`` is a per-particle
        mask, ``True`` where the folded coords lie in ``[0, 1)`` on every
        axis. For fully-periodic cells, ``in_cell`` is trivially ``True``
        after folding; for cells with non-periodic axes, particles that
        leaked out of the box on those axes are flagged ``False``.

        Used by neighbor-list spatial hashing; complementary to
        [`wrap`][kups.core.cell.Cell.wrap] which uses the ``[-0.5, 0.5)``
        convention.
        """
        folded = jnp.where(jnp.array(self.periodic), r_frac % 1, r_frac)
        in_cell = jnp.all((folded >= 0) & (folded < 1), axis=-1)
        return folded, in_cell

    def minimum_image_shifts(self, deltas: Array) -> Array:
        """Per-axis minimum-image shifts for fractional separations.

        Returns ``round(deltas)`` on periodic axes (the closest integer
        cell offset that wraps the separation to its minimum image) and
        ``0`` on non-periodic axes.
        """
        return jnp.where(jnp.array(self.periodic), jnp.round(deltas), 0.0)

    def __mul__(self, other: Array | float | int) -> Self:
        return bind(self, lambda c: c.frame).set(self.frame * other)

    @overload
    @staticmethod
    def from_pbc(frame: Frame, pbc: Periodic3D) -> PeriodicCell: ...
    @overload
    @staticmethod
    def from_pbc(frame: Frame, pbc: Vacuum) -> VacuumCell: ...
    @overload
    @staticmethod
    def from_pbc[Q: tuple[bool, bool, bool]](frame: Frame, pbc: Q) -> Cell[Q]: ...
    @staticmethod
    def from_pbc(frame: Frame, pbc: tuple[bool, bool, bool]) -> Cell[Any]:
        """Construct the cell flavor matching ``pbc``.

        Returns [`PeriodicCell`][kups.core.cell.PeriodicCell] for
        ``(True, True, True)``, [`VacuumCell`][kups.core.cell.VacuumCell]
        for ``(False, False, False)``, and a generic ``Cell[P]`` carrying
        the runtime mask for slab and wire geometries (``P`` is inferred
        from the literal tuple).
        """
        match pbc:
            case (True, True, True):
                return PeriodicCell(frame)
            case (False, False, False):
                return VacuumCell(frame)
            case _:
                return Cell(frame, periodic=pbc)

class CoordinateSpace(Enum):
    """Enumeration for coordinate systems.

    Attributes:
        REAL: Cartesian coordinates in Angstroms.
        FRACTIONAL: Scaled coordinates in [0, 1) relative to frame vectors.
    """

    REAL = "real"
    FRACTIONAL = "fractional"

class Frame(Protocol):
    """3D parallelepiped geometry, no periodicity attached.

    A Frame is a pure geometric container — three basis vectors that span
    a parallelepiped in 3D space. It does not commit to whether those
    vectors represent periodic translations or just bounding-box edges;
    that distinction is supplied by the [Cell][kups.core.cell.Cell] type
    that wraps a Frame.

    In crystallography these basis vectors are conventionally called
    "lattice vectors". They are exposed here under the name
    [`vectors`][kups.core.cell.Frame.vectors] because the crystallographic
    label implies a periodic interpretation that does not apply to all
    Frame uses (e.g. the bounding box of a vacuum simulation).

    Concrete implementations:

    - [OrthogonalFrame][kups.core.cell.OrthogonalFrame]: 3 DOF (lengths).
    - [TriclinicFrame][kups.core.cell.TriclinicFrame]: 6 DOF (lower-triangular).
    """

    @property
    def vectors(self) -> Array:
        """Basis vectors of the parallelepiped, shape ``(..., 3, 3)``.

        Rows are the basis vectors. Lower-triangular by convention so that
        ``v[0]`` lies along x, ``v[1]`` in the xy-plane, ``v[2]`` general.
        Crystallography calls this matrix the *lattice vectors*.
        """
        ...

    @property
    def inverse_vectors(self) -> Array:
        """Matrix inverse of [`vectors`][kups.core.cell.Frame.vectors],
        used to convert real-space coordinates to fractional, shape
        ``(..., 3, 3)``."""
        ...

    @property
    def volume(self) -> Array:
        """Volume of the parallelepiped, shape ``(...)``."""
        ...

    @property
    def perpendicular_lengths(self) -> Array:
        """Perpendicular distance between opposing faces, per axis,
        shape ``(..., 3)``. Used by neighbor-list cutoff checks and by
        [min_multiplicity][kups.core.cell.min_multiplicity]."""
        ...

    def to_fractional(self, r: Array) -> Array:
        """Convert real-space coordinates to fractional, shape ``(..., 3)``."""
        ...

    def to_real(self, r_frac: Array) -> Array:
        """Convert fractional coordinates to real-space, shape ``(..., 3)``."""
        ...

    def __mul__(self, other: Array | float | int) -> Self:
        """Uniformly scale all basis vectors by ``other``."""
        ...

    def tile(self, multiplicities: tuple[int, int, int]) -> Self:
        """Per-axis integer scaling. Used by
        [make_supercell][kups.core.cell.make_supercell] to build supercells."""
        ...

    @classmethod
    def from_matrix(cls, vecs: Array) -> Self:
        """Construct from a ``(..., 3, 3)`` basis matrix.

        Projects the input onto the frame's parameter space — entries
        not represented by the parameterisation are discarded
        (``OrthogonalFrame`` keeps the diagonal; ``TriclinicFrame``
        keeps the lower-triangular block). Used to wrap a generic
        ``∂E/∂h`` matrix back into the same frame type as an input cell.
        """
        ...

    def materialize(self) -> MaterializedFrame:
        """Return a [MaterializedFrame][kups.core.cell.MaterializedFrame]
        with ``vectors``, ``inverse_vectors`` and ``volume`` evaluated
        and stored as concrete arrays.

        Use this to avoid recomputing the inverse and determinant when
        the same frame is queried many times, or to lift these arrays
        across a JIT boundary so downstream callers don't need to know
        the frame's parametrisation.

        Requires at least one leading batch dim so that ``vectors``
        ``(B, ..., 3, 3)``, ``inverse_vectors`` ``(B, ..., 3, 3)`` and
        ``volume`` ``(B, ...)`` share a leading axis (see
        [MaterializedFrame][kups.core.cell.MaterializedFrame]).
        """
        ...

@dataclass
class MaterializedFrame(Sliceable):
    """[Frame][kups.core.cell.Frame] that stores its basis matrix, inverse,
    and volume as concrete arrays.

    Produced by [`Frame.materialize`][kups.core.cell.Frame.materialize].
    Useful when the same frame is queried repeatedly (no recomputation
    of the inverse or determinant) or when the arrays need to cross a
    JIT boundary independently of the frame's original parametrisation.

    The stored matrices are assumed to follow the lower-triangular
    convention shared by every other Frame implementation; coordinate
    transforms use [triangular_3x3_matmul][kups.core.utils.math.triangular_3x3_matmul]
    accordingly. Manually constructing this frame with non-triangular
    vectors violates the convention.

    All three fields are pytree leaves and must share a leading batch
    dim ``B`` to satisfy the [Sliceable][kups.core.data.Sliceable]
    contract — unbatched single-frame inputs must be wrapped with an
    explicit ``[None]`` axis before being passed in.

    Attributes:
        vectors: Basis matrix, shape ``(B, 3, 3)``.
        inverse_vectors: Matrix inverse of ``vectors``, shape ``(B, 3, 3)``.
        volume: Absolute determinant of ``vectors``, shape ``(B,)``.
    """

    vectors: Array
    inverse_vectors: Array
    volume: Array

    @classmethod
    def from_matrix(cls, vecs: Array) -> Self:
        vecs = jnp.asarray(vecs)
        det, inv = triangular_3x3_det_and_inverse(vecs)
        return cls(vectors=vecs, inverse_vectors=inv, volume=jnp.abs(det))

    @property
    def perpendicular_lengths(self) -> Array:
        v = self.vectors
        a, b, c = v[..., 0, :], v[..., 1, :], v[..., 2, :]
        Lx = self.volume / jnp.linalg.norm(jnp.cross(b, c), axis=-1)
        Ly = self.volume / jnp.linalg.norm(jnp.cross(a, c), axis=-1)
        Lz = self.volume / jnp.linalg.norm(jnp.cross(a, b), axis=-1)
        return jnp.stack([Lx, Ly, Lz], axis=-1)

    def to_fractional(self, r: Array) -> Array:
        return triangular_3x3_matmul(self.inverse_vectors, r)

    def to_real(self, r_frac: Array) -> Array:
        return triangular_3x3_matmul(self.vectors, r_frac)

    def tile(self, multiplicities: tuple[int, int, int]) -> Self:
        m = jnp.asarray(multiplicities)
        return type(self)(
            vectors=self.vectors * m[:, None],
            inverse_vectors=self.inverse_vectors / m[None, :],
            volume=self.volume * jnp.prod(m),
        )

    def __mul__(self, other: Array | float | int) -> Self:
        scale = jnp.asarray(other)
        return type(self)(
            vectors=self.vectors * scale[..., None, None],
            inverse_vectors=self.inverse_vectors / scale[..., None, None],
            volume=self.volume * scale**3,
        )

    def materialize(self) -> Self:
        return self

@dataclass
class OrthogonalFrame(Sliceable):
    """Axis-aligned [Frame][kups.core.cell.Frame] with 3 degrees of freedom.

    Parameterized by the three side lengths. Exploits the diagonal
    structure for cheaper volume, inverse, and coordinate-transform
    operations than the general triclinic path. Use this when the
    simulation domain has perpendicular axes (cubic, tetragonal, or
    orthorhombic crystals; standard rectangular MD boxes).

    Attributes:
        lengths: Box side lengths ``[Lx, Ly, Lz]`` in Angstroms,
            shape ``(..., 3)``.
    """

    lengths: Array

    @classmethod
    def from_matrix(cls, vecs: Array) -> Self:
        """Construct from a diagonal basis matrix, shape ``(..., 3, 3)``.

        Projects the input matrix onto the orthogonal subspace by taking
        its diagonal — off-diagonal entries are discarded, matching the
        3-parameter ``(Lx, Ly, Lz)`` representation.
        """
        vecs = jnp.asarray(vecs)
        return cls(jnp.diagonal(vecs, axis1=-2, axis2=-1))

    @property
    def vectors(self) -> Array:
        return self.lengths[..., :, None] * jnp.eye(3)

    @property
    def inverse_vectors(self) -> Array:
        return (1.0 / self.lengths)[..., :, None] * jnp.eye(3)

    @property
    def volume(self) -> Array:
        return jnp.prod(self.lengths, axis=-1)

    @property
    def perpendicular_lengths(self) -> Array:
        return self.lengths

    def to_fractional(self, r: Array) -> Array:
        return r / self.lengths

    def to_real(self, r_frac: Array) -> Array:
        return r_frac * self.lengths

    def tile(self, multiplicities: tuple[int, int, int]) -> Self:
        return type(self)(self.lengths * jnp.asarray(multiplicities))

    def __mul__(self, other: Array | float | int) -> Self:
        return type(self)(self.lengths * jnp.asarray(other)[..., None])

    def materialize(self) -> MaterializedFrame:
        return MaterializedFrame(
            vectors=self.vectors,
            inverse_vectors=self.inverse_vectors,
            volume=self.volume,
        )

@dataclass
class PeriodicCell(Cell[Periodic3D]):
    """Cell that is periodic along all three axes.

    The frame is interpreted as a *unit cell* — a tile of a periodic
    crystal or fluid. ``wrap`` folds coordinates into the primary unit
    cell; the neighbor list applies the minimum-image convention; the
    Ewald summation requires this cell type.
    """

    periodic: Periodic3D = field(default=(True, True, True), init=False, static=True)

@dataclass
class TriclinicFrame(Sliceable):
    """General triclinic [Frame][kups.core.cell.Frame] with 6 degrees of freedom.

    Stores the 6 independent elements of the lower-triangular basis matrix.
    Vectors are a linear function of these parameters, making them suitable
    for gradient-based optimization (e.g. NPT cell-vector relaxation).

    Use this when the simulation domain has non-orthogonal axes
    (monoclinic / triclinic crystals, sheared MD boxes). For axis-aligned
    domains, [OrthogonalFrame][kups.core.cell.OrthogonalFrame] is cheaper.

    Attributes:
        tril: Lower-triangular elements ``[L00, L10, L11, L20, L21, L22]``,
            shape ``(..., 6)``. The basis matrix is::

                [[L00,   0,   0],
                 [L10, L11,   0],
                 [L20, L21, L22]]
    """

    tril: Array

    @classmethod
    def from_matrix(cls, vecs: Array) -> TriclinicFrame:
        """Construct from a lower-triangular basis matrix, shape ``(..., 3, 3)``."""
        vecs = jnp.asarray(vecs)
        return cls(vecs[..., *np.tril_indices(3)])

    @classmethod
    def from_lengths_and_angles(cls, lengths: Array, angles: Array) -> TriclinicFrame:
        """Construct from crystallographic parameters.

        Args:
            lengths: Lattice lengths ``[a, b, c]`` in Angstroms, shape ``(..., 3)``.
            angles: Lattice angles ``[alpha, beta, gamma]`` in degrees, shape ``(..., 3)``.
                alpha = angle(b, c), beta = angle(a, c), gamma = angle(a, b).
        """
        return cls.from_matrix(_build_vectors(lengths, angles))

    @property
    def vectors(self) -> Array:
        zero = jnp.zeros_like(self.tril[..., :1])
        return jnp.stack(
            [
                jnp.concatenate([self.tril[..., 0:1], zero, zero], axis=-1),
                jnp.concatenate([self.tril[..., 1:3], zero], axis=-1),
                self.tril[..., 3:6],
            ],
            axis=-2,
        )

    @property
    def inverse_vectors(self) -> Array:
        return triangular_3x3_det_and_inverse(self.vectors)[1]

    @property
    def volume(self) -> Array:
        return jnp.abs(self.tril[..., 0] * self.tril[..., 2] * self.tril[..., 5])

    @property
    def lengths(self) -> Array:
        return jnp.linalg.norm(self.vectors, axis=-1)

    @property
    def angles(self) -> Array:
        v = self.vectors
        a, b, c = v[..., 0, :], v[..., 1, :], v[..., 2, :]
        la, lb, lc = (
            jnp.linalg.norm(a, axis=-1),
            jnp.linalg.norm(b, axis=-1),
            jnp.linalg.norm(c, axis=-1),
        )
        cos_alpha = jnp.clip(jnp.sum(b * c, axis=-1) / (lb * lc), -1.0, 1.0)
        cos_beta = jnp.clip(jnp.sum(a * c, axis=-1) / (la * lc), -1.0, 1.0)
        cos_gamma = jnp.clip(jnp.sum(a * b, axis=-1) / (la * lb), -1.0, 1.0)
        return jnp.degrees(
            jnp.stack(
                [jnp.arccos(cos_alpha), jnp.arccos(cos_beta), jnp.arccos(cos_gamma)],
                axis=-1,
            )
        )

    @property
    def perpendicular_lengths(self) -> Array:
        v = self.vectors
        a, b, c = v[..., 0, :], v[..., 1, :], v[..., 2, :]
        Lx = self.volume / jnp.linalg.norm(jnp.cross(b, c), axis=-1)
        Ly = self.volume / jnp.linalg.norm(jnp.cross(a, c), axis=-1)
        Lz = self.volume / jnp.linalg.norm(jnp.cross(a, b), axis=-1)
        return jnp.stack([Lx, Ly, Lz], axis=-1)

    def to_fractional(self, r: Array) -> Array:
        return triangular_3x3_matmul(self.inverse_vectors, r)

    def to_real(self, r_frac: Array) -> Array:
        return triangular_3x3_matmul(self.vectors, r_frac)

    def tile(self, multiplicities: tuple[int, int, int]) -> Self:
        m = jnp.asarray(multiplicities)
        scale = jnp.array([m[0], m[1], m[1], m[2], m[2], m[2]])
        return type(self)(self.tril * scale)

    def __mul__(self, other: Array | float | int) -> Self:
        return type(self)(self.tril * jnp.asarray(other)[..., None])

    def materialize(self) -> MaterializedFrame:
        vecs = self.vectors
        det, inv = triangular_3x3_det_and_inverse(vecs)
        return MaterializedFrame(vectors=vecs, inverse_vectors=inv, volume=jnp.abs(det))

class TriclinicMap(Protocol):
    """Mapping that takes positions in an arbitrary frame and returns them
    in the lower-triangular triclinic frame produced by
    [to_lower_triangular][kups.core.cell.to_lower_triangular]."""

    def __call__(self, r: Array, /) -> Array: ...

@dataclass
class VacuumCell(Cell[Vacuum]):
    """Cell with all three axes open.

    The frame is interpreted as the *bounding parallelepiped* of a finite
    simulation domain — a cluster, an isolated molecule, a gas-phase
    sample. ``wrap`` is a no-op; long-range electrostatics use direct
    pairwise sums (no Ewald). The frame is still required because the
    neighbor-list machinery needs a spatial-partitioning hint.
    """

    periodic: Vacuum = field(default=(False, False, False), init=False, static=True)