Repository navigation
Add doubly stochastic matrix constraint #3372
Description
Activity
functions { /* Returns a doubly stochastic matrix by transforming pre-parameters based on Section 3.1 of http://proceedings.mlr.press/v84/linderman18a/linderman18a.pdf and has the side-effect of modifying target to account for the Jacobian based on Section B.2 of http://proceedings.mlr.press/v84/linderman18a/linderman18a-supp.pdf Originally coded by Ben Goodrich and posted at https://discourse.mc-stan.org/t/doubly-stochastic-matrices-in-stan/9987 @param B matrix (square) of pre-parameters all on (0,1) @return matrix that is doubly stochastic @throws if B is not square, is 0x0, or its elements are not all on (0,1) */ matrix stick_break_jacobian(matrix B) { int Nm1 = rows(B); int N = Nm1 + 1; matrix[N, N] X; // doubly stochastic output row_vector[N] c; // running column sums for rows 1:(N-1) if (cols(B) != Nm1) reject("B is not a square matrix"); if (Nm1 == 0) return(rep_matrix(1.0, 1, 1)); c = rep_row_vector(0.0, N); // One pass over rows 1:(N-1). Row 1 naturally falls out of the same equations. for (m in 1:Nm1) { real r = 0.0; // running sum in row m real top_right = (m - 1) - c[1]; for (n in 1:Nm1) { real l = fmax(0.0, 1 - N + n - r + top_right); real u = fmin(1.0 - r, 1.0 - c[n]); real d = u - l; real x = l + B[m, n] * d; jacobian += log(d); // n=1 of row 1 contributes log(1)=0 X[m, n] = x; r += x; top_right -= c[n + 1]; c[n] += x; } X[m, N] = 1.0 - r; c[N] += 1.0 - r; } for (n in 1:N) X[N, n] = 1.0 - c[n]; return X; } } data { int<lower = 1> N; } parameters { matrix<lower = 0, upper = 1>[N - 1, N - 1] B; } transformed parameters { matrix[N, N] X = stick_break_jacobian(B); } model { // implies: X ~ jointly uniform over Birkhoff polytope /* UNCOMMENT THIS SECTION TO VERIFY X IS A DOUBLY STOCHASTIC MATRIX vector[N] rowsums = X * rep_vector(1, N); row_vector[N] colsums = rep_row_vector(1, N) * X; for (n in 1:N) { if (fabs(1 - rowsums[n]) > 1e-15) reject("rows not normalized correctly"); if (fabs(1 - colsums[n]) > 1e-15) reject("columns not normalized correctly"); } */ }
Discussed at StanCon, this is an extension of row/column stochastic matrices where the constraint applies in both directions simultaneously.
@SteveBronder said he had some use cases, and @spinkney said he knew a transform we could implement