| |
| """Load one released endpoint and average its five seed models. |
| |
| The checkpoint directory is self-contained. Fingerprints are computed here. |
| PaDEL 3D columns, the GIN embedding, and the MiniMol embedding are inputs: |
| their encoder weights are not part of this release. |
| |
| from predict import EndpointModel |
| model = EndpointModel.from_pretrained(".", "pgp_broccatelli") |
| model.predict(["CCO", "c1ccccc1"]) |
| """ |
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
|
|
|
|
| def _read_json(path: Path) -> dict: |
| return json.loads(path.read_text()) |
|
|
|
|
| def invert_y(raw: np.ndarray, state: dict) -> np.ndarray: |
| y = np.asarray(raw, dtype=float).reshape(-1) * float(state["scale"]) + float(state["mean"]) |
| if state["log"]: |
| y = np.power(10.0, y) - 1.0 |
| return y + float(state["offset"]) |
|
|
|
|
| def _mlp(x: np.ndarray, head: dict) -> np.ndarray: |
| """One hidden ReLU layer, matching the exported scikit-learn head.""" |
| hidden = np.maximum(0.0, x @ head["coef_0"] + head["intercept_0"]) |
| out = hidden @ head["coef_1"] + head["intercept_1"] |
| out = np.asarray(out, dtype=float) |
| if out.ndim == 1 or out.shape[-1] == 1: |
| return out.reshape(-1) |
| if head["kind"] != "classifier": |
| raise RuntimeError("regressor head has more than one output") |
| shifted = out - out.max(axis=1, keepdims=True) |
| exp = np.exp(shifted) |
| return (exp / exp.sum(axis=1, keepdims=True))[:, 1] |
|
|
|
|
| def head_predict(minimol: np.ndarray, head: dict) -> np.ndarray: |
| pred = np.full(len(minimol), np.nan, dtype=float) |
| ok = np.isfinite(minimol).all(axis=1) |
| if not ok.any(): |
| return pred |
| scaled = (minimol[ok] - head["x_mean"]) / head["x_scale"] |
| raw = _mlp(scaled, head) |
| if head["kind"] == "classifier": |
| if str(head["activation"]) == "logistic" and head["coef_1"].shape[-1] == 1: |
| raw = 1.0 / (1.0 + np.exp(-raw)) |
| pred[ok] = raw |
| return pred |
| pred[ok] = invert_y(raw, head) |
| return pred |
|
|
|
|
| def mix(catboost: np.ndarray, head: np.ndarray, weight: float) -> np.ndarray: |
| blended = (1.0 - weight) * catboost + weight * head |
| missing = ~np.isfinite(head) |
| blended[missing] = catboost[missing] |
| return blended |
|
|
|
|
| def fingerprints(smiles: list[str], descriptor_names: list[str]) -> np.ndarray: |
| """Morgan count r=2 1024, Avalon count 1024, ErG 315, then the named RDKit descriptors.""" |
| from rdkit import Chem, RDLogger |
| from rdkit.Avalon.pyAvalonTools import GetAvalonCountFP |
| from rdkit.Chem import DataStructs, rdReducedGraphs |
| from rdkit.Chem.rdMolDescriptors import GetHashedMorganFingerprint |
| from rdkit.ML.Descriptors.MoleculeDescriptors import MolecularDescriptorCalculator |
|
|
| RDLogger.DisableLog("rdApp.*") |
| calculator = MolecularDescriptorCalculator(descriptor_names) |
| rows = [] |
| for one in smiles: |
| mol = Chem.MolFromSmiles(one) if isinstance(one, str) else None |
| if mol is None: |
| rows.append(None) |
| continue |
| morgan = np.zeros((0,), dtype=np.int32) |
| DataStructs.ConvertToNumpyArray( |
| GetHashedMorganFingerprint(mol, nBits=1024, radius=2), morgan |
| ) |
| avalon = np.zeros((0,), dtype=np.int32) |
| DataStructs.ConvertToNumpyArray(GetAvalonCountFP(mol, nBits=1024), avalon) |
| erg = np.asarray(rdReducedGraphs.GetErGFingerprint(mol), dtype=float) |
| desc = np.asarray(calculator.CalcDescriptors(mol), dtype=float) |
| rows.append(np.concatenate([morgan, avalon, erg, desc]).astype(float)) |
| width = next(len(row) for row in rows if row is not None) |
| out = np.full((len(smiles), width), np.nan, dtype=float) |
| for i, row in enumerate(rows): |
| if row is not None: |
| out[i] = row |
| return out |
|
|
|
|
| def _as_2d(block, width: int, name: str, n_rows: int) -> np.ndarray: |
| array = np.asarray(block, dtype=float) |
| if array.ndim != 2 or array.shape != (n_rows, width): |
| raise ValueError(f"{name} must have shape ({n_rows}, {width}), got {array.shape}") |
| return array |
|
|
|
|
| class EndpointModel: |
| def __init__(self, root: Path, endpoint: str): |
| self.root = Path(root) |
| self.endpoint = endpoint |
| self.config = _read_json(self.root / endpoint / "config.json") |
| self.spec = _read_json(self.root / "feature_spec.json") |
| self.endpoint_spec = self.spec["endpoints"][endpoint] |
| if self.config["endpoint"] != endpoint: |
| raise RuntimeError(f"{endpoint} config says {self.config['endpoint']}") |
|
|
| @classmethod |
| def from_pretrained(cls, root: str | Path, endpoint: str) -> "EndpointModel": |
| return cls(Path(root), endpoint) |
|
|
| def _padel_block(self, padel, n_rows: int) -> np.ndarray: |
| n_keep = int(self.config["n_3d"]) |
| if n_keep == 0: |
| if padel is not None: |
| raise ValueError(f"{self.endpoint} does not use PaDEL columns") |
| return np.zeros((n_rows, 0), dtype=float) |
| if padel is None: |
| raise ValueError( |
| f"{self.endpoint} expects PaDEL 3D columns. " |
| "Pass an array of shape (n, 431) in feature_spec padel_3d column order, " |
| "or an array already reduced to this endpoint's columns." |
| ) |
| if hasattr(padel, "columns"): |
| names = list(self.endpoint_spec["padel_columns"]) |
| missing = [name for name in names if name not in padel.columns] |
| if missing: |
| raise ValueError(f"PaDEL frame is missing {missing[:5]}") |
| return np.asarray(padel[names], dtype=float) |
| array = np.asarray(padel, dtype=float) |
| if array.ndim != 2 or array.shape[0] != n_rows: |
| raise ValueError(f"padel must have {n_rows} rows, got {getattr(array, 'shape', None)}") |
| canonical = len(self.spec["padel_3d"]["columns"]) |
| if array.shape[1] == canonical: |
| index = self.endpoint_spec["padel_index"] |
| return np.ascontiguousarray(array[:, index]) |
| if array.shape[1] == n_keep: |
| return array |
| raise ValueError( |
| f"padel second dimension must be {canonical} or {n_keep}, got {array.shape[1]}" |
| ) |
|
|
| def features(self, smiles: list[str], padel=None, gin=None, minimol=None) -> np.ndarray: |
| fp = fingerprints(smiles, self.spec["rdkit_descriptors"]) |
| if fp.shape[1] != int(self.config["n_fingerprint"]): |
| raise RuntimeError(f"fingerprint width {fp.shape[1]} != {self.config['n_fingerprint']}") |
| blocks = [fp, self._padel_block(padel, len(smiles))] |
| if self.config["gin"]: |
| if gin is None: |
| raise ValueError(f"{self.endpoint} expects a GIN embedding of shape (n, 300)") |
| blocks.append(_as_2d(gin, 300, "gin", len(smiles))) |
| if self.config["minimol_concat"]: |
| if minimol is None: |
| raise ValueError(f"{self.endpoint} expects a MiniMol embedding of shape (n, 512)") |
| blocks.append(_as_2d(minimol, 512, "minimol", len(smiles))) |
| x = np.ascontiguousarray(np.hstack(blocks)) |
| if x.shape[1] != int(self.config["n_features"]): |
| raise RuntimeError(f"feature width {x.shape[1]} != {self.config['n_features']}") |
| return x |
|
|
| def predict_features(self, x: np.ndarray, minimol=None) -> np.ndarray: |
| from catboost import CatBoostClassifier, CatBoostRegressor |
|
|
| x = np.asarray(x, dtype=float) |
| if x.ndim != 2 or x.shape[1] != int(self.config["n_features"]): |
| raise ValueError(f"x must have shape (n, {self.config['n_features']})") |
| weight = float(self.config["blend_weight"]) |
| if weight > 0.0: |
| minimol = _as_2d(minimol, 512, "minimol", len(x)) |
| task = self.config["task"] |
| klass = CatBoostRegressor if task == "regression" else CatBoostClassifier |
| seeds = [] |
| for seed in self.config["seeds"]: |
| folder = self.root / self.endpoint / f"seed{seed}" |
| model = klass() |
| model.load_model(str(folder / "catboost.cbm")) |
| if task == "regression": |
| pred = invert_y(model.predict(x), _read_json(folder / "y_scaler.json")) |
| else: |
| pred = np.asarray(model.predict_proba(x)[:, 1], dtype=float) |
| if weight > 0.0: |
| head = dict(np.load(folder / "minimol_head.npz")) |
| pred = mix(pred, head_predict(minimol, head), weight) |
| seeds.append(pred) |
| return np.mean(seeds, axis=0) |
|
|
| def predict(self, smiles, padel=None, gin=None, minimol=None) -> np.ndarray: |
| if isinstance(smiles, str): |
| smiles = [smiles] |
| smiles = list(smiles) |
| x = self.features(smiles, padel=padel, gin=gin, minimol=minimol) |
| return self.predict_features(x, minimol=minimol) |
|
|
