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

@@ -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,

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,

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,

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,20 +117,23 @@ 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
         self.step_offset = step_offset
         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):