From cad4649d19001f62ef592dedf09f3ac53744962a Mon Sep 17 00:00:00 2001
From: Frederick Muriuki Muriithi
Date: Thu, 21 Oct 2021 09:00:16 +0300
Subject: Implement `find_identical_traits` function

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

* gn3/partial_correlations.py: implement function `find_identical_traits`
* tests/unit/test_partial_correlations.py: implement tests for function
  `find_identical_traits`

  Migrate `web.webqtl.correlation.correlationFunction.findIdenticalTraits`
  function in GN1 to here, adding in tests to ensure the migration works in a
  bug-compatible version with the original.
---
 gn3/partial_correlations.py             | 38 ++++++++++++++++++++++++++++++++-
 tests/unit/test_partial_correlations.py | 33 +++++++++++++++++++++++++++-
 2 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/gn3/partial_correlations.py b/gn3/partial_correlations.py
index c556d10..1fb0ccc 100644
--- a/gn3/partial_correlations.py
+++ b/gn3/partial_correlations.py
@@ -6,7 +6,7 @@ GeneNetwork1.
 """
 
 from functools import reduce
-from typing import Any, Sequence
+from typing import Any, Tuple, Sequence
 
 def control_samples(controls: Sequence[dict], sampleslist: Sequence[str]):
     """
@@ -86,3 +86,39 @@ def fix_samples(primary_trait: dict, control_traits: Sequence[dict]) -> Sequence
         control_vals_vars[0],
         tuple(primary_trait[sample]["variance"] for sample in primary_samples),
         control_vals_vars[1])
+
+def find_identical_traits(
+        primary_name: str, primary_value: float, control_names: Tuple[str, ...],
+        control_values: Tuple[float, ...]) -> Tuple[str, ...]:
+    """
+    Find traits that have the same value when the values are considered to
+    3 decimal places.
+
+    This is a migration of the
+    `web.webqtl.correlation.correlationFunction.findIdenticalTraits` function in
+    GN1.
+    """
+    def __merge_identicals__(
+            acc: Tuple[str, ...],
+            ident: Tuple[str, Tuple[str, ...]]) -> Tuple[str, ...]:
+        return acc + ident[1]
+
+    def __dictify_controls__(acc, control_item):
+        ckey = "{:.3f}".format(control_item[0])
+        return {**acc, ckey: acc.get(ckey, tuple()) + (control_item[1],)}
+
+    return (reduce(## for identical control traits
+        __merge_identicals__,
+        (item for item in reduce(# type: ignore[var-annotated]
+            __dictify_controls__, zip(control_values, control_names),
+            {}).items() if len(item[1]) > 1),
+        tuple())
+            or
+            reduce(## If no identical control traits, try primary and controls
+                __merge_identicals__,
+                (item for item in reduce(# type: ignore[var-annotated]
+                    __dictify_controls__,
+                    zip((primary_value,) + control_values,
+                        (primary_name,) + control_names), {}).items()
+                 if len(item[1]) > 1),
+                tuple()))
diff --git a/tests/unit/test_partial_correlations.py b/tests/unit/test_partial_correlations.py
index c591c8f..60e54c1 100644
--- a/tests/unit/test_partial_correlations.py
+++ b/tests/unit/test_partial_correlations.py
@@ -4,7 +4,8 @@ from unittest import TestCase
 from gn3.partial_correlations import (
     fix_samples,
     control_samples,
-    dictify_by_samples)
+    dictify_by_samples,
+    find_identical_traits)
 
 sampleslist = ["B6cC3-1", "BXD1", "BXD12", "BXD16", "BXD19", "BXD2"]
 control_traits = (
@@ -178,3 +179,33 @@ class TestPartialCorrelations(TestCase):
              (None,),
              (None, None, None, None, None, None, None, None, None, None, None,
               None, None)))
+
+    def test_find_identical_traits(self):
+        """
+        Test `gn3.partial_correlations.find_identical_traits`.
+
+        Given:
+            - the name of a primary trait
+            - the value of a primary trait
+            - a sequence of names of control traits
+            - a sequence of values of control traits
+        When:
+            - the arguments above are passed to the `find_identical_traits`
+              function
+        Then:
+            - Return ALL trait names that have the same value when up to three
+              decimal places are considered
+        """
+        for primn, primv, contn, contv, expected in (
+                ("pt", 12.98395, ("ct0", "ct1", "ct2"),
+                 (0.1234, 2.3456, 3.4567), tuple()),
+                ("pt", 12.98395, ("ct0", "ct1", "ct2"),
+                 (12.98354, 2.3456, 3.4567), ("pt", "ct0")),
+                ("pt", 12.98395, ("ct0", "ct1", "ct2", "ct3"),
+                 (0.1234, 2.3456, 0.1233, 4.5678), ("ct0", "ct2"))
+        ):
+            with self.subTest(
+                    primary_name=primn, primary_value=primv,
+                    control_names=contn, control_values=contv):
+                self.assertEqual(
+                    find_identical_traits(primn, primv, contn, contv), expected)
-- 
cgit v1.2.3