diverse feat, perf and fixes

results/input_configs/test_config.yml

@@ -19,7 +19,7 @@ project:
     length: "10 TO *"  # *=open end
     custom_data_location: "tests/head_10.tsv"  # use "" if no custom data is desired
     out_dir: "results/"
-    df_coi: ["accession", "reviewed", "ec", "organism_id", "length", "xref_brenda", "xref_pdb", "sequence"]  # long inline lists also possible or write like in 'query_terms'
+    df_coi: ["accession", "reviewed", "ec", "organism_id", "length", "xref_pdb", "xref_alphafolddb", "xref_brenda", "sequence"]  # long inline lists also possible or write like in 'query_terms'
 
   plm:
     plm_model: "esm1b" # [esm1b, esm2, prott5, prostt5]

scripts/get_pdb_af_from_fasta.py

@@ -1,52 +1,122 @@
-import requests, time, pathlib, re
+import pathlib
+import re
+import time
 
-# === Configuration ===
-FASTA_FILE = "tests/head_10.fasta"   # your input FASTA filename
-OUT_DIR = pathlib.Path("pdbs")   # output folder for PDB files
+import requests
+
+# ===========================
+# Configuration
+# ===========================
+FASTA_FILE = "tests/head_10.fasta"
+OUT_DIR = pathlib.Path("pdbs")
 OUT_DIR.mkdir(exist_ok=True)
-API_URL = "https://alphafold.ebi.ac.uk/api/prediction/{}"
+
+ALPHAFOLD_API = "https://alphafold.ebi.ac.uk/api/prediction/{}"
+ESM_API = "https://api.esmatlas.com/fetchPredictedStructure/{}"
 
 
+def extract_ids(fasta_path):
+    """
+    Extract identifiers from FASTA headers.
 
-def extract_uniprot_ids(fasta_path):
-    """Extract UniProt IDs from FASTA headers (>ID|...)."""
+    Expected formats:
+        >P12345|...
+        >A0A023GPI8|...
+        >MGYP002537940442|...
+    """
     ids = []
-    pattern = re.compile(r"^>([A-Z0-9]+)\|")
-    with open(fasta_path, "r") as f:
+
+    pattern = re.compile(r"^>([^|\s]+)")
+
+    with open(fasta_path) as f:
         for line in f:
             if line.startswith(">"):
                 m = pattern.match(line)
                 if m:
                     ids.append(m.group(1))
+
     return ids
 
-def download_pdb(uniprot_id):
-    """Download only the .pdb file for a given UniProt ID from AlphaFold DB."""
+
+def download_from_alphafold(uniprot_id):
+    """Download a PDB from AlphaFoldDB."""
+
+    pdb_path = OUT_DIR / f"{uniprot_id}.pdb"
+
+    if pdb_path.exists():
+        print(f"[SKIP] {pdb_path.name} already exists")
+        return
+
+    r = requests.get(ALPHAFOLD_API.format(uniprot_id), timeout=15)
+    r.raise_for_status()
+
+    data = r.json()
+
+    if not data or "pdbUrl" not in data[0]:
+        print(f"[WARN] {uniprot_id}: no AlphaFold model found")
+        return
+
+    pdb_url = data[0]["pdbUrl"]
+
+    with requests.get(pdb_url, stream=True, timeout=60) as s:
+        s.raise_for_status()
+
+        with open(pdb_path, "wb") as fh:
+            for chunk in s.iter_content(8192):
+                fh.write(chunk)
+
+    print(f"[AFDB] Saved {pdb_path.name}")
+
+
+def download_from_esm(mgyp_id):
+    """Download a PDB directly from the ESM Metagenomic Atlas."""
+
+    pdb_path = OUT_DIR / f"{mgyp_id}.pdb"
+
+    if pdb_path.exists():
+        print(f"[SKIP] {pdb_path.name} already exists")
+        return
+
+    url = ESM_API.format(mgyp_id)
+
+    r = requests.get(url, timeout=60)
+
+    if r.status_code == 404:
+        print(f"[WARN] {mgyp_id}: no ESM model found")
+        return
+
+    r.raise_for_status()
+
+    with open(pdb_path, "w") as fh:
+        fh.write(r.text)
+
+    print(f"[ESM ] Saved {pdb_path.name}")
+
+
+def download_pdb(identifier):
+    """
+    Automatically choose the correct database.
+
+    MGYP*  -> ESM Metagenomic Atlas
+    others -> AlphaFoldDB
+    """
+
     try:
