Fix multi integrator

src/graphrelax/relaxer.py

@@ -449,9 +449,10 @@ def get_energy_breakdown(self, pdb_string: str) -> dict:
                 # Set this force to group i
                 force.setForceGroup(i)
 
-            # Recreate simulation with force groups
+            # Recreate simulation with force groups - add new integrator
+            integrator2 = openmm.LangevinIntegrator(0, 0.01, 0.0)
             simulation = openmm_app.Simulation(
-                pdb.topology, system, integrator, platform
+                pdb.topology, system, integrator2, platform
             )
             simulation.context.setPositions(pdb.positions)
 

tests/test_relaxer_integration.py

@@ -130,6 +130,23 @@ def test_get_energy_breakdown_has_total(
 
         assert "total_energy" in result
 
+    def test_get_energy_breakdown_computes_components(
+        self, relaxer, small_peptide_pdb_string
+    ):
+        """Test that energy breakdown computes all components."""
+        result = relaxer.get_energy_breakdown(small_peptide_pdb_string)
+
+        assert "bond_energy" in result or "HarmonicBondForce" in result
+        assert "angle_energy" in result or "HarmonicAngleForce" in result
+        assert "nonbonded_energy" in result or "NonbondedForce" in result
+
+        assert result["total_energy"] != 0.0
+
+        # test more than just total_energy
+        assert (
+            len(result) > 1
+        ), "Energy breakdown should have individual components"
+
 
 @pytest.mark.integration
 class TestRelaxerConfig: