diff --git a/apis/python/src/tiledbvcf/allele_frequency.py b/apis/python/src/tiledbvcf/allele_frequency.py index 588e1f997..6b7e16834 100644 --- a/apis/python/src/tiledbvcf/allele_frequency.py +++ b/apis/python/src/tiledbvcf/allele_frequency.py @@ -11,4 +11,4 @@ def read_allele_frequency(dataset_uri: str, region: str) -> pandas.DataFrame: import tiledbvcf ds = tiledbvcf.Dataset(uri=dataset_uri, mode="r") - return ds.read_variant_stats(region).to_pandas() + return ds.read_variant_stats(region) diff --git a/apis/python/src/tiledbvcf/dataset.py b/apis/python/src/tiledbvcf/dataset.py index c4a6a7d04..102a79617 100644 --- a/apis/python/src/tiledbvcf/dataset.py +++ b/apis/python/src/tiledbvcf/dataset.py @@ -2,10 +2,11 @@ import shutil import warnings from collections import namedtuple -from typing import List +from typing import Generator, List import pandas as pd import pyarrow as pa +import pyarrow.compute as pc from . import libtiledbvcf @@ -100,6 +101,43 @@ class Dataset(object): TileDB configuration, alternative to `cfg.tiledb_config`. """ + class Region(object): + """ + Represents a 1-based inclusive region. + + Parameters + ---------- + region + A string in the form ":-". + """ + + def __init__(self, region: str): + try: + contig, interval = region.split(":") + start, end = map(int, interval.split("-")) + except Exception: + raise Exception( + '"region" parameter must have format ":-"' + ) + if contig == "": + raise Exception("Region contig cannot be empty") + if start <= 0: + raise Exception("Regions must be 1-based") + if end < start: + raise Exception(f'"{interval}" is not a valid region interval') + self.contig = contig + self.start = start + self.end = end + + def __lt__(self, region): + return self.to_tuple() < region.to_tuple() + + def __str__(self): + return f"{self.contig}:{self.start}-{self.end}" + + def to_tuple(self): + return (self.contig, self.start, self.end) + def __init__( self, uri: str, @@ -219,6 +257,22 @@ def _set_write_cfg(self, cfg): pass self.writer.set_tiledb_config(",".join(tiledb_config_list)) + # parses and sorts regions, then generates consolidated regions one at a time + def _prepare_regions(self, regions: List[str]) -> Generator[Region, None, None]: + if not regions: + return + prev_region, *parsed_regions = sorted(map(self.Region, regions)) + for r in parsed_regions: + if prev_region.contig != r.contig: + yield prev_region + prev_region = r + elif r.start <= prev_region.end + 1: + prev_region.end = max(r.end, prev_region.end) + else: # the regions are non-overlapping + yield prev_region + prev_region = r + yield prev_region + def read_arrow( self, attrs: List[str] = DEFAULT_ATTRS, @@ -271,6 +325,11 @@ def read_arrow( if isinstance(regions, str): regions = [regions] + if isinstance(regions, list): + regions = map(str, self._prepare_regions(regions)) + else: + regions = "" + if isinstance(samples, str): samples = [samples] @@ -278,7 +337,6 @@ def read_arrow( self.reader.set_export_to_disk(False) self._set_samples(samples, samples_file) - regions = "" if regions is None else regions self.reader.set_regions(",".join(regions)) self.reader.set_attributes(attrs) self.reader.set_check_samples_exist(not skip_check_samples) @@ -294,36 +352,177 @@ def read_arrow( def read_variant_stats( self, region: str = None, + drop_ref: bool = False, + regions: List[str] = None, + scan_all_samples: bool = False, ) -> pd.DataFrame: """ Read variant stats from the dataset into a Pandas DataFrame Parameters ---------- + drop_ref + Omit "ref" alleles from the results + regions + Genomic regions to be queried. + scan_all_samples + Scan all samples when computing internal allele frequency. + region + **DEPRECATED** - Genomic region to be queried. + """ + # TODO: deprecated region and parse regions like read() + if not (region or regions): + raise Exception('"region" or "regions" parameter is required') + if region and regions: + raise Exception('"region" and "regions" parameters are mutually exclusive') + if region: + warnings.warn( + '"region" parameter is deprecated, use "regions" instead', + DeprecationWarning, + ) + regions = [region] + + kwargs = { + "drop_ref": drop_ref, + "regions": regions, + "scan_all_samples": scan_all_samples, + } + return self.read_variant_stats_arrow(**kwargs).to_pandas( + split_blocks=True, self_destruct=True + ) + + def read_variant_stats_arrow( + self, + region: str = None, + drop_ref: bool = False, + regions: List[str] = None, + scan_all_samples: bool = False, + ) -> pa.Table: + """ + Read variant stats from the dataset into a PyArrow Table + + Parameters + ---------- + drop_ref + Omit "ref" alleles from the results + regions + Genomic regions to be queried. + scan_all_samples + Scan all samples when computing internal allele frequency. region - Genomic region to be queried. + **DEPRECATED** - Genomic region to be queried. """ + # TODO: deprecated region and parse regions like read() + if not (region or regions): + raise Exception('"region" or "regions" parameter is required') + if region and regions: + raise Exception('"region" and "regions" parameters are mutually exclusive') + if region: + warnings.warn( + '"region" parameter is deprecated, use "regions" instead', + DeprecationWarning, + ) + regions = [region] if self.mode != "r": raise Exception("Dataset not open in read mode") - self.reader.set_regions(region) - return self.reader.get_variant_stats_results() + + self.reader.reset() + self.reader.set_scan_all_samples(scan_all_samples) + + # generates stats, sorts the results, and adds contig column one region at a time + def variant_stats_generator(regions): + for r in regions: + self.reader.set_regions(str(r)) + stats = self.reader.get_variant_stats_results() + stats = stats.sort_by([("pos", "ascending"), ("alleles", "ascending")]) + n = stats.num_rows + contig_col = [r.contig] * n + yield stats.add_column(0, "contig", [contig_col]) + + # drop reference alleles from results + consolidated_regions = self._prepare_regions(regions) + stats_tbl = pa.concat_tables(variant_stats_generator(consolidated_regions)) + if drop_ref: + expr = pc.field("alleles") != "ref" + return stats_tbl.filter(expr) + return stats_tbl def read_allele_count( self, region: str = None, + regions: List[str] = None, ) -> pd.DataFrame: """ Read allele count from the dataset into a Pandas DataFrame Parameters ---------- + regions + Genomic regions to be queried. region - Genomic region to be queried. + **DEPRECATED** - Genomic region to be queried. """ + # TODO: deprecated region and parse regions like read() + if not (region or regions): + raise Exception('"region" or "regions" parameter is required') + if region and regions: + raise Exception('"region" and "regions" parameters are mutually exclusive') + if region: + warnings.warn( + '"region" parameter is deprecated, use "regions" instead', + DeprecationWarning, + ) + regions = [region] if self.mode != "r": raise Exception("Dataset not open in read mode") - self.reader.set_regions(region) - return self.reader.get_allele_count_results() + + return self.read_allele_count_arrow(regions=regions).to_pandas( + split_blocks=True, self_destruct=True + ) + + def read_allele_count_arrow( + self, + region: str = None, + regions: List[str] = None, + ) -> pa.Table: + """ + Read allele count from the dataset into a Pandas DataFrame + + Parameters + ---------- + regions + Genomic regions to be queried. + region + **DEPRECATED** - Genomic region to be queried. + """ + # TODO: deprecated region and parse regions like read() + if not (region or regions): + raise Exception('"region" or "regions" parameter is required') + if region and regions: + raise Exception('"region" and "regions" parameters are mutually exclusive') + if region: + warnings.warn( + '"region" parameter is deprecated, use "regions" instead', + DeprecationWarning, + ) + regions = [region] + if self.mode != "r": + raise Exception("Dataset not open in read mode") + + # generates counts and adds contig column one region at a time + def allele_count_generator(regions): + for r in regions: + self.reader.set_regions(str(r)) + counts = self.reader.get_allele_count_results() + contigs = counts.sort_by( + [("pos", "ascending"), ("ref", "ascending"), ("alt", "ascending")] + ) + n = counts.num_rows + contig_col = [r.contig] * n + yield counts.add_column(0, "contig", [contig_col]) + + consolidated_regions = self._prepare_regions(regions) + return pa.concat_tables(allele_count_generator(consolidated_regions)) def read( self, @@ -377,6 +576,10 @@ def read( if isinstance(regions, str): regions = [regions] + if isinstance(regions, list): + regions = map(str, self._prepare_regions(regions)) + else: + regions = "" if isinstance(samples, str): samples = [samples] @@ -384,7 +587,6 @@ def read( self.reader.set_export_to_disk(False) self._set_samples(samples, samples_file) - regions = "" if regions is None else regions self.reader.set_regions(",".join(regions)) self.reader.set_attributes(attrs) self.reader.set_check_samples_exist(not skip_check_samples) @@ -446,6 +648,10 @@ def export( if isinstance(regions, str): regions = [regions] + if isinstance(regions, list): + regions = map(str, self._prepare_regions(regions)) + else: + regions = "" if isinstance(samples, str): samples = [samples] @@ -453,7 +659,6 @@ def export( self.reader.set_export_to_disk(True) self._set_samples(samples, samples_file) - regions = "" if regions is None else regions self.reader.set_regions(",".join(regions)) self.reader.set_check_samples_exist(not skip_check_samples) self.reader.set_enable_progress_estimation(enable_progress_estimation) @@ -501,13 +706,17 @@ def read_iter( if isinstance(regions, str): regions = [regions] + if isinstance(regions, list): + regions = map(str, self._prepare_regions(regions)) + else: + regions = "" if isinstance(samples, str): samples = [samples] self.reader.reset() if not self.read_completed(): - yield self.read(attrs, samples, regions, samples_file, bed_file) + yield self.read(attrs, samples, list(regions), samples_file, bed_file) while not self.read_completed(): yield self.continue_read() @@ -527,7 +736,7 @@ def continue_read(self, release_buffers: bool = True) -> pd.DataFrame: """ table = self.continue_read_arrow(release_buffers=release_buffers) - return table.to_pandas() + return table.to_pandas(split_blocks=True, self_destruct=True) def continue_read_arrow(self, release_buffers: bool = True) -> pa.Table: """ @@ -593,14 +802,18 @@ def count( if isinstance(regions, str): regions = [regions] + if isinstance(regions, list): + regions = map(str, self._prepare_regions(regions)) + else: + regions = "" if isinstance(samples, str): samples = [samples] + elif samples is None: + samples = "" self.reader.reset() self.reader.set_export_to_disk(False) - samples = "" if samples is None else samples - regions = "" if regions is None else regions self.reader.set_samples(",".join(samples)) self.reader.set_regions(",".join(regions)) diff --git a/apis/python/src/tiledbvcf/sample_qc.py b/apis/python/src/tiledbvcf/sample_qc.py index 298e24793..8da0f93f8 100644 --- a/apis/python/src/tiledbvcf/sample_qc.py +++ b/apis/python/src/tiledbvcf/sample_qc.py @@ -26,4 +26,6 @@ def sample_qc( pandas.DataFrame """ - return clib.sample_qc(dataset_uri, samples=samples, config=config).to_pandas() + return clib.sample_qc(dataset_uri, samples=samples, config=config).to_pandas( + split_blocks=True, self_destruct=True + ) diff --git a/apis/python/tests/test_tiledbvcf.py b/apis/python/tests/test_tiledbvcf.py index f79d95edf..26e2a27f8 100755 --- a/apis/python/tests/test_tiledbvcf.py +++ b/apis/python/tests/test_tiledbvcf.py @@ -3,6 +3,7 @@ import subprocess import os import pandas as pd +import pyarrow as pa import re import glob import shutil @@ -1294,22 +1295,361 @@ def test_ingest_with_stats_v3( ]["info_TILEDB_IAF"].iloc[0][0] == 0.9375 ) - df = test_stats_v3_ingestion.read_variant_stats("chr1:1-10000") - assert df.shape == (13, 5) + + ###################### + # read_variant_stats # + ###################### + + # test errors + no_parameter_error = '"region" or "regions" parameter is required' + exclusive_parameter_error = ( + '"region" and "regions" parameters are mutually exclusive' + ) + format_error = '"region" parameter must have format ":-"' + empty_contig_error = "Region contig cannot be empty" + base_1_error = "Regions must be 1-based" + interval_error = '"100-1" is not a valid region interval' + with pytest.raises(Exception, match=no_parameter_error): + test_stats_v3_ingestion.read_variant_stats() + with pytest.raises(Exception, match=no_parameter_error): + test_stats_v3_ingestion.read_variant_stats_arrow() + with pytest.raises(Exception, match=exclusive_parameter_error): + test_stats_v3_ingestion.read_variant_stats("chr1:1-100", regions=["chr1:1-100"]) + with pytest.raises(Exception, match=exclusive_parameter_error): + test_stats_v3_ingestion.read_variant_stats_arrow( + "chr1:1-100", regions=["chr1:1-100"] + ) + with pytest.raises(Exception, match=format_error): + test_stats_v3_ingestion.read_variant_stats(regions=[""]) + with pytest.raises(Exception, match=format_error): + test_stats_v3_ingestion.read_variant_stats_arrow(regions=[""]) + with pytest.raises(Exception, match=format_error): + test_stats_v3_ingestion.read_variant_stats(regions=["chr1"]) + with pytest.raises(Exception, match=format_error): + test_stats_v3_ingestion.read_variant_stats_arrow(regions=["chr1"]) + with pytest.raises(Exception, match=format_error): + test_stats_v3_ingestion.read_variant_stats(regions=["chr1:-"]) + with pytest.raises(Exception, match=format_error): + test_stats_v3_ingestion.read_variant_stats_arrow(regions=["chr1:-"]) + with pytest.raises(Exception, match=empty_contig_error): + test_stats_v3_ingestion.read_variant_stats(regions=[":1-100"]) + with pytest.raises(Exception, match=empty_contig_error): + test_stats_v3_ingestion.read_variant_stats_arrow(regions=[":1-100"]) + with pytest.raises(Exception, match=base_1_error): + test_stats_v3_ingestion.read_variant_stats(regions=["chr1:0-100"]) + with pytest.raises(Exception, match=base_1_error): + test_stats_v3_ingestion.read_variant_stats_arrow(regions=["chr1:0-100"]) + with pytest.raises(Exception, match=interval_error): + test_stats_v3_ingestion.read_variant_stats(regions=["chr1:100-1"]) + with pytest.raises(Exception, match=interval_error): + test_stats_v3_ingestion.read_variant_stats_arrow(regions=["chr1:100-1"]) + + # test types and deprecated region parameter + region1 = "chr1:1-10000" + df = test_stats_v3_ingestion.read_variant_stats(region1) + tbl = test_stats_v3_ingestion.read_variant_stats_arrow(region1) + assert isinstance(df, pd.DataFrame) + assert isinstance(tbl, pa.Table) + assert df.shape == (13, 6) + assert df.equals(tbl.to_pandas()) + df = test_stats_v3_ingestion.read_variant_stats(regions=[region1]) + tbl = test_stats_v3_ingestion.read_variant_stats_arrow(regions=[region1]) + assert isinstance(df, pd.DataFrame) + assert isinstance(tbl, pa.Table) + assert df.shape == (13, 6) + assert df.equals(tbl.to_pandas()) + + # test a region on a different contig + region2 = "chr2:1-10000" + df = test_stats_v3_ingestion.read_variant_stats(regions=[region2]) + tbl = test_stats_v3_ingestion.read_variant_stats_arrow(regions=[region2]) + assert df.shape == (2, 6) + assert df.equals(tbl.to_pandas()) + + # test multiple regions from different contigs and their ordering + regions = [region1, region2] + contigs = ["chr1"] * 13 + ["chr2"] * 2 + df = test_stats_v3_ingestion.read_variant_stats(regions=regions) + assert df.shape == (15, 6) + assert contigs == list(df["contig"].values) + df2 = test_stats_v3_ingestion.read_variant_stats(regions=reversed(regions)) + assert df.equals(df2) + tbl = test_stats_v3_ingestion.read_variant_stats_arrow(regions=regions) + tbl2 = test_stats_v3_ingestion.read_variant_stats_arrow(regions=reversed(regions)) + assert tbl.equals(tbl2) + assert df.equals(tbl.to_pandas()) + assert df2.equals(tbl2.to_pandas()) + + # test overlapping regions on different contigs and their order + region1 = "chr1:1-1" + df = test_stats_v3_ingestion.read_variant_stats(regions=[region1]) + assert df.shape == (2, 6) + region2 = "chr1:1-2" + df = test_stats_v3_ingestion.read_variant_stats(regions=[region2]) + assert df.shape == (5, 6) + region3 = "chr1:3-4" + df = test_stats_v3_ingestion.read_variant_stats(regions=[region3]) + assert df.shape == (6, 6) + region4 = "chr1:2-5" + df = test_stats_v3_ingestion.read_variant_stats(regions=[region4]) + assert df.shape == (11, 6) + regions_chr1 = [region1, region2, region3, region4] + df = test_stats_v3_ingestion.read_variant_stats(regions=regions_chr1) + df2 = test_stats_v3_ingestion.read_variant_stats(regions=reversed(regions_chr1)) + assert df.shape == (13, 6) + assert df.equals(df2) + region5 = "chr2:1-1" + df = test_stats_v3_ingestion.read_variant_stats(regions=[region5]) + assert df.shape == (1, 6) + region6 = "chr2:3-3" + df = test_stats_v3_ingestion.read_variant_stats(regions=[region6]) + assert df.shape == (1, 6) + regions_chr2 = [region5, region6] + df = test_stats_v3_ingestion.read_variant_stats(regions=regions_chr2) + df2 = test_stats_v3_ingestion.read_variant_stats(regions=reversed(regions_chr2)) + assert df.shape == (2, 6) + assert df.equals(df2) + regions = regions_chr1 + regions_chr2 + df = test_stats_v3_ingestion.read_variant_stats(regions=regions) + df2 = test_stats_v3_ingestion.read_variant_stats(regions=reversed(regions)) + assert df.shape == (15, 6) + assert contigs == list(df["contig"].values) + assert df.equals(df2) + regions = regions_chr2 + regions_chr1 + df = test_stats_v3_ingestion.read_variant_stats(regions=regions) + df2 = test_stats_v3_ingestion.read_variant_stats(regions=reversed(regions)) + assert df.shape == (15, 6) + assert contigs == list(df["contig"].values) + assert df.equals(df2) + + # test scan_all_samples + ac = [8, 8, 5, 6, 5, 4, 4, 4, 4, 1, 15, 1, 2, 2, 2] + an = [16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 3, 3, 2, 2] + af = [ + 0.5, + 0.5, + 0.3125, + 0.375, + 0.3125, + 0.25, + 0.25, + 0.25, + 0.25, + 0.0625, + 0.9375, + 0.33333334, + 0.6666667, + 1.0, + 1.0, + ] + df = test_stats_v3_ingestion.read_variant_stats(regions=regions) + assert ac == list(df["ac"].values) + assert an == list(df["an"].values) + assert af == list(df["af"].values) + ac = [8, 8, 5, 6, 5, 4, 4, 4, 4, 1, 15, 1, 2, 2, 2] + an = [16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16, 16] + af = [ + 0.5, + 0.5, + 0.3125, + 0.375, + 0.3125, + 0.25, + 0.25, + 0.25, + 0.25, + 0.0625, + 0.9375, + 0.0625, + 0.125, + 0.125, + 0.125, + ] + df = test_stats_v3_ingestion.read_variant_stats( + regions=regions, + scan_all_samples=True, + ) + assert ac == list(df["ac"].values) + assert an == list(df["an"].values) + assert af == list(df["af"].values) + + # test drop_ref + alleles = [ + "T,C", + "ref", + "G,GTTTA", + "G,T", + "ref", + "C,A", + "C,G", + "C,T", + "ref", + "G,GTTTA", + "ref", + "C,T", + "ref", + "G,GTTTA", + "G,GTTTA", + ] + df = test_stats_v3_ingestion.read_variant_stats(regions=regions) + assert alleles == list(df["alleles"].values) + alleles = [ + "T,C", + "G,GTTTA", + "G,T", + "C,A", + "C,G", + "C,T", + "G,GTTTA", + "C,T", + "G,GTTTA", + "G,GTTTA", + ] + df = test_stats_v3_ingestion.read_variant_stats( + regions=regions, + drop_ref=True, + ) + assert alleles == list(df["alleles"].values) + + ###################### + # read_allele_count # + ###################### + + # test errors + with pytest.raises(Exception, match=no_parameter_error): + test_stats_v3_ingestion.read_allele_count() + with pytest.raises(Exception, match=no_parameter_error): + test_stats_v3_ingestion.read_allele_count_arrow() + with pytest.raises(Exception, match=exclusive_parameter_error): + test_stats_v3_ingestion.read_allele_count("chr1:1-100", regions=["chr1:1-100"]) + with pytest.raises(Exception, match=exclusive_parameter_error): + test_stats_v3_ingestion.read_allele_count_arrow( + "chr1:1-100", regions=["chr1:1-100"] + ) + with pytest.raises(Exception, match=format_error): + test_stats_v3_ingestion.read_allele_count(regions=[""]) + with pytest.raises(Exception, match=format_error): + test_stats_v3_ingestion.read_allele_count_arrow(regions=[""]) + with pytest.raises(Exception, match=format_error): + test_stats_v3_ingestion.read_allele_count(regions=["chr1"]) + with pytest.raises(Exception, match=format_error): + test_stats_v3_ingestion.read_allele_count_arrow(regions=["chr1"]) + with pytest.raises(Exception, match=format_error): + test_stats_v3_ingestion.read_allele_count(regions=["chr1:-"]) + with pytest.raises(Exception, match=format_error): + test_stats_v3_ingestion.read_allele_count_arrow(regions=["chr1:-"]) + with pytest.raises(Exception, match=empty_contig_error): + test_stats_v3_ingestion.read_allele_count(regions=[":1-100"]) + with pytest.raises(Exception, match=empty_contig_error): + test_stats_v3_ingestion.read_allele_count_arrow(regions=[":1-100"]) + with pytest.raises(Exception, match=base_1_error): + test_stats_v3_ingestion.read_allele_count(regions=["chr1:0-100"]) + with pytest.raises(Exception, match=base_1_error): + test_stats_v3_ingestion.read_allele_count_arrow(regions=["chr1:0-100"]) + with pytest.raises(Exception, match=interval_error): + test_stats_v3_ingestion.read_allele_count(regions=["chr1:100-1"]) + with pytest.raises(Exception, match=interval_error): + test_stats_v3_ingestion.read_allele_count_arrow(regions=["chr1:100-1"]) + + # test allele count + + # test types and deprecated region parameter + region1 = "chr1:1-10000" + pos = (0, 1, 1, 2, 2, 2, 3) + count = (8, 5, 3, 4, 2, 2, 1) + df = test_stats_v3_ingestion.read_allele_count(region1) + tbl = test_stats_v3_ingestion.read_allele_count_arrow(region1) + assert isinstance(df, pd.DataFrame) + assert isinstance(tbl, pa.Table) + assert df.shape == (7, 7) + assert df.equals(tbl.to_pandas()) + assert sum(df["pos"] == pos) == 7 + assert sum(df["count"] == count) == 7 + df = test_stats_v3_ingestion.read_allele_count(regions=[region1]) + tbl = test_stats_v3_ingestion.read_allele_count_arrow(regions=[region1]) + assert isinstance(df, pd.DataFrame) + assert isinstance(tbl, pa.Table) + assert df.shape == (7, 7) + assert df.equals(tbl.to_pandas()) + assert sum(df["pos"] == pos) == 7 + assert sum(df["count"] == count) == 7 + + # test a region on a different contig + region2 = "chr2:1-10000" + df = test_stats_v3_ingestion.read_allele_count(regions=[region2]) + tbl = test_stats_v3_ingestion.read_allele_count_arrow(regions=[region2]) + assert df.shape == (2, 7) + assert df.equals(tbl.to_pandas()) + + # test multiple regions from different contigs and their ordering + regions = [region1, region2] + contigs = ["chr1"] * 7 + ["chr2"] * 2 + df = test_stats_v3_ingestion.read_allele_count(regions=regions) + assert df.shape == (9, 7) + assert contigs == list(df["contig"].values) + df2 = test_stats_v3_ingestion.read_allele_count(regions=reversed(regions)) + assert df.equals(df2) + tbl = test_stats_v3_ingestion.read_allele_count_arrow(regions=regions) + tbl2 = test_stats_v3_ingestion.read_allele_count_arrow(regions=reversed(regions)) + assert tbl.equals(tbl2) + assert df.equals(tbl.to_pandas()) + assert df2.equals(tbl2.to_pandas()) + + # test overlapping regions on different contigs and their order + region1 = "chr1:1-1" + df = test_stats_v3_ingestion.read_allele_count(regions=[region1]) + assert df.shape == (1, 7) + region2 = "chr1:1-2" + df = test_stats_v3_ingestion.read_allele_count(regions=[region2]) + assert df.shape == (3, 7) + region3 = "chr1:3-4" + df = test_stats_v3_ingestion.read_allele_count(regions=[region3]) + assert df.shape == (4, 7) + region4 = "chr1:2-5" + df = test_stats_v3_ingestion.read_allele_count(regions=[region4]) + assert df.shape == (6, 7) + regions_chr1 = [region1, region2, region3, region4] + df = test_stats_v3_ingestion.read_allele_count(regions=regions_chr1) + df2 = test_stats_v3_ingestion.read_allele_count(regions=reversed(regions_chr1)) + assert df.shape == (7, 7) + assert df.equals(df2) + region5 = "chr2:1-1" + df = test_stats_v3_ingestion.read_allele_count(regions=[region5]) + assert df.shape == (1, 7) + region6 = "chr2:3-3" + df = test_stats_v3_ingestion.read_allele_count(regions=[region6]) + assert df.shape == (1, 7) + regions_chr2 = [region5, region6] + df = test_stats_v3_ingestion.read_allele_count(regions=regions_chr2) + df2 = test_stats_v3_ingestion.read_allele_count(regions=reversed(regions_chr2)) + assert df.shape == (2, 7) + assert df.equals(df2) + regions = regions_chr1 + regions_chr2 + df = test_stats_v3_ingestion.read_allele_count(regions=regions) + df2 = test_stats_v3_ingestion.read_allele_count(regions=reversed(regions)) + assert df.shape == (9, 7) + assert contigs == list(df["contig"].values) + assert df.equals(df2) + regions = regions_chr2 + regions_chr1 + df = test_stats_v3_ingestion.read_allele_count(regions=regions) + df2 = test_stats_v3_ingestion.read_allele_count(regions=reversed(regions)) + assert df.shape == (9, 7) + assert contigs == list(df["contig"].values) + assert df.equals(df2) + + ######################### + # read_allele_frequency # + ######################### + + region = "chr1:1-10000" df = tiledbvcf.allele_frequency.read_allele_frequency( - os.path.join(tmp_path, "stats_test"), "chr1:1-10000" + os.path.join(tmp_path, "stats_test"), region ) assert df.pos.is_monotonic_increasing df["an_check"] = (df.ac / df.af).round(0).astype("int32") assert df.an_check.equals(df.an) - df = test_stats_v3_ingestion.read_variant_stats("chr1:1-10000") - assert df.shape == (13, 5) - df = df.to_pandas() - df = test_stats_v3_ingestion.read_allele_count("chr1:1-10000") - assert df.shape == (7, 6) - df = df.to_pandas() - assert sum(df["pos"] == (0, 1, 1, 2, 2, 2, 3)) == 7 - assert sum(df["count"] == (8, 5, 3, 4, 2, 2, 1)) == 7 + df = test_stats_v3_ingestion.read_variant_stats(region) + assert df.shape == (13, 6) @pytest.mark.skipif( @@ -1396,7 +1736,7 @@ def test_ingest_with_stats_v2(tmp_path): ) ds = tiledbvcf.Dataset(uri=os.path.join(tmp_path, "stats_test"), mode="r") df = ds.read_variant_stats("chr1:1-10000") - assert df.shape == (13, 5) + assert df.shape == (13, 6) df = tiledbvcf.allele_frequency.read_allele_frequency( os.path.join(tmp_path, "stats_test"), "chr1:1-10000" ) @@ -1404,11 +1744,9 @@ def test_ingest_with_stats_v2(tmp_path): df["an_check"] = (df.ac / df.af).round(0).astype("int32") assert df.an_check.equals(df.an) df = ds.read_variant_stats("chr1:1-10000") - assert df.shape == (13, 5) - df = df.to_pandas() + assert df.shape == (13, 6) df = ds.read_allele_count("chr1:1-10000") - assert df.shape == (7, 6) - df = df.to_pandas() + assert df.shape == (7, 7) assert sum(df["pos"] == (0, 1, 1, 2, 2, 2, 3)) == 7 assert sum(df["count"] == (8, 5, 3, 4, 2, 2, 1)) == 7 diff --git a/libtiledbvcf/src/read/reader.cc b/libtiledbvcf/src/read/reader.cc index af9dfc914..67298988c 100644 --- a/libtiledbvcf/src/read/reader.cc +++ b/libtiledbvcf/src/read/reader.cc @@ -980,7 +980,13 @@ void Reader::read_from_variant_stats( int* ac, int* an, float_t* af) { - af_filter_->retrieve_variant_stats(pos, allele, allele_offsets, ac, an, af); + if (params_.scan_all_samples) { + size_t n = dataset_->sample_names().size(); + af_filter_->retrieve_variant_stats( + n, pos, allele, allele_offsets, ac, an, af); + } else { + af_filter_->retrieve_variant_stats(pos, allele, allele_offsets, ac, an, af); + } } void Reader::read_from_allele_count( diff --git a/libtiledbvcf/src/stats/allele_count.cc b/libtiledbvcf/src/stats/allele_count.cc index 70f6989d5..135f368b6 100644 --- a/libtiledbvcf/src/stats/allele_count.cc +++ b/libtiledbvcf/src/stats/allele_count.cc @@ -532,19 +532,8 @@ AlleleCountReader::allele_count_buffer_sizes() { } void AlleleCountReader::prepare_allele_count(Region region) { - // TODO: fix this ALLELE_COUNT_ARRAY so it links against the static singleton - // defined in allele_count.h - const std::string ALLELE_COUNT_ARRAY = "allele_count"; - - // TODO: fix this COLUMN_NAME so it links against the static singleton defined - // in allele_count.h - std::vector COLUMN_NAME = { - "pos", "ref", "alt", "filter", "gt", "count"}; - // TODO: likewise with this enum - enum Columns { POS, REF, ALT, FILTER, GT, COUNT }; - - ManagedQuery mq(array_, ALLELE_COUNT_ARRAY, TILEDB_UNORDERED); - mq.select_columns(COLUMN_NAME); + ManagedQuery mq(array_, AlleleCount::ALLELE_COUNT_ARRAY, TILEDB_UNORDERED); + mq.select_columns(AlleleCount::COLUMN_NAME); mq.select_point("contig", region.seq_name); mq.select_ranges("pos", {{region.min, region.max}}); // TODO: use regions set in reader to assign subarray diff --git a/libtiledbvcf/src/stats/allele_count.h b/libtiledbvcf/src/stats/allele_count.h index d009328ad..51f70d2b4 100644 --- a/libtiledbvcf/src/stats/allele_count.h +++ b/libtiledbvcf/src/stats/allele_count.h @@ -188,6 +188,9 @@ class AlleleCount { void flush(bool finalize = false); private: + // give the AlleleCountReader access to AlleleCount's private members + friend class AlleleCountReader; + //=================================================================== //= private static //=================================================================== diff --git a/libtiledbvcf/src/stats/variant_stats_reader.cc b/libtiledbvcf/src/stats/variant_stats_reader.cc index f576c79b0..6e23c9d4f 100644 --- a/libtiledbvcf/src/stats/variant_stats_reader.cc +++ b/libtiledbvcf/src/stats/variant_stats_reader.cc @@ -96,144 +96,6 @@ inline void AFMap::advance_to_ref_block(uint32_t pos) { } } -std::tuple AFMap::af( - uint32_t pos, const std::string& allele) { - // calculate af = ac / an - std::unordered_map< - uint32_t, - std::pair>>::const_iterator - pos_map_iterator = ac_map_.find(pos); - - // Return -1.0 if the allele was not called - if (pos_map_iterator == ac_map_.end()) { - return {-1.0, 0, 0}; - } - - const auto pos_map = pos_map_iterator->second; - - // We don't know that allele was called in this sample. Ask nicely for the - // allele count. - auto next_allele = pos_map.second.find(allele); - - // First multiply by 1.0 to force a float type, then look up AC from the - // above iterator. Substitute 0 for the AC value if it is absent in pos_map. - // Divide by AN (pos_map.first) to get AF. - if (next_allele == pos_map.second.end()) { - return {0, 0, pos_map.first}; - } - return { - 1.0 * next_allele->second / pos_map.first, - next_allele->second, - pos_map.first}; -} - -std::tuple AFMap::af( - uint32_t pos, const std::string& allele, size_t num_samples) { - // calculate af = ac / an - std::unordered_map< - uint32_t, - std::pair>>::const_iterator - pos_map_iterator = ac_map_.find(pos); - - // Return -1.0 if the allele was not called - if (pos_map_iterator == ac_map_.end()) { - return {-1.0, 0, 0}; - } - - const std::pair>& pos_map = - pos_map_iterator->second; - - // We don't know that allele was called in this sample. Ask nicely for the - // allele count. - auto next_allele = pos_map.second.find(allele); - - // First multiply by 1.0 to force a float type, then look up AC from the - // above iterator. Substitute 0 for the AC value if it is absent in pos_map. - // Divide by AN (pos_map.first) to get AF. - if (next_allele == pos_map.second.end()) { - return {0, 0, num_samples * 2}; - } - return { - 1.0 * next_allele->second / num_samples / 2, - next_allele->second, - num_samples * 2}; -} - -std::tuple AFMap::af_v3( - uint32_t pos, const std::string& allele) { - // calculate af = ac / an - std::unordered_map< - uint32_t, - std::pair>>::const_iterator - pos_map_iterator = ac_map_.find(pos); - - // Return -1.0 if the allele was not called - if (pos_map_iterator == ac_map_.end()) { - return {-1.0, 0, 0}; - } - - const auto pos_map = pos_map_iterator->second; - - // We don't know that allele was called in this sample. Ask nicely for the - // allele count. - auto next_allele = pos_map.second.find(allele); - bool is_ref = allele == "ref"; - - // First multiply by 1.0 to force a float type, then look up AC from the - // above iterator. Substitute 0 for the AC value if it is absent in pos_map. - // Divide by AN (pos_map.first) to get AF. - if (next_allele == pos_map.second.end()) { - return { - is_ref ? 1.0 * ac_sum_ / (pos_map.first + an_sum_) : 0, - is_ref ? ac_sum_ : 0, - pos_map.first + an_sum_}; - } - return { - 1.0 * (next_allele->second + (is_ref ? ac_sum_ : 0)) / - (pos_map.first + an_sum_), - next_allele->second + (is_ref ? ac_sum_ : 0), - (pos_map.first + an_sum_)}; -} - -std::tuple AFMap::af_v3( - uint32_t pos, const std::string& allele, size_t num_samples) { - // calculate af = ac / an - std::unordered_map< - uint32_t, - std::pair>>::const_iterator - pos_map_iterator = ac_map_.find(pos); - - // Return -1.0 if the allele was not called - if (pos_map_iterator == ac_map_.end()) { - return {-1.0, 0, 0}; - } - - const std::pair>& pos_map = - pos_map_iterator->second; - - // We don't know that allele was called in this sample. Ask nicely for the - // allele count. - decltype(pos_map.second)::const_iterator next_allele = - pos_map.second.find(allele); - - bool is_ref = allele == "ref"; - - // First multiply by 1.0 to force a float type, then look up AC from the - // above iterator. Substitute 0 for the AC value if it is absent in pos_map. - // Divide by AN (pos_map.first) to get AF. - if (next_allele == pos_map.second.end()) { - return { - is_ref ? 1.0 * ac_sum_ / (num_samples * 2 + an_sum_) : 0, - is_ref ? ac_sum_ : 0, - num_samples / 2 + an_sum_}; - } - return { - 1.0 * (next_allele->second + (is_ref ? ac_sum_ : 0)) / - (num_samples * 2 + an_sum_), - next_allele->second + (is_ref ? ac_sum_ : 0), - num_samples * 2 + an_sum_}; -} - void AFMap::finalize_ref_block_cache() { std::sort(ref_block_cache_.begin(), ref_block_cache_.end(), RefBlockComp()); ref_block_by_end_.clear(); @@ -245,72 +107,46 @@ void AFMap::finalize_ref_block_cache() { selected_ref_block_end_ = ref_block_by_end_.begin(); } -// The following two methods are declared to eliminate a branchpoint in a loop -// where they're called: instead of wrapping the AF version in an if statement, -// use indirection or generate per-version code paths for the loop. -static inline std::tuple call_af_v2( - AFMap* map, uint32_t& pos, const std::string& allele) { - return map->af(pos, allele); -} - -static inline std::tuple call_af_v3( - AFMap* map, uint32_t& pos, const std::string& allele) { - map->advance_to_ref_block(pos); - return map->af_v3(pos, allele); -} - inline void AFMap::retrieve_variant_stats( - uint32_t* pos, - char* allele, + const AFComputer& af_computer, + uint32_t* pos_buffer, + char* allele_buffer, int32_t* allele_offsets, - int* ac, - int* an, - float_t* afb) { - std::tuple (&call_af)( - AFMap*, uint32_t&, const std::string&) = - array_version > 2 ? call_af_v3 : call_af_v2; - { - size_t row = 0; - for (std::pair< - uint32_t, - std::pair>> - position_to_allele : ac_map_) { - for (size_t remaining_to_insert = position_to_allele.second.second.size(); - remaining_to_insert > 0; - remaining_to_insert--) { - pos[row] = position_to_allele.first; - row++; - if (row > allele_cardinality_) { - throw std::runtime_error( - "[VariantStatsReader] export buffer size computed incorrectly: " - "too low"); - } + int* ac_buffer, + int* an_buffer, + float_t* af_buffer) { + size_t row = 0; + for (const auto& [pos, an_ac] : stats_map_) { + for (size_t remaining_to_insert = an_ac.second.size(); + remaining_to_insert > 0; + remaining_to_insert--) { + pos_buffer[row] = pos; + row++; + if (row > allele_cardinality_) { + throw std::runtime_error( + "[VariantStatsReader] export buffer size computed incorrectly: " + "too low"); } } - if (row < allele_cardinality_) { - throw std::runtime_error( - "[VariantStatsReader] export buffer size computed incorrectly: too " - "high"); - } + } + if (row < allele_cardinality_) { + throw std::runtime_error( + "[VariantStatsReader] export buffer size computed incorrectly: too " + "high"); } - qsort(pos, allele_cardinality_, sizeof(uint32_t), pos_comparator); + qsort(pos_buffer, allele_cardinality_, sizeof(uint32_t), pos_comparator); allele_offsets[0] = 0; - for (size_t row = 0; row < allele_cardinality_;) { - std::pair> an_to_allele = - ac_map_[pos[row]]; - for (std::pair allele_to_count : an_to_allele.second) { - auto [this_af, this_ac, this_an] = - call_af(this, pos[row], allele_to_count.first); - ac[row] = this_ac; - an[row] = this_an; - afb[row] = this_af; + for (row = 0; row < allele_cardinality_;) { + ac_map_t& ac_map = stats_map_[pos_buffer[row]].second; + for (auto& [allele, _] : ac_map) { + auto [af, ac, an] = af_computer.af(pos_buffer[row], allele); + ac_buffer[row] = ac; + an_buffer[row] = an; + af_buffer[row] = af; std::memcpy( - allele + allele_offsets[row], - allele_to_count.first.c_str(), - allele_to_count.first.size()); - allele_offsets[row + 1] = - allele_offsets[row] + allele_to_count.first.size(); + allele_buffer + allele_offsets[row], allele.c_str(), allele.size()); + allele_offsets[row + 1] = allele_offsets[row] + allele.size(); row++; } } @@ -362,19 +198,52 @@ uint32_t VariantStatsReader::array_version() { } void VariantStatsReader::retrieve_variant_stats( - uint32_t* pos, - char* allele, + uint32_t* pos_buffer, + char* allele_buffer, int32_t* allele_offsets, - int* ac, - int* an, - float_t* af) { + int* ac_buffer, + int* an_buffer, + float_t* af_buffer) { // there is no thread-safe implementation of this yet if (async_query_) { throw std::runtime_error( "[VariantStatsReader] can not retrieve variant stats when async quries " "are enabled"); } - af_map_.retrieve_variant_stats(pos, allele, allele_offsets, ac, an, af); + AFMap::AFComputerSingle af_computer(&af_map_); + af_map_.retrieve_variant_stats( + af_computer, + pos_buffer, + allele_buffer, + allele_offsets, + ac_buffer, + an_buffer, + af_buffer); +} + +void VariantStatsReader::retrieve_variant_stats( + const size_t num_samples, + uint32_t* pos_buffer, + char* allele_buffer, + int32_t* allele_offsets, + int* ac_buffer, + int* an_buffer, + float_t* af_buffer) { + // there is no thread-safe implementation of this yet + if (async_query_) { + throw std::runtime_error( + "[VariantStatsReader] can not retrieve variant stats when async quries " + "are enabled"); + } + AFMap::AFComputerAll af_computer(&af_map_, num_samples); + af_map_.retrieve_variant_stats( + af_computer, + pos_buffer, + allele_buffer, + allele_offsets, + ac_buffer, + an_buffer, + af_buffer); } void VariantStatsReader::compute_af() { @@ -416,12 +285,13 @@ std::tuple VariantStatsReader::pass( // TODO: replace placeholder return values if necessary return {false, 0.0, 0, 0}; } - auto [af, ac, an] = - (array_version() > 2) ? - (scan_all_samples ? af_map_.af_v3(pos, allele, num_samples) : - af_map_.af_v3(pos, allele)) : - (scan_all_samples ? af_map_.af(pos, allele, num_samples) : - af_map_.af(pos, allele)); + std::unique_ptr af_computer; + if (scan_all_samples) { + af_computer = std::make_unique(&af_map_, num_samples); + } else { + af_computer = std::make_unique(&af_map_); + } + auto [af, ac, an] = af_computer->af(pos, allele); // Fail the filter if allele was not called if (af < 0.0) { diff --git a/libtiledbvcf/src/stats/variant_stats_reader.h b/libtiledbvcf/src/stats/variant_stats_reader.h index 1a272295d..3e9ab910b 100644 --- a/libtiledbvcf/src/stats/variant_stats_reader.h +++ b/libtiledbvcf/src/stats/variant_stats_reader.h @@ -41,10 +41,193 @@ namespace tiledb::vcf { /** * @brief A class to store allele counts and compute allele frequency. - * */ class AFMap { public: + // allele frequency type: + typedef std::tuple af_t; + + /** + * An abstract class that defines the interface for computing allele frequency + * stats. + */ + class AFComputer { + public: + /** + * @param map The AFMap allele frequencies are being computed for + */ + AFComputer(const AFMap* map) + : af_map_(map) { + // use a function pointer to avoid a branch point when computing AF + if (map->array_version <= 2) + af_ptr = &AFComputer::af_v2; + else + af_ptr = &AFComputer::af_v3; + } + + virtual ~AFComputer() = default; + + /** + * Computes the allele frequency for an allele at the given position. + * + * @param pos Position of the allele + * @param allele Allele value + * @return Allele frequency + */ + af_t af(const uint32_t pos, const std::string& allele) const { + return (this->*af_ptr)(pos, allele); + } + + protected: + const AFMap* af_map_; + + private: + /** A pointer to the member function actually used to compute allele + * frequency. */ + af_t (AFComputer::*af_ptr)( + const uint32_t pos, const std::string& allele) const; + + /** Computes the allele frequency for version 1 and 2 stats arrays. */ + virtual af_t af_v2(const uint32_t pos, const std::string& allele) const = 0; + + /** Computes the allele frequency for version 3+ stats arrays. */ + virtual af_t af_v3(const uint32_t pos, const std::string& allele) const = 0; + }; + + /** + * A class for computing allele frequency stats for single alleles. + */ + class AFComputerSingle : public AFComputer { + public: + /** + * @param map The AFMap allele frequencies are being computed for + */ + AFComputerSingle(const AFMap* map) + : AFComputer(map) { + } + + private: + inline float compute_af_v3(const int ac, const int an) const { + return 1.0 * ac / an; + } + + /** Computes the allele frequency for version 1 and 2 stats arrays. */ + af_t af_v2(const uint32_t pos, const std::string& allele) const { + // Return -1.0 if the allele was not called + pos_stats_map_t::const_iterator pos_map_itr = + af_map_->stats_map_.find(pos); + if (pos_map_itr == af_map_->stats_map_.end()) { + return {-1.0, 0, 0}; + } + // Get the allele count + const auto& [allele_an, ac_map] = pos_map_itr->second; + const ac_map_t::const_iterator& ac_itr = ac_map.find(allele); + // Compute AF, substituting 0 for the AF and AC values if allele not in + // pos's an_ac + if (ac_itr == ac_map.end()) { + return {0, 0, allele_an}; + } + const int allele_ac = ac_itr->second; + const float af = 1.0 * allele_ac / allele_an; + return {af, allele_ac, allele_an}; + } + + /** Computes the allele frequency for version 3+ stats arrays. */ + af_t af_v3(const uint32_t pos, const std::string& allele) const { + // Return -1.0 if the allele was not called + pos_stats_map_t::const_iterator pos_map_itr = + af_map_->stats_map_.find(pos); + if (pos_map_itr == af_map_->stats_map_.end()) { + return {-1.0, 0, 0}; + } + // Get the allele count + const auto& [allele_an, ac_map] = pos_map_itr->second; + const ac_map_t::const_iterator& ac_itr = ac_map.find(allele); + bool is_ref = allele == "ref"; + // Compute AF + uint32_t ac = is_ref ? af_map_->ac_sum_ : 0; + uint32_t an = allele_an + af_map_->an_sum_; + if (ac_itr == ac_map.end()) { + float af = is_ref ? compute_af_v3(af_map_->ac_sum_, an) : 0; + uint32_t unmapped_an = allele_an + af_map_->an_sum_; + return {af, ac, unmapped_an}; + } + const int allele_ac = ac + ac_itr->second; + float af = compute_af_v3(allele_ac, an); + return {af, allele_ac, an}; + } + }; + + /** + * A class for computing allelel frequencies stats relativekto all samples in + * the dataset. + */ + class AFComputerAll : public AFComputer { + public: + /** + * @param map The AFMap allele frequencies are being computed for + * @param n The total number of samples in the dataset + */ + AFComputerAll(const AFMap* map, const size_t n) + : AFComputer(map) + , num_samples(n) { + } + + private: + const size_t num_samples; + + inline float compute_af_v3(const int ac, const int an) const { + return 1.0 * ac / an; + } + + /** Computes the allele frequency for version 1 and 2 stats arrays. */ + af_t af_v2(const uint32_t pos, const std::string& allele) const { + // Return -1.0 if the allele was not called + pos_stats_map_t::const_iterator pos_map_itr = + af_map_->stats_map_.find(pos); + if (pos_map_itr == af_map_->stats_map_.end()) { + return {-1.0, 0, 0}; + } + // Get the allele count + const auto& [allele_an, ac_map] = pos_map_itr->second; + const ac_map_t::const_iterator& ac_itr = ac_map.find(allele); + // Compute AF, substituting 0 for the AF and AC values if allele not in + // pos's an_ac + const uint32_t an = num_samples * 2; + if (ac_itr == ac_map.end()) { + return {0, 0, an}; + } + const int allele_ac = ac_itr->second; + const float af = 1.0 * allele_ac / num_samples / 2; + return {af, allele_ac, an}; + } + + /** Computes the allele frequency for version 3+ stats arrays. */ + af_t af_v3(const uint32_t pos, const std::string& allele) const { + // Return -1.0 if the allele was not called + pos_stats_map_t::const_iterator pos_map_itr = + af_map_->stats_map_.find(pos); + if (pos_map_itr == af_map_->stats_map_.end()) { + return {-1.0, 0, 0}; + } + // Get the allele count + const auto& [allele_an, ac_map] = pos_map_itr->second; + const ac_map_t::const_iterator& ac_itr = ac_map.find(allele); + bool is_ref = allele == "ref"; + // Compute AF + uint32_t ac = is_ref ? af_map_->ac_sum_ : 0; + uint32_t an = num_samples * 2 + af_map_->an_sum_; + if (ac_itr == ac_map.end()) { + float af = is_ref ? compute_af_v3(af_map_->ac_sum_, an) : 0; + uint32_t unmapped_an = num_samples / 2 + af_map_->an_sum_; + return {af, ac, unmapped_an}; + } + const int allele_ac = ac + ac_itr->second; + float af = compute_af_v3(allele_ac, an); + return {af, allele_ac, an}; + } + }; + /** * @brief Insert an Allele Count into the map. * @@ -65,13 +248,13 @@ class AFMap { if (pos >= min_pos) { // Should this insertion be tallied for contributing to transport // (Arrow) buffer size? - bool should_tally = - !ac_map_.contains(pos) || !ac_map_[pos].second.contains(allele); + bool should_tally = !stats_map_.contains(pos) || + !stats_map_[pos].second.contains(allele); // add ac to the an for this position - ac_map_[pos].first += ac; + stats_map_[pos].first += ac; // add ac to the ac for this position, allele - ac_map_[pos].second[allele] += ac; + stats_map_[pos].second[allele] += ac; if (false) { LOG_TRACE( @@ -79,8 +262,8 @@ class AFMap { pos, allele, ac, - ac_map_[pos].second[allele], - ac_map_[pos].first); + stats_map_[pos].second[allele], + stats_map_[pos].first); } if (should_tally) { allele_aggregate_size_ += allele.length(); @@ -114,56 +297,12 @@ class AFMap { return std::make_tuple(allele_cardinality_, allele_aggregate_size_); } - /** - * @brief Compute the Allele Frequency for an allele at the given position. - * - * @param pos Position of the allele - * @param allele Allele value - * @return float Allele Frequency - */ - inline std::tuple af( - uint32_t pos, const std::string& allele); - - /** - * @brief Compute the Allele Frequency for an allele at the given position, - * accounting for the full population of the dataset - * - * @param pos Position of the allele - * @param allele Allele value - * @param num_samples Number of samples in the entire dataset - * @return float Allele Frequency - */ - inline std::tuple af( - uint32_t pos, const std::string& allele, size_t num_samples); - - /** - * @brief Compute the Allele Frequency for an allele at the given position. - * - * @param pos Position of the allele - * @param allele Allele value - * @return float Allele Frequency - */ - inline std::tuple af_v3( - uint32_t pos, const std::string& allele); - - /** - * @brief Compute the Allele Frequency for an allele at the given position, - * accounting for the full population of the dataset - * - * @param pos Position of the allele - * @param allele Allele value - * @param num_samples Number of samples in the entire dataset - * @return float Allele Frequency - */ - inline std::tuple af_v3( - uint32_t pos, const std::string& allele, size_t num_samples); - /** * @brief Clear the map. * */ void clear() { - ac_map_.clear(); + stats_map_.clear(); ref_block_cache_.clear(); ac_sum_ = 0; an_sum_ = 0; @@ -171,22 +310,25 @@ class AFMap { } /** - * @brief Populate buffers + * @brief Populates variant stats buffers using the given AFComputer to + * compute allele frequencies * - * @param pos buffer of int representing position - * @param allele variable buffer of char representing allele - * @param allele_offsets allele offsets - * @param ac buffer of int representing allele count - * @param an buffer of int representing allele number - * @param af buffer of float representing internal allele frequency + * @param af_computer The AFComputer to use when computing allele frequencies + * @param pos_buffer buffer of int representing position + * @param allele_buffer variable buffer of char representing allele + * @param allele_offsets buffer of offsets for allele_buffer + * @param ac_buffer buffer of int representing allele count + * @param an_buffer buffer of int representing allele number + * @param af_buffer buffer of float representing internal allele frequency */ void retrieve_variant_stats( - uint32_t* pos, - char* allele, + const AFComputer& af_computer, + uint32_t* pos_buffer, + char* allele_buffer, int32_t* allele_offsets, - int* ac, - int* an, - float_t* af); + int* ac_buffer, + int* an_buffer, + float_t* af_buffer); uint32_t array_version = 2; @@ -194,6 +336,11 @@ class AFMap { uint32_t min_pos = 0; private: + /** position stats map: pos -> (an, map: allele -> ac) */ + typedef std::unordered_map ac_map_t; + typedef std::pair an_ac_t; + typedef std::unordered_map pos_stats_map_t; + struct RefBlock { uint32_t start; uint32_t end; @@ -202,30 +349,23 @@ class AFMap { }; /** - * @brief The RefBlockComp class serves as a comparator for RefBlocks to be - * sorted. - * + * The RefBlockComp class serves as a comparator for RefBlocks to be sorted. */ class RefBlockComp { public: /** - * @brief sort ref blocks directly in the cache, ascending by start - * + * Sort ref blocks directly in the cache, ascending by start */ bool operator()(const RefBlock& a, const RefBlock& b) const; /** - * @brief sort ref block pointers to the cache, ascending by end - * + * Sort ref block pointers to the cache, ascending by end */ bool operator()(const RefBlock* a, const RefBlock* b) const; }; - /** Allele Count map: pos -> (an, map: allele -> ac) */ - std::unordered_map< - uint32_t, - std::pair>> - ac_map_; + /** allele Count map */ + pos_stats_map_t stats_map_; /** track running sum of AC when computing IAF for GVCF */ uint64_t ac_sum_ = 0; @@ -303,22 +443,44 @@ class VariantStatsReader { } /** - * @brief Populate buffers + * @brief Populates variant stats buffers + * + * @param num_samples Number of samples in the entire dataset + * @param pos_buffer buffer of int representing position + * @param allele_buffer variable buffer of char representing allele + * @param allele_offsets buffer of offsets for allele_buffer + * @param ac_buffer buffer of int representing allele count + * @param an_buffer buffer of int representing allele number + * @param af_buffer buffer of float representing internal allele frequency + */ + void retrieve_variant_stats( + uint32_t* pos_buffer, + char* allele_buffer, + int32_t* allele_offsets, + int* ac_buffer, + int* an_buffer, + float_t* af_buffer); + + /** + * @brief Scans all samples when computing internal allele frequency and + * populates variant stats buffers * - * @param pos buffer of int representing position - * @param allele variable buffer of char representing allele - * @param allele_offsets allele offsets - * @param ac buffer of int representing allele count - * @param an buffer of int representing allele number - * @param af buffer of float representing internal allele frequency + * @param num_samples Number of samples in the entire dataset + * @param pos_buffer buffer of int representing position + * @param allele_buffer variable buffer of char representing allele + * @param allele_offsets buffer of offsets for allele_buffer + * @param ac_buffer buffer of int representing allele count + * @param an_buffer buffer of int representing allele number + * @param af_buffer buffer of float representing internal allele frequency */ void retrieve_variant_stats( - uint32_t* pos, - char* allele, + const size_t num_samples, + uint32_t* pos_buffer, + char* allele_buffer, int32_t* allele_offsets, - int* ac, - int* an, - float_t* af); + int* ac_buffer, + int* an_buffer, + float_t* af_buffer); /** * @brief Accessor for buffer size metrics when retrieving variant stats @@ -362,6 +524,9 @@ class VariantStatsReader { * * @param pos Position of the allele * @param allele Allele value + * @param scan_all_samples Scan all samples when computing internal allele + * frequency + * @param num_samples Number of samples in the entire dataset * @return std::pair (Allele passed the filter, AF value) */ std::tuple pass(