pocketnet 0.1.0__py3-none-any.whl

pocketnet/__init__.py

@@ -0,0 +1,8 @@
+"""PocketNet: ligand-binding pocket prediction from protein structures.
+
+>>> from pocketnet import PocketPredictor
+>>> pockets = PocketPredictor().predict("examples/3POZ.pdb")
+>>> pockets[0]["center"], pockets[0]["score"]
+"""
+__version__ = "0.1.0"
+from .api import PocketPredictor, predict  # noqa: E402,F401

pocketnet/api.py

@@ -0,0 +1,92 @@
+"""High-level interface: a five-model PocketNet ensemble that turns a PDB file into ranked pockets."""
+from __future__ import annotations
+import json, os, tempfile
+from pathlib import Path
+import numpy as np, torch
+from . import feat, pockets
+from .model import PocketNet
+from .pdbio import read_pdb
+
+_HERE = Path(__file__).resolve().parent
+WEIGHTS = sorted((_HERE / "weights").glob("v3_fold*.pt"))
+MAX_ATOMS = 14000   # structures above this size were not used in training; larger inputs are accepted but untested
+
+
+def _load(path, device):
+    sd = torch.load(path, map_location=device, weights_only=True)
+    hidden = sd["emb.weight"].shape[1]
+    layers = len([k for k in sd if k.startswith("layers.") and k.endswith("msg.0.weight")])
+    m = PocketNet(hidden, layers).to(device)
+    m.load_state_dict(sd)
+    return m.eval()
+
+
+class PocketPredictor:
+    """Averages the sigmoid scores of the five fold models and extracts at most ten ranked pockets.
+
+    Parameters
+    ----------
+    device : "cuda", "cpu" or None (cuda if available).
+    weights : optional list of checkpoint paths (default: the five bundled fold models).
+    """
+
+    def __init__(self, device: str | None = None, weights=None):
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.models = [_load(w, self.device) for w in (weights or WEIGHTS)]
+        cal = json.loads((_HERE / "calibration.json").read_text())
+        self._cal_x, self._cal_y = np.asarray(cal["x"]), np.asarray(cal["y"])
+
+    @torch.no_grad()
+    def score_points(self, structure):
+        """Return (surface_points [P,3], ensemble_score [P]) for a PDB file path or a PDB-format string."""
+        P = self._read(structure)
+        dev = self.device
+        ty, el = feat.atom_types(P["res"], P["name"])
+        x = torch.tensor(P["xyz"], device=dev); e = torch.tensor(el, device=dev); t = torch.tensor(ty, device=dev)
+        pts, nrm = feat.sas_points(x, e)
+        pf = feat.point_feats(pts, nrm, x, e)
+        s = 0
+        for m in self.models:
+            if dev == "cuda":
+                with torch.autocast("cuda", dtype=torch.bfloat16):
+                    lo, _ = m(t, x, pts, pf)
+            else:
+                lo, _ = m(t, x, pts, pf)
+            s = s + torch.sigmoid(lo.float())
+        return pts.cpu().numpy(), (s / len(self.models)).cpu().numpy(), P
+
+    def predict(self, structure, max_pockets: int = 10, residue_radius: float = 8.0):
+        """Ranked pockets for a PDB file (path or text).
+
+        Each pocket is a dict with ``rank``, ``score`` (sum of its 24 highest point scores), ``probability``
+        (isotonic calibration of the score on held-out proteins; indicative only), ``center`` (x, y, z in angstrom),
+        ``n_points`` and ``residues`` (chain:number of residues with an atom within ``residue_radius`` of the centre).
+        """
+        pts, s, P = self.score_points(structure)
+        found = pockets.extract_with_members(pts, s, topn=max_pockets, **pockets.TUNED)
+        out = []
+        for rank, c in enumerate(found, 1):
+            ctr = np.asarray(c["center"], float)
+            near = np.linalg.norm(P["xyz"] - ctr[None], axis=1) <= residue_radius
+            res = sorted({f"{P['chain'][j]}:{P['rid'][j]}" for j in np.where(near)[0]})
+            out.append(dict(rank=rank, score=float(c["score"]), probability=float(np.interp(c["score"], self._cal_x, self._cal_y)),
+                            center=[round(float(v), 3) for v in ctr], n_points=int(len(c["idx"])), residues=res))
+        return out
+
+    def _read(self, structure):
+        text = str(structure)
+        if "\n" in text or text.lstrip().startswith(("ATOM", "HEADER", "REMARK", "CRYST1", "HETATM")):
+            with tempfile.NamedTemporaryFile("w", suffix=".pdb", delete=False) as f:
+                f.write(text); name = f.name
+            try: P, _ = read_pdb(name)
+            finally: os.unlink(name)
+        else:
+            P, _ = read_pdb(text)
+        if len(P["xyz"]) < 30:
+            raise ValueError("fewer than 30 protein atoms were read (only ATOM records of standard residues are used; mmCIF is not supported)")
+        return P
+
+
+def predict(structure, device=None, **kw):
+    """One-shot convenience wrapper around :class:`PocketPredictor`."""
+    return PocketPredictor(device).predict(structure, **kw)

