kups.core.cell ¶

class BaseFrame(Frame, abc.ABC):
    """Abstract base implementing every [Frame][kups.core.cell.Frame]
    operation that derives purely from
    [`vectors`][kups.core.cell.Frame.vectors].

    Concrete frames subclass ``BaseFrame`` and supply only their
    parameterisation — the ``vectors`` primitive plus
    [`from_matrix`][kups.core.cell.Frame.from_matrix],
    [`tile`][kups.core.cell.Frame.tile] and ``__mul__``. The lower-triangular
    convention shared by all frames lets the basis inverse, volume,
    perpendicular lengths, and coordinate transforms be computed here once via
    the triangular fast paths in
    [kups.core.utils.math][kups.core.utils.math]. Subclasses may override any
    of these with cheaper specialised forms (e.g. ``OrthogonalFrame``'s
    diagonal paths).

    Mirrors the [Lens][kups.core.lens.Lens] /
    [BaseLens][kups.core.lens.BaseLens] split: ``Frame`` is the structural
    interface, ``BaseFrame`` the shared implementation.
    """

    @classmethod
    @abc.abstractmethod
    @override
    def from_matrix(cls, vecs: Array) -> Self: ...

    @abc.abstractmethod
    @override
    def tile(self, multiplicities: tuple[int, int, int]) -> Self: ...

    @abc.abstractmethod
    @override
    def __mul__(self, other: Array | float | int) -> Self: ...

    @property
    @override
    def matrix(self) -> SquareMatrix:
        # Conservative default: any frame without an explicit structure is
        # correct (just not fast) under the general dense paths.
        return GeneralSquareMatrix(self.vectors)

    @property
    @override
    def inverse_vectors(self) -> Array:
        return self.matrix.inverse().array

    @property
    @override
    def reference_vectors(self) -> Array:
        # No stored reference: atom DOFs are fractional. DeformedFrame overrides.
        return jnp.broadcast_to(
            jnp.eye(3, dtype=self.vectors.dtype), self.vectors.shape
        )

    @property
    @override
    def volume(self) -> Array:
        return jnp.abs(self.matrix.det())

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

    @override
    def to_fractional(self, r: Array) -> Array:
        return self.matrix.inverse().matmul(r)

    @override
    def to_real(self, r_frac: Array) -> Array:
        return self.matrix.matmul(r_frac)

    @override
    @jit
    def materialize(self) -> MaterializedFrame:
        matrix = self.matrix
        det, inv = matrix.det(), matrix.inverse()
        return MaterializedFrame(self.matrix, inv, jnp.abs(det))

    @override
    @jit
    def parameter_gradient(self, vectors_grad: Array) -> Self:
        (grad,) = jax.vjp(lambda f: f.vectors, self)[1](vectors_grad)
        return grad

    @override
    @jit
    def vectors_gradient(self, parameter_grad: Frame) -> Array:
        def single(frame: Frame, grad: Frame) -> Array:
            theta, unravel = ravel_pytree(frame)
            g_theta, _ = ravel_pytree(grad)
            jac = jax.jacfwd(lambda p: unravel(p).vectors)(theta)  # (3, 3, n)
            # vec(∂E/∂h) = (Jᵀ)⁺·∂E/∂θ; lstsq on the wide ``Jᵀ`` gives the
            # min-norm solution, i.e. the pseudo-inverse.
            vec_grad = jnp.linalg.lstsq(jac.reshape(9, theta.shape[0]).T, g_theta)[0]
            return vec_grad.reshape(3, 3)

        fn = single
        for _ in range(self.vectors.ndim - 2):
            fn = jax.vmap(fn)
        with no_post_init():
            return fn(self, parameter_grad)

