admix-kit 0.1.4__py3-none-any.whl

admix/__init__.py

@@ -0,0 +1,17 @@
+from ._logging import logger
+from .dataset import Dataset
+from . import data, dataset, simulate, ancestry, plot, assoc, estimate, tools, io, cli
+from .version import __version__
+
+__all__ = [
+    "data",
+    "simulate",
+    "ancestry",
+    "plot",
+    "assoc",
+    "estimate",
+    "tools",
+    "io",
+    "cli",
+    "Dataset",
+]

admix/_logging.py

@@ -0,0 +1,3 @@
+import structlog
+
+logger = structlog.get_logger()

admix/ancestry/__init__.py

@@ -0,0 +1,4 @@
+from ._wrapper import lampld_wrapper
+from ._local import lanc, lanc_single_chrom
+
+__all__ = ["lampld_wrapper"]

admix/ancestry/_local.py

@@ -0,0 +1,107 @@
+import numpy as np
+from typing import List
+import admix
+
+"""
+This will be the interface to call various local ancestry inference methods
+"""
+
+
+def lanc_single_chrom(
+    sample_dset: admix.Dataset,
+    ref_dsets: List[admix.Dataset],
+    method="lampld",
+    n_proto=6,
+    window_size=300,
+):
+    """Local ancestry inference for data sets of single chromosome
+
+    Parameters
+    ----------
+    sample_dset : admix.Dataset
+        dataset of the sample to be inferred
+    ref_dsets : List[admix.Dataset]
+        dataset of the reference panel
+    method : str, optional
+        method used to perform local ancestry inference, by default "lampld"
+    n_proto : int, optional
+        number of prototypes used in lampld, by default 6
+    window_size : int, optional
+    """
+    from pylampld import LampLD
+
+    assert method in ["lampld"]
+    assert len(sample_dset.snp.CHROM.unique()), "Datasets are not of single chromosome"
+    assert admix.dataset.is_aligned(
+        [sample_dset, *ref_dsets]
+    ), "Datasets are not aligned"
+
+    n_anc = len(ref_dsets)
+    n_snp = sample_dset.n_snp
+
+    model = LampLD(n_snp=n_snp, n_anc=n_anc, n_proto=n_proto, window_size=window_size)
+    ref_list = [
+        np.vstack(
+            [
+                dset.geno[:, :, 0].T,
+                dset.geno[:, :, 1].T,
+            ]
+        )
+        for dset in ref_dsets
+    ]
+
+    model.set_pos(sample_dset.snp.POS.values)
+    model.fit(ref_list)
+
+    est = np.dstack(
+        [
+            model.infer_lanc(sample_dset.geno[:, :, 0].T.compute()),
+            model.infer_lanc(sample_dset.geno[:, :, 1].T.compute()),
+        ]
+    )
+    return np.swapaxes(est.astype(np.int8), 0, 1)
+
+
+def lanc(
+    sample_dset: admix.Dataset,
+    ref_dsets: List[admix.Dataset],
+    method="lampld",
+    n_proto=6,
+    window_size=300,
+) -> np.ndarray:
+    """Local ancestry inference
+
+    Parameters
+    ----------
+    sample_dset : admix.Dataset
+        dataset of the sample to be inferred
+    ref_dsets : List[admix.Dataset]
+        dataset of the reference panel
+    method : str, optional
+        method used to perform local ancestry inference, by default "lampld"
+    n_proto : int, optional
+        number of prototypes used in lampld, by default 6
+    window_size : int, optional
+        window size used in local ancestry inference, by default 100
+
+    Returns
+    -------
+    np.ndarray
+        Estimated local ancestries
+    """
+
+    assert method in ["lampld"]
+    chrom_list = sample_dset.snp.CHROM.unique()
+    lanc = []
+    for chrom in chrom_list:
+        chrom_sample_dset = sample_dset[(sample_dset.snp.CHROM == chrom).values]
+        chrom_ref_dsets = [dset[(dset.snp.CHROM == chrom).values] for dset in ref_dsets]
+        chrom_lanc = lanc_single_chrom(
+            chrom_sample_dset,
+            chrom_ref_dsets,
+            method=method,
+            n_proto=n_proto,
+            window_size=window_size,
+        )
+        lanc.append(chrom_lanc)
+    return np.concatenate(lanc, axis=0)

admix/ancestry/_wrapper.py

@@ -0,0 +1,48 @@
+import numpy as np
+import pandas as pd
+import tempfile
+from os.path import join
+import subprocess
+from ..io import read_digit_mat, write_digit_mat
+
+
+def lampld_wrapper(
+    ref_hap_list,
+    sample_hap,
+    snp_pos,
+    window_size=300,
+    n_proto=6,
+    bin_dir="/u/project/pasaniuc/kangchen/admix-tools/cache/LAMPLD-v1.3/bin/",
+):
+    tmp = tempfile.TemporaryDirectory()
+    tmp_dir = tmp.name
+
+    assert len(ref_hap_list) in [2, 3], "Currently only support 2-way / 3-way ancestry"
+    for ref_i in range(len(ref_hap_list)):
+        write_digit_mat(join(tmp_dir, f"ref{ref_i}.hap"), ref_hap_list[ref_i])
+    write_digit_mat(join(tmp_dir, "sample.hap"), sample_hap)
+    np.savetxt(join(tmp_dir, "pos.txt"), snp_pos, fmt="%s")
+
+    if len(ref_hap_list) == 2:
+        bin_path = join(bin_dir, "haplanc2way")
+    elif len(ref_hap_list) == 3:
+        bin_path = join(bin_dir, "haplanc")
+    else:
+        raise NotImplementedError
+    cmd = " ".join(
+        [
+            bin_path,
+            str(window_size),
+            str(n_proto),
+            join(tmp_dir, "pos.txt"),
+            *[join(tmp_dir, f"ref{ref_i}.hap") for ref_i in range(len(ref_hap_list))],
+            join(tmp_dir, "sample.hap"),
+            join(tmp_dir, "out.txt"),
+            str(1),
+        ]
+    )
+    print(cmd)
+    subprocess.check_output(cmd, shell=True, stderr=subprocess.STDOUT)
+    infer_lanc = read_digit_mat(join(tmp_dir, "out.txt"))
+    tmp.cleanup()
+    return infer_lanc