Support membrane protein in a3fe via Gromacs-based System Preparation

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,29 @@ class Config:
         extra = "forbid"
         validate_assignment = True
 
+    def get_membrane_equilibration_config(self) -> dict:
+        """Get the membrane configuration for equilibration runs."""
+        if self.membrane_protein:
+            return {
+            "pcoupltype": "semiisotropic",
+            "pcoupl": "C-rescale",           
+            "tau-p": "5.0",                  
+            "nstpcouple": "100",             
+            "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_tot_simtime(self, n_runs: int, leg_type: _LegType) -> int:
         """
         Get the total simulation time for the ensemble equilibration.
@@ -578,21 +607,29 @@ def heat_and_preequil_input(
     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",
+    if not cfg.membrane_protein:
+        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)
+    #membrane protein does not need heavy atom restraints.
+    else:
+        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,
+        )
+    equil4 = run_process(equil3, protocol, extra_options=cfg.get_membrane_equilibration_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=cfg.get_membrane_equilibration_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
@@ -681,7 +718,7 @@ def run_ensemble_equilibration(
         work_dir = output_dir
     else:
         work_dir = None
-    final_system = run_process(pre_equilibrated_system, protocol, work_dir=work_dir)
+    final_system = run_process(pre_equilibrated_system, protocol, work_dir=work_dir, extra_options=cfg.get_membrane_equilibration_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 +736,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 +751,18 @@ 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/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_equilibration_config`` function 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.