-        r = requests.get(API_URL.format(uniprot_id), timeout=15)
-        r.raise_for_status()
-        data = r.json()
-        if not data or "pdbUrl" not in data[0]:
-            print(f"[WARN] {uniprot_id}: no PDB found")
-            return
-        pdb_url = data[0]["pdbUrl"]
-        pdb_path = OUT_DIR / f"{uniprot_id}.pdb"
-        if pdb_path.exists():
-            print(f"[SKIP] {pdb_path.name} already exists")
-            return
-        with requests.get(pdb_url, stream=True, timeout=60) as s:
-            s.raise_for_status()
-            with open(pdb_path, "wb") as fh:
-                for chunk in s.iter_content(8192):
-                    fh.write(chunk)
-        print(f"[OK] Saved {pdb_path.name}")
+        if identifier.upper().startswith("MGYP"):
+            download_from_esm(identifier)
+        else:
+            download_from_alphafold(identifier)
+
     except Exception as e:
-        print(f"[ERR] {uniprot_id}: {e}")
+        print(f"[ERR] {identifier}: {e}")
 
 
 if __name__ == "__main__":
-    ids = extract_uniprot_ids(FASTA_FILE)
-    print(f"Found {len(ids)} UniProt IDs.")
-    for uid in ids:
-        download_pdb(uid)
+
+    ids = extract_ids(FASTA_FILE)
+
+    print(f"Found {len(ids)} identifiers.")
+
+    for identifier in ids:
+        download_pdb(identifier)
         time.sleep(0.2)

selectzyme/backend/predict.py

@@ -4,11 +4,15 @@
 import os
 from types import SimpleNamespace
 
+import torch
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 
 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 "CLEAN_EC_pred" not in df.columns and "CLEAN_probability" not in df.columns:
         if config["project"]["plm"]["plm_model"] == "esm1b" and config["project"]["predictions"]["ec"] == True:
             from selectzyme.backend.predict_ec import run_clean_inference_with_embeddings
 
@@ -21,24 +25,30 @@ def run_predictions(df, X, config, analysis_path):
                 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",
+                device=device,
             )
 
             df = df.merge(df_clean[["accession", "CLEAN_EC_pred", "CLEAN_probability"]], on="accession", how="left")
 
+        else:
+            logging.info("CLEAN predictions are only available for ESM1b embeddings. Skipping prediction. Other reasons can be column names CLEAN_EC_pred and CLEAN_probability already exist in the dataframe or config['project']['predictions']['ec'] is set to False.")
+
     # 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():
+    if config["project"]["predictions"]["solubility"] == True and "netSolP_solubility" not in df.columns and "netSolP_usability" not in df.columns:
         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/
+            MODELS_PATH="/scratch/global_1/fmoorhof/NetSolP/models/",  # NetSolP-1.0/PredictionServer/models
             NUM_THREADS=os.cpu_count(),
             PREDICTION_TYPE="SU",  # Either Solubility(S), Usability(U) or Both
         )
 
         df = get_preds(df=df, args=args)
+
+    else:
+        logging.info("NetSolP predictions already exist (columns named netSolP_solubility and netSolP_usability). Skipping prediction.")
 
     return df

selectzyme/backend/predict_ec.py

