Support membrane protein in a3fe via Gromacs-based System Preparation

a3fe/_version.py

@@ -1 +1 @@
-__version__ = "0.3.3"
+__version__ = "0.3.4"

a3fe/run/system_prep.py

@@ -64,6 +64,8 @@ class SystemPreparationConfig(_BaseModel):
         , the restraints generated for the first repeat are used. This allows meaningful
         comparison between repeats for the bound leg. If False, the unique restraints are
         generated for each repeat.
+    membrane_protein: bool
+        If True, the protein is a membrane protein requiring special NPT settings.
     """
 
     slurm: bool = _Field(True)
@@ -175,6 +177,10 @@ class SystemPreparationConfig(_BaseModel):
             ],
         },
     }
+    membrane_protein: bool = _Field(
+        False,
+        description="Whether the protein is a membrane protein requiring special configuration settings.",
+    )
 
     class Config:
         """
@@ -184,6 +190,46 @@ class Config:
         extra = "forbid"
         validate_assignment = True
 
+    def get_membrane_nvt_config(self) -> dict:
+        """Get the temperature coupling configuration for membraneNVT equilibration runs."""
+        if self.membrane_protein:
+            return {
+                "tcoupl": "v-rescale",
+                "tc-grps": "Protein non-Protein",
+                "tau-t": "1.0 1.0",
+                "ref-t": "303.15 303.15",
+                "annealing": "no no",
+            }
+        return {}
+
+    def get_membrane_npt_config(self) -> dict:
+        """Get the pressure coupling configuration for membrane equilibration runs."""
+        if self.membrane_protein:
+            return {
+                "pcoupltype": "semiisotropic",
+                "tau-p": "5.0",
+                "compressibility": "4.5e-5 4.5e-5",
+                "ref-p": "1.0 1.0",
+                "rcoulomb": "1.2",
+                "rvdw": "1.2",
+                "rlist": "1.2",
+                "rvdw-switch": "1.0",
+                "vdw-modifier": "Force-switch",
+                "nstcomm": "100",
+                "comm-mode": "linear",
+                "nstxout-compressed": "5000",
+            }
+        return {}
+
+    def get_membrane_bound_equil_config(self) -> dict:
+        """Get the temperature and pressure coupling configuration for membrane equilibration runs."""
+        if self.membrane_protein:
+            return {
+                **self.get_membrane_nvt_config(),
+                **self.get_membrane_npt_config(),
+            }
+        return {}
+
     def get_tot_simtime(self, n_runs: int, leg_type: _LegType) -> int:
         """
         Get the total simulation time for the ensemble equilibration.
@@ -551,7 +597,11 @@ def heat_and_preequil_input(
         temperature_end=cfg.end_temp * _BSS.Units.Temperature.kelvin,
         restraint="all",
     )
-    equil1 = run_process(minimised_system, protocol)
+    # Only use membrane NVT config for bound leg (which has protein)
+    membrane_nvt_config = (
+        cfg.get_membrane_nvt_config() if leg_type == _LegType.BOUND else {}
+    )
+    equil1 = run_process(minimised_system, protocol, extra_options=membrane_nvt_config)
 
     # If this is the bound leg, carry out step with backbone restraints
     if leg_type == _LegType.BOUND:
@@ -563,7 +613,7 @@ def heat_and_preequil_input(
             temperature=cfg.end_temp * _BSS.Units.Temperature.kelvin,
             restraint="backbone",
         )
-        equil2 = run_process(equil1, protocol)
+        equil2 = run_process(equil1, protocol, extra_options=membrane_nvt_config)
 
     else:  # Free leg - skip the backbone restraint step
         equil2 = equil1
@@ -573,26 +623,40 @@ def heat_and_preequil_input(
         runtime=cfg.runtime_nvt * _BSS.Units.Time.picosecond,
         temperature=cfg.end_temp * _BSS.Units.Temperature.kelvin,
     )
-    equil3 = run_process(equil2, protocol)
+    equil3 = run_process(equil2, protocol, extra_options=membrane_nvt_config)
 
     print(
         f"NPT equilibration for {cfg.runtime_npt} ps while restraining non-solvent heavy atoms"
     )
-    protocol = _BSS.Protocol.Equilibration(
-        runtime=cfg.runtime_npt * _BSS.Units.Time.picosecond,
-        pressure=1 * _BSS.Units.Pressure.atm,
-        temperature=cfg.end_temp * _BSS.Units.Temperature.kelvin,
-        restraint="heavy",
-    )
-    equil4 = run_process(equil3, protocol)
+
+    # Create basic npt protocol
+    protocol_args = {
+        "runtime": cfg.runtime_npt * _BSS.Units.Time.picosecond,
+        "pressure": 1 * _BSS.Units.Pressure.atm,
+        "temperature": cfg.end_temp * _BSS.Units.Temperature.kelvin,
+    }
+    # membrane protein does not need heavy atom restraints.
+    if not cfg.membrane_protein:
+        protocol_args["restraint"] = "heavy"
+    protocol = _BSS.Protocol.Equilibration(**protocol_args)
+
+    # NPT configuration: bound leg needs full config, free leg only needs pressure coupling
+    if leg_type == _LegType.BOUND:
+        membrane_npt_config = cfg.get_membrane_bound_equil_config()
+    else:
+        membrane_npt_config = cfg.get_membrane_npt_config()
+
+    equil4 = run_process(equil3, protocol, extra_options=membrane_npt_config)
 
     print(f"NPT equilibration for {cfg.runtime_npt_unrestrained} ps without restraints")
     protocol = _BSS.Protocol.Equilibration(
         runtime=cfg.runtime_npt_unrestrained * _BSS.Units.Time.picosecond,
         pressure=1 * _BSS.Units.Pressure.atm,
         temperature=cfg.end_temp * _BSS.Units.Temperature.kelvin,
     )
-    preequilibrated_system = run_process(equil4, protocol)
+    preequilibrated_system = run_process(
+        equil4, protocol, extra_options=membrane_npt_config
+    )
 
     # Save, renaming the velocity property to foo so avoid saving velocities. Saving the
     # velocities sometimes causes issues with the size of the floats overflowing the RST7
@@ -679,9 +743,17 @@ def run_ensemble_equilibration(
     )
     if leg_type == _LegType.BOUND:
         work_dir = output_dir
+        membrane_equil_config = cfg.get_membrane_bound_equil_config()
     else:
         work_dir = None
-    final_system = run_process(pre_equilibrated_system, protocol, work_dir=work_dir)
+        membrane_equil_config = cfg.get_membrane_npt_config()
+
+    final_system = run_process(
+        pre_equilibrated_system,
+        protocol,
+        work_dir=work_dir,
+        extra_options=membrane_equil_config,
+    )
 
     # Save the coordinates only, renaming the velocity property to foo so avoid saving velocities. Saving the
     # velocities sometimes causes issues with the size of the floats overflowing the RST7
@@ -699,6 +771,7 @@ def run_process(
     system: _BSS._SireWrappers._system.System,
     protocol: _BSS.Protocol._protocol.Protocol,
     work_dir: _Optional[str] = None,
+    extra_options: _Optional[dict] = None,
 ) -> _BSS._SireWrappers._system.System:
     """
     Run a process with GROMACS, raising informative
