From 1d0495d552ae78f8aca5e80b7d3df9e50a26325d Mon Sep 17 00:00:00 2001
From: Emily Shinkle <eshinkle@lanl.gov>
Date: Fri, 18 Oct 2024 14:43:50 -0600
Subject: [PATCH] clarify language on units for custom MD algorithms, remove
 default units

---
 hippynn/molecular_dynamics/__init__.py |  4 +--
 hippynn/molecular_dynamics/md.py       | 38 +++++++++-----------------
 2 files changed, 15 insertions(+), 27 deletions(-)

diff --git a/hippynn/molecular_dynamics/__init__.py b/hippynn/molecular_dynamics/__init__.py
index 66e8feb6..5dd3f3b6 100644
--- a/hippynn/molecular_dynamics/__init__.py
+++ b/hippynn/molecular_dynamics/__init__.py
@@ -7,7 +7,7 @@
 
 """
 
-from .md import MolecularDynamics, Variable, NullUpdater, VelocityVerlet, LangevinDynamics
+from .md import MolecularDynamics, Variable, NullUpdater, VelocityVerlet, LangevinDynamics, VariableUpdater
 
 
-__all__ = ["MolecularDynamics", "Variable", "NullUpdater", "VelocityVerlet", "LangevinDynamics"]
+__all__ = ["MolecularDynamics", "Variable", "NullUpdater", "VelocityVerlet", "LangevinDynamics", "VariableUpdater"]
diff --git a/hippynn/molecular_dynamics/md.py b/hippynn/molecular_dynamics/md.py
index edb76b8c..46274dd5 100644
--- a/hippynn/molecular_dynamics/md.py
+++ b/hippynn/molecular_dynamics/md.py
@@ -196,23 +196,17 @@ class VelocityVerlet(VariableUpdater):
     def __init__(
         self,
         force_db_name: str,
-        units_force: float = None,
-        units_acc: float = None,
+        units_force: float,
+        units_acc: float,
     ):
         """
         :param force_db_name: key which will correspond to the force on the corresponding Variable
             in the HIPNN model output dictionary
-        :param units_force: amount of eV equal to one in the units used for force output
-            of HIPNN model (eg. if force output in kcal/mol/A, units_force =
-            ase.units.kcal/ase.units.mol/ase.units.Ang ~=.0434 since .0434 kcal ~= 1 eV),
-            by default ase.units.eV = 1, defaults to ase.units.eV / ase.units.Ang
-        :param units_acc: amount of Ang/fs^2 equal to one in the units used for acceleration
-            in the corresponding Variable, by default units.Ang/(1.0 ** 2) = 1, defaults to ase.units.Ang/(1.0**2)
+        :param units_force: force units via ase.units (eg. ase.units.kcal/ase.units.mol/ase.units.Ang),
+        :param units_acc: acceleration units via ase.units (eg. ase.units.Ang/(ase.units.ps**2)),
+
+        mass of attached Variable must be in amu
         """
-        if units_force is None:
-            units_force = ase.units.eV / ase.units.Ang
-        if units_acc is None:
-            units_acc = ase.units.Ang / (1.0**2)
         self.force_key = force_db_name
         self.force_factor = units_force / units_acc
 
@@ -256,8 +250,8 @@ def __init__(
         force_db_name: str,
         temperature: float,
         frix: float,
-        units_force: float = None,
-        units_acc: float = None,
+        units_force: float,
+        units_acc: float,
         seed: Optional[int] = None,
     ):
         """
@@ -265,19 +259,12 @@ def __init__(
             in the HIPNN model output dictionary
         :param temperature: temperature for Langevin algorithm
         :param frix: friction coefficient for Langevin algorithm
-        :param units_force: amount of eV equal to one in the units used for force output
-            of HIPNN model (eg. if force output in kcal/mol/A, units_force =
-            ase.units.kcal/ase.units.mol/ase.units.Ang ~=.0434 since .0434 kcal ~= 1 eV),
-            by default ase.units.eV = 1, defaults to ase.units.eV / ase.units.Ang
-        :param units_acc: amount of Ang/fs^2 equal to one in the units used for acceleration
-            in the corresponding Variable, by default units.Ang/(1.0 ** 2) = 1, defaults to ase.units.Ang/(1.0**2)
+        :param units_force: force units via ase.units (eg. ase.units.kcal/ase.units.mol/ase.units.Ang),
+        :param units_acc: acceleration units via ase.units (eg. ase.units.Ang/(ase.units.ps**2)),
         :param seed: used to set seed for reproducibility, defaults to None
-        """
-        if units_force is None:
-            units_force = ase.units.eV / ase.units.Ang
-        if units_acc is None:
-            units_acc = ase.units.Ang / (1.0**2)
 
+        mass of attached Variable must be in amu
+        """
         self.force_key = force_db_name
         self.force_factor = units_force / units_acc
         self.temperature = temperature
@@ -301,6 +288,7 @@ def pre_step(self, dt:float):
                 self.variable.data["unwrapped_position"] = self.variable.data["unwrapped_position"] + self.variable.data["velocity"] * dt
             except KeyError:
                 self.variable.data["unwrapped_position"] = copy(self.variable.data["position"])
+                
     def post_step(self, dt: float, model_outputs: dict):
         """
         Updates to variables performed during each step of MD simulation after HIPNN model evaluation