Merge pull request #38 from steinmig/general-cv-model

General CV model capability

Author: steinmig

nff/io/bias_calculators.py

@@ -23,7 +23,7 @@ class BiasBase(NeuralFF):
     """Basic Calculator class with neural force field
 
     Args:
-        model: the deural force field model
+        model: the neural force field model
         cv_def: list of Collective Variable (CV) definitions
             [["cv_type", [atom_indices], np.array([minimum, maximum]), bin_width], [possible second dimension]]
         equil_temp: float temperature of the simulation (important for extended system dynamics)
@@ -790,7 +790,7 @@ class WTMeABF(eABF):
             [["cv_type", [atom_indices], np.array([minimum, maximum]), bin_width], [possible second dimension]]
         equil_temp: float temperature of the simulation (important for extended system dynamics)
         dt: time step of the extended dynamics (has to be equal to that of the real system dyn!)
-        friction_per_ps: friction for the Lagevin dyn of extended system
+        friction_per_ps: friction for the Langevin dyn of extended system
         (has to be equal to that of the real system dyn!)
         nfull: numer of samples need for full application of bias force
         hill_height: unscaled height of the MetaD Gaussian hills in eV
@@ -958,7 +958,7 @@ class AttractiveBias(NeuralFF):
        Designed to be used with UQ as CV
 
     Args:
-        model: the deural force field model
+        model: the neural force field model
         cv_def: list of Collective Variable (CV) definitions
             [["cv_type", [atom_indices], np.array([minimum, maximum]), bin_width], [possible second dimension]]
         gamma: coupling strength, regulates strength of attraction

nff/md/colvars.py

@@ -43,8 +43,10 @@ class ColVar(torch.nn.Module):
         "projection_channelnormal",
         "Sp",
         "Sd",
-        "adjecencey_matrix",
+        "adjecencey_matrix",  # for backwards compatibility
+        "adjacency_matrix",
         "energy_gap",
+        "neural_cv"
     ]
 
     def __init__(self, info_dict: dict):
@@ -76,7 +78,7 @@ def __init__(self, info_dict: dict):
             self.ro = self.info_dict["acidhyd"]
             self.r1 = self.info_dict["waterhyd"]
 
-        elif self.info_dict["name"] == "adjecencey_matrix":
+        elif self.info_dict["name"] in ["adjecencey_matrix", "adjacency_matrix"]:
             self.model = self.info_dict["model"]
             self.device = self.info_dict["device"]
             self.bond_length = self.info_dict["bond_length"]
@@ -108,6 +110,22 @@ def __init__(self, info_dict: dict):
             self.model = self.model.to(self.device)
             self.model.eval()
 
+        elif self.info_dict["name"] == "neural_cv":
+            # expects a 'model' entry that behaves like a nn.Module and outputs a CV tensor
+            # a device for the model must be specified
+            # a 'descriptor_generation' entry can be specified for coordinate transformations / featurizations
+            # independent of the model by giving a Callable[[torch.Tensor], torch.Tensor]
+            # where the input tensor are the Cartesian coordinates of the atoms
+            # and the output tensor is the input tensor for the model
+            # if none is specified, the Cartesian coordinates are fed directly to the model
+            self.device = self.info_dict["device"]
+            self.descriptor_generation = self.info_dict.get("descriptor_generation")
+            if self.descriptor_generation is None:
+                self.descriptor_generation = lambda xyz_tensor: xyz_tensor
+            self.model = self.info_dict["model"]
+            self.model = self.model.to(self.device)
+            self.model.eval()
+
     def _get_com(self, indices: int | list[int]) -> torch.Tensor:
         """Get center of mass (com) of group of atoms
 
@@ -285,7 +303,7 @@ def minimal_distance(self, index_list: list[list[int]]) -> torch.Tensor:
 
     def projecting_centroidvec(self) -> torch.Tensor:
         """Projection of a position vector onto a reference vector
-        Atomic indices are used to determine the coordiantes of the vectors.
+        Atomic indices are used to determine the coordinates of the vectors.
         """
         vector_pos = self.xyz[self.vector_inds]
         vector = vector_pos[1] - vector_pos[0]
@@ -573,6 +591,14 @@ def forward(self, atoms: Atoms) -> tuple[np.ndarray, np.ndarray]:
         elif self.info_dict["name"] == "energy_gap":
             cv, cv_grad = self.energy_gap(self.info_dict["enkey_1"], self.info_dict["enkey_2"])
 
+        elif self.info_dict["name"] == "neural_cv":
+            desc = self.descriptor_generation(self.xyz)
+            cv = self.model(desc)
+            cv_grad, = torch.autograd.grad(cv.sum(), self.xyz, create_graph=False, retain_graph=False)
+
+        else:
+            raise RuntimeError(f"CV {self.info_dict['name']} not implemented!")
+
         return cv.detach().cpu().numpy(), cv_grad.detach().cpu().numpy()
 
 

tutorials/17_eABF_and_WTMeABF.ipynb

@@ -89,7 +89,9 @@
     "    \"index_list\": [0, 1, 2, 3],\n",
     "}\n",
     "\n",
-    "atoms = AtomsBatch.from_atoms(start_geom)\n",
+    "device = \"cpu\"\n",
+    "\n",
+    "atoms = AtomsBatch.from_atoms(start_geom, device=device)\n",
     "\n",
     "# use the distance of the scan for center of umbrella window\n",
     "CV = ColVar(info_dict)\n",
@@ -121,8 +123,6 @@
     }
    ],
    "source": [
-    "device = \"cpu\"\n",
-    "\n",
     "cv_defs = [\n",
     "    {\n",
     "        \"definition\": info_dict,\n",
@@ -419,7 +419,7 @@
    "outputs": [],
    "source": [
     "# Let's start exactly where the other simulation started\n",
-    "atoms = AtomsBatch.from_atoms(start_geom)\n",
+    "atoms = AtomsBatch.from_atoms(start_geom, device=device)\n",
     "\n",
     "# use the distance of the scan for center of umbrella window\n",
     "CV = ColVar(info_dict)\n",