diff --git a/small_RNA-seq/small_RNA-seq.snakefile b/small_RNA-seq/small_RNA-seq.snakefile
index 691ba509b003373298bc901830fdc317e903f5f7..72b4555759e27ab2b609ccbfa79f0f85164f7dd3 100644
--- a/small_RNA-seq/small_RNA-seq.snakefile
+++ b/small_RNA-seq/small_RNA-seq.snakefile
@@ -506,7 +506,7 @@ read_graphs = [
 if contrasts_dict["ip"]:
     fold_heatmaps = expand(
         OPJ(output_dir, "figures", "fold_heatmaps", "{small_type}_{fold_type}_heatmap.{fig_format}"),
-        small_type=["pisimi"], fold_type=["mean_log2_RPM_fold", "log2FoldChange"], fig_format=FIG_FORMATS)
+        small_type=["pisimi", "prot_si"], fold_type=["mean_log2_RPM_fold", "log2FoldChange"], fig_format=FIG_FORMATS)
     ip_fold_boxplots = expand(
         OPJ(output_dir, "figures", "all_{contrast_type}",
             "{contrast_type}_{small_type}_{fold_type}_{gene_list}_boxplots.{fig_format}"),
@@ -515,7 +515,7 @@ if contrasts_dict["ip"]:
 else:
     fold_heatmaps = expand(
         OPJ(output_dir, "figures", "fold_heatmaps", "{small_type}_{fold_type}_heatmap.{fig_format}"),
-        small_type=["pisimi"], fold_type=["log2FoldChange"], fig_format=FIG_FORMATS)
+        small_type=["pisimi", "prot_si"], fold_type=["log2FoldChange"], fig_format=FIG_FORMATS)
     ip_fold_boxplots = []
 
 exploratory_graphs = [
@@ -1699,6 +1699,9 @@ def plot_fold_heatmap(data, gene_colours=None, labels_dict=None):
     g.cax.yaxis.set_label_position("left")
 
 
+#class DefaultCounter(defaultdict, Counter):
+#    pass
+
 rule make_fold_heatmap:
     """With a full list of genes, this rule can use a lot of memory."""
     input:
@@ -1716,10 +1719,18 @@ rule make_fold_heatmap:
         gene_colours.name = f"{wildcards.small_type}_{wildcards.fold_type}"
         #all_folds.drop(["cosmid", "name", "small_type"], axis=1)
         #all_folds[all_folds.columns.difference(["cosmid", "name", "small_type"])]
-        save_plot(
-            output.fold_heatmap, plot_fold_heatmap, all_folds.drop(["cosmid", "name", "small_type"], axis=1),
-            gene_colours, merge_with(tuple, small_type2colour, Counter(all_folds.small_type)),
-            tight=False)
+        try:
+            #labels_dict = merge_with(tuple, small_type2colour, DefaultCounter(all_folds.small_type))
+            # https://stackoverflow.com/a/47396625/1878788
+            small_type_counts = Counter(small_type2colour.keys())
+            small_type_counts.update(all_folds.small_type)
+            labels_dict = {small_type : (small_type2colour[small_type], count - 1) for (small_type, count) in small_type_counts.items()}
+            save_plot(
+                output.fold_heatmap, plot_fold_heatmap, all_folds.drop(["cosmid", "name", "small_type"], axis=1),
+                gene_colours, labels_dict, tight=False)
+        except ValueError as e:
+            print(labels_dict)
+            raise
 
 
 #def plot_norm_counts(counts_data, summaries):
@@ -2042,7 +2053,7 @@ rule merge_bigwig_reps:
                 for bw in bws:
                     bw.close()
 
-# TODO: use log2
+
 rule make_bigwig_ratios:
     """This rule tries to make a bigwig file displaying the ratio between a condition and a reference, from the corresponding two normalized bigwig files."""
     input: