From b0dc259094e5248e4ef251c4f0216d86c3969096 Mon Sep 17 00:00:00 2001 From: Emily Shinkle Date: Mon, 24 Jun 2024 13:05:21 -0600 Subject: [PATCH] update progress bar in molecular_dynamics/md.py for consistency, fix spacing, fix bug in examples/molecular_dynamics.py --- examples/molecular_dynamics.py | 4 ++-- hippynn/layers/physics.py | 3 +-- hippynn/molecular_dynamics/md.py | 8 +++----- 3 files changed, 6 insertions(+), 9 deletions(-) diff --git a/examples/molecular_dynamics.py b/examples/molecular_dynamics.py index 3ab4d6f1..18a38028 100644 --- a/examples/molecular_dynamics.py +++ b/examples/molecular_dynamics.py @@ -148,8 +148,8 @@ def update(self, diff_steps, data): def print(self, diff_steps=None, data=None): time_per_atom_step = self.update(diff_steps, data) """Function to print the potential, kinetic and total energy""" - atoms.set_positions(data["position_position"][-1]) - atoms.set_velocities(data["position_velocity"][-1]) + atoms.set_positions(np.array(data["position_position"][-1])) + atoms.set_velocities(np.array(data["position_velocity"][-1])) print( "Performance:", round(1e6 * time_per_atom_step, 1), diff --git a/hippynn/layers/physics.py b/hippynn/layers/physics.py index cc9e457e..81e9b2ac 100644 --- a/hippynn/layers/physics.py +++ b/hippynn/layers/physics.py @@ -205,8 +205,7 @@ def forward(self, pair_dist, radius): class LocalDampingCosine(AlphaScreening): """ Local damping using complement of the hipnn cutoff function. ('glue-on' method) - g = 1 if pair_dist > R_cutoff - 1 - [cos(pi/2 * dist * R_cutoff)]^2 otherwise + g = 1 if pair_dist > R_cutoff, 1 - [cos(pi/2 * dist * R_cutoff)]^2 otherwise """ def __init__(self, alpha): """ diff --git a/hippynn/molecular_dynamics/md.py b/hippynn/molecular_dynamics/md.py index 72e06f5b..86857c9a 100644 --- a/hippynn/molecular_dynamics/md.py +++ b/hippynn/molecular_dynamics/md.py @@ -3,10 +3,9 @@ import numpy as np import torch - -from tqdm.autonotebook import trange import ase +from ..tools import progress_bar from ..graphs import Predictor from ..layers.pairs.periodic import wrap_systems_torch @@ -359,8 +358,7 @@ def __init__( device: torch.device = None, dtype: torch.dtype = None, ): - """_summary_ - + """ :param variables: list of Variable objects which will be tracked during simulation :type variables: list[Variable] :param model: HIPNN Predictor @@ -506,7 +504,7 @@ def run(self, dt: float, n_steps: int, record_every: int = None): :type record_every: int, optional """ - for i in trange(n_steps): + for i in progress_bar(range(n_steps)): model_outputs = self._step(dt) if record_every is not None and (i + 1) % record_every == 0: self._update_data(model_outputs)