Digital Trials pipeline configured for WES

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.

nf_metabol_screen_adaptive/app/convert.py

@@ -0,0 +1,134 @@
+#!/usr/bin/env python3
+
+import argparse
+import pandas as pd
+import sys
+from pathlib import Path
+
+def create_smiles_file(drug_csv_path, metabolite_file_path, output_path="smiles.smi"):
+    """
+    Create smiles.smi file by combining drug CSV and metabolite file.
+    
+    Args:
+        drug_csv_path: Path to CSV file with 'SMILES' header
+        metabolite_file_path: Path to tab-separated file with SMILES and metabolite IDs
+        output_path: Output file path (default: smiles.smi)
+    """
+    
+    # Convert to Path objects for better handling
+    drug_csv_path = Path(drug_csv_path)
+    metabolite_file_path = Path(metabolite_file_path)
+    output_path = Path(output_path)
+    
+    # Check if input files exist
+    if not drug_csv_path.exists():
+        raise FileNotFoundError(f"Drug CSV file not found: {drug_csv_path}")
+    
+    if not metabolite_file_path.exists():
+        raise FileNotFoundError(f"Metabolite file not found: {metabolite_file_path}")
+    
+    # Read drug CSV
+    try:
+        drug_df = pd.read_csv(drug_csv_path)
+    except Exception as e:
+        raise Exception(f"Error reading drug CSV file: {e}")
+    
+    # Validate drug CSV has SMILES column
+    if 'SMILES' not in drug_df.columns:
+        raise ValueError("Drug CSV file must have 'SMILES' column")
+    
+    # Check if metabolite file contains "NO_METABOLITES"
+    try:
+        with open(metabolite_file_path, 'r') as f:
+            first_line = f.readline().strip()
+    except Exception as e:
+        raise Exception(f"Error reading metabolite file: {e}")
+    
+    # Handle no metabolites case
+    if first_line == "NO_METABOLITES":
+        print("No metabolites found, creating output file with drugs only")
+        
+        # Create drug dataframe with drug_N format
+        drug_output = drug_df.copy()
+        drug_output['ID'] = [f'drug_{i}' for i in range(len(drug_df))]
+        drug_output = drug_output[['SMILES', 'ID']]
+        
+        # Write only drugs to output file
+        try:
+            drug_output.to_csv(output_path, sep='\t', header=False, index=False)
+        except Exception as e:
+            raise Exception(f"Error writing output file: {e}")
+        
+        print(f"Successfully created {output_path}")
+        print(f"  - {len(drug_df)} drugs")
+        print(f"  - 0 metabolites")
+        return
+    
+    # Read metabolite file normally (tab-separated, no header)
+    try:
+        metabolite_df = pd.read_csv(metabolite_file_path, sep='\t', header=None, names=['SMILES', 'ID'])
+    except Exception as e:
+        raise Exception(f"Error reading metabolite file: {e}")
+    
+    # Validate metabolite file has at least 2 columns
+    if metabolite_df.shape[1] < 2:
+        raise ValueError("Metabolite file must have at least 2 columns (SMILES and ID)")
+    
+    # Create drug dataframe with drug_N format
+    drug_output = drug_df.copy()
+    drug_output['ID'] = [f'drug_{i}' for i in range(len(drug_df))]
+    drug_output = drug_output[['SMILES', 'ID']]
+    
+    # Use metabolite dataframe as-is (already has SMILES and ID columns)
+    metabolite_output = metabolite_df[['SMILES', 'ID']]
+    
+    # Combine drug and metabolite dataframes
+    combined_df = pd.concat([drug_output, metabolite_output], ignore_index=True)
+    
+    # Write to output file (tab-separated, no header)
+    try:
+        combined_df.to_csv(output_path, sep='\t', header=False, index=False)
+    except Exception as e:
+        raise Exception(f"Error writing output file: {e}")
+    
+    print(f"Successfully created {output_path}")
+    print(f"  - {len(drug_df)} drugs")
+    print(f"  - {len(metabolite_df)} metabolites")
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Create smiles.smi file by combining drug CSV and metabolite file",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  python3 create_smiles.py drug.csv metabolites.txt
+  python3 create_smiles.py drug.csv metabolites.txt -o compounds.smi
+        """
+    )
+    
+    parser.add_argument(
+        'drug_csv', 
+        help='Path to CSV file with SMILES column containing drug compounds'
+    )
+    
+    parser.add_argument(
+        'metabolite_file', 
+        help='Path to tab-separated file with SMILES and metabolite IDs'
+    )
+    
+    parser.add_argument(
+        '-o', '--output', 
+        default='smiles.smi',
+        help='Output file path (default: smiles.smi)'
+    )
+    
+    args = parser.parse_args()
+    
+    try:
+        create_smiles_file(args.drug_csv, args.metabolite_file, args.output)
+    except Exception as e:
+        print(f"Error: {e}", file=sys.stderr)
+        sys.exit(1)
+
+if __name__ == "__main__":
+    main()

nf_metabol_screen_adaptive/app/get_round_2.py

@@ -0,0 +1,225 @@
+import argparse
+import pandas as pd
+import numpy as np
+from collections import defaultdict
+from itertools import chain
+from pathlib import Path
+
+import zarr
+from zarr.storage import LocalStore
+
+# Load reference data
+seq_exist = pd.read_csv('/app/MANE_all_transcipts.csv')
+
+
+def get_round_2(threshold, workdir, round, drug_csv: Path):
+    workdir = Path(workdir)
+    fasta_name = "patient_0"
+    drug_name = drug_csv.stem
+    name_out = f"{drug_name}_{fasta_name}"
+
+    # Load all zarr score stores produced by screen.py.
+    results_zarr = list(workdir.glob('*_results.zarr'))
+    if not results_zarr:
+        raise FileNotFoundError("No results Zarr stores found")
+
+    print("Loading Zarr results...")
+    # Lazy zarr handles per file; we read the (proteins, compounds) score matrix
+    # in protein chunks rather than loading it fully — keeps peak memory bounded
+    # regardless of compound count.
+    zarr_groups = []
+    compound_ids_per = []
+    protein_ids_per = []
+    for z_path in results_zarr:
+        store = LocalStore(z_path, read_only=True)
+        group = zarr.open(store, mode="r")
+        zarr_groups.append(group)
+        # numpy 2.4+ returns StringDType for these arrays; convert via tolist()
+        # because direct .astype(str) raises "cannot cast StringDType to StrDType".
+        compound_ids_per.append(np.array(group["compound_ids"][:].tolist(), dtype=str))
+        protein_ids_per.append(np.array(group["protein_ids"][:].tolist(), dtype=str))
+
+    # All zarrs share the same compound axis (drug + same metabolites).
+    all_compound_ids = compound_ids_per[0]
+    for cids in compound_ids_per[1:]:
+        if cids.shape != all_compound_ids.shape or not np.array_equal(cids, all_compound_ids):
+            raise ValueError("All *_results.zarr stores must share the same compound_ids axis")
+    all_protein_ids = np.concatenate(protein_ids_per) if len(protein_ids_per) > 1 else protein_ids_per[0]
+    n_proteins = sum(g["scores"].shape[0] for g in zarr_groups)
+    n_compounds = all_compound_ids.shape[0]
+
+    # ---- Compound-axis indices (tiny) ----
+    drug_mask_cmpd = np.array([s.startswith("drug") for s in all_compound_ids], dtype=bool)
+    if not drug_mask_cmpd.any():
+        raise ValueError("No 'drug_*' compounds present in compound_ids")
+    drug_cols_idx = np.where(drug_mask_cmpd)[0]
+    drug0_search = np.where(all_compound_ids == 'drug_0')[0]
+    if drug0_search.size == 0:
+        raise ValueError("'drug_0' not found in compound_ids")
+    drug0_idx = int(drug0_search[0])
+
+    # ---- Stream the score matrix in protein chunks, writing significant_interactions
+    #      incrementally and accumulating per-protein summary stats. Avoids both the
+    #      full score matrix load and the giant in-memory above-threshold DataFrame. ----
+    # Per chunk peak ≈ CHUNK_PROTEINS * n_compounds * 4 bytes; for 1000×38k ≈ 150 MB.
+    CHUNK_PROTEINS = 1000
+    drug_position = np.zeros(n_proteins, dtype=np.int64)
+    drug_0_col_full = np.empty(n_proteins, dtype=np.float32)
+    prot_count = defaultdict(int)    # protein_id -> count of above-threshold compounds
+    prot_sum = defaultdict(float)    # protein_id -> sum of above-threshold scores
+
+    # Open significant_interactions.tsv before the loop (round 2 only); write rows per chunk.
+    sig_path = workdir / f'{name_out}_significant_interactions.tsv'
+    sig_writer = None
+    smile_dict = None
+    if round == 2:
+        test_smiles = pd.read_csv(workdir / 'smiles.smi', sep='\t', header=None)
+        smile_dict = dict(zip(test_smiles[1], test_smiles[0]))
+        sig_writer = open(sig_path, 'w')
+        sig_writer.write("drug/metabolite\tsmile\ttranscipt\tconplex_score\n")
+
+    cursor = 0
+    for group in zarr_groups:
+        scores_arr = group["scores"]
+        n_this = scores_arr.shape[0]
+        for start in range(0, n_this, CHUNK_PROTEINS):
+            end = min(start + CHUNK_PROTEINS, n_this)
+            chunk = np.asarray(scores_arr[start:end])  # (chunk_size, n_compounds)
+            global_start = cursor + start
+            global_end = cursor + end
+            chunk_size = end - start
+
+            # drug score columns + max-per-protein for this chunk
+            chunk_max_drug = chunk[:, drug_cols_idx].max(axis=1)
+
+            # drug_position[i] = # compounds with score strictly > max_drug_score_i.
+            # Matches production's stable-sort + first-drug-position convention.
+            drug_position[global_start:global_end] = (chunk > chunk_max_drug[:, None]).sum(axis=1)
+
+            # drug_0 column
+            drug_0_col_full[global_start:global_end] = chunk[:, drug0_idx].astype(np.float32, copy=False)
+
+            # Above-threshold rows for this chunk
+            pi_local, ci = np.where(chunk > threshold)
+            if pi_local.size:
+                scores_vec = chunk[pi_local, ci].astype(np.float32, copy=False)
+
+                # Per-protein aggregates accumulated via bincount on local indices.
+                counts_local = np.bincount(pi_local, minlength=chunk_size)
+                sums_local = np.bincount(pi_local, weights=scores_vec.astype(np.float64), minlength=chunk_size)
+                for pi in np.flatnonzero(counts_local):
+                    pid = all_protein_ids[global_start + pi]
+                    prot_count[pid] += int(counts_local[pi])
+                    prot_sum[pid] += float(sums_local[pi])
+
+                # Write significant_interactions rows for this chunk directly to the file.
+                if sig_writer is not None:
+                    compound_ids_vec = all_compound_ids[ci]
+                    protein_ids_vec = all_protein_ids[global_start + pi_local]
+                    smiles_vec = np.array(
+                        [smile_dict.get(c, "UNKNOWN") for c in compound_ids_vec]
+                    )
+                    pd.DataFrame({
+                        "drug/metabolite": compound_ids_vec,
+                        "smile":           smiles_vec,
+                        "transcipt":       protein_ids_vec,
+                        "conplex_score":   scores_vec,
+                    }).to_csv(sig_writer, sep='\t', index=False, header=False)
+        cursor += n_this
+
+    if sig_writer is not None:
+        sig_writer.close()
+
+    drug_0_score = pd.DataFrame({
+        "Drug":       "drug_0",
+        "Transcript": all_protein_ids,
+        "Score":      drug_0_col_full,
+    })
+
+    if round == 2:
+        drug_0_score.to_csv(workdir / f'{name_out}_drug_scores.tsv', sep='\t', index=False)
+
+    # ---- Per-protein summary (only proteins with any above-threshold compound) ----
+    drug_pos_map = dict(zip(all_protein_ids, drug_position))
+
+    if round == 1 and not prot_count:
+        # Equivalent to old "transcipts_summary.empty" early-out for round 1.
+        pd.DataFrame().to_csv(workdir / 'round_1.csv', index=False)
+        with open(workdir / 'round_1.fasta', 'w') as f:
+            f.write("all_data_is_filtered_out\n")
+        return 'STOP NO TRANSCIPTS ABOVE THRESHOLD'
+
+    if prot_count:
+        proteins_with_above = list(prot_count.keys())
+        # mean dtype kept as float32 to match the precision of the previous
+        # groupby-on-float32 path (CSV output is bit-identical that way).
+        means = np.array(
+            [prot_sum[p] / prot_count[p] for p in proteins_with_above],
+            dtype=np.float32,
+        )
+        summary = pd.DataFrame({
+            'transcipt_name': proteins_with_above,
+            'number_of_iteracting_compounds': [prot_count[p] for p in proteins_with_above],
+            'mean_binding_above_threshold': means,
+        })
+    else:
+        summary = pd.DataFrame(columns=[
+            'transcipt_name', 'number_of_iteracting_compounds', 'mean_binding_above_threshold'
+        ])
+    summary['drug_position'] = summary['transcipt_name'].map(drug_pos_map)
+
+    path_out = workdir / f'round_{round}.csv'
+    if len(summary) == 0:
+        with path_out.open("w") as f:
+            f.write("NO TRANSCIPTS ABOVE THRESHOLD")
+        return 'STOP NO TRANSCIPTS ABOVE THRESHOLD'
+
+    summary['if_drug_above_threshold'] = (
+        summary['number_of_iteracting_compounds'] > summary['drug_position']
+    )
+
+    # MANE symbol lookup; strip mutated-protein "_2" suffix.
+    clean_transcript_names = [str(t).split('_')[0] for t in summary['transcipt_name']]
+    seq_lookup = seq_exist.set_index('transcipt')['symbol']
+    summary['protein_name'] = [
+        seq_lookup.get(t, "") for t in clean_transcript_names
+    ]
+
+    summary.to_csv(path_out, index=False)
+
+    # ---- Round 1 also writes the round-2 FASTA ----
+    if round == 1:
+        name_2_filtered = list(summary['protein_name'])
+        transcipts_2 = list(chain.from_iterable(
+            list(seq_exist[seq_exist['symbol'] == sym]['transcipt'])
+            for sym in name_2_filtered
+        ))
+        transcipts_2 = list(np.unique(transcipts_2))
+        existing = set(summary['transcipt_name'])
+        transcipts_2 = [t for t in transcipts_2 if t not in existing]
+
+        fasta_new = []
+        for t in transcipts_2:
+            row = seq_exist[seq_exist['transcipt'] == t]
+            if not row.empty:
+                fasta_new.append(f">{t}")
+                fasta_new.append(row['seq'].iloc[0])
+
+        with open(workdir / 'round_2.fasta', 'w') as f:
+            for line in fasta_new:
+                f.write(f"{line}\n")
+
+    print(f"Round {round} processing complete.")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Process and screen fragments.")
+    parser.add_argument("--threshold", required=True, type=float,
+                       help="Threshold for 1. round of conplex scores")
+    parser.add_argument("--workdir", required=True, type=Path, help="workdir")
+    parser.add_argument("--round", required=True, type=int, help="Round 1 or 2")
+    parser.add_argument("--drug-csv", required=True, type=Path,
+                       help="csv file with drug smiles")
+    args = parser.parse_args()
+
+    get_round_2(args.threshold, args.workdir, args.round, args.drug_csv)

nf_metabol_screen_adaptive/app/old_get_round_20.py

@@ -0,0 +1,143 @@
+import argparse
+import pandas as pd
+import numpy as np
+from itertools import chain
+from pathlib import PosixPath as Path
+
+seq_exist = pd.read_csv('/app/MANE_all_transcipts.csv')
+
+def get_round_2(threshold, workdir, round, drug_csv : Path):
+    workdir = Path(workdir)
+    
+    # Get output names
+    # Get patient fasta name
+    # fasta_files = list(workdir.glob('*_variants_transcript_id_mutations.fasta'))
+    # if not fasta_files:
+    #     raise FileNotFoundError("No variants transcript mutations fasta file found")
+    # fasta_name = fasta_files[0].stem.replace('_variants_transcript_id_mutations', '')
+    fasta_name = "patient_0"
+    
+    # Get test drug
+    drug_name = drug_csv.stem
+    
+    name_out = f"{drug_name}_{fasta_name}"
+    
+    # Get all work dir complex files
+    results_files = list(workdir.glob('*_results.tsv'))
+    if not results_files:
+        raise FileNotFoundError("No results TSV files found")
+    
+    round_1_score_list = [str(f) for f in results_files]
+    
+    # Read 1st round
+    transcipts_1 = [
+        pd.read_csv(i, sep='\t', header=None).sort_values([2], ascending=False)
+        for i in round_1_score_list
+        ]
+    
+    # Get position of drug interaction vs all metabolites
+    drug_pos = [
+        [
+            n
+            for n, j in enumerate(list(i[0]))
+            if j[:4] == 'drug'
+        ] 
+        for i in transcipts_1
+    ]
+    
+    # Drug score
+    drug_0_score = pd.concat([i[i[0] =='drug_0'] for i in transcipts_1]).rename({0:'Drug', 1:'Transcript', 2:'Score'}, axis=1)
+    
+    # Filter all below threshold
+    transcipts_1 = [i[i[2] > threshold] for i in transcipts_1]
+    
+    # Save transcript_complex above threshold to one file
+    if round == 2:
+        inter_import = pd.concat(transcipts_1, ignore_index=True).rename({0:'drug/metabolite',1:'transcipt',2:'conplex_score'},axis=1)
+        # Add drug
+        test_smiles = pd.read_csv(workdir / 'smiles.smi', sep='\t', header=None)
+        smi_ = [test_smiles[test_smiles[1] == i].iloc[0][0] for i in list(inter_import['drug/metabolite'])]
+        inter_import['smile'] = smi_
+        inter_import[['drug/metabolite','smile','transcipt','conplex_score']].to_csv(
+            workdir / f'{name_out}_significant_interactions.tsv', sep='\t', index=False
+        )
+        drug_0_score.to_csv(workdir / f'{name_out}_drug_scores.tsv', sep='\t', index=False)
+    
+    transcipts_1 = pd.DataFrame([
+        (Path(j).stem.replace('_results', ''), i.shape[0], i[2].mean(), k[0]) 
+        for i, j, k in zip(transcipts_1, round_1_score_list, drug_pos) 
+        if i.shape[0] != 0
+    ])
+    
+    # Check if any interaction is above threshold for first round 
+    if round == 1:
+        if transcipts_1.shape == (0,0):
+            transcipts_1.to_csv(workdir / 'round_1.csv', index=False)
+            with open(workdir / 'round_1.fasta', 'w') as f:
+                f.write("all_data_is_filtered_out\n")
+            return 'STOP NO TRANSCIPTS ABOVE THRESHOLD'
+
+    transcipts_1 = transcipts_1.rename({
+        0:'transcipt_name',
+        1:'number_of_iteracting_compounds', 
+        2:'mean_binding_above_threshold', 
+        3:'drug_position'
+    }, axis='columns')
+    
+    path_out = workdir / f'round_{round}.csv'
+    
+    if len(transcipts_1) == 0:
+        with path_out.open("w") as f:
+            f.write("NO TRANSCIPTS ABOVE THRESHOLD")
+            return 'STOP NO TRANSCIPTS ABOVE THRESHOLD'
+        
+    
+    transcipts_1['if_drug_above_threshold'] = transcipts_1.iloc[:,1] > transcipts_1.iloc[:,3]
+    
+    # If protein is mutated it has _2 in name, removes it
+    transcipt_names = [i.split('_')[0] for i in list(transcipts_1['transcipt_name'])]
+    temp = [
+        seq_exist[seq_exist['transcipt'] == i]
+        for i in transcipt_names
+    ]
+    transcipts_1['protein_name'] =[
+        i.iloc[0]['symbol'] if len(i) > 0 else ""
+        for i in temp
+    ]
+    
+    # Save data on first round
+    transcipts_1.to_csv(path_out, index=False)
+    
+    # Get fasta for second round
+    if round == 1:
+        # Get all transcripts of proteins above threshold
+        name_2_filttered = list(transcipts_1['protein_name'])
+        transcipts_2 = [list(seq_exist[seq_exist['symbol'] == i]['transcipt']) for i in name_2_filttered]
+        transcipts_2 = list(chain(*transcipts_2))
+        transcipts_2 = list(np.unique((transcipts_2)))
+        
+        # Filter out transcripts already ran through complex
+        transcipts_2 = list(np.array(transcipts_2)[
+            [i not in list(transcipts_1['transcipt_name']) for i in transcipts_2]
+        ])
+        
+        fasta_new = [
+            ['>'+i, seq_exist[seq_exist['transcipt'] == i]['seq'].iloc[0]] 
+            for i in transcipts_2
+        ]
+        fasta_new = list(chain(*fasta_new))
+        
+        # Write fasta to run
+        with open(workdir / 'round_2.fasta', 'w') as f:
+            for line in fasta_new:
+                f.write(f"{line}\n")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Process and screen fragments.")
+    parser.add_argument("--threshold", required=True, type=float, help="Threshold for 1. round.of conplex scores")
+    parser.add_argument("--workdir", required=True, type=Path, help="workdir")
+    parser.add_argument("--round", required=True, type=int, help="Round 1 or 2")
+    parser.add_argument("--drug-csv", required=True, type=Path, help="csv file with drug smiles")
+    args = parser.parse_args()
+
+    get_round_2(args.threshold, args.workdir, args.round, args.drug_csv)

nf_metabol_screen_adaptive/app/old_screen.py

@@ -0,0 +1,127 @@
+# 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")

nf_metabol_screen_adaptive/app/screen.py

@@ -0,0 +1,144 @@
+# 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 zarr.codecs import ZstdCodec
+
+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_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
+    
+    # Run inference
+    outputs = session.run([output_name], {input_name: input_vector})
+    
+    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"][:]
+
+# Create output zarr store
+path_outdir = Path("./")
+zarr_name = path_zarr.stem + "_results.zarr"
+path_output_zarr = path_outdir / zarr_name
+
+# Create zarr store
+store_output = LocalStore(path_output_zarr, read_only=False)
+group_output = zarr.open(store_output, mode="w")
+
+out_compress =[
+    ZstdCodec( level=13 )
+]
+
+# Store the data
+# Similarity matrix: (n_proteins, n_compounds)
+z1 = group_output.create_array(
+    "scores",
+    dtype="float32",
+    shape=my_results.shape,
+    chunks=(min(1000, my_results.shape[0]), my_results.shape[1]),
+    compressors=out_compress
+)
+z1[:] = my_results
+
+# Compound IDs
+z2 = group_output.create_array(
+    "compound_ids",
+    dtype="str",
+    shape=(len(df_metabolite),),
+    compressors=out_compress
+)
+z2[:] = df_metabolite["compound_id"].values
+
+# Protein IDs
+z3 = group_output.create_array(
+    "protein_ids",
+    dtype="str",
+    shape=(len(pids),),
+    compressors=out_compress
+)
+z3[:] = pids
+
+# Store metadata
+group_output.attrs['n_proteins'] = my_results.shape[0]
+group_output.attrs['n_compounds'] = my_results.shape[1]
+group_output.attrs['source_zarr'] = str(path_zarr)
+
+print(f"Results saved to {path_output_zarr}")
+print(f"Shape: {my_results.shape[0]} proteins x {my_results.shape[1]} compounds")
+print("PROCESS COMPLETED SUCCESSFULLY")