Add get_samples(), node_table_get_row() and table_collection_sort()

Author: LynxJinyangii

RcppTskit/NEWS.md

@@ -40,12 +40,18 @@ 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 `rtsk_node_table_get_row()` and `TableCollection$node_table_get_row()`
+  to retrieve node-table rows by 0-based row ID.
+- Added `rtsk_table_collection_sort()` and `TableCollection$sort()` to sort
+  table collections with 0-based `edge_start` semantics.
 - 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.
+- Added `rtsk_treeseq_get_samples()` and `TreeSequence$samples()` to retrieve
+  sample node IDs from a tree sequence.
 - TODO
 
 ### Changed

RcppTskit/R/Class-TableCollection.R

@@ -103,6 +103,47 @@ TableCollection <- R6Class(
       TreeSequence$new(xptr = ts_xptr)
     },
 
+    #' @description Sort this table collection in place.
+    #' @param edge_start integer scalar edge-table start row (0-based).
+    #' @param no_check_integrity logical; when \code{TRUE}, pass
+    #'   \code{TSK_NO_CHECK_INTEGRITY} to \code{tskit C}.
+    #' @details See the \code{tskit Python} equivalent at
+    #'   \url{https://tskit.dev/tskit/docs/latest/python-api.html#tskit.TableCollection.sort}.
+    #' @return No return value; called for side effects.
+    #' @examples
+    #' ts_file <- system.file("examples/test.trees", package = "RcppTskit")
+    #' tc <- tc_load(ts_file)
+    #' tc$sort()
+    sort = function(edge_start = 0L, no_check_integrity = FALSE) {
+      if (
+        is.null(edge_start) ||
+          length(edge_start) != 1L ||
+          is.na(as.integer(edge_start))
+      ) {
+        stop("edge_start must be a non-NA integer scalar (0-based)!")
+      }
+      if (as.integer(edge_start) < 0L) {
+        stop("edge_start must be >= 0 (0-based)!")
+      }
+      if (
+        !is.logical(no_check_integrity) ||
+          length(no_check_integrity) != 1L ||
+          is.na(no_check_integrity)
+      ) {
+        stop("no_check_integrity must be TRUE/FALSE!")
+      }
+      options <- if (isTRUE(no_check_integrity)) {
+        as.integer(rtsk_const_tsk_no_check_integrity())
+      } else {
+        0L
+      }
+      rtsk_table_collection_sort(
+        tc = self$xptr,
+        edge_start = as.integer(edge_start),
+        options = options
+      )
+    },
+
     #' @description Get the number of provenances in a table collection.
     #' @return A signed 64 bit integer \code{bit64::integer64}.
     #' @examples
@@ -202,6 +243,27 @@ TableCollection <- R6Class(
       rtsk_table_collection_get_num_nodes(self$xptr)
     },
 
+    #' @description Get one row from the nodes table.
+    #' @param row_id integer scalar node row ID (0-based).
+    #' @details The ID is 0-based, matching \code{tskit C/Python} semantics.
+    #' @return A named list with fields \code{id}, \code{flags}, \code{time},
+    #'   \code{population}, \code{individual}, and \code{metadata}.
+    #' @examples
+    #' ts_file <- system.file("examples/test.trees", package = "RcppTskit")
+    #' tc <- tc_load(ts_file)
+    #' tc$node_table_get_row(0L)
+    node_table_get_row = function(row_id) {
+      if (
+        is.null(row_id) || length(row_id) != 1L || is.na(as.integer(row_id))
+      ) {
+        stop("row_id must be a non-NA integer scalar (0-based)!")
+      }
+      if (as.integer(row_id) < 0L) {
+        stop("row_id must be >= 0 (0-based)!")
+      }
+      rtsk_node_table_get_row(self$xptr, row_id = as.integer(row_id))
+    },
+
     #' @description Add a row to the nodes table.
     #' @param flags integer flags for the new node.
     #' @param time numeric time value for the new node.

RcppTskit/R/Class-TreeSequence.R

@@ -305,6 +305,18 @@ TreeSequence <- R6Class(
       rtsk_treeseq_get_num_samples(self$xptr)
     },
 
+    #' @description Get sample node IDs in this tree sequence.
+    #' @return An integer vector with sample node IDs (0-based).
+    #' @details See the \code{tskit Python} equivalent at
+    #'   \url{https://tskit.dev/tskit/docs/latest/python-api.html#tskit.TreeSequence.samples}.
+    #' @examples
+    #' ts_file <- system.file("examples/test.trees", package = "RcppTskit")
+    #' ts <- ts_load(ts_file)
+    #' ts$samples()
+    samples = function() {
+      rtsk_treeseq_get_samples(self$xptr)
+    },
+
     #' @description Get the number of nodes 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

