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.

2 files changed, 118 insertions, 1 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
diff --git a/tests/unit/test_heatmaps.py b/tests/unit/test_heatmaps.py
index cfdde1e..f3a81c5 100644
--- a/tests/unit/test_heatmaps.py
+++ b/tests/unit/test_heatmaps.py
@@ -5,7 +5,8 @@ from gn3.heatmaps import (
     get_lrs_from_chr,
     export_trait_data,
     compute_traits_order,
-    retrieve_strains_and_values)
+    retrieve_strains_and_values,
+    process_traits_data_for_heatmap)
 
 strainlist = ["B6cC3-1", "BXD1", "BXD12", "BXD16", "BXD19", "BXD2"]
 trait_data = {
@@ -199,3 +200,107 @@ class TestHeatmap(TestCase):
                  [13.2, 1.9, 2.22, 53.21]]]:
             with self.subTest(trait=trait, chromosome=chromosome):
                 self.assertEqual(get_lrs_from_chr(trait, chromosome), expected)
+
+    def test_process_traits_data_for_heatmap(self):
+        self.assertEqual(
+            process_traits_data_for_heatmap(
+                {"1": {
+                    "ID": "T1",
+                    "chromosomes": {
+                        1: {"Chr": 1,
+                            "loci": [
+                                {
+                                    "Locus": "rs31443144",  "cM": 1.500, "Mb": 3.010,
+                                    "LRS": 0.500, "Additive": -0.074, "pValue": 1.000
+                                },
+                                {
+                                    "Locus": "rs6269442", "cM": 1.500, "Mb": 3.492,
+                                    "LRS": 0.500, "Additive": -0.074, "pValue": 1.000
+                                },
+                                {
+                                    "Locus": "rs32285189", "cM": 1.630, "Mb": 3.511,
+                                    "LRS": 0.500, "Additive": -0.074, "pValue": 1.000
+                                },
+                                {
+                                    "Locus": "rs258367496", "cM": 1.630, "Mb": 3.660,
+                                    "LRS": 0.500, "Additive": -0.074, "pValue": 1.000
+                                },
+                                {
+                                    "Locus": "rs32430919", "cM": 1.750, "Mb": 3.777,
+                                    "LRS": 0.500, "Additive": -0.074, "pValue": 1.000
+                                },
+                                {
+                                    "Locus": "rs36251697", "cM": 1.880, "Mb": 3.812,
+                                    "LRS": 0.500, "Additive": -0.074, "pValue": 1.000
+                                },
+                                {
+                                    "Locus": "rs30658298", "cM": 2.010, "Mb": 4.431,
+                                    "LRS": 0.500, "Additive": -0.074, "pValue": 1.000
+                                }]},
+                        2: {"Chr": 2,
+                            "loci": [
+                                {
+                                    "Locus": "rs51852623", "cM": 2.010, "Mb": 4.447,
+                                    "LRS": 0.500, "Additive": -0.074, "pValue": 1.000
+                                },
+                                {
+                                    "Locus": "rs31879829", "cM": 2.140, "Mb": 4.519,
+                                    "LRS": 0.500, "Additive": -0.074, "pValue": 1.000
+                                },
+                                {
+                                    "Locus": "rs36742481", "cM": 2.140, "Mb": 4.776,
+                                    "LRS": 0.500, "Additive": -0.074, "pValue": 1.000
+                                }]}}},
+                 "2": {
+                     "ID": "T1",
+                     "chromosomes": {
+                         1: {"Chr": 1,
+                             "loci": [
+                                 {
+                                     "Locus": "rs31443144",  "cM": 1.500, "Mb": 3.010,
+                                     "LRS": 0.500, "Additive": -0.074, "pValue": 1.000
+                                 },
+                                 {
+                                     "Locus": "rs6269442", "cM": 1.500, "Mb": 3.492,
+                                     "LRS": 0.500, "Additive": -0.074, "pValue": 1.000
+                                 },
+                                 {
+                                     "Locus": "rs32285189", "cM": 1.630, "Mb": 3.511,
+                                     "LRS": 0.500, "Additive": -0.074, "pValue": 1.000
+                                 },
+                                 {
+                                     "Locus": "rs258367496", "cM": 1.630, "Mb": 3.660,
+                                     "LRS": 0.500, "Additive": -0.074, "pValue": 1.000
+                                 },
+                                 {
+                                     "Locus": "rs32430919", "cM": 1.750, "Mb": 3.777,
+                                     "LRS": 0.500, "Additive": -0.074, "pValue": 1.000
+                                 },
+                                 {
+                                     "Locus": "rs36251697", "cM": 1.880, "Mb": 3.812,
+                                     "LRS": 0.500, "Additive": -0.074, "pValue": 1.000
+                                 },
+                                 {
+                                     "Locus": "rs30658298", "cM": 2.010, "Mb": 4.431,
+                                     "LRS": 0.500, "Additive": -0.074, "pValue": 1.000
+                                 }]},
+                         2: {"Chr": 2,
+                             "loci": [
+                                 {
+                                     "Locus": "rs51852623", "cM": 2.010, "Mb": 4.447,
+                                     "LRS": 0.500, "Additive": -0.074, "pValue": 1.000
+                                 },
+                                 {
+                                     "Locus": "rs31879829", "cM": 2.140, "Mb": 4.519,
+                                     "LRS": 0.500, "Additive": -0.074, "pValue": 1.000
+                                 },
+                                 {
+                                     "Locus": "rs36742481", "cM": 2.140, "Mb": 4.776,
+                                     "LRS": 0.579, "Additive": -0.074, "pValue": 1.000
+                                 }]}}}},
+                ["2", "1"],
+                [1, 2]),
+            [[[0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5],
+              [0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5]],
+             [[0.5, 0.579, 0.5],
+              [0.5, 0.5, 0.5]]])