pocketnet/calibration.json

@@ -0,0 +1 @@
+{"x": [0.106, 0.699, 1.292, 1.885, 2.478, 3.071, 3.664, 4.257, 4.851, 5.444, 6.037, 6.63, 7.223, 7.816, 8.409, 9.002, 9.595, 10.189, 10.782, 11.375, 11.968, 12.561, 13.154, 13.747, 14.34, 14.934, 15.527, 16.12, 16.713, 17.306, 17.899, 18.492, 19.085, 19.679, 20.272, 20.865, 21.458, 22.051, 22.644, 23.237], "y": [0.02, 0.032, 0.032, 0.032, 0.045, 0.045, 0.046, 0.046, 0.054, 0.057, 0.061, 0.12, 0.12, 0.136, 0.144, 0.144, 0.169, 0.169, 0.216, 0.216, 0.216, 0.3, 0.329, 0.348, 0.348, 0.349, 0.503, 0.503, 0.55, 0.551, 0.551, 0.551, 0.636, 0.68, 0.734, 0.768, 0.819, 0.912, 0.912, 0.98], "n_pockets": 5324, "n_proteins": 1000, "rank_hit_rate": {"1": 0.8316532258064516, "2": 0.6538849646821393, "3": 0.4606986899563319, "4": 0.23684210526315788, "5": 0.1338709677419355, "6": 0.08018867924528301}}

pocketnet/cli.py

@@ -0,0 +1,31 @@
+"""Command line interface: ``pocketnet predict protein.pdb -o results/``."""
+import argparse, csv, json, os, sys
+from . import __version__
+from .api import PocketPredictor
+
+
+def main(argv=None):
+    ap = argparse.ArgumentParser(prog="pocketnet", description="Predict ligand-binding pockets with PocketNet.")
+    ap.add_argument("--version", action="version", version=f"pocketnet {__version__}")
+    sub = ap.add_subparsers(dest="cmd", required=True)
+    p = sub.add_parser("predict", help="rank pockets of one or more PDB files")
+    p.add_argument("pdb", nargs="+", help="input PDB file(s) (ATOM records; mmCIF is not supported)")
+    p.add_argument("-o", "--out", default="pocketnet_out", help="output directory (default: pocketnet_out)")
+    p.add_argument("--device", choices=["cpu", "cuda"], default=None, help="default: cuda if available, else cpu")
+    p.add_argument("--max-pockets", type=int, default=10)
+    a = ap.parse_args(argv)
+    os.makedirs(a.out, exist_ok=True)
+    pred = PocketPredictor(a.device)
+    for path in a.pdb:
+        res = pred.predict(path, max_pockets=a.max_pockets); stem = os.path.basename(path)
+        with open(os.path.join(a.out, stem + "_pockets.csv"), "w", newline="") as f:
+            w = csv.writer(f); w.writerow(["rank", "score", "probability", "center_x", "center_y", "center_z", "n_points", "residues"])
+            for r in res: w.writerow([r["rank"], f"{r['score']:.3f}", f"{r['probability']:.3f}", *r["center"], r["n_points"], " ".join(r["residues"])])
+        with open(os.path.join(a.out, stem + "_pockets.json"), "w") as f: json.dump(res, f, indent=1)
+        print(f"{stem}: {len(res)} pockets")
+        for r in res[:5]: print(f"  {r['rank']:>2}  score {r['score']:6.2f}  p~{r['probability']:.2f}  centre {r['center']}  ({len(r['residues'])} residues)")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

pocketnet/feat.py

