import argparse import os import MDAnalysis as mda import numpy as np import pandas as pd import torch as th import torch.multiprocessing from torch_geometric.loader import DataLoader from onescience.datapipes.genscore.data import VSDataset from models.model.model import GatedGCN, GenScore, GraphTransformer from onescience.metrics.genscore.utils import run_an_eval_epoch torch.multiprocessing.set_sharing_strategy("file_system") def _default_model_path() -> str: return os.path.join( os.path.dirname(__file__), "..", "..", "..", "..", "examples", "biosciences", "genscore", "trained_models", "GT_0.0_1.pth", ) def parse_args(): parser = argparse.ArgumentParser(description="Run GenScore protein-ligand scoring.") parser.add_argument("-p", "--prot", required=True, help="Input protein or pocket file (.pdb).") parser.add_argument("-l", "--lig", required=True, help="Input ligand file (.sdf/.mol2).") parser.add_argument( "-m", "--model", default=_default_model_path(), help="Path to a trained GenScore checkpoint.", ) parser.add_argument( "-e", "--encoder", default="gt", choices=["gt", "gatedgcn"], help="Protein and ligand graph encoder.", ) parser.add_argument("-o", "--outprefix", default="out", help="Output file prefix.") parser.add_argument( "-gen_pocket", "--gen_pocket", action="store_true", default=False, help="Generate a pocket from the input protein.", ) parser.add_argument( "-c", "--cutoff", default=10.0, type=float, help="Pocket and interaction cutoff distance.", ) parser.add_argument("-rl", "--reflig", default=None, help="Reference ligand for pocket generation.") parser.add_argument( "-pl", "--parallel", default=False, action="store_true", help="Build ligand graphs in parallel.", ) parser.add_argument( "-ac", "--atom_contribution", default=False, action="store_true", help="Compute atom-level score contributions.", ) parser.add_argument( "-rc", "--res_contribution", default=False, action="store_true", help="Compute residue-level score contributions.", ) parser.add_argument("--batch_size", type=int, default=128) parser.add_argument("--num_workers", type=int, default=10) args = parser.parse_args() if args.gen_pocket and args.reflig is None: raise ValueError("If pocket generation is enabled, --reflig must be provided.") if args.atom_contribution and args.res_contribution: raise ValueError("Only one of atom_contribution and res_contribution is supported.") return args def _build_encoder(encoder, kwargs): if encoder == "gt": ligmodel = GraphTransformer( in_channels=kwargs["num_node_featsl"], edge_features=kwargs["num_edge_featsl"], num_hidden_channels=kwargs["hidden_dim0"], activ_fn=th.nn.SiLU(), transformer_residual=True, num_attention_heads=4, norm_to_apply="batch", dropout_rate=0.15, num_layers=6, ) protmodel = GraphTransformer( in_channels=kwargs["num_node_featsp"], edge_features=kwargs["num_edge_featsp"], num_hidden_channels=kwargs["hidden_dim0"], activ_fn=th.nn.SiLU(), transformer_residual=True, num_attention_heads=4, norm_to_apply="batch", dropout_rate=0.15, num_layers=6, ) else: ligmodel = GatedGCN( in_channels=kwargs["num_node_featsl"], edge_features=kwargs["num_edge_featsl"], num_hidden_channels=kwargs["hidden_dim0"], residual=True, dropout_rate=0.15, equivstable_pe=False, num_layers=6, ) protmodel = GatedGCN( in_channels=kwargs["num_node_featsp"], edge_features=kwargs["num_edge_featsp"], num_hidden_channels=kwargs["hidden_dim0"], residual=True, dropout_rate=0.15, equivstable_pe=False, num_layers=6, ) return ligmodel, protmodel def scoring( prot, lig, modpath, cut=10.0, gen_pocket=False, reflig=None, encoder="gt", atom_contribution=False, res_contribution=False, explicit_H=False, use_chirality=True, parallel=False, **kwargs, ): data = VSDataset( ligs=lig, prot=prot, cutoff=cut, gen_pocket=gen_pocket, reflig=reflig, explicit_H=explicit_H, use_chirality=use_chirality, parallel=parallel, ) test_loader = DataLoader( dataset=data, batch_size=kwargs["batch_size"], shuffle=False, num_workers=kwargs["num_workers"], ) ligmodel, protmodel = _build_encoder(encoder, kwargs) model = GenScore( ligmodel, protmodel, in_channels=kwargs["hidden_dim0"], hidden_dim=kwargs["hidden_dim"], n_gaussians=kwargs["n_gaussians"], dropout_rate=kwargs["dropout_rate"], dist_threhold=kwargs["dist_threhold"], ).to(kwargs["device"]) checkpoint = th.load(modpath, map_location=th.device(kwargs["device"])) model.load_state_dict(checkpoint["model_state_dict"]) if atom_contribution: preds, at_contrs, _ = run_an_eval_epoch( model, test_loader, pred=True, atom_contribution=True, res_contribution=False, dist_threhold=kwargs["dist_threhold"], device=kwargs["device"], ) atids = [f"{a.GetSymbol()}{a.GetIdx()}" for a in data.ligs[0].GetAtoms()] return data.ids, preds, atids, at_contrs if res_contribution: preds, _, res_contrs = run_an_eval_epoch( model, test_loader, pred=True, atom_contribution=False, res_contribution=True, dist_threhold=kwargs["dist_threhold"], device=kwargs["device"], ) universe = mda.Universe(data.prot) resids = [ f"{chain_id}_{resname}{resid}" for chain_id, resname, resid in zip( universe.residues.chainIDs, universe.residues.resnames, universe.residues.resids, ) ] return data.ids, preds, resids, res_contrs preds = run_an_eval_epoch( model, test_loader, pred=True, dist_threhold=kwargs["dist_threhold"], device=kwargs["device"], ) return data.ids, preds def _runtime_args(args): return { "batch_size": args.batch_size, "dist_threhold": 5.0, "device": "cuda" if th.cuda.is_available() else "cpu", "num_workers": args.num_workers, "num_node_featsp": 41, "num_node_featsl": 41, "num_edge_featsp": 5, "num_edge_featsl": 10, "hidden_dim0": 128, "hidden_dim": 128, "n_gaussians": 10, "dropout_rate": 0.15, } def main(): inargs = parse_args() runtime = _runtime_args(inargs) common = { "prot": inargs.prot, "lig": inargs.lig, "modpath": inargs.model, "cut": inargs.cutoff, "gen_pocket": inargs.gen_pocket, "reflig": inargs.reflig, "encoder": inargs.encoder, "explicit_H": False, "use_chirality": True, "parallel": inargs.parallel, **runtime, } if inargs.atom_contribution: ids, scores, atids, at_contrs = scoring(atom_contribution=True, **common) df = pd.DataFrame(at_contrs).T df.columns = ids df.index = atids df = df[df.apply(np.sum, axis=1) != 0].T dfx = pd.DataFrame(zip(*(ids, scores)), columns=["id", "score"]) dfx.index = dfx.id df = pd.concat([dfx["score"], df], axis=1) df.sort_values("score", ascending=False, inplace=True) df.to_csv(f"{inargs.outprefix}_at.csv") elif inargs.res_contribution: ids, scores, resids, res_contrs = scoring(res_contribution=True, **common) df = pd.DataFrame(res_contrs).T df.columns = ids df.index = resids df = df[df.apply(np.sum, axis=1) != 0].T dfx = pd.DataFrame(zip(*(ids, scores)), columns=["id", "score"]) dfx.index = dfx.id df = pd.concat([dfx["score"], df], axis=1) df.sort_values("score", ascending=False, inplace=True) df.to_csv(f"{inargs.outprefix}_res.csv") else: ids, scores = scoring(**common) df = pd.DataFrame(zip(*(ids, scores)), columns=["id", "score"]) df.sort_values("score", ascending=False, inplace=True) df.to_csv(f"{inargs.outprefix}.csv", index=False) if __name__ == "__main__": main()