xenosite-predict 0.2.1__py3-none-any.whl

xenosite/predict/__init__.py

@@ -0,0 +1,65 @@
+"""Public user API for ``xenosite.predict``.
+
+Call :func:`predict` with a SMILES string or an existing :class:`Molecule`.
+Results append; parse/canonicalize happens once when several models run.
+
+See the package README for backends, environment variables, and versions.
+"""
+
+from .api import list_models, predict
+from .parallel import apredict, apredict_many, predict_many
+from .registry import register_model
+from .errors import (
+    BackendNotConfigured,
+    InvalidMolecule,
+    ModelNotAvailable,
+    OpenBabelNotAvailable,
+    UnknownModel,
+    WeightsDownloadError,
+    WeightsNotFound,
+)
+from .weights import download_weights, ensure_weights
+from .types import (
+    AtomBondResult,
+    AtomResult,
+    Atoms,
+    BondResult,
+    Bonds,
+    Metabolite,
+    ModelResult,
+    MolAtomPairResult,
+    MolAtomResult,
+    MolBondResult,
+    Molecule,
+    Result,
+)
+
+__all__ = [
+    "predict",
+    "predict_many",
+    "apredict",
+    "apredict_many",
+    "list_models",
+    "register_model",
+    "Molecule",
+    "Atoms",
+    "Bonds",
+    "Result",
+    "MolBondResult",
+    "MolAtomResult",
+    "MolAtomPairResult",
+    "AtomResult",
+    "BondResult",
+    "AtomBondResult",
+    "Metabolite",
+    "ModelResult",
+    "InvalidMolecule",
+    "UnknownModel",
+    "BackendNotConfigured",
+    "WeightsNotFound",
+    "WeightsDownloadError",
+    "ModelNotAvailable",
+    "OpenBabelNotAvailable",
+    "download_weights",
+    "ensure_weights",
+]

xenosite/predict/__main__.py

