Skip to content

Commit 0cc52d6

Browse files
FridrichMethodFridrichMethod
committed
refactor: drive rbfe ligand pairing by SMILES via complex_pdb_for
Replace the regex-based _match_ligand PDB-to-ligand pairing with a single mapping function complex_pdb_for(name) -> Path defined in the config cell. Ligand names from the .smi file are now the source of truth: both prep passes loop over smiles_dict and resolve each name to its complex PDB. Mapping logic lives in one editable place; drop the now-unused import re.
1 parent d837321 commit 0cc52d6

1 file changed

Lines changed: 9 additions & 214 deletions

File tree

‎examples/openfe/rbfe.ipynb‎

Lines changed: 9 additions & 214 deletions
Original file line numberDiff line numberDiff line change
@@ -14,252 +14,47 @@
1414
},
1515
{
1616
"cell_type": "code",
17-
"execution_count": 1,
17+
"execution_count": null,
1818
"id": "fc97de03",
1919
"metadata": {},
20-
"outputs": [
21-
{
22-
"data": {
23-
"application/vnd.jupyter.widget-view+json": {
24-
"model_id": "60677d08ac4a40fdbf84bfae3f780378",
25-
"version_major": 2,
26-
"version_minor": 0
27-
},
28-
"text/plain": []
29-
},
30-
"metadata": {},
31-
"output_type": "display_data"
32-
}
33-
],
34-
"source": [
35-
"import multiprocessing as mp\n",
36-
"import os\n",
37-
"import re\n",
38-
"from dataclasses import asdict, dataclass\n",
39-
"from pathlib import Path\n",
40-
"\n",
41-
"import MDAnalysis as mda\n",
42-
"import mdtraj as md\n",
43-
"import openfe\n",
44-
"from MDAnalysis.analysis import align\n",
45-
"from openfe import (\n",
46-
" AlchemicalNetwork,\n",
47-
" ChemicalSystem,\n",
48-
" ProteinComponent,\n",
49-
" SmallMoleculeComponent,\n",
50-
" SolventComponent,\n",
51-
" Transformation,\n",
52-
")\n",
53-
"from openfe.protocols.openmm_rfe import RelativeHybridTopologyProtocol\n",
54-
"from openfe.protocols.openmm_septop import SepTopProtocol\n",
55-
"from openfe.protocols.openmm_utils.charge_generation import bulk_assign_partial_charges\n",
56-
"from openfe.protocols.openmm_utils.omm_settings import OpenFFPartialChargeSettings\n",
57-
"from openfe.setup import RBFEAlchemicalNetworkPlanner, RHFEAlchemicalNetworkPlanner\n",
58-
"from openfe.setup.atom_mapping import KartografAtomMapper, LomapAtomMapper\n",
59-
"from openfe.setup.atom_mapping.lomap_scorers import (\n",
60-
" atomic_number_score,\n",
61-
" default_lomap_score,\n",
62-
" ecr_score,\n",
63-
" mcsr_score,\n",
64-
" mncar_score,\n",
65-
")\n",
66-
"from openfe.setup.ligand_network_planning import (\n",
67-
" generate_lomap_network,\n",
68-
" generate_maximal_network,\n",
69-
" generate_minimal_redundant_network,\n",
70-
" generate_minimal_spanning_network,\n",
71-
" generate_network_from_indices,\n",
72-
" generate_network_from_names,\n",
73-
" generate_radial_network,\n",
74-
")\n",
75-
"from openfe.utils.atommapping_network_plotting import plot_atommapping_network\n",
76-
"from openff.units import unit\n",
77-
"from rdkit import Chem\n",
78-
"from rdkit.Chem import AllChem\n",
79-
"from rdkit.Chem.Descriptors3D import Asphericity\n",
80-
"\n",
81-
"from mdpp.plots import draw_mols, make_atom_labels_3d, view_mol_3d, view_traj_3d\n",
82-
"from mdpp.prep import assign_topology, fix_pdb"
83-
]
20+
"outputs": [],
21+
"source": "import multiprocessing as mp\nimport os\nfrom dataclasses import asdict, dataclass\nfrom pathlib import Path\n\nimport MDAnalysis as mda\nimport mdtraj as md\nimport openfe\nfrom MDAnalysis.analysis import align\nfrom openfe import (\n AlchemicalNetwork,\n ChemicalSystem,\n ProteinComponent,\n SmallMoleculeComponent,\n SolventComponent,\n Transformation,\n)\nfrom openfe.protocols.openmm_rfe import RelativeHybridTopologyProtocol\nfrom openfe.protocols.openmm_septop import SepTopProtocol\nfrom openfe.protocols.openmm_utils.charge_generation import bulk_assign_partial_charges\nfrom openfe.protocols.openmm_utils.omm_settings import OpenFFPartialChargeSettings\nfrom openfe.setup import RBFEAlchemicalNetworkPlanner, RHFEAlchemicalNetworkPlanner\nfrom openfe.setup.atom_mapping import KartografAtomMapper, LomapAtomMapper\nfrom openfe.setup.atom_mapping.lomap_scorers import (\n atomic_number_score,\n default_lomap_score,\n ecr_score,\n mcsr_score,\n mncar_score,\n)\nfrom openfe.setup.ligand_network_planning import (\n generate_lomap_network,\n generate_maximal_network,\n generate_minimal_redundant_network,\n generate_minimal_spanning_network,\n generate_network_from_indices,\n generate_network_from_names,\n generate_radial_network,\n)\nfrom openfe.utils.atommapping_network_plotting import plot_atommapping_network\nfrom openff.units import unit\nfrom rdkit import Chem\nfrom rdkit.Chem import AllChem\nfrom rdkit.Chem.Descriptors3D import Asphericity\n\nfrom mdpp.plots import draw_mols, make_atom_labels_3d, view_mol_3d, view_traj_3d\nfrom mdpp.prep import assign_topology, fix_pdb"
8422
},
8523
{
8624
"cell_type": "markdown",
8725
"id": "deb54778",
8826
"metadata": {},
89-
"source": [
90-
"## Configuration\n",
91-
"\n",
92-
"Define input paths (complex PDBs, SMILES file), chain identifiers, and simulation parameters (temperature, pH, lambda windows, production length). All intermediate files are written under a working directory."
93-
]
27+
"source": "## Configuration\n\nDefine input paths (complex PDBs, SMILES file), chain identifiers, and simulation parameters (temperature, pH, lambda windows, production length). The `complex_pdb_for` function maps each ligand name -- taken from the SMILES file, which is the single source of truth -- to its complex PDB file; to change how names resolve to PDBs, edit only this function. All intermediate files are written under a working directory."
9428
},
9529
{
9630
"cell_type": "code",
97-
"execution_count": 2,
31+
"execution_count": null,
9832
"id": "0c51edf0",
9933
"metadata": {},
10034
"outputs": [],
101-
"source": [
102-
"COMPLEX_DIR = Path(\"pdbs\")\n",
103-
"SMILES_PATH = Path(\"ligands_amp_fep.smi\")\n",
104-
"PROTEIN_CHAIN_ID = \"A\"\n",
105-
"LIGAND_CHAIN_ID = \"B\"\n",
106-
"\n",
107-
"WORKING_DIR = Path(\"tmp\")\n",
108-
"WORKING_DIR.mkdir(exist_ok=True)\n",
109-
"\n",
110-
"ALIGN_DIR = WORKING_DIR / \"align\"\n",
111-
"ALIGN_DIR.mkdir(exist_ok=True)\n",
112-
"PROTEIN_PATH = WORKING_DIR / \"protein.pdb\"\n",
113-
"FIXED_PROTEIN_PATH = WORKING_DIR / \"protein_fixed.pdb\"\n",
114-
"LIGAND_DIR = WORKING_DIR / \"ligands\"\n",
115-
"LIGAND_DIR.mkdir(exist_ok=True)\n",
116-
"CHARGED_LIGAND_DIR = WORKING_DIR / \"charged_ligands\"\n",
117-
"CHARGED_LIGAND_DIR.mkdir(exist_ok=True)\n",
118-
"LIGAND_NETWORK_PATH = WORKING_DIR / \"ligand_network.graphml\"\n",
119-
"TRANSFORMATION_DIR = WORKING_DIR / \"transformations\"\n",
120-
"TRANSFORMATION_DIR.mkdir(exist_ok=True)\n",
121-
"SEPTOP_TRANSFORMATION_DIR = WORKING_DIR / \"septop_transformations\"\n",
122-
"SEPTOP_TRANSFORMATION_DIR.mkdir(exist_ok=True)\n",
123-
"\n",
124-
"# Simulation parameters\n",
125-
"LAMBDA_WINDOWS = 51\n",
126-
"N_REPLICAS = LAMBDA_WINDOWS\n",
127-
"PH = 7.0\n",
128-
"PRODUCTION_LENGTH = 20.0\n",
129-
"REPEATS = 1\n",
130-
"TEMPERATURE = 298.15\n",
131-
"\n",
132-
"LOMAP_MAPPER_SEED = \"O=C(CCCC[C,N])OP(OC[C@H]1O[C@@H](n2cnc3c2ncnc3N)[C@H](O)[C@@H]1O)([O-])=O\"\n",
133-
"\n",
134-
"\n",
135-
"@dataclass\n",
136-
"class LomapAtomMapperConfig:\n",
137-
" \"\"\"Configuration for the Lomap atom mapper.\"\"\"\n",
138-
"\n",
139-
" time: int = 20\n",
140-
" threed: bool = True\n",
141-
" max3d: float = 1.0\n",
142-
" element_change: bool = False # Do not allow mappings that change an atoms element\n",
143-
" seed: str = LOMAP_MAPPER_SEED # SMARTS string for MCS search\n",
144-
" shift: bool = False\n",
145-
"\n",
146-
"\n",
147-
"@dataclass\n",
148-
"class KartografAtomMapperConfig:\n",
149-
" \"\"\"Configuration for the Kartograf atom mapper.\"\"\"\n",
150-
"\n",
151-
" atom_max_distance: float = 0.9\n",
152-
" atom_map_hydrogens: bool = True\n",
153-
" map_hydrogens_on_hydrogens_only: bool = True\n",
154-
" map_exact_ring_matches_only: bool = True\n",
155-
" allow_partial_fused_rings: bool = False\n",
156-
" allow_bond_breaks: bool = False"
157-
]
35+
"source": "COMPLEX_DIR = Path(\"pdbs\")\nSMILES_PATH = Path(\"ligands_amp_fep.smi\")\nPROTEIN_CHAIN_ID = \"A\"\nLIGAND_CHAIN_ID = \"B\"\n\n\ndef complex_pdb_for(name: str) -> Path:\n \"\"\"Resolve a ligand name (the SMILES `_Name`) to its complex PDB path.\n\n Ligand names from `SMILES_PATH` are the single source of truth; this\n function is the only place that encodes how each name maps to a complex\n PDB file. To change the mapping logic (a different filename template, an\n explicit lookup table, multiple input directories, etc.), edit only this\n function.\n\n Args:\n name: Ligand name, i.e. the title field of an entry in `SMILES_PATH`.\n\n Returns:\n Path to the complex PDB file for that ligand.\n \"\"\"\n return COMPLEX_DIR / f\"FLE03_{name}_3a7r_model_0.pdb\"\n\n\nWORKING_DIR = Path(\"tmp\")\nWORKING_DIR.mkdir(exist_ok=True)\n\nALIGN_DIR = WORKING_DIR / \"align\"\nALIGN_DIR.mkdir(exist_ok=True)\nPROTEIN_PATH = WORKING_DIR / \"protein.pdb\"\nFIXED_PROTEIN_PATH = WORKING_DIR / \"protein_fixed.pdb\"\nLIGAND_DIR = WORKING_DIR / \"ligands\"\nLIGAND_DIR.mkdir(exist_ok=True)\nCHARGED_LIGAND_DIR = WORKING_DIR / \"charged_ligands\"\nCHARGED_LIGAND_DIR.mkdir(exist_ok=True)\nLIGAND_NETWORK_PATH = WORKING_DIR / \"ligand_network.graphml\"\nTRANSFORMATION_DIR = WORKING_DIR / \"transformations\"\nTRANSFORMATION_DIR.mkdir(exist_ok=True)\nSEPTOP_TRANSFORMATION_DIR = WORKING_DIR / \"septop_transformations\"\nSEPTOP_TRANSFORMATION_DIR.mkdir(exist_ok=True)\n\n# Simulation parameters\nLAMBDA_WINDOWS = 51\nN_REPLICAS = LAMBDA_WINDOWS\nPH = 7.0\nPRODUCTION_LENGTH = 20.0\nREPEATS = 1\nTEMPERATURE = 298.15\n\nLOMAP_MAPPER_SEED = \"O=C(CCCC[C,N])OP(OC[C@H]1O[C@@H](n2cnc3c2ncnc3N)[C@H](O)[C@@H]1O)([O-])=O\"\n\n\n@dataclass\nclass LomapAtomMapperConfig:\n \"\"\"Configuration for the Lomap atom mapper.\"\"\"\n\n time: int = 20\n threed: bool = True\n max3d: float = 1.0\n element_change: bool = False # Do not allow mappings that change an atoms element\n seed: str = LOMAP_MAPPER_SEED # SMARTS string for MCS search\n shift: bool = False\n\n\n@dataclass\nclass KartografAtomMapperConfig:\n \"\"\"Configuration for the Kartograf atom mapper.\"\"\"\n\n atom_max_distance: float = 0.9\n atom_map_hydrogens: bool = True\n map_hydrogens_on_hydrogens_only: bool = True\n map_exact_ring_matches_only: bool = True\n allow_partial_fused_rings: bool = False\n allow_bond_breaks: bool = False"
15836
},
15937
{
16038
"cell_type": "markdown",
16139
"id": "ce5e9dbf",
16240
"metadata": {},
163-
"source": [
164-
"## Prepare structures from complex PDB files\n",
165-
"\n",
166-
"Each complex PDB file contains a protein and ligand. We select the reference protein as the complex bound to the most aspherical ligand (highest Asphericity descriptor), align all other complexes to it, and assign correct bond orders from the SMILES file using `assign_topology`. The reference protein is expected to already be fixed and protonated (via `fix_pdb`) from a prior run."
167-
]
168-
},
169-
{
170-
"cell_type": "code",
171-
"execution_count": null,
172-
"id": "70473571",
173-
"metadata": {},
174-
"outputs": [],
175-
"source": [
176-
"def _match_ligand(stem: str, smi_d: dict[str, Chem.Mol]) -> str | None:\n",
177-
" \"\"\"Return the first SMILES-dict key that appears as a whole word in *stem*.\n",
178-
"\n",
179-
" Matching uses word-boundary assertions so that e.g. ``\"lig1\"`` matches\n",
180-
" the stem ``\"complex_lig1_chainA\"`` but not ``\"complex_lig10_chainA\"``.\n",
181-
"\n",
182-
" Args:\n",
183-
" stem: PDB filename stem (without extension) to search in.\n",
184-
" smi_d: Mapping of ligand names to RDKit template molecules.\n",
185-
"\n",
186-
" Returns:\n",
187-
" The matched ligand name, or ``None`` if no key matches.\n",
188-
" \"\"\"\n",
189-
" return next(\n",
190-
" (n for n in smi_d if re.search(rf\"(?<![A-Za-z\\d]){re.escape(n)}(?![A-Za-z\\d])\", stem)),\n",
191-
" None,\n",
192-
" )"
193-
]
41+
"source": "## Prepare structures from complex PDB files\n\nEach complex PDB file contains a protein and ligand. We iterate over the ligand names from the SMILES file and resolve each to its complex PDB via `complex_pdb_for`. We select the reference protein as the complex bound to the most aspherical ligand (highest Asphericity descriptor), align all other complexes to it, and assign correct bond orders from the SMILES file using `assign_topology`. The reference protein is expected to already be fixed and protonated (via `fix_pdb`) from a prior run."
19442
},
19543
{
19644
"cell_type": "code",
19745
"execution_count": null,
19846
"id": "51b8c74c",
19947
"metadata": {},
20048
"outputs": [],
201-
"source": [
202-
"# Build SMILES template dict from .smi file\n",
203-
"smi_suppl = Chem.SmilesMolSupplier(str(SMILES_PATH), sanitize=True)\n",
204-
"smiles_dict = {mol.GetProp(\"_Name\"): mol for mol in smi_suppl if mol is not None}\n",
205-
"\n",
206-
"# Pass 1: extract ligands and compute Asphericity to select reference structure\n",
207-
"ligand_asphericity = {}\n",
208-
"ligand_volume = {}\n",
209-
"for pdb_file in sorted(COMPLEX_DIR.glob(\"*.pdb\")):\n",
210-
" # Pair ligand name with PDB filename; if no match, skip\n",
211-
" lig_name = _match_ligand(pdb_file.stem, smiles_dict)\n",
212-
" if lig_name is None:\n",
213-
" continue\n",
214-
" ligand_pdb = ALIGN_DIR / f\"{lig_name}.pdb\"\n",
215-
" u = mda.Universe(pdb_file)\n",
216-
" u.select_atoms(f\"chainID {LIGAND_CHAIN_ID}\").write(ligand_pdb)\n",
217-
" rdkit_mol = Chem.MolFromPDBFile(str(ligand_pdb), sanitize=True, removeHs=True)\n",
218-
" rdkit_mol = assign_topology(rdkit_mol, smiles_dict[lig_name])\n",
219-
" ligand_asphericity[pdb_file] = Asphericity(rdkit_mol)\n",
220-
" ligand_volume[pdb_file] = AllChem.ComputeMolVolume(rdkit_mol)\n",
221-
"\n",
222-
"ref_pdb = max(ligand_asphericity, key=ligand_asphericity.get)\n",
223-
"print(f\"{'File':<40} {'Asphericity':>12} {'Volume (A^3)':>12}\")\n",
224-
"for pdb_file in sorted(ligand_asphericity):\n",
225-
" marker = \" <-- ref\" if pdb_file == ref_pdb else \"\"\n",
226-
" print(\n",
227-
" f\"{pdb_file.name:<40} {ligand_asphericity[pdb_file]:>12.4f}\"\n",
228-
" f\" {ligand_volume[pdb_file]:>12.2f}{marker}\"\n",
229-
" )"
230-
]
49+
"source": "# Build SMILES template dict from .smi file\nsmi_suppl = Chem.SmilesMolSupplier(str(SMILES_PATH), sanitize=True)\nsmiles_dict = {mol.GetProp(\"_Name\"): mol for mol in smi_suppl if mol is not None}\n\n# Pass 1: extract ligands and compute Asphericity to select reference structure.\n# Ligand names from the .smi file are the standard; complex_pdb_for resolves\n# each name to its complex PDB file.\nligand_asphericity = {}\nligand_volume = {}\nfor lig_name, template in smiles_dict.items():\n pdb_file = complex_pdb_for(lig_name)\n ligand_pdb = ALIGN_DIR / f\"{lig_name}.pdb\"\n u = mda.Universe(pdb_file)\n u.select_atoms(f\"chainID {LIGAND_CHAIN_ID}\").write(ligand_pdb)\n rdkit_mol = Chem.MolFromPDBFile(str(ligand_pdb), sanitize=True, removeHs=True)\n rdkit_mol = assign_topology(rdkit_mol, template)\n ligand_asphericity[lig_name] = Asphericity(rdkit_mol)\n ligand_volume[lig_name] = AllChem.ComputeMolVolume(rdkit_mol)\n\nref_name = max(ligand_asphericity, key=ligand_asphericity.get)\nref_pdb = complex_pdb_for(ref_name)\nprint(f\"{'Ligand':<40} {'Asphericity':>12} {'Volume (A^3)':>12}\")\nfor lig_name in sorted(ligand_asphericity):\n marker = \" <-- ref\" if lig_name == ref_name else \"\"\n print(\n f\"{lig_name:<40} {ligand_asphericity[lig_name]:>12.4f}\"\n f\" {ligand_volume[lig_name]:>12.2f}{marker}\"\n )"
23150
},
23251
{
23352
"cell_type": "code",
23453
"execution_count": null,
23554
"id": "f3bbbd84",
23655
"metadata": {},
23756
"outputs": [],
238-
"source": [
239-
"# Pass 2: align all complexes to reference, extract ligands and reference protein\n",
240-
"ref_u = mda.Universe(ref_pdb)\n",
241-
"\n",
242-
"for pdb_file in sorted(COMPLEX_DIR.glob(\"*.pdb\")):\n",
243-
" lig_name = _match_ligand(pdb_file.stem, smiles_dict)\n",
244-
" if lig_name is None:\n",
245-
" continue\n",
246-
" mob_u = mda.Universe(pdb_file)\n",
247-
" old_rmsd, new_rmsd = align.alignto(\n",
248-
" mob_u.select_atoms(\"protein and backbone\"),\n",
249-
" ref_u.select_atoms(\"protein and backbone\"),\n",
250-
" )\n",
251-
" ligand_pdb = ALIGN_DIR / f\"{lig_name}_aligned.pdb\"\n",
252-
" mob_u.select_atoms(f\"chainID {LIGAND_CHAIN_ID}\").write(ligand_pdb)\n",
253-
" rdkit_mol = Chem.MolFromPDBFile(str(ligand_pdb), removeHs=True)\n",
254-
" rdkit_mol = assign_topology(rdkit_mol, smiles_dict[lig_name])\n",
255-
" rdkit_mol.SetProp(\"_Name\", lig_name)\n",
256-
" Chem.MolToMolFile(rdkit_mol, str(LIGAND_DIR / f\"{lig_name}.sdf\"))\n",
257-
"\n",
258-
" print(f\"{lig_name} (RMSD: {old_rmsd:.3f} -> {new_rmsd:.3f} A)\")\n",
259-
" view_mol_3d(rdkit_mol, width=400, height=300, show=True)\n",
260-
" if pdb_file == ref_pdb:\n",
261-
" mob_u.select_atoms(f\"chainID {PROTEIN_CHAIN_ID}\").write(PROTEIN_PATH)"
262-
]
57+
"source": "# Pass 2: align all complexes to reference, extract ligands and reference protein\nref_u = mda.Universe(ref_pdb)\n\nfor lig_name, template in smiles_dict.items():\n pdb_file = complex_pdb_for(lig_name)\n mob_u = mda.Universe(pdb_file)\n old_rmsd, new_rmsd = align.alignto(\n mob_u.select_atoms(\"protein and backbone\"),\n ref_u.select_atoms(\"protein and backbone\"),\n )\n ligand_pdb = ALIGN_DIR / f\"{lig_name}_aligned.pdb\"\n mob_u.select_atoms(f\"chainID {LIGAND_CHAIN_ID}\").write(ligand_pdb)\n rdkit_mol = Chem.MolFromPDBFile(str(ligand_pdb), removeHs=True)\n rdkit_mol = assign_topology(rdkit_mol, template)\n rdkit_mol.SetProp(\"_Name\", lig_name)\n Chem.MolToMolFile(rdkit_mol, str(LIGAND_DIR / f\"{lig_name}.sdf\"))\n\n print(f\"{lig_name} (RMSD: {old_rmsd:.3f} -> {new_rmsd:.3f} A)\")\n view_mol_3d(rdkit_mol, width=400, height=300, show=True)\n if lig_name == ref_name:\n mob_u.select_atoms(f\"chainID {PROTEIN_CHAIN_ID}\").write(PROTEIN_PATH)"
26358
},
26459
{
26560
"cell_type": "code",

0 commit comments

Comments
 (0)