kups.potential.mliap.torch ¶

class AtomGraphInput(TypedDict):
    """Universal input schema shared by all torch MLFF backends.

    Mirrors the JAX [AtomGraphInput][kups.potential.mliap.tojax.AtomGraphInput].
    Shapes use ``N`` atoms, ``B`` systems, and ``E`` edges (each padded by one
    extra atom/system to work around backends that cannot handle empty graphs).
    """

    pos: Array  # (N, 3)
    atomic_numbers: Array  # (N,)
    cell: Array  # (B, 3, 3)
    pbc: Array  # (B, 3)
    edge_index: Array  # (2, E)
    cell_offsets: Array  # (E, 3) integer multiples of cell vectors
    batch: Array  # (N,)
    charge: Array  # (B,)
    spin: Array  # (B,)

class IsTorchMliapParticles(IsRadiusGraphPoints, HasAtomicNumbers, Protocol):
    """Particle protocol for torch MLFF models."""

    ...

class MACEModule(torch.nn.Module):
    """Adapter: ``AtomGraphInput`` → MACE PyG-style input → energy + gradients.

    Wraps a MACE ``nn.Module`` and translates the universal graph input into
    the (``node_attrs``, ``positions``, ``edge_index``, ``batch``, ``ptr``,
    ``shifts``, ``cell``) tuple that MACE expects. Returns gradients of energy
    w.r.t. positions (and optionally cell vectors).

    Attributes:
        mace: Underlying MACE ``nn.Module``.
        species_to_index: Buffer mapping atomic number ``Z`` → MACE species
            index (0..``num_species``-1).
        num_species: Number of species the MACE model was trained on.
        compute_cell_gradients: Whether to compute cell gradients (stress).
    """

    species_to_index: torch.Tensor

    def __init__(
        self,
        mace_model: torch.nn.Module,
        species_to_index: torch.Tensor,
        num_species: int,
        compute_cell_gradients: bool = False,
    ) -> None:
        """Initialise ``MACEModule``.

        Args:
            mace_model: Underlying MACE ``nn.Module``.
            species_to_index: Tensor mapping ``Z`` → MACE index. Indexed by
                atomic number; entries for unsupported ``Z`` are ignored.
            num_species: Number of species the MACE model was trained on.
            compute_cell_gradients: Whether to compute cell gradients.
        """
        super().__init__()
        self.mace = mace_model
        self.mace.eval()
        self.register_buffer("species_to_index", species_to_index.to(dtype=torch.int64))
        self.num_species = num_species
        self.compute_cell_gradients = compute_cell_gradients

    @override
    def forward(self, input: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:  # type: ignore
        """Run MACE on a universal ``AtomGraphInput`` and return gradients.

        Args:
            input: Dict matching the universal ``AtomGraphInput`` schema.

        Returns:
            Dict with ``"energy"`` ``(B,)``, ``"position_gradients"`` ``(N, 3)``,
            and optionally ``"cell_gradients"`` ``(B, 3, 3)``.
        """
        # MACE is loaded at a fixed precision (float32 or float64 — see
        # ``load_mace(dtype=...)``); align every float input to that dtype
        # rather than whatever JAX hands us.
        model_dtype = next(self.mace.parameters()).dtype

        pos = input["pos"].to(model_dtype)
        species = input["atomic_numbers"]
        cell = input["cell"].to(model_dtype)
        batch = input["batch"]
        edge_index = input["edge_index"]
        cell_offsets = input["cell_offsets"].to(model_dtype)

        n_atoms = pos.shape[0]
        n_sys = cell.shape[0]

        # kUPS's neighbor list is a fixed-size buffer; ``indices_in`` maps
        # unused slots to the OOB sentinel ``len(keys) == n_atoms``. Drop
        # padded edges before indexing into atom-shaped tensors.
        valid_edge = (edge_index < n_atoms).all(dim=0)
        edge_index = edge_index[:, valid_edge]
        cell_offsets = cell_offsets[valid_edge]

        counts = torch.bincount(batch, minlength=n_sys)
        ptr = torch.cat([batch.new_zeros(1), counts.cumsum(0)])
        # ``species_to_index`` is registered as a CPU buffer at construction
        # time. ``TorchModuleWrapper`` calls us once with mock tensors on the
        # device of the wrapped MACE model (typically cuda) without first
        # calling ``self.to(device)``, so the indexing would straddle devices.
        # Co-locate the lookup table with the input on every call.
        species_to_index = self.species_to_index.to(species.device)
        node_attrs = F.one_hot(species_to_index[species], self.num_species).to(
            pos.dtype
        )
        cell_per_edge = cell[batch[edge_index[0]]]
        # cell_offsets (E,3) integer multiples → absolute Å via per-edge cell:
        # shifts[e, j] = Σ_i cell_offsets[e, i] * cell_per_edge[e, i, j]
        shifts = (cell_offsets.to(pos.dtype).unsqueeze(1) @ cell_per_edge).squeeze(1)

        out = self.mace(
            {
                "node_attrs": node_attrs,
                "positions": pos,
                "edge_index": edge_index,
                "batch": batch,
                "ptr": ptr,
                "shifts": shifts,
                # MACE's ``prepare_graph`` reads ``unit_shifts`` (the integer
                # cell-offset multiples) when ``compute_virials=True`` to build
                # the strain-perturbed graph in ``get_symmetric_displacement``.
                "unit_shifts": cell_offsets.to(pos.dtype),
                "cell": cell,
            },
            compute_force=True,
            compute_virials=self.compute_cell_gradients,
        )

        forces = out["forces"]
        result: dict[str, torch.Tensor] = {
            "energy": out["energy"].detach(),
            "position_gradients": (-forces).detach(),
        }
        if self.compute_cell_gradients:
            # MACE's ``virials`` = -sym(pos_virial + cell^T @ ∂E/∂h).
            # Negate to get the symmetric-strain virial, then invert.
            virial = -out["virials"]
            cell_grad = lattice_gradient_from_virial(
                forces=forces,
                positions=pos,
                batch=batch,
                cell=cell,
                virial=virial,
            )
            result["cell_gradients"] = cell_grad.detach()
        return result

@dataclass
class TorchMliap:
    """Container for a torch MLFF wired into JAX.

    Attributes:
        cutoff: Per-system cutoff radius [Å].
        wrapper: ``TorchModuleWrapper`` over the MLFF module.
        compute_cell_gradients: Whether the module returns ``"cell_gradients"``.
    """

    cutoff: Table[SystemId, Array]
    wrapper: TorchModuleWrapper = field(static=True)
    compute_cell_gradients: bool = field(static=True, default=False)

    @staticmethod
    def from_module(
        module: torch.nn.Module,
        cutoff: float,
        compute_cell_gradients: bool = False,
    ) -> "TorchMliap":
        """Wrap a torch.nn.Module that returns energy and gradients.

        Args:
            module: torch ``nn.Module`` satisfying ``TorchMliapForward``.
            cutoff: Interaction cutoff radius [Å].
            compute_cell_gradients: Whether the module returns
                ``"cell_gradients"`` for stress computation.

        Returns:
            Configured ``TorchMliap`` ready for use with the kUPS interface.
        """
        wrapper = TorchModuleWrapper(module, requires_grad=True)
        return TorchMliap(
            cutoff=Table((SystemId(0),), jnp.array([cutoff], float)),
            wrapper=wrapper,
            compute_cell_gradients=compute_cell_gradients,
        )

    def call(self, input: AtomGraphInput) -> dict[str, Array]:
        """Call the wrapped module on a prepared ``AtomGraphInput``."""
        return self.wrapper(input)

class TorchMliapForward(Protocol):
    """Forward contract for a torch MLFF module.

    The module must accept an ``AtomGraphInput`` dict and return a dict with:

    - ``"energy"``: ``(B,)`` per-system total energies.
    - ``"position_gradients"``: ``(N, 3)`` :math:`\\partial E / \\partial r`.
    - ``"cell_gradients"``: ``(B, 3, 3)`` :math:`\\partial E / \\partial h`,
      required only when ``compute_cell_gradients=True``.

    Outputs are gradients (not forces); adapters around models that natively
    produce forces/virials negate appropriately inside the module.
    """

    def __call__(self, input: AtomGraphInput) -> dict[str, Array]: ...

class UMAModule(torch.nn.Module):
    """Adapter: ``AtomGraphInput`` → fairchem ``AtomicData`` → energy + gradients.

    Wraps a fairchem ``MLIPPredictUnit`` and translates the universal graph
    input into the ``AtomicData`` object UMA expects. Returns gradients of
    energy w.r.t. positions (and optionally w.r.t. cell vectors).

    The wrapped predict-unit holds its own torch module and manages its own
    device placement; this adapter intentionally does not register it as a
    submodule (it is not an ``nn.Module``).

    Attributes:
        predict_unit: fairchem ``MLIPPredictUnit`` (held by reference).
        task_name: UMA inference head to route every system to.
        compute_cell_gradients: Whether to also return ``"cell_gradients"``.

    Note:
        UMA's ``stress`` is the symmetrized strain virial ``V_ij /
        volume`` from a joint symmetric strain on positions and cell
        (cf. ``fairchem.core.models.uma.outputs.compute_forces_and_stress``).
        We subtract the (fully known) position virial and apply the ``h^-T``
        factor to recover the raw lattice gradient ``∂E/∂h`` see
        ``lattice_gradient_from_virial``.
    """

    def __init__(
        self,
        predict_unit: Any,
        task_name: UMATaskName | str = "omat",
        compute_cell_gradients: bool = False,
    ) -> None:
        """Initialise ``UMAModule``.

        Args:
            predict_unit: fairchem ``MLIPPredictUnit`` (already loaded onto a
                device).
            task_name: UMA inference head (e.g. ``"omat"``, ``"omol"``).
            compute_cell_gradients: Whether to compute cell gradients (stress).
        """
        super().__init__()
        # PredictUnit is not an nn.Module; keep as plain attribute so
        # ``module.to(device)`` does not try to traverse it.
        self.predict_unit = predict_unit
        self.task_name = str(task_name)
        self.compute_cell_gradients = compute_cell_gradients

    @override
    def forward(self, input: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:  # type: ignore
        """Run UMA on a universal ``AtomGraphInput`` and return gradients.

        Args:
            input: Dict matching the universal ``AtomGraphInput`` schema.

        Returns:
            Dict with ``"energy"`` ``(B,)``, ``"position_gradients"`` ``(N, 3)``,
            and optionally ``"cell_gradients"`` ``(B, 3, 3)``.
        """
        from fairchem.core.datasets.atomic_data import (  # pyright: ignore[reportMissingImports]
            AtomicData,
        )

        # Keep AtomicData inputs at the device JAX/DLPack handed us. We must
        # NOT pre-move them to ``predict_unit.device``: the first call into
        # ``predict_unit.predict`` triggers ``_lazy_init`` which calls
        # ``prepare_for_inference(data, ...)`` while the underlying model is
        # still on cpu (``move_to_device`` runs after ``prepare_for_inference``
        # inside ``_lazy_init``). Pre-moving data to cuda makes UMA's MOLE
        # merge embed cuda indices into cpu weights and crash. The
        # predict-unit itself does ``data.to(self.device).clone()`` internally,
        # so any device-transition cost is paid there.
        #
        # We DO cast every float input to the predict-unit's inference dtype
        # (typically float32) before constructing AtomicData. Otherwise turbo
        # mode's MOLE merge — which runs inside ``_lazy_init.prepare_for_inference``,
        # *before* predict_unit's own dtype cast — pairs float64 data with
        # float32 model weights inside an einsum.
        inf_dtype = self.predict_unit.inference_settings.base_precision_dtype
        pos = input["pos"].to(dtype=inf_dtype)
        species = input["atomic_numbers"].to(dtype=torch.int64)
        cell = input["cell"].to(dtype=inf_dtype)
        batch = input["batch"].to(dtype=torch.int64).clone()
        edge_index = input["edge_index"].to(dtype=torch.int64)
        cell_offsets = -input["cell_offsets"].to(dtype=inf_dtype)
        pbc = input["pbc"]
        charge = input["charge"].to(dtype=inf_dtype)
        spin = input["spin"].to(dtype=inf_dtype)

        n_atoms = pos.shape[0]
        n_sys = cell.shape[0]

        # ``TorchModuleWrapper._get_output_info`` calls us once with all-zero
        # mock tensors to infer output shapes. UMA's turbo path merges its
        # MOLE experts on the *first* real ``predict()`` call (in
        # ``_lazy_init``) and then asserts every subsequent call has the same
        # composition. If we let the mock call through, the merge bakes in a
        # "100% H" composition that conflicts with the real structure. Detect
        # the mock by ``atomic_numbers.sum() == 0`` and return dummies of the
        # right shape/dtype without invoking ``predict_unit`` — that way the
        # first real call is what triggers ``_lazy_init`` with the real
        # composition.
        if int(species.abs().sum().item()) == 0:
            out_dtype = self.predict_unit.inference_settings.base_precision_dtype
            dev = pos.device
            result: dict[str, torch.Tensor] = {
                "energy": torch.zeros(n_sys, dtype=out_dtype, device=dev),
                "position_gradients": torch.zeros(
                    n_atoms, 3, dtype=out_dtype, device=dev
                ),
            }
            if self.compute_cell_gradients:
                result["cell_gradients"] = torch.zeros(
                    n_sys, 3, 3, dtype=out_dtype, device=dev
                )
            return result

        # kUPS's neighbor list is a fixed-size buffer; ``indices_in`` maps
        # unused slots to the OOB sentinel ``len(keys) == n_atoms``. Drop
        # those padded edges before handing the graph to UMA — otherwise the
        # species/position gather hits true out-of-bounds indices and CUDA
        # asserts.
        valid_edge = (edge_index < n_atoms).all(dim=0)
        edge_index = edge_index[:, valid_edge]
        cell_offsets = cell_offsets[valid_edge]

        # Pin the last atom to the last system so ``batch.max() + 1 == n_sys``
        # (validates AtomicData). For ``sorted_by_system`` real input this is
        # a no-op (the last atom is already in the last system).
        batch[-1] = n_sys - 1

        # First-real-call path: ``predict_unit._lazy_init`` runs
        # ``prepare_for_inference`` *before* moving the model to its target
        # device (``self.predict_unit.device``). If we pass cuda tensors at
        # that point, MOLE merge tries to gather cuda indices through cpu
        # weights and crashes. Move data to cpu only for this one call —
        # ``predict_unit.predict`` re-moves to its own device inside, so the
        # actual model forward still runs on cuda. From the second call
        # onward, data can stay on whatever device JAX handed us.
        if not getattr(self.predict_unit, "lazy_model_intialized", True):
            pos = pos.cpu()
            species = species.cpu()
            cell = cell.cpu()
            batch = batch.cpu()
            edge_index = edge_index.cpu()
            cell_offsets = cell_offsets.cpu()
            pbc = pbc.cpu()
            charge = charge.cpu()
            spin = spin.cpu()

        natoms = torch.bincount(batch, minlength=n_sys)
        edge_batch = batch[edge_index[0]]
        nedges = torch.bincount(edge_batch, minlength=n_sys)
        fixed = torch.zeros(n_atoms, dtype=torch.int64, device=pos.device)
        tags = torch.zeros(n_atoms, dtype=torch.int64, device=pos.device)

        data = AtomicData(
            pos=pos,
            atomic_numbers=species,
            cell=cell,
            pbc=pbc,
            natoms=natoms,
            edge_index=edge_index,
            cell_offsets=cell_offsets.to(pos.dtype),
            nedges=nedges,
            charge=charge,
            spin=spin,
            fixed=fixed,
            tags=tags,
            batch=batch,
            sid=[""] * n_sys,
            dataset=[self.task_name] * n_sys,
        )

        preds = self.predict_unit.predict(data)

        forces = preds["forces"]
        # ``predict_unit`` places outputs on its own device, which may differ
        # from the input device (cpu mock vs cuda predict-unit, single- vs
        # multi-gpu). Pin our post-processing tensors to the output device so
        # ``stress * volume`` and ``lattice_gradient_from_virial`` stay on a
        # single device.
        out_device = forces.device
        result: dict[str, torch.Tensor] = {
            "energy": preds["energy"].detach(),
            "position_gradients": (-forces).detach(),
        }
        if self.compute_cell_gradients:
            stress = preds["stress"]
            # UMA flattens stress to (B, 9); reshape to (B, 3, 3) if needed.
            if stress.dim() == 2 and stress.shape[-1] == 9:
                stress = stress.view(-1, 3, 3)
            cell_d = cell.to(out_device)
            volume = torch.linalg.det(cell_d).abs()
            virial = stress * volume.view(-1, 1, 1)
            cell_grad = lattice_gradient_from_virial(
                forces=forces,
                positions=pos.to(out_device),
                batch=batch.to(out_device),
                cell=cell_d,
                virial=virial,
            )
            result["cell_gradients"] = cell_grad.detach()
        return result