Enable MD workflows for any ASE calculator

nff/io/ase.py

@@ -1,5 +1,6 @@
 """ASE wrapper for the Neural Force Field."""
 
+from typing import List
 import copy
 
 import numpy as np
@@ -85,7 +86,6 @@ def convert_props_units(self, target_unit):
 
         self.props = const.convert_units(self.props, conversion_factor)
         self.props.update({"units": target_unit})
-        return
 
     def get_mol_nbrs(self, r_cut=95):
         """Dense directed neighbor list for each molecule, in case that's needed
@@ -248,6 +248,7 @@ def get_batch(self):
         if self.pbc.any():
             self.props["cell"] = torch.Tensor(np.array(self.cell))
             self.props["lattice"] = self.cell.tolist()
+            self.props["pbc"] = self.pbc.tolist()
 
         self.props["nxyz"] = torch.Tensor(self.get_nxyz())
         if self.props.get("num_atoms") is None:
@@ -259,9 +260,11 @@ def get_batch(self):
         if self.mol_idx is not None:
             self.props["mol_idx"] = self.mol_idx
 
+        self.props["device"] = self.device
+
         return self.props
 
-    def get_list_atoms(self):
+    def get_list_atoms(self) -> List[Atoms]:
         """Returns a list of ASE Atoms objects, each representing a molecule in the system.
 
         Returns:
@@ -285,7 +288,7 @@ def get_list_atoms(self):
                 cells = torch.split(torch.Tensor(self.props["lattice"]), 3)
             else:
                 cells = torch.unsqueeze(torch.Tensor(np.array(self.cell)), 0).repeat(len(mol_split_idx), 1, 1)
