Skip to content

Commit

Permalink
update docstring format in molecular_dynamics/md.py for compatibility…
Browse files Browse the repository at this point in the history 
… with sphinx
  • Loading branch information
@shinkle-lanl
shinkle-lanl committed Jun 24, 2024
1 parent c26cdde commit a020912
Showing 1 changed file with 86 additions and 106 deletions.
192 changes: 86 additions & 106 deletions hippynn/molecular_dynamics/md.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,24 +29,22 @@ def __init__(
dtype: torch.dtype = None,
) -> None:
"""
Parameters
----------
name : str
name for variable
data : dict[str, torch.Tensor]
dictionary of tracked data in the form `value_name: value`
model_input_map : dict[str, str], optional
dictionary of correspondences between data tracked by Variable
:param name: name for variable
:type name: str
:param data: dictionary of tracked data in the form `value_name: value`
:type data: dict[str, torch.Tensor]
:param model_input_map: dictionary of correspondences between data tracked by Variable
and inputs to the HIP-NN model in the form
`hipnn-db_name: variable-data-key`, by default dict()
updater : _VariableUpdater, optional
object which will update the data of the Variable
over the course of the MD simulation, by default None
device : torch.device, optional
device on which to keep data, by default None
dtype : torch.dtype, optional
dtype for float type data, by default None
"""
`hipnn-db_name: variable-data-key`, defaults to dict()
: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
: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
:type dtype: torch.dtype, optional
"""
self.name = name
self.data = data
self.model_input_map = model_input_map
Expand Down Expand Up @@ -173,24 +171,20 @@ def pre_step(self, dt):
"""Updates to variables performed during each step of MD simulation
before HIPNN model evaluation
Parameters
----------
dt : float
timestep
"""
:param dt: timestep
:type dt: float
"""
pass

def post_step(self, dt, model_outputs):
"""Updates to variables performed during each step of MD simulation
after HIPNN model evaluation
Parameters
----------
dt : float
timestep
model_outputs : dict
dictionary of HIPNN model outputs
"""
:param dt: timestep
:type dt: float
:param model_outputs: dictionary of HIPNN model outputs
:type model_outputs: dict
"""
pass


Expand Down Expand Up @@ -219,32 +213,28 @@ def __init__(
units_acc: float = ase.units.Ang / (1.0**2),
):
"""
Parameters
----------
force_db_name : str
key which will correspond to the force on the corresponding Variable
:param force_db_name: key which will correspond to the force on the corresponding Variable
in the HIPNN model output dictionary
units_force : float, optional
amount of eV equal to one in the units used for force output
:type force_db_name: str
: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, units_force =
ase.units.kcal = 2.6114e22 since 2.6114e22 kcal = 1 eV),
by default ase.units.eV = 1
units_acc : float, optional
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
"""
by default ase.units.eV = 1, defaults to ase.units.eV
:type units_force: float, optional
: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)
:type units_acc: float, optional
"""
self.force_key = force_db_name
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
Parameters
----------
dt : float
timestep
"""
:param dt: timestep
:type dt: float
"""
self.variable.data["velocity"] = self.variable.data["velocity"] + 0.5 * dt * self.variable.data["acceleration"]
self.variable.data["position"] = self.variable.data["position"] + self.variable.data["velocity"] * dt
try:
Expand All @@ -256,13 +246,11 @@ def post_step(self, dt, model_outputs):
"""Updates to variables performed during each step of MD simulation
after HIPNN model evaluation
Parameters
----------
dt : float
timestep
model_outputs : dict
dictionary of HIPNN model outputs
"""
:param dt: timestep
:type dt: float
:param model_outputs: dictionary of HIPNN model outputs
:type model_outputs: dict
"""
self.variable.data["force"] = model_outputs[self.force_key].to(self.variable.device)
if len(self.variable.data["force"].shape) == len(self.variable.data["mass"].shape):
self.variable.data["acceleration"] = self.variable.data["force"].detach() / self.variable.data["mass"] * self.force_factor
Expand Down Expand Up @@ -290,26 +278,24 @@ def __init__(
seed: int = None,
):
"""
Parameters
----------
force_db_name : str
key which will correspond to the force on the corresponding Variable
:param force_db_name: key which will correspond to the force on the corresponding Variable
in the HIPNN model output dictionary
temperature : float
temperature for Langevin algorithm
frix : float
friction coefficient for Langevin algorithm
units_force : float, optional
amount of eV equal to one in the units used for force output
:type force_db_name: str
:param temperature: temperature for Langevin algorithm
:type temperature: float
:param frix: friction coefficient for Langevin algorithm
:type frix: float
: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, units_force =
ase.units.kcal = 2.6114e22 since 2.6114e22 kcal = 1 eV),
by default ase.units.eV = 1
units_acc : float, optional
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
seed : int, optional
used to set seed for reproducibility, by default None
"""
by default ase.units.eV = 1, defaults to ase.units.eV
:type units_force: float, optional
: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)
:type units_acc: float, optional
:param seed: used to set seed for reproducibility, defaults to None
:type seed: int, optional
"""

