"""PDB in -> per-residue BindCraft patch-coverage targetability in the B-factor column. This is the whole pipeline in one place, and the only file `app.py` talks to. It reproduces `PeSTo/processing/apply_model_bindcraft.py` exactly -- same preprocessing, same element-only feature encoding, same 64-nearest-neighbour graph, same sigmoid -- with one difference in the output: scores are written back onto the *input* file's chain ids and residue numbers instead of PeSTo's internal renumbering. Run it as a script for a batch job: python inference.py 1M47.pdb -o 1M47_targetability.pdb python inference.py *.pdb --outdir scored/ --chains A """ import argparse import os import sys from dataclasses import dataclass, field import numpy as np import torch as pt from model import Model, config_model from preprocessing import ( collate_batch_features, concatenate_chains, encode_bfactor, encode_features, encode_structure, extract_topology_knn, map_scores_to_original, preprocess_structure, read_pdb, ) DEFAULT_WEIGHTS = os.path.join(os.path.dirname(os.path.abspath(__file__)), "weights.pt") # graph connectivity the checkpoint was trained with -- not a tunable NUM_NN = 64 # Guard rail. The 32 message-passing layers dominate, and they are linear in atoms, # so on a 2-vCPU Space this is tens of minutes at the cap -- past any patience, and # past the request timeout. Raise it locally, or on GPU hardware, via the env var. MAX_ATOMS = int(os.environ.get("BINDSCOUT_MAX_ATOMS", 20000)) # torch opens one thread per core by default, and for a graph this small that is # actively harmful: on a 210-residue target, 2 threads cost 29 s of CPU where 8 cost # 92 s and the 64-thread default thrashed for half an hour. Parallel efficiency here # is negative past a couple of threads, so cap it. A 2-vCPU Space is already below # the cap and is left alone; the environment overrides either way. _threads = os.environ.get("BINDSCOUT_THREADS") if _threads: pt.set_num_threads(int(_threads)) elif pt.get_num_threads() > 4: pt.set_num_threads(4) @dataclass class Prediction: """Everything one scored structure yields.""" pdb_text: str # input structure, B-factors replaced by the score records: list = field(repr=False) # per-residue rows in the input's own numbering num_residues_scored: int = 0 num_atoms: int = 0 chains: tuple = () def protein_records(self): return [r for r in self.records if r['is_protein']] def scores(self): return np.array([r['score'] for r in self.protein_records()], dtype=np.float32) def to_csv(self, filepath): with open(filepath, 'w') as fs: fs.write("chain,resid,icode,resname,aa,targetability\n") for r in self.protein_records(): fs.write("{chain},{resid},{icode},{resname},{aa},{score:.4f}\n".format(**r)) return filepath def top_residues(self, n=20): return sorted(self.protein_records(), key=lambda r: -r['score'])[:n] class TargetabilityPredictor: """Holds the loaded network so a Space can score many structures without reloading.""" def __init__(self, weights_path=DEFAULT_WEIGHTS, device=None): self.device = pt.device(device) if device is not None else pt.device( "cuda" if pt.cuda.is_available() else "cpu" ) self.model = Model(config_model) state = pt.load(weights_path, map_location="cpu") # strict: an architecture/checkpoint mismatch must fail here rather than # produce confident-looking scores from partly random weights self.model.load_state_dict(state) self.model = self.model.eval().to(self.device) self.weights_path = weights_path def to(self, device): """Move the loaded model. Needed for ZeroGPU, where no CUDA device exists at import time and one appears only inside the decorated request handler.""" self.device = pt.device(device) self.model = self.model.to(self.device) return self @pt.no_grad() def predict(self, pdb_filepath, chains=None, out_filepath=None): """Score one structure. `chains` is an optional iterable of author chain ids. Everything not selected is removed *before* the graph is built, so the model genuinely does not see it -- selecting the target chain of a complex gives the same answer as feeding a target-only file, which is how the model was trained and applied. """ structure = read_pdb(pdb_filepath) if chains: wanted = {c.strip() for c in chains if c.strip()} # chain_name is ":" at this stage m = np.array([cn.split(':')[0] in wanted for cn in structure['chain_name']]) if not m.any(): present = sorted({cn.split(':')[0] for cn in structure['chain_name']}) raise ValueError( f"none of the requested chains {sorted(wanted)} are in the file; " f"it contains {present}" ) structure = {k: v[m] for k, v in structure.items()} # clean, split into subunits, drop non-atomic and duplicate subunits subunits = preprocess_structure(structure) if not subunits: raise ValueError( "nothing left to score after preprocessing -- the file may be a " "CA-only trace, or contain no polymer chain" ) # all remaining subunits are scored together, as one graph structure = concatenate_chains(subunits) n_atoms = structure['xyz'].shape[0] if n_atoms > MAX_ATOMS: raise ValueError( f"{n_atoms} atoms exceeds the {MAX_ATOMS}-atom limit of this deployment; " f"select fewer chains, or run inference.py locally" ) # encode: coordinates + residue membership, and the element one-hot only. # encode_features returns (element, resname, atomname); the BindCraft head was # trained on the 30-wide element block alone, matching config_model['em']['N0']. X, M = encode_structure(structure) q = encode_features(structure)[0] ids_topk = extract_topology_knn(X, NUM_NN) # pack into a batch of one and install the sink node at index 0 X, ids_topk, q, M = collate_batch_features([[X, ids_topk, q, M]], max_num_nn=NUM_NN) z = self.model( X.to(self.device), ids_topk.to(self.device), q.to(self.device), M.float().to(self.device), ) p = pt.sigmoid(z[:, 0]).cpu().numpy() # per-residue score -> per-atom, then back onto the input's own numbering structure = encode_bfactor(structure, p) st, records = map_scores_to_original( pdb_filepath, structure['xyz'], structure['bfactor'], out_filepath ) return Prediction( pdb_text=st.make_pdb_string(), records=records, num_residues_scored=int(p.shape[0]), num_atoms=int(n_atoms), chains=tuple(sorted({r['chain'] for r in records})), ) def main(argv=None): ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) ap.add_argument("pdb", nargs="+", help="input PDB/mmCIF file(s)") ap.add_argument("-o", "--output", default=None, help="output PDB (single input only)") ap.add_argument("--outdir", default=None, help="output directory for multiple inputs") ap.add_argument("--chains", default=None, help="comma-separated chain ids to score (default: all)") ap.add_argument("--weights", default=DEFAULT_WEIGHTS) ap.add_argument("--csv", action="store_true", help="also write a per-residue CSV next to each output") args = ap.parse_args(argv) if args.output and len(args.pdb) > 1: ap.error("--output takes a single input; use --outdir for several") chains = args.chains.split(",") if args.chains else None predictor = TargetabilityPredictor(args.weights) for pdb_filepath in args.pdb: if args.output: out = args.output else: stem = os.path.splitext(os.path.basename(pdb_filepath))[0] outdir = args.outdir or os.path.dirname(os.path.abspath(pdb_filepath)) os.makedirs(outdir, exist_ok=True) out = os.path.join(outdir, f"{stem}_targetability.pdb") pred = predictor.predict(pdb_filepath, chains=chains, out_filepath=out) s = pred.scores() print( f"{pdb_filepath} -> {out} " f"({pred.num_residues_scored} residues, chains {''.join(pred.chains)}, " f"max={s.max():.3f}, mean={s.mean():.3f})" ) if args.csv: print(f" {pred.to_csv(os.path.splitext(out)[0] + '.csv')}") return 0 if __name__ == "__main__": sys.exit(main())