add aceff-2.0

doc/userguide.md

@@ -198,6 +198,7 @@ Pretrained and your own local models are supported.
 | ---- | ----- |
 | `aceff-1.0` | Pretrained [AceFF-1.0](https://huggingface.co/Acellera/AceFF-1.0) [TensorNet](https://doi.org/10.48550/arXiv.2306.06482) model. |
 | `aceff-1.1` | Pretrained [AceFF-1.1](https://huggingface.co/Acellera/AceFF-1.0) [TensorNet](https://doi.org/10.48550/arXiv.2306.06482) model. |
+| `aceff-2.0` | Pretrained [AceFF-2.0](https://huggingface.co/Acellera/AceFF-2.0) [TensorNet2](https://arxiv.org/abs/2601.00581) model. |
 | `torchmdnet` |  Custom TorchMD-Net models specified with the `modelPath` argument. |
 
 
@@ -229,4 +230,4 @@ OpenMM-ML is based on a plugin architecture, allowing other packages to provide
 packages listed above are the ones for which OpenMM-ML has built in support.  Other packages can interface to it by
 defining two classes that subclass `MLPotentialImpl` and `MLPotentialImplFactory`, then registering them by specifying
 an [entry point](https://packaging.python.org/en/latest/specifications/entry-points/) in the group `openmmml.potentials`.
-Consult the documentation for other packages to see whether they provide interfaces for OpenMM-ML.
+Consult the documentation for other packages to see whether they provide interfaces for OpenMM-ML.

openmmml/models/torchmdnetpotential.py

@@ -73,10 +73,12 @@ class TorchMDNetPotentialImpl(MLPotentialImpl):
 
     Pretained AceFF models can be used directly:
 
-    >>> potential = MLPotential('aceff-1.0')
+    >>> potential = MLPotential('aceff-2.0')
 
     >>> potential = MLPotential('aceff-1.1')
 
+    >>> potential = MLPotential('aceff-1.0')
+
     """
 
     def __init__(self, 
@@ -144,6 +146,9 @@ def addForces(self,
             elif self.name == 'aceff-1.1':
                 repo_id="Acellera/AceFF-1.1"
                 filename="aceff_v1.1.ckpt"
+            elif self.name == 'aceff-2.0':
+                repo_id="Acellera/AceFF-2.0"
+                filename="aceff_v2.0.ckpt"
             else:
                 raise ValueError(f'Model name {self.name} does not exist.')
 
@@ -170,7 +175,11 @@ def addForces(self,
             batch = torch.tensor(batch, dtype=torch.long)
 
         # TensorNet models can use CUDA graphs and the default is to use them.
-        use_cudagraphs = args.get('cudaGraphs', True if isinstance(model.representation_model, torchmdnet.models.tensornet.TensorNet) else False)
+        use_cudagraphs = args.get('cudaGraphs', 
+                                True if (isinstance(model.representation_model, torchmdnet.models.tensornet.TensorNet)
+                                or isinstance(model.representation_model, torchmdnet.models.tensornet2.TensorNet2))
+                                else False
+                            )
 
         class TorchMDNetForce(torch.nn.Module):
             def __init__(self, model, numbers, charge, atoms, batch, lengthScale, energyScale):

setup.py

@@ -57,6 +57,7 @@
             'torchmdnet = openmmml.models.torchmdnetpotential:TorchMDNetPotentialImplFactory',
             'aceff-1.0 = openmmml.models.torchmdnetpotential:TorchMDNetPotentialImplFactory',
             'aceff-1.1 = openmmml.models.torchmdnetpotential:TorchMDNetPotentialImplFactory',
+            'aceff-2.0 = openmmml.models.torchmdnetpotential:TorchMDNetPotentialImplFactory',
         ]
     }
 )