Add generic fn for computing k and gwa values

4 files changed, 81 insertions, 4 deletions

diff --git a/gn3/api/gemma.py b/gn3/api/gemma.py
index d82f5a5..d08c11b 100644
--- a/gn3/api/gemma.py
+++ b/gn3/api/gemma.py
@@ -12,6 +12,7 @@ from gn3.commands import queue_cmd
 from gn3.commands import run_cmd
 from gn3.file_utils import get_hash_of_files
 from gn3.file_utils import jsonfile_to_dict
+from gn3.computations.gemma import compute_k_values
 from gn3.computations.gemma import do_paths_exist
 from gn3.computations.gemma import generate_hash_of_string
 from gn3.computations.gemma import generate_pheno_txt_file
diff --git a/gn3/computations/gemma.py b/gn3/computations/gemma.py
index 7b5f91d..3c6bd29 100644
--- a/gn3/computations/gemma.py
+++ b/gn3/computations/gemma.py
@@ -8,6 +8,7 @@ from typing import List
 from typing import Optional
 from typing import ValuesView
 from gn3.commands import compose_gemma_cmd
+from gn3.file_utils import get_hash_of_files
 
 
 def generate_hash_of_string(unhashed_str: str) -> str:
@@ -68,3 +69,32 @@ def generate_gemma_computation_cmd(
                              gemma_wrapper_kwargs=gemma_wrapper_kwargs,
                              gemma_kwargs=_kwargs,
                              gemma_args=["-gk", ">", output_file])
+
+
+# pylint: disable=R0913
+def generate_gemma_cmd(gemma_cmd: str,
+                       output_dir: str,
+                       token: str,
+                       gemma_kwargs: Dict,
+                       gemma_wrapper_kwargs: Dict = None,
+                       chromosomes: str = None) -> Dict:
+    """Compute k values"""
+    _hash = get_hash_of_files(
+        [v for k, v in gemma_kwargs.items() if k in ["g", "p", "a", "c"]])
+    if chromosomes:  # Only reached when calculating k-values
+        gemma_wrapper_kwargs = {"loco": f"--input {chromosomes}"}
+        _hash += f"-{generate_hash_of_string(chromosomes)[:6]}"
+    _output_filename = f"{_hash}-output.json"
+    return {
+        "output_file":
+        _output_filename,
+        "gemma_cmd":
+        compose_gemma_cmd(gemma_wrapper_cmd=gemma_cmd,
+                          gemma_wrapper_kwargs=gemma_wrapper_kwargs,
+                          gemma_kwargs=gemma_kwargs,
+                          gemma_args=[
+                              "-gk", ">",
+                              (f"{output_dir}/"
+                               f"{token}/{_output_filename}")
+                          ])
+    }
diff --git a/tests/integration/test_gemma.py b/tests/integration/test_gemma.py
index 822f088..f4ac685 100644
--- a/tests/integration/test_gemma.py
+++ b/tests/integration/test_gemma.py
@@ -162,7 +162,7 @@ class GemmaAPITest(unittest.TestCase):
         self.assertEqual(response.get_json(), {"status": "test"})
 
     @mock.patch("gn3.api.gemma.queue_cmd")
-    @mock.patch("gn3.api.gemma.get_hash_of_files")
+    @mock.patch("gn3.computations.gemma.get_hash_of_files")
     @mock.patch("gn3.api.gemma.jsonfile_to_dict")
     @mock.patch("gn3.api.gemma.do_paths_exist")
     @mock.patch("gn3.api.gemma.redis.Redis")
@@ -195,7 +195,7 @@ class GemmaAPITest(unittest.TestCase):
         })
 
     @mock.patch("gn3.api.gemma.queue_cmd")
-    @mock.patch("gn3.api.gemma.get_hash_of_files")
+    @mock.patch("gn3.computations.gemma.get_hash_of_files")
     @mock.patch("gn3.api.gemma.jsonfile_to_dict")
     @mock.patch("gn3.api.gemma.do_paths_exist")
     @mock.patch("gn3.api.gemma.redis.Redis")
@@ -223,9 +223,9 @@ class GemmaAPITest(unittest.TestCase):
                  "-g /tmp/test-data/genofile.txt "
                  "-p /tmp/test-data/phenofile.txt "
                  "-a /tmp/test-data/snpfile.txt "
-                 "-gk > /tmp/test-data/hash-k-output.json"))
+                 "-gk > /tmp/test-data/hash-3R77Mz-k-output.json"))
         self.assertEqual(response.get_json(), {
-            "output_file": "hash-k-output.json",
+            "output_file": "hash-3R77Mz-k-output.json",
             "status": "queued",
             "unique_id": "my-unique-id"
         })
diff --git a/tests/unit/computations/test_gemma.py b/tests/unit/computations/test_gemma.py
index 2d2fbf9..d3fb7aa 100644
--- a/tests/unit/computations/test_gemma.py
+++ b/tests/unit/computations/test_gemma.py
@@ -2,6 +2,7 @@
 import unittest
 
 from unittest import mock
+from gn3.computations.gemma import compute_k_values
 from gn3.computations.gemma import generate_hash_of_string
 from gn3.computations.gemma import generate_pheno_txt_file
 from gn3.computations.gemma import generate_gemma_computation_cmd
@@ -50,3 +51,48 @@ class TestGemma(unittest.TestCase):
                 "test.txt -a genofile_snps.txt "
                 "-gk > /tmp/gn2/"
                 "k_output_gUFhGu4rLG7k+CXLPk1OUg.txt"))
+
+    @mock.patch("gn3.computations.gemma.get_hash_of_files")
+    def test_compute_k_values_without_loco(self, mock_get_hash):
+        """Test computing k valuse without loco"""
+        mock_get_hash.return_value = "my-hash"
+        self.assertEqual(
+            compute_k_values(gemma_cmd="gemma-wrapper",
+                             output_dir="/tmp",
+                             token="my-token",
+                             gemma_kwargs={
+                                 "g": "genofile",
+                                 "p": "phenofile",
+                                 "a": "snpsfile"
+                             }), {
+                                 "output_file":
+                                 "my-hash-k-output.json",
+                                 "gemma_cmd":
+                                 ("gemma-wrapper --json -- -g genofile "
+                                  "-p phenofile -a snpsfile "
+                                  "-gk > /tmp/my-token/my-hash-k-output.json")
+                             })
+
+    @mock.patch("gn3.computations.gemma.get_hash_of_files")
+    def test_compute_k_values_with_loco(self, mock_get_hash):
+        """Test computing k valuse with loco"""
+        mock_get_hash.return_value = "my-hash"
+        self.assertEqual(
+            compute_k_values(gemma_cmd="gemma-wrapper",
+                             output_dir="/tmp",
+                             token="my-token",
+                             chromosomes="1,2,3,4,5",
+                             gemma_kwargs={
+                                 "g": "genofile",
+                                 "p": "phenofile",
+                                 "a": "snpsfile"
+                             }), {
+                                 "output_file":
+                                 "my-hash-r+gF5a-k-output.json",
+                                 "gemma_cmd": ("gemma-wrapper --json "
+                                               "--loco --input 1,2,3,4,5 "
+                                               "-- -g genofile "
+                                               "-p phenofile -a snpsfile "
+                                               "-gk > /tmp/my-token/"
+                                               "my-hash-r+gF5a-k-output.json")
+                             })