self.force_key = force_db_name
self.force_factor = units_force / units_acc
Expand All @@ -324,11 +310,9 @@ def pre_step(self, dt):
"""Updates to variables performed during each step of MD simulation
before HIPNN model evaluation
Parameters
----------
dt : float
timestep
"""
:param dt: timestep
:type dt: float
"""

self.variable.data["position"] = self.variable.data["position"] + self.variable.data["velocity"] * dt

Expand All @@ -341,13 +325,12 @@ def post_step(self, dt, model_outputs):
"""Updates to variables performed during each step of MD simulation
after HIPNN model evaluation
Parameters
----------
dt : float
timestep
model_outputs : dict
dictionary of HIPNN model outputs
"""
:param dt: timestep
:type dt: float
:param model_outputs: dictionary of HIPNN model outputs
:type model_outputs: dict
"""

self.variable.data["force"] = model_outputs[self.force_key].to(self.variable.device)

if len(self.variable.data["force"].shape) != len(self.variable.data["mass"].shape):
Expand Down Expand Up @@ -376,18 +359,17 @@ def __init__(
device: torch.device = None,
dtype: torch.dtype = None,
):
"""
Parameters
----------
variables : list[Variable]
list of Variable objects which will be tracked during simulation
model : Predictor
HIPNN Predictor
device : torch.device, optional
device to move variables and model to, by default None
dtype : torch.dtype, optional
dtype to convert all float type variable data and model parameters to, by default None
"""
"""_summary_
:param variables: list of Variable objects which will be tracked during simulation
:type variables: list[Variable]
:param model: HIPNN Predictor
:type model: Predictor
:param device: device to move variables and model to, defaults to None
:type device: torch.device, optional
:param dtype: dtype to convert all float type variable data and model parameters to, defaults to None
:type dtype: torch.dtype, optional
"""

self.variables = variables
self.model = model
Expand Down Expand Up @@ -513,19 +495,17 @@ def _update_data(self, model_outputs: dict):


def run(self, dt: float, n_steps: int, record_every: int = None):
"""
Run `n_steps` of MD algorithm.
Parameters
----------
dt : float
timestep
n_steps : int
number of steps to execute
record_every : int, optional
frequency at which to store the data at a step in memory,
record_every = 1 means every step will be stored, by default None
"""
"""Run `n_steps` of MD algorithm.
:param dt: timestep
:type dt: float
:param n_steps: number of steps to execute
:type n_steps: int
:param record_every: frequency at which to store the data at a step in memory,
record_every = 1 means every step will be stored, defaults to None
:type record_every: int, optional
"""

for i in trange(n_steps):
model_outputs = self._step(dt)
if record_every is not None and (i + 1) % record_every == 0:
Expand Down

0 comments on commit a020912

Please sign in to comment.