@dataclass
class Cell(Sliceable, Generic[_P]):
    """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)

    def materialize(self) -> Self:
        return bind(self, lambda c: c.frame).set(self.frame.materialize())

    @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"

@dataclass
class DeformedFrame(BaseFrame, Sliceable):
    """[Frame][kups.core.cell.Frame] parameterised as a differentiable deformation
    of a fixed reference frame: ``vectors = base @ deformation``.

    ``base`` is held fixed (stop-gradient); only ``deformation``'s parameters are
    optimised. Right-multiplication deforms Cartesian space, so ``vectors`` stays
    lower-triangular and atoms at fixed fractional coordinates ride the cell. The
    deformation frame chooses the parameterisation: a
    [LogTriclinicFrame][kups.core.cell.LogTriclinicFrame] gives the unconstrained
    matrix-exponential map of ASE's ``FrechetCellFilter`` (build with
    [from_frame][kups.core.cell.DeformedFrame.from_frame]); a
    [TriclinicFrame][kups.core.cell.TriclinicFrame] gives the linear deformation
    gradient of ASE's ``UnitCellFilter``. ``vectors`` is generally nonlinear in the
    parameters, so the gradient maps use the general inverse-Jacobian path on
    [BaseFrame][kups.core.cell.BaseFrame].

    Attributes:
        base: Reference frame, held fixed (stop-gradient); the identity
            deformation reproduces it.
        deformation: Differentiable deformation gradient (the optimised frame).
    """

    base: Frame
    deformation: Frame

    @classmethod
    def from_frame(
        cls,
        frame: Frame,
        *,
        cell_factor: float | Array = 1.0,
        deformation: type[LogTriclinicFrame | MatrixLogFrame] = LogTriclinicFrame,
    ) -> Self:
        """ASE ``FrechetCellFilter``-style: an exponential-map deformation anchored
        at ``frame`` (identity deformation, so ``vectors == frame.vectors``).

        ``deformation`` selects the deformation parameterisation:
        [LogTriclinicFrame][kups.core.cell.LogTriclinicFrame] (default) keeps
        ``base = frame`` with a lower-triangular ``(6,)`` matrix-log deformation;
        [MatrixLogFrame][kups.core.cell.MatrixLogFrame] gives a full-``3x3``
        deformation, re-anchoring ``base`` as a ``MatrixLogFrame`` built from
        ``frame.vectors`` (via the exact triangular log) so every leaf is
        ``(..., 3, 3)`` and the resulting frame represents rotated / sheared cells.
        """
        eye = jnp.broadcast_to(
            jnp.eye(3, dtype=frame.vectors.dtype), frame.vectors.shape
        )
        deform = deformation.from_matrix(eye, cell_factor=cell_factor)
        return cls(frame, deform)

    @classmethod
    @override
    def from_matrix(cls, vecs: Array) -> Self:
        return cls.from_frame(TriclinicFrame.from_matrix(vecs))

    @property
    @override
    def matrix(self) -> SquareMatrix:
        return self._base.matmul(self.deformation.matrix)

    @property
    def _base(self) -> SquareMatrix:
        return jax.lax.stop_gradient(self.base.matrix)

    @property
    @override
    def reference_vectors(self) -> Array:
        return self._base.array

    @property
    @override
    def vectors(self) -> Array:
        # Right-multiply: deform Cartesian space (v -> v @ deformation), so atoms at
        # fixed fractional coordinates ride the cell.
        return self._base.matmul(self.deformation.matrix).array

    @override
    def tile(self, multiplicities: tuple[int, int, int]) -> Self:
        # Per-axis row scaling acts on the left of ``vectors`` with no closed form
        # in the right deformation factor, so apply it to the reference base.
        return type(self)(self.base.tile(multiplicities), self.deformation)

    @override
    def __mul__(self, other: Array | float | int) -> Self:
        # Uniform scaling commutes with the right factor: scale the base.
        return type(self)(self.base * other, self.deformation)

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).

    The parameterised implementations subclass
    [BaseFrame][kups.core.cell.BaseFrame], which implements every operation
    derivable from [`vectors`][kups.core.cell.Frame.vectors]:

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

    [MaterializedFrame][kups.core.cell.MaterializedFrame] stores the basis
    matrix, inverse and volume directly and implements this interface without
    going through ``BaseFrame``.
    """

    @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 matrix(self) -> SquareMatrix:
        """Static structure of [`vectors`][kups.core.cell.Frame.vectors] used to
        dispatch the geometric primitives (inverse, volume, transforms) onto the
        matching fast path. ``LOWER_TRIANGULAR`` / ``DIAGONAL`` keep the triangular
        and diagonal forms; ``GENERAL`` falls back to dense ``3x3`` algebra."""
        ...

    @property
    def reference_vectors(self) -> Array:
        """Fixed reference basis for deformation-relative atom coordinates,
        shape ``(..., 3, 3)``. The identity for directly-parameterised frames
        (so atom DOFs are fractional); a frame that parameterises a deformation
        of a stored reference (e.g. [DeformedFrame][kups.core.cell.DeformedFrame])
        returns that reference, conditioning the DOFs against it."""
        ...

    @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]).
        """
        ...

    def parameter_gradient(self, vectors_grad: Array) -> Self:
        """Pull a cartesian ``∂E/∂h`` matrix back onto the frame's parameter
        space, returning a same-type frame whose parameter leaves hold
        ``∂E/∂θ``.

        Implements ``∂E/∂θ = Jᵀ·∂E/∂h`` with ``J = ∂vectors/∂θ`` via the
        vector-Jacobian product of [`vectors`][kups.core.cell.Frame.vectors].
        General for any parameterisation — the vjp linearises at the current
        parameters, so it is correct even when ``vectors`` is a nonlinear
        function of the parameters. Inverse of
        [`vectors_gradient`][kups.core.cell.Frame.vectors_gradient].
        """
        ...

    def vectors_gradient(self, parameter_grad: Frame) -> Array:
        """Map a parameter-space gradient frame (``parameter_grad`` holding
        ``∂E/∂θ`` in its parameter leaves) to the cartesian ``∂E/∂h`` matrix,
        shape ``(..., 3, 3)``.

        Implements ``vec(∂E/∂h) = (Jᵀ)⁺·∂E/∂θ`` with ``J = ∂vectors/∂θ`` — the
        pseudo-inverse of the Jacobian, evaluated at this frame's parameters.
        For a parameterisation whose ``vectors`` map is linear with orthonormal
        Jacobian (``TriclinicFrame``, ``OrthogonalFrame``) this reduces to
        reconstructing the matrix from the gradient frame; the general path
        covers nonlinear parameterisations. Inverse of
        [`parameter_gradient`][kups.core.cell.Frame.parameter_gradient].
        """
        ...

class LinearFrame(BaseFrame, abc.ABC):
    """Abstract base for frames whose ``vectors`` map is linear with orthonormal
    Jacobian (e.g. ``OrthogonalFrame``, ``TriclinicFrame``).

    For these frames, the parameter-space gradient to cartesian gradient map
    reduces to reconstructing the matrix from the parameter leaves, so we
    can skip the general pseudo-inverse path in
    [BaseFrame.vectors_gradient][kups.core.cell.BaseFrame.vectors_gradient] and
    just return the reconstructed matrix directly.
    """

    @override
    def vectors_gradient(self, parameter_grad: Frame) -> Array:
        return self.from_matrix(parameter_grad.vectors).vectors

    @override
    def parameter_gradient(self, vectors_grad: Array) -> Self:
        return self.from_matrix(vectors_grad)

@dataclass
class LogTriclinicFrame(BaseFrame, Sliceable):
    """Triclinic [Frame][kups.core.cell.Frame] parameterised in the matrix log.

    Stores the 6 lower-triangular elements of a matrix ``A``; the basis vectors
    are ``expm(A)``. Because ``A`` is lower-triangular, ``expm(A)`` is too (its
    diagonal is ``exp`` of ``A``'s diagonal), so the lower-triangular convention
    shared by every other frame is preserved.

    The exponential map is unconstrained: any ``A`` yields a basis with strictly
    positive volume ``exp(tr A)``, so gradient-based cell relaxation in ``A`` never
    has to guard against a collapsing or inverting cell. Unlike
    [TriclinicFrame][kups.core.cell.TriclinicFrame], ``vectors`` is a *nonlinear*
    function of the parameters, so the parameter/cartesian gradient maps use the
    general inverse-Jacobian path on [BaseFrame][kups.core.cell.BaseFrame] rather
    than a constant-Jacobian shortcut.

    Attributes:
        tril: Lower-triangular elements ``[A00, A10, A11, A20, A21, A22]`` of
            ``A``, shape ``(..., 6)``.
        cell_factor: Stiffening factor (ASE's ``exp_cell_factor``) so that
            ``vectors = expm(A / cell_factor)`` and the gradient w.r.t. ``A``
            carries a ``1 / cell_factor`` factor. A stop-gradient leaf broadcast
            to ``tril``'s shape (equal per slot) so every leaf shares ``tril``'s
            leading dimension inside batched containers; read the per-system value
            through ``factor``.
    """

    tril: Array
    cell_factor: Array

    @classmethod
    def from_frame(cls, frame: Frame, *, cell_factor: float | Array = 1.0) -> Self:
        return cls.from_matrix(frame.vectors, cell_factor=cell_factor)

    @classmethod
    @override
    def from_matrix(cls, vecs: Array, *, cell_factor: float | Array = 1.0) -> Self:
        """Construct from a lower-triangular basis matrix ``L`` (positive
        diagonal), shape ``(..., 3, 3)``, storing ``A = cell_factor * logm(L)`` so
        that ``vectors == expm(A / cell_factor) == L``."""
        vecs = jnp.asarray(vecs)
        tril = triangular_3x3_logm(vecs)[..., *np.tril_indices(3)]
        cf = _broadcast_cell_factor(cell_factor, tril)
        return cls(tril * cf, cf)

    @property
    @override
    def matrix(self) -> SquareMatrix:
        # expm of a lower-triangular A is lower-triangular.
        return LowerTriangularSquareMatrix(self.vectors)

    @property
    @override
    def vectors(self) -> Array:
        # fixed preconditioner, never optimised.
        cf = jax.lax.stop_gradient(self.cell_factor)
        A = triangular_3x3_from_tril(self.tril / cf)
        return triangular_3x3_expm(A=A)

    @property
    @override
    def volume(self) -> Array:
        # det(expm(A / cf)) = exp(tr(A / cf)); A's diagonal is tril[0], tril[2], tril[5].
        return jnp.exp(
            (self.tril[..., 0] + self.tril[..., 2] + self.tril[..., 5])
            / self.cell_factor[..., 0]
        )

    @override
    def tile(self, multiplicities: tuple[int, int, int]) -> Self:
        # Per-axis row scaling has no closed form in A (diag(m) does not commute
        # with A), so scale the basis matrix and re-take the log.
        m = jnp.asarray(multiplicities)
        return type(self).from_matrix(
            self.vectors * m[:, None], cell_factor=self.cell_factor
        )

    @override
    def __mul__(self, other: Array | float | int) -> Self:
        # Uniform scaling commutes: expm(B + log(s) I) = s * expm(B) for B = A / cf,
        # so add cf * log(s) to A's diagonal elements (tril indices 0, 2, 5).
        log_scale = self.cell_factor * jnp.log(jnp.asarray(other))
        diagonal = jnp.array([1.0, 0.0, 1.0, 0.0, 0.0, 1.0])
        return type(self)(self.tril + log_scale * diagonal, self.cell_factor)

@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.

    Coordinate transforms dispatch on the stored
    [`structure`][kups.core.cell.MaterializedFrame.structure] (default
    ``GENERAL``, so manually-constructed frames with arbitrary vectors are
    correct); [`materialize`][kups.core.cell.Frame.materialize] propagates the
    source frame's structure so triangular / diagonal sources keep the fast path.

    The three array 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,)``.
        structure: Static structure flag dispatching the coordinate transforms.
    """

    matrix: SquareMatrix
    inverse_matrix: SquareMatrix
    volume: Array

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

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

    @classmethod
    def from_matrix(
        cls,
        vecs: Array,
        *,
        structure: Callable[[Array], SquareMatrix] = GeneralSquareMatrix,
    ) -> Self:
        matrix = structure(jnp.asarray(vecs))
        det, inv = matrix.det(), matrix.inverse()
        return cls(matrix, inv, jnp.abs(det))

    @property
    def reference_vectors(self) -> Array:
        return jnp.broadcast_to(
            jnp.eye(3, dtype=self.vectors.dtype), self.vectors.shape
        )

    @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 self.inverse_matrix.matmul(r)

    def to_real(self, r_frac: Array) -> Array:
        return self.matrix.matmul(r_frac)

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

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

    def materialize(self) -> Self:
        return self

    def parameter_gradient(self, vectors_grad: Array) -> Self:
        (grad,) = jax.vjp(lambda f: f.vectors, self)[1](vectors_grad)
        return grad

    def vectors_gradient(self, parameter_grad: Frame) -> Array:
        # Stored ``vectors`` is the parameterisation (identity Jacobian).
        return parameter_grad.vectors

@dataclass
class MatrixLogFrame(BaseFrame, Sliceable):
    """Full-3×3 matrix-log [Frame][kups.core.cell.Frame]: ``vectors = expm(A / cf)``.

    The ``GENERAL`` sibling of [LogTriclinicFrame][kups.core.cell.LogTriclinicFrame]:
    stores a full ``(..., 3, 3)`` generator ``A`` and exponentiates the *whole*
    matrix (not just a lower-triangular block), so ``vectors`` can be a rotated /
    sheared cell. The matrix exponential is the principal one; its vjp is the
    ``expm_frechet`` adjoint, so the inherited
    [`parameter_gradient`][kups.core.cell.Frame.parameter_gradient] /
    [`vectors_gradient`][kups.core.cell.Frame.vectors_gradient] are exact with no
    custom rule.

    Used as both the base and the deformation of a
    [DeformedFrame][kups.core.cell.DeformedFrame] (build with
    [from_frame][kups.core.cell.DeformedFrame.from_frame] passing
    ``deformation=MatrixLogFrame``) to express the full-``3x3`` Frechet basis: every
    leaf is ``(..., 3, 3)``, so a per-system
    [Batched][kups.core.data.batched.Batched] slice yields all-``(3, 3)`` leaves
    that share a leading axis (a ``(6,)`` tril base would fail).

    Attributes:
        log_matrix: Generator ``A``, shape ``(..., 3, 3)``.
        cell_factor: Stiffening factor (ASE's ``exp_cell_factor``) so that
            ``vectors = expm(A / cell_factor)``. A stop-gradient leaf broadcast to
            ``log_matrix``'s shape (equal per slot) so it shares the leading axis;
            read the per-system value through
            [`factor`][kups.core.cell.MatrixLogFrame.factor].
    """

    log_matrix: Array
    cell_factor: Array

    @classmethod
    @override
    def from_matrix(cls, vecs: Array, *, cell_factor: float | Array = 1.0) -> Self:
        """Construct from an arbitrary basis matrix via the general (eig-based)
        [logm][kups.core.utils.math.logm], shape ``(..., 3, 3)``.

        For an *arbitrary* matrix; inaccurate for repeated eigenvalues (defective
        matrices). Use [from_lower_triangular][kups.core.cell.MatrixLogFrame.from_lower_triangular]
        for the reference cell, whose diagonal can repeat.
        """
        vecs = jnp.asarray(vecs)
        log_matrix = logm(vecs).real
        cf = _broadcast_cell_factor(cell_factor, log_matrix)
        return cls(log_matrix * cf, cf)

    @classmethod
    def from_lower_triangular(
        cls, vecs: Array, *, cell_factor: float | Array = 1.0
    ) -> Self:
        """Construct from a lower-triangular basis matrix (positive diagonal) via the
        exact [triangular_3x3_logm][kups.core.utils.math.triangular_3x3_logm].

        The closed-form triangular log is exact even for repeated diagonal entries
        (cubic / tetragonal cells), unlike the eig-based
        [from_matrix][kups.core.cell.MatrixLogFrame.from_matrix].
        """
        vecs = jnp.asarray(vecs)
        log_matrix = triangular_3x3_logm(vecs)
        cf = _broadcast_cell_factor(cell_factor, log_matrix)
        return cls(log_matrix * cf, cf)

    @property
    @override
    def vectors(self) -> Array:
        # fixed preconditioner, never optimised
        cf = jax.lax.stop_gradient(self.cell_factor)
        fn = jax.scipy.linalg.expm
        for _ in range(self.log_matrix.ndim - 2):
            fn = jax.vmap(fn)
        return fn(self.log_matrix / cf)

    @override
    def tile(self, multiplicities: tuple[int, int, int]) -> Self:
        m = jnp.asarray(multiplicities)
        return type(self).from_matrix(
            self.vectors * m[:, None], cell_factor=self.cell_factor
        )

    @override
    def __mul__(self, other: Array | float | int) -> Self:
        return type(self).from_matrix(
            self.vectors * jnp.asarray(other), cell_factor=self.cell_factor
        )