test_allele

limix · Jan 24, 2019 · 7c4ff10 · 7c4ff10
1 parent 48c5797
commit 7c4ff10
Show file tree

Hide file tree

Showing 12 changed files with 77 additions and 1 deletion.
diff --git a/limix/qc/_allele.py b/limix/qc/_allele.py
@@ -31,14 +31,32 @@ def compute_maf(X):
         [0.49  0.49  0.445 0.495 0.5   0.45  0.48  0.48  0.47  0.435]
     """
     import dask.array as da
+    import xarray as xr
+    from pandas import DataFrame
     from numpy import isnan, logical_not, minimum, nansum
 
     if isinstance(X, da.Array):
         s0 = da.nansum(X, axis=0).compute()
         denom = 2 * (X.shape[0] - da.isnan(X).sum(axis=0)).compute()
+    elif isinstance(X, DataFrame):
+        s0 = X.sum(axis=0, skipna=True)
+        denom = 2 * logical_not(X.isna()).sum(axis=0)
+    elif isinstance(X, xr.DataArray):
+        if "sample" in X.dims:
+            kwargs = {"dim": "sample"}
+        else:
+            kwargs = {"axis": 0}
+        s0 = X.sum(skipna=True, **kwargs)
+        denom = 2 * logical_not(isnan(X)).sum(**kwargs)
     else:
         s0 = nansum(X, axis=0)
         denom = 2 * logical_not(isnan(X)).sum(axis=0)
+
     s0 = s0 / denom
     s1 = 1 - s0
-    return minimum(s0, s1)
+    maf = minimum(s0, s1)
+
+    if hasattr(maf, "name"):
+        maf.name = "maf"
+
+    return maf
diff --git a/limix/qc/test/test_allele.py b/limix/qc/test/test_allele.py
@@ -0,0 +1,58 @@
+import dask.array as da
+import xarray as xr
+from numpy import array_equal, ndarray
+from numpy.random import RandomState
+from numpy.testing import assert_, assert_allclose
+from pandas import DataFrame, Series
+
+from limix.qc import compute_maf
+
+
+def test_compute_maf_numpy():
+    random = RandomState(0)
+    X = random.randint(0, 3, size=(100, 10))
+
+    maf = compute_maf(X)
+    assert_allclose(maf, [0.49, 0.49, 0.445, 0.495, 0.5, 0.45, 0.48, 0.48, 0.47, 0.435])
+
+
+def test_compute_maf_dataframe():
+    random = RandomState(0)
+
+    X = random.randint(0, 3, size=(100, 10))
+    columns = [f"snp{i}" for i in range(X.shape[1])]
+    maf = compute_maf(DataFrame(X, columns=columns))
+
+    assert_(isinstance(maf, Series))
+    assert_(maf.name == "maf")
+    assert_(array_equal(maf.index, columns))
+    assert_allclose(maf, [0.49, 0.49, 0.445, 0.495, 0.5, 0.45, 0.48, 0.48, 0.47, 0.435])
+
+
+def test_compute_maf_dask_array():
+    random = RandomState(0)
+
+    X = da.from_array(random.randint(0, 3, size=(100, 10)), chunks=2)
+    maf = compute_maf(X)
+
+    assert_(isinstance(maf, ndarray))
+    assert_allclose(maf, [0.49, 0.49, 0.445, 0.495, 0.5, 0.45, 0.48, 0.48, 0.47, 0.435])
+
+
+def test_compute_maf_dataarray():
+    random = RandomState(0)
+
+    X = random.randint(0, 3, size=(100, 10))
+    samples = [f"snp{i}" for i in range(X.shape[0])]
+    candidates = [f"snp{i}" for i in range(X.shape[1])]
+    X = xr.DataArray(
+        X,
+        dims=["sample", "candidate"],
+        coords={"sample": samples, "candidate": candidates},
+    )
+    maf = compute_maf(X)
+
+    assert_(isinstance(maf, xr.DataArray))
+    assert_(maf.name == "maf")
+    assert_(array_equal(maf.candidate, candidates))
+    assert_allclose(maf, [0.49, 0.49, 0.445, 0.495, 0.5, 0.45, 0.48, 0.48, 0.47, 0.435])
diff --git a/limix/qc/test/test_impute.py b/limix/qc/test/test_impute.py
diff --git a/limix/qc/test/test_kinship.py b/limix/qc/test/test_kinship.py
diff --git a/limix/qc/test/test_ld.py b/limix/qc/test/test_ld.py
diff --git a/limix/qc/test/test_lik.py b/limix/qc/test/test_lik.py
diff --git a/limix/qc/test/test_linalg.py b/limix/qc/test/test_linalg.py
diff --git a/limix/qc/test/test_missing.py b/limix/qc/test/test_missing.py
diff --git a/limix/qc/test/test_regress.py b/limix/qc/test/test_regress.py
diff --git a/limix/qc/test/test_test b/limix/qc/test/test_test
diff --git a/limix/qc/test/test_trans.py b/limix/qc/test/test_trans.py
diff --git a/limix/qc/test/test_unique.py b/limix/qc/test/test_unique.py