Process data into format usable by heatmaps

* gn3/heatmaps.py: implement `process_traits_data_for_heatmap` function, that
  will process the data into a form usable by heatmaps.

* tests/unit/test_heatmaps.py: check that the function processes the data into
  the correct form.

1 files changed, 12 insertions, 0 deletions

diff --git a/gn3/heatmaps.py b/gn3/heatmaps.py
index 991ddec..0c00d6c 100644
--- a/gn3/heatmaps.py
+++ b/gn3/heatmaps.py
@@ -277,6 +277,9 @@ def get_nearest_marker(traits_list, genotype):
     return [marker_finder(trait) for trait in traits_list]
 
 def get_lrs_from_chr(trait, chr_name):
+    """
+    Retrieve the LRS values for a specific chromosome in the given trait.
+    """
     chromosome = trait["chromosomes"].get(chr_name)
     if chromosome:
         return [
@@ -284,6 +287,15 @@ def get_lrs_from_chr(trait, chr_name):
             sorted(chromosome["loci"], key=lambda loc: loc["Locus"])]
     return [None]
 
+def process_traits_data_for_heatmap(data, trait_names, chromosome_names):
+    """
+    Process the traits data in a format useful for generating heatmap diagrams.
+    """
+    hdata = [
+        [get_lrs_from_chr(data[trait], chr_name) for trait in trait_names]
+        for chr_name in chromosome_names]
+    return hdata
+
 # # Grey + Blue + Red
 # def generate_heatmap():
 #     cols = 20