From 68a5773331b22bece9e0cb99686f0d30e4652e9a Mon Sep 17 00:00:00 2001 From: erliuu Date: Fri, 12 Jun 2026 17:10:00 -0700 Subject: [PATCH 1/2] [r][cpp] Add functionality for IterableMatrix Bernoulli sampling and tests --- r/NAMESPACE | 1 + r/NEWS.md | 1 + r/R/RcppExports.R | 12 ++++ r/R/transforms.R | 52 ++++++++++++++ r/man/IterableMatrix-methods.Rd | 5 +- r/man/bernoulli_sample.Rd | 22 ++++++ r/src/RcppExports.cpp | 42 +++++++++++ .../matrixIterators/BernoulliSample.h | 72 +++++++++++++++++++ r/src/matrix_utils.cpp | 28 ++++++++ r/tests/testthat/test-matrix_transforms.R | 59 +++++++++++++++ 10 files changed, 293 insertions(+), 1 deletion(-) create mode 100644 r/man/bernoulli_sample.Rd create mode 100644 r/src/bpcells-cpp/matrixIterators/BernoulliSample.h diff --git a/r/NAMESPACE b/r/NAMESPACE index 39e8b00f..017dd5ce 100644 --- a/r/NAMESPACE +++ b/r/NAMESPACE @@ -13,6 +13,7 @@ export(add_rows) export(all_matrix_inputs) export(apply_by_col) export(apply_by_row) +export(bernoulli_sample) export(binarize) export(call_macs_peaks) export(call_peaks_macs) diff --git a/r/NEWS.md b/r/NEWS.md index ac6ca2a3..87d26e10 100644 --- a/r/NEWS.md +++ b/r/NEWS.md @@ -8,6 +8,7 @@ - Add `tile_width` and `normalization` arguments to `write_insertion_bedgraph()` to allow for more flexible bedgraph creation (pull request #299) - Export `write_insertion_bed()`, which originally was only a helper for peak calling (pull request #302). - Add support for reading 64 bit Anndata matrices, as well as reading signed integer matrices. Note, signed integers are converted to float/double matrices when read in. (pull request #330) +- Add `bernoulli_sample()`, a lazy `IterableMatrix` transform that keeps each non-zero entry independently with probability `prob` using a reproducible, stateless hash of `(seed, row, col)`. ## Bug-fixes - Fix error in documentation examples for `plot_embedding()`, resulting from the way documentation examples use nested function calls (pull request #316). diff --git a/r/R/RcppExports.R b/r/R/RcppExports.R index e79f0d11..8a63cb2b 100644 --- a/r/R/RcppExports.R +++ b/r/R/RcppExports.R @@ -685,6 +685,18 @@ iterate_matrix_mask_double_cpp <- function(mat, mask, invert) { .Call(`_BPCells_iterate_matrix_mask_double_cpp`, mat, mask, invert) } +iterate_matrix_bernoulli_sample_uint32_t_cpp <- function(mat, prob, seed) { + .Call(`_BPCells_iterate_matrix_bernoulli_sample_uint32_t_cpp`, mat, prob, seed) +} + +iterate_matrix_bernoulli_sample_float_cpp <- function(mat, prob, seed) { + .Call(`_BPCells_iterate_matrix_bernoulli_sample_float_cpp`, mat, prob, seed) +} + +iterate_matrix_bernoulli_sample_double_cpp <- function(mat, prob, seed) { + .Call(`_BPCells_iterate_matrix_bernoulli_sample_double_cpp`, mat, prob, seed) +} + iterate_matrix_rank_uint32_t_cpp <- function(matrix) { .Call(`_BPCells_iterate_matrix_rank_uint32_t_cpp`, matrix) } diff --git a/r/R/transforms.R b/r/R/transforms.R index 097b2bd9..3f6cacda 100644 --- a/r/R/transforms.R +++ b/r/R/transforms.R @@ -1099,3 +1099,55 @@ regress_out <- function(mat, latent_data, prediction_axis = c("row", "col")) { vars_to_regress = vars_to_regress ) } + +################# +# Bernoulli sample +################# + +setClass("MatrixBernoulliSample", + contains = "IterableMatrix", + slots = c( + matrix = "IterableMatrix", + prob = "numeric", + seed = "integer" + ), + prototype = list( + matrix = NULL, + prob = 1, + seed = 0L + ) +) + +setMethod("matrix_type", "MatrixBernoulliSample", function(x) matrix_type(x@matrix)) +setMethod("iterate_matrix", "MatrixBernoulliSample", function(x) { + iter_fn <- get(sprintf("iterate_matrix_bernoulli_sample_%s_cpp", matrix_type(x))) + iter_fn(iterate_matrix(x@matrix), x@prob, x@seed) +}) + +setMethod("short_description", "MatrixBernoulliSample", function(x) { + c( + short_description(x@matrix), + sprintf("Bernoulli sample non-zero entries (prob=%.3g, seed=%d)", x@prob, x@seed) + ) +}) + +#' Randomly drop non-zero matrix entries given probability and seed +#' +#' Keep each non-zero entry independently with probability `prob`. The +#' keep/drop decision is a deterministic function of `(seed, row, col)` +#' +#' @param mat IterableMatrix +#' @param prob Numeric value in (0, 1]. Keep probability per non-zero entry +#' @param seed Reproducibility seed +#' @return IterableMatrix +#' @export +bernoulli_sample <- function(mat, prob, seed) { + assert_is(mat, "IterableMatrix") + assert_is_numeric(prob) + assert_len(prob, 1) + assert_true(!is.na(prob) && prob > 0 && prob <= 1) + assert_is_wholenumber(seed) + assert_len(seed, 1) + if (prob == 1) return(mat) + wrapMatrix("MatrixBernoulliSample", mat, prob = prob, seed = as.integer(seed)) +} diff --git a/r/man/IterableMatrix-methods.Rd b/r/man/IterableMatrix-methods.Rd index f737299c..be522632 100644 --- a/r/man/IterableMatrix-methods.Rd +++ b/r/man/IterableMatrix-methods.Rd @@ -248,7 +248,8 @@ Generic methods and built-in functions for IterableMatrix objects \item \code{e1 * e2}: Multiply by a constant, or multiply rows by a vector length nrow(mat) -\item \code{e1 + e2}: Add a constant, or row-wise addition with a vector length nrow(mat) +\item \code{e1 + e2}: Add a constant, row-wise addition with a vector length nrow(mat), +or element-wise addition of two IterableMatrix objects \item \code{e1 / e2}: Divide by a constant, or divide rows by a vector length nrow(mat) @@ -420,6 +421,8 @@ mat + 5 ## Adding row-wise by a vector of length `nrow(mat)` mat + 1:nrow(mat) +## Element-wise addition of two IterableMatrix objects +mat + (mat * 2) ####################################################################### ## `e1 / e2` example diff --git a/r/man/bernoulli_sample.Rd b/r/man/bernoulli_sample.Rd new file mode 100644 index 00000000..da55cce0 --- /dev/null +++ b/r/man/bernoulli_sample.Rd @@ -0,0 +1,22 @@ +% Generated by roxygen2: do not edit by hand +% Please edit documentation in R/transforms.R +\name{bernoulli_sample} +\alias{bernoulli_sample} +\title{Randomly drop non-zero matrix entries given probability and seed} +\usage{ +bernoulli_sample(mat, prob, seed) +} +\arguments{ +\item{mat}{IterableMatrix} + +\item{prob}{Numeric value in (0, 1]. Keep probability per non-zero entry} + +\item{seed}{Reproducibility seed} +} +\value{ +IterableMatrix +} +\description{ +Keep each non-zero entry independently with probability \code{prob}. The +keep/drop decision is a deterministic function of \verb{(seed, row, col)} +} diff --git a/r/src/RcppExports.cpp b/r/src/RcppExports.cpp index ac7014e5..c4aad300 100644 --- a/r/src/RcppExports.cpp +++ b/r/src/RcppExports.cpp @@ -2287,6 +2287,45 @@ BEGIN_RCPP return rcpp_result_gen; END_RCPP } +// iterate_matrix_bernoulli_sample_uint32_t_cpp +SEXP iterate_matrix_bernoulli_sample_uint32_t_cpp(SEXP mat, double prob, int seed); +RcppExport SEXP _BPCells_iterate_matrix_bernoulli_sample_uint32_t_cpp(SEXP matSEXP, SEXP probSEXP, SEXP seedSEXP) { +BEGIN_RCPP + Rcpp::RObject rcpp_result_gen; + Rcpp::RNGScope rcpp_rngScope_gen; + Rcpp::traits::input_parameter< SEXP >::type mat(matSEXP); + Rcpp::traits::input_parameter< double >::type prob(probSEXP); + Rcpp::traits::input_parameter< int >::type seed(seedSEXP); + rcpp_result_gen = Rcpp::wrap(iterate_matrix_bernoulli_sample_uint32_t_cpp(mat, prob, seed)); + return rcpp_result_gen; +END_RCPP +} +// iterate_matrix_bernoulli_sample_float_cpp +SEXP iterate_matrix_bernoulli_sample_float_cpp(SEXP mat, double prob, int seed); +RcppExport SEXP _BPCells_iterate_matrix_bernoulli_sample_float_cpp(SEXP matSEXP, SEXP probSEXP, SEXP seedSEXP) { +BEGIN_RCPP + Rcpp::RObject rcpp_result_gen; + Rcpp::RNGScope rcpp_rngScope_gen; + Rcpp::traits::input_parameter< SEXP >::type mat(matSEXP); + Rcpp::traits::input_parameter< double >::type prob(probSEXP); + Rcpp::traits::input_parameter< int >::type seed(seedSEXP); + rcpp_result_gen = Rcpp::wrap(iterate_matrix_bernoulli_sample_float_cpp(mat, prob, seed)); + return rcpp_result_gen; +END_RCPP +} +// iterate_matrix_bernoulli_sample_double_cpp +SEXP iterate_matrix_bernoulli_sample_double_cpp(SEXP mat, double prob, int seed); +RcppExport SEXP _BPCells_iterate_matrix_bernoulli_sample_double_cpp(SEXP matSEXP, SEXP probSEXP, SEXP seedSEXP) { +BEGIN_RCPP + Rcpp::RObject rcpp_result_gen; + Rcpp::RNGScope rcpp_rngScope_gen; + Rcpp::traits::input_parameter< SEXP >::type mat(matSEXP); + Rcpp::traits::input_parameter< double >::type prob(probSEXP); + Rcpp::traits::input_parameter< int >::type seed(seedSEXP); + rcpp_result_gen = Rcpp::wrap(iterate_matrix_bernoulli_sample_double_cpp(mat, prob, seed)); + return rcpp_result_gen; +END_RCPP +} // iterate_matrix_rank_uint32_t_cpp SEXP iterate_matrix_rank_uint32_t_cpp(SEXP matrix); RcppExport SEXP _BPCells_iterate_matrix_rank_uint32_t_cpp(SEXP matrixSEXP) { @@ -2774,6 +2813,9 @@ static const R_CallMethodDef CallEntries[] = { {"_BPCells_iterate_matrix_mask_uint32_t_cpp", (DL_FUNC) &_BPCells_iterate_matrix_mask_uint32_t_cpp, 3}, {"_BPCells_iterate_matrix_mask_float_cpp", (DL_FUNC) &_BPCells_iterate_matrix_mask_float_cpp, 3}, {"_BPCells_iterate_matrix_mask_double_cpp", (DL_FUNC) &_BPCells_iterate_matrix_mask_double_cpp, 3}, + {"_BPCells_iterate_matrix_bernoulli_sample_uint32_t_cpp", (DL_FUNC) &_BPCells_iterate_matrix_bernoulli_sample_uint32_t_cpp, 3}, + {"_BPCells_iterate_matrix_bernoulli_sample_float_cpp", (DL_FUNC) &_BPCells_iterate_matrix_bernoulli_sample_float_cpp, 3}, + {"_BPCells_iterate_matrix_bernoulli_sample_double_cpp", (DL_FUNC) &_BPCells_iterate_matrix_bernoulli_sample_double_cpp, 3}, {"_BPCells_iterate_matrix_rank_uint32_t_cpp", (DL_FUNC) &_BPCells_iterate_matrix_rank_uint32_t_cpp, 1}, {"_BPCells_iterate_matrix_rank_float_cpp", (DL_FUNC) &_BPCells_iterate_matrix_rank_float_cpp, 1}, {"_BPCells_iterate_matrix_rank_double_cpp", (DL_FUNC) &_BPCells_iterate_matrix_rank_double_cpp, 1}, diff --git a/r/src/bpcells-cpp/matrixIterators/BernoulliSample.h b/r/src/bpcells-cpp/matrixIterators/BernoulliSample.h new file mode 100644 index 00000000..7482c191 --- /dev/null +++ b/r/src/bpcells-cpp/matrixIterators/BernoulliSample.h @@ -0,0 +1,72 @@ +// Copyright 2026 BPCells contributors +// +// Licensed under the Apache License, Version 2.0 or the MIT license +// , at your +// option. This file may not be copied, modified, or distributed +// except according to those terms. + +#pragma once + +#include +#include +#include + +#include "MatrixIterator.h" + +namespace BPCells { + +// Keep each non-zero entry independently with probability `prob`, with the +// keep/drop decision a deterministic function of (seed, col, row) via splitmix64. +template class BernoulliSample : public MatrixLoaderWrapper { + private: + uint64_t seed; + uint32_t threshold; + size_t loaded = 0; + + // Hash algorithm + static inline uint64_t splitmix64(uint64_t x) { + x += 0x9e3779b97f4a7c15ULL; + x = (x ^ (x >> 30)) * 0xbf58476d1ce4e5b9ULL; + x = (x ^ (x >> 27)) * 0x94d049bb133111ebULL; + return x ^ (x >> 31); + } + + public: + BernoulliSample(std::unique_ptr> &&loader, double prob, uint64_t seed) + : MatrixLoaderWrapper(std::move(loader)), seed(seed) { + if (!(prob > 0.0 && prob <= 1.0)) + throw std::runtime_error("BernoulliSample: prob must be in (0, 1]"); + double scaled = prob * 4294967296.0; // 2^32 + threshold = (prob >= 1.0 || scaled >= 4294967295.5) + ? UINT32_MAX + : static_cast(scaled); + } + + bool load() override { + loaded = 0; + + while (loaded == 0) { + if (!this->loader->load()) return false; + uint32_t *row_data = this->loader->rowData(); + T *val_data = this->loader->valData(); + size_t cap = this->loader->capacity(); + + // Calculate per-column hash + uint64_t col_mix = splitmix64(seed ^ (uint64_t(this->loader->currentCol()) << 32)); + + for (size_t i = 0; i < cap; i++) { + // Calculate final hash with row number + uint32_t h = static_cast(splitmix64(col_mix ^ row_data[i])); + row_data[loaded] = row_data[i]; + val_data[loaded] = val_data[i]; + loaded += h < threshold; + } + } + return true; + } + + uint32_t capacity() const override { return loaded; } +}; + +} // end namespace BPCells diff --git a/r/src/matrix_utils.cpp b/r/src/matrix_utils.cpp index 23fd51f1..a533c4fa 100644 --- a/r/src/matrix_utils.cpp +++ b/r/src/matrix_utils.cpp @@ -10,6 +10,7 @@ #define RCPP_NO_SUGAR #include #include +#include "bpcells-cpp/matrixIterators/BernoulliSample.h" #include "bpcells-cpp/matrixIterators/CSparseMatrix.h" #include "bpcells-cpp/matrixIterators/ColwiseRank.h" #include "bpcells-cpp/matrixIterators/MatrixAddition.h" @@ -345,6 +346,33 @@ SEXP iterate_matrix_mask_double_cpp(SEXP mat, SEXP mask, bool invert) { ); } +// [[Rcpp::export]] +SEXP iterate_matrix_bernoulli_sample_uint32_t_cpp(SEXP mat, double prob, int seed) { + return make_unique_xptr>( + take_unique_xptr>(mat), + prob, + static_cast(static_cast(seed)) + ); +} + +// [[Rcpp::export]] +SEXP iterate_matrix_bernoulli_sample_float_cpp(SEXP mat, double prob, int seed) { + return make_unique_xptr>( + take_unique_xptr>(mat), + prob, + static_cast(static_cast(seed)) + ); +} + +// [[Rcpp::export]] +SEXP iterate_matrix_bernoulli_sample_double_cpp(SEXP mat, double prob, int seed) { + return make_unique_xptr>( + take_unique_xptr>(mat), + prob, + static_cast(static_cast(seed)) + ); +} + // [[Rcpp::export]] SEXP iterate_matrix_rank_uint32_t_cpp(SEXP matrix) { return make_unique_xptr>(take_unique_xptr>(matrix) diff --git a/r/tests/testthat/test-matrix_transforms.R b/r/tests/testthat/test-matrix_transforms.R index 09bf4e20..65aa58e7 100644 --- a/r/tests/testthat/test-matrix_transforms.R +++ b/r/tests/testthat/test-matrix_transforms.R @@ -418,3 +418,62 @@ test_that("matrix subtraction works", { bp1_t <- t(as(t(m1), "IterableMatrix")) expect_error(bp1 - bp1_t) }) + +test_that("bernoulli_sample works", { + set.seed(125124) + m <- generate_sparse_matrix(50, 40, fraction_nonzero = 0.4, max_val = 100) + bp <- as(m, "IterableMatrix") + + # prob == 1: no-op, returns the same object unchanged + expect_identical(bernoulli_sample(bp, 1, 0), bp) + + # Keep fraction: ~prob of non-zeros survive, within 4 sigma on ~2M entries + m_big <- generate_sparse_matrix(2000, 2000, fraction_nonzero = 0.5, max_val = 10) + bp_big <- as(m_big, "IterableMatrix") + nnz_in <- length(m_big@x) + out <- as(bernoulli_sample(bp_big, 0.5, 17L), "dgCMatrix") + nnz_out <- length(out@x) + sigma <- sqrt(0.25 * nnz_in) + expect_lt(abs(nnz_out - 0.5 * nnz_in), 4 * sigma) + + # Reproduciblity: same seed materializes identical output + y1 <- as(bernoulli_sample(bp, 0.3, 42L), "dgCMatrix") + y2 <- as(bernoulli_sample(bp, 0.3, 42L), "dgCMatrix") + expect_identical(y1, y2) + + # Reverse column order matches forward iteration (exercises seekCol via MatrixSubset) + sampled <- bernoulli_sample(bp, 0.3, 42L) + reversed <- sampled[, rev(seq_len(ncol(bp)))] + expect_identical(as(reversed, "dgCMatrix"), as(sampled, "dgCMatrix")[, rev(seq_len(ncol(bp)))]) + + # Test hash keying to output columns (duplicated subset columns get independent draws) + col1 <- ceiling(ncol(bp) / 2) + dup_subset <- methods::new( + "MatrixSubset", + matrix = bp, + col_selection = rep(as.integer(col1), 3), + dim = c(nrow(bp), 3L), + dimnames = list(rownames(bp), NULL), + zero_dims = c(FALSE, FALSE) + ) + dup <- bernoulli_sample(dup_subset, 0.5, 99L) + dup_mat <- as(dup, "dgCMatrix") + expect_false(identical(dup_mat[, 1], dup_mat[, 2])) + expect_false(identical(dup_mat[, 1], dup_mat[, 3])) + expect_false(identical(dup_mat[, 2], dup_mat[, 3])) + + # Sampling is value-independent: Shared keep/drop decisions regardless of type + for (tp in c("uint32_t", "float", "double")) { + bp_t <- convert_matrix_type(bp, tp) + out <- bernoulli_sample(bp_t, 0.5, 13L) + expect_identical(matrix_type(out), tp) + expect_equal( + as(out, "dgCMatrix")@i, + as(bernoulli_sample(bp, 0.5, 13L), "dgCMatrix")@i + ) + expect_equal( + as(out, "dgCMatrix")@p, + as(bernoulli_sample(bp, 0.5, 13L), "dgCMatrix")@p + ) + } +}) From 9e6b6fe91bc30dcfdcbba1f578a4bceaa82ccbee Mon Sep 17 00:00:00 2001 From: erliuu Date: Wed, 24 Jun 2026 00:30:15 -0700 Subject: [PATCH 2/2] [r][cpp] Match splitmix64 to standard implementation, add transpose storage order invariance, rename to sample_bernoulli --- r/NAMESPACE | 2 +- r/NEWS.md | 2 +- r/R/RcppExports.R | 12 +++--- r/R/transforms.R | 4 +- ...ernoulli_sample.Rd => sample_bernoulli.Rd} | 6 +-- r/src/RcppExports.cpp | 27 ++++++------ .../matrixIterators/BernoulliSample.h | 40 +++++++++++------- r/src/matrix_utils.cpp | 42 +++++++++++-------- r/tests/testthat/test-matrix_transforms.R | 42 ++++++++----------- 9 files changed, 94 insertions(+), 83 deletions(-) rename r/man/{bernoulli_sample.Rd => sample_bernoulli.Rd} (86%) diff --git a/r/NAMESPACE b/r/NAMESPACE index 017dd5ce..ca3c3cc2 100644 --- a/r/NAMESPACE +++ b/r/NAMESPACE @@ -13,7 +13,6 @@ export(add_rows) export(all_matrix_inputs) export(apply_by_col) export(apply_by_row) -export(bernoulli_sample) export(binarize) export(call_macs_peaks) export(call_peaks_macs) @@ -114,6 +113,7 @@ export(rowSums) export(rowVars) export(rowVars.IterableMatrix) export(rowVars.default) +export(sample_bernoulli) export(sctransform_pearson) export(select_cells) export(select_chromosomes) diff --git a/r/NEWS.md b/r/NEWS.md index 87d26e10..ef04ef6e 100644 --- a/r/NEWS.md +++ b/r/NEWS.md @@ -8,7 +8,7 @@ - Add `tile_width` and `normalization` arguments to `write_insertion_bedgraph()` to allow for more flexible bedgraph creation (pull request #299) - Export `write_insertion_bed()`, which originally was only a helper for peak calling (pull request #302). - Add support for reading 64 bit Anndata matrices, as well as reading signed integer matrices. Note, signed integers are converted to float/double matrices when read in. (pull request #330) -- Add `bernoulli_sample()`, a lazy `IterableMatrix` transform that keeps each non-zero entry independently with probability `prob` using a reproducible, stateless hash of `(seed, row, col)`. +- Add `sample_bernoulli()`, a lazy `IterableMatrix` transform that keeps each non-zero entry independently with probability `prob` using a reproducible, stateless hash of `(seed, row, col)`. Transform is invariant to storage order (`transpose_storage_order()`). ## Bug-fixes - Fix error in documentation examples for `plot_embedding()`, resulting from the way documentation examples use nested function calls (pull request #316). diff --git a/r/R/RcppExports.R b/r/R/RcppExports.R index 8a63cb2b..382599ee 100644 --- a/r/R/RcppExports.R +++ b/r/R/RcppExports.R @@ -685,16 +685,16 @@ iterate_matrix_mask_double_cpp <- function(mat, mask, invert) { .Call(`_BPCells_iterate_matrix_mask_double_cpp`, mat, mask, invert) } -iterate_matrix_bernoulli_sample_uint32_t_cpp <- function(mat, prob, seed) { - .Call(`_BPCells_iterate_matrix_bernoulli_sample_uint32_t_cpp`, mat, prob, seed) +iterate_matrix_bernoulli_sample_uint32_t_cpp <- function(mat, prob, seed, transpose) { + .Call(`_BPCells_iterate_matrix_bernoulli_sample_uint32_t_cpp`, mat, prob, seed, transpose) } -iterate_matrix_bernoulli_sample_float_cpp <- function(mat, prob, seed) { - .Call(`_BPCells_iterate_matrix_bernoulli_sample_float_cpp`, mat, prob, seed) +iterate_matrix_bernoulli_sample_float_cpp <- function(mat, prob, seed, transpose) { + .Call(`_BPCells_iterate_matrix_bernoulli_sample_float_cpp`, mat, prob, seed, transpose) } -iterate_matrix_bernoulli_sample_double_cpp <- function(mat, prob, seed) { - .Call(`_BPCells_iterate_matrix_bernoulli_sample_double_cpp`, mat, prob, seed) +iterate_matrix_bernoulli_sample_double_cpp <- function(mat, prob, seed, transpose) { + .Call(`_BPCells_iterate_matrix_bernoulli_sample_double_cpp`, mat, prob, seed, transpose) } iterate_matrix_rank_uint32_t_cpp <- function(matrix) { diff --git a/r/R/transforms.R b/r/R/transforms.R index 3f6cacda..4280ef70 100644 --- a/r/R/transforms.R +++ b/r/R/transforms.R @@ -1121,7 +1121,7 @@ setClass("MatrixBernoulliSample", setMethod("matrix_type", "MatrixBernoulliSample", function(x) matrix_type(x@matrix)) setMethod("iterate_matrix", "MatrixBernoulliSample", function(x) { iter_fn <- get(sprintf("iterate_matrix_bernoulli_sample_%s_cpp", matrix_type(x))) - iter_fn(iterate_matrix(x@matrix), x@prob, x@seed) + iter_fn(iterate_matrix(x@matrix), x@prob, x@seed, x@transpose) }) setMethod("short_description", "MatrixBernoulliSample", function(x) { @@ -1141,7 +1141,7 @@ setMethod("short_description", "MatrixBernoulliSample", function(x) { #' @param seed Reproducibility seed #' @return IterableMatrix #' @export -bernoulli_sample <- function(mat, prob, seed) { +sample_bernoulli <- function(mat, prob, seed) { assert_is(mat, "IterableMatrix") assert_is_numeric(prob) assert_len(prob, 1) diff --git a/r/man/bernoulli_sample.Rd b/r/man/sample_bernoulli.Rd similarity index 86% rename from r/man/bernoulli_sample.Rd rename to r/man/sample_bernoulli.Rd index da55cce0..83353242 100644 --- a/r/man/bernoulli_sample.Rd +++ b/r/man/sample_bernoulli.Rd @@ -1,10 +1,10 @@ % Generated by roxygen2: do not edit by hand % Please edit documentation in R/transforms.R -\name{bernoulli_sample} -\alias{bernoulli_sample} +\name{sample_bernoulli} +\alias{sample_bernoulli} \title{Randomly drop non-zero matrix entries given probability and seed} \usage{ -bernoulli_sample(mat, prob, seed) +sample_bernoulli(mat, prob, seed) } \arguments{ \item{mat}{IterableMatrix} diff --git a/r/src/RcppExports.cpp b/r/src/RcppExports.cpp index c4aad300..ca8d835f 100644 --- a/r/src/RcppExports.cpp +++ b/r/src/RcppExports.cpp @@ -2288,41 +2288,44 @@ BEGIN_RCPP END_RCPP } // iterate_matrix_bernoulli_sample_uint32_t_cpp -SEXP iterate_matrix_bernoulli_sample_uint32_t_cpp(SEXP mat, double prob, int seed); -RcppExport SEXP _BPCells_iterate_matrix_bernoulli_sample_uint32_t_cpp(SEXP matSEXP, SEXP probSEXP, SEXP seedSEXP) { +SEXP iterate_matrix_bernoulli_sample_uint32_t_cpp(SEXP mat, double prob, int seed, bool transpose); +RcppExport SEXP _BPCells_iterate_matrix_bernoulli_sample_uint32_t_cpp(SEXP matSEXP, SEXP probSEXP, SEXP seedSEXP, SEXP transposeSEXP) { BEGIN_RCPP Rcpp::RObject rcpp_result_gen; Rcpp::RNGScope rcpp_rngScope_gen; Rcpp::traits::input_parameter< SEXP >::type mat(matSEXP); Rcpp::traits::input_parameter< double >::type prob(probSEXP); Rcpp::traits::input_parameter< int >::type seed(seedSEXP); - rcpp_result_gen = Rcpp::wrap(iterate_matrix_bernoulli_sample_uint32_t_cpp(mat, prob, seed)); + Rcpp::traits::input_parameter< bool >::type transpose(transposeSEXP); + rcpp_result_gen = Rcpp::wrap(iterate_matrix_bernoulli_sample_uint32_t_cpp(mat, prob, seed, transpose)); return rcpp_result_gen; END_RCPP } // iterate_matrix_bernoulli_sample_float_cpp -SEXP iterate_matrix_bernoulli_sample_float_cpp(SEXP mat, double prob, int seed); -RcppExport SEXP _BPCells_iterate_matrix_bernoulli_sample_float_cpp(SEXP matSEXP, SEXP probSEXP, SEXP seedSEXP) { +SEXP iterate_matrix_bernoulli_sample_float_cpp(SEXP mat, double prob, int seed, bool transpose); +RcppExport SEXP _BPCells_iterate_matrix_bernoulli_sample_float_cpp(SEXP matSEXP, SEXP probSEXP, SEXP seedSEXP, SEXP transposeSEXP) { BEGIN_RCPP Rcpp::RObject rcpp_result_gen; Rcpp::RNGScope rcpp_rngScope_gen; Rcpp::traits::input_parameter< SEXP >::type mat(matSEXP); Rcpp::traits::input_parameter< double >::type prob(probSEXP); Rcpp::traits::input_parameter< int >::type seed(seedSEXP); - rcpp_result_gen = Rcpp::wrap(iterate_matrix_bernoulli_sample_float_cpp(mat, prob, seed)); + Rcpp::traits::input_parameter< bool >::type transpose(transposeSEXP); + rcpp_result_gen = Rcpp::wrap(iterate_matrix_bernoulli_sample_float_cpp(mat, prob, seed, transpose)); return rcpp_result_gen; END_RCPP } // iterate_matrix_bernoulli_sample_double_cpp -SEXP iterate_matrix_bernoulli_sample_double_cpp(SEXP mat, double prob, int seed); -RcppExport SEXP _BPCells_iterate_matrix_bernoulli_sample_double_cpp(SEXP matSEXP, SEXP probSEXP, SEXP seedSEXP) { +SEXP iterate_matrix_bernoulli_sample_double_cpp(SEXP mat, double prob, int seed, bool transpose); +RcppExport SEXP _BPCells_iterate_matrix_bernoulli_sample_double_cpp(SEXP matSEXP, SEXP probSEXP, SEXP seedSEXP, SEXP transposeSEXP) { BEGIN_RCPP Rcpp::RObject rcpp_result_gen; Rcpp::RNGScope rcpp_rngScope_gen; Rcpp::traits::input_parameter< SEXP >::type mat(matSEXP); Rcpp::traits::input_parameter< double >::type prob(probSEXP); Rcpp::traits::input_parameter< int >::type seed(seedSEXP); - rcpp_result_gen = Rcpp::wrap(iterate_matrix_bernoulli_sample_double_cpp(mat, prob, seed)); + Rcpp::traits::input_parameter< bool >::type transpose(transposeSEXP); + rcpp_result_gen = Rcpp::wrap(iterate_matrix_bernoulli_sample_double_cpp(mat, prob, seed, transpose)); return rcpp_result_gen; END_RCPP } @@ -2813,9 +2816,9 @@ static const R_CallMethodDef CallEntries[] = { {"_BPCells_iterate_matrix_mask_uint32_t_cpp", (DL_FUNC) &_BPCells_iterate_matrix_mask_uint32_t_cpp, 3}, {"_BPCells_iterate_matrix_mask_float_cpp", (DL_FUNC) &_BPCells_iterate_matrix_mask_float_cpp, 3}, {"_BPCells_iterate_matrix_mask_double_cpp", (DL_FUNC) &_BPCells_iterate_matrix_mask_double_cpp, 3}, - {"_BPCells_iterate_matrix_bernoulli_sample_uint32_t_cpp", (DL_FUNC) &_BPCells_iterate_matrix_bernoulli_sample_uint32_t_cpp, 3}, - {"_BPCells_iterate_matrix_bernoulli_sample_float_cpp", (DL_FUNC) &_BPCells_iterate_matrix_bernoulli_sample_float_cpp, 3}, - {"_BPCells_iterate_matrix_bernoulli_sample_double_cpp", (DL_FUNC) &_BPCells_iterate_matrix_bernoulli_sample_double_cpp, 3}, + {"_BPCells_iterate_matrix_bernoulli_sample_uint32_t_cpp", (DL_FUNC) &_BPCells_iterate_matrix_bernoulli_sample_uint32_t_cpp, 4}, + {"_BPCells_iterate_matrix_bernoulli_sample_float_cpp", (DL_FUNC) &_BPCells_iterate_matrix_bernoulli_sample_float_cpp, 4}, + {"_BPCells_iterate_matrix_bernoulli_sample_double_cpp", (DL_FUNC) &_BPCells_iterate_matrix_bernoulli_sample_double_cpp, 4}, {"_BPCells_iterate_matrix_rank_uint32_t_cpp", (DL_FUNC) &_BPCells_iterate_matrix_rank_uint32_t_cpp, 1}, {"_BPCells_iterate_matrix_rank_float_cpp", (DL_FUNC) &_BPCells_iterate_matrix_rank_float_cpp, 1}, {"_BPCells_iterate_matrix_rank_double_cpp", (DL_FUNC) &_BPCells_iterate_matrix_rank_double_cpp, 1}, diff --git a/r/src/bpcells-cpp/matrixIterators/BernoulliSample.h b/r/src/bpcells-cpp/matrixIterators/BernoulliSample.h index 7482c191..a67e26d9 100644 --- a/r/src/bpcells-cpp/matrixIterators/BernoulliSample.h +++ b/r/src/bpcells-cpp/matrixIterators/BernoulliSample.h @@ -17,19 +17,28 @@ namespace BPCells { // Keep each non-zero entry independently with probability `prob`, with the -// keep/drop decision a deterministic function of (seed, col, row) via splitmix64. -template class BernoulliSample : public MatrixLoaderWrapper { +// keep/drop decision a deterministic function of (seed, col, row) via splitmix64 and invariant to storage order. +template class BernoulliSample : public MatrixLoaderWrapper { private: uint64_t seed; uint32_t threshold; size_t loaded = 0; - // Hash algorithm - static inline uint64_t splitmix64(uint64_t x) { - x += 0x9e3779b97f4a7c15ULL; - x = (x ^ (x >> 30)) * 0xbf58476d1ce4e5b9ULL; - x = (x ^ (x >> 27)) * 0x94d049bb133111ebULL; - return x ^ (x >> 31); + // splitmix64 via Daniel Lemire's testingRNG: + // https://github.com/lemire/testingRNG/blob/master/source/splitmix64.h + static constexpr uint64_t GOLDEN_GAMMA = UINT64_C(0x9E3779B97F4A7C15); + + static inline uint64_t splitmix64_r(uint64_t *seed) { + uint64_t z = (*seed += GOLDEN_GAMMA); + z = (z ^ (z >> 30)) * UINT64_C(0xBF58476D1CE4E5B9); + z = (z ^ (z >> 27)) * UINT64_C(0x94D049BB133111EB); + return z ^ (z >> 31); + } + + // returns the value of splitmix64 "offset" steps from seed + static inline uint64_t splitmix64_stateless(uint64_t seed, uint64_t offset) { + seed += offset * GOLDEN_GAMMA; + return splitmix64_r(&seed); } public: @@ -45,20 +54,21 @@ template class BernoulliSample : public MatrixLoaderWrapper { bool load() override { loaded = 0; + const uint64_t nrow = this->loader->rows(); + const uint64_t ncol = this->loader->cols(); while (loaded == 0) { if (!this->loader->load()) return false; uint32_t *row_data = this->loader->rowData(); T *val_data = this->loader->valData(); size_t cap = this->loader->capacity(); - - // Calculate per-column hash - uint64_t col_mix = splitmix64(seed ^ (uint64_t(this->loader->currentCol()) << 32)); - + uint32_t col = this->loader->currentCol(); + for (size_t i = 0; i < cap; i++) { - // Calculate final hash with row number - uint32_t h = static_cast(splitmix64(col_mix ^ row_data[i])); - row_data[loaded] = row_data[i]; + uint32_t row = row_data[i]; + uint64_t offset = !Transpose ? (col * nrow + row) : (row * ncol + col); + uint32_t h = static_cast(splitmix64_stateless(seed, offset)); + row_data[loaded] = row; val_data[loaded] = val_data[i]; loaded += h < threshold; } diff --git a/r/src/matrix_utils.cpp b/r/src/matrix_utils.cpp index a533c4fa..caed0dab 100644 --- a/r/src/matrix_utils.cpp +++ b/r/src/matrix_utils.cpp @@ -347,30 +347,36 @@ SEXP iterate_matrix_mask_double_cpp(SEXP mat, SEXP mask, bool invert) { } // [[Rcpp::export]] -SEXP iterate_matrix_bernoulli_sample_uint32_t_cpp(SEXP mat, double prob, int seed) { - return make_unique_xptr>( - take_unique_xptr>(mat), - prob, - static_cast(static_cast(seed)) - ); +SEXP iterate_matrix_bernoulli_sample_uint32_t_cpp(SEXP mat, double prob, int seed, bool transpose) { + uint64_t s = static_cast(static_cast(seed)); + if (transpose) + return make_unique_xptr>( + take_unique_xptr>(mat), prob, s); + else + return make_unique_xptr>( + take_unique_xptr>(mat), prob, s); } // [[Rcpp::export]] -SEXP iterate_matrix_bernoulli_sample_float_cpp(SEXP mat, double prob, int seed) { - return make_unique_xptr>( - take_unique_xptr>(mat), - prob, - static_cast(static_cast(seed)) - ); +SEXP iterate_matrix_bernoulli_sample_float_cpp(SEXP mat, double prob, int seed, bool transpose) { + uint64_t s = static_cast(static_cast(seed)); + if (transpose) + return make_unique_xptr>( + take_unique_xptr>(mat), prob, s); + else + return make_unique_xptr>( + take_unique_xptr>(mat), prob, s); } // [[Rcpp::export]] -SEXP iterate_matrix_bernoulli_sample_double_cpp(SEXP mat, double prob, int seed) { - return make_unique_xptr>( - take_unique_xptr>(mat), - prob, - static_cast(static_cast(seed)) - ); +SEXP iterate_matrix_bernoulli_sample_double_cpp(SEXP mat, double prob, int seed, bool transpose) { + uint64_t s = static_cast(static_cast(seed)); + if (transpose) + return make_unique_xptr>( + take_unique_xptr>(mat), prob, s); + else + return make_unique_xptr>( + take_unique_xptr>(mat), prob, s); } // [[Rcpp::export]] diff --git a/r/tests/testthat/test-matrix_transforms.R b/r/tests/testthat/test-matrix_transforms.R index 65aa58e7..1736e680 100644 --- a/r/tests/testthat/test-matrix_transforms.R +++ b/r/tests/testthat/test-matrix_transforms.R @@ -419,61 +419,53 @@ test_that("matrix subtraction works", { expect_error(bp1 - bp1_t) }) -test_that("bernoulli_sample works", { +test_that("sample_bernoulli works", { set.seed(125124) m <- generate_sparse_matrix(50, 40, fraction_nonzero = 0.4, max_val = 100) bp <- as(m, "IterableMatrix") # prob == 1: no-op, returns the same object unchanged - expect_identical(bernoulli_sample(bp, 1, 0), bp) + expect_identical(sample_bernoulli(bp, 1, 0), bp) # Keep fraction: ~prob of non-zeros survive, within 4 sigma on ~2M entries m_big <- generate_sparse_matrix(2000, 2000, fraction_nonzero = 0.5, max_val = 10) bp_big <- as(m_big, "IterableMatrix") nnz_in <- length(m_big@x) - out <- as(bernoulli_sample(bp_big, 0.5, 17L), "dgCMatrix") + out <- as(sample_bernoulli(bp_big, 0.5, 17L), "dgCMatrix") nnz_out <- length(out@x) sigma <- sqrt(0.25 * nnz_in) expect_lt(abs(nnz_out - 0.5 * nnz_in), 4 * sigma) # Reproduciblity: same seed materializes identical output - y1 <- as(bernoulli_sample(bp, 0.3, 42L), "dgCMatrix") - y2 <- as(bernoulli_sample(bp, 0.3, 42L), "dgCMatrix") + y1 <- as(sample_bernoulli(bp, 0.3, 42L), "dgCMatrix") + y2 <- as(sample_bernoulli(bp, 0.3, 42L), "dgCMatrix") expect_identical(y1, y2) # Reverse column order matches forward iteration (exercises seekCol via MatrixSubset) - sampled <- bernoulli_sample(bp, 0.3, 42L) + sampled <- sample_bernoulli(bp, 0.3, 42L) reversed <- sampled[, rev(seq_len(ncol(bp)))] expect_identical(as(reversed, "dgCMatrix"), as(sampled, "dgCMatrix")[, rev(seq_len(ncol(bp)))]) - # Test hash keying to output columns (duplicated subset columns get independent draws) - col1 <- ceiling(ncol(bp) / 2) - dup_subset <- methods::new( - "MatrixSubset", - matrix = bp, - col_selection = rep(as.integer(col1), 3), - dim = c(nrow(bp), 3L), - dimnames = list(rownames(bp), NULL), - zero_dims = c(FALSE, FALSE) - ) - dup <- bernoulli_sample(dup_subset, 0.5, 99L) - dup_mat <- as(dup, "dgCMatrix") - expect_false(identical(dup_mat[, 1], dup_mat[, 2])) - expect_false(identical(dup_mat[, 1], dup_mat[, 3])) - expect_false(identical(dup_mat[, 2], dup_mat[, 3])) - # Sampling is value-independent: Shared keep/drop decisions regardless of type + # Sampling is value-independent for (tp in c("uint32_t", "float", "double")) { bp_t <- convert_matrix_type(bp, tp) - out <- bernoulli_sample(bp_t, 0.5, 13L) + out <- sample_bernoulli(bp_t, 0.5, 13L) expect_identical(matrix_type(out), tp) expect_equal( as(out, "dgCMatrix")@i, - as(bernoulli_sample(bp, 0.5, 13L), "dgCMatrix")@i + as(sample_bernoulli(bp, 0.5, 13L), "dgCMatrix")@i ) expect_equal( as(out, "dgCMatrix")@p, - as(bernoulli_sample(bp, 0.5, 13L), "dgCMatrix")@p + as(sample_bernoulli(bp, 0.5, 13L), "dgCMatrix")@p ) } + + # Storage-order invariant + bp_tso <- transpose_storage_order(bp) + expect_identical( + as(sample_bernoulli(bp, 0.5, 7L), "dgCMatrix"), + as(sample_bernoulli(bp_tso, 0.5, 7L), "dgCMatrix") + ) })