From 4c50d84efd00d0b0a0dc3b385c81f487efb336fe Mon Sep 17 00:00:00 2001
From: Jody Phelan <jody.phelan@lshtm.ac.uk>
Date: Fri, 21 Jun 2024 16:39:18 +0200
Subject: [PATCH] fix for #366

---
 tbprofiler/collate.py | 28 ++++++++++++++++++++++++++--
 1 file changed, 26 insertions(+), 2 deletions(-)

diff --git a/tbprofiler/collate.py b/tbprofiler/collate.py
index 85de3d0..ce890dc 100644
--- 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}