@@ -0,0 +1,51 @@
+"""``python -m xenosite.predict download`` — fetch ONNX weights via env URL."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+from .errors import WeightsDownloadError
+from .weights import ENV_ONNX_URL, download_weights
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(
+        prog="python -m xenosite.predict",
+        description=(
+            "Download ONNX inference graphs. The archive URL comes from "
+            f"{ENV_ONNX_URL} unless --url is passed."
+        ),
+    )
+    sub = parser.add_subparsers(dest="cmd", required=True)
+    dl = sub.add_parser("download", help="download and extract ONNX weights")
+    dl.add_argument(
+        "--url",
+        default=None,
+        help=f"tarball URL or local path (default: ${ENV_ONNX_URL})",
+    )
+    dl.add_argument(
+        "--dest",
+        type=Path,
+        default=None,
+        help="extract directory (default: cache or $XENOSITE_MODELS_WEIGHTS)",
+    )
+    dl.add_argument(
+        "--force",
+        action="store_true",
+        help="re-download even if *.onnx files are already present",
+    )
+    args = parser.parse_args(argv)
+    if args.cmd == "download":
+        try:
+            dest = download_weights(url=args.url, dest=args.dest, force=args.force)
+        except WeightsDownloadError as exc:
+            print(exc, file=sys.stderr)
+            return 1
+        print(dest)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

xenosite/predict/_private.py

@@ -0,0 +1,36 @@
+"""Private integration API for in-repo callers (e.g. ``xenosite-api``).
+
+Import as ``xenosite.predict._private``. Not part of the stable public API —
+signature and exports may change without a major version bump.
+
+Forest metabolite attachment is implemented once in :func:`~xenosite.predict.forest.attach_metabolites`.
+This module re-exports that function (as ``add_metabolites``) plus discovery helpers.
+
+Typical **xenosite-api** usage after legacy HTTP adapters populate scores::
+
+    from xenosite.predict.types import Molecule as PredictMolecule
+    from xenosite.predict._private import add_metabolites, metabolite_supported
+
+    mol = PredictMolecule.model_validate(api_molecule.model_dump())
+    if metabolites_requested:
+        add_metabolites(mol, mapped_smiles=mapped_smiles)
+    # copy ``result.metabolite`` (and optional map fields) back onto the API model
+"""
+
+from __future__ import annotations
+
+from .forest import (
+    attach_metabolites,
+    metabolite_supported,
+    supported_metabolite_models,
+)
+
+# Alias for service layers; same function object as attach_metabolites.
+add_metabolites = attach_metabolites
+
+__all__ = [
+    "add_metabolites",
+    "attach_metabolites",
+    "metabolite_supported",
+    "supported_metabolite_models",
+]

xenosite/predict/api.py

@@ -0,0 +1,147 @@
+"""``predict`` / ``list_models`` — backend-agnostic user API."""
+
+from __future__ import annotations
+
+from typing import Any, Iterable, Mapping, Optional, Union
+
+from .backends import PredictBackend, resolve_backend, resolve_for_model
+from .errors import BackendNotConfigured
+from .molecule import as_molecule
+from .registry import Spec, ensure_builtins, load_runner, normalize_models, registered
+from ._private import add_metabolites
+from .types import Molecule
+
+ModelsArg = Union[str, Spec, Iterable[str | Spec]]
+BackendArg = Union[str, PredictBackend, None]
+BackendMap = Optional[Mapping[Spec, str | PredictBackend]]
+
+
+def predict(
+    inp: str | Molecule,
+    model: Optional[str] = None,
+    models: Optional[ModelsArg] = None,
+    *,
+    backend: BackendArg = None,
+    backends: BackendMap = None,
+    env: Optional[Mapping[str, str]] = None,
+    metabolites: bool = False,
+    metabolites_min_score: Optional[float] = None,
+    mapped_smiles: bool = False,
+    _parameter: Optional[Mapping[str, Any]] = None,
+) -> Molecule:
+    """Run one or more models and return a :class:`Molecule` with appended results.
+
+    Parameters
+    ----------
+    inp:
+        SMILES string or an existing :class:`Molecule` (results are appended).
+    model:
+        Single model name (default version). Ignored if ``models`` is set.
+    models:
+        Names and/or ``(name, version)`` pairs. Parse/canonicalize once.
+    backend:
+        Pin every model in this call: ``"onnx"``, ``"http"``, ``"legacy"``,
+        a URL, or a :class:`PredictBackend`. ``None`` uses the env picker.
+    backends:
+        Per-``(name, version)`` override (wins over ``backend``).
+    env:
+        Environment mapping for the picker. ``None`` uses ``os.environ``.
+        Tests should pass ``env={}`` or rely on the autouse clearer.
+    metabolites:
+        When ``True``, attach every metabolite structure the forest ruleset
+        generates for the substrate, annotated with predictor site scores and
+        sorted by score (descending). See ``xenosite.predict.forest``.
+    metabolites_min_score:
+        When set, drop metabolites whose site score is below this threshold.
+        Default ``None`` includes all forest products.
+    mapped_smiles:
+        When ``True`` (with ``metabolites=True``), add ``mapped_smiles`` to each
+        forest metabolite — canonical SMILES with ``:N`` atom-map numbers tracing
+        heavy atoms back to the parent (1-based; new atoms unmapped). ``map_idx``
+        is always populated when forest metabolites are attached.
+    _parameter:
+        Internal per-call options (not part of the public HTTP API). Runners
+        read ``molecule._parameter``; e.g. ``ndealk_site_mode`` is ``legacy``
+        for golden parity tests and ``principled`` (default) for production;
+        ``quinone_omp_mode`` is ``legacy`` (deterministic sorted BFS) for golden
+        tests and ``principled`` (any qualifying tied shortest path; binary) for
+        production; ``mean`` averages all shortest-path indicators (fractional);
+        ``symmetry_group_mode`` is ``openbabel`` for golden parity and ``rdkit``
+        (default) for production bond-class deduplication and score pooling
+        (mean of active scores per class); ``bond_nrings_mode`` is ``legacy``
+        (DFS back-edge atom counts) for golden/ob dumps and ``principled``
+        (RDKit ``RingInfo.NumAtomRings`` per BondTD endpoint) for production.
+
+    Notes
+    -----
+    One molecule at a time (no batch API). For many molecules use
+    :func:`xenosite.predict.predict_many` or :func:`xenosite.predict.apredict_many`.
+    Import does not open ONNX, HTTP, or OpenBabel.
+    First use downloads ONNX weights when XENOSITE_ONNX_URL is set.
+    """
+    ensure_builtins()
+    if models is None:
+        models = model
+    specs = normalize_models(models)
+    _, molecule = as_molecule(inp)
+    if _parameter:
+        molecule._parameter = dict(_parameter)
+
+    for spec in specs:
+        be = resolve_for_model(spec, backend=backend, backends=backends, env=env)
+        runner = load_runner(*spec)
+        runner.predict_molecule(molecule, be)
+    if metabolites:
+        add_metabolites(
+            molecule,
+            min_score=metabolites_min_score,
+            mapped_smiles=mapped_smiles,
+        )
+    return molecule
+
+
+def list_models(
+    *,
+    backend: BackendArg = None,
+    env: Optional[Mapping[str, str]] = None,
+) -> list[dict]:
+    """What this process can actually run (backend-aware), not a fictional union.
+
+    Each item is ``{"name", "version", "available", "backend", "reason"}``.
+    """
+    ensure_builtins()
+    try:
+        be = resolve_backend(backend, env=env)
+        available = set(be.available_models())
+        bname = be.name
+    except BackendNotConfigured as exc:
+        available = set()
+        bname = None
+        default_reason = str(exc)
+    else:
+        default_reason = ""
+
+    out = []
+    for info in registered():
+        spec = (info.name, info.version)
+        ok = spec in available and not info.blocked_reason
+        reason = info.blocked_reason or ("" if ok else (default_reason or "not on this backend"))
+        if ok and bname == "onnx" and info.name not in ("bioactivation",):
+            from .features import _ob
+
+            if not _ob.installed():
+                ok = False
+                reason = "OpenBabel is required for ONNX descriptors (uv add openbabel)"
+        out.append(
+            {
+                "name": info.name,
+                "version": info.version,
+                "available": bool(ok),
+                "backend": bname,
+                "reason": reason,
+                "heads": list(info.heads),
+                "two_stage": info.two_stage,
+                "pipeline": info.pipeline,
+            }
+        )
+    return out

xenosite/predict/backends/__init__.py

@@ -0,0 +1,133 @@
+"""Backend protocol, picker, and env handling.
+
+Picker (explicit env wins; first match):
+
+1. ``XENOSITE_BACKEND`` is an ``http://`` / ``https://`` URL → HTTP backend
+   (deployed xenosite-api). Optional ``XENOSITE_API_KEY`` as Bearer.
+2. Else ``XENOSITE_MODELS_WEIGHTS`` → local ONNX directory.
+3. Else auto-detect ``./weights/onnx/v0`` (or a flat ``./weights/onnx`` tree) → local ONNX.
+4. Else user cache (``$XDG_CACHE_HOME/xenosite/onnx/v0``) if ``*.onnx`` exist.
+5. Else, when ``XENOSITE_ONNX_URL`` is set, fetch that archive into the cache
+   (an INFO line reports when weights are found or downloaded).
+6. Else raise :class:`BackendNotConfigured`.
+
+The archive URL is never compiled into this package; set ``XENOSITE_ONNX_URL``.
+Tests must pass ``backend=`` / ``env={}`` and must not inherit a developer
+shell. ``conftest.py`` clears ``XENOSITE_*`` unless a test opts in.
+An isolated ``env={}`` does not auto-download.
+
+Per-``(model, version)`` override: ``predict(..., backend=...)`` applies to all
+models in the call; ``predict(..., backends={(name, version): backend})`` pins
+individual models (so ONNX epoxidation can coexist with HTTP bioactivation).
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any, Mapping, Optional, Protocol, runtime_checkable
+
+from ..errors import BackendNotConfigured
+from ..weights import ENV_ONNX_URL, ENV_WEIGHTS, resolve_onnx_dir
+
+ENV_BACKEND = "XENOSITE_BACKEND"
+ENV_API_KEY = "XENOSITE_API_KEY"
+ENV_LEGACY_URL = "XENOSITE_LEGACY_TEST_URL"
+
+Spec = tuple[str, str]
+
+
+@runtime_checkable
+class PredictBackend(Protocol):
+    """Swappable backend. Native results need not look like :class:`Molecule`."""
+
+    name: str
+
+    def available_models(self) -> list[Spec]:
+        """``(name, version)`` pairs this process can actually run on this backend."""
+
+    def predict_native(self, smiles: str, model: str, version: str) -> Any:
+        """Run one model; return backend-native output for the adapter."""
+
+
+def is_url(value: str) -> bool:
+    return value.startswith("http://") or value.startswith("https://")
+
+
+def resolve_backend(
+    backend: Optional[str | PredictBackend] = None,
+    *,
+    env: Optional[Mapping[str, str]] = None,
+    cwd: Optional[Path] = None,
+    auto_download: Optional[bool] = None,
+) -> PredictBackend:
+    """Resolve a backend. ``env=None`` uses ``os.environ``; tests should pass a dict.
+
+    ``auto_download`` defaults on only when ``env is None`` and
+    ``XENOSITE_ONNX_URL`` is set. Isolated ``env`` mappings do not fetch.
+    """
+    if isinstance(backend, PredictBackend) and not isinstance(backend, str):
+        return backend
+
+    from .http import HttpBackend
+    from .legacy import LegacyTestBackend
+    from .onnx import OnnxBackend
+
+    if isinstance(backend, str):
+        if is_url(backend):
+            return HttpBackend(backend)
+        key = backend.lower()
+        if key in {"onnx", "local"}:
+            weights = resolve_onnx_dir(
+                env=env, cwd=cwd, auto_download=auto_download, fallback=True
+            )
+            if weights is None:
+                raise BackendNotConfigured(
+                    "backend='onnx' needs local *.onnx files, "
+                    f"{ENV_WEIGHTS}, or {ENV_ONNX_URL} (auto-downloaded on first use)."
+                )
+            return OnnxBackend(weights)
+        if key in {"legacy", "legacy-test", "test-api"}:
+            url = (env or os.environ).get(ENV_LEGACY_URL, "http://127.0.0.1:8099")
+            return LegacyTestBackend(url)
+        if key == "http":
+            e = env if env is not None else os.environ
+            url = e.get(ENV_BACKEND, "")
+            if not is_url(url):
+                raise BackendNotConfigured(
+                    "backend='http' requires XENOSITE_BACKEND to be an http(s) URL"
+                )
+            return HttpBackend(url, api_key=e.get(ENV_API_KEY))
+        raise BackendNotConfigured(f"Unknown backend {backend!r}")
+
+    e = dict(os.environ if env is None else env)
+    cwd = cwd or Path.cwd()
+
+    url = e.get(ENV_BACKEND, "").strip()
+    if is_url(url):
+        return HttpBackend(url, api_key=e.get(ENV_API_KEY))
+
+    weights = resolve_onnx_dir(env=env, cwd=cwd, auto_download=auto_download)
+    if weights is not None:
+        return OnnxBackend(weights)
+
+    raise BackendNotConfigured(
+        "No predictor backend configured. Set XENOSITE_BACKEND to an http(s) "
+        "xenosite-api URL, XENOSITE_MODELS_WEIGHTS to an ONNX directory, put "
+        "*.onnx files under ./weights/onnx/v0, or set XENOSITE_ONNX_URL "
+        "(weights download on first predict())."
+    )
+
+
+def resolve_for_model(
+    spec: Spec,
+    *,
+    backend: Optional[str | PredictBackend] = None,
+    backends: Optional[Mapping[Spec, str | PredictBackend]] = None,
+    env: Optional[Mapping[str, str]] = None,
+    cwd: Optional[Path] = None,
+) -> PredictBackend:
+    """Per-model override, then the call-level backend, then env picker."""
+    if backends and spec in backends:
+        return resolve_backend(backends[spec], env=env, cwd=cwd)
+    return resolve_backend(backend, env=env, cwd=cwd)