@@ -93,9 +93,9 @@ def clean_max_sep_predictions(CLEAN_model, sequence_label_esm_emb_dict, emb_trai
             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()
+        model_emb_inference  = CLEAN_model(esm_emb_inference.to(device))
     inference_dist = get_dist_map_test(
-        emb_train, model_emb_inference, ec_id_dict_train, id_ec_inference_dummy, "cpu", torch.float32)
+        emb_train, model_emb_inference, ec_id_dict_train, id_ec_inference_dummy, device, 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
@@ -107,7 +107,7 @@ def run_clean_inference_with_embeddings(
     ec_csv_path: str,
     model_ckpt_path: str,
     out_csv: Optional[str] = None,
-    device: str = "cpu",
+    device = "cpu",
     gmm: str = "",
 ) -> pd.DataFrame:
     """Run CLEAN inference using already-computed ESM1b embeddings.
@@ -129,7 +129,7 @@ def run_clean_inference_with_embeddings(
     CLEAN_model.eval()
 
     # load emb_train
-    emb_train = torch.load(emb_train_path, map_location="cpu")
+    emb_train = torch.load(emb_train_path, map_location=device)
 
     # load ec id mapping
     _, ec_id_dict_train = get_ec_id_dict(ec_csv_path)

selectzyme/backend/predict_solubility.py

@@ -9,8 +9,11 @@
 from PredictionServer.data import BatchConverter, FastaBatchedDataset  # part of local netsolp install
 from PredictionServer.predict import get_preds_split, sigmoid
 
+from selectzyme.backend.utils import run_time
 
-def get_preds(df, args):
+
+@run_time
+def get_preds(df, args, pred_frequency=2):
     alphabet_path = os.path.join(args.MODELS_PATH, "ESM1b_alphabet.pkl")  # args.MODELS_PATH
 
     with open(alphabet_path, "rb") as f:
@@ -24,19 +27,19 @@ def get_preds(df, args):
 
     if "S" in args.PREDICTION_TYPE:
         preds_per_split = []
-        for i in range(5):
+        for i in range(pred_frequency):
             pred_df = get_preds_split(i, embed_dataloader, args, "Solubility", df)
             preds_i = sigmoid(np.stack(pred_df.preds.to_numpy()))
             preds_per_split.append(preds_i)
-        avg_pred = sum(preds_per_split) / 5
+        avg_pred = sum(preds_per_split) / pred_frequency
         df["netSolP_solubility"] = pd.Series(avg_pred)
     if "U" in args.PREDICTION_TYPE:
         preds_per_split = []
-        for i in range(5):
+        for i in range(pred_frequency):
             pred_df = get_preds_split(i, embed_dataloader, args, "Usability", df)
             preds_i = sigmoid(np.stack(pred_df.preds.to_numpy()))
             preds_per_split.append(preds_i)
-        avg_pred = sum(preds_per_split) / 5
+        avg_pred = sum(preds_per_split) / pred_frequency
         df["netSolP_usability"] = pd.Series(avg_pred)
 
     df.rename(columns={"sid": "accession", "fasta": "sequence"}, inplace=True)

selectzyme/pages/callbacks.py

@@ -137,6 +137,11 @@ def _process_selection(df, shared_table, point):
     accession = point["customdata"]
     selected_row = df[df["accession"] == accession].iloc[0]
     selected_row[df.columns.get_loc("selected")] = True
+    if "xref_alphafolddb" in df.columns and selected_row["xref_brenda"] != "unknown":
+        selected_row["xref_alphafolddb"] = (f"https://alphafold.com/entry/AF-{selected_row['xref_alphafolddb'].split(';')[0]}-F1")
+    else:
+        selected_row["Alphafold URL (if accession is from UniProt)"] = (f"https://alphafold.ebi.ac.uk/entry/{selected_row['accession']}")
+
     if all(col in df.columns for col in {"xref_brenda", "accession", "organism_id"}) and selected_row["xref_brenda"] != "unknown":
         selected_row["BRENDA URL"] = (
             f"https://www.brenda-enzymes.org/enzyme.php?ecno={selected_row['xref_brenda'].split(';')[0]}&UniProtAcc={selected_row['accession']}&OrganismID={selected_row['organism_id']}"