Merge pull request #343 from OpenBioSim/backport_341_342

Backport fixes from #341 and #342

python/BioSimSpace/Process/_amber.py

@@ -304,6 +304,7 @@ def _setup(self, **kwargs):
 
                     # Set the simulation box.
                     system.setBox(*_cubic(box_length))
+                    reference_system.setBox(*_cubic(box_length))
 
             # Apply SOMD1 compatibility to the perturbation.
             if (
@@ -322,6 +323,7 @@ def _setup(self, **kwargs):
         else:
             # Check for perturbable molecules and convert to the chosen end state.
             system = self._checkPerturbable(system)
+            reference_system = self._checkPerturbable(self._reference_system)
 
         # RST file (coordinates).
         try:

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
 

python/BioSimSpace/Process/_namd.py

@@ -2115,6 +2115,9 @@ def _createRestrainedSystem(self, system, restraint):
         # Copy the original system.
         s = system.copy()
 
+        # Convert to the chosen end state.
+        s = self._checkPerturbable(s)
+
         # Keyword restraint.
         if isinstance(restraint, str):
             # Loop over all molecules by number.

python/BioSimSpace/Process/_openmm.py

@@ -254,6 +254,7 @@ def _setup(self):
 
         # Check for perturbable molecules and convert to the chosen end state.
         system = self._checkPerturbable(system)
+        reference_system = self._checkPerturbable(self._reference_system)
 
         # Create the input files...
 
@@ -274,7 +275,7 @@ def _setup(self):
                 file = _os.path.splitext(self._ref_file)[0]
                 _IO.saveMolecules(
                     file,
-                    self._reference_system,
+                    reference_system,
                     "rst7",
                     property_map=self._property_map,
                 )

tests/Process/test_amber.py

@@ -243,6 +243,17 @@ def test_backbone_restraint_mask_rna(rna_system):
     assert "   restraintmask=\"@P,C5',C3',O3',O5'\"," in config
 
 
+@pytest.mark.skipif(has_amber is False, reason="Requires AMBER to be installed.")
+def test_perturbable_restraint(perturbable_system):
+    """Test a free energy perturbation protocol."""
+
+    # Create a short minimisation prototocol 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, check_data=False):
     """Helper function to run various simulation protocols."""
 

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

tests/Process/test_namd.py

@@ -90,6 +90,17 @@ def test_production(namd_system, restraint):
     run_process(namd_system, protocol)
 
 
+@pytest.mark.skipif(has_namd is False, reason="Requires NAMD to be installed.")
+def test_perturbable_restraint(perturbable_system):
+    """Test a free energy perturbation protocol."""
+
+    # Create a short minimisation prototocol 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(namd_system, protocol):
     """Helper function to run various simulation protocols."""
 

tests/Process/test_openmm.py

@@ -111,6 +111,16 @@ def test_parmed_triclinic():
     run_process(system, protocol)
 
 
+def test_perturbable_restraint(perturbable_system):
+    """Test a free energy perturbation protocol."""
+
+    # Create a short minimisation prototocol 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):
     """Helper function to run various simulation protocols."""