Need to index the energy components file name.

src/somd2/runner/_base.py

@@ -1163,13 +1163,16 @@ def _checkpoint(
                     df.iloc[-self._energy_per_block :],
                 )
 
-    def _save_energy_components(self, context):
+    def _save_energy_components(self, index, context):
         """
         Internal function to save the energy components for each force group to file.
 
         Parameters
         ----------
 
+        index : int
+            The index of the window or replica.
+
         context : openmm.Context
             The current OpenMM context.
         """
@@ -1199,11 +1202,11 @@ def _save_energy_components(self, context):
 
         # Write to file.
         if self._nrg_sample == 0:
-            with open(self._filenames["energy_components"], "w") as f:
+            with open(self._filenames[index]["energy_components"], "w") as f:
                 f.write(header + "\n")
                 f.write(record + "\n")
         else:
-            with open(self._filenames["energy_components"], "a") as f:
+            with open(self._filenames[index]["energy_components"], "a") as f:
                 f.write(record + "\n")
 
         # Increment the sample number.

src/somd2/runner/_runner.py

@@ -488,7 +488,7 @@ def generate_lam_vals(lambda_base, increment=0.001):
                 try:
                     # Save the energy contribution for each force.
                     if self._config.save_energy_components:
-                        self._save_energy_components(dynamics.context())
+                        self._save_energy_components(index, dynamics.context())
 
                     # Commit the current system.
                     system = dynamics.commit()
@@ -539,7 +539,7 @@ def generate_lam_vals(lambda_base, increment=0.001):
 
                     # Save the energy contribution for each force.
                     if self._config.save_energy_components:
-                        self._save_energy_components(dynamics.context())
+                        self._save_energy_components(index, dynamics.context())
 
                     # Commit the current system.
                     system = dynamics.commit()