foldenv 0.1.1__py3-none-any.whl

foldenv/__init__.py

@@ -0,0 +1,50 @@
+"""foldenv — `get_structural_context(uniprot_id, position)`.
+
+A residue's physical neighborhood in the AlphaFold-predicted fold — RSA, secondary
+structure, packing/contacts, a PLM embedding, and a pLDDT confidence flag — to help
+downstream tools separate buried structural residues (spandrels) from functional surface
+residues.
+
+Decisions and their rationale live in `decisions.yaml`.
+"""
+
+__version__ = "0.1.0"
+import os as _os
+import sys as _sys
+
+# PLM embedding on Apple MPS relies on CPU fallback for a few ops the Metal backend does
+# not implement. This must be set before torch initializes its MPS backend, so we do it at
+# import — but only on macOS (where MPS exists) and only via setdefault, so non-Mac hosts
+# and callers who never touch torch see no global env mutation.
+if _sys.platform == "darwin":
+    _os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
+
+from . import analysis, config, tool, validation
+from .context import (
+    clear_cache,
+    get_contacts,
+    get_dssp,
+    get_sequence,
+    get_structural_context,
+    get_structure,
+    structural_profile,
+)
+from .fetch import NoAlphaFoldModelError, StructureRecord, fetch_structure
+
+__all__ = [
+    "__version__",
+    "config",
+    "validation",
+    "analysis",
+    "tool",
+    "get_structural_context",
+    "get_structure",
+    "get_dssp",
+    "get_contacts",
+    "get_sequence",
+    "structural_profile",
+    "clear_cache",
+    "fetch_structure",
+    "StructureRecord",
+    "NoAlphaFoldModelError",
+]

foldenv/analysis.py

