9e75f44f1a
Source-only snapshot of the cluster branch for WES execution. Large reference files (HPA/MANE/ensemble FASTA, model weights, ~597 MB) are omitted: they are baked into the container images at build time and mounted from the dreamdock-data PVC at runtime, and exceed the Gitea request size limit. Pipeline entry point is main.nf, which orchestrates the biotransformer, conplex and tissue modules as a single workflow. Ligand inputs are read from the eureka workspace; protein_zarr and chembl_db come from the dreamdock-data PVC.
128 lines
3.6 KiB
Python
128 lines
3.6 KiB
Python
# Custom Rust modules
|
|
import mol_fingerprint as mf
|
|
import bitexpand
|
|
|
|
import pandas as pd
|
|
import numpy as np
|
|
from pathlib import Path
|
|
from argparse import ArgumentParser
|
|
|
|
import zarr
|
|
from zarr.storage import LocalStore
|
|
|
|
from sklearn.preprocessing import normalize
|
|
from parallelbar import progress_map
|
|
|
|
import onnxruntime as ort
|
|
print('CUDA available:', 'CUDAExecutionProvider' in ort.get_available_providers())
|
|
|
|
|
|
|
|
def parse_args():
|
|
parser = ArgumentParser()
|
|
parser.add_argument("-i","--input-tsv", type=Path, required=True, help="Input TSV file with molecules to screen. It must not have any labels and has only two columns, 1st molecule_label, 2nd molecule_smiles")
|
|
parser.add_argument("-z","--zarr-protein", type=Path, required=True, help="Input Zarr DB that holds the protein ids and their normalized projection vectors")
|
|
return parser
|
|
|
|
parser = parse_args()
|
|
|
|
args = parser.parse_args()
|
|
|
|
# path_fp_basedir : Path = args.input_dir
|
|
# my_id :int = args.chunk_id
|
|
|
|
path_metabolites : Path = args.input_tsv
|
|
if not path_metabolites.is_file():
|
|
raise ValueError(f"{path_metabolites} IS MISSING OR NOT A FILE")
|
|
|
|
df_metabolite = pd.read_csv(
|
|
path_metabolites,
|
|
sep="\t",
|
|
header=None,
|
|
names=["smiles" ,"compound_id"]
|
|
)
|
|
|
|
path_zarr = args.zarr_protein
|
|
store_zarr = LocalStore(path_zarr , read_only=True)
|
|
group_zarr = zarr.open(store_zarr, mode="r")
|
|
protein_vecs = (group_zarr["normalized"][:])
|
|
|
|
|
|
path_onnx = Path("/app/drug_projector.onnx")
|
|
|
|
def load_model(model_path):
|
|
"""Load ONNX model with GPU support"""
|
|
model_path = Path(model_path)
|
|
|
|
# Set up providers for GPU execution
|
|
# providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
|
|
providers = ['CPUExecutionProvider']
|
|
|
|
# Create inference session
|
|
session = ort.InferenceSession(str(model_path), providers=providers)
|
|
|
|
# Verify GPU is being used
|
|
print(f"Available providers: {session.get_providers()}")
|
|
|
|
return session
|
|
|
|
session = load_model(path_onnx)
|
|
|
|
def run_inference(input_vector):
|
|
"""Run inference on the model"""
|
|
# Get input/output names
|
|
input_name = session.get_inputs()[0].name
|
|
output_name = session.get_outputs()[0].name
|
|
|
|
# Ensure input is the correct shape and type
|
|
# if input_vector.shape != (2048,):
|
|
# raise ValueError(f"Input shape should be (2048,), got {input_vector.shape}")
|
|
|
|
# Add batch dimension if needed
|
|
input_data = input_vector
|
|
|
|
# Run inference
|
|
outputs = session.run([output_name], {input_name: input_data})
|
|
|
|
# Return the output vector (remove batch dimension)
|
|
return outputs[0]
|
|
|
|
def normalize_vectors_sklearn(vectors):
|
|
"""Normalize using sklearn - often fastest for large arrays"""
|
|
return normalize(vectors, norm='l2', axis=1)
|
|
|
|
def project_drugs(smiles : list[str]) -> np.ndarray :
|
|
fps = bitexpand.expand_bits(np.array([
|
|
np.frombuffer(fp,dtype=np.uint8)
|
|
for fp in mf.generate_fingerprints(smiles)
|
|
]))
|
|
return normalize_vectors_sklearn(run_inference(fps))
|
|
|
|
pvecs = project_drugs(df_metabolite["smiles"].values.tolist())
|
|
|
|
my_results = np.dot(protein_vecs,pvecs.T)
|
|
|
|
pids = group_zarr["ids"][:]
|
|
|
|
path_outdir = Path("./")
|
|
def write_protein(i : int):
|
|
df_out = df_metabolite[["compound_id"]].copy()
|
|
df_out["protein_id"] = pids[i]
|
|
df_out["score"] = my_results[i]
|
|
# df_out["protein_sequence"] = group_zarr["sequences"][i]
|
|
path_out = path_outdir / f"{pids[i]}_results.tsv"
|
|
|
|
df_out.to_csv(
|
|
path_out,
|
|
sep="\t",
|
|
index=False,
|
|
header=False
|
|
)
|
|
|
|
tasks = list(range(my_results.shape[0]))
|
|
progress_map(
|
|
write_protein, tasks
|
|
)
|
|
|
|
print("PROCESS COMPLETED SUCCESSFULLY")
|