From e1c2ba99cd9beafc647770e6d7fb87ebcd37606b Mon Sep 17 00:00:00 2001 From: Emily Shinkle Date: Mon, 24 Jun 2024 12:01:47 -0600 Subject: [PATCH] update docstring format in molecular_dynamics/md.py for compatibility with sphinx (#79) * update docstring format in molecular_dynamics/md.py for compatibility with sphinx --- hippynn/molecular_dynamics/md.py | 192 ++++++++++++++----------------- 1 file changed, 86 insertions(+), 106 deletions(-) diff --git a/hippynn/molecular_dynamics/md.py b/hippynn/molecular_dynamics/md.py index fe7a06b1..72e06f5b 100644 --- a/hippynn/molecular_dynamics/md.py +++ b/hippynn/molecular_dynamics/md.py @@ -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 @@ -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 @@ -219,20 +213,18 @@ 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 @@ -240,11 +232,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["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: @@ -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 @@ -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 @@ -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 @@ -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): @@ -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 @@ -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: