Update to C++ accelerated IAF computation (#538)

* add dummy condition when IAF field is detected

* add option for dataset-wide cohort for IAF computation

* actually compute IAF from full dataset sample count when requested

* add test for full sample scan

* count 'ref' in variant stats array

* examples appear to work

* disable async variant stats reader and fix race

* unify ref,alt representation

* check stats table version

* make clumsy types more parsimonious

* clarify logic, flip default async value

* update stats check to require a minimum version, update error message

* temporary Quarto rendering fix

---------

Co-authored-by: George Powley <george.powley@gmail.com>

Author: awenocur

.github/workflows/quarto-render.yml

@@ -17,7 +17,7 @@ jobs:
     - name: "Quarto render"
       shell: bash
       run: |
-        pip install quartodoc
+        pip install quartodoc "pydantic<2"
         # create a symlink to the tiledbvcf python package, so it doesn't have to be installed
         ln -s apis/python/src/tiledbvcf
         quartodoc build

apis/python/src/tiledbvcf/binding/libtiledbvcf.cc

@@ -39,6 +39,7 @@ PYBIND11_MODULE(libtiledbvcf, m) {
       .def("set_output_path", &Reader::set_output_path)
       .def("set_output_dir", &Reader::set_output_dir)
       .def("set_af_filter", &Reader::set_af_filter)
+      .def("set_scan_all_samples", &Reader::set_scan_all_samples)
       .def("read", &Reader::read, py::arg("release_buffs") = true)
       .def("get_results_arrow", &Reader::get_results_arrow)
       .def("completed", &Reader::completed)

apis/python/src/tiledbvcf/binding/reader.cc

@@ -201,12 +201,26 @@ void Reader::set_af_filter(const std::string& af_filter) {
       reader, tiledb_vcf_reader_set_af_filter(reader, af_filter.c_str()));
 }
 
+void Reader::set_scan_all_samples(const bool scan_all_samples) {
+  auto reader = ptr.get();
+  check_error(
+      reader, tiledb_vcf_reader_set_scan_all_samples(reader, scan_all_samples));
+}
+
 void Reader::read(const bool release_buffs) {
   auto reader = ptr.get();
   bool af_filter_enabled = false;
   if (tiledb_vcf_reader_get_af_filter_exists(reader, &af_filter_enabled) ==
       TILEDB_VCF_ERR)
     throw std::runtime_error("TileDB-VCF-Py: Error finding AF filter.");
+  if (!af_filter_enabled) {
+    for (std::string attribute : attributes_) {
+      if (attribute == "info_TILEDB_IAF") {
+        tiledb_vcf_reader_set_af_filter(reader, ">=0");
+        af_filter_enabled = true;
+      }
+    }
+  }
   if (af_filter_enabled) {
     // add alleles buffer only if not already present
     bool add_alleles = true;

apis/python/src/tiledbvcf/binding/reader.h

@@ -154,6 +154,9 @@ class Reader {
   /** Set internal allele frequency filtering expression */
   void set_af_filter(const std::string& af_filter);
 
+  /** Set whether to scan all samples in the dataset when computing frequency */
+  void set_scan_all_samples(const bool scan_all_samples);
+
   /** Set the TileDB query buffer memory percentage */
   void set_buffer_percentage(float buffer_percentage);
 

apis/python/src/tiledbvcf/dataset.py

@@ -212,6 +212,7 @@ def read_arrow(
         bed_file: str = None,
         skip_check_samples: bool = False,
         set_af_filter: str = "",
+        scan_all_samples: bool = False,
         enable_progress_estimation: bool = False,
     ) -> pa.Table:
         """
@@ -239,6 +240,8 @@ def read_arrow(
         set_af_filter
             Filter variants by internal allele frequency. For example, to include
             variants with AF > 0.1, set this to ">0.1".
+        scan_all_samples
+            Scan all samples when computing internal allele frequency.
         enable_progress_estimation
             **DEPRECATED** - This parameter will be removed in a future release.
 
@@ -264,6 +267,7 @@ def read_arrow(
         self.reader.set_attributes(attrs)
         self.reader.set_check_samples_exist(not skip_check_samples)
         self.reader.set_af_filter(set_af_filter)
+        self.reader.set_scan_all_samples(scan_all_samples)
         self.reader.set_enable_progress_estimation(enable_progress_estimation)
 
         if bed_file is not None:
@@ -280,6 +284,7 @@ def read(
         bed_file: str = None,
         skip_check_samples: bool = False,
         set_af_filter: str = "",
+        scan_all_samples: bool = False,
         enable_progress_estimation: bool = False,
     ) -> pd.DataFrame:
         """
@@ -334,6 +339,7 @@ def read(
         self.reader.set_attributes(attrs)
         self.reader.set_check_samples_exist(not skip_check_samples)
         self.reader.set_af_filter(set_af_filter)
+        self.reader.set_scan_all_samples(scan_all_samples)
         self.reader.set_enable_progress_estimation(enable_progress_estimation)
 
         if bed_file is not None:
@@ -369,6 +375,11 @@ def export(
             URI of a BED file of genomic regions to be read.
         skip_check_samples
             Skip checking if the samples in `samples_file` exist in the dataset.
+        set_af_filter
+            Filter variants by internal allele frequency. For example, to include
+            variants with AF > 0.1, set this to ">0.1".
+        scan_all_samples
+            Scan all samples when computing internal allele frequency.
         enable_progress_estimation
             **DEPRECATED** - This parameter will be removed in a future release.
         merge

apis/python/tests/test_tiledbvcf.py

@@ -1168,6 +1168,17 @@ def test_ingest_with_stats(tmp_path):
         data_frame[data_frame["sample_name"] == "second"]["info_TILEDB_IAF"].iloc[0][0]
         == 0.0625
     )
+    data_frame = ds.read(
+        samples=sample_names,
+        attrs=["contig", "pos_start", "id", "qual", "info_TILEDB_IAF", "sample_name"],
+        scan_all_samples=True,
+    )
+    assert (
+        data_frame[
+            (data_frame["sample_name"] == "second") & (data_frame["pos_start"] == 4)
+        ]["info_TILEDB_IAF"].iloc[0][0]
+        == 0.0625
+    )
 
 
 def test_ingest_mode_separate(tmp_path):

libtiledbvcf/src/c_api/tiledbvcf.cc

@@ -942,6 +942,18 @@ int32_t tiledb_vcf_reader_set_af_filter(
   return TILEDB_VCF_OK;
 }
 
+int32_t tiledb_vcf_reader_set_scan_all_samples(
+    tiledb_vcf_reader_t* reader, bool scan_all_samples) {
+  if (sanity_check(reader) == TILEDB_VCF_ERR)
+    return TILEDB_VCF_ERR;
+
+  if (SAVE_ERROR_CATCH(
+          reader, reader->reader_->set_scan_all_samples(scan_all_samples)))
+    return TILEDB_VCF_ERR;
+
+  return TILEDB_VCF_OK;
+}
+
 /* ********************************* */
 /*             BED FILE              */
 /* ********************************* */

libtiledbvcf/src/c_api/tiledbvcf.h

@@ -918,6 +918,13 @@ TILEDBVCF_EXPORT int32_t tiledb_vcf_reader_set_output_path(
 TILEDBVCF_EXPORT int32_t tiledb_vcf_reader_set_af_filter(
     tiledb_vcf_reader_t* reader, const char* af_filter);
 
+/**
+ * Sets whether to scan all samples for IAF computation
+ *
+ */
+TILEDBVCF_EXPORT int32_t tiledb_vcf_reader_set_scan_all_samples(
+    tiledb_vcf_reader_t* reader, bool scan_all_samples);
+
 /**
  * Sets export output directory
  * @param reader VCF reader object

libtiledbvcf/src/read/reader.cc

@@ -28,6 +28,7 @@
 #include <iomanip>
 #include <random>
 #include <thread>
+#include <vector>
 
 #include "dataset/attribute_buffer_set.h"
 #include "read/bcf_exporter.h"
@@ -1282,8 +1283,17 @@ bool Reader::process_query_results_v4() {
 
       read_state_.query_results.af_values.clear();
       int allele_index = 0;
+      size_t num_samples = 0;
+      if (params_.scan_all_samples) {
+        num_samples = dataset_->sample_names().size();
+      }
+      bool is_ref = true;
       for (auto&& allele : alleles) {
-        auto [allele_passes, af] = af_filter_->pass(real_start, allele);
+        auto [allele_passes, af] = af_filter_->pass(
+            real_start,
+            is_ref ? "ref" : alleles[0] + "," + allele,
+            params_.scan_all_samples,
+            num_samples);
 
         // If the allele is in GT, consider it in the pass computation
         // TODO: when supporting greater than diploid organisms, expand the
@@ -1303,6 +1313,7 @@ bool Reader::process_query_results_v4() {
         read_state_.query_results.af_values.push_back(af);
 
         LOG_TRACE("  pass = {}", pass);
+        is_ref = false;
       }
 
       // If all alleles do not pass the af filter, continue
@@ -2562,6 +2573,10 @@ void Reader::set_af_filter(const std::string& af_filter) {
   }
 }
 
+void Reader::set_scan_all_samples(bool scan_all_samples) {
+  params_.scan_all_samples = scan_all_samples;
+}
+
 void Reader::set_tiledb_query_config() {
   assert(read_state_.query != nullptr);
   assert(buffers_a != nullptr);

libtiledbvcf/src/read/reader.h

@@ -139,6 +139,9 @@ struct ExportParams {
   //   and VALUE = float
   // If empty, AF filtering is not applied
   std::string af_filter = "";
+
+  // Should all samples be scanned when computing internal allele frequency?
+  bool scan_all_samples = false;
 };
 
 /* ********************************* */
@@ -443,6 +446,12 @@ class Reader {
    */
   void set_af_filter(const std::string& af_filter);
 
+  /**
+   * sets whether to scan all samples in the dataset when computing IAF
+   * @param af_filter setting
+   */
+  void set_scan_all_samples(bool scan_all_samples);
+
   /**
    * Sets disabling of progress estimation in verbose mode
    * @param enable_progress_estimation setting