diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 492e299f..c9e4a820 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -34,7 +34,7 @@ jobs: uv pip install ".[test,docs]" --system - name: Install extras for tutorial generation - run: uv pip install ".[graphpes,mace,metatomic]" --system + run: uv pip install ".[mace,metatomic]" --system - name: Copy tutorials run: | diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 7f3402b8..56c2fc8b 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -61,7 +61,6 @@ jobs: - { python: '3.14', resolution: highest } model: - { name: fairchem, test_path: "tests/models/test_fairchem.py" } - - { name: graphpes, test_path: "tests/models/test_graphpes_framework.py" } - { name: mace, test_path: "tests/models/test_mace.py" } - { name: mace, test_path: "tests/test_elastic.py" } - { name: mace, test_path: "tests/test_optimizers_vs_ase.py" } diff --git a/README.md b/README.md index 70ffba22..de7ebd91 100644 --- a/README.md +++ b/README.md @@ -14,7 +14,7 @@ era. By rewriting the core primitives of atomistic simulation in Pytorch, it all orders of magnitude acceleration of popular machine learning potentials. * Automatic batching and GPU memory management allowing significant simulation speedup -* Support for MACE, Fairchem, SevenNet, ORB, MatterSim, graph-pes, metatomic, and Nequix MLIP models +* Support for MACE, Fairchem, SevenNet, ORB, MatterSim, metatomic, and Nequix MLIP models * Support for classical lennard jones, morse, and soft-sphere potentials * Molecular dynamics integration schemes like NVE, NVT Langevin, and NPT Langevin * Relaxation of atomic positions and cell with gradient descent and FIRE diff --git a/docs/conf.py b/docs/conf.py index b42066ac..b162d0c8 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -68,7 +68,6 @@ "metatomic", "orb", "sevennet", - "graphpes", ] # use type hints diff --git a/docs/tutorials/index.rst b/docs/tutorials/index.rst index 1e3fd150..7c236293 100644 --- a/docs/tutorials/index.rst +++ b/docs/tutorials/index.rst @@ -18,6 +18,5 @@ versions of the tutorials can also be found in the `torch-sim /examples/tutorial autobatching_tutorial low_level_tutorial hybrid_swap_tutorial - using_graphpes_tutorial metatomic_tutorial integrator_tests_analysis diff --git a/examples/tutorials/using_graphpes_tutorial.py b/examples/tutorials/using_graphpes_tutorial.py deleted file mode 100644 index b192407d..00000000 --- a/examples/tutorials/using_graphpes_tutorial.py +++ /dev/null @@ -1,89 +0,0 @@ -# %% -# /// script -# dependencies = [ -# "torch_sim_atomistic[graphpes]" -# ] -# /// - - -# %% [markdown] -""" -# Integrating TorchSim with `graph-pes` - -This brief tutorial demonstrates how to use models trained with the -[graph-pes](https://github.com/jla-gardner/graph-pes) package to drive -MD simulations and geometry optimizations in TorchSim. - -## Step 1: loading a model - -As an output of the `graph-pes-train` command, you receive a path -to a `.pt` file containing your trained model. To use this model -with TorchSim, pass the path to this `.pt` file, or the model itself, -to the `GraphPESWrapper` constructor. - -Below, we create a dummy TensorNet model with random weights as a demonstration: -""" - -# %% -from graph_pes.models import TensorNet, load_model - -# if you had a model saved to disk, you could load it like this: -# model = load_model("path/to/model.pt") - -# here, we just create a TensorNet model with random weights -model = TensorNet(cutoff=5.0) - -print(f"Number of parameters: {sum(p.numel() for p in model.parameters()):,}") - -# %% [markdown] -""" -## Step 2: wrapping the model for use with TorchSim - -We provide the `GraphPESWrapper` class to wrap a `graph-pes` model for use with TorchSim. -If you intend to drive simulations that require stresses, you will need to specify the -`compute_stress` argument to `True`. -""" - -# %% -from torch_sim.models.graphpes import GraphPESWrapper - -# wrap the model for use with TorchSim -ts_model = GraphPESWrapper(model, compute_stress=False) - -# or, alternatively, pass a model path directly: -# ts_model = GraphPESWrapper("path/to/model.pt", compute_stress=False) - -# %% [markdown] -""" -## Step 3: driving MD with the model - -Now that we have a model, we can drive MD simulations with it. For this, we will use the -`integrate` function. -""" - -# %% -from ase.build import molecule -import torch_sim as ts -from load_atoms import view - -# NVT at 300K -atoms = molecule("H2O") - -final_state = ts.integrate( - system=atoms, - model=ts_model, - integrator=ts.Integrator.nvt_langevin, - n_steps=50, - temperature=300, - timestep=0.001, -) - -final_atoms = final_state.to_atoms()[0] -view(final_atoms, show_bonds=True) - -# %% [markdown] -""" -Of course, this is a very simple example. However, you are now equipped to -use any `graph-pes` model that you have trained to drive any of the functionality -exposed by TorchSim! -""" diff --git a/pyproject.toml b/pyproject.toml index ffcd1d32..3645db77 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,7 +56,6 @@ mattersim = ["mattersim>=1.2.2"] metatomic = ["metatomic-torchsim>=0.1.1", "metatomic-ase>=0.1.0", "upet>=0.2.0"] orb = ["orb-models>=0.6.2"] sevenn = ["sevenn[torchsim]>=0.12.1"] -graphpes = ["graph-pes>=0.1", "mace-torch>=0.3.12"] nequip = ["nequip>=0.17.1"] nequix = ["nequix[torch-sim]>=0.4.5"] fairchem = ["fairchem-core>=2.7", "scipy<1.17.0"] @@ -152,14 +151,6 @@ markers = [ # make these dependencies mutually exclusive since they use incompatible e3nn versions # see https://docs.astral.sh/uv/concepts/projects/config/#conflicting-dependencies for more details conflicts = [ - [ - { extra = "fairchem" }, - { extra = "graphpes" }, - ], - [ - { extra = "fairchem" }, - { extra = "graphpes" }, - ], [ { extra = "fairchem" }, { extra = "mace" }, @@ -168,22 +159,6 @@ conflicts = [ { extra = "fairchem" }, { extra = "mace" }, ], - [ - { extra = "graphpes" }, - { extra = "mattersim" }, - ], - [ - { extra = "graphpes" }, - { extra = "nequip" }, - ], - [ - { extra = "graphpes" }, - { extra = "nequix" }, - ], - [ - { extra = "graphpes" }, - { extra = "sevenn" }, - ], [ { extra = "mace" }, { extra = "mattersim" }, diff --git a/tests/models/test_graphpes_framework.py b/tests/models/test_graphpes_framework.py deleted file mode 100644 index 7487914a..00000000 --- a/tests/models/test_graphpes_framework.py +++ /dev/null @@ -1,157 +0,0 @@ -import traceback - -import pytest -import torch -from ase.build import bulk, molecule -from ase.calculators.calculator import Calculator - -import torch_sim as ts -from tests.conftest import DEVICE -from tests.models.conftest import ( - make_model_calculator_consistency_test, - make_validate_model_outputs_test, -) -from torch_sim.testing import CONSISTENCY_SIMSTATES - - -try: - from graph_pes.atomic_graph import AtomicGraph, to_batch - from graph_pes.models import LennardJones, SchNet, TensorNet - - from torch_sim.models.graphpes_framework import GraphPESWrapper -except (ImportError, OSError, RuntimeError, AttributeError, ValueError): - pytest.skip( - f"graph-pes not installed: {traceback.format_exc()}", - allow_module_level=True, - ) - -DTYPE = torch.float32 - - -def test_graphpes_isolated(): - # test that the raw model and torch-sim wrapper give the same results - # for an isolated, unbatched structure - - water_atoms = molecule("H2O") - water_atoms.center(vacuum=10.0) - - gp_model = SchNet(cutoff=5.5) - gp_graph = AtomicGraph.from_ase(water_atoms, cutoff=5.5) - gp_energy = gp_model.predict_energy(gp_graph) - - ts_model = GraphPESWrapper( - gp_model, - device=DEVICE, - dtype=DTYPE, - compute_forces=True, - compute_stress=False, - ) - ts_output = ts_model(ts.io.atoms_to_state([water_atoms], DEVICE, DTYPE)) - assert set(ts_output) == {"energy", "forces"} - assert ts_output["energy"].shape == (1,) - - assert gp_energy.item() == pytest.approx(ts_output["energy"].item(), abs=1e-5) - - -def test_graphpes_periodic(): - # test that the raw model and torch-sim wrapper give the same results - # for a periodic, unbatched structure - - bulk_atoms = bulk("Al", "hcp", a=4.05) - assert bulk_atoms.pbc.all() - - gp_model = TensorNet(cutoff=5.5) - gp_graph = AtomicGraph.from_ase(bulk_atoms, cutoff=5.5) - gp_forces = gp_model.predict_forces(gp_graph) - - ts_model = GraphPESWrapper( - gp_model, - device=DEVICE, - dtype=DTYPE, - compute_forces=True, - compute_stress=True, - ) - ts_output = ts_model(ts.io.atoms_to_state([bulk_atoms], DEVICE, DTYPE)) - assert set(ts_output) == {"energy", "forces", "stress"} - assert ts_output["energy"].shape == (1,) - assert ts_output["forces"].shape == (len(bulk_atoms), 3) - assert ts_output["stress"].shape == (1, 3, 3) - - torch.testing.assert_close(ts_output["forces"].to("cpu"), gp_forces) - - -def test_batching(): - # test that the raw model and torch-sim wrapper give the same results - # when batching is done via torch-sim's atoms_to_state function - - water = molecule("H2O") - methane = molecule("CH4") - systems = [water, methane] - for s in systems: - s.center(vacuum=10.0) - - gp_model = SchNet(cutoff=5.5) - gp_graphs = [AtomicGraph.from_ase(s, cutoff=5.5) for s in systems] - - gp_energies = gp_model.predict_energy(to_batch(gp_graphs)) - - ts_model = GraphPESWrapper( - gp_model, - device=DEVICE, - dtype=DTYPE, - compute_forces=True, - compute_stress=True, - ) - ts_output = ts_model(ts.io.atoms_to_state(systems, DEVICE, DTYPE)) - - assert set(ts_output) == {"energy", "forces", "stress"} - assert ts_output["energy"].shape == (2,) - assert ts_output["forces"].shape == (sum(len(s) for s in systems), 3) - assert ts_output["stress"].shape == (2, 3, 3) - - assert gp_energies[0].item() == pytest.approx(ts_output["energy"][0].item(), abs=1e-5) - - -@pytest.mark.parametrize("dtype", [torch.float32, torch.float64]) -def test_graphpes_dtype(dtype: torch.dtype): - water = molecule("H2O") - - model = SchNet() - - ts_wrapper = GraphPESWrapper(model, device=DEVICE, dtype=dtype, compute_stress=False) - ts_output = ts_wrapper(ts.io.atoms_to_state([water], DEVICE, dtype)) - assert ts_output["energy"].dtype == dtype - assert ts_output["forces"].dtype == dtype - - -@pytest.fixture -def graphpes_lj_model() -> LennardJones: - return LennardJones(sigma=0.5) - - -@pytest.fixture -def ts_lj_model(graphpes_lj_model: LennardJones) -> GraphPESWrapper: - return GraphPESWrapper( - graphpes_lj_model, device=DEVICE, dtype=DTYPE, compute_stress=False - ) - - -@pytest.fixture -def ase_lj_calculator(graphpes_lj_model: LennardJones) -> Calculator: - return graphpes_lj_model.to(DEVICE, DTYPE).ase_calculator(skin=0.0) - - -test_graphpes_consistency = make_model_calculator_consistency_test( - test_name="graphpes-lj", - model_fixture_name="ts_lj_model", - calculator_fixture_name="ase_lj_calculator", - sim_state_names=CONSISTENCY_SIMSTATES, - device=DEVICE, - dtype=DTYPE, -) - -test_graphpes_model_outputs = make_validate_model_outputs_test( - model_fixture_name="ts_lj_model", - device=DEVICE, - dtype=DTYPE, -) diff --git a/torch_sim/models/graphpes.py b/torch_sim/models/graphpes.py deleted file mode 100644 index 737f36dc..00000000 --- a/torch_sim/models/graphpes.py +++ /dev/null @@ -1,16 +0,0 @@ -"""Deprecated module for importing GraphPESWrapper, AtomicGraph, and GraphPESModel. - -This module is deprecated. Please use the ts.models.graphpes_framework module instead. -""" - -import warnings - -from .graphpes_framework import AtomicGraph, GraphPESModel, GraphPESWrapper # noqa: F401 - - -warnings.warn( - "Importing from the ts.models.graphpes module is deprecated. " - "Please use the ts.models.graphpes_framework module instead.", - DeprecationWarning, - stacklevel=2, -) diff --git a/torch_sim/models/graphpes_framework.py b/torch_sim/models/graphpes_framework.py deleted file mode 100644 index 16c5d82c..00000000 --- a/torch_sim/models/graphpes_framework.py +++ /dev/null @@ -1,191 +0,0 @@ -"""An interface for using arbitrary GraphPESModels in ts. - -This module provides a TorchSim wrapper of the GraphPES models for computing -energies, forces, and stresses of atomistic systems. It serves as a wrapper around -the graph_pes library, integrating it with the torch-sim framework to enable seamless -simulation of atomistic systems with machine learning potentials. - -The GraphPESWrapper class adapts GraphPESModels to the ModelInterface protocol, -allowing them to be used within the broader torch-sim simulation framework. - -Notes: - This implementation requires graph_pes to be installed and accessible. - It supports various model configurations through model instances or model paths. -""" - -import traceback -import warnings -from pathlib import Path -from typing import Any - -import torch - -import torch_sim as ts -from torch_sim.models.interface import ModelInterface -from torch_sim.neighbors import torchsim_nl - - -try: - from graph_pes import AtomicGraph, GraphPESModel - from graph_pes.atomic_graph import PropertyKey - from graph_pes.models import load_model - -except ImportError as exc: - warnings.warn(f"GraphPES import failed: {traceback.format_exc()}", stacklevel=2) - PropertyKey = str - - class GraphPESWrapper(ModelInterface): - """GraphPESModel wrapper for torch-sim. - - This class is a placeholder for the GraphPESWrapper class. - It raises an ImportError if graph_pes is not installed. - """ - - def __init__(self, err: ImportError = exc, *_args: Any, **_kwargs: Any) -> None: - """Dummy init for type checking.""" - raise err - - def forward(self, *_args: Any, **_kwargs: Any) -> Any: - """Unreachable — __init__ always raises.""" - raise NotImplementedError - - class AtomicGraph: # noqa: D101 - def __init__(self, *args: Any, **kwargs: Any) -> None: # noqa: D107,ARG002 - raise ImportError("graph_pes must be installed to use this model.") - - class GraphPESModel(torch.nn.Module): # noqa: D101 - pass - - -def state_to_atomic_graph(state: ts.SimState, cutoff: torch.Tensor) -> AtomicGraph: - """Convert a SimState object into an AtomicGraph object. - - Args: - state: SimState object containing atomic positions, cell, and atomic numbers - cutoff: Cutoff radius for the neighbor list - - Returns: - AtomicGraph object representing the batched structures - """ - # graph-pes models internally trim the neighbor list to the - # model's cutoff value. To ensure no strange edge effects whereby - # edges that are exactly `cutoff` long are included/excluded, - # we bump cutoff + 1e-5 up slightly - nl, _system_mapping, shifts = torchsim_nl( - state.positions, - state.row_vector_cell, - state.pbc, - cutoff + 1e-5, - state.system_idx, - ) - n_atoms_per_system = torch.bincount(state.system_idx) - ptr = torch.zeros(state.n_systems + 1, dtype=torch.long, device=state.device) - ptr[1:] = n_atoms_per_system.cumsum(dim=0) - n_sys = state.n_systems - total_charge = torch.zeros(n_sys, device=state.device) - total_spin = torch.zeros(n_sys, device=state.device) - return AtomicGraph( - Z=state.atomic_numbers.long(), - R=state.positions, - cell=state.row_vector_cell, - neighbour_list=nl.long(), - neighbour_cell_offsets=shifts, - properties={}, - cutoff=cutoff.item(), - other={ - "total_charge": total_charge, - "total_spin": total_spin, - }, - batch=state.system_idx, - ptr=ptr, - ) - - -class GraphPESWrapper(ModelInterface): - """Wrapper for GraphPESModel in TorchSim. - - This class provides a TorchSim wrapper around GraphPESModel instances, - allowing them to be used within the broader torch-sim simulation framework. - - The graph-pes package allows for the training of existing model architectures, - including SchNet, PaiNN, MACE, NequIP, TensorNet, EDDP and more. - You can use any of these, as well as your own custom architectures, with this wrapper. - See the the graph-pes repo for more details: https://github.com/jla-gardner/graph-pes - - Args: - model: GraphPESModel instance, or a path to a model file - device: Device to run the model on - dtype: Data type for the model - compute_forces: Whether to compute forces - compute_stress: Whether to compute stress - - Example: - >>> from torch_sim.models.graphpes import GraphPESWrapper - >>> from graph_pes.models import load_model - >>> model = load_model("path/to/model.pt") - >>> wrapper = GraphPESWrapper(model) - >>> state = ts.SimState( - ... positions=torch.randn(10, 3), - ... cell=torch.eye(3), - ... atomic_numbers=torch.randint(1, 104, (10,)), - ... ) - >>> wrapper(state) - """ - - def __init__( - self, - model: GraphPESModel | str | Path, - device: torch.device | None = None, - dtype: torch.dtype = torch.float64, - *, - compute_forces: bool = True, - compute_stress: bool = True, - ) -> None: - """Initialize the GraphPESWrapper. - - Args: - model: GraphPESModel instance, or a path to a model file - device: Device to run the model on - dtype: Data type for the model - compute_forces: Whether to compute forces - compute_stress: Whether to compute stress - """ - super().__init__() - self._device = device or torch.device( - "cuda" if torch.cuda.is_available() else "cpu" - ) - self._dtype = dtype - - _model = model if isinstance(model, GraphPESModel) else load_model(model) - self._gp_model = _model.to(device=self.device, dtype=self.dtype) - - self._compute_forces = compute_forces - self._compute_stress = compute_stress - - self._properties: list[PropertyKey] = ["energy"] - if self.compute_forces: - self._properties.append("forces") - if self.compute_stress: - self._properties.append("stress") - - cutoff_val = self._gp_model.cutoff - if isinstance(cutoff_val, torch.Tensor) and cutoff_val.item() < 0.5: - self._memory_scales_with = "n_atoms" - - def forward(self, state: ts.SimState, **_kwargs: object) -> dict[str, torch.Tensor]: - """Forward pass for the GraphPESWrapper. - - Args: - state: SimState object containing atomic positions, cell, and atomic numbers - **_kwargs: Unused; accepted for interface compatibility. - - Returns: - Dictionary containing the computed energies, forces, and stresses - (where applicable) - """ - cutoff = self._gp_model.cutoff - if not isinstance(cutoff, torch.Tensor): - raise TypeError("GraphPES model cutoff must be a tensor") - atomic_graph = state_to_atomic_graph(state, cutoff) - preds = self._gp_model.predict(atomic_graph, self._properties) # ty: ignore[call-non-callable] - return {k: v.detach() for k, v in preds.items()}