diff --git a/hippynn/molecular_dynamics/__init__.py b/hippynn/molecular_dynamics/__init__.py index d54a269b..3bcc1722 100644 --- a/hippynn/molecular_dynamics/__init__.py +++ b/hippynn/molecular_dynamics/__init__.py @@ -2,4 +2,4 @@ Molecular dynamics driver with great flexibility and customizability regarding which quantities which are evolved and what algorithms are used to evolve them. Calls a hippynn `Predictor` on current state during each MD step. """ -from .md import Variable, NullUpdater, VelocityVerlet, LangevinDynamics, MolecularDynamics \ No newline at end of file +from .md import * \ No newline at end of file diff --git a/hippynn/molecular_dynamics/md.py b/hippynn/molecular_dynamics/md.py index 86857c9a..b16cc552 100644 --- a/hippynn/molecular_dynamics/md.py +++ b/hippynn/molecular_dynamics/md.py @@ -23,7 +23,7 @@ def __init__( name: str, data: dict[str, torch.Tensor], model_input_map: dict[str, str] = dict(), - updater: _VariableUpdater = None, + updater: VariableUpdater = None, device: torch.device = None, dtype: torch.dtype = None, ) -> None: @@ -38,7 +38,7 @@ def __init__( :type model_input_map: dict[str, str], optional :param updater: object which will update the data of the Variable over the course of the MD simulation, defaults to None - :type updater: _VariableUpdater, optional + :type updater: VariableUpdater, optional :param device: device on which to keep data, defaults to None :type device: torch.device, optional :param dtype: dtype for float type data, defaults to None @@ -136,14 +136,11 @@ def _(self, arg: torch.dtype): self.dtype = arg -class _VariableUpdater: +class VariableUpdater: """ - Parent class for algorithms that make updates to the data of a Variable during - each step of an MD simulation. + Parent class for algorithms that make updates to the data of a Variable during each step of an MD simulation. - Subclasses should redefine __init__, pre_step, post_step, and - required_variable_data as needed. The inputs to pre_step and post_step - should not be changed. + Subclasses should redefine __init__, pre_step, post_step, and required_variable_data as needed. The inputs to pre_step and post_step should not be changed. """ # A list of keys which must appear in Variable.data for any Variable that will be updated by objects of this class. @@ -167,8 +164,7 @@ def variable(self, variable): self._variable = variable def pre_step(self, dt): - """Updates to variables performed during each step of MD simulation - before HIPNN model evaluation + """Updates to variables performed during each step of MD simulation before HIPNN model evaluation :param dt: timestep :type dt: float @@ -176,8 +172,7 @@ def pre_step(self, dt): pass def post_step(self, dt, model_outputs): - """Updates to variables performed during each step of MD simulation - after HIPNN model evaluation + """Updates to variables performed during each step of MD simulation after HIPNN model evaluation :param dt: timestep :type dt: float @@ -187,7 +182,7 @@ def post_step(self, dt, model_outputs): pass -class NullUpdater(_VariableUpdater): +class NullUpdater(VariableUpdater): """ Makes no change to the variable data at each step of MD. """ @@ -198,7 +193,7 @@ def pre_step(self, dt): def post_step(self, dt, model_outputs): pass -class VelocityVerlet(_VariableUpdater): +class VelocityVerlet(VariableUpdater): """ Implements the Velocity Verlet algorithm """ @@ -228,8 +223,7 @@ def __init__( self.force_factor = units_force / units_acc def pre_step(self, dt): - """Updates to variables performed during each step of MD simulation - before HIPNN model evaluation + """Updates to variables performed during each step of MD simulation before HIPNN model evaluation :param dt: timestep :type dt: float @@ -242,8 +236,7 @@ def pre_step(self, dt): pass def post_step(self, dt, model_outputs): - """Updates to variables performed during each step of MD simulation - after HIPNN model evaluation + """Updates to variables performed during each step of MD simulation after HIPNN model evaluation :param dt: timestep :type dt: float @@ -260,7 +253,7 @@ def post_step(self, dt, model_outputs): self.variable.data["velocity"] = self.variable.data["velocity"] + 0.5 * dt * self.variable.data["acceleration"] -class LangevinDynamics(_VariableUpdater): +class LangevinDynamics(VariableUpdater): """ Implements the Langevin algorithm """ @@ -306,8 +299,7 @@ def __init__( torch.manual_seed(seed) def pre_step(self, dt): - """Updates to variables performed during each step of MD simulation - before HIPNN model evaluation + """Updates to variables performed during each step of MD simulation before HIPNN model evaluation :param dt: timestep :type dt: float @@ -321,8 +313,7 @@ def pre_step(self, dt): pass def post_step(self, dt, model_outputs): - """Updates to variables performed during each step of MD simulation - after HIPNN model evaluation + """Updates to variables performed during each step of MD simulation after HIPNN model evaluation :param dt: timestep :type dt: float @@ -386,7 +377,7 @@ def variables(self, variables): variables = [variables] for variable in variables: if variable.updater is None: - raise ValueError(f"Variable with name {variable.name} does not have a _VariableUpdater set.") + raise ValueError(f"Variable with name {variable.name} does not have a VariableUpdater set.") variable_names = [variable.name for variable in variables] if len(variable_names) != len(set(variable_names)):