@@ -0,0 +1,123 @@
+"""Biological analysis: do RSA + contacts separate functional from structural?
+
+For a protein and a set of known functional residues, report each site's burial (RSA) and
+packing (contact_count) together with its **percentile rank** within the protein, plus a
+descriptive verdict. This is deliberately descriptive, not a benchmark — curated functional
+sets are tiny (a handful of residues), so percentiles and medians are more honest than an AUC.
+
+Key finding this exists to make explicit: enzyme catalytic
+residues sit in **buried active-site clefts**, so they look structurally like buried-structural
+"spandrels" on RSA/contacts alone — the structural axis is *complementary* to conservation,
+not a standalone functional-site classifier.
+"""
+from __future__ import annotations
+
+import bisect
+from dataclasses import dataclass
+
+from . import config as _config
+from . import context as _ctx
+
+# Known functional residues in UniProt numbering (map literature/Ambler numbers first!).
+# Example data for the descriptive helper below, not a curated benchmark — pass `sites=` for
+# anything else. The three TP53 entries are not one category, and the difference matters here.
+#
+# TEM-1 β-lactamase (P62593): the class A catalytic triad, S70/K73/E166 in *Ambler* numbering
+# (Ambler et al. 1991) → UniProt 68/71/164. Ambler numbering carries insertions relative to the
+# sequence, so the offset is not derivable from the sequence alone; the identity self-check in
+# functional_site_stats is what stops a wrong one from reporting a neighbouring residue.
+#
+# TP53 (P04637): already UniProt numbering. **R248 and R273 are DNA-contact residues; R175 is
+# not** — it is a conformational mutant that distorts the DNA-binding domain fold rather than
+# losing a contact (the contact-vs-structural split from the DNA-binding domain crystal
+# structure). All three are cancer *hotspots*, which is a mutation-frequency claim from the TP53
+# mutation databases rather than a structural one.
+#
+# R175 therefore sits on both sides of the comparison this module exists to make: the finding in
+# the docstring is that catalytic residues look structurally like buried spandrels, and R175 is a
+# functional site whose mechanism *is* structural. Kept, because it is a genuine hotspot and this
+# helper is descriptive — but read its row knowing that, and do not treat the set as three
+# equivalent "DNA-contact" sites.
+KNOWN_FUNCTIONAL_SITES: dict[str, dict[int, str]] = {
+    "P62593": {68: "S", 71: "K", 164: "E"},
+    "P04637": {175: "R", 248: "R", 273: "R"},
+}
+
+
+@dataclass
+class SiteStat:
+    position: int
+    aa: str
+    rsa: float | None
+    rsa_percentile: float | None    # fraction of residues as-or-more buried (RSA ≤ this)
+    contact_count: int
+    contact_percentile: float       # fraction of residues with ≤ this many contacts
+    buried: bool                    # rsa < config buried_threshold
+
+
+def _percentile(sorted_vals: list, v) -> float:
+    """Fraction of values ≤ v (0..1); ties inclusive (bisect_right)."""
+    return bisect.bisect_right(sorted_vals, v) / len(sorted_vals) if sorted_vals else float("nan")
+
+
+def functional_site_stats(
+    uniprot_id: str,
+    sites: dict[int, str] | None = None,
+    cfg: dict | None = None,
+) -> list[SiteStat]:
+    """Structural signature (RSA + contacts, with percentile ranks) for known functional sites."""
+    cfg = cfg or _config.load()
+    sites = sites if sites is not None else KNOWN_FUNCTIONAL_SITES.get(uniprot_id, {})
+    if not sites:
+        raise ValueError(f"no known functional sites for {uniprot_id}; pass `sites=`")
+
+    profile = _ctx.structural_profile(uniprot_id, cfg)
+    rsas = sorted(p["rsa"] for p in profile.values() if p["rsa"] is not None)
+    ccs = sorted(p["contact_count"] for p in profile.values())
+    buried_thr = cfg["rsa"]["buried_threshold"]
+
+    out = []
+    for pos, expect in sorted(sites.items()):
+        if pos not in profile:
+            raise ValueError(
+                f"position {pos} out of range 1..{len(profile)} for {uniprot_id}"
+            )
+        p = profile[pos]
+        # Self-check the numbering mapping: the residue in the structure must match the
+        # expected identity. Guards against a wrong Ambler→UniProt offset silently reporting
+        # a neighboring residue (the central M7 risk). Pass expect="" / "X" to skip.
+        if expect and expect not in ("X", "?") and p["aa"] != expect:
+            raise ValueError(
+                f"{uniprot_id} site {pos}: expected {expect} but structure has {p['aa']} — "
+                "check the functional-site numbering (literature vs UniProt)."
+            )
+        rsa = p["rsa"]
+        out.append(
+            SiteStat(
+                position=pos,
+                aa=p["aa"],
+                rsa=rsa,
+                rsa_percentile=_percentile(rsas, rsa) if rsa is not None else None,
+                contact_count=p["contact_count"],
+                contact_percentile=_percentile(ccs, p["contact_count"]),
+                buried=(rsa is not None and rsa < buried_thr),
+            )
+        )
+    return out
+
+
+def summarize(uniprot_id: str, cfg: dict | None = None) -> dict:
+    """Descriptive summary: are the known functional sites buried, and how packed vs typical?"""
+    stats = functional_site_stats(uniprot_id, cfg=cfg)
+    n_buried = sum(s.buried for s in stats)
+    present_rsa = [s.rsa for s in stats if s.rsa is not None]
+    return {
+        "uniprot_id": uniprot_id,
+        "n_sites": len(stats),
+        "n_buried": n_buried,
+        "all_buried": n_buried == len(stats),
+        # mean over sites that actually have a DSSP RSA (don't fold None → 0.0)
+        "mean_rsa": (sum(present_rsa) / len(present_rsa)) if present_rsa else float("nan"),
+        "mean_contact_percentile": sum(s.contact_percentile for s in stats) / len(stats),
+        "sites": stats,
+    }

foldenv/config.py

