fix for #366

jodyphelan · Jun 21, 2024 · 4c50d84 · abhi18av · Jun 26, 2024 · 4c50d84
1 parent 3428202
commit 4c50d84
Showing 1 changed file with 26 additions and 2 deletions.
diff --git a/tbprofiler/collate.py b/tbprofiler/collate.py
@@ -169,8 +169,32 @@ def generate_itol_config(rows: List[dict], drugs: list, prefix: str) -> None:
     prefix : str
         Prefix for output files
     """
-    all_lineage_cols = {"lineage1":"#104577","lineage2":"#ab2323","lineage3":"#18a68c","lineage4":"#f68e51","lineage5":"#7cb5d2","lineage6":"#fde05e","lineage7":"#bc94b7","lineage8":"#ccc9e7","lineage9":"#bd9391","Animal strains":"#f8e0c8","Other":"#000000","Not called": "#ffffff"}
-    lineage_aggregation = {"": "Not called","M.caprae":"Animal strains","M.bovis":"Animal strains","M.orygis":"Animal strains"}
+    all_lineage_cols = {
+        "lineage1":"#104577",
+        "lineage2":"#ab2323",
+        "lineage3":"#18a68c",
+        "lineage4":"#f68e51",
+        "lineage5":"#7cb5d2",
+        "lineage6":"#fde05e",
+        "lineage7":"#bc94b7",
+        "lineage8":"#ccc9e7",
+        "lineage9":"#bd9391",
+        "Animal strains":"#f8e0c8",
+        "Other":"#000000",
+        "Not called": "#ffffff"
+    }
+
+    lineage_aggregation = {
+        "": "Not called",
+        "La1": "Animal strains",
+        "La2": "Animal strains",
+        "La3": "Animal strains",
+        "M.canetti":"Other",
+        "M.caprae":"Animal strains",
+        "M.bovis":"Animal strains",
+        "M.orygis":"Animal strains"
+    }
+
     lineage_dict = {r['sample']:lineage_aggregation.get(r["main_lineage"],r["main_lineage"]) if ";" not in r["main_lineage"] else "Other" for r in rows}
     lineages_present = set(lineage_dict.values())
     lineage_cols = {key:val for key,val in all_lineage_cols.items() if key in lineages_present}