Merge pull request #6 from Multiomics-Analytics-Group/new_data

Inference mode for scripts, add argument parsing to scripts

Author: marcoreverenna

notebooks/heatmaps_gridsearch.ipynb

@@ -18,8 +18,8 @@
     }
    ],
    "source": [
-    "r\"\"\" Heatmaps notebook.\n",
-    " _____  _______  _    _ \n",
+    "r\"\"\"Heatmaps notebook.\n",
+    " _____  _______  _    _\n",
     "|  __ \\|__   __|| |  | |\n",
     "| |  | |  | |   | |  | |\n",
     "| |  | |  | |   | |  | |\n",

notebooks/prot_optimization_dbg.ipynb

@@ -7,8 +7,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "r\"\"\" Protease optimization\n",
-    " _____  _______  _    _ \n",
+    "r\"\"\"Protease optimization\n",
+    " _____  _______  _    _\n",
     "|  __ \\|__   __|| |  | |\n",
     "| |  | |  | |   | |  | |\n",
     "| |  | |  | |   | |  | |\n",

notebooks/prot_optimization_greedy.ipynb

@@ -7,8 +7,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "r\"\"\" Protease optimization\n",
-    " _____  _______  _    _ \n",
+    "r\"\"\"Protease optimization\n",
+    " _____  _______  _    _\n",
     "|  __ \\|__   __|| |  | |\n",
     "| |  | |  | |   | |  | |\n",
     "| |  | |  | |   | |  | |\n",

src/preprocessing.py

@@ -84,6 +84,7 @@ def normalize_sequence(sequence):
 def remove_modifications(psm_column):
     """
     Remove any content within parentheses, including the parentheses, from a given string.
+    Remove UNIMOD modifications and normalize I to L.
 
     Parameters:
     - psm_column (str): The string containing modifications in parentheses (e.g., "A(ox)BC(mod)D"). If the value is null, it returns None.
@@ -93,12 +94,31 @@ def remove_modifications(psm_column):
     """
 
     if pd.notnull(psm_column):
-        return re.sub(
+        ret = re.sub(
             r"\(.*?\)", "", psm_column
         )  # Replace any content in parentheses with an empty string
+        ret = re.sub(
+            r"\[.*?\]", "", ret
+        )  # replace UNIMOD modifications in square brackets
+        ret = normalize_sequence(ret)
+        return ret
     return None
 
 
+# ! needs to move once it is a package
+def test_remove_modifications():
+    assert remove_modifications("A(ox)BC(mod)D") == "ABCD"
+    assert remove_modifications("A[UNIMOD:21]BC[UNIMOD:35]D") == "ABCD"
+    assert remove_modifications("A(ox)[UNIMOD:21]BC(mod)[UNIMOD:35]D") == "ABCD"
+    assert remove_modifications(None) is None
+    assert remove_modifications("ACD") == "ACD"
+    assert remove_modifications("A(I)BCD") == "ABCD"
+    assert remove_modifications("A(ox)B(I)C(mod)D") == "ABCD"
+    assert remove_modifications("A(ox)[UNIMOD:21]B(I)C(mod)[UNIMOD:35]D") == "ABCD"
+    assert remove_modifications("AI BCD") == "AL BCD"
+    assert remove_modifications("A(ox)I B(mod)CD") == "AL BCD"
+
+
 def clean_dataframe(df):
     """
     Clean and preprocess a DataFrame for analysis by removing '(ox)' substrings from sequences in the 'seq' column.

src/script_dbg.py

@@ -20,6 +20,7 @@
 # !pip install kaleido # to export plotly figures as png
 # !pip install --upgrade nbformat # to avoid plotly error
 
+import argparse
 import json
 import logging
 import os
@@ -29,7 +30,6 @@
 import pandas as pd
 
 # import libraries
-
 import alignment as align
 import clustering as clus
 import compute_statistics as comp_stat
@@ -40,16 +40,25 @@
 import mapping as map
 import preprocessing as prep
 
+repo_folder = Path(__file__).resolve().parents[1]
+
+parser = argparse.ArgumentParser(description="Protein Assembly Script")
+parser.add_argument("--input_csv", type=str, help="Input file")
+parser.add_argument(
+    "--folder_outputs", default="outputs", type=str, help="Outputs folder"
+)
+parser.add_argument("--training", action="store_true", help="Training mode")
+
 logging.basicConfig(
     level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s"
 )
 logger = logging.getLogger(__name__)
 
 BASE_DIR = Path(__file__).resolve().parents[1]
 JSON_DIR = BASE_DIR / "json"
-INPUT_DIR = BASE_DIR / "inputs"
-FASTA_DIR = BASE_DIR / "fasta"
-OUTPUTS_DIR = BASE_DIR / "outputs"
+# INPUT_DIR = BASE_DIR / "inputs"
+# FASTA_DIR = BASE_DIR / "fasta"
+# OUTPUTS_DIR = BASE_DIR / "outputs"
 
 
 def get_sample_metadata(run, chain="", json_path=JSON_DIR / "sample_metadata.json"):
@@ -68,17 +77,19 @@ def get_sample_metadata(run, chain="", json_path=JSON_DIR / "sample_metadata.jso
     raise ValueError(f"No metadata found for run '{run}' with chain '{chain}'.")
 
 
-def main():
+def main(input_csv: str, folder_outputs: str = "outputs", training: bool = False):
     """Main function to run the assembly script."""
 
-    logger.info("Starting protein assembly pipeline.")
+    input_csv = Path(input_csv)
 
-    run = "ma3"
+    logger.info("Starting protein assembly pipeline.")
 
-    meta = get_sample_metadata(run, chain="light")
-    protein = meta["protein"]
-    chain = meta["chain"]
-    proteases = meta["proteases"]
+    run = input_csv.stem
+    if training:
+        meta = get_sample_metadata(run, chain="light")
+        protein = meta["protein"]
+        # chain = meta["chain"]
+        proteases = meta["proteases"]
 
     ass_method = "dbg"
 
@@ -91,34 +102,37 @@ def main():
 
     logger.info("Parameters loaded.")
 
-    folder_outputs = f"../outputs/{run}{chain}"
-    prep.create_directory(folder_outputs)
-    combination_folder_out = os.path.join(
-        folder_outputs,
-        f"comb_{ass_method}_c{conf}_ks{kmer_size}_ts{size_threshold}_mo{min_overlap}_mi{min_identity}_mm{max_mismatches}",
+    folder_outputs = Path(folder_outputs) / run
+    folder_outputs.mkdir(parents=True, exist_ok=True)
+
+    combination_folder_out = (
+        folder_outputs
+        / f"comb_{ass_method}_c{conf}_ks{kmer_size}_ts{size_threshold}_mo{min_overlap}_mi{min_identity}_mm{max_mismatches}"
     )
     prep.create_subdirectories_outputs(combination_folder_out)
 
     logger.info(f"Output folders created at: {combination_folder_out}")
 
     # Data cleaning
     logger.info("Starting data cleaning...")
-
-    protein_norm = prep.normalize_sequence(protein)
-    df = pd.read_csv(f"../inputs/{run}.csv")
-    df["protease"] = df["experiment_name"].apply(
-        lambda name: prep.extract_protease(name, proteases)
-    )
+    if training:
+        protein_norm = prep.normalize_sequence(protein)
+    df = pd.read_csv(input_csv)
+    if training:
+        df["protease"] = df["experiment_name"].apply(
+            lambda name: prep.extract_protease(name, proteases)
+        )
     df = prep.clean_dataframe(df)
     df["cleaned_preds"] = df["preds"].apply(prep.remove_modifications)
     cleaned_psms = df["cleaned_preds"].tolist()
     filtered_psms = prep.filter_contaminants(
-        cleaned_psms, run, "../fasta/contaminants.fasta"
+        cleaned_psms, run, repo_folder / "fasta/contaminants.fasta"
     )
     df = df[df["cleaned_preds"].isin(filtered_psms)]
-    df["mapped"] = df["cleaned_preds"].apply(
-        lambda x: "True" if x in protein_norm else "False"
-    )
+    if training:
+        df["mapped"] = df["cleaned_preds"].apply(
+            lambda x: "True" if x in protein_norm else "False"
+        )
     df = df[df["conf"] > conf]
     df.reset_index(drop=True, inplace=True)
     final_psms = df["cleaned_preds"].tolist()
@@ -145,16 +159,17 @@ def main():
         f"{combination_folder_out}/contigs/{ass_method}_contig_{conf}_{run}.fasta",
         "fasta",
     )
-    mapped_contigs = map.process_protein_contigs_scaffold(
-        assembled_contigs, protein_norm, max_mismatches, min_identity
-    )
-    df_contigs = map.create_dataframe_from_mapped_sequences(data=mapped_contigs)
-    comp_stat.compute_assembly_statistics(
-        df=df_contigs,
-        sequence_type="contigs",
-        output_folder=f"{combination_folder_out}/statistics",
-        reference=protein_norm,
-    )
+    if training:
+        mapped_contigs = map.process_protein_contigs_scaffold(
+            assembled_contigs, protein_norm, max_mismatches, min_identity
+        )
+        df_contigs = map.create_dataframe_from_mapped_sequences(data=mapped_contigs)
+        comp_stat.compute_assembly_statistics(
+            df=df_contigs,
+            sequence_type="contigs",
+            output_folder=f"{combination_folder_out}/statistics",
+            reference=protein_norm,
+        )
     assembled_scaffolds = dbg.create_scaffolds(assembled_contigs, min_overlap)
     assembled_scaffolds = list(set(assembled_scaffolds))
     assembled_scaffolds = sorted(assembled_scaffolds, key=len, reverse=True)
@@ -180,22 +195,23 @@ def main():
         f"{combination_folder_out}/scaffolds/{ass_method}_scaffold_{conf}_{run}.fasta",
         "fasta",
     )
-    mapped_scaffolds = map.process_protein_contigs_scaffold(
-        assembled_contigs=assembled_scaffolds,
-        target_protein=protein_norm,
-        max_mismatches=max_mismatches,
-        min_identity=min_identity,
-    )
+    if training:
+        mapped_scaffolds = map.process_protein_contigs_scaffold(
+            assembled_contigs=assembled_scaffolds,
+            target_protein=protein_norm,
+            max_mismatches=max_mismatches,
+            min_identity=min_identity,
+        )
 
-    df_scaffolds_mapped = map.create_dataframe_from_mapped_sequences(
-        data=mapped_scaffolds
-    )
-    comp_stat.compute_assembly_statistics(
-        df=df_scaffolds_mapped,
-        sequence_type="scaffolds",
-        output_folder=f"{combination_folder_out}/statistics",
-        reference=protein_norm,
-    )
+        df_scaffolds_mapped = map.create_dataframe_from_mapped_sequences(
+            data=mapped_scaffolds
+        )
+        comp_stat.compute_assembly_statistics(
+            df=df_scaffolds_mapped,
+            sequence_type="scaffolds",
+            output_folder=f"{combination_folder_out}/statistics",
+            reference=protein_norm,
+        )
 
     # Clustering
     scaffolds_folder_out = f"{combination_folder_out}/scaffolds"
@@ -241,4 +257,9 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    args = parser.parse_args()
+    main(
+        input_csv=args.input_csv,
+        folder_outputs=args.folder_outputs,
+        training=args.training,
+    )