1 files changed, 19 insertions, 0 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()))