Spaces:

arvban
/

bindscout

Running on Zero

bindscout / inference.py

BindScout: BindCraft patch-coverage targetability

d847871 verified 20 days ago

9 kB

"""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 "<author id>:<model index>" 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())

	"""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 "<author id>:<model index>" 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())