membrane CLI support

docs/conf.py

@@ -193,7 +193,7 @@
     else:
         repo = git.Repo.clone_from(
             "https://github.com/OpenFreeEnergy/ExampleNotebooks.git",
-            branch="2026.01.26",
+            branch="membrane_tutorial",  # TODO: update this after merging to main
             to_path=example_notebooks_path,
         )
 except Exception as e:

docs/tutorials/index.rst

@@ -12,8 +12,9 @@ Relative Free Energies
 ----------------------
 
 - :any:`Python API Showcase <showcase_notebook>`: Start here! An introduction to OpenFE's Python API and approach to performing a relative binding free energy calculation.
-- :any:`RBFE with the Python API <rbfe_python_tutorial>`: A step-by-step tutorial for using the Python API to calculate relative binding free energies for TYK2.
-- :ref:`RBFE with the CLI <rbfe_cli_tutorial>`: A step-by-step tutorial for using the OpenFE command line interface (CLI) to calculate relative binding free energies for TYK2.
+- :any:`RBFE using the Python API <rbfe_python_tutorial>`: A step-by-step tutorial for using the Python API to calculate relative binding free energies for TYK2.
+- :ref:`RBFE using the CLI <rbfe_cli_tutorial>`: A step-by-step tutorial for using the OpenFE command line interface (CLI) to calculate relative binding free energies for TYK2.
+- :any:`RBFE with membrane systems <rbfe_membrane_protein>`: A step-by-step guide to setting up an RBFE calculation with special considerations for membrane systems.
 
 Absolute Free Energies
 ----------------------
@@ -50,6 +51,7 @@ Generating Partial Charges
     showcase_notebook
     rbfe_python_tutorial
     rbfe_cli_tutorial
+    rbfe_membrane_protein
     abfe_tutorial
     abfe_analysis_tutorial
     ahfe_tutorial

docs/tutorials/rbfe_membrane_protein.nblink

@@ -0,0 +1,3 @@
+{
+   "path": "../ExampleNotebooks/membranes/rbfe_membrane_protein.ipynb"
+}

src/openfe/setup/alchemical_network_planner/relative_alchemical_network_planner.py

@@ -10,12 +10,12 @@
     LigandAtomMapping,
     LigandNetwork,
     ProteinComponent,
+    ProteinMembraneComponent,
     Protocol,
     SmallMoleculeComponent,
     SolventComponent,
     Transformation,
 )
-from openff.units import unit
 
 from ...protocols.openmm_rfe.equil_rfe_methods import RelativeHybridTopologyProtocol
 from .. import LomapAtomMapper
