Add option to save OpenMM system to XML file.

Author: lohedges

src/somd2/config/_config.py

@@ -165,6 +165,7 @@ def __init__(
         auto_fix_minimise=True,
         save_crash_report=False,
         save_energy_components=True,
+        save_xml=False,
         page_size=None,
         timeout="300 s",
     ):
@@ -507,6 +508,11 @@ def __init__(
         save_energy_components: bool
             Whether to save the energy contribution for each force when checkpointing.
 
+        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
             will automatically set the page size.
@@ -606,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
@@ -2418,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

@@ -1198,6 +1198,7 @@ def increment_filename(base_filename, suffix):
         )
         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")

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

src/somd2/runner/_runner.py

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