@@ -51,6 +51,10 @@ tskit_version <- function() {
     .Call(`_RcppTskit_tskit_version`)
 }
 
+rtsk_const_tsk_no_check_integrity <- function() {
+    .Call(`_RcppTskit_rtsk_const_tsk_no_check_integrity`)
+}
+
 rtsk_treeseq_load <- function(filename, options = 0L) {
     .Call(`_RcppTskit_rtsk_treeseq_load`, filename, options)
 }
@@ -95,6 +99,10 @@ rtsk_treeseq_get_num_samples <- function(ts) {
     .Call(`_RcppTskit_rtsk_treeseq_get_num_samples`, ts)
 }
 
+rtsk_treeseq_get_samples <- function(ts) {
+    .Call(`_RcppTskit_rtsk_treeseq_get_samples`, ts)
+}
+
 rtsk_treeseq_get_num_nodes <- function(ts) {
     .Call(`_RcppTskit_rtsk_treeseq_get_num_nodes`, ts)
 }
@@ -215,6 +223,10 @@ rtsk_table_collection_drop_index <- function(tc, options = 0L) {
     invisible(.Call(`_RcppTskit_rtsk_table_collection_drop_index`, tc, options))
 }
 
+rtsk_table_collection_sort <- function(tc, edge_start = 0L, options = 0L) {
+    invisible(.Call(`_RcppTskit_rtsk_table_collection_sort`, tc, edge_start, options))
+}
+
 rtsk_table_collection_summary <- function(tc) {
     .Call(`_RcppTskit_rtsk_table_collection_summary`, tc)
 }
@@ -231,6 +243,10 @@ rtsk_node_table_add_row <- function(tc, flags = 0L, time = 0, population = -1L,
     .Call(`_RcppTskit_rtsk_node_table_add_row`, tc, flags, time, population, individual, metadata)
 }
 
+rtsk_node_table_get_row <- function(tc, row_id) {
+    .Call(`_RcppTskit_rtsk_node_table_get_row`, tc, row_id)
+}
+
 rtsk_edge_table_add_row <- function(tc, left, right, parent, child, metadata = NULL) {
     .Call(`_RcppTskit_rtsk_edge_table_add_row`, tc, left, right, parent, child, metadata)
 }

RcppTskit/inst/include/RcppTskit_public.hpp

@@ -9,6 +9,7 @@
 // RcppTskit.cpp
 Rcpp::IntegerVector kastore_version();
 Rcpp::IntegerVector tskit_version();
+int rtsk_const_tsk_no_check_integrity();
 
 // sync default options with .cpp!
 SEXP rtsk_treeseq_load(std::string &filename, int options = 0);
@@ -30,6 +31,7 @@ SEXP rtsk_treeseq_get_num_populations(SEXP ts);
 SEXP rtsk_treeseq_get_num_migrations(SEXP ts);
 SEXP rtsk_treeseq_get_num_individuals(SEXP ts);
 SEXP rtsk_treeseq_get_num_samples(SEXP ts);
+Rcpp::IntegerVector rtsk_treeseq_get_samples(SEXP ts);
 SEXP rtsk_treeseq_get_num_nodes(SEXP ts);
 SEXP rtsk_treeseq_get_num_edges(SEXP ts);
 SEXP rtsk_treeseq_get_num_trees(SEXP ts);
@@ -61,6 +63,7 @@ Rcpp::String rtsk_table_collection_get_file_uuid(SEXP tc);
 bool rtsk_table_collection_has_index(SEXP tc, int options = 0);
 void rtsk_table_collection_build_index(SEXP tc, int options = 0);
 void rtsk_table_collection_drop_index(SEXP tc, int options = 0);
+void rtsk_table_collection_sort(SEXP tc, int edge_start = 0, int options = 0);
 Rcpp::List rtsk_table_collection_summary(SEXP tc);
 Rcpp::List rtsk_table_collection_metadata_length(SEXP tc);
 int rtsk_individual_table_add_row(
@@ -71,6 +74,7 @@ int rtsk_individual_table_add_row(
 int rtsk_node_table_add_row(
     SEXP tc, int flags = 0, double time = 0, int population = -1,
     int individual = -1, Rcpp::Nullable<Rcpp::RawVector> metadata = R_NilValue);
+Rcpp::List rtsk_node_table_get_row(SEXP tc, int row_id);
 int rtsk_edge_table_add_row(
     SEXP tc, double left, double right, int parent, int child,
     Rcpp::Nullable<Rcpp::RawVector> metadata = R_NilValue);