| 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() |
|
|
