Python API

class LancData:
    """The genotype and local ancestry data for a single chromosome/dataset.

    Attributes:
        pgen (PgenReader): A pgenlib PgenReader object.
        pvar (PvarReader): A pgenlib PVarReader object.
        lanc (FlatLanc): A FlatLanc object with local ancestry data.
        ancestries (list[str]): An ordered list of ancestry names. The integer codes in
            the .lanc file and `self.lanc` correspond to indices in this list (e.g.
            0 -> ancestries[0]).
        plink_prefix (str): The prefix for the corresponding plink2 fileset.
    """

    def __init__(
        self,
        plink_prefix: str,
        lanc_file: str,
        ancestries: list[str] | None = None,
    ):
        """Constructs a LancData from plink2 files.

        Args:
            plink_prefix (str): The prefix for a plink2 fileset.
            lanc_file (str): The path to a .lanc file.
            ancestries (Optional[list[str]): An optional list of ordered ancestry
                names corresponding to the .lanc file.
        """
        with ExitStack() as stack:
            pgen = PgenReader(bytes(plink_prefix + ".pgen", "utf8"))
            stack.callback(pgen.close)
            pvar = PvarReader(bytes(plink_prefix + ".pvar", "utf8"))
            stack.callback(pvar.close)
            lanc = _read_lanc(lanc_file)
            n_variants = pvar.get_variant_ct()
            if lanc.n_variants != n_variants:
                raise ValueError("PLINK and .lanc files have different numbers of variants")
            if lanc.offsets.shape[0] - 1 != pgen.get_raw_sample_ct():
                raise ValueError("PLINK and .lanc files have different numbers of samples")

            if ancestries is None:
                all_values = np.concatenate([lanc.left_haps, lanc.right_haps])
                ancestries = [str(i) for i in np.unique(all_values)]
            elif len(ancestries) <= int(max(lanc.left_haps.max(), lanc.right_haps.max())):
                raise ValueError("Ancestry names do not cover all values in the .lanc file")
            stack.pop_all()

        self.pgen = pgen
        self.pvar = pvar
        self.lanc = lanc
        self.ancestries = ancestries
        self.plink_prefix = plink_prefix
        self._closed = False

    def close(self) -> None:
        """Release the underlying PLINK readers.

        Calling ``close`` more than once is safe. Query methods cannot be used
        after the readers have been closed.
        """
        if not self._closed:
            self.pgen.close()
            self.pvar.close()
            self._closed = True

    def __enter__(self) -> LancData:
        if self._closed:
            raise RuntimeError("LancData is closed")
        return self

    def __exit__(self, exc_type, exc_value, traceback) -> None:
        self.close()

    def _ensure_open(self) -> None:
        if self._closed:
            raise RuntimeError("LancData is closed")

    def get_info(self, indices: NDArray[np.uint32]) -> DataFrame:
        """Query info for a set of variants.

        Args:
            indices: The variant indices in pvar order (0-based), shape (V,)

        Returns:
            pandas.DataFrame: One row per variant with the following columns:

                - CHR (str): Chromosome name. \n
                - BP (int): 1-based genomic position. \n
                - REF (str): Reference allele. \n
                - ALT (str): Alternate allele. \n
                - ID (str): Variant identifier. \n
        """

        self._ensure_open()
        return _get_info(self.pvar, _validate_indices(indices, self.pvar.get_variant_ct()))

    def get_lanc(self, indices: NDArray[np.unsignedinteger]) -> NDArray[np.uint8]:
        """Query phased local ancestry.

        Args:
            indices: The variant indices in pvar order (0-based), shape (V,)

        Returns:
            An array of ancestries, shape (N, V, 2)
        """

        self._ensure_open()
        return self.lanc.get_lanc(_validate_indices(indices, self.pvar.get_variant_ct()))

    def get_lanc_dosage(self, indices: NDArray[np.uint32]) -> NDArray[np.int32]:
        """Query local ancestry dosage.

        Args:
            indices: An array of variant indices in pvar order (0-based), shape (V,)

        Returns:
            An array of local ancestry dosages, shape (N, V, K) (where K is the
                number of ancestries)
        """

        self._ensure_open()
        indices = _validate_indices(indices, self.pvar.get_variant_ct())
        lanc = np.asarray(self.get_lanc(indices), dtype=np.uint8)
        ancestries = np.arange(len(self.ancestries), dtype=np.uint8)
        left_haps_mask = (lanc[:, :, 0:1] == ancestries[None, None, :]).astype(np.int32)
        right_haps_mask = (lanc[:, :, 1:2] == ancestries[None, None, :]).astype(np.int32)
        return left_haps_mask + right_haps_mask

    def get_geno(self, indices: NDArray[np.uint32]) -> NDArray[np.int32]:
        """Query phased genotypes.

        Args:
            indices: An array of variant indices (0-based)

        Returns:
            An array of phased genotypes, shape (N, V, 2)
        """

        self._ensure_open()
        return _get_geno(
            self.pgen,
            _validate_indices(indices, self.pvar.get_variant_ct()),
        )

    def get_lanc_geno(self, indices: NDArray[np.unsignedinteger]) -> NDArray[np.int32]:
        """Query genotypes deconvoluted/masked by ancestry.

        Args:
            indices: An array of variant indices (0-based)

        Returns:
            An array of genotypes masked by ancestry, shape (N, V, 2)
        """
        self._ensure_open()
        indices = _validate_indices(indices, self.pvar.get_variant_ct())
        geno = np.asarray(self.get_geno(indices), dtype=np.int32)
        lanc = np.asarray(self.lanc.get_lanc(indices), dtype=np.uint8)
        ancestries = np.arange(len(self.ancestries), dtype=np.uint8)
        left_haps_mask = (lanc[:, :, 0:1] == ancestries[None, None, :]).astype(np.int32)
        right_haps_mask = (lanc[:, :, 1:2] == ancestries[None, None, :]).astype(np.int32)
        geno_masked = left_haps_mask * geno[:, :, 0:1] + right_haps_mask * geno[:, :, 1:2]
        return geno_masked

