add doubly stochastic matrices

[spinkney]

https://discourse.mc-stan.org/t/doubly-stochastic-matrices-in-stan/9987?u=spinkney

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
   
     @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_lp(matrix B) {
    int Nm1 = rows(B);
    int N = Nm1 + 1;
    matrix[N, N] X;  // the doubly stochastic matrix to output
    row_vector[N] c; // keeps track of the column sums of X
    real r = Nm1 > 0 ? B[1,1] : not_a_number(); // keeps track of the sum of a row of X
    if (cols(B) != Nm1) reject("B is not a square matrix")
    if (Nm1 == 0) return(rep_matrix(1.0, 1, 1));
    X[1,1] = r;
    c[1] = r;
    for (n in 2:Nm1) {
      real B_ = B[1, n];
      real x = B_ * (1 - r);
      if (B_ < 0) reject("B is negative in row 1 and column ", n);
      if (B_ > 1) reject("B is greater than 1.0 in row 1 and column ", n);
      X[1, n] = x;
      r += x;
      c[n] = x;
    }
    X[1, N] = 1 - r;
    c[N] = 1 - r;
    for (m in 2:Nm1) {
      real top_right = sum(c) - c[1]; // keeps track of the top right block of X
      r = 0;
      for (n in 1:Nm1) {
        real l = fmax(0, 1 - N  + n - r + top_right);
        real u = fmin(1 - r, 1 - c[n]);
        real d = u - l;
        real B_ = B[m, n];
        real x = l + B_ * d;
        if (B_ < 0) reject("B is negative in row ", m, " and column ", n);
        if (B_ > 1) reject("B is greater than 1.0 in row ", m, " and column ", n);
/*      
        UNCOMMENT THIS SECTION TO VERIFY THE ACCUMULATIONS USED IN THIS FUNCTION
        if (fabs(r - sum(sub_row(X, m, 1, n - 1))) > 1e-15) 
          reject("wrong row sum");
        if (fabs(c[n] - sum(sub_col(X, 1, n, m - 1))) > 1e-15)
          reject("wrong column sum");
        if (fabs(top_right - sum(block(X, 1, n + 1, m - 1, N - n))) > 1e-15) 
          reject("wrong sum of top right block")
*/      
        target += log(d); // due to Jacobian of map from B to X
        X[m,n] = x;
        r += x;
        top_right -= c[n + 1];
        c[n] += x;
      }
      X[m, N] = 1 - r;
      c[N] += 1 - r;
    }
    for (n in 1:N) X[N, n] = 1 - c[n];
    return X;
  }
}
data {
  int<lower = 1> N;
}
parameters {
  matrix<lower = 0, upper = 1>[N - 1, N - 1] B;
}
model {
  // implies: X ~ jointly uniform over Birkhoff polytope
  matrix[N, N] X = stick_break_lp(B);
/*
  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");
  }
*/
}