diff --git a/keggcharter.py b/keggcharter.py
index 205af17..6277b4f 100644
--- a/keggcharter.py
+++ b/keggcharter.py
@@ -18,7 +18,7 @@
 
 from keggpathway_map import KEGGPathwayMap, expand_by_list_column
 
-__version__ = "0.4.0"
+__version__ = "0.4.1"
 
 
 def get_arguments():
@@ -438,7 +438,7 @@ def main():
     if args.transcriptomic_columns:
         args.transcriptomic_columns = args.transcriptomic_columns.split(',')
 
-    ko_column = 'KO (KEGGCharter)'  # TODO - set ko_column to user defined value
+    ko_column = 'KO (KEGGCharter)' if not hasattr(args, 'ko_column') or not args.ko_column else args.ko_column
 
     if args.resume:
         data = pd.read_csv(f'{args.output}/data_for_charting.tsv', sep='\t')
diff --git a/keggpathway_map.py b/keggpathway_map.py
index 8e92f1b..f60cf54 100644
--- a/keggpathway_map.py
+++ b/keggpathway_map.py
@@ -355,7 +355,7 @@ def most_abundant_taxa(self, data, columns, taxa_column, number_of_taxa=10):
         return data.index.tolist()[:number_of_taxa]
 
     def genomic_potential_taxa(
-            self, data, samples, ko_column, taxon_to_mmap_to_orthologs, mmaps2taxa=None,
+            self, data, samples, ko_column, taxon_to_mmap_to_orthologs, mmaps2taxa,
             taxa_column='Taxonomic lineage (GENUS)', output_basename=None, number_of_taxa=10, grey_taxa='Other taxa'):
         """
         Represents the genomic potential of the dataset for a certain taxa level,
@@ -408,12 +408,12 @@ def genomic_potential_taxa(
                             box2taxon[box] = [grey_taxa]
         else:
             # if input_taxonomy
-            dic_colors = {grey_taxa : "#7c7272"}
+            dic_colors = {grey_taxa: "#7c7272"}
             df = data[data.any(axis=1)]
             for ortholog in df[ko_column]:
                 if ortholog in self.ko_boxes.keys():
                     for box in self.ko_boxes[ortholog]:
-                        if box not in box2taxon.keys():
+                        if box in box2taxon.keys():
                             box2taxon[box].append(grey_taxa)
                         else:
                             box2taxon[box] = [grey_taxa]