class FlatLanc:
    """Stores .lanc file ancestry data in a flattened structure for fast querying.

    Attributes:
        right_haps (NDArray[uint8]): Concatenated right haplotypes for all samples, shape (H,)
        left_haps (NDArray[uint8]): Concatenated left haplotypes for all samples, shape (H,)
        breakpoints (NDArray[uint32]): Concatenated breakpoints for all samples, shape (H,)
        offsets (NDArray[uint32]): Cumulative end indices separating samples, shape (N,)
    """

    def __init__(
        self,
        left_haps: NDArray[np.uint8],
        right_haps: NDArray[np.uint8],
        breakpoints: NDArray[np.uint32],
        offsets: NDArray[np.uint32],
        n_variants: int | None = None,
    ):
        arrays = (left_haps, right_haps, breakpoints, offsets)
        if any(np.asarray(array).ndim != 1 for array in arrays):
            raise ValueError("FlatLanc arrays must be one-dimensional")
        if not (len(left_haps) == len(right_haps) == len(breakpoints)):
            raise ValueError("FlatLanc tract arrays must have equal lengths")
        if len(offsets) == 0 or offsets[0] != 0 or offsets[-1] != len(breakpoints):
            raise ValueError("FlatLanc offsets must delimit the tract arrays")
        if np.any(np.diff(offsets) < 0):
            raise ValueError("FlatLanc offsets must be non-decreasing")
        if n_variants is not None and n_variants <= 0:
            raise ValueError("FlatLanc variant count must be positive")

        self.left_haps = left_haps
        self.right_haps = right_haps
        self.breakpoints = breakpoints
        self.offsets = offsets
        self.n_variants = n_variants

    def get_lanc(self, indices: NDArray[np.unsignedinteger]) -> NDArray[np.uint8]:
        """Query phased local ancestry.

        Args:
            indices: The variant indices in (0-based)

        Returns:
            An array of ancestries, shape (N, V, 2)
        """

        if self.n_variants is not None:
            indices = _validate_indices(indices, self.n_variants)
        else:
            indices = np.asarray(indices)
            if indices.ndim != 1:
                raise ValueError("Variant indices must be a one-dimensional array")
            if not np.issubdtype(indices.dtype, np.integer):
                raise TypeError("Variant indices must have an integer dtype")
            if np.issubdtype(indices.dtype, np.signedinteger) and np.any(indices < 0):
                raise IndexError("Variant indices must be non-negative")
            if np.any(indices >= self.breakpoints[-1]):
                raise IndexError("Variant index is outside the .lanc variant range")

        idx_order = np.argsort(indices)
        idx_ordered = np.ascontiguousarray(indices[idx_order])
        idx_inverse = np.argsort(idx_order)
        left, right = _get_lanc(
            self.left_haps,
            self.right_haps,
            self.breakpoints,
            self.offsets,
            idx_ordered,
        )
        return np.stack((left[:, idx_inverse], right[:, idx_inverse]), axis=-1)