kups.potential.mliap.torch.mace ¶

def load_mace_wrapper(
    model_path: str | Path,
    device: str = "cuda",
    compute_force: bool = True,
    compute_virials: bool = False,
    dtype: Literal["float32", "float64"] = "float32",
) -> TorchModuleWrapper:
    """Load a PyTorch MACE model and wrap it for JAX computation.

    Args:
        model_path: Path to the MACE .model file
        device: Device to load the model onto (default: "cuda")
        compute_force: Whether to compute forces (default: True)
        compute_virials: Whether to compute virials for stress (default: False)
        dtype: Model precision - "float32" (default) or "float64"

    Returns:
        TorchModuleWrapper containing MACEModule.
    """
    if device == "cuda" and not torch.cuda.is_available():
        raise RuntimeError(
            "Device 'cuda' requested but CUDA is not available. "
            "Use device='cpu' or ensure CUDA is properly installed."
        )

    path = Path(model_path)
    if not path.exists():
        raise FileNotFoundError(f"MACE model not found: {model_path}")

    model = torch.load(path, weights_only=False, map_location=device)
    model.eval()

    if dtype == "float64":
        model = model.double()
    else:
        model = model.float()

    module = MACEModule(
        model, compute_force=compute_force, compute_virials=compute_virials
    )
    return TorchModuleWrapper(module, requires_grad=compute_force)

def make_torch_mace_potential[
    State,
    P: IsTorchMACEParticles,
    S: HasUnitCell,
    NNList: NearestNeighborList,
](
    particles_view: View[State, Table[ParticleId, P]],
    systems_view: View[State, Table[SystemId, S]],
    neighborlist_view: View[State, NNList],
    model: View[State, TorchMACEModel] | TorchMACEModel,
    cutoffs_view: View[State, Table[SystemId, Array]],
    compute_virials: bool = False,
    patch_idx_view: Any | None = None,
    out_cache_lens: Any | None = None,
) -> Any:
    """Create kUPS potential from PyTorch MACE model.

    Forces are computed by PyTorch natively (not JAX autodiff).
    Hessians are NOT supported (returns EmptyType).

    Args:
        particles_view: Extracts particle data from state
        systems_view: Extracts system data (unit cell) from state
        neighborlist_view: Extracts neighbor list from state
        model: TorchMACEModel instance or view to model in state
        cutoffs_view: Extracts cutoffs as ``Indexed[SystemId, Array]``
        compute_virials: Whether to compute virials for stress (default: False)
        patch_idx_view: Cached output index structure (optional)
        out_cache_lens: Cache location lens (optional)

    Returns:
        kUPS ``Potential`` backed by the PyTorch MACE model.
    """
    if isinstance(model, TorchMACEModel):
        model_view: View[State, TorchMACEModel] = constant(model)
    else:
        model_view = model

    if compute_virials:

        def virial_fn(inp):
            return torch_mace_model_fn(inp, compute_virials=True)

        return make_mliap_potential(
            model_fn=virial_fn,
            particles_view=particles_view,
            systems_view=systems_view,
            neighborlist_view=neighborlist_view,
            model_view=model_view,
            cutoffs_view=cutoffs_view,
            patch_idx_view=patch_idx_view,
            out_cache_lens=out_cache_lens,
        )

    def forces_fn(inp):
        return torch_mace_model_fn(inp, compute_virials=False)

    return make_mliap_potential(
        model_fn=forces_fn,
        particles_view=particles_view,
        systems_view=systems_view,
        neighborlist_view=neighborlist_view,
        model_view=model_view,
        cutoffs_view=cutoffs_view,
        patch_idx_view=patch_idx_view,
        out_cache_lens=out_cache_lens,
    )

def torch_mace_model_fn[
    P: IsTorchMACEParticles,
    S: HasUnitCell,
](
    inp: TorchMACEInput[P, S],
    *,
    compute_virials: bool = False,
) -> (
    WithPatch[PotentialOut[Array, EmptyType], IdPatch]
    | WithPatch[PotentialOut[PositionAndUnitCell, EmptyType], IdPatch]
):
    """Model function for PyTorch MACE models.

    Returns ``PotentialOut`` with forces (and optionally virials) computed by
    PyTorch.

    Args:
        inp: Graph potential input containing the MACE model and graph.
        compute_virials: Whether to include virial gradients in the output.

    Returns:
        ``WithPatch`` containing ``PotentialOut`` with energy, gradients, and
        an identity patch.
    """
    result = _call_mace_wrapper(inp)

    if compute_virials:
        assert result.virials is not None, "Model must have compute_virials=True"
        gradients = PositionAndUnitCell(
            positions=Table(inp.graph.particles.keys, -result.forces),
            unitcell=Table(
                inp.graph.systems.keys,
                TriclinicUnitCell.from_matrix(result.virials),
            ),
        )
        return WithPatch(
            PotentialOut(Table.arange(result.energy, label=SystemId), gradients, EMPTY),
            IdPatch(),
        )
    else:
        return WithPatch(
            PotentialOut(
                Table.arange(result.energy, label=SystemId), -result.forces, EMPTY
            ),
            IdPatch(),
        )