@@ -0,0 +1,106 @@
+"""Protein featurisation for pocket detection. Everything heavy runs on GPU (torch)."""
+import math, numpy as np, torch
+
+AA = ['ALA','ARG','ASN','ASP','CYS','GLN','GLU','GLY','HIS','ILE','LEU','LYS','MET','PHE','PRO','SER','THR','TRP','TYR','VAL']
+RMAP = {'HIE':'HIS','HID':'HIS','HIP':'HIS','HSD':'HIS','HSE':'HIS','HSP':'HIS','CYX':'CYS','CYM':'CYS','ASH':'ASP','GLH':'GLU','LYN':'LYS','MSE':'MET'}
+ATOMS = {  # standard heavy-atom names per residue (PDB naming)
+ 'ALA':'N CA C O CB','ARG':'N CA C O CB CG CD NE CZ NH1 NH2','ASN':'N CA C O CB CG OD1 ND2','ASP':'N CA C O CB CG OD1 OD2',
+ 'CYS':'N CA C O CB SG','GLN':'N CA C O CB CG CD OE1 NE2','GLU':'N CA C O CB CG CD OE1 OE2','GLY':'N CA C O',
+ 'HIS':'N CA C O CB CG ND1 CD2 CE1 NE2','ILE':'N CA C O CB CG1 CG2 CD1','LEU':'N CA C O CB CG CD1 CD2',
+ 'LYS':'N CA C O CB CG CD CE NZ','MET':'N CA C O CB CG SD CE','PHE':'N CA C O CB CG CD1 CD2 CE1 CE2 CZ',
+ 'PRO':'N CA C O CB CG CD','SER':'N CA C O CB OG','THR':'N CA C O CB OG1 CG2','TRP':'N CA C O CB CG CD1 CD2 NE1 CE2 CE3 CZ2 CZ3 CH2',
+ 'TYR':'N CA C O CB CG CD1 CD2 CE1 CE2 CZ OH','VAL':'N CA C O CB CG1 CG2'}
+VOCAB = {}
+for r in AA:
+    for a in ATOMS[r].split(): VOCAB[(r, a)] = len(VOCAB) + 1       # 0 = unknown
+for r in AA: VOCAB[(r, 'OXT')] = len(VOCAB) + 1
+N_TYPES = len(VOCAB) + 2                                           # +unknown(0) +other-residue(last)
+OTHER = N_TYPES - 1
+ALIAS = {'OC1':'O','OC2':'OXT','O1':'O','O2':'OXT'}
+EL = {'C':0,'N':1,'O':2,'S':3,'P':4}
+VDW = torch.tensor([1.70, 1.55, 1.52, 1.80, 1.80, 1.80])          # C N O S P other
+
+def atom_types(res, name):
+    out = np.zeros(len(res), np.int64); el = np.zeros(len(res), np.int64)
+    for i, (r, a) in enumerate(zip(res, name)):
+        r = RMAP.get(r, r); a = ALIAS.get(a, a)
+        out[i] = VOCAB.get((r, a), 0) if r in ATOMS else OTHER
+        el[i] = EL.get(a[0], 5)
+    return out, el
+
+def fib_sphere(n, device):
+    i = torch.arange(n, device=device, dtype=torch.float32) + 0.5
+    phi = torch.acos(1 - 2 * i / n); th = math.pi * (1 + 5 ** 0.5) * i
+    return torch.stack([torch.cos(th) * torch.sin(phi), torch.sin(th) * torch.sin(phi), torch.cos(phi)], 1)
+
+@torch.no_grad()
+def sas_points(xyz, el, probe=1.4, n_sphere=48, chunk=4096):
+    """Shrake-Rupley points on the solvent-accessible surface. Returns pts [P,3], normals [P,3] (pointing out)."""
+    dev = xyz.device; r = VDW.to(dev)[el] + probe
+    sph = fib_sphere(n_sphere, dev)
+    pts = (xyz[:, None, :] + r[:, None, None] * sph[None]).reshape(-1, 3)
+    own = torch.arange(len(xyz), device=dev).repeat_interleave(n_sphere)
+    keep = torch.ones(len(pts), dtype=torch.bool, device=dev)
+    for s in range(0, len(pts), chunk):
+        d = torch.cdist(pts[s:s+chunk], xyz)                                # [c,N]
+        inside = d < (r[None] - 1e-3)
+        inside[torch.arange(inside.shape[0], device=dev), own[s:s+chunk]] = False
+        keep[s:s+chunk] = ~inside.any(1)
+    pts = pts[keep]; nrm = (pts - xyz[own[keep]]); nrm = nrm / nrm.norm(dim=1, keepdim=True)
+    return pts, nrm
+
+@torch.no_grad()
+def occupancy_grid(xyz, el, cell=0.8, pad=18.0):
+    dev = xyz.device; lo = xyz.min(0).values - pad; hi = xyz.max(0).values + pad
+    dims = ((hi - lo) / cell).ceil().long() + 1
+    grid = torch.zeros(tuple(dims.tolist()), dtype=torch.bool, device=dev)
+    rr = (VDW.to(dev)[el] + 0.2)                                          # slightly inflated
+    m = int(math.ceil(2.1 / cell)); rng = torch.arange(-m, m + 1, device=dev)
+    off = torch.stack(torch.meshgrid(rng, rng, rng, indexing='ij'), -1).reshape(-1, 3)   # [K,3]
+    for s in range(0, len(xyz), 2048):
+        a = xyz[s:s+2048]; c = ((a - lo) / cell).round().long()            # [n,3]
+        idx = c[:, None, :] + off[None]; pos = lo + idx.float() * cell
+        ok = ((pos - a[:, None, :]).norm(dim=-1) <= rr[s:s+2048, None])
+        idx = idx[ok]; grid[idx[:, 0], idx[:, 1], idx[:, 2]] = True
+    return grid, lo, cell
+
+@torch.no_grad()
+def point_feats(pts, nrm, xyz, el, aa_sel=None, n_rays=64, step=0.8, rmax=16.0):
+    """Rotation-invariant per-point descriptors: ray-cast burial (all dirs + cones about the normal), shell counts."""
+    dev = pts.device
+    grid, lo, cell = occupancy_grid(xyz, el)
+    rays = fib_sphere(n_rays, dev)                                         # [R,3]
+    ns = int(rmax / step); t = (torch.arange(1, ns + 1, device=dev, dtype=torch.float32) * step)
+    P = len(pts); fh = torch.empty(P, n_rays, device=dev)                  # first-hit distance
+    for s in range(0, P, 2048):
+        p = pts[s:s+2048]                                                  # [c,3]
+        q = p[:, None, None, :] + rays[None, :, None, :] * t[None, None, :, None]    # [c,R,S,3]
+        idx = ((q - lo) / cell).round().long()
+        idx = torch.minimum(torch.maximum(idx, torch.zeros_like(idx)), torch.tensor(grid.shape, device=dev) - 1)
+        occ = grid[idx[..., 0], idx[..., 1], idx[..., 2]]                  # [c,R,S]
+        first = torch.where(occ.any(-1), occ.float().argmax(-1), torch.full_like(occ[..., 0], ns, dtype=torch.long))
+        fh[s:s+2048] = torch.where(first >= ns, torch.full_like(first, rmax * 2, dtype=torch.float32), t[first.clamp(max=ns - 1)])
+    cosang = nrm @ rays.T                                                  # [P,R]
+    f = []
+    for rng in (4.0, 8.0, 12.0, 16.0):
+        b = (fh <= rng).float(); f.append(b.mean(1))
+        for cth in (0.866, 0.5, 0.0):                                      # cones 30/60/90 deg about normal
+            m = (cosang >= cth).float(); f.append((b * m).sum(1) / m.sum(1).clamp(min=1))
+    f.append(fh.clamp(max=rmax).mean(1) / rmax)
+    f.append(fh.clamp(max=rmax).min(1).values / rmax)
+    # opposite-hemisphere enclosure (pocket depth): distance to first hit along -normal cone
+    m = (cosang <= -0.5).float(); f.append(((fh.clamp(max=rmax) * m).sum(1) / m.sum(1).clamp(min=1)) / rmax)
+    # shell counts by element class
+    d = torch.cdist(pts, xyz); eo = torch.nn.functional.one_hot(el, 6).float()    # [N,6]
+    for lo_, hi_ in ((0, 4.0), (4.0, 6.0), (6.0, 8.0), (8.0, 12.0)):
+        w = ((d >= lo_) & (d < hi_)).float(); f.append((w.sum(1, keepdim=True) / 30.0)[:, 0])
+        c = w @ eo; f.extend([c[:, k] / 12.0 for k in range(4)])           # C N O S
+    return torch.stack(f, 1)                                               # [P, ~32]
+N_PF = 4 * 4 + 3 + 4 * 5
+
+def knn(x, y, k, chunk=2048):
+    """for each row of x the k nearest rows of y -> (dist [n,k], idx [n,k])"""
+    ds, ix = [], []
+    for s in range(0, len(x), chunk):
+        d = torch.cdist(x[s:s+chunk], y); v, i = d.topk(min(k, y.shape[0]), largest=False); ds.append(v); ix.append(i)
+    return torch.cat(ds), torch.cat(ix)