xenosite/predict/backends/adapters.py

@@ -0,0 +1,180 @@
+"""Shared adapters: backend-native scores → user-API :class:`Molecule` results.
+
+Ported in spirit from ``xenosite-api`` ``v0/adapters.py``. One helper set, not
+per-model copies. Numeric scores must match; 0-based RDKit indices; adapter
+quirks (UGT replacing ``results``, quinone ``{}`` → 0.0) are not replayed
+except insofar as scores stay the same.
+"""
+
+from __future__ import annotations
+
+from typing import Iterable, Optional, Sequence
+
+import numpy as np
+
+from ..types import (
+    AtomBondResult,
+    AtomResult,
+    BondResult,
+    Metabolite,
+    MolAtomPairResult,
+    MolAtomResult,
+    MolBondResult,
+    Molecule,
+)
+
+
+def append_mol_bond(
+    molecule: Molecule,
+    *,
+    model: str,
+    version: str,
+    mol: float,
+    bond: Sequence[float],
+) -> None:
+    molecule.results.append(
+        MolBondResult(model=model, version=version, mol=float(mol), bond=[float(x) for x in bond])
+    )
+
+
+def append_mol_atom(
+    molecule: Molecule,
+    *,
+    model: str,
+    version: str,
+    mol: float,
+    atom: Sequence[float],
+    metabolite: Optional[list[Metabolite]] = None,
+) -> None:
+    molecule.results.append(
+        MolAtomResult(
+            model=model,
+            version=version,
+            mol=float(mol),
+            atom=[float(x) for x in atom],
+            metabolite=metabolite,
+        )
+    )
+
+
+def append_atom(
+    molecule: Molecule,
+    *,
+    model: str,
+    version: str,
+    atom: Sequence[float],
+) -> None:
+    molecule.results.append(
+        AtomResult(model=model, version=version, atom=[float(x) for x in atom])
+    )
+
+
+def append_bond(
+    molecule: Molecule,
+    *,
+    model: str,
+    version: str,
+    bond: Sequence[float],
+) -> None:
+    molecule.results.append(
+        BondResult(model=model, version=version, bond=[float(x) for x in bond])
+    )
+
+
+def append_atom_bond(
+    molecule: Molecule,
+    *,
+    model: str,
+    version: str,
+    atom: Sequence[float],
+    bond: Sequence[float],
+) -> None:
+    molecule.results.append(
+        AtomBondResult(
+            model=model,
+            version=version,
+            atom=[float(x) for x in atom],
+            bond=[float(x) for x in bond],
+        )
+    )
+
+
+def append_atom_pair(
+    molecule: Molecule,
+    *,
+    model: str,
+    version: str,
+    mol: float,
+    atom: Sequence[float],
+    pair: Sequence[float],
+    pair_idx: Sequence[tuple[int, int]],
+) -> None:
+    pp = canonicalize_pair_idx(list(pair_idx), list(pair))
+    molecule.results.append(
+        MolAtomPairResult(
+            model=model,
+            version=version,
+            mol=float(mol),
+            atom=[float(x) for x in atom],
+            pair=pp["pair"],
+            pair_idx=pp["pair_idx"],
+        )
+    )
+
+
+def canonicalize_pair_idx(
+    pair_idx: Sequence[tuple[int, int]], pair: Sequence[float]
+) -> dict:
+    ix = [(tuple(sorted(idxs)), float(x)) for idxs, x in zip(pair_idx, pair)]
+    ix.sort()
+    return {"pair_idx": [i for i, _ in ix], "pair": [x for _, x in ix]}
+
+
+def canonical_bond_site_pair(a: int, b: int) -> tuple[int, int]:
+    """Legacy ndealk/isozyme site keys use ascending atom ids (``2-1`` → ``1-2``)."""
+    if a <= b:
+        return a, b
+    return b, a
+
+
+def reorder_by_bond(
+    scores: Sequence[float],
+    current: Sequence[Iterable[int]],
+    new: Sequence[tuple[int, int]],
+    *,
+    fill: float = 0.0,
+) -> list[float]:
+    """Map bond scores onto ``molecule.bonds.idx`` (frozenset of atom ids)."""
+    lookup = {frozenset(b): i for i, b in enumerate(current)}
+    out: list[float] = []
+    for b in new:
+        key = frozenset(b)
+        out.append(float(scores[lookup[key]]) if key in lookup else fill)
+    return out
+
+
+def or_combine(values: Sequence[float]) -> float:
+    """``1 - prod(1 - p)`` aggregation used by quinone atoms and bioactivation."""
+    if not values:
+        return 0.0
+    arr = np.asarray(values, dtype=float)
+    return float(1.0 - np.prod(1.0 - arr))
+
+
+def safe_atom_index(v) -> int:
+    """OpenBabel 1-based atom id → 0-based RDKit index (``v0/adapters.py``)."""
+    try:
+        return int(v) - 1
+    except (TypeError, ValueError):
+        return v  # type: ignore[return-value]
+
+
+def legacy_atom_vector(site_map: dict, n_atoms: int, *, one_based: bool = True) -> list[float]:
+    """Per-atom scores from a legacy site map → 0-based RDKit vector.
+
+    Delegates to :func:`xenosite.predict.numbering.legacy_site_to_atom_vector`
+    (handles gapped legacy OB 2.4 keys on ``[nH]`` SMILES).
+    """
+    from ..numbering import legacy_site_to_atom_vector
+
+    return legacy_site_to_atom_vector(site_map, n_atoms)

