functions{ vector[] make_stuff(vector mu, matrix L, vector b, vector s, vector u) { int K = rows(mu); vector[K] d; vector[K] z; vector[K] out[2]; for (k in 1:K) { int km1 = k - 1; if (s[k] != 0) { real z_star = (b[k] - (mu[k] + ((k > 1) ? L[k,1:km1] * head(z, km1) : 0))) / L[k,k]; real v; real u_star = Phi(z_star); if (s[k] == -1) { v = u_star * u[k]; d[k] = u_star; } else { d[k] = 1 - u_star; v = u_star + d[k] * u[k]; } z[k] = inv_Phi(v); } else { z[k] = inv_Phi(u[k]); d[k] = 1; } } out[1] = z; out[2] = d; return out; } /*Functions to create all three covariance (sub) matrices (fourth is just cov_quad_exp) */ real K22(real s, real sp, real alpha, real rho){ real sq_delta = square(s-sp); real sq_rho = square(rho); real cov = square(alpha) * exp(-sq_delta/(2*sq_rho)) * 1/square(sq_rho) * ( square(sq_delta)/square(sq_rho) - (6*sq_delta/sq_rho) + 3 ); return(cov); } matrix K22_cov_mat(real[] s, real alpha, real rho){ int K = size(s); matrix[K, K] X; real diag_entry = 3*square(alpha/square(rho)); for (i in 1:K) { X[i, i] = diag_entry; for (j in 1:K) { X[i, j] = K22(s[i], s[j], alpha, rho); X[j, i] = X[i, j]; } } return(X); } real K20(real t, real s, real alpha, real rho){ real cov; // computation in log-space for numerical stability real sq_delta = square(t-s); real sq_rho = square(rho); cov = square(alpha) * exp(-sq_delta/(2*sq_rho)) * ( sq_delta / square(sq_rho) - 1/sq_rho); return(cov); } matrix K20_cov_mat(real[] x, real[] s, real alpha, real rho){ int N = size(x); int K = size(s); matrix[N, K] Y; for (i in 1:N) { for (j in 1:K) { Y[i, j] = K20(x[i], s[j], alpha, rho); } } return(Y); } matrix ffprime_cov_mat(real [] x, real[] s, real alpha, real rho){ int N = size(x); int K = size(s); int J = N + K; matrix[N, N] A; matrix[N, K] B; matrix[K, K] C; matrix[J, J] M; // | A, B| // | B',C| A = cov_exp_quad(x, alpha, rho); B = K20_cov_mat(x, s, alpha, rho); C = K22_cov_mat(s, alpha, rho); /*Filling in the matrix*/ M[1:N, 1:N] = A; M[1:N, (N+1):J] = B; M[(N + 1):J, (N+1):J] = C; M[(N + 1):J, 1:N] = B'; return(M); } matrix full_cov_mat(real [] x, real[] s, real[] x_pred, real alpha, real rho){ int N = size(x); int K = size(s); int N_pred = size(x_pred); int W = N + K + N_pred; matrix[W, W] M; matrix[N_pred, K] Q = K20_cov_mat(x_pred, s, alpha, rho); /* building sub-matrices */ // | (f, f) (f,f'') (f, f_p) | // | (f'',f) (f'', f'') (f'', f_p) | // | (f_p, f) (f_p, f'') (f_p, f_p) | M[1:(N + K), 1:(N + K)] = ffprime_cov_mat(x, s, alpha, rho); M[(N + K + 1):W, 1:N] = cov_exp_quad(x_pred, x, alpha, rho); M[(N + K + 1):W, (N + 1):(N + K)] = Q; M[1:N, (N + K + 1):W] = cov_exp_quad(x, x_pred, alpha, rho); M[(N + 1):(N + K), (N + K + 1):W] = Q'; M[(N + 1 + K):W, (N + 1 + K):W] = cov_exp_quad(x_pred, x_pred, alpha, rho); return(M); } } data { int K; real x[K]; vector[K] y; } transformed data { vector[K*2] b = rep_vector(0, 2*K); // lower or upper bound real alpha = 1.0; real rho = 1.0; // s[k] == 0 implies no constraint; otherwise // s[k] == -1 -> b[k] is an upper bound // s[k] == +1 -> b[k] is a lower bound vector[K*2] s = append_row(rep_vector(1, K), rep_vector(0, K)); matrix[K*2,K*2] Sigma = append_row( append_col(K22_cov_mat(x, alpha, rho), K20_cov_mat(x, x, alpha, rho)) , append_col( K20_cov_mat(x, x, alpha, rho), cov_exp_quad(x, alpha, rho)) ) + diag_matrix(rep_vector(1e-5,2*K)); matrix[K*2,K*2] L = cholesky_decompose(Sigma); } parameters { vector[K*2] u; real sigma; real mu; } transformed parameters { vector[K*2] y_g = mu + L * make_stuff(rep_vector(mu, 2*K), L, b, s, u)[1]; } model { target += log(make_stuff(rep_vector(mu, K), L, b, s, u)[2]); // Jacobian adjustments // implicit: u ~ uniform(0,1) sigma ~ cauchy(0, 1); y ~ normal(y_g[K+1:2*K], sigma); } generated quantities { matrix[K*2,K*2] Sigma_gen = Sigma; }