pocketnet/model.py

@@ -0,0 +1,71 @@
+import math, torch, torch.nn as nn, torch.nn.functional as F
+from . import feat
+
+class RBF(nn.Module):
+    def __init__(self, n, lo, hi):
+        super().__init__(); self.register_buffer('mu', torch.linspace(lo, hi, n)); self.g = (hi - lo) / n
+    def forward(self, d): return torch.exp(-((d[..., None] - self.mu) / self.g) ** 2)
+
+class Layer(nn.Module):
+    def __init__(self, h, ne):
+        super().__init__()
+        self.msg = nn.Sequential(nn.Linear(2 * h + ne, h), nn.SiLU(), nn.Linear(h, h))
+        self.gate = nn.Linear(2 * h + ne, h)
+        self.upd = nn.Sequential(nn.Linear(2 * h, h), nn.SiLU(), nn.Linear(h, h))
+        self.n1 = nn.LayerNorm(h)
+    def forward(self, h, ei, ea):
+        i, j = ei; z = torch.cat([h[i], h[j], ea], 1)
+        m = self.msg(z) * torch.sigmoid(self.gate(z))
+        agg = torch.zeros_like(h, dtype=m.dtype).index_add_(0, i, m) / 12.0
+        return self.n1(h + self.upd(torch.cat([h, agg], 1)))
+
+class PocketNet(nn.Module):
+    def __init__(self, h=128, layers=6, k_atom=24, k_pt=32, r_atom=6.0, r_pt=10.0, drop=0.1):
+        super().__init__()
+        self.k_atom, self.k_pt, self.r_atom, self.r_pt = k_atom, k_pt, r_atom, r_pt
+        self.emb = nn.Embedding(feat.N_TYPES, h); self.bur = nn.Linear(4, h)
+        self.rbf = RBF(16, 0, r_atom); self.layers = nn.ModuleList([Layer(h, 16) for _ in range(layers)])
+        self.atom_head = nn.Linear(h, 1)
+        hp = 64
+        self.proj = nn.Linear(h, hp)
+        self.rax = RBF(12, -r_pt, r_pt); self.rrad = RBF(10, 0, r_pt)
+        self.att = nn.Sequential(nn.Linear(22 + hp, 128), nn.SiLU(), nn.Linear(128, 4))     # 4 attention heads
+        self.val = nn.Linear(hp + 22, hp)
+        self.pfn = nn.Sequential(nn.Linear(feat.N_PF, 128), nn.SiLU(), nn.Linear(128, 128))
+        self.out = nn.Sequential(nn.Linear(4 * hp + 128 + hp, 256), nn.SiLU(), nn.Dropout(drop), nn.Linear(256, 128), nn.SiLU(), nn.Linear(128, 1))
+
+    @torch.no_grad()
+    def graph(self, ty, xyz, pts, pf):
+        """Per-protein geometry: everything that does not need gradients."""
+        d, idx = feat.knn(xyz, xyz, self.k_atom + 1)
+        d, idx = d[:, 1:], idx[:, 1:]; m = d < self.r_atom
+        i = torch.arange(len(xyz), device=xyz.device)[:, None].expand_as(idx)
+        cnt = torch.stack([(d < 4).sum(1), (d < 6).sum(1), (d < 3).sum(1), (d < 5).sum(1)], 1).float() / 24.0
+        d1, i1 = feat.knn(pts, xyz, 1); nrm = pts - xyz[i1[:, 0]]; nrm = nrm / nrm.norm(dim=1, keepdim=True).clamp(min=1e-6)
+        dp, ip = feat.knn(pts, xyz, self.k_pt)
+        rel = xyz[ip] - pts[:, None, :]
+        ax = (rel * nrm[:, None, :]).sum(-1); rad = (dp ** 2 - ax ** 2).clamp(min=0).sqrt()
+        return dict(ty=ty, cnt=cnt, ei=torch.stack([i[m], idx[m]]), d=d[m], i1=i1[:, 0], ip=ip, ax=ax, rad=rad, dp=dp, pf=pf, na=len(xyz), npt=len(pts))
+
+    def collate(self, gs):
+        off = 0; ei = []; i1 = []; ip = []; o = 0
+        for g in gs:
+            ei.append(g['ei'] + off); i1.append(g['i1'] + off); ip.append(g['ip'] + off); off += g['na']
+        cat = lambda k: torch.cat([g[k] for g in gs])
+        return dict(ty=cat('ty'), cnt=cat('cnt'), ei=torch.cat(ei, 1), d=cat('d'), i1=torch.cat(i1), ip=torch.cat(ip), ax=cat('ax'), rad=cat('rad'),
+                    dp=cat('dp'), pf=cat('pf'), sizes=[g['npt'] for g in gs])
+
+    def forward_graph(self, g):
+        h = self.emb(g['ty']) + self.bur(g['cnt']); ea = self.rbf(g['d'])
+        for l in self.layers: h = l(h, g['ei'], ea)
+        a_logit = self.atom_head(h)[:, 0]; hp = self.proj(h)
+        ge = torch.cat([self.rax(g['ax']), self.rrad(g['rad'])], -1)
+        m = g['dp'] < self.r_pt; hj = hp[g['ip']]
+        a = self.att(torch.cat([ge, hj], -1)).masked_fill(~m[..., None], -1e4).softmax(1)
+        v = self.val(torch.cat([hj, ge], -1))
+        pooled = torch.einsum('pkh,pkv->phv', a, v).reshape(len(g['pf']), -1)
+        z = torch.cat([pooled, self.pfn(g['pf']), hp[g['i1']]], 1)
+        return self.out(z)[:, 0], a_logit
+
+    def forward(self, ty, xyz, pts, pf):
+        return self.forward_graph(self.collate([self.graph(ty, xyz, pts, pf)]))

