File size: 9,408 Bytes
f2ad668 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 | """CCD conformer loading utilities.
Loads idealized conformer coordinates from a CCD pickle file containing RDKit molecules.
Conformer priority follows AF3 Section 2.8: Computed > Ideal > first available.
"""
from __future__ import annotations
import os
import pickle
from pathlib import Path
import numpy as np
from huggingface_hub import hf_hub_download
from .esmfold2_constants import RES_TYPE_TO_CCD
if os.environ.get("ESMCFOLD_CCD_PATH"):
CCD_PICKLE_PATH = Path(os.environ["ESMCFOLD_CCD_PATH"])
else:
CCD_PICKLE_PATH = None
# Lazily loaded CCD dictionary
_CCD_MOLECULES: dict | None = None
# Caches
_CCD_CONFORMERS: dict[str, dict[str, np.ndarray]] = {}
_CCD_ATOM_CACHE: dict[str, list[tuple[str, str, int]]] = {}
_CCD_BONDS_CACHE: dict[str, list[tuple[str, str]]] = {}
_CCD_LEAVING_ATOMS_CACHE: dict[str, set[str]] = {}
_IDEALIZED_POS_CACHE: dict[tuple[int, str], np.ndarray | None] = {}
_LIGAND_IDEALIZED_POS_CACHE: dict[tuple[str, str], np.ndarray | None] = {}
def load_ccd(cache_dir: Path | str | None = None) -> dict:
"""Load CCD molecules from pickle file, downloading if needed.
Args:
cache_dir: Directory to cache the downloaded CCD pickle.
If None, uses CCD_PICKLE_PATH env var or downloads to ~/.cache/esmcfold/.
"""
global _CCD_MOLECULES
if _CCD_MOLECULES is not None:
return _CCD_MOLECULES
# Determine pickle path
if CCD_PICKLE_PATH is not None and CCD_PICKLE_PATH.exists():
pkl_path = CCD_PICKLE_PATH
elif cache_dir is not None:
cache_dir = Path(cache_dir)
cache_dir.mkdir(parents=True, exist_ok=True)
pkl_path = cache_dir / "ccd.pkl"
else:
try:
pkl_path = Path(
hf_hub_download(repo_id="biohub/ESMFold2", filename="ccd.pkl")
)
except Exception as e:
raise FileNotFoundError(
f"Failed to download CCD pickle file from Hugging Face repository: {e}"
)
if not pkl_path.exists():
raise FileNotFoundError(
f"CCD pickle file not found: {pkl_path}. Please set the ESMCFOLD_CCD_PATH environment variable to the path of a valid CCD pickle file or download the file from the Hugging Face repository."
)
print(f"Loading CCD dictionary from {pkl_path}")
with open(pkl_path, "rb") as f:
_CCD_MOLECULES = pickle.load(f)
if _CCD_MOLECULES is None:
_CCD_MOLECULES = {}
return _CCD_MOLECULES
def _get_ccd_molecules() -> dict:
"""Get CCD molecules, loading lazily on first call."""
global _CCD_MOLECULES
if _CCD_MOLECULES is None:
return load_ccd()
return _CCD_MOLECULES
def _get_ccd_mol_with_significant_h(comp_id: str):
"""Get CCD molecule with only chemically significant hydrogens.
Returns (mol, conformer) tuple or (None, None) if not available.
"""
ccd = _get_ccd_molecules()
if comp_id not in ccd:
return None, None
mol = ccd[comp_id]
if mol.GetNumConformers() == 0:
return None, None
# Find the "Computed" conformer (RDKit ETKDGv3), fall back to "Ideal"
conf_idx = 0
for i, c in enumerate(mol.GetConformers()):
props = c.GetPropsAsDict()
if props.get("name") == "Computed":
conf_idx = i
break
else:
for i, c in enumerate(mol.GetConformers()):
props = c.GetPropsAsDict()
if props.get("name") == "Ideal":
conf_idx = i
break
from rdkit import Chem
mol_no_h = Chem.RemoveHs(mol, sanitize=False)
if mol_no_h.GetNumConformers() == 0:
return None, None
return mol_no_h, mol_no_h.GetConformer(
min(conf_idx, mol_no_h.GetNumConformers() - 1)
)
def get_ccd_conformer(comp_id: str) -> dict[str, np.ndarray] | None:
"""Get idealized conformer as dict of atom_name -> position [3].
Conformer priority: Computed > Ideal > first available.
"""
if comp_id in _CCD_CONFORMERS:
cached = _CCD_CONFORMERS[comp_id]
return cached if cached else None
mol, conf = _get_ccd_mol_with_significant_h(comp_id)
if mol is None or conf is None:
_CCD_CONFORMERS[comp_id] = {}
return None
conformer: dict[str, np.ndarray] = {}
for atom in mol.GetAtoms():
props = atom.GetPropsAsDict()
atom_name = props.get("name")
if not isinstance(atom_name, str) or not atom_name:
continue
idx = atom.GetIdx()
pos = conf.GetAtomPosition(idx)
conformer[atom_name] = np.array([pos.x, pos.y, pos.z], dtype=np.float32)
_CCD_CONFORMERS[comp_id] = conformer
return conformer if conformer else None
def get_idealized_atom_pos(res_type: int, atom_name: str) -> np.ndarray | None:
"""Get idealized position for a standard residue atom.
Uses res_type index to look up CCD component, then returns position.
Returns None if not found.
"""
cache_key = (res_type, atom_name)
if cache_key in _IDEALIZED_POS_CACHE:
return _IDEALIZED_POS_CACHE[cache_key]
comp_id = RES_TYPE_TO_CCD.get(res_type)
if comp_id:
ccd_conformer = get_ccd_conformer(comp_id)
if ccd_conformer and atom_name in ccd_conformer:
pos = ccd_conformer[atom_name]
_IDEALIZED_POS_CACHE[cache_key] = pos
return pos
_IDEALIZED_POS_CACHE[cache_key] = None
return None
def get_ligand_idealized_atom_pos(res_name: str, atom_name: str) -> np.ndarray | None:
"""Get idealized position for a ligand/modified residue atom.
Returns None if not found.
"""
cache_key = (res_name, atom_name)
if cache_key in _LIGAND_IDEALIZED_POS_CACHE:
return _LIGAND_IDEALIZED_POS_CACHE[cache_key]
ccd_conformer = get_ccd_conformer(res_name)
if ccd_conformer and atom_name in ccd_conformer:
pos = ccd_conformer[atom_name]
_LIGAND_IDEALIZED_POS_CACHE[cache_key] = pos
return pos
_LIGAND_IDEALIZED_POS_CACHE[cache_key] = None
return None
def get_ligand_ccd_atoms_with_charges(
comp_id: str,
) -> list[tuple[str, str, int]] | None:
"""Get list of (atom_name, element, charge) for a CCD component.
Uses RDKit RemoveHs(sanitize=False) to keep chemically significant hydrogens.
Returns None if CCD data not available.
"""
if comp_id in _CCD_ATOM_CACHE:
cached = _CCD_ATOM_CACHE[comp_id]
return cached if cached else None
mol, _ = _get_ccd_mol_with_significant_h(comp_id)
if mol is None:
_CCD_ATOM_CACHE[comp_id] = []
return None
atoms: list[tuple[str, str, int]] = []
for atom in mol.GetAtoms():
props = atom.GetPropsAsDict()
atom_name = props.get("name")
if not isinstance(atom_name, str) or not atom_name:
continue
element = atom.GetSymbol()
charge = atom.GetFormalCharge()
atoms.append((atom_name, element, charge))
_CCD_ATOM_CACHE[comp_id] = atoms
return atoms if atoms else None
def get_ligand_ccd_bonds(comp_id: str) -> list[tuple[str, str]] | None:
"""Get list of (atom1_name, atom2_name) bonds for a CCD component.
Returns None if CCD data not available.
"""
if comp_id in _CCD_BONDS_CACHE:
cached = _CCD_BONDS_CACHE[comp_id]
return cached if cached else None
mol, _ = _get_ccd_mol_with_significant_h(comp_id)
if mol is None:
_CCD_BONDS_CACHE[comp_id] = []
return None
# Get included atom names
included_atoms = set()
for atom in mol.GetAtoms():
props = atom.GetPropsAsDict()
atom_name = props.get("name")
if isinstance(atom_name, str) and atom_name:
included_atoms.add(atom_name)
bonds: list[tuple[str, str]] = []
for bond in mol.GetBonds():
a1 = bond.GetBeginAtom()
a2 = bond.GetEndAtom()
n1 = a1.GetPropsAsDict().get("name")
n2 = a2.GetPropsAsDict().get("name")
if (
isinstance(n1, str)
and isinstance(n2, str)
and n1
and n2
and n1 in included_atoms
and n2 in included_atoms
):
bonds.append((n1, n2))
_CCD_BONDS_CACHE[comp_id] = bonds
return bonds if bonds else None
def get_ccd_leaving_atoms(comp_id: str) -> set[str]:
"""Get set of atom names marked as leaving atoms in CCD.
Leaving atoms are removed during polymerization (e.g., OP3 in nucleotides).
"""
if comp_id in _CCD_LEAVING_ATOMS_CACHE:
return _CCD_LEAVING_ATOMS_CACHE[comp_id]
ccd = _get_ccd_molecules()
if comp_id not in ccd:
_CCD_LEAVING_ATOMS_CACHE[comp_id] = set()
return set()
mol = ccd[comp_id]
leaving_atoms = set()
for atom in mol.GetAtoms():
if atom.HasProp("leaving_atom"):
if atom.GetProp("leaving_atom") == "1":
name = atom.GetProp("name") if atom.HasProp("name") else ""
if name:
leaving_atoms.add(name)
_CCD_LEAVING_ATOMS_CACHE[comp_id] = leaving_atoms
return leaving_atoms
|