Retrieve the strains with valid values

Issue:
https://github.com/genenetwork/gn-gemtext-threads/blob/main/topics/gn1-migration-to-gn2/clustering.gmi

* gn3/computations/heatmap.py: add function to get strains with values
* tests/unit/computations/test_heatmap.py: new tests

  Add function to get the strains whose values are not `None` from the
  `trait_data` object passed in.

  This migrates
  https://github.com/genenetwork/genenetwork1/blob/master/web/webqtl/heatmap/Heatmap.py#L215-221
  into a separate function that can handle that and be tested independently of
  any other code.

2 files changed, 32 insertions, 1 deletions

diff --git a/gn3/computations/heatmap.py b/gn3/computations/heatmap.py
index 1c86261..5a3c619 100644
--- a/gn3/computations/heatmap.py
+++ b/gn3/computations/heatmap.py
@@ -203,3 +203,22 @@ def compute_heatmap_order(
         return __order_maker(__order_maker(norder, slnk_dt[0]), slnk_dt[1])
 
     return __order_maker(neworder, slink_data)
+
+def retrieve_strains_and_values(strainlist, trait_data):
+    """
+    Get the strains and their corresponding values from `strainlist` and
+    `trait_data`.
+
+    This migrates the code in
+    https://github.com/genenetwork/genenetwork1/blob/master/web/webqtl/heatmap/Heatmap.py#L215-221
+    """
+    def __strains_and_values(acc, i):
+        if trait_data[i] is None:
+            return acc
+        if len(acc) == 0:
+            return ((strainlist[i], ), (trait_data[i], ))
+        _strains = acc[0]
+        _vals = acc[1]
+        return (_strains + (strainlist[i], ), _vals + (trait_data[i], ))
+    return reduce(
+        __strains_and_values, range(len(strainlist)), (tuple(), tuple()))
diff --git a/tests/unit/computations/test_heatmap.py b/tests/unit/computations/test_heatmap.py
index 14807bb..686288d 100644
--- a/tests/unit/computations/test_heatmap.py
+++ b/tests/unit/computations/test_heatmap.py
@@ -3,7 +3,8 @@ from unittest import TestCase
 from gn3.computations.heatmap import (
     cluster_traits,
     export_trait_data,
-    compute_heatmap_order)
+    compute_heatmap_order,
+    retrieve_strains_and_values)
 
 strainlist = ["B6cC3-1", "BXD1", "BXD12", "BXD16", "BXD19", "BXD2"]
 trait_data = {
@@ -164,3 +165,14 @@ class TestHeatmap(TestCase):
             with self.subTest(xoffset=xoff):
                 self.assertEqual(
                     compute_heatmap_order(slinked, xoffset=xoff), expected)
+
+    def test_retrieve_strains_and_values(self):
+        """Test retrieval of strains and values."""
+        for slist, tdata, expected in [
+                [["s1", "s2", "s3", "s4"], [9, None, 5, 4],
+                 (("s1", "s3", "s4"), (9, 5, 4))],
+                [["s1", "s2", "s3", "s4", "s5"], [6, None, None, 4, None],
+                 (("s1", "s4"), (6, 4))]]:
+            with self.subTest(strainlist=slist, traitdata=tdata):
+                self.assertEqual(
+                    retrieve_strains_and_values(slist, tdata), expected)