diff --git a/r/NAMESPACE b/r/NAMESPACE index 39e8b00f..ca3c3cc2 100644 --- a/r/NAMESPACE +++ b/r/NAMESPACE @@ -113,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 ac6ca2a3..ef04ef6e 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 `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 e79f0d11..382599ee 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, transpose) { + .Call(`_BPCells_iterate_matrix_bernoulli_sample_uint32_t_cpp`, mat, prob, seed, transpose) +} + +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, transpose) { + .Call(`_BPCells_iterate_matrix_bernoulli_sample_double_cpp`, mat, prob, seed, transpose) +} + 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..4280ef70 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, x@transpose) +}) + +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 +sample_bernoulli <- 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/sample_bernoulli.Rd b/r/man/sample_bernoulli.Rd new file mode 100644 index 00000000..83353242 --- /dev/null +++ b/r/man/sample_bernoulli.Rd @@ -0,0 +1,22 @@ +% Generated by roxygen2: do not edit by hand +% Please edit documentation in R/transforms.R +\name{sample_bernoulli} +\alias{sample_bernoulli} +\title{Randomly drop non-zero matrix entries given probability and seed} +\usage{ +sample_bernoulli(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..ca8d835f 100644 --- a/r/src/RcppExports.cpp +++ b/r/src/RcppExports.cpp @@ -2287,6 +2287,48 @@ 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, 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::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, 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::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, 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::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 +} // 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 +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, 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 new file mode 100644 index 00000000..a67e26d9 --- /dev/null +++ b/r/src/bpcells-cpp/matrixIterators/BernoulliSample.h @@ -0,0 +1,82 @@ +// 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 and invariant to storage order. +template class BernoulliSample : public MatrixLoaderWrapper { + private: + uint64_t seed; + uint32_t threshold; + size_t loaded = 0; + + // 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: + 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; + 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(); + uint32_t col = this->loader->currentCol(); + + for (size_t i = 0; i < cap; 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; + } + } + 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..caed0dab 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,39 @@ 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, 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, 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, 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_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..1736e680 100644 --- a/r/tests/testthat/test-matrix_transforms.R +++ b/r/tests/testthat/test-matrix_transforms.R @@ -418,3 +418,54 @@ test_that("matrix subtraction works", { bp1_t <- t(as(t(m1), "IterableMatrix")) expect_error(bp1 - bp1_t) }) + +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(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(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(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 <- 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)))]) + + + # Sampling is value-independent + for (tp in c("uint32_t", "float", "double")) { + bp_t <- convert_matrix_type(bp, tp) + out <- sample_bernoulli(bp_t, 0.5, 13L) + expect_identical(matrix_type(out), tp) + expect_equal( + as(out, "dgCMatrix")@i, + as(sample_bernoulli(bp, 0.5, 13L), "dgCMatrix")@i + ) + expect_equal( + as(out, "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") + ) +})