[r][cpp] add Bernoulli Sampling

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)

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).

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)
 }

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))
+}

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},

r/src/bpcells-cpp/matrixIterators/BernoulliSample.h

@@ -0,0 +1,82 @@
+// Copyright 2026 BPCells contributors
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or https://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+#pragma once
+
+#include <cstdint>
+#include <memory>
+#include <stdexcept>
+
+#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 <typename T, bool Transpose = false> class BernoulliSample : public MatrixLoaderWrapper<T> {
+  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<MatrixLoader<T>> &&loader, double prob, uint64_t seed)
+        : MatrixLoaderWrapper<T>(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<uint32_t>(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<uint32_t>(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

r/src/matrix_utils.cpp

@@ -10,6 +10,7 @@
 #define RCPP_NO_SUGAR
 #include <Rcpp.h>
 #include <RcppEigen.h>
+#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<uint64_t>(static_cast<uint32_t>(seed));
+    if (transpose)
+        return make_unique_xptr<BernoulliSample<uint32_t, true>>(
+            take_unique_xptr<MatrixLoader<uint32_t>>(mat), prob, s);
+    else
+        return make_unique_xptr<BernoulliSample<uint32_t, false>>(
+            take_unique_xptr<MatrixLoader<uint32_t>>(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<uint64_t>(static_cast<uint32_t>(seed));
+    if (transpose)
+        return make_unique_xptr<BernoulliSample<float, true>>(
+            take_unique_xptr<MatrixLoader<float>>(mat), prob, s);
+    else
+        return make_unique_xptr<BernoulliSample<float, false>>(
+            take_unique_xptr<MatrixLoader<float>>(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<uint64_t>(static_cast<uint32_t>(seed));
+    if (transpose)
+        return make_unique_xptr<BernoulliSample<double, true>>(
+            take_unique_xptr<MatrixLoader<double>>(mat), prob, s);
+    else
+        return make_unique_xptr<BernoulliSample<double, false>>(
+            take_unique_xptr<MatrixLoader<double>>(mat), prob, s);
+}
+
 // [[Rcpp::export]]
 SEXP iterate_matrix_rank_uint32_t_cpp(SEXP matrix) {
     return make_unique_xptr<ColwiseRank<uint32_t>>(take_unique_xptr<MatrixLoader<uint32_t>>(matrix)

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")
+    )
+})