Add doubly stochastic matrix constraint

[WardBrian]

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

[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
  
     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");
  }
*/
}