pocketnet/pdbio.py

@@ -0,0 +1,21 @@
+import numpy as np
+from . import feat
+def read_pdb(path, lig_codes=None):
+    """-> dict(res,name,rid,chain,xyz) of protein heavy atoms (first model, altloc A/blank) and list of ligand coord arrays."""
+    prot, ligs = [], {}
+    with open(path) as f:
+        for ln in f:
+            if ln.startswith('ENDMDL'): break
+            rec = ln[:6]
+            if rec not in ('ATOM  ', 'HETATM'): continue
+            if ln[16] not in ' A1': continue
+            el = (ln[76:78].strip() or ln[12:16].strip()[:1]).upper()
+            if el in ('H', 'D'): continue
+            res = ln[17:20].strip(); xyz = (float(ln[30:38]), float(ln[38:46]), float(ln[46:54]))
+            if rec == 'ATOM  ' or (res == 'MSE'):
+                prot.append((ln[12:16].strip(), res, int(ln[22:26]) if ln[22:26].strip().lstrip('-').isdigit() else 0, ln[21], el, *xyz))
+            elif lig_codes and res in lig_codes:
+                ligs.setdefault((res, ln[21], ln[22:27]), []).append(xyz)
+    P = dict(name=np.array([p[0] for p in prot]), res=np.array([p[1] for p in prot]), rid=np.array([p[2] for p in prot]),
+             chain=np.array([p[3] for p in prot]), xyz=np.array([p[5:8] for p in prot], np.float32))
+    return P, [np.array(v, np.float32) for v in ligs.values()]

