From ec7b17222fe6f7ba7b16d7e1e6a3312f8cc65927 Mon Sep 17 00:00:00 2001
From: Ran Bi <biran@alum.mit.edu>
Date: Sun, 8 Dec 2024 21:28:56 -0800
Subject: [PATCH] feat(ai2bmd): option to exclude solvent in output

add option to control if solvent atoms should be written to output
trajectory file. default is 'True'.

fixes #32.
---
 scripts/ai2bmd        | 11 +++++++++++
 src/AIMD/arguments.py |  6 ++++++
 src/AIMD/simulator.py |  1 +
 src/utils/utils.py    |  9 +++++++--
 4 files changed, 25 insertions(+), 2 deletions(-)

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):