@@ -334,7 +334,7 @@ def _build_transformation(
             protocol_settings.nonbonded_method = "nocutoff"
 
         transformation_protocol = transformation_protocol.__class__(settings=protocol_settings)
-
+        transformation_protocol.validate(stateA=stateA, stateB=stateB, mapping=ligand_mapping_edge)
         return Transformation(
             stateA=stateA,
             stateB=stateB,
@@ -347,7 +347,7 @@ def __call__(
         self,
         ligands: Iterable[SmallMoleculeComponent],
         solvent: SolventComponent,
-        protein: ProteinComponent,
+        protein: ProteinComponent | ProteinMembraneComponent,
         cofactors: Optional[Iterable[SmallMoleculeComponent]] = None,
     ) -> AlchemicalNetwork:
         """plan the alchemical network for RBFE calculations with the given ligands, protein and solvent.
@@ -358,7 +358,7 @@ def __call__(
             ligands that shall be used for the alchemical network.
         solvent : SolventComponent
             solvent for solvated and complex simulations
-        protein : ProteinComponent
+        protein : ProteinComponent | ProteinMembraneComponent
             protein for complex simulations
         cofactors : Iterable[SmallMoleculeComponent]
             any cofactors in the system, can be empty list

src/openfe/setup/chemicalsystem_generator/easy_chemicalsystem_generator.py

@@ -7,6 +7,7 @@
     ChemicalSystem,
     Component,
     ProteinComponent,
+    ProteinMembraneComponent,
     SmallMoleculeComponent,
     SolventComponent,
 )
@@ -20,9 +21,9 @@
 class EasyChemicalSystemGenerator(AbstractChemicalSystemGenerator):
     def __init__(
         self,
-        solvent: Optional[SolventComponent] = None,
-        protein: Optional[ProteinComponent] = None,
-        cofactors: Optional[Iterable[SmallMoleculeComponent]] = None,
+        solvent: SolventComponent | None = None,
+        protein: ProteinComponent | None = None,
+        cofactors: Iterable[SmallMoleculeComponent] | None = None,
         do_vacuum: bool = False,
     ):
         """
@@ -45,8 +46,8 @@ def __init__(
         ----------
         solvent : SolventComponent, optional
             if a SolventComponent is given, solvated chemical systems will be generated, by default None
-        protein : ProteinComponent, optional
-            if a ProteinComponent is given, complex chemical systems will be generated, by default None
+        protein : ProteinComponent | ProteinMembraneComponent, optional
+            if a ProteinComponent or ProteinMembraneComponent is given, complex chemical systems will be generated, by default None
         cofactors : Iterable[SmallMoleculeComponent], optional
             any cofactors in the system.  will be put in any systems containing
             the protein
@@ -112,7 +113,9 @@ def __call__(self, component: SmallMoleculeComponent) -> Iterable[ChemicalSystem
             }
             for i, c in enumerate(self.cofactors):
                 components.update({f"{RFEComponentLabels.COFACTOR.value}{i + 1}": c})
-            if self.solvent is not None:
+
+            # ProteinMembraneComponent has its own solvent.
+            if self.solvent is not None and not isinstance(self.protein, ProteinMembraneComponent):
                 components.update({RFEComponentLabels.SOLVENT.value: self.solvent})
             chem_sys = ChemicalSystem(components=components, name=component.name + "_complex")
             yield chem_sys

src/openfe/tests/conftest.py

@@ -279,11 +279,14 @@ def benzene_transforms():
 
 
 @pytest.fixture(scope="session")
-def T4_protein_component():
+def T4_protein_component_pdb() -> str:
     with resources.as_file(resources.files("openfe.tests.data")) as d:
-        fn = str(d / "181l_only.pdb")
+        return str(d / "181l_only.pdb")
 
-    return gufe.ProteinComponent.from_pdb_file(fn, name="T4_protein")
+
+@pytest.fixture(scope="session")
+def T4_protein_component(T4_protein_component_pdb) -> gufe.ProteinComponent:
+    return gufe.ProteinComponent.from_pdb_file(T4_protein_component_pdb, name="T4_protein")
 
 
 @pytest.fixture(scope="session")
@@ -293,9 +296,8 @@ def a2a_protein_membrane_pdb():
 
 
 @pytest.fixture(scope="session")
-def a2a_protein_membrane_component(a2a_protein_membrane_pdb):
-    with gzip.open(a2a_protein_membrane_pdb, "rb") as f:
-        yield openfe.ProteinMembraneComponent.from_pdb_file(f, name="a2a")
+def a2a_protein_membrane_component(a2a_protein_membrane_pdb) -> openfe.ProteinMembraneComponent:
+    yield openfe.ProteinMembraneComponent.from_pdb_file(a2a_protein_membrane_pdb, name="a2a")
 
 
 @pytest.fixture(scope="session")
@@ -357,7 +359,7 @@ def fepplus_network():
 
 
 @pytest.fixture()
-def CN_molecule():
+def CN_molecule() -> list[SmallMoleculeComponent]:
     """
     A basic CH3NH2 molecule for quick testing.
     """

src/openfecli/commands/plan_rbfe_network.py

@@ -12,6 +12,7 @@
     OUTPUT_DIR,
     OVERWRITE,
     PROTEIN,
+    PROTEIN_MEMBRANE,
     YAML_OPTIONS,
 )
 from openfecli.utils import print_duration, write
@@ -44,8 +45,8 @@ def plan_rbfe_network_main(
         ligands of the system
     solvent : SolventComponent
         Solvent component used for solvation
-    protein : ProteinComponent
-        protein component for complex simulations, to which the ligands are bound
+    protein : ProteinComponent | ProteinMembraneComponent
+        ProteinComponent or ProteinMembraneComponent for complex simulations, to which the ligands are bound.
     cofactors : Iterable[SmallMoleculeComponent]
         any cofactors alongside the protein, can be empty list
     n_protocol_repeats: int
@@ -123,7 +124,10 @@ def plan_rbfe_network_main(
     ),
 )
 @MOL_DIR.parameter(required=True, help=MOL_DIR.kwargs["help"] + " Any number of sdf paths.")
-@PROTEIN.parameter(multiple=False, required=True, default=None, help=PROTEIN.kwargs["help"])
+@PROTEIN.parameter(multiple=False, required=False, default=None, help=PROTEIN.kwargs["help"])
+@PROTEIN_MEMBRANE.parameter(
+    multiple=False, required=False, default=None, help=PROTEIN_MEMBRANE.kwargs["help"]
+)
 @COFACTORS.parameter(multiple=True, required=False, default=None, help=COFACTORS.kwargs["help"])
 @YAML_OPTIONS.parameter(multiple=False, required=False, default=None, help=YAML_OPTIONS.kwargs["help"])  # fmt: skip
 @OUTPUT_DIR.parameter(help=OUTPUT_DIR.kwargs["help"] + " Defaults to `./alchemicalNetwork`.", default="alchemicalNetwork")  # fmt: skip
@@ -133,7 +137,8 @@ def plan_rbfe_network_main(
 @print_duration
 def plan_rbfe_network(
     molecules: list[str],
-    protein: str,
+    protein: str | None,
+    protein_membrane: str | None,
     cofactors: tuple[str],
     yaml_settings: str,
     output_dir: str,
@@ -142,14 +147,12 @@ def plan_rbfe_network(
     overwrite_charges: bool,
 ):
     """
-    Plan a relative binding free energy network, saved as JSON files for use by
-    the quickrun command.
+    Plan a relative binding free energy AlchemicalNetwork, saved as JSON files for use by the quickrun command.
 
     This tool is an easy way to set up a RBFE calculation campaign.
-    The JSON files this outputs can be used to run each leg of the campaign.
-    openfe.
-    The generated Network will be stored in a folder containing for each
-    transformation a JSON file, that can be run with quickrun.
+
+    The generated AlchemicalNetwork will be stored in --output-directory along with JSON files for each alchemical transformation
+    that can be used to execute the campaign using ``openfe quickrun``.
 
     .. note::
 
@@ -169,9 +172,7 @@ def plan_rbfe_network(
     * Protocol is the OpenMM-based relative hybrid topology protocol, with
       default settings.
 
-    These choices can be customized by creating a settings yaml file,
-    which is passed in via the ``-s settings.yaml`` option,
-    which is detailed in the Options section.
+    These choices can be customized by creating a settings yaml file, which is passed in via the ``-s settings.yaml`` option.
     For more advanced setups, please consider using the Python layer of openfe.
     """
     write("RBFE-NETWORK PLANNER")
@@ -187,9 +188,18 @@ def plan_rbfe_network(
 
     small_molecules = MOL_DIR.get(molecules)
     write("\t\tSmall Molecules: " + " ".join([str(sm) for sm in small_molecules]))
-
-    protein = PROTEIN.get(protein)
-    write("\t\tProtein: " + str(protein))
+    if protein and protein_membrane:
+        raise ValueError("Only --protein (-p) or --protein-membrane may be provided, not both.")
+    elif protein:
+        protein_component = PROTEIN.get(protein)
+        write("\t\tProteinComponent: " + str(protein_component))
+    elif protein_membrane:
+        protein_component = PROTEIN_MEMBRANE.get(protein_membrane)
+        write("\t\tProteinMembraneComponent: " + str(protein_component))
+    else:
+        raise ValueError(
+            "Either --protein or --protein-membrane must be provided. Run ``openfe plan-rbfe-network -h`` for more information."
+        )
 
     if cofactors is not None:
         cofactors = sum((COFACTORS.get(c) for c in cofactors), start=[])
@@ -216,6 +226,9 @@ def plan_rbfe_network(
     # TODO:  write nice parameter
     write("\tNetwork Generation: " + str(ligand_network_planner))
 
+    write(f"Validating {type(protein_component).__name__}...")
+    protein_component.validate()
+
     write("\tPartial Charge Generation: " + str(partial_charge.partial_charge_method))
     if overwrite_charges:
         write("\tOverwriting partial charges")
@@ -230,7 +243,7 @@ def plan_rbfe_network(
         ligand_network_planner=ligand_network_planner,
         small_molecules=small_molecules,
         solvent=solvent,
-        protein=protein,
+        protein=protein_component,
         cofactors=cofactors,
         n_protocol_repeats=n_protocol_repeats,
         partial_charge_settings=partial_charge,

src/openfecli/parameters/__init__.py

@@ -8,4 +8,4 @@
 from .output import OUTPUT_FILE_AND_EXT
 from .output_dir import OUTPUT_DIR
 from .plan_network_options import YAML_OPTIONS
-from .protein import PROTEIN
+from .protein import PROTEIN, PROTEIN_MEMBRANE

src/openfecli/parameters/protein.py

@@ -1,46 +1,60 @@
 # This code is part of OpenFE and is licensed under the MIT license.
 # For details, see https://github.com/OpenFreeEnergy/openfe
 
-from plugcli.params import NOT_PARSED, MultiStrategyGetter, Option
+import sys
+from typing import Iterable
 
+import click
+from plugcli.params import Option
 
-def _load_protein_from_pdb(user_input, context):
-    if ".pdb" not in str(user_input):  # this silences some stderr spam
-        return NOT_PARSED
+from openfe import ProteinComponent, ProteinMembraneComponent
 
-    from gufe import ProteinComponent
-
-    try:
-        return ProteinComponent.from_pdb_file(user_input)
-    except ValueError:  # any exception should try other strategies
-        return NOT_PARSED
+_PDB_EXT = [".pdb"]
+_PDBX_EXT = [".cif", ".pdbx"]
 
 
-def _load_protein_from_pdbx(user_input, context):
-    if not any(
-        [ext in str(user_input) for ext in [".pdb", ".cif", ".pdbx"]]
-    ):  # this silences some stderr spam
-        return NOT_PARSED
+def _contains_any_substring(input: str, substrings: Iterable[str]) -> bool:
+    return any([substring in input for substring in substrings])
 
-    from gufe import ProteinComponent
 
+def _load_protein_from_file(input_file, protein_class: ProteinComponent | ProteinMembraneComponent):
+    valid_ext = _PDB_EXT + _PDBX_EXT
+    if not _contains_any_substring(input_file, valid_ext):
+        raise ValueError(
+            f"To load a {protein_class.__name__}, the file extension must contain one of: {valid_ext}."
+        )
     try:
-        return ProteinComponent.from_pdbx_file(user_input)
-    except ValueError:  # any exception should try other strategies
-        return NOT_PARSED
+        if _contains_any_substring(input_file, _PDB_EXT):
+            return protein_class.from_pdb_file(input_file)
+        elif _contains_any_substring(input_file, _PDBX_EXT):
+            return protein_class.from_pdbx_file(input_file)
+    except ValueError:
+        click.secho(f"Unable to load a {protein_class.__name__} from {input_file}.", err=True)
+        sys.exit(1)
 
 
-get_molecule = MultiStrategyGetter(
-    strategies=[
-        _load_protein_from_pdb,
-        _load_protein_from_pdbx,
-    ],
-    error_message="Unable to generate a molecule from '{user_input}'.",
-)
+# TODO: these functions are shims to work with plugcli. We should consider migrating to just click.
+def _get_protein(user_input, context):
+    return _load_protein_from_file(user_input, ProteinComponent)
+
+
+def _get_protein_membrane(user_input, context):
+    return _load_protein_from_file(user_input, ProteinMembraneComponent)
+
 
 PROTEIN = Option(
     "-p",
     "--protein",
-    help=("ProteinComponent. Can be provided as an PDB or as a PDBx/mmCIF file.  string."),
-    getter=get_molecule,
+    help=(
+        "Path to a PDB or PDBx/mmCIF file containing a protein. Mutually exclusive with --protein-membrane."
+    ),
+    getter=_get_protein,
+)
+
+PROTEIN_MEMBRANE = Option(
+    "--protein-membrane",
+    help=(
+        'Path to a PDB or PDBx/mmCIF file containing a fully solvated protein-membrane system. Mutually exclusive with --protein. See "Combining System Components into a Single PDB File"'
+    ),
+    getter=_get_protein_membrane,
 )