Solubility usability prediction NetSolP

.gitignore

@@ -8,7 +8,12 @@ datasets/
 .vscode/
 reports/
 .ruff_cache/
+results/
+results2/
+scripts/
 calculation/
+CLEAN/
+NetSolP-1.0/
 
 # C extensions
 *.so

README.md

@@ -37,16 +37,29 @@ docker start CONTAINERID
 docker exec -it CONTAINERID /bin/bash
 ```
 
-## EC number prediction with CLEAN
+## Additional predictors based on ESM1b
+Note: Deployment yet untested on absence of additional predictors. Tried to import only when called but not sure if this was enough.
+### EC number prediction with CLEAN
+Install CLEAN, download weights and place them into the hard coded locations.
 ```
-git clone https://github.com/tttianhao/CLEAN.git
+git clone https://github.com/fmoorhof/CLEAN.git
 cd CLEAN/app
 # get model weights (see below)
 python build.py install  # requires activated venv/conda
 cd ../../SelectZyme
 ```
 Download, unzip [these files](https://drive.google.com/file/d/1kwYd4VtzYuMvJMWXy6Vks91DSUAOcKpZ/view?usp=sharing) and move the contents to `data/pretrained`.
 
+### Solubility prediction with NetSolP
+Install NetSolP, download weights and place them into the hard coded locations.
+```
+git clone https://github.com/fmoorhof/NetSolP-1.0
+pip install -r NetSolP-1.0/PredictionServer/requirements_gpu.txt --extra-index-url https://download.pytorch.org/whl/cu126
+pip install -e NetSolP-1.0/
+wget https://services.healthtech.dtu.dk/services/NetSolP-1.0/netsolp-1.0.ALL.tar.gz
+tar -xzf netsolp-1.0.ALL.tar.gz NetSolP-1.0/
+mv NetSolP-1.0/models NetSolP-1.0/PredictionServer/models
+```
 
 ## Test the install
 Run some unit tests to see if SelectZyme got setup properly on your system.

app.py

@@ -14,7 +14,7 @@
 import selectzyme.pages.eda as eda
 from selectzyme.backend.embed import load_embeddings
 from selectzyme.backend.ml import dimred_caller, perform_hdbscan_clustering
-from selectzyme.backend.predict import run_clean_inference_with_embeddings
+from selectzyme.backend.predict import run_predictions
 from selectzyme.backend.utils import export_data, parse_and_preprocess
 from selectzyme.frontend.mst_plotting import MinimumSpanningTree
 from selectzyme.frontend.single_linkage_plotting import create_dendrogram
@@ -28,6 +28,7 @@ def main(app, config):
     os.makedirs(analysis_path, exist_ok=True)
 
     df = parse_and_preprocess(config, existing_file=analysis_path + "/data.csv")
+
     X = load_embeddings(df, config["project"]["plm"]["plm_model"], embedding_file=os.path.join(analysis_path, "X.npz"))
 
     # Clustering
@@ -46,22 +47,7 @@ def main(app, config):
         config["project"]["dimred"]["random_state"],
     )
 
-    # predict EC numbers with CLEAN
-    if config["project"]["plm"]["plm_model"] == "esm1b":
-        data_path = "../CLEAN/app/data/"
-        desired_split = "100"
-        df_clean = run_clean_inference_with_embeddings(
-            sequence_label_esm_emb_dict={row["accession"]: X[i] for i, row in df.iterrows()},
-            emb_train_path=f'{data_path}pretrained/{desired_split}.pt',
-            ec_csv_path=f'{data_path}split{desired_split}.csv',
-            model_ckpt_path=f'{data_path}pretrained/split{desired_split}.pth',
-            out_csv=analysis_path + "/clean_predictions.csv",
-            gmm=f'{data_path}pretrained/gmm_ensumble.pkl',
-            device="cpu",
-        )
-
-        df = df.merge(df_clean[["accession", "CLEAN_EC_pred", "CLEAN_probability"]], on="accession", how="left")
-
+    df = run_predictions(df, X, config, analysis_path)
 
     # save intermediates for external minimal dash version
     export_data(df, X_red, _mst, _linkage, analysis_path=analysis_path)

results/input_configs/test_config.yml

@@ -33,6 +33,10 @@ project:
     method: "UMAP"  # [UMAP, TSNE, openTSNE, PCA]
     n_neighbors: 15  # recommended between 5-50
 
+  predictions:
+    solubility: True  # predict solubility and usability with NetSolP
+    ec: True  # predict EC numbers with CLEAN
+
   plot_customizations:  # define marker size and shape
     objective: "cluster"  # define the legend objective for MST and SLC
 

selectzyme/backend/predict.py

@@ -2,206 +2,43 @@
 
 import logging
 import os
-import pickle
-from typing import Dict, Optional
-
-import numpy as np
-import pandas as pd
-import torch
-from CLEAN.distance_map import get_dist_map_test
-from CLEAN.evaluate import infer_confidence_gmm, maximum_separation
-from CLEAN.model import LayerNormNet
-from CLEAN.utils import get_ec_id_dict
-
-logger = logging.getLogger(__name__)
-
-
-def get_max_sep_predictions_dict(inference_df, gmm):
-    """Convert a CLEAN distance matrix into ranked EC predictions.
-
-    CLEAN produces a distance table where each column corresponds to an
-    inference sequence and each row index corresponds to an EC label. This
-    helper selects the 10 closest EC candidates per sequence, applies the
-    maximum-separation rule to keep only the most plausible prefix of those
-    candidates, and formats the result as strings like ``EC:1.1.1.1/0.1234``.
-
-    Parameters
-    ----------
-    inference_df : pandas.DataFrame
-        Distance matrix returned by CLEAN, with sequence IDs as columns and EC
-        labels as row indices.
-    gmm : str
-        Path to a Gaussian mixture model used to transform raw
-        distances into confidence scores.
-
-    Returns
-    -------
-    dict[str, list[str]]
-        Mapping from sequence ID to a list of EC predictions ordered from best
-        to worst.
-    """
-    gmm_lst = pickle.load(open(gmm, 'rb'))
-
-    max_sep_predictions = {}
-    for sequence_label in inference_df.columns:
-        smallest_10_dist_df = inference_df[sequence_label].nsmallest(10)
-        dist_lst = list(smallest_10_dist_df)
-        max_sep_i = maximum_separation(dist_lst, True, False)
-        ec = []
-        for i in range(max_sep_i+1):
-            EC_i = smallest_10_dist_df.index[i]
-            dist_i = smallest_10_dist_df[i]
-            dist_i = infer_confidence_gmm(dist_i, gmm_lst)
-            dist_str = "{:.4f}".format(dist_i)
-            ec.append('EC:' + str(EC_i) + '/' + dist_str)
-        max_sep_predictions[sequence_label] = ec
-    return max_sep_predictions
-
-
-def clean_max_sep_predictions(CLEAN_model, sequence_label_esm_emb_dict, emb_train, ec_id_dict_train, gmm_path, device):
-    """Run CLEAN inference on precomputed ESM embeddings.
-
-    This is the in-process inference step from CLEAN's original workflow. It
-    avoids recomputing ESM1b embeddings by taking your already-generated
-    embedding dictionary, projecting it through the CLEAN model, computing the
-    distance map against CLEAN's training embeddings, and converting that map
-    into final EC predictions with the maximum-separation heuristic.
-
-    Parameters
-    ----------
-    CLEAN_model : torch.nn.Module
-        Loaded CLEAN model used to transform ESM embeddings.
-    sequence_label_esm_emb_dict : dict[str, torch.Tensor]
-        Mapping from sequence ID to its precomputed ESM1b embedding.
-    emb_train : torch.Tensor
-        CLEAN's cached training embeddings (typically ``100.pt`` or similar).
-    ec_id_dict_train : dict
-        EC-to-index mapping loaded from the CLEAN training CSV.
-    gmm_path : str
-        Path to a Gaussian mixture model used to transform raw
-        distances into confidence scores.
-    device : str
-        Torch device to run the CLEAN model on.
-
-    Returns
-    -------
-    dict[str, list[str]]
-        Mapping from sequence ID to formatted EC predictions.
-    """
-    esm_emb_inference = torch.cat(
-        [torch.from_numpy(sequence_label_esm_emb_dict[label]).unsqueeze(0) if isinstance(sequence_label_esm_emb_dict[label], np.ndarray) else sequence_label_esm_emb_dict[label].unsqueeze(0)
-            for label in sequence_label_esm_emb_dict])
-    id_ec_inference_dummy = {seq_label:[] for seq_label in sequence_label_esm_emb_dict}
-    with torch.no_grad():
-        model_emb_inference  = CLEAN_model(esm_emb_inference.to(device)).to("cpu").clone()
-    inference_dist = get_dist_map_test(
-        emb_train, model_emb_inference, ec_id_dict_train, id_ec_inference_dummy, "cpu", torch.float32)
-    inference_df = pd.DataFrame.from_dict(inference_dist)
-    max_sep_predictions_dict = get_max_sep_predictions_dict(inference_df, gmm_path)
-    return max_sep_predictions_dict
-
-
-def run_clean_inference_with_embeddings(
-    sequence_label_esm_emb_dict: Dict[str, torch.Tensor],
-    emb_train_path: str,
-    ec_csv_path: str,
-    model_ckpt_path: str,
-    out_csv: Optional[str] = None,
-    device: str = "cpu",
-    gmm: str = "",
-) -> pd.DataFrame:
-    """Run CLEAN inference using already-computed ESM1b embeddings.
-
-    Parameters:
-    - sequence_label_esm_emb_dict: mapping from sequence id/label -> torch.Tensor (embedding)
-    - emb_train_path: path to `100.pt` (emb_train) used by CLEAN training
-    - ec_csv_path: path to the CSV used to build EC id mapping (train csv)
-    - model_ckpt_path: path to the CLEAN model checkpoint (.pth)
-    - out_csv: optional path to write predictions CSV
-    - device: 'cpu' or 'cuda'
-    - gmm: path to the Gaussian mixture model used to transform raw distances into confidence scores
-    Returns a pandas.DataFrame with columns `Seq_ID` and `Prediction`.
-    """
-    # build model and load checkpoint
-    CLEAN_model = LayerNormNet(512, 128, device, torch.float32)
-    checkpoint = torch.load(model_ckpt_path, map_location=device)
-    CLEAN_model.load_state_dict(checkpoint)
-    CLEAN_model.eval()
-
-    # load emb_train
-    emb_train = torch.load(emb_train_path, map_location="cpu")
-
-    # load ec id mapping
-    _, ec_id_dict_train = get_ec_id_dict(ec_csv_path)
-
-    # call CLEAN inference routine
-    preds = clean_max_sep_predictions(
-        CLEAN_model, sequence_label_esm_emb_dict, emb_train, ec_id_dict_train, gmm, device
-    )
-
-    # format to DataFrame: split each prediction "EC:3.1.1.1/0.0230" into
-    # CLEAN_EC_pred -> "3.1.1.1" and CLEAN_probability -> "0.0230".
-    records = []
-    for k, v in preds.items():
-        ecs = []
-        probs = []
-        for item in v:
-            parts = item.split("/", 1)
-            left = parts[0]
-            if left.startswith("EC:"):
-                left = left[len("EC:"):]
-            prob = parts[1] if len(parts) > 1 else ""
-            ecs.append(left)
-            probs.append(prob)
-        records.append({
-            "accession": k,
-            "CLEAN_EC_pred": "; ".join(ecs),
-            "CLEAN_probability": "; ".join(probs),
-        })
-
-    df = pd.DataFrame(records)
-
-    if out_csv:
-        os.makedirs(os.path.dirname(os.path.abspath(out_csv)), exist_ok=True)
-        df.to_csv(out_csv, index=False)
-
-    return df
-
-
-if __name__ == "__main__":
-    # mock some sequence embeddings for testing
-    import sys
-    from pathlib import Path
-
-    for candidate in [Path.cwd().resolve(), *Path.cwd().resolve().parents]:
-        if (candidate / "selectzyme").exists():
-            if str(candidate) not in sys.path:
-                sys.path.insert(0, str(candidate))
-            break
-
-    from selectzyme.backend.embed import gen_embedding
-    from selectzyme.backend.parsing import ParseLocalFiles
-
-    df = ParseLocalFiles("scripts/2_old+new.fasta").parse_fasta()
-    X = gen_embedding(sequences=df["sequence"].tolist(), plm_model="esm1b")
-    emb = {row["accession"]: X[i] for i, row in df.iterrows()}
-
-    # Example usage
-    pretraining_path = "../CLEAN/app/data/pretrained/"
-    ec_csv_path = "../CLEAN/app/data/"
-    desired_split = "100"
-    out_csv = "clean_predictions.csv"
-
-    df_clean = run_clean_inference_with_embeddings(
-        sequence_label_esm_emb_dict=emb,
-        emb_train_path=f'{pretraining_path}{desired_split}.pt',
-        ec_csv_path=f'{ec_csv_path}split{desired_split}.csv',
-        model_ckpt_path=f'{pretraining_path}split{desired_split}.pth',
-        out_csv=out_csv,
-        gmm=f'{pretraining_path}gmm_ensumble.pkl',
-        device="cpu",
-    )
-
+from types import SimpleNamespace
+
+
+def run_predictions(df, X, config, analysis_path):
+
+    # predict EC numbers with CLEAN if they dont exist yet
+    if df.CLEAN_EC_pred.isnull().all() and df.CLEAN_probability.isnull().all():
+        if config["project"]["plm"]["plm_model"] == "esm1b" and config["project"]["predictions"]["ec"] == True:
+            from selectzyme.backend.predict_ec import run_clean_inference_with_embeddings
+
+            data_path = "../CLEAN/app/data/"
+            desired_split = "100"
+            df_clean = run_clean_inference_with_embeddings(
+                sequence_label_esm_emb_dict={row["accession"]: X[i] for i, row in df.iterrows()},
+                emb_train_path=f'{data_path}pretrained/{desired_split}.pt',
+                ec_csv_path=f'{data_path}split{desired_split}.csv',
+                model_ckpt_path=f'{data_path}pretrained/split{desired_split}.pth',
+                out_csv=analysis_path + "/clean_predictions.csv",
+                gmm=f'{data_path}pretrained/gmm_ensumble.pkl',
+                device="cpu",
+            )
+
+            df = df.merge(df_clean[["accession", "CLEAN_EC_pred", "CLEAN_probability"]], on="accession", how="left")
+
+    # predict solubility and usability with NetSolP if they dont exist yet
+    if config["project"]["predictions"]["solubility"] == True and df.netSolP_solubility.isnull().all() and df.netSolP_usability.isnull().all():
+        from selectzyme.backend.predict_solubility import get_preds
+        logging.info("Running NetSolP predictions. This might take a long time.")
+
+        # hard coded configurations to pass to NetSolP package
+        args = SimpleNamespace(
+            MODEL_TYPE="ESM1b",
+            MODELS_PATH="NetSolP-1.0/PredictionServer/models",  # /scratch/global_1/fmoorhof/NetSolP/models/
+            NUM_THREADS=os.cpu_count(),
+            PREDICTION_TYPE="SU",  # Either Solubility(S), Usability(U) or Both
+        )
+
+        df = get_preds(df=df, args=args)
 
-    print(f"Wrote predictions to: {os.path.abspath(out_csv)}")
-    df = df.merge(df_clean[["accession", "CLEAN_EC_pred", "CLEAN_probability"]], on="accession", how="left")
+    return df