Skip to content

Commit

Permalink
Documentation fixes and README.md update for MACE (#72)
Browse files Browse the repository at this point in the history 
* Fixed documentation

* Updated README to include reference to MACE
  • Loading branch information
@JMorado
JMorado committed Mar 5, 2024
1 parent 4cc85f4 commit 67feb6b
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 13 deletions.
4 changes: 2 additions & 2 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,8 @@
This is a high level API for using machine learning models in OpenMM simulations. With just a few lines of code, you
can set up a simulation that uses a standard, pretrained model to represent some or all of the interactions in a system.

In the current release, the only supported potential functions are [ANI-1ccx](https://www.nature.com/articles/s41467-019-10827-4)
and [ANI-2x](https://pubs.acs.org/doi/full/10.1021/acs.jctc.0c00121), using the implementations in [TorchANI](https://github.com/aiqm/torchani). They are
In the current release, the supported potential functions are [ANI-1ccx](https://www.nature.com/articles/s41467-019-10827-4)
and [ANI-2x](https://pubs.acs.org/doi/full/10.1021/acs.jctc.0c00121), using the implementations in [TorchANI](https://github.com/aiqm/torchani), as well as [MACE](https://arxiv.org/abs/2206.07697) models, including the pre-trained [MACE-OFF23 models](https://arxiv.org/abs/2312.15211), utilizing the [MACE](https://github.com/ACEsuit/mace) implementation. They are
suitable for small molecules involving a limited set of elements and no charges. Future releases will add new potential
functions that support a much wider range of molecules.

Expand Down
26 changes: 15 additions & 11 deletions openmmml/models/macepotential.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -50,35 +50,39 @@ class MACEPotentialImpl(MLPotentialImpl):
This implementation supports both MACE-OFF23 and locally trained MACE models.
To use one of the pre-trained MACE-OFF23 models, specify the model name. For example:
>>> potential = MLPotential('mace-off23-small')
Other available MACE-OFF23 models include 'mace-off23-medium' and 'mace-off23-large'.
To use a locally trained MACE model, provide the path to the model file. For example:
>>> potential = MLPotential('mace', modelPath='MACE.model')
During system creation, you can optionally specify the precision of the model using the
`precision` keyword argument. Supported options are 'single' and 'double'. For example:
``precision`` keyword argument. Supported options are 'single' and 'double'. For example:
>>> system = potential.createSystem(topology, precision='single')
By default, the implementation uses the precision of the loaded MACE model.
According to the MACE documentation, 'single' precision is recommended for MD (faster but
less accurate), while 'double' precision is recommended for geometry optimization.
Additionally, you can request computation of the full atomic energy, including the atom
self-energy, instead of the default interaction energy, by setting `returnEnergyType` to
self-energy, instead of the default interaction energy, by setting ``returnEnergyType`` to
'energy'. For example:
>>> system = potential.createSystem(topology, returnEnergyType='energy')
The default is to compute the interaction energy, which can be made explicit by setting
`returnEnergyType='interaction_energy'`.
``returnEnergyType='interaction_energy'``.
Attributes
----------
name : str
The name of the MACE model.
modelPath : str
The path to the locally trained MACE model if `name` is 'mace'.
The path to the locally trained MACE model if ``name`` is 'mace'.
"""

def __init__(self, name: str, modelPath) -> None:
Expand All @@ -91,7 +95,7 @@ def __init__(self, name: str, modelPath) -> None:
The name of the MACE model.
Options include 'mace-off23-small', 'mace-off23-medium', 'mace-off23-large', and 'mace'.
modelPath : str, optional
The path to the locally trained MACE model if `name` is 'mace'.
The path to the locally trained MACE model if ``name`` is 'mace'.
"""
self.name = name
self.modelPath = modelPath
Expand All @@ -116,12 +120,12 @@ def addForces(
system : openmm.System
The system to which the force will be added.
atoms : iterable of int
The indices of the atoms to include in the model. If `None`, all atoms are included.
The indices of the atoms to include in the model. If ``None``, all atoms are included.
forceGroup : int
The force group to which the force should be assigned.
precision : str, optional
The precision of the model. Supported options are 'single' and 'double'.
If `None`, the default precision of the model is used.
If ``None``, the default precision of the model is used.
returnEnergyType : str, optional
Whether to return the interaction energy or the energy including the self-energy.
Default is 'interaction_energy'. Supported options are 'interaction_energy' and 'energy'.
Expand Down Expand Up @@ -252,7 +256,7 @@ def __init__(
nodeAttrs : torch.Tensor
The one-hot encoded atomic numbers.
atoms : iterable of int
The indices of the atoms. If `None`, all atoms are included.
The indices of the atoms. If ``None``, all atoms are included.
periodic : bool
Whether the system is periodic.
dtype : torch.dtype
Expand Down Expand Up @@ -294,9 +298,9 @@ def _getNeighborPairs(
Notes
-----
This method calculates the shifts and edge indices by determining neighbor pairs (`neighbors`)
and respective wrapped distances (`wrappedDeltas`) using `NNPOps.neighbors.getNeighborPairs`.
After obtaining the `neighbors` and `wrappedDeltas`, the pairs with negative indices (r>cutoff)
This method calculates the shifts and edge indices by determining neighbor pairs (``neighbors``)
and respective wrapped distances (``wrappedDeltas``) using ``NNPOps.neighbors.getNeighborPairs``.
After obtaining the ``neighbors`` and ``wrappedDeltas``, the pairs with negative indices (r>cutoff)
are filtered out, and the edge indices and shifts are finally calculated.
Parameters
Expand Down

0 comments on commit 67feb6b

Please sign in to comment.