-        Atoms_list = []
+        atoms_list = []
 
         for i, molecule_xyz in enumerate(positions):
             atoms = Atoms(
@@ -299,9 +302,9 @@ def get_list_atoms(self):
             # of any of the atoms
             atoms.set_masses(masses[i])
 
-            Atoms_list.append(atoms)
+            atoms_list.append(atoms)
 
-        return Atoms_list
+        return atoms_list
 
     def update_num_atoms(self):
         """Update the number of atoms in the system.

nff/io/ase_calcs.py

@@ -7,11 +7,12 @@
 given geometry. Both calculators are used to perform molecular dynamics simulations
 and geometry optimizations using NFF AtomsBatch objects.
 """
+from __future__ import annotations
 
 import os
 import sys
 from collections import Counter
-from typing import List, Union
+from typing import List, Union, TYPE_CHECKING
 
 import numpy as np
 import torch
@@ -33,6 +34,8 @@
 from nff.utils.cuda import batch_detach, batch_to
 from nff.utils.geom import batch_compute_distance, compute_distances
 from nff.utils.scatter import compute_grad
+if TYPE_CHECKING:
+    from nff.io.potential import Potential
 
 HARTREE_TO_EV = HARTREE_TO_KCAL_MOL / EV_TO_KCAL_MOL
 
@@ -53,7 +56,7 @@ class NeuralFF(Calculator):
 
     def __init__(
         self,
-        model,
+        model: Potential,
         device="cpu",
         jobdir=None,
         en_key="energy",
@@ -134,15 +137,10 @@ def calculate(
         Calculator.calculate(self, atoms, self.properties, system_changes)
 
         # run model
-        # atomsbatch = AtomsBatch(atoms)
-        # batch_to(atomsbatch.get_batch(), self.device)
         batch = batch_to(atoms.get_batch(), self.device)
 
         # add keys so that the readout function can calculate these properties
-        # print("Properties: ", self.properties, "\n\n")
-        # print("en_key:", self.en_key)
         grad_key = self.en_key + "_grad"
-        # print("grad_key:", grad_key)
         batch[self.en_key] = []
         batch[grad_key] = []
 
@@ -158,7 +156,6 @@ def calculate(
             kwargs.update(self.model_kwargs)
 
         prediction = self.model(batch, **kwargs)
-        # print(prediction.keys())
 
         # change energy and force to numpy array
         conversion_factor: dict = const.conversion_factors.get((self.model_units, self.prediction_units), const.DEFAULT)
@@ -1159,9 +1156,7 @@ def calculate(
             system_changes (default from ase)
         """
 
-        # print("calculating ...")
         self.step += 1
-        # print("step ", self.step, self.step*0.0005)
         if not any(isinstance(self.model, i) for i in UNDIRECTED):
             check_directed(self.model, atoms)
 

nff/io/bias_calculators.py

@@ -47,7 +47,7 @@ class BiasBase(NeuralFF):
     def __init__(
         self,
         model,
-        cv_defs: list[dict],
+        cv_defs: List[dict],
         equil_temp: float = 300.0,
         device="cpu",
         en_key="energy",
@@ -290,6 +290,8 @@ def calculate(
         requires_stress = "stress" in self.properties
         if requires_stress:
             kwargs["requires_stress"] = True
+        if "forces" in self.properties:
+            kwargs["requires_forces"] = True
         if getattr(self, "model_kwargs", None) is not None:
             kwargs.update(self.model_kwargs)
 
@@ -298,8 +300,12 @@ def calculate(
         # change energy and force to numpy array and eV
         model_energy = prediction[self.en_key].detach().cpu().numpy() * (1 / const.EV_TO_KCAL_MOL)
 
-        if grad_key in prediction:
-            model_grad = prediction[grad_key].detach().cpu().numpy() * (1 / const.EV_TO_KCAL_MOL)
+        gradient = prediction.get(grad_key)
+        forces = prediction.get("forces")
+        if gradient is not None:
+            model_grad = gradient.detach().cpu().numpy() * (1 / const.EV_TO_KCAL_MOL)
+        elif forces is not None:
+            model_grad = - forces.detach().cpu().numpy() * (1 / const.EV_TO_KCAL_MOL)
         else:
             raise KeyError(grad_key)
 
@@ -389,7 +395,7 @@ class with neural force field
     def __init__(
         self,
         model,
-        cv_defs: list[dict],
+        cv_defs: List[dict],
         dt: float,
         friction_per_ps: float,
         equil_temp: float = 300.0,
@@ -570,7 +576,7 @@ class with neural force field
     def __init__(
         self,
         model,
-        cv_defs: list[dict],
+        cv_defs: List[dict],
         dt: float,
         friction_per_ps: float,
         amd_parameter: float,
@@ -801,7 +807,7 @@ class WTMeABF(eABF):
     def __init__(
         self,
         model,
-        cv_defs: list[dict],
+        cv_defs: List[dict],
         dt: float,
         friction_per_ps: float,
         equil_temp: float = 300.0,
@@ -982,7 +988,7 @@ class AttractiveBias(NeuralFF):
     def __init__(
         self,
         model,
-        cv_defs: list[dict],
+        cv_defs: List[dict],
         gamma=1.0,
         device="cpu",
         en_key="energy",

nff/io/potential.py

@@ -0,0 +1,71 @@
+from __future__ import annotations
+from copy import deepcopy
+from typing import Optional, Callable
+
+from ase.calculators.calculator import Calculator, all_changes
+from ase.symbols import Symbols
+import numpy as np
+import torch
+
+from nff.io.ase_calcs import AtomsBatch
+
+
+
+class Potential(torch.nn.Module):
+    pass
+
+
+class AsePotential(Potential):
+
+    def __init__(self, calculator: Calculator, embedding_fun: Optional[Callable[[AtomsBatch], torch.Tensor]] = None) \
+            -> None:
+        super().__init__()
+        self.calculator = calculator
+        self.embedding_fun = embedding_fun
+
+    def __call__(self, batch: dict, **kwargs):
+        properties = ["energy"]
+        if kwargs.get("requires_stress", False):
+            properties.append("stress")
+        if kwargs.get("requires_forces", False):
+            properties.append("forces")
+        if kwargs.get("requires_dipole", False):
+            properties.append("dipole")
+        if kwargs.get("requires_charges", False):
+            properties.append("charges")
+        if kwargs.get("requires_embedding", False):
+            if self.embedding_fun is None:
+                raise RuntimeError("Required embedding but no embedding function provided.")
+            embedding = self.embedding_fun(batch)
+        else:
+            embedding = None
+
+        nxyz = batch.get("nxyz")
+        if nxyz is None:
+            raise RuntimeError("Batch is missing 'nxyz' key.")
+        pbc = batch.get("pbc")
+        if pbc is not None:
+            pbc = np.array(pbc, dtype=bool)
+            cell = np.array(batch.get("cell")).reshape(3, 3)
+        else:
+            cell = None
+        atoms_batch = AtomsBatch(
+            symbols=Symbols(nxyz[:,0].detach().cpu().numpy()),
+            positions=nxyz[:, 1:4].detach().cpu().numpy(),
+            pbc=pbc,
+            cell=cell,
+            device=batch.get("device", "cpu")
+        )
+
+        self.calculator.calculate(atoms_batch, properties=properties, system_changes=all_changes)
+        results = deepcopy(self.calculator.results)
+        for key, value in results.items():
+            if isinstance(value, str):
+                continue
+            if not hasattr(value, "__iter__"):
+                results[key] = torch.tensor([value], device=atoms_batch.device)
+            else:
+                results[key] = torch.tensor(value, device=atoms_batch.device)
+        if embedding is not None:
+            results["embedding"] = embedding
+        return results

nff/nn/models/mace.py

@@ -265,6 +265,15 @@ def load_foundations(
             NffScaleMACE: NffScaleMACE foundational model.
         """
         mace_model_path = get_mace_mp_model_path(model)
+        return cls.load_foundations_path(mace_model_path, map_location, default_dtype)
+
+    @classmethod
+    def load_foundations_path(
+        cls,
+        mace_model_path: str,
+        map_location: str = "cpu",
+        default_dtype: Literal["", "float32", "float64"] = "float32",
+    ) -> NffScaleMACE:
         mace_model = torch.load(mace_model_path, map_location=map_location)
         init_params = get_init_kwargs_from_model(mace_model)
         model_dtype = get_model_dtype(mace_model)

nff/utils/constants.py

@@ -191,26 +191,18 @@
 
 
 ELEC_CONFIG = {
-    "1": [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-    "6": [2, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-    "7": [2, 2, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-    "8": [2, 2, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-    "9": [2, 2, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-    "11": [2, 2, 6, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-    "14": [2, 2, 6, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-    "16": [2, 2, 6, 2, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-    "17": [2, 2, 5, 2, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
-    "86": [2, 2, 6, 2, 6, 2, 10, 6, 2, 10, 6, 2, 14, 10, 6, 2, 6, 10, 4],
+    1: [1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    6: [2, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    7: [2, 2, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    8: [2, 2, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    9: [2, 2, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    11: [2, 2, 6, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    14: [2, 2, 6, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    16: [2, 2, 6, 2, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    17: [2, 2, 5, 2, 5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
+    86: [2, 2, 6, 2, 6, 2, 10, 6, 2, 10, 6, 2, 14, 10, 6, 2, 6, 10, 4],
 }
 
-ELEC_CONFIG = {int(key): val for key, val in ELEC_CONFIG.items()}
-
-
-# with open(os.path.join(os.path.dirname(os.path.abspath(__file__)),
-#                        "elec_config.json"), "r") as f:
-#     ELEC_CONFIG = json.load(f)
-#     ELEC_CONFIG = {int(key): val for key, val in ELEC_CONFIG.items()}
-
 
 def convert_units(props, conversion_dict):
     """Converts dictionary of properties to the desired units.