From 11632a565a6f901eca852a5a40a6f9fd3170152a Mon Sep 17 00:00:00 2001
From: Frederick Muriuki Muriithi
Date: Wed, 15 Sep 2021 12:08:56 +0300
Subject: 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.
---
 gn3/heatmaps.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

(limited to 'gn3')

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
-- 
cgit 1.4.1