pocketnet/pockets.py

@@ -0,0 +1,45 @@
+"""Turn per-point scores into ranked pockets; DCA/DCC metrics (P2Rank convention)."""
+import numpy as np
+from scipy.cluster.hierarchy import linkage, fcluster
+
+def extract(pts, s, thr_q=0.97, link=5.0, topn=10, max_pts=600, method='average', top_k=12, min_pts=3):
+    """pts [P,3] numpy, s [P] scores -> list of (score, center[3]) sorted desc."""
+    if len(pts) < 3: return []
+    k = max(min_pts, min(max_pts, int(round(len(pts) * (1 - thr_q)))))
+    sel = np.argsort(-s)[:k]; p, v = pts[sel], s[sel]
+    if len(p) < 2: return [(float(v.sum()), p[0])]
+    lab = fcluster(linkage(p, method), link, 'distance'); out = []
+    for c in np.unique(lab):
+        m = lab == c
+        if m.sum() < min_pts: continue
+        vc, pc = v[m], p[m]; o = np.argsort(-vc)[:top_k]
+        w = vc[o] - vc.min() + 1e-6 if False else np.maximum(vc[o], 1e-6)
+        out.append((float(vc[o].sum()), (pc[o] * w[:, None]).sum(0) / w.sum()))
+    out.sort(key=lambda t: -t[0]); return out[:topn]
+
+def evaluate_one(pockets, lig, n_lig=1, thr=4.0):
+    """DCA success@n, @n+2 and DCC; pockets list of (score, center); lig [A,3]."""
+    if not pockets: return dict(top1=0, top3=0, topany=0, dcc1=0, dist=np.inf, found=0)
+    c = np.stack([p[1] for p in pockets]); d = np.linalg.norm(c[:, None, :] - lig[None], axis=-1).min(1)
+    dcc = np.linalg.norm(c - lig.mean(0), axis=-1)
+    return dict(top1=float(d[:n_lig].min() <= thr), top3=float(d[:n_lig + 2].min() <= thr), topany=float(d.min() <= thr),
+                dcc1=float(dcc[:n_lig].min() <= thr), dist=float(d[0]), found=len(pockets))
+
+
+TUNED = dict(thr_q=0.97, link=6.5, max_pts=300, top_k=24)
+
+def extract_with_members(pts, s, thr_q=0.97, link=6.5, topn=10, max_pts=300, method='average', top_k=24, min_pts=3):
+    """Same clustering as `extract`, but also returns each pocket's member point indices (into pts).
+    -> list of dict(score, center[3], idx[int array], top_idx[int array])."""
+    if len(pts) < 3: return []
+    k = max(min_pts, min(max_pts, int(round(len(pts) * (1 - thr_q)))))
+    sel = np.argsort(-s)[:k]; p, v = pts[sel], s[sel]
+    if len(p) < 2: return [dict(score=float(v.sum()), center=p[0], idx=sel, top_idx=sel)]
+    lab = fcluster(linkage(p, method), link, 'distance'); out = []
+    for c in np.unique(lab):
+        m = lab == c
+        if m.sum() < min_pts: continue
+        vc, pc, ic = v[m], p[m], sel[m]; o = np.argsort(-vc)[:top_k]
+        w = np.maximum(vc[o], 1e-6)
+        out.append(dict(score=float(vc[o].sum()), center=(pc[o] * w[:, None]).sum(0) / w.sum(), idx=ic, top_idx=ic[o]))
+    out.sort(key=lambda t: -t['score']); return out[:topn]

