diff --git a/results/input_configs/test_config.yml b/results/input_configs/test_config.yml index d59baa8..ed4724d 100644 --- a/results/input_configs/test_config.yml +++ b/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] diff --git a/scripts/get_pdb_af_from_fasta.py b/scripts/get_pdb_af_from_fasta.py index 7bf4e9e..20c0841 100644 --- a/scripts/get_pdb_af_from_fasta.py +++ b/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) diff --git a/selectzyme/backend/predict.py b/selectzyme/backend/predict.py index b4e6476..a0195ef 100644 --- a/selectzyme/backend/predict.py +++ b/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 \ No newline at end of file diff --git a/selectzyme/backend/predict_ec.py b/selectzyme/backend/predict_ec.py index 650575d..0fe948f 100644 --- a/selectzyme/backend/predict_ec.py +++ b/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) diff --git a/selectzyme/backend/predict_solubility.py b/selectzyme/backend/predict_solubility.py index 2c4656a..f4d6ab6 100644 --- a/selectzyme/backend/predict_solubility.py +++ b/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) diff --git a/selectzyme/pages/callbacks.py b/selectzyme/pages/callbacks.py index 236df68..cff71cc 100644 --- a/selectzyme/pages/callbacks.py +++ b/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']}"