From 9380130c830e7ba26b8910ea674e46cd3c4ffad7 Mon Sep 17 00:00:00 2001
From: lmgueguen <lm.gueguen@orange.fr>
Date: Tue, 25 Aug 2026 16:21:27 -0400
Subject: [PATCH] Updated the heatmap object and the parser to allow to include
 loose hits with the color black

---
 app/Heatmap.py  | 18 ++++++++++--------
 app/MainBase.py |  4 +++-
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/app/Heatmap.py b/app/Heatmap.py
index 82d1cfc..e27dc73 100644
--- a/app/Heatmap.py
+++ b/app/Heatmap.py
@@ -14,10 +14,11 @@ class Heatmap(object):
     This is a program that genreates a heatmap of multiple RGI analyses.
     """
 
-    def __init__(self, input, classification, frequency, output, cluster, display, debug):
+    def __init__(self, input, classification, frequency, include_loose,output, cluster, display, debug):
         self.input = input
         self.classification = classification
         self.frequency = frequency
+        self.include_loose = include_loose
         self.output = output
         self.cluster = cluster
         self.display = display
@@ -332,7 +333,7 @@ def run(self):
                         hsp = max(value.keys(), key=(lambda key: value[key]['bit_score']))
 
                         # Flag to exclude loose hits
-                        if value[hsp]["type_match"] != "Loose":
+                        if value[hsp]["type_match"] != "Loose" or self.include_loose:
                             topmodel = value[hsp]["model_name"]
                             tophits[topmodel] = value[hsp]["type_match"]
 
@@ -409,7 +410,7 @@ def run(self):
         genelist = sorted(genelist)
 
         # Create a dictionary that will convert type of hit to num. value
-        conversion = {"Perfect": 2, "Strict": 1}
+        conversion = {"Perfect": 3, "Strict": 2, "Loose": 1}
 
         # Apply conversion so hit criteria is number based
         for sample in genes:
@@ -428,7 +429,7 @@ def run(self):
 
         # Fixed colourmap values (purple, teal, yellow)
         cmap_values = [0, 1, 2, 3]
-        custom_cmap = matplotlib.colors.ListedColormap(['#4c0057', '#00948f', '#feed00'])
+        custom_cmap = matplotlib.colors.ListedColormap(['#000000', '#4c0057', '#00948f', '#feed00'])
         norm = matplotlib.colors.BoundaryNorm(cmap_values, custom_cmap.N)
 
         # If the classification option chosen:
@@ -848,15 +849,16 @@ def run(self):
                     print('Output file %s: AMR genes are listed in alphabetical order '
                     'and samples have been clustered hierarchically (see SciPy documentation). '
                     'Yellow represents a perfect hit, teal represents a strict hit, purple '
-                    'represents no hit.' %(file_name))
+                    'represents a loose hit, black represents no hit.' %(file_name))
                 elif self.cluster == 'genes':
                     print('Output file %s: AMR genes have been clustered hierarchically. '
                     'Yellow represents a perfect hit, teal represents a strict hit, purple '
-                    'represents no hit.' %(file_name))
+                    'represents a loose hit, black represents no hit.' %(file_name))
                 elif self.cluster == 'both':
                     print('Output file %s: AMR genes and samples have been clustered hierarchically '
                     '(see SciPy documentation). Yellow represents a perfect hit, teal represents a strict hit, purple '
-                    'represents no hit.' %(file_name))
+                    'represents a loose hit, black represents no hit.' %(file_name))
                 else:
                     print('Output file %s: Yellow represents a perfect hit, '
-                    'teal represents a strict hit, purple represents no hit.' %(file_name))
+                    'teal represents a strict hit, purple represents a loose hit,'
+                    ' black represents  no hit.' %(file_name))
diff --git a/app/MainBase.py b/app/MainBase.py
index 93a7131..8ef0e08 100644
--- a/app/MainBase.py
+++ b/app/MainBase.py
@@ -388,6 +388,8 @@ def heatmap_args(self):
                             help="The option to organize resistance genes based on a category.")
         parser.add_argument('-f', '--frequency', dest="frequency", action="store_true",
                             help="Represent samples based on resistance profile.")
+        parser.add_argument('-l', '--include_loose', dest="include_loose", action="store_true",
+                            help="Include loose predictions.")
         parser.add_argument('-o', '--output', dest="output", default="RGI_heatmap",
                             help="Name for the output EPS and PNG files.\nThe number of files run will automatically \nbe appended to the end of the file name.(default={})".format('RGI_heatmap'))
         parser.add_argument('-clus', '--cluster', dest="cluster", choices=("samples", "genes", "both"),
@@ -400,7 +402,7 @@ def heatmap_args(self):
         return parser
 
     def heatmap_run(self, args):
-        obj = Heatmap(args.input, args.classification, args.frequency,
+        obj = Heatmap(args.input, args.classification, args.frequency, args.include_loose,
                       args.output, args.cluster, args.display, args.debug)
         obj.run()