Make MD output settings optional

Allow for conventional MD output writing to be skipped if set as None in the settings.
OpenFreeEnergy · Dec 18, 2024 · 3f7122a · 3f7122a
2 parents fe10c39 + 0ced170
commit 3f7122a
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 31 deletions.
diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py
@@ -341,9 +341,10 @@ def _run_MD(simulation: openmm.app.Simulation,
             mdtraj_top.subset(selection_indices),
         )
 
-        traj.save_pdb(
-            shared_basepath / output_settings.minimized_structure
-        )
+        if output_settings.minimized_structure:
+            traj.save_pdb(
+                shared_basepath / output_settings.minimized_structure
+            )
         # equilibrate
         # NVT equilibration
         if equil_steps_nvt:
@@ -431,26 +432,33 @@ def _run_MD(simulation: openmm.app.Simulation,
             mc_steps=1,
         )
 
-        simulation.reporters.append(XTCReporter(
-            file=str(shared_basepath / output_settings.production_trajectory_filename),
-            reportInterval=write_interval,
-            atomSubset=selection_indices))
-        simulation.reporters.append(openmm.app.CheckpointReporter(
-            file=str(shared_basepath / output_settings.checkpoint_storage_filename),
-            reportInterval=checkpoint_interval))
-        simulation.reporters.append(openmm.app.StateDataReporter(
-            str(shared_basepath / output_settings.log_output),
-            checkpoint_interval,
-            step=True,
-            time=True,
-            potentialEnergy=True,
-            kineticEnergy=True,
-            totalEnergy=True,
-            temperature=True,
-            volume=True,
-            density=True,
-            speed=True,
-        ))
+        if output_settings.production_trajectory_filename:
+            simulation.reporters.append(XTCReporter(
+                file=str(
+                    shared_basepath /
+                    output_settings.production_trajectory_filename),
+                reportInterval=write_interval,
+                atomSubset=selection_indices))
+        if output_settings.checkpoint_storage_filename:
+            simulation.reporters.append(openmm.app.CheckpointReporter(
+                file=str(
+                    shared_basepath /
+                    output_settings.checkpoint_storage_filename),
+                reportInterval=checkpoint_interval))
+        if output_settings.log_output:
+            simulation.reporters.append(openmm.app.StateDataReporter(
+                str(shared_basepath / output_settings.log_output),
+                checkpoint_interval,
+                step=True,
+                time=True,
+                potentialEnergy=True,
+                kineticEnergy=True,
+                totalEnergy=True,
+                temperature=True,
+                volume=True,
+                density=True,
+                speed=True,
+            ))
         t0 = time.time()
         simulation.step(prod_steps)
         t1 = time.time()
@@ -616,10 +624,13 @@ def run(self, *, dry=False, verbose=True,
             )
 
         # f. Save pdb of entire system
-        with open(shared_basepath / output_settings.preminimized_structure, "w") as f:
-            openmm.app.PDBFile.writeFile(
-                stateA_topology, stateA_positions, file=f, keepIds=True
-            )
+        if output_settings.preminimized_structure:
+            with open(
+                    shared_basepath /
+                    output_settings.preminimized_structure, "w") as f:
+                openmm.app.PDBFile.writeFile(
+                    stateA_topology, stateA_positions, file=f, keepIds=True
+                )
 
         # 10. Get platform
         restrict_cpu = forcefield_settings.nonbonded_method.lower() == 'nocutoff'

diff --git a/openfe/protocols/openmm_utils/omm_settings.py b/openfe/protocols/openmm_utils/omm_settings.py
@@ -637,16 +637,16 @@ class Config:
         arbitrary_types_allowed = True
 
     # reporter settings
-    production_trajectory_filename = 'simulation.xtc'
+    production_trajectory_filename: Optional[str] = 'simulation.xtc'
     """Path to the storage file for analysis. Default 'simulation.xtc'."""
     trajectory_write_interval: FloatQuantity['picosecond'] = 20 * unit.picosecond
     """
     Frequency to write the xtc file. Default 5000 * unit.timestep.
     """
-    preminimized_structure = 'system.pdb'
+    preminimized_structure: Optional[str] = 'system.pdb'
     """Path to the pdb file of the full pre-minimized system. 
     Default 'system.pdb'."""
-    minimized_structure = 'minimized.pdb'
+    minimized_structure: Optional[str] = 'minimized.pdb'
     """Path to the pdb file of the system after minimization. 
     Only the specified atom subset is saved. Default 'minimized.pdb'."""
     equil_nvt_structure: Optional[str] = 'equil_nvt.pdb'
@@ -655,7 +655,7 @@ class Config:
     equil_npt_structure: Optional[str] = 'equil_npt.pdb'
     """Path to the pdb file of the system after NPT equilibration. 
     Only the specified atom subset is saved. Default 'equil_npt.pdb'."""
-    log_output = 'simulation.log'
+    log_output: Optional[str] = 'simulation.log'
     """
     Filename for writing the log of the MD simulation, including timesteps,
     energies, density, etc.