fix bug with no seed set

lanl · May 9, 2024 · 341c5b7 · 341c5b7
1 parent bdf3173
commit 341c5b7
Showing 1 changed file with 11 additions and 3 deletions.
diff --git a/hippynn/molecular_dynamics/md.py b/hippynn/molecular_dynamics/md.py
@@ -305,7 +305,8 @@ def __init__(
         self.frix = frix
         self.kB = 0.001987204 * self.force_factor
 
-        torch.manual_seed(seed)
+        if seed is not None:
+            torch.manual_seed(seed)
 
     def _pre_step(self, dt):
         """Updates to variables performed during each step of MD simulation
@@ -316,6 +317,7 @@ def _pre_step(self, dt):
         dt : float
             timestep
         """
+
         self.variable.values["position"] = (
             self.variable.values["position"] + self.variable.values["velocity"] * dt
         )
@@ -335,14 +337,20 @@ def _post_step(self, dt, model_outputs):
             self.variable.device
         )
 
+        if len(self.variable.values["force"].shape) != len(
+            self.variable.values["mass"].shape
+        ):
+            self.variable.values["mass"] = self.variable.values["mass"][..., None]
+
         self.variable.values["acceleration"] = (
             self.variable.values["force"].detach()
             / self.variable.values["mass"]
             * self.force_factor
         )
 
-        self.variable.values["velocity"] = self.variable.values["velocity"] + (
-            +dt * self.variable.values["acceleration"]
+        self.variable.values["velocity"] = (
+            self.variable.values["velocity"]
+            + dt * self.variable.values["acceleration"]
             - self.frix * self.variable.values["velocity"] * dt
             + torch.sqrt(
                 2