diff --git a/Project.toml b/Project.toml index 3481647..770428b 100644 --- a/Project.toml +++ b/Project.toml @@ -2,7 +2,7 @@ name = "RandomBooleanMatrices" uuid = "9ae346a0-3d16-5633-ad70-ddb60ab77eac" license = "MIT" authors = ["mkborregaard big(1)) + prob = ExactProblem(p, length(rowsums), _pascal(length(colsums)), table, L) + c = copy(conjugate) + total = _count(prob, c, Cursor(1, 1, 0, p[2], c[1])) + table[copy(conjugate)] = total + ExactCounts(prob, rowperm, colsums, conjugate, total) +end + +# --------------------------------------------------------------------------- +# Sampling +# --------------------------------------------------------------------------- + +""" + _sample_plan!(plan, ex, rng) + +Draw a uniform fixed-margin matrix as a `plan`, where `plan[row, block]` is the +number of ones the sorted `row` places in column-block `block`. Each branch of +the recursion is taken in proportion to the number of completions it leaves, so +the resulting matrix is exactly uniform. +""" +function _sample_plan!(plan::Matrix{Int}, ex::ExactCounts, rng) + prob = ex.problem + c = copy(ex.conjugate) + fill!(plan, 0) + _sample!(plan, prob, c, Cursor(1, 1, 0, prob.p[2], c[1]), ex.total, rng) + plan +end + +function _sample!(plan::Matrix{Int}, prob::ExactProblem, c::Vector{Int}, cur::Cursor, total::BigInt, rng) + if cur.placed == prob.p[cur.row] + cur.row == prob.nrows + 1 && return + return _sample!(plan, prob, c, _next_row(prob, c, cur), total, rng) + end + width, range = _block_range(prob, c, cur) + target = rand(rng, zero(total):(total - one(total))) + cumulative = zero(BigInt) + for s in range + c[cur.block] -= s + completions = _count(prob, c, _next_block(prob, c, cur, s)) + cumulative += prob.binomial[width+1, s+1] * completions + if cumulative > target + plan[cur.row, cur.block] = s + _sample!(plan, prob, c, _next_block(prob, c, cur, s), completions, rng) + c[cur.block] += s + return + end + c[cur.block] += s + end + error("exact sampler reached an inconsistent state") +end + +# Turn a sampled plan into the actual matrix: for each row and block, pick which +# of the equal-residual columns receive the row's ones, uniformly at random. +function _plan_to_columns!(columns, plan::Matrix{Int}, ex::ExactCounts, rng) + residual = copy(ex.colsums) + for col in columns + empty!(col) + end + @inbounds for row in 1:ex.problem.nrows + for level in 1:ex.problem.L-1 + block = findall(==(level), residual) + for col in sample(rng, block, plan[row, level], replace = false) + push!(columns[col], ex.rowperm[row]) + residual[col] -= 1 + end + end + end + columns +end + +""" + _exact_sample!(columns, ex, rng) + +Draw one uniform fixed-margin matrix into `columns` (row indices per column). +""" +function _exact_sample!(columns, ex::ExactCounts, rng) + plan = Matrix{Int}(undef, ex.problem.nrows, ex.problem.L) + _sample_plan!(plan, ex, rng) + _plan_to_columns!(columns, plan, ex, rng) +end + +function _exact!(m::SparseMatrixCSC{Bool, Int}, rng = Random.GLOBAL_RNG) + rowsums, colsums = _margins(m) + columns = [Int[] for _ in colsums] + _exact_sample!(columns, _exact_counts(rowsums, colsums), rng) + _writecols!(m, columns) +end + +# Generator draw: the cached `ExactCounts` is the generator state, so repeated +# draws reuse the (expensive) count instead of recomputing it. +function _draw!(m::SparseMatrixCSC{Bool, Int}, rng, ex::ExactCounts) + columns = [Int[] for _ in 1:size(m, 2)] + _exact_sample!(columns, ex, rng) + _writecols!(m, columns) +end diff --git a/src/sis.jl b/src/sis.jl new file mode 100644 index 0000000..831cbe1 --- /dev/null +++ b/src/sis.jl @@ -0,0 +1,337 @@ +# Sequential importance sampling of fixed-margin binary matrices, after +# +# Harrison & Miller (2013), "Importance sampling for weighted binary random +# matrices with specified margins". +# +# The matrix is built one column at a time. Each column is a binary vector with a +# prescribed number of ones, drawn from a fast approximation of the conditional +# distribution that (a) respects the column total, (b) keeps the remaining +# problem completable (Gale–Ryser), and (c) prefers configurations in proportion +# to a combinatorial estimate of how many completions they admit. A dynamic +# program turns those per-row preferences and feasibility bounds into exact +# forward sampling probabilities for the column. This implements the uniform case +# (every fixed-margin matrix targeted with equal weight), the natural counterpart +# to the curveball sampler. + +# --------------------------------------------------------------------------- +# Combinatorial approximation of the per-row odds of a one +# --------------------------------------------------------------------------- + +""" + ColumnScorer + +A combinatorial approximation that scores how likely each row is to carry a one +in the column currently being sampled. Concrete scorers are callable: given a +row's residual sum they return that row's (unconstrained) probability of a one. +New approximations (e.g. Greenhill–McKay for sparse margins) can be added as +further subtypes without touching the dynamic program. +""" +abstract type ColumnScorer end + +""" + Canfield(nrows, ncols, total, weightA) + +The Canfield–Greenhill–McKay (2008) approximation used by Harrison & Miller for +near-regular margins. `ncols` and `total` describe the columns and ones that +remain *after* the current column, and `weightA` collects the curvature term of +the asymptotic enumeration so that scoring a row is a single `exp`. +""" +struct Canfield <: ColumnScorer + nrows::Int + ncols::Int + total::Int + weightA::Float64 +end + +function Canfield(nrows::Int, ncols::Int, total::Int, sumsq::Int) + weightA = if total == 0 || nrows * ncols == total + 0.0 + else + η = nrows * ncols / (total * (nrows * ncols - total)) + ν = η * (sumsq - total^2 / ncols) + η * (1 - ν) / 2 + end + Canfield(nrows, ncols, total, weightA) +end + +# Probability that a row with residual sum `r` carries a one in this column. +# Rows that are empty (r == 0) or saturated (r == ncols + 1) fall out at 0 and 1. +function (s::Canfield)(r::Int) + odds = r * exp(s.weightA * (1 - 2 * (r - s.total / s.nrows))) + odds / (s.ncols + 1 - r + odds) +end + +# --------------------------------------------------------------------------- +# The residual problem: margins still to be filled, evolving column by column +# --------------------------------------------------------------------------- + +""" + SISResidual + +The part of the problem still to be sampled. Rows are tracked by `residual` sum +and kept in descending `order`; columns are summarised by their `conjugate` +(for feasibility) together with the running `total` of ones, sum of squares +`sumsq` (for the scorer) and number `ncols` still to place. +""" +mutable struct SISResidual + residual::Vector{Int} # residual row sums, indexed by original row + order::Vector{Int} # original row indices, sorted by descending residual + pos::Vector{Int} # inverse of `order`: where each row sits within it + conjugate::Vector{Int} # conjugate of the columns not yet sampled + total::Int # ones still to place + sumsq::Int # Σ colsum² over columns not yet sampled + ncols::Int # columns not yet sampled + nrows::Int +end + +function SISResidual(rowsums::Vector{Int}, colsums::Vector{Int}) + nrows = length(rowsums) + order = sortperm(rowsums, rev = true) + SISResidual(copy(rowsums), order, invperm(order), _conjugate(colsums, nrows), + sum(colsums), sum(abs2, colsums), length(colsums), nrows) +end + +# Remove a column of sum `k` from the column-side statistics, leaving the +# residual describing the *other* remaining columns (what a completion must fit). +function _advance!(res::SISResidual, k::Int) + @inbounds for level in 1:k + res.conjugate[level] -= 1 + end + res.total -= k + res.sumsq -= k^2 + res.ncols -= 1 + res +end + +# Decrement the rows that received a one and restore the descending order. Each +# of those rows dropped by exactly one, so it only has to slide past the block of +# rows still holding its previous value — an O(ones placed) repair rather than a +# full re-sort. Repairing the lowest-valued rows first keeps every swap local. +function _place!(res::SISResidual, rows) + order, pos, residual = res.order, res.pos, res.residual + @inbounds for r in rows + residual[r] -= 1 + end + @inbounds for idx in lastindex(rows):-1:firstindex(rows) + r = rows[idx] + i = pos[r] + value = residual[r] + j = i + while j < res.nrows && residual[order[j+1]] > value + j += 1 + end + if j > i + displaced = order[j] + order[i], order[j] = displaced, r + pos[r], pos[displaced] = j, i + end + end + res +end + +# --------------------------------------------------------------------------- +# Per-column dynamic program +# --------------------------------------------------------------------------- + +""" + SISWorkspace + +Reusable scratch space for sampling one column: the per-row probability `p` of a +one and the (possibly weighted) `one`-weight that biases it, the feasibility band +`lo:hi` on the partial column sum after each row, the forward transition +probabilities `S`, and two rolling buffers for the backward dynamic-programming +weights. For the uniform target `one == p`; a weight model scales `one` by its +per-row factor while leaving the zero-weight `1 - p` alone. +""" +struct SISWorkspace + p::Vector{Float64} + one::Vector{Float64} + lo::Vector{Int} + hi::Vector{Int} + S::Matrix{Float64} + gcur::Vector{Float64} + gnext::Vector{Float64} +end + +SISWorkspace(nrows::Int, maxcol::Int) = + SISWorkspace(Vector{Float64}(undef, nrows), Vector{Float64}(undef, nrows), + Vector{Int}(undef, nrows), Vector{Int}(undef, nrows), + Matrix{Float64}(undef, maxcol + 1, nrows), + Vector{Float64}(undef, maxcol + 2), Vector{Float64}(undef, maxcol + 2)) + +# Score every row for the current column: `p` is the combinatorial probability of +# a one, and `one` is that probability scaled by the weight model's factor (so the +# odds of a one become p·v : 1-p). A weight factor of `Inf` flags a row the +# remaining weights force to a one, which we encode as p = 1. +function _score!(ws::SISWorkspace, res::SISResidual, scorer::ColumnScorer, + model::WeightModel, pos::Int) + @inbounds for i in 1:res.nrows + row = res.order[i] + ρ = res.residual[row] + p = scorer(ρ) + v = _weightfactor(model, row, ρ, pos) + if isinf(v) + ws.p[i] = 1.0 + ws.one[i] = 1.0 + else + ws.p[i] = p + ws.one[i] = p * v + end + end + ws +end + +# Feasibility band for the partial column sum s_i = x_1 + … + x_i (rows in +# descending order). Gale–Ryser bounds it below; the column total and row count +# bound it above. `res.conjugate` is the conjugate of the *other* columns here. +function _band!(ws::SISWorkspace, res::SISResidual, k::Int) + m = res.nrows + rowsum = 0 # prefix of residual row sums + conjsum = 0 # prefix of the conjugate of the other columns + lo = 0 + @inbounds for i in 1:m + rowsum += res.residual[res.order[i]] + conjsum += res.conjugate[i] + lo = max(lo, rowsum - conjsum, k - (m - i)) + ws.lo[i] = lo + ws.hi[i] = min(k, i) + end + ws +end + +# Backward dynamic program (Harrison & Geman 2009): turn the per-row odds and the +# feasibility band into forward transition probabilities. `g[s]` is the total +# weight of feasible completions reaching column total `k` from partial sum `s`; +# `S[s+1, i]` is the resulting probability that row `i` is a one given partial sum +# `s` before it. Returns `false` if the column has no feasible completion — only +# possible when structural zeros forbid enough ones — so the caller can reject. +function _transitions!(ws::SISWorkspace, k::Int) + m = length(ws.p) + gnext, gcur, S = ws.gnext, ws.gcur, ws.S + fill!(gnext, 0.0) + fill!(gcur, 0.0) + gnext[k+1] = 1.0 # base case: all rows placed, sum == k + lonext, hinext = k + 1, k + 1 # index range where gnext is nonzero + locur, hicur = 1, 0 # gcur is empty (all zero) + @inbounds for i in m:-1:1 + lo = i == 1 ? 0 : ws.lo[i-1] # band on the partial sum before row i + hi = i == 1 ? 0 : ws.hi[i-1] + one = ws.one[i] # weight of a one (odds one : zero = one : 1-p) + zero = 1 - ws.p[i] # weight of a zero + for t in locur:hicur # clear only gcur's stale band + gcur[t] = 0.0 + end + total = 0.0 + for s in lo:hi + w1 = one * gnext[s+2] # row i is a one -> partial sum s+1 + w0 = zero * gnext[s+1] # row i is a zero -> partial sum s + weight = w0 + w1 + gcur[s+1] = weight + S[s+1, i] = weight > 0 ? w1 / weight : 0.0 + total += weight + end + total > 0 || return false # no completion: a structural-zero dead-end + invtotal = inv(total) # rescale to avoid underflow + for s in lo:hi + gcur[s+1] *= invtotal + end + locur, hicur = lo + 1, hi + 1 + gnext, gcur = gcur, gnext + lonext, hinext, locur, hicur = locur, hicur, lonext, hinext + end + true +end + +# Forward pass: walk the rows, drawing each from its transition probability. +# Accumulate the log proposal probability `logq` of the column and, for a weighted +# target, the log target weight `logp = Σ log w` over the ones placed. +function _draw_column!(rows, ws::SISWorkspace, order, k::Int, model::WeightModel, label::Int, rng) + s = 0 + logq = 0.0 + logp = 0.0 + @inbounds for i in eachindex(ws.p) + p = ws.S[s+1, i] + if rand(rng) < p + row = order[i] + push!(rows, row) + s += 1 + logq += log(p) + logp += _logweight(model, row, label) + s == k && break + else + logq += log(1 - p) + end + end + logq, logp +end + +""" + _sample_column!(rows, res, k, ws, model, pos, label, rng) + +Sample the column with label `label` (at sampling position `pos`) and total `k` +into `rows` (original row indices), update the residual problem, and return +`(logq, logp, feasible)`: the log proposal probability, the log target weight, and +whether a completion existed (`false` is a structural-zero dead-end to reject). +""" +function _sample_column!(rows, res::SISResidual, k::Int, ws::SISWorkspace, + model::WeightModel, pos::Int, label::Int, rng) + k == 0 && return 0.0, 0.0, true + _advance!(res, k) + _score!(ws, res, Canfield(res.nrows, res.ncols, res.total, res.sumsq), model, pos) + _band!(ws, res, k) + _transitions!(ws, k) || return 0.0, 0.0, false + logq, logp = _draw_column!(rows, ws, res.order, k, model, label, rng) + _place!(res, rows) + logq, logp, true +end + +# --------------------------------------------------------------------------- +# Driver +# --------------------------------------------------------------------------- + +""" + _sis(rowsums, colsums, model, rng) + +Sample a fixed-margin binary matrix by sequential importance sampling under the +target described by `model` (uniform, or a [`WeightMatrix`]). Returns the row +indices of each column with the log proposal probability `logq` and log target +weight `logp` of the whole matrix; the importance weight of the draw is +`exp(logp - logq)`. Columns are visited in the order the model prescribes (by +decreasing sum), which the authors find improves the approximation. With +structural zeros a draw can dead-end; it is returned with `logq = Inf` so that its +importance weight `exp(logp - logq)` is zero. +""" +function _sis(rowsums::Vector{Int}, colsums::Vector{Int}, model::WeightModel, rng) + res = SISResidual(rowsums, colsums) + ws = SISWorkspace(res.nrows, maximum(colsums, init = 0)) + columns = [Int[] for _ in colsums] + logq = 0.0 + logp = 0.0 + for (pos, label) in enumerate(_columnorder(model, colsums)) + dlogq, dlogp, feasible = _sample_column!(columns[label], res, colsums[label], ws, model, pos, label, rng) + feasible || return columns, Inf, logp + logq += dlogq + logp += dlogp + end + columns, logq, logp +end + +_sis(rowsums::Vector{Int}, colsums::Vector{Int}, rng) = + _sis(rowsums, colsums, UniformWeights(), rng) + +function _sis!(m::SparseMatrixCSC{Bool, Int}, rng = Random.GLOBAL_RNG) + rowsums, colsums = _margins(m) + columns, _, _ = _sis(rowsums, colsums, UniformWeights(), rng) + _writecols!(m, columns) +end + +""" + SISSampler() + +Generator state for the SIS method. SIS draws are independent, so the generator +keeps no precomputed state — each draw resamples the stored matrix from its +margins. +""" +struct SISSampler end + +_draw!(m::SparseMatrixCSC{Bool, Int}, rng, ::SISSampler) = _sis!(m, rng) diff --git a/src/weights.jl b/src/weights.jl new file mode 100644 index 0000000..fb44298 --- /dev/null +++ b/src/weights.jl @@ -0,0 +1,169 @@ +# Weighted targets for sequential importance sampling. +# +# The uniform sampler in sis.jl targets every fixed-margin matrix equally. +# Harrison & Miller (2013) generalise this to the weighted distribution +# +# P*(z) ∝ ∏_{ij} w[i,j]^z[i,j] over fixed-margin binary z, +# +# of which the uniform case (w ≡ 1) is a special case. The generalisation only +# changes the per-row odds of a one in the column being sampled: those odds are +# multiplied by a factor vᵢ built from the row's remaining weights (eq. 15). The +# factor uses the elementary symmetric polynomials of the weights over the +# not-yet-sampled columns, all carried in logs to keep their range in check. +# +# A draw is then no longer (near-)uniform, so it carries an importance weight +# f(z) = ∏ w^z / Q*(z); callers reweight Monte-Carlo estimates by it. +# +# Two refinements from the paper are included. Structural zeros (w[i,j] == 0, which +# forbid a one at that position) are handled: the symmetric polynomials and vᵢ drop +# those positions automatically (a zero weight is log -∞), and a draw that the +# zeros leave no way to complete is rejected with importance weight zero. And the +# weights are Sinkhorn-balanced into the canonical w̄ ∈ Λ(w) of eq. (16), which P* +# is invariant to but which makes the proposal scale-free and lower-variance. +# (Not yet done: the paper's variance-based column ordering.) + +""" + WeightModel + +How the SIS proposal biases a one beyond the combinatorial [`ColumnScorer`]. +`UniformWeights` adds nothing (the uniform target); `WeightMatrix` carries an +entrywise weight matrix and the factors needed to bias the proposal toward it. +New targets dispatch [`_weightfactor`](@ref) and [`_logweight`](@ref) here +without touching the dynamic program. +""" +abstract type WeightModel end + +struct UniformWeights <: WeightModel end + +""" + WeightMatrix(w, rowsums, colsums; canonicalize = true) + +The target P*(z) ∝ ∏ w[i,j]^z[i,j], for nonnegative weights `w` (zeros forbid a +one at that position). Stores the log weights (for the importance weight) and, +with columns in sampling order, the per-row count of available (positive) +positions and the log elementary symmetric polynomials +`loge[k+1, i, pos] = log eₖ(positive weights of row i from position pos onward)`, +from which the per-row factor vᵢ is read. The proposal is built from the +Sinkhorn-canonical `w̄` unless `canonicalize = false`; the importance weight always +uses the original `w`. +""" +struct WeightMatrix <: WeightModel + logw::Matrix{Float64} # log original weights [row, col] (the importance weight) + logwbar::Matrix{Float64} # log canonical weights [row, pos] (columns in sampling order) + navail::Matrix{Int} # navail[row, pos] = #positive weights at positions pos:end + order::Vector{Int} # column sampling order (column labels) + loge::Array{Float64, 3} # loge[k+1, row, pos] = log eₖ(weights at positions pos:end) +end + +function WeightMatrix(w::AbstractMatrix, rowsums::Vector{Int}, colsums::Vector{Int}; + canonicalize::Bool = true) + all(>=(0), w) || throw(ArgumentError("weights must be nonnegative")) + any(>(0), w) || throw(ArgumentError("weights must have a positive entry")) + logw = Matrix{Float64}(log.(float.(w))) # original weights, for the weight + wbar = canonicalize ? _canonical(w) : Matrix{Float64}(float.(w)) + order = sortperm(colsums, rev = true) + logwbar = log.(wbar)[:, order] # proposal weights, in sampling order + navail = _suffixcount(logwbar) + loge = _logesym(logwbar, maximum(rowsums, init = 0)) + WeightMatrix(logw, logwbar, navail, order, loge) +end + +# log(exp(a) + exp(b)), with -Inf absorbing as the log of zero. +function _logaddexp(a::Float64, b::Float64) + a == -Inf && return b + b == -Inf && return a + m = max(a, b) + m + log1p(exp(-abs(a - b))) +end + +# Number of positive weights in each row over column suffixes: `navail[i, pos]` +# counts positions pos:n. A weight is positive exactly when its log is finite. +function _suffixcount(logwbar::Matrix{Float64}) + m, n = size(logwbar) + navail = zeros(Int, m, n + 1) + @inbounds for pos in n:-1:1, i in 1:m + navail[i, pos] = navail[i, pos+1] + (logwbar[i, pos] > -Inf) + end + navail +end + +# Elementary symmetric polynomials of each row's weights over column suffixes, in +# logs. Built from the empty suffix backwards using eₖ(j:n) = eₖ(j+1:n) + +# w[j]·eₖ₋₁(j+1:n); `loge[k+1, i, pos]` covers columns at positions pos:n. A zero +# weight (log -∞) contributes nothing, so structural zeros drop out naturally. +function _logesym(logwbar::Matrix{Float64}, maxk::Int) + m, n = size(logwbar) + loge = fill(-Inf, maxk + 1, m, n + 1) + loge[1, :, :] .= 0.0 # e₀ ≡ 1 over any suffix + @inbounds for pos in n:-1:1, i in 1:m + lw = logwbar[i, pos] + for k in 1:maxk + loge[k+1, i, pos] = _logaddexp(loge[k+1, i, pos+1], lw + loge[k, i, pos+1]) + end + end + loge +end + +""" + _canonical(w; tol = 1e-10, maxiter = 1000) + +Sinkhorn-balance `w` into the canonical `w̄ ∈ Λ(w)` of Harrison & Miller (eq. 16): +the scaling `w̄ = αᵢ βⱼ wᵢⱼ` whose row and column sums equal the number of positive +entries in that row/column. `P*` is invariant to this scaling, but it makes the +proposal scale-invariant and tends to reduce its variance. +""" +function _canonical(w::AbstractMatrix; tol::Float64 = 1e-10, maxiter::Int = 1000) + m, n = size(w) + wbar = Matrix{Float64}(float.(w)) + ni = [count(>(0), view(wbar, i, :)) for i in 1:m] # positive entries per row + mj = [count(>(0), view(wbar, :, j)) for j in 1:n] # positive entries per column + for _ in 1:maxiter + for i in 1:m # scale each row to sum nᵢ + s = sum(view(wbar, i, :)) + s > 0 || continue + f = ni[i] / s + @inbounds for j in 1:n + wbar[i, j] *= f + end + end + drift = 0.0 + for j in 1:n # scale each column to sum mⱼ + s = sum(view(wbar, :, j)) + s > 0 || continue + drift = max(drift, abs(s - mj[j])) + f = mj[j] / s + @inbounds for i in 1:m + wbar[i, j] *= f + end + end + drift < tol && break + end + wbar +end + +# --------------------------------------------------------------------------- +# Interface consumed by the sampler (sis.jl). Plain-Int arguments keep these +# independent of the residual/workspace types. +# --------------------------------------------------------------------------- + +# Order in which to visit the columns (by decreasing sum). +_columnorder(::UniformWeights, colsums) = sortperm(colsums, rev = true) +_columnorder(model::WeightMatrix, colsums) = model.order + +# Factor multiplying a row's odds of a one in the column at sampling position +# `pos`, given the row's residual sum `ρ`. `0` forbids a one (structural zero); +# `Inf` forces one (the remaining positive weights leave no alternative). +_weightfactor(::UniformWeights, row, ρ, pos) = 1.0 +function _weightfactor(model::WeightMatrix, row, ρ, pos) + ρ == 0 && return 1.0 + logw = model.logwbar[row, pos] + logw == -Inf && return 0.0 # structural zero: no one here + logden = model.loge[ρ+1, row, pos+1] # log eρ over the columns after pos + logden == -Inf && return Inf # eρ == 0: the row is forced to a one here + lognum = model.loge[ρ, row, pos+1] # log eρ₋₁ + (model.navail[row, pos+1] - ρ + 1) / ρ * exp(logw + lognum - logden) +end + +# Contribution of a placed one at (row, column) to log ∏ w^z. +_logweight(::UniformWeights, row, label) = 0.0 +_logweight(model::WeightMatrix, row, label) = model.logw[row, label] diff --git a/test/runtests.jl b/test/runtests.jl index ea7c235..81399cc 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -3,8 +3,37 @@ using SparseArrays using Random using Test +mean(x) = sum(x) / length(x) + Random.seed!(1337) +# the row and column sums of a matrix, as plain vectors +margins(m) = (vec(sum(m, dims = 1)), vec(sum(m, dims = 2))) + +# brute-force count of the m×n binary matrices with the given margins +function countmatrices(rowsums, colsums) + m, n = length(rowsums), length(colsums) + count = 0 + for bits in 0:(2^(m * n) - 1) + A = [(bits >> ((j - 1) * m + (i - 1))) & 1 for i in 1:m, j in 1:n] + vec(sum(A, dims = 2)) == rowsums && vec(sum(A, dims = 1)) == colsums && (count += 1) + end + count +end + +# brute-force weighted count κ = Σ_z ∏ w^z over the binary matrices with the margins +function weightedcount(rowsums, colsums, w) + m, n = length(rowsums), length(colsums) + κ = 0.0 + for bits in 0:(2^(m * n) - 1) + A = [(bits >> ((j - 1) * m + (i - 1))) & 1 for i in 1:m, j in 1:n] + if vec(sum(A, dims = 2)) == rowsums && vec(sum(A, dims = 1)) == colsums + κ += prod(w[i, j]^A[i, j] for i in 1:m, j in 1:n) + end + end + κ +end + @testset "curveball" begin m = sprand(Bool, 8, 6, 0.2) m_old = copy(m) @@ -28,4 +57,114 @@ Random.seed!(1337) m4 = rand(rmg) @test m3 != m4 + + # the trade count is controllable; zero trades leaves the matrix untouched + m5 = sprand(Bool, 8, 6, 0.3) + @test randomize_matrix!(copy(m5), method = curveball, trades = 0) == m5 + @test matrixrandomizer(m5, method = curveball, trades = 7).state.trades == 7 + + # the generator warm-starts from an independent draw and preserves margins + csm5, rsm5 = margins(m5) + @test margins(rand(matrixrandomizer(m5, method = curveball, trades = 10))) == (csm5, rsm5) +end + +@testset "sis" begin + m = sprand(Bool, 8, 6, 0.3) + csm, rsm = margins(m) + randomize_matrix!(m, method = sis) + + @test margins(m) == (csm, rsm) + + m2 = rand(0:1, 6, 5) + csm, rsm = margins(m2) + rmg = matrixrandomizer(m2, method = sis) + m3 = rand(rmg) + m4 = rand(rmg) + + @test margins(m3) == (csm, rsm) # margins preserved + @test margins(m4) == (csm, rsm) + @test m3 != m4 # independent draws + + # a fully determined problem has a single solution, which must be returned + determined = sparse(Bool[1 1 1; 1 1 0; 1 0 0]) # unique matrix with margins (3,2,1) + @test countmatrices([3, 2, 1], [3, 2, 1]) == 1 + @test rand(matrixrandomizer(determined, method = sis)) == determined +end + +@testset "exact" begin + # counting agrees with brute-force enumeration + for (rs, cs) in [([2, 1, 1], [1, 2, 1]), ([2, 2, 1, 1], [2, 1, 2, 1]), + ([3, 1, 1, 1], [2, 2, 1, 1]), ([2, 2, 2], [2, 2, 2])] + @test RandomBooleanMatrices._exact_counts(rs, cs).total == countmatrices(rs, cs) + end + + m = sprand(Bool, 6, 5, 0.4) + csm, rsm = margins(m) + randomize_matrix!(m, method = exact) + @test margins(m) == (csm, rsm) + + m2 = rand(MersenneTwister(2718), 0:1, 5, 6) + csm, rsm = margins(m2) + rmg = matrixrandomizer(m2, MersenneTwister(2719), method = exact) # counts once, reused + draws = [rand(rmg) for _ in 1:10] + + @test all(d -> margins(d) == (csm, rsm), draws) + @test length(unique(draws)) > 1 # independent draws vary + + # a fully determined problem returns its single solution + determined = sparse(Bool[1 1 1; 1 1 0; 1 0 0]) + @test rand(matrixrandomizer(determined, method = exact)) == determined + + # the exact sampler is uniform: every matrix with these margins appears, and + # frequencies are close to equal (seeded rng keeps this reproducible) + rs, cs = [2, 2, 1, 1], [2, 1, 2, 1] + seed = sparse(Bool[1 0 1 0; 1 1 0 0; 0 0 1 0; 0 0 0 1]) + @test margins(seed) == (cs, rs) + n = countmatrices(rs, cs) + rmg = matrixrandomizer(seed, MersenneTwister(20240620), method = exact) + samples = [rand(rmg) for _ in 1:200n] + counts = collect(values(Dict(s => count(==(s), samples) for s in unique(samples)))) + @test length(counts) == n # full support covered + @test maximum(counts) < 2 * minimum(counts) # roughly uniform +end + +@testset "weighted sis" begin + m0 = sparse(Bool[1 0 1; 1 1 0; 0 1 0]) # margins (2,2,1) / (2,2,1) + cs, rs = margins(m0) + w = [1.0 2.0 0.5; 1.5 1.0 2.0; 0.5 1.0 1.0] + + sampler = importance_sampler(m0, w, MersenneTwister(11)) + draws = [rand(sampler) for _ in 1:100_000] + + # every draw is a fixed-margin matrix + @test all(d -> margins(d.matrix) == (cs, rs), draws) + @test draws[1].matrix != draws[2].matrix # independent draws + + # the mean importance weight is an unbiased estimate of the weighted count κ + weights = [exp(d.logweight) for d in draws] + @test isapprox(mean(weights), weightedcount(rs, cs, w), rtol = 0.03) + + # uniform weights reduce to the plain (unweighted) matrix count + uniform = importance_sampler(m0, ones(size(m0)), MersenneTwister(11)) + uweights = [exp(rand(uniform).logweight) for _ in 1:50_000] + @test isapprox(mean(uweights), countmatrices(rs, cs), rtol = 0.02) + + # canonicalisation changes the proposal but not the (unbiased) estimate + off = importance_sampler(m0, w, MersenneTwister(11); canonicalize = false) + offweights = [exp(rand(off).logweight) for _ in 1:100_000] + @test isapprox(mean(offweights), weightedcount(rs, cs, w), rtol = 0.03) + + # structural zeros: a zero weight forbids a one there, and κ is still recovered + wz = [0.0 2.0 1.0; 1.5 0.0 2.0; 1.0 0.5 0.0] # zero diagonal + derange = sparse(Bool[0 1 0; 0 0 1; 1 0 0]) # margins (1,1,1), off-diagonal + zsampler = importance_sampler(derange, wz, MersenneTwister(3)) + zdraws = [rand(zsampler) for _ in 1:100_000] + @test all(d -> all(i -> !d.matrix[i, i], 1:3), zdraws) # never a one on a zero + zweights = [exp(d.logweight) for d in zdraws] + @test isapprox(mean(zweights), weightedcount([1, 1, 1], [1, 1, 1], wz), rtol = 0.05) + + # input validation + @test_throws DimensionMismatch importance_sampler(m0, ones(2, 2)) + @test_throws ArgumentError importance_sampler(m0, fill(-1.0, 3, 3)) # negative weight + @test_throws ArgumentError importance_sampler(m0, zeros(3, 3)) # no positive entry end