pocketnet-0.1.0.dist-info/METADATA

@@ -0,0 +1,68 @@
+Metadata-Version: 2.4
+Name: pocketnet
+Version: 0.1.0
+Summary: Ligand-binding pocket prediction from protein structures with a message-passing / attention-pooling network (five-model ensemble, weights included).
+Author: Stalin Arulsamy, Kiruthiga Natarajan
+License: MIT
+Project-URL: Homepage, https://github.com/precisionmatics/pocketnet
+Project-URL: Issues, https://github.com/precisionmatics/pocketnet/issues
+Keywords: binding site prediction,drug discovery,protein structure,graph neural network,bioinformatics
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Classifier: Topic :: Scientific/Engineering :: Chemistry
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=1.23
+Requires-Dist: scipy>=1.9
+Requires-Dist: torch>=2.0
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Dynamic: license-file
+
+# PocketNet
+
+Ligand-binding pocket prediction from a protein structure. A gated message-passing network over the atom graph feeds an
+attention-pooling read-out over the 32 nearest atoms of each solvent-accessible surface point; surface points are scored,
+clustered and ranked into pockets. The package bundles a five-model ensemble (about 20 MB of weights); no download is needed at run time.
+
+```bash
+pip install pocketnet
+pocketnet predict protein.pdb -o results/
+```
+
+```python
+from pocketnet import PocketPredictor
+
+pred = PocketPredictor()               # uses the GPU if available, otherwise CPU
+for p in pred.predict("protein.pdb"):  # at most 10 pockets, best first
+    print(p["rank"], p["score"], p["probability"], p["center"], p["residues"][:5])
+```
+
+Each pocket has `rank`, `score` (sum of its 24 highest surface-point scores), `probability` (isotonic calibration of that score on held-out
+proteins; indicative only), `center` (x, y, z in angstrom), `n_points` and `residues` (`chain:number` within 8 A of the centre).
+The CLI writes `<input>_pockets.csv` and `.json`. A run takes about 1 s per protein on a GPU and a few seconds on a CPU.
+
+## Input
+PDB-format files with `ATOM` records (standard residues; hydrogens and alternate locations other than A are ignored, first model only).
+mmCIF is not supported. Structures above about 14,000 atoms were not seen in training.
+
+## Performance (summary)
+Evaluated with P2Rank's ligand definitions and a 4 A distance-to-any-ligand-atom criterion, all methods capped at 10 pockets, on COACH420, HOLO4K
+and Chen11 after sequence- and structure-identity filtering against the training set. PocketNet is on par with P2Rank for the first-ranked pocket,
+ahead of P2Rank at larger pocket budgets, and behind GrASP at top-n on some sets. Full tables, confidence intervals, McNemar tests and caveats are in the manuscript.
+
+Limitations: PUResNet was run through our own wrapper; VN-EGNN, EquiPocket, DeepSurf and Kalasanty were not benchmarked; the model was trained on
+PLAS-20K complexes (simulation-derived protein-ligand structures), so performance on very different targets (membrane proteins, large assemblies, cryptic sites) is untested.
+
+## Repository layout
+`pocketnet/` package and weights, `tests/`, `examples/3POZ.pdb`, `training/` (data preparation and training), `benchmarks/` (scoring on the P2Rank datasets).
+
+## Citation
+Arulsamy S, Natarajan K. PocketNet (manuscript in preparation). Please check the repository for an updated citation.
+
+## License
+MIT. See `LICENSE`.

