diff --git a/scripts/ai2bmd b/scripts/ai2bmd index 8445f52..5b0ac62 100755 --- a/scripts/ai2bmd +++ b/scripts/ai2bmd @@ -115,6 +115,17 @@ def parse_args(argv): dest='solvent', help="Do not use solvent", ) + parser.add_argument( + "--write-solvent", + action='store_true', + help="Write coordinates of solvent atoms in output", + ) + parser.add_argument( + "--no-write-solvent", + action='store_false', + dest='write_solvent', + help="Do not write coordinates of solvent atoms in output", + ) parser.add_argument( "--preprocess-method", type=str, diff --git a/src/AIMD/arguments.py b/src/AIMD/arguments.py index ab94142..3cfb7d8 100644 --- a/src/AIMD/arguments.py +++ b/src/AIMD/arguments.py @@ -99,6 +99,12 @@ def init(argv=None): default=True, help="Use solvent or not", ) + parser.add_argument( + "--write-solvent", + action=argparse.BooleanOptionalAction, + default=True, + help="Write coordinates of solvent atoms in output", + ) parser.add_argument( "--preprocess-method", type=str, diff --git a/src/AIMD/simulator.py b/src/AIMD/simulator.py index 24989cf..1409899 100644 --- a/src/AIMD/simulator.py +++ b/src/AIMD/simulator.py @@ -124,6 +124,7 @@ def simulate( observer = MDObserver( a=self.prot, + q=self.qmatoms, md=MolDyn, traj=moldyn_traj, rng=rng_pool, diff --git a/src/utils/utils.py b/src/utils/utils.py index bbbb8d8..eae55c6 100644 --- a/src/utils/utils.py +++ b/src/utils/utils.py @@ -11,6 +11,8 @@ from ase.io.trajectory import TrajectoryWriter from ase.md.md import MolecularDynamics +from AIMD import arguments + def record_time(func): def wrapper(*args, **kwargs): @@ -115,8 +117,9 @@ class MDObserver: MolecularDynamics object, notified at every step. """ - def __init__(self, a: Atoms, md: MolecularDynamics, traj: TrajectoryWriter, rng: RNGPool, step_offset: int, temp_k: int): + def __init__(self, a: Atoms, q: Atoms, md: MolecularDynamics, traj: TrajectoryWriter, rng: RNGPool, step_offset: int, temp_k: int): self.a = a + self.q = q self.md = md self.traj = traj self.rng = rng @@ -124,11 +127,13 @@ def __init__(self, a: Atoms, md: MolecularDynamics, traj: TrajectoryWriter, rng: self.copy = None self.temp_k = temp_k + self.atoms = a if arguments.get().write_solvent else q + def get_md_step(self): return self.step_offset + self.md.nsteps def save_traj_copy(self): - self.copy = self.a.copy() + self.copy = self.atoms.copy() @delay_work def write_traj(self):