xenosite/predict/backends/http.py

@@ -0,0 +1,60 @@
+"""HTTP backend against a deployed xenosite-api (not the legacy Flask app).
+
+``XENOSITE_BACKEND`` is the API origin (no trailing path required).
+``XENOSITE_API_KEY`` is sent as ``Authorization: Bearer …`` when set.
+
+This backend never loads ONNX. Live parity tests must pin ONNX vs the
+legacy test-API, not vs production HTTP.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+from urllib.parse import quote
+
+import httpx
+
+from ..errors import UnknownModel
+from ..types import Molecule
+
+# xenosite-api v0 routes (query: ?smiles=)
+_V0_ROUTES: dict[tuple[str, str], str] = {
+    ("epoxidation", "0"): "/v0/epoxidation",
+    ("quinone", "0"): "/v0/quinone",
+    ("ugt", "0"): "/v0/ugt",
+    ("ndealk", "0"): "/v0/ndealk",
+    ("isozyme", "0"): "/v0/isozyme",
+    ("phase1", "0"): "/v0/phase1",
+    ("bioactivation", "0"): "/v0/bioactivation",
+    ("reactivity", "0"): "/v0/reactivity",
+}
+
+
+class HttpBackend:
+    """GET ``{origin}{route}?smiles=`` and return a Molecule-shaped dict."""
+
+    name = "http"
+
+    def __init__(self, origin: str, api_key: Optional[str] = None, *, timeout: float = 60.0):
+        self.origin = origin.rstrip("/")
+        self.api_key = api_key
+        self.timeout = timeout
+
+    def _headers(self) -> dict[str, str]:
+        h = {"accept": "application/json"}
+        if self.api_key:
+            h["authorization"] = f"Bearer {self.api_key}"
+        return h
+
+    def available_models(self) -> list[tuple[str, str]]:
+        return list(_V0_ROUTES)
+
+    def predict_native(self, smiles: str, model: str, version: str) -> Any:
+        route = _V0_ROUTES.get((model, version))
+        if route is None:
+            raise UnknownModel(f"HTTP backend has no route for {model!r} {version!r}")
+        url = f"{self.origin}{route}"
+        with httpx.Client(timeout=self.timeout, headers=self._headers()) as client:
+            r = client.get(url, params={"smiles": smiles})
+            r.raise_for_status()
+            return r.json()