Merge pull request #142 from OpenBioSim/feature_force_groups

Feature force groups

Author: lohedges

src/somd2/config/_config.py

@@ -164,7 +164,8 @@ def __init__(
         pert_file=None,
         auto_fix_minimise=True,
         save_crash_report=False,
-        save_energy_components=False,
+        save_energy_components=True,
+        save_xml=False,
         page_size=None,
         timeout="300 s",
     ):
@@ -506,7 +507,11 @@ def __init__(
 
         save_energy_components: bool
             Whether to save the energy contribution for each force when checkpointing.
-            This is useful when debugging crashes.
+
+        save_xml: bool
+            Whether to write an XML file for the OpenMM system to the output
+            directory on startup. This can be useful for debugging or for
+            use with other tools that can read OpenMM XML files.
 
         page_size: int
             The page size for trajectory handling in megabytes. If None, then Sire
@@ -607,6 +612,7 @@ def __init__(
         self.auto_fix_minimise = auto_fix_minimise
         self.save_crash_report = save_crash_report
         self.save_energy_components = save_energy_components
+        self.save_xml = save_xml
         self.timeout = timeout
         self.num_energy_neighbours = num_energy_neighbours
         self.null_energy = null_energy
@@ -2419,6 +2425,16 @@ def save_energy_components(self, save_energy_components):
             raise ValueError("'save_energy_components' must be of type 'bool'")
         self._save_energy_components = save_energy_components
 
+    @property
+    def save_xml(self):
+        return self._save_xml
+
+    @save_xml.setter
+    def save_xml(self, save_xml):
+        if not isinstance(save_xml, bool):
+            raise ValueError("'save_xml' must be of type 'bool'")
+        self._save_xml = save_xml
+
     @property
     def page_size(self):
         return self._page_size

src/somd2/runner/_base.py

@@ -555,9 +555,6 @@ def __init__(self, system, config):
             self._config.checkpoint_frequency / self._config.energy_frequency
         )
 
-        # Zero the energy sample.
-        self._nrg_sample = 0
-
         # GCMC specific validation.
         if self._config.gcmc:
             if self._config.platform not in ["cuda", "opencl"]:
@@ -1197,10 +1194,11 @@ def increment_filename(base_filename, suffix):
         filenames["trajectory"] = str(output_directory / f"traj_{lam}.dcd")
         filenames["trajectory_chunk"] = str(output_directory / f"traj_{lam}_")
         filenames["energy_components"] = str(
-            output_directory / f"energy_components_{lam}.txt"
+            output_directory / f"energy_components_{lam}.csv"
         )
         filenames["gcmc_ghosts"] = str(output_directory / f"gcmc_ghosts_{lam}.txt")
         filenames["sampler_stats"] = str(output_directory / f"sampler_stats_{lam}.pkl")
+        filenames["xml"] = str(output_directory / f"system_{lam}.xml")
         if restart:
             filenames["config"] = str(
                 output_directory / increment_filename("config", "yaml")
@@ -2024,9 +2022,10 @@ def _backup_checkpoint(self, index):
 
         return index, None
 
-    def _save_energy_components(self, index, context):
+    def _save_energy_components(self, index, context, time_ns):
         """
-        Internal function to save the energy components for each force group to file.
+        Internal function to save the energy components for each force group to a
+        CSV file.
 
         Parameters
         ----------
@@ -2036,44 +2035,38 @@ def _save_energy_components(self, index, context):
 
         context : openmm.Context
             The current OpenMM context.
+
+        time_ns : float
+            The current simulation time in nanoseconds.
         """
 
-        from copy import deepcopy
+        import csv as _csv
         import openmm
 
-        # Get the current context and system.
-        system = deepcopy(context.getSystem())
-
-        # Add each force to a unique group.
-        for i, f in enumerate(system.getForces()):
-            f.setForceGroup(i)
-
-        # Create a new context.
-        new_context = openmm.Context(system, deepcopy(context.getIntegrator()))
-        new_context.setPositions(context.getState(getPositions=True).getPositions())
-
-        header = f"{'# Sample':>10}"
-        record = f"{self._nrg_sample:>10}"
-
-        # Process the records.
-        for i, f in enumerate(system.getForces()):
-            state = new_context.getState(getEnergy=True, groups={i})
-            name = f.getName()
-            name_len = len(name)
-            header += f"{f.getName():>{name_len + 2}}"
-            record += f"{state.getPotentialEnergy().value_in_unit(openmm.unit.kilocalories_per_mole):>{name_len + 2}.2f}"
-
-        # Write to file.
-        if self._nrg_sample == 0:
-            with open(self._filenames[index]["energy_components"], "w") as f:
-                f.write(header + "\n")
-                f.write(record + "\n")
-        else:
-            with open(self._filenames[index]["energy_components"], "a") as f:
-                f.write(record + "\n")
+        filepath = self._filenames[index]["energy_components"]
+        file_exists = _Path(filepath).exists()
+
+        # Use the named force groups already assigned by sire_to_openmm_system,
+        # sorted alphabetically for a consistent column order across runs.
+        energies = {}
+        for name, grp in sorted(context._force_group_map.items()):
+            state = context.getState(getEnergy=True, groups=(1 << grp))
+            energies[name] = state.getPotentialEnergy().value_in_unit(
+                openmm.unit.kilocalories_per_mole
+            )
+
+        columns = ["time"] + list(energies.keys())
+        row = {"time": round(time_ns, 6)} | {
+            name: round(nrg, 4) for name, nrg in energies.items()
+        }
 
-        # Increment the sample number.
-        self._nrg_sample += 1
+        with open(filepath, "a", newline="") as f:
+            writer = _csv.DictWriter(f, fieldnames=columns)
+            if not file_exists:
+                # Write a comment line with units before the header.
+                f.write("# time: ns, energy: kcal/mol\n")
+                writer.writeheader()
+            writer.writerow(row)
 
     def _restore_backup_files(self):
         """

src/somd2/runner/_repex.py

@@ -53,6 +53,7 @@ def __init__(
         gcmc_kwargs=None,
         output_directory=None,
         perturbed_system=None,
+        xml_filenames=None,
     ):
         """
         Constructor.
@@ -84,6 +85,10 @@ def __init__(
         perturbed_system: :class: `System <sire.system.System>`
             The perturbed end-state system used to seed starting coordinates for
             lambda > 0.5 replicas. If None, the perturbed state is not used.
+
+        xml_filenames: list of str
+            A list of file paths for the OpenMM XML output, one per replica.
+            If None, XML files are not written.
         """
 
         # Warn if the number of replicas is not a multiple of the number of GPUs.
@@ -117,6 +122,7 @@ def __init__(
             gcmc_kwargs=gcmc_kwargs,
             output_directory=output_directory,
             perturbed_system=perturbed_system,
+            xml_filenames=xml_filenames,
         )
 
     def __setstate__(self, state):
@@ -168,6 +174,7 @@ def _create_dynamics(
         gcmc_kwargs=None,
         output_directory=None,
         perturbed_system=None,
+        xml_filenames=None,
     ):
         """
         Create the dynamics objects.
@@ -199,6 +206,10 @@ def _create_dynamics(
         perturbed_system: :class: `System <sire.system.System>`
             The perturbed end-state system used to seed starting coordinates for
             lambda > 0.5 replicas. If None, the perturbed state is not used.
+
+        xml_filenames: list of str
+            A list of file paths for the OpenMM XML output, one per replica.
+            If None, XML files are not written.
         """
 
         from math import floor
@@ -315,6 +326,13 @@ def _create_dynamics(
             # Append the dynamics object.
             self._dynamics.append(dynamics)
 
+            # Write the OpenMM XML file to the output directory.
+            if xml_filenames is not None:
+                _logger.info(
+                    f"Writing OpenMM XML for lambda {lam:.5f} on device {device}"
+                )
+                dynamics.to_xml(xml_filenames[i])
+
             # Track memory footprint for this device.
             info = device_mem[device]
             info["count"] += 1
@@ -740,6 +758,11 @@ def __init__(self, system, config):
 
         # Create the dynamics cache.
         if not self._is_restart:
+            xml_filenames = (
+                [self._filenames[i]["xml"] for i in range(len(self._lambda_values))]
+                if self._config.save_xml
+                else None
+            )
             self._dynamics_cache = DynamicsCache(
                 self._system,
                 self._lambda_values,
@@ -749,6 +772,7 @@ def __init__(self, system, config):
                 gcmc_kwargs=self._gcmc_kwargs,
                 perturbed_system=self._perturbed_system,
                 output_directory=self._config.output_directory,
+                xml_filenames=xml_filenames,
             )
         else:
             _logger.debug("Restarting from file")
@@ -1397,11 +1421,6 @@ def _run_block(
             energies = dynamics._current_energy_array()
 
         except Exception as e:
-            try:
-                # Save the energy components for debugging purposes.
-                self._save_energy_components(index, dynamics.context())
-            except:
-                pass
             return False, index, e
 
         # Return the index and the energies.
@@ -1647,11 +1666,6 @@ def _equilibrate(self, index):
             )
 
         except Exception as e:
-            try:
-                # Save the energy components for debugging purposes.
-                self._save_energy_components(index, dynamics.context())
-            except:
-                pass
             return False, index, e
 
         return True, index, None
@@ -1767,6 +1781,12 @@ def _checkpoint(self, index, lambdas, block, num_blocks, is_final_block=False):
             # Commit the current system.
             system = dynamics.commit()
 
+            # Save the energy contribution for each force.
+            if self._config.save_energy_components:
+                self._save_energy_components(
+                    index, dynamics.context(), system.time().to("ns")
+                )
+
             # If performing GCMC, then we need to flag the ghost waters.
             if gcmc_sampler is not None:
                 system = gcmc_sampler._flag_ghost_waters(system)

src/somd2/runner/_runner.py

@@ -592,10 +592,6 @@ def generate_lam_vals(lambda_base, increment=0.001):
                     system.set_time(_sr.u("0ps"))
 
             except Exception as e:
-                try:
-                    self._save_energy_components(index, dynamics.context())
-                except:
-                    pass
                 raise RuntimeError(f"Equilibration failed: {e}")
 
         # Work out the lambda values for finite-difference gradient analysis.
@@ -645,6 +641,11 @@ def generate_lam_vals(lambda_base, increment=0.001):
         # Create the dynamics object.
         dynamics = system.dynamics(**dynamics_kwargs)
 
+        # Write the OpenMM XML file to the output directory.
+        if self._config.save_xml and not is_restart:
+            _logger.info(f"Writing OpenMM XML for {_lam_sym} = {lambda_value:.5f}")
+            dynamics.to_xml(self._filenames[index]["xml"])
+
         # Reset the GCMC sampler. This resets the sampling statistics and clears
         # the associated OpenMM forces.
         if gcmc_sampler is not None:
@@ -914,23 +915,21 @@ def generate_lam_vals(lambda_base, increment=0.001):
                             save_crash_report=self._config.save_crash_report,
                         )
                 except Exception as e:
-                    try:
-                        self._save_energy_components(index, dynamics.context())
-                    except:
-                        pass
                     raise RuntimeError(
                         f"Dynamics block {block + 1} for {_lam_sym} = {lambda_value:.5f} failed: {e}"
                     )
 
                 # Checkpoint.
                 try:
-                    # Save the energy contribution for each force.
-                    if self._config.save_energy_components:
-                        self._save_energy_components(index, dynamics.context())
-
                     # Commit the current system.
                     system = dynamics.commit()
 
+                    # Save the energy contribution for each force.
+                    if self._config.save_energy_components:
+                        self._save_energy_components(
+                            index, dynamics.context(), system.time().to("ns")
+                        )
+
                     # If performing GCMC, then we need to flag the ghost waters.
                     if gcmc_sampler is not None:
                         system = gcmc_sampler._flag_ghost_waters(system)
@@ -1056,13 +1055,15 @@ def generate_lam_vals(lambda_base, increment=0.001):
                         save_crash_report=self._config.save_crash_report,
                     )
 
-                    # Save the energy contribution for each force.
-                    if self._config.save_energy_components:
-                        self._save_energy_components(index, dynamics.context())
-
                     # Commit the current system.
                     system = dynamics.commit()
 
+                    # Save the energy contribution for each force.
+                    if self._config.save_energy_components:
+                        self._save_energy_components(
+                            index, dynamics.context(), system.time().to("ns")
+                        )
+
                     # Record the end time.
                     block_end = _timer()
 
@@ -1101,10 +1102,6 @@ def generate_lam_vals(lambda_base, increment=0.001):
                         f"{_lam_sym} = {lambda_value:.5f} complete, speed = {speed:.2f} ns day-1"
                     )
                 except Exception as e:
-                    try:
-                        self._save_energy_components(index, dynamics.context())
-                    except:
-                        pass
                     raise RuntimeError(
                         f"Final dynamics block for {lam_sym} = {lambda_value:.5f} failed: {e}"
                     )
@@ -1232,10 +1229,6 @@ def generate_lam_vals(lambda_base, increment=0.001):
                         save_crash_report=self._config.save_crash_report,
                     )
             except Exception as e:
-                try:
-                    self._save_energy_components(index, dynamics.context())
-                except:
-                    pass
                 raise RuntimeError(
                     f"Dynamics for {_lam_sym} = {lambda_value:.5f} failed: {e}"
                 )