@@ -0,0 +1,68 @@
+"""Load the locked decisions (D1–D6) from `decisions.yaml`.
+
+The YAML holds the reasonable defaults; `load()` returns a nested dict and lets callers
+override any leaf without touching the file. Keeping the decisions in data (not code) means
+a results writeup can point at one file for "which cutoff / table / mask did you use?".
+"""
+from __future__ import annotations
+
+import copy
+import os
+from pathlib import Path
+from typing import Any, Mapping
+
+import yaml
+
+_DECISIONS_PATH = Path(__file__).with_name("decisions.yaml")
+
+
+def _default_cache_dir() -> Path:
+    """Default on-disk cache location, resolved fresh on each call.
+
+    Overridable via the ``FOLDENV_CACHE_DIR`` env var or the ``cache.dir`` config leaf.
+    Defaults under the current working directory so the package works the same whether
+    installed into site-packages or run from a checkout. Resolving here (not at import
+    time) means the env var and the working directory are read when ``load()`` runs, not
+    when the module is first imported.
+    """
+    return Path(os.environ.get("FOLDENV_CACHE_DIR", Path.cwd() / ".foldenv_cache"))
+
+
+def _deep_update(base: dict, overrides: Mapping[str, Any]) -> dict:
+    """Recursively merge `overrides` into `base` (mutating and returning `base`)."""
+    for key, value in overrides.items():
+        if isinstance(value, Mapping) and isinstance(base.get(key), dict):
+            _deep_update(base[key], value)
+        else:
+            base[key] = value
+    return base
+
+
+def load(overrides: Mapping[str, Any] | None = None) -> dict:
+    """Return the decisions dict, with `cache.dir` resolved to a concrete path.
+
+    Args:
+        overrides: nested dict merged over the file defaults, e.g.
+            ``{"contacts": {"primary": "cb"}}``.
+    """
+    with open(_DECISIONS_PATH) as f:
+        cfg = yaml.safe_load(f)
+    if overrides:
+        _deep_update(cfg, copy.deepcopy(dict(overrides)))
+
+    # Resolve the caching default here so downstream code always sees a real directory.
+    if cfg["cache"].get("dir") is None:
+        cfg["cache"]["dir"] = str(_default_cache_dir())
+
+    # Allow the AlphaFold API base to be overridden by env (e.g. to point tests at an
+    # unreachable host for deterministic, offline CI). An explicit `overrides` value wins.
+    if (overrides is None or "alphafold" not in overrides) and os.environ.get(
+        "FOLDENV_ALPHAFOLD_API_BASE"
+    ):
+        cfg["alphafold"]["api_base"] = os.environ["FOLDENV_ALPHAFOLD_API_BASE"]
+    return cfg
+
+
+def decisions_path() -> Path:
+    """Path to the backing YAML (for logging which config a run used)."""
+    return _DECISIONS_PATH

foldenv/constants.py

@@ -0,0 +1,68 @@
+"""Amino-acid tables and the PLM-encoder registry.
+
+Vendored so ``foldenv`` has no import dependency on any external package for these
+facts. The amino-acid lookups are standard reference data; ``PLM_ENCODERS`` maps each
+supported embedding backend to its HuggingFace Hub id.
+"""
+
+# Single letter, three letter, and full amino acid names.
+aa_names = (
+    ('A', 'ALA', 'alanine'),
+    ('R', 'ARG', 'arginine'),
+    ('N', 'ASN', 'asparagine'),
+    ('D', 'ASP', 'aspartic acid'),
+    ('C', 'CYS', 'cysteine'),
+    ('E', 'GLU', 'glutamic acid'),
+    ('Q', 'GLN', 'glutamine'),
+    ('G', 'GLY', 'glycine'),
+    ('H', 'HIS', 'histidine'),
+    ('I', 'ILE', 'isoleucine'),
+    ('L', 'LEU', 'leucine'),
+    ('K', 'LYS', 'lysine'),
+    ('M', 'MET', 'methionine'),
+    ('F', 'PHE', 'phenylalanine'),
+    ('P', 'PRO', 'proline'),
+    ('S', 'SER', 'serine'),
+    ('T', 'THR', 'threonine'),
+    ('W', 'TRP', 'tryptophan'),
+    ('Y', 'TYR', 'tyrosine'),
+    ('V', 'VAL', 'valine'),
+    # Extended AAs
+    ('B', 'ASX', 'asparagine or aspartic acid'),
+    ('Z', 'GLX', 'glutamine or glutamic acid'),
+    ('X', 'XAA', 'Any'),
+    ('J', 'XLE', 'Leucine or isoleucine'),
+)
+
+# Indices of standard amino acids in `aa_names`.
+standard_indices = tuple(range(20))
+
+# Single letter codes of standard amino acids.
+standard_aas = tuple(aa_names[i][0] for i in standard_indices)
+AAs = tuple(sorted(standard_aas))
+
+# aa_to_idx and idx_to_aa
+aa2idx = dict(zip(AAs, standard_indices))
+idx2aa = {v: k for k, v in aa2idx.items()}
+
+# dictionaries for aa name conversion
+one2three = dict(aa_names[i][:2] for i in standard_indices)
+three2one = {v: k for k, v in one2three.items()}
+
+
+# PLM encoders → HuggingFace Hub ids
+PLM_ENCODERS = {
+    "esm": "facebook/esm2_t36_3B_UR50D",
+    "ankh": "ElnaggarLab/ankh-large",
+    "esm_35M": "facebook/esm2_t12_35M_UR50D",
+    "esm_650M": "facebook/esm2_t33_650M_UR50D",
+    "ankh_base": "ElnaggarLab/ankh-base",
+    "protbert": "Rostlab/prot_bert",
+    "prott5_xl_half": "Rostlab/prot_t5_xl_half_uniref50-enc",
+    "prostt5": "Rostlab/ProstT5",
+    "esmc_6b": "EvolutionaryScale/esmc-6b-2024-12",  # ESM Cambrian 6B, 2560-dim (needs transformers>=4.57)
+    "ankh3_large": "ElnaggarLab/ankh3-large",        # Ankh3-large, T5 encoder, 1536-dim
+    "ankh3_xl": "ElnaggarLab/ankh3-xl",              # Ankh3-XL, T5 encoder, 2560-dim
+    "saprot": "westlake-repl/SaProt_650M_AF2",       # SaProt 650M, ESM2-650M arch + SA (AA+3Di) vocab, 1280-dim
+    "saprot_1.3b": "westlake-repl/SaProt_1.3B_AFDB_OMG_NCBI",  # SaProt 1.3B, ESM arch, 1280-dim
+}

