From c003b2f6a9f5b24c61c81328beb603ce24538814 Mon Sep 17 00:00:00 2001
From: alexwhan <wha022@csiro.au>
Date: Fri, 23 Jan 2026 13:12:11 +1100
Subject: [PATCH 1/2] add second integrator in get_energy_breakdown

---
 src/graphrelax/relaxer.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/graphrelax/relaxer.py b/src/graphrelax/relaxer.py
index 3407ef3..e6e50fb 100644
--- a/src/graphrelax/relaxer.py
+++ b/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)
 

From dd912557b0b8093c92f1abbde0efa9b1658ec3a6 Mon Sep 17 00:00:00 2001
From: alexwhan <wha022@csiro.au>
Date: Fri, 23 Jan 2026 13:51:30 +1100
Subject: [PATCH 2/2] added test for whether get_energy_breakdown computes
 components

---
 tests/test_relaxer_integration.py | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/tests/test_relaxer_integration.py b/tests/test_relaxer_integration.py
index d62d92d..72f97c6 100644
--- a/tests/test_relaxer_integration.py
+++ b/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: