add function .variants()

Author: LynxJinyangii

RcppTskit/NEWS.md

@@ -40,6 +40,12 @@ and releases adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html
 - Added `rtsk_mutation_table_add_row()` and
   `TableCollection$mutation_table_add_row()` to append mutation rows from
   \code{R}, mirroring `tsk_mutation_table_add_row()`.
+- Added low-level variant iterators
+  (`rtsk_variant_iterator_init()`/`rtsk_variant_iterator_next()`) and a
+  user-facing `TreeSequence$variants()` method to iterate over decoded
+  site-by-site variants from \code{R}, aligned with `tskit` Python API
+  semantics for `samples`, `isolated_as_missing`, `alleles`, and
+  `left`/`right` intervals.
 - TODO
 
 ### Changed

RcppTskit/R/Class-TreeSequence.R

@@ -161,6 +161,90 @@ TreeSequence <- R6Class(
       )
     },
 
+    #' @description Iterate over sites as decoded variants.
+    #' @param samples Optional integer vector of sample node IDs to decode.
+    #' @param isolated_as_missing Logical; decode isolated samples as missing
+    #'   data (\code{TRUE}, default) or as ancestral state (\code{FALSE}).
+    #' @param alleles Optional character vector of allele states; when set,
+    #'   genotypes are indexed to this allele order.
+    #' @param impute_missing_data Deprecated alias for
+    #'   \code{!isolated_as_missing}.
+    #' @param copy Logical; currently only \code{TRUE} is supported.
+    #' @param left Left genomic coordinate (inclusive).
+    #' @param right Right genomic coordinate (exclusive). \code{NULL} means
+    #'   sequence length.
+    #' @details See the \code{tskit Python} equivalent at
+    #'   \url{https://tskit.dev/tskit/docs/latest/python-api.html#tskit.TreeSequence.variants}.
+    #' @return A simple iterator object with methods \code{next()} and
+    #'   \code{next_variant()} that each return either a variant list or
+    #'   \code{NULL} at end.
+    #' @examples
+    #' ts_file <- system.file("examples/test.trees", package = "RcppTskit")
+    #' ts <- ts_load(ts_file)
+    #' it <- ts$variants()
+    #' v1 <- it$next_variant()
+    #' v2 <- it$next_variant()
+    #' is.list(v1)
+    #' is.list(v2)
+    variants = function(
+      samples = NULL,
+      isolated_as_missing = TRUE,
+      alleles = NULL,
+      impute_missing_data = NULL,
+      copy = TRUE,
+      left = 0,
+      right = NULL
+    ) {
+      if (!is.logical(copy) || length(copy) != 1 || is.na(copy)) {
+        stop("copy must be TRUE/FALSE!")
+      }
+      if (!copy) {
+        stop("copy = FALSE is not supported yet!")
+      }
+      if (!is.null(impute_missing_data)) {
+        if (
+          !is.logical(impute_missing_data) ||
+            length(impute_missing_data) != 1 ||
+            is.na(impute_missing_data)
+        ) {
+          stop("impute_missing_data must be TRUE/FALSE or NULL!")
+        }
+        mapped <- !impute_missing_data
+        if (
+          !missing(isolated_as_missing) &&
+            !identical(isolated_as_missing, mapped)
+        ) {
+          stop(
+            "isolated_as_missing and impute_missing_data are inconsistent!"
+          )
+        }
+        warning(
+          "impute_missing_data is deprecated; use isolated_as_missing",
+          call. = FALSE
+        )
+        isolated_as_missing <- mapped
+      }
+
+      iter_xptr <- rtsk_variant_iterator_init(
+        ts = self$xptr,
+        samples = samples,
+        isolated_as_missing = isolated_as_missing,
+        alleles = alleles,
+        left = left,
+        right = if (is.null(right)) NA_real_ else right
+      )
+
+      env <- new.env(parent = emptyenv())
+      env$iter_xptr <- iter_xptr
+      next_fun <- function() {
+        rtsk_variant_iterator_next(env$iter_xptr)
+      }
+      structure(
+        list(`next` = next_fun, next_variant = next_fun),
+        class = "rtsk_variant_iterator"
+      )
+    },
+
     #' @description Get the number of provenances in a tree sequence.
     #' @return A signed 64 bit integer \code{bit64::integer64}.
     #' @details See the \code{tskit Python} equivalent at

RcppTskit/R/RcppExports.R

@@ -1,6 +1,14 @@
 # Generated by using Rcpp::compileAttributes() -> do not edit by hand
 # Generator token: 10BE3573-1514-4C36-9D1C-5A225CD40393
 
+rtsk_variant_iterator_init <- function(ts, samples = NULL, isolated_as_missing = TRUE, alleles = NULL, left = 0.0, right = NA_real_) {
+    .Call(`_RcppTskit_rtsk_variant_iterator_init`, ts, samples, isolated_as_missing, alleles, left, right)
+}
+
+rtsk_variant_iterator_next <- function(iterator) {
+    .Call(`_RcppTskit_rtsk_variant_iterator_next`, iterator)
+}
+
 test_validate_options <- function(options, supported) {
     .Call(`_RcppTskit_test_validate_options`, options, supported)
 }