@@ -713,13 +786,20 @@ def run_process(
     work_dir : str, optional
         The working directory to run the process in. If none,
         a temporary directory will be created.
+    extra_options : dict, optional
+        Extra options to pass to the GROMACS process, such as
+        membrane-specific configuration. If None, no extra options are passed.
 
     Returns
     -------
     system : _BSS._SireWrappers._system.System
         System after the process has been run.
     """
-    process = _BSS.Process.Gromacs(system, protocol, work_dir=work_dir)
+    if extra_options is None:
+        extra_options = {}
+    process = _BSS.Process.Gromacs(
+        system, protocol, work_dir=work_dir, extra_options=extra_options
+    )
     process.start()
     process.wait()
     import time

docs/CHANGELOG.rst

@@ -1,6 +1,9 @@
 ===============
 Change Log
 ===============
+0.3.4
+====================
+- Added ability to run membrane proteins. This is done by setting the ``membrane_protein`` option to ``True`` in the ``SystemPreparationConfig`` object.
 
 0.3.3
 ====================

docs/guides.rst

@@ -276,4 +276,29 @@ Since A3FE 0.2.0, ABFE calculations with charged ligands are supported using a c
     cutoff type = PME
     cutoff distance = 10 * angstrom
 
-The default `template_config.cfg` uses reaction field instead of PME. This is faster (around twice as fast for some of our systems) and has been shown to give equivalent results for neutral ligands in RBFE calculations - see https://pubs.acs.org/doi/full/10.1021/acs.jcim.0c01424 .
+The default `template_config.cfg` uses reaction field instead of PME. This is faster (around twice as fast for some of our systems) and has been shown to give equivalent results for neutral ligands in RBFE calculations - see https://pubs.acs.org/doi/full/10.1021/acs.jcim.0c01424 .
+
+ABFE with Membrane Proteins
+***************************
+
+Since A3FE 0.3.4, ABFE calculations with membrane proteins are supported via gromacs for minimization, NVT, NPT, and ensemble equilibration. The only change in the input required is that the use of barostat_membrane, rather than barostat, should be specified in ``template_config.cfg`` as e.g.:
+
+.. code-block:: bash
+
+    ### Barostat ###
+    barostat = False
+    barostat_membrane = True
+
+By modifying the ``SystemPreparationConfig`` object as described above, we can now try running membrane protein abfe calculation. 
+Note that this is expected to change the template_config.cfg file in the same time.
+
+.. code-block:: python
+
+    import a3fe as a3
+    cfg = a3.SystemPreparationConfig(membrane_protein=True)
+    calc = a3.Calculation()
+    calc.setup(bound_leg_sysprep_config = cfg, free_leg_sysprep_config = cfg)
+
+However, users must provide the ``bound_solv`` and ``free_solv`` files externally, as automated parameterization and solvation are not yet supported. 
+Additionally, the configuration files (gromacs.mdp) for these preparation steps specified in ``get_membrane_nvt_config`` and ``get_membrane_npt_config`` functions of ``system_prep.py`` file in which is general and not highly specialized; 
+Thus, it is recommended that users perform these pre-equilibration steps externally. Future updates will address these limitations.

docs/warnings.rst

@@ -1,4 +1,7 @@
 Warnings
 ========
 
- - We do no recommend running ABFE calculations with membrane proteins using the current version of ``a3fe``. This is because SOMD (Sire/ OpenMM Molecular Dynamics, part of Sire) is used for the free energy calculations (after setup with GROMACS), and currently SOMD uses an isotropic barostat. We plan to implement support for ABFE with GROMACS in the future.
+ - ABFE calculations with membrane proteins are supported via GROMACS for minimization, NVT, NPT, and ensemble equilibration. 
+   However, the configuration file (gromacs.mdp) used for these preparation steps is general and non-specialized. 
+   It is currently recommended that users perform the pre-equilibration procedures externally. 
+   Nevertheless, future releases will include support for these steps and we also plan to implement support ABFE calculations with GROMACS in the future.