Merge pull request #342 from OpenBioSim/fix_339

Fix issue #339
OpenBioSim · Sep 25, 2024 · 71927d1 · 71927d1
2 parents 8fe069e + b40f4c9
commit 71927d1
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 7 deletions.
diff --git a/python/BioSimSpace/Process/_gromacs.py b/python/BioSimSpace/Process/_gromacs.py
@@ -2131,16 +2131,31 @@ def _add_position_restraints(self):
                         )
 
                         with open(restraint_file, "w") as file:
-                            # Write the header.
-                            file.write("[ position_restraints ]\n")
-                            file.write(";  i funct       fcx        fcy        fcz\n")
-
-                            # Write restraints for each atom.
-                            for atom_idx in restrained_atoms:
+                            if isinstance(self._protocol, _FreeEnergyMixin):
+                                # Write the header.
+                                file.write("[ position_restraints ]\n")
                                 file.write(
-                                    f"{atom_idx+1:4}    1       {force_constant}       {force_constant}       {force_constant}\n"
+                                    ";  i funct       fcx_A       fcy_A       fcz_A       fcx_B       fcy_B       fcz_B\n"
+                                )
+                                # Write restraints for each atom.
+                                for atom_idx in restrained_atoms:
+                                    file.write(
+                                        f"{atom_idx+1:4}    1       {force_constant}       {force_constant}       {force_constant}       {force_constant}       {force_constant}       {force_constant}\n"
+                                    )
+
+                            else:
+                                # Write the header.
+                                file.write("[ position_restraints ]\n")
+                                file.write(
+                                    ";  i funct       fcx        fcy        fcz\n"
                                 )
 
+                                # Write restraints for each atom.
+                                for atom_idx in restrained_atoms:
+                                    file.write(
+                                        f"{atom_idx+1:4}    1       {force_constant}       {force_constant}       {force_constant}\n"
+                                    )
+
                         # Work out the offset.
                         offset = num_restraint - 1
 

diff --git a/tests/Process/test_gromacs.py b/tests/Process/test_gromacs.py
@@ -116,6 +116,17 @@ def test_restraints(perturbable_system, restraint):
     process = BSS.Process.Gromacs(perturbable_system, protocol)
 
 
+@pytest.mark.skipif(has_gromacs is False, reason="Requires GROMACS to be installed.")
+def test_perturbable_restraint(perturbable_system):
+    """Test a free energy perturbation protocol."""
+
+    # Create a short minimisation protocol with a restraint.
+    protocol = BSS.Protocol.Minimisation(steps=100, restraint="heavy")
+
+    # Run the process, check that it finished without error, and returns a system.
+    run_process(perturbable_system, protocol)
+
+
 def run_process(system, protocol, **kwargs):
     """Helper function to run various simulation protocols."""