pocketnet-0.1.0.dist-info/RECORD

@@ -0,0 +1,19 @@
+pocketnet/__init__.py,sha256=wZPZeHbjtOzI5j8HoXrKVpH5b5Md2h8kiWhAC6hu3rw,309
+pocketnet/api.py,sha256=s1W-RBPWsDuuX7XI5i5FtmROfqhyBQFOUnnfa4MUoZE,4330
+pocketnet/calibration.json,sha256=RQR-rxoujGHm-u671pF65N5y4FeF6_VokAT2S-ApNzI,796
+pocketnet/cli.py,sha256=KWXpEEcAquYVArLJDnLIV_RmMmdgg59C6hGQvWle5uY,1873
+pocketnet/feat.py,sha256=tEBT3lmap906eRkznDybC98F5XrST9z4yNLpTtVTgG4,6703
+pocketnet/model.py,sha256=XROG-ZqY4pPdkpRwjiInAhGcOcke6MLQQvYKXk09yNQ,4079
+pocketnet/pdbio.py,sha256=5NS6Af553HrK57HfnPpb4wmS2ouPAJ4yFOj22O2kDlw,1253
+pocketnet/pockets.py,sha256=hghRxeIJ-WfXx7TISMFvbqcU0RGtVQQbDLhYpLczaIc,2609
+pocketnet/weights/v3_fold0.pt,sha256=V3OoMWGlY-ZxqHYO4P9iTkwtdiLL2s-JRBjsuqq4GyY,4179301
+pocketnet/weights/v3_fold1.pt,sha256=7PrBR_Kw0aEZ8KblYL_Qw_47vdizqOiWf367G56iG-0,4179301
+pocketnet/weights/v3_fold2.pt,sha256=FrUadMugGKFuB8JT6HJT4dTER6HRJ1OLNLrwmaTvTWo,4179301
+pocketnet/weights/v3_fold3.pt,sha256=l2F_VKuvYurAXtuvQem45sz4wLn92JOppCzGocXQBjw,4179301
+pocketnet/weights/v3_fold4.pt,sha256=41RXCp2XMm7LU92nHu1Me0B0Avi7r87Gc30_wSEpb8c,4179301
+pocketnet-0.1.0.dist-info/licenses/LICENSE,sha256=gZKJ2nJt23zP1Aa5tSgJ2SJUJ-WJCZXK6uykw-oTceY,1096
+pocketnet-0.1.0.dist-info/METADATA,sha256=BDtGpeWnZdZdNiFaSSsIH5HVQxvprHjESPsGS3AjpQk,3605
+pocketnet-0.1.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+pocketnet-0.1.0.dist-info/entry_points.txt,sha256=KxHWO3hWkLewIoFz0GjDcNttjLpTMEE7EyUU7Kl6aHo,49
+pocketnet-0.1.0.dist-info/top_level.txt,sha256=NO2AV5ueTNUKCjUZXO2Y5ckyLHQGM7CYpomhxSzOQ6U,10
+pocketnet-0.1.0.dist-info/RECORD,,

pocketnet-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

pocketnet-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+pocketnet = pocketnet.cli:main

pocketnet-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,9 @@
+MIT License
+
+Copyright (c) 2026 Stalin Arulsamy and Kiruthiga Natarajan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

pocketnet-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+pocketnet