foldenv/contacts.py

@@ -0,0 +1,170 @@
+"""Contacts via KD-tree (Biopython NeighborSearch).
+
+For a queried residue: count neighbors under the D1 cutoff (Cα–Cα ≤ 8 Å primary, or Cβ–Cβ ≤
+5 Å), excluding self, and return the top-N nearest as `nearest_contacts`. Partners whose
+pLDDT is below the D2 mask threshold are dropped as candidates (AF disordered loops otherwise
+produce spurious contacts). Glycine has no Cβ → fall back to Cα (config
+`contacts.glycine_cb_fallback`). pLDDT lives in the AF B-factor column.
+
+`contact_count` is a burial/packing proxy: buried residues ~10–20 Cα-neighbors, surface ~4–8.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from .constants import three2one
+
+
+@dataclass
+class Contact:
+    resnum: int          # partner residue number (UniProt/AF numbering)
+    aa: str              # 1-letter amino acid
+    distance: float      # Å between representative atoms
+
+
+@dataclass
+class ContactResult:
+    contact_count: int                 # neighbors under cutoff (self excluded, mask applied)
+    nearest_contacts: list[Contact]    # up to n_nearest, ascending distance
+    atom_mode: str                     # "CA" or "CB" — which atom actually anchored the query
+    cutoff: float                      # Å cutoff used
+
+
+def _one_letter(resname: str) -> str:
+    return three2one.get(resname.strip().upper(), "X")
+
+
+def _plddt(residue) -> float | None:
+    """AF pLDDT for a residue = its B-factor (shared across the residue's atoms)."""
+    if "CA" in residue:
+        return float(residue["CA"].get_bfactor())
+    for atom in residue:
+        return float(atom.get_bfactor())
+    return None
+
+
+def _rep_atom(residue, mode: str, glycine_cb_fallback: str):
+    """Representative atom for contacts: CA in ca-mode; CB in cb-mode (→ CA for glycine /
+    any residue missing CB when `glycine_cb_fallback == 'ca'`). Returns (atom, atom_name)."""
+    if mode == "ca":
+        return (residue["CA"], "CA") if "CA" in residue else (None, None)
+    if "CB" in residue:
+        return residue["CB"], "CB"
+    if glycine_cb_fallback == "ca" and "CA" in residue:
+        return residue["CA"], "CA"
+    return None, None
+
+
+def _is_amino_acid(residue) -> bool:
+    """Standard residue (not HETATM/water). AF models are all-standard, but be defensive."""
+    return not residue.id[0].strip()
+
+
+def _chain(structure, chain_id: str | None):
+    model = next(structure.get_models())
+    return model[chain_id] if chain_id is not None else next(model.get_chains())
+
+
+def build_contact_index(
+    structure,
+    *,
+    chain_id: str | None = None,
+    mode: str = "ca",
+    plddt_mask_below: float = 50.0,
+    glycine_cb_fallback: str = "ca",
+):
+    """One KD-tree of representative atoms for a chain, reusable across positions.
+
+    The tree covers every standard residue passing the pLDDT mask (self is NOT excluded
+    here — `compute_contacts` filters the query residue out per call by object identity). The
+    agent queries many positions of one protein, so caching this (per protein+mode+mask) turns
+    per-position work from O(N log N) tree builds into a single search. Returns a Biopython
+    `NeighborSearch`, or None if no atoms qualify.
+    """
+    from Bio.PDB import NeighborSearch
+
+    atoms = []
+    for r in _chain(structure, chain_id):
+        if not _is_amino_acid(r):
+            continue
+        plddt = _plddt(r)
+        if plddt is not None and plddt < plddt_mask_below:
+            continue
+        atom, _ = _rep_atom(r, mode, glycine_cb_fallback)
+        if atom is not None:
+            atoms.append(atom)
+    return NeighborSearch(atoms) if atoms else None
+
+
+def compute_contacts(
+    structure,
+    resnum: int,
+    *,
+    chain_id: str | None = None,
+    mode: str = "ca",
+    ca_cutoff: float = 8.0,
+    cb_cutoff: float = 5.0,
+    n_nearest: int = 5,
+    plddt_mask_below: float = 50.0,
+    glycine_cb_fallback: str = "ca",
+    index=None,
+) -> ContactResult:
+    """Contacts of residue `resnum` (1-based UniProt numbering) in one chain.
+
+    Args:
+        mode: "ca" (Cα–Cα ≤ ca_cutoff) or "cb" (Cβ–Cβ ≤ cb_cutoff).
+        plddt_mask_below: drop partner residues with pLDDT below this (D2).
+        chain_id: chain to search; None → first chain (AF monomers are single-chain).
+        index: a prebuilt `build_contact_index` result to reuse (must share mode / mask /
+            glycine_cb_fallback). If None, a one-off tree is built for this call.
+
+    Note: in cb-mode the glycine fallback measures a glycine partner at its Cα, so
+    `nearest_contacts` distances can mix Cβ–Cβ and Cβ–Cα; `atom_mode` reflects only the query
+    atom. Default config is ca-mode, where this does not arise.
+
+    Raises:
+        KeyError: `resnum` not present in the chain.
+        ValueError: query residue lacks a representative atom, or bad `mode`.
+    """
+    if mode not in ("ca", "cb"):
+        raise ValueError(f"mode must be 'ca' or 'cb', got {mode!r}")
+
+    chain = _chain(structure, chain_id)
+    cutoff = ca_cutoff if mode == "ca" else cb_cutoff
+
+    query_res = next(
+        (r for r in chain if _is_amino_acid(r) and r.id[1] == resnum), None
+    )
+    if query_res is None:
+        raise KeyError(f"residue {resnum} not found in chain {chain.id}")
+    q_atom, q_name = _rep_atom(query_res, mode, glycine_cb_fallback)
+    if q_atom is None:
+        raise ValueError(f"residue {resnum} has no {mode.upper()} atom for contacts")
+
+    if index is None:
+        index = build_contact_index(
+            structure, chain_id=chain_id, mode=mode,
+            plddt_mask_below=plddt_mask_below, glycine_cb_fallback=glycine_cb_fallback,
+        )
+
+    contacts: list[Contact] = []
+    if index is not None:
+        for atom in index.search(q_atom.coord, cutoff, level="A"):
+            r = atom.get_parent()
+            # Exclude self by residue id (het, resnum, icode) rather than object identity: the
+            # cached KD-tree can be built from a different structure *instance* than query_res
+            # (e.g. cache.in_memory=false), so `is` would miss it and count self as a contact.
+            # Comparing .id still distinguishes insertion codes within the chain.
+            if r.id == query_res.id:
+                continue
+            contacts.append(
+                Contact(resnum=r.id[1], aa=_one_letter(r.resname), distance=float(atom - q_atom))
+            )
+        contacts.sort(key=lambda c: c.distance)
+
+    return ContactResult(
+        contact_count=len(contacts),
+        nearest_contacts=contacts[:n_nearest],
+        atom_mode=q_name,
+        cutoff=cutoff,
+    )