Clearing the way to use haplotype_frequencies in the haplotype_frequencies code.

Author: jonbrenas

malariagen_data/anoph/h12.py

@@ -1,4 +1,3 @@
-from collections import Counter
 from typing import Optional, Tuple, Dict, Mapping
 
 import allel  # type: ignore
@@ -7,7 +6,7 @@
 import bokeh.plotting
 
 from .hap_data import AnophelesHapData
-from ..util import hash_columns, check_types, CacheMiss
+from ..util import check_types, CacheMiss, haplotype_frequencies
 from . import base_params
 from . import h12_params, gplt_params, hap_params
 
@@ -515,21 +514,11 @@ def plot_h12_gwss(
             return fig
 
 
-def haplotype_frequencies(h):
-    """Compute haplotype frequencies, returning a dictionary that maps
-    haplotype hash values to frequencies."""
-    n = h.shape[1]
-    hashes = hash_columns(np.asarray(h))
-    counts = Counter(hashes)
-    freqs = {key: count / n for key, count in counts.items()}
-    return freqs
-
-
 def garud_h12(ht):
     """Compute Garud's H12."""
 
     # Compute haplotype frequencies.
-    frq_counter = haplotype_frequencies(ht)
+    frq_counter, _, _ = haplotype_frequencies(ht)
 
     # Convert to array of sorted frequencies.
     f = np.sort(np.fromiter(frq_counter.values(), dtype=float))[::-1]

malariagen_data/anoph/h1x.py

@@ -6,10 +6,9 @@
 import bokeh.plotting
 
 from .hap_data import AnophelesHapData
-from ..util import check_types, CacheMiss
+from ..util import check_types, CacheMiss, haplotype_frequencies
 from . import base_params
 from . import h12_params, gplt_params, hap_params
-from .h12 import haplotype_frequencies
 
 
 class AnophelesH1XAnalysis(
@@ -342,8 +341,8 @@ def haplotype_joint_frequencies(ha, hb):
     """Compute the joint frequency of haplotypes in two difference
     cohorts. Returns a dictionary mapping haplotype hash values to
     the product of frequencies in each cohort."""
-    frqa = haplotype_frequencies(ha)
-    frqb = haplotype_frequencies(hb)
+    frqa, _, _ = haplotype_frequencies(ha)
+    frqb, _, _ = haplotype_frequencies(hb)
     keys = set(frqa.keys()) | set(frqb.keys())
     joint_freqs = {key: frqa.get(key, 0) * frqb.get(key, 0) for key in keys}
     return joint_freqs

malariagen_data/util.py

@@ -4,6 +4,7 @@
 import re
 import sys
 import warnings
+from collections import Counter
 from enum import Enum
 from math import prod
 from functools import wraps
@@ -1606,6 +1607,18 @@ def hash_columns(x):
     return out
 
 
+def haplotype_frequencies(h):
+    """Compute haplotype frequencies, returning a dictionary that maps
+    haplotype hash values to frequencies."""
+    n = h.shape[1]
+    hashes = hash_columns(np.asarray(h))
+    count = Counter(hashes)
+    freqs = {key: count / n for key, count in count.items()}
+    counts = {key: count for key, count in count.items()}
+    nobs = {key: n for key, count in count.items()}
+    return freqs, counts, nobs
+
+
 def distributed_client():
     from distributed import get_client
 

tests/anoph/test_h12.py

@@ -87,8 +87,10 @@ def test_haplotype_frequencies():
         dtype="i1",
     )
 
-    f = haplotype_frequencies(h1)
+    f, c, o = haplotype_frequencies(h1)
     assert isinstance(f, dict)
+    assert isinstance(c, dict)
+    assert isinstance(o, dict)
     vals = np.array(list(f.values()))
     vals.sort()
     assert np.all(vals >= 0)