RcppTskit/inst/include/RcppTskit_public.hpp

@@ -18,6 +18,12 @@ void rtsk_table_collection_dump(SEXP tc, std::string &filename,
                                 int options = 0);
 SEXP rtsk_treeseq_copy_tables(SEXP ts, int options = 0);
 SEXP rtsk_treeseq_init(SEXP tc, int options = 0);
+SEXP rtsk_variant_iterator_init(
+    SEXP ts, Rcpp::Nullable<Rcpp::IntegerVector> samples = R_NilValue,
+    bool isolated_as_missing = true,
+    Rcpp::Nullable<Rcpp::CharacterVector> alleles = R_NilValue,
+    double left = 0.0, double right = NA_REAL);
+SEXP rtsk_variant_iterator_next(SEXP iterator);
 
 SEXP rtsk_treeseq_get_num_provenances(SEXP ts);
 SEXP rtsk_treeseq_get_num_populations(SEXP ts);

RcppTskit/man/TreeSequence.Rd

@@ -10,6 +10,33 @@ Rcpp::Rostream<true>&  Rcpp::Rcout = Rcpp::Rcpp_cout_get();
 Rcpp::Rostream<false>& Rcpp::Rcerr = Rcpp::Rcpp_cerr_get();
 #endif
 
+// rtsk_variant_iterator_init
+SEXP rtsk_variant_iterator_init(const SEXP ts, const Rcpp::Nullable<Rcpp::IntegerVector> samples, const bool isolated_as_missing, const Rcpp::Nullable<Rcpp::CharacterVector> alleles, const double left, const double right);
+RcppExport SEXP _RcppTskit_rtsk_variant_iterator_init(SEXP tsSEXP, SEXP samplesSEXP, SEXP isolated_as_missingSEXP, SEXP allelesSEXP, SEXP leftSEXP, SEXP rightSEXP) {
+BEGIN_RCPP
+    Rcpp::RObject rcpp_result_gen;
+    Rcpp::RNGScope rcpp_rngScope_gen;
+    Rcpp::traits::input_parameter< const SEXP >::type ts(tsSEXP);
+    Rcpp::traits::input_parameter< const Rcpp::Nullable<Rcpp::IntegerVector> >::type samples(samplesSEXP);
+    Rcpp::traits::input_parameter< const bool >::type isolated_as_missing(isolated_as_missingSEXP);
+    Rcpp::traits::input_parameter< const Rcpp::Nullable<Rcpp::CharacterVector> >::type alleles(allelesSEXP);
+    Rcpp::traits::input_parameter< const double >::type left(leftSEXP);
+    Rcpp::traits::input_parameter< const double >::type right(rightSEXP);
+    rcpp_result_gen = Rcpp::wrap(rtsk_variant_iterator_init(ts, samples, isolated_as_missing, alleles, left, right));
+    return rcpp_result_gen;
+END_RCPP
+}
+// rtsk_variant_iterator_next
+SEXP rtsk_variant_iterator_next(const SEXP iterator);
+RcppExport SEXP _RcppTskit_rtsk_variant_iterator_next(SEXP iteratorSEXP) {
+BEGIN_RCPP
+    Rcpp::RObject rcpp_result_gen;
+    Rcpp::RNGScope rcpp_rngScope_gen;
+    Rcpp::traits::input_parameter< const SEXP >::type iterator(iteratorSEXP);
+    rcpp_result_gen = Rcpp::wrap(rtsk_variant_iterator_next(iterator));
+    return rcpp_result_gen;
+END_RCPP
+}
 // test_validate_options
 int test_validate_options(const int options, const int supported);
 RcppExport SEXP _RcppTskit_test_validate_options(SEXP optionsSEXP, SEXP supportedSEXP) {
@@ -742,6 +769,8 @@ END_RCPP
 }
 
 static const R_CallMethodDef CallEntries[] = {
+    {"_RcppTskit_rtsk_variant_iterator_init", (DL_FUNC) &_RcppTskit_rtsk_variant_iterator_init, 6},
+    {"_RcppTskit_rtsk_variant_iterator_next", (DL_FUNC) &_RcppTskit_rtsk_variant_iterator_next, 1},
     {"_RcppTskit_test_validate_options", (DL_FUNC) &_RcppTskit_test_validate_options, 2},
     {"_RcppTskit_test_rtsk_wrap_tsk_size_t_as_integer64", (DL_FUNC) &_RcppTskit_test_rtsk_wrap_tsk_size_t_as_integer64, 2},
     {"_RcppTskit_kastore_version", (DL_FUNC) &_RcppTskit_kastore_version, 0},