Truncated Multivariate Normal Functions

Functions
real	multi_normal_cholesky_truncated_lpdf (vector u, vector mu, matrix L, vector lb, vector ub, vector lb_ind, vector ub_ind)
vector	multi_normal_cholesky_truncated_rng (vector u, vector mu, matrix L, vector lb, vector ub, vector lb_ind, vector ub_ind)
void	multi_normal_cholesky_truncated_lp (vector u, vector mu, matrix L, vector lb, vector ub, vector lb_ind, vector ub_ind)

Detailed Description

The truncated multivariate normal with mean vector \( \mu \) and variance-covariance matrix \( \Sigma \). Additionally, \( L \) is the Cholesky factor defined by \( \text{Chol}(\Sigma) = LL^T \).

The following derivation first appeared in an unpublished manuscript by Ben Goodrich, circa 2017.

Let a \( k \)-vector random variable \( \mathbf{y} \) be distributed multivariate normal then

\[ \mathbf{y} \stackrel{d}{=} \mathbf{\mu} + \mathbf{Lz(u)}, \]

where \( \mathbf{u} \) is a vector of uniform variates on \([0, 1]\) and \( \mathbf{z(u)} = \Phi^{-1}(\mathbf{u}) \).

Using the above definition, we may partition the standard uniform line so that the constaints are satisfied, however, we must supply a Jacobian adjustment due to the transformation of a uniform variate from \([0, 1]\) to the constrained space.

Given realiziations of the uniform variate \( \mathbf{u} = u_1, \ldots, u_k \) and a \( k \)-array of length-2 vector bounds

\[ \mathbf{b[lb, ub]} = [lb_1, ub_1], \ldots, [lb_k, ub_k], \]

we constrain \( \mathbf{u} \) to lie within the given bounds. That is, for each \( k \) a lower and upper bound must satisfy \( u_k^* = \Phi(b_k - \mu_k + L_{k, 1:k-1} z_{k - 1}) \). The random variate \( \mathbf{u} \) is then constrained to fall within those bounds. The new uniform variate is constrained to lie within the given bounds \( v_k \sim \mathcal{U}(u_k^*[1], u_k^*[2]) \) by

\[ v_k = u_k^*[1] + (u_k^*[2] - u_k^*[1]) u_k. \]

This implies that \( \frac{\partial u_k}{\partial v_k} = (u_k^*[2] - u_k^*[1]) \) and the absolute value of the log Jacobian is \(\ln(u_k^*[2] - u_k^*[1]) \).

The final realizations of the truncated multivariate normal vector is given by

\[ \mathbf{r} = \mu + L \Phi^{-1}(\mathbf{v}). \]

 
real multi_normal_cholesky_truncated_lpdf(vector u, vector mu, matrix L, vector lb,
                                          vector ub, vector lb_ind, vector ub_ind) {
  int K = rows(u);
  vector[K] z;
  real lp = 0;
  
  for (k in 1 : K) {
    // if kth u is unbounded
    // else kth u has at least one bound
    if (lb_ind[k] == 0 && ub_ind[k] == 0) 
      z[k] = inv_Phi(u[k]);
    else {
      int km1 = k - 1;
      real v;
      real z_star;
      real logd;
      row_vector[2] log_ustar = [negative_infinity(), 0]; // [-Inf, 0] = log([0,1])
      real constrain = mu[k] + ((k > 1) ? L[k, 1 : km1] * head(z, km1) : 0);
      
      // obtain log of upper and lower bound (if applicable)
      if (lb_ind[k] == 1) 
        log_ustar[1] = normal_lcdf((lb[k] - constrain) / L[k, k] | 0.0, 1.0);
      if (ub_ind[k] == 1) 
        log_ustar[2] = normal_lcdf((ub[k] - constrain) / L[k, k] | 0.0, 1.0);
      
      // update log gradient and z  
      logd = log_diff_exp(log_ustar[2], log_ustar[1]);
      v = exp(log_sum_exp(log_ustar[1], log(u[k]) + logd)); // v = ustar[1] + (ustar[2] - ustar[1]) * u[k] ~ U(ustar[1], ustar[2])
      z[k] = inv_Phi(v); // z ~ TN
      lp += logd; // increment by log gradient
    }
  }
  return lp;
}
 
vector multi_normal_cholesky_truncated_rng(vector u, vector mu, matrix L, vector lb,
                                           vector ub, vector lb_ind, vector ub_ind) {
  int K = rows(u);
  vector[K] z;
  
  for (k in 1 : K) {
    // if kth u is unbounded
    // else kth u has at least one bound
    if (lb_ind[k] == 0 && ub_ind[k] == 0) 
      z[k] = inv_Phi(u[k]);
    else {
      int km1 = k - 1;
      real v;
      real z_star;
      real logd;
      row_vector[2] log_ustar = [negative_infinity(), 0]; // [-Inf, 0] = log([0,1])
      real constrain = mu[k] + ((k > 1) ? L[k, 1 : km1] * head(z, km1) : 0);
      
      // obtain log of upper and lower bound (if applicable)
      if (lb_ind[k] == 1) 
        log_ustar[1] = normal_lcdf((lb[k] - constrain) / L[k, k] | 0.0, 1.0);
      if (ub_ind[k] == 1) 
        log_ustar[2] = normal_lcdf((ub[k] - constrain) / L[k, k] | 0.0, 1.0);
      
      // update log gradient and z  
      logd = log_diff_exp(log_ustar[2], log_ustar[1]);
      v = exp(log_sum_exp(log_ustar[1], log(u[k]) + logd)); // v = ustar[1] + (ustar[2] - ustar[1]) * u[k] ~ U(ustar[1], ustar[2])
      z[k] = inv_Phi(v); // z ~ TN
    }
  }
  return mu + L * z;
}
 
void multi_normal_cholesky_truncated_lp(vector u, vector mu, matrix L, vector lb,
                                        vector ub, vector lb_ind, vector ub_ind) {
  int K = rows(u);
  vector[K] z;
  
  for (k in 1 : K) {
    // if kth u is unbounded
    // else kth u has at least one bound
    if (lb_ind[k] == 0 && ub_ind[k] == 0) 
      z[k] = inv_Phi(u[k]);
    else {
      int km1 = k - 1;
      real v;
      real z_star;
      real logd;
      row_vector[2] log_ustar = [negative_infinity(), 0]; // [-Inf, 0] = log([0,1])
      real constrain = mu[k] + ((k > 1) ? L[k, 1 : km1] * head(z, km1) : 0);
      
      // obtain log of upper and lower bound (if applicable)
      if (lb_ind[k] == 1) 
        log_ustar[1] = normal_lcdf((lb[k] - constrain) / L[k, k] | 0.0, 1.0);
      if (ub_ind[k] == 1) 
        log_ustar[2] = normal_lcdf((ub[k] - constrain) / L[k, k] | 0.0, 1.0);
      
      // update log gradient and z  
      logd = log_diff_exp(log_ustar[2], log_ustar[1]);
      v = exp(log_sum_exp(log_ustar[1], log(u[k]) + logd)); // v = ustar[1] + (ustar[2] - ustar[1]) * u[k] ~ U(ustar[1], ustar[2])
      z[k] = inv_Phi(v); // z ~ TN
      target += logd; // increment by log gradient
    }
  }
}
 

Function Documentation

multi_normal_cholesky_truncated_lp()

void multi_normal_cholesky_truncated_lp	(	vector	u,
		vector	mu,
		matrix	L,
		vector	lb,
		vector	ub,
		vector	lb_ind,
		vector	ub_ind
	)

Truncated Multivariate Increment Log-Probability

Copyright

Ben Goodrich, Ethan Alt, Sean Pinkney 2017, 2021, 2021
https://groups.google.com/g/stan-users/c/GuWUJogum1o/m/LvxjlUBnBwAJ
Accessed and modified Apr. 18, 2021

Parameters

u	Vector number on [0,1], not checked but function will return NaN
mu	Vector
L	Cholesky Factor Corr
lb	Vector of lower bounds
ub	Vector of upper bounds
lb_ind	Vector indicator if there is an lower bound
ub_ind	Vector indicator if there is an upper bound

multi_normal_cholesky_truncated_lpdf()

real multi_normal_cholesky_truncated_lpdf	(	vector	u,
		vector	mu,
		matrix	L,
		vector	lb,
		vector	ub,
		vector	lb_ind,
		vector	ub_ind
	)

Truncated Multivariate Log Probability Density

Copyright

Ben Goodrich, Ethan Alt, Sean Pinkney 2017, 2021, 2021
https://groups.google.com/g/stan-users/c/GuWUJogum1o/m/LvxjlUBnBwAJ
Accessed and modified Apr. 18, 2021

Parameters

u	Vector number on [0,1], not checked but function will return NaN
mu	Vector
L	Cholesky Factor Corr
lb	Vector of lower bounds
ub	Vector of upper bounds
lb_ind	Vector indicator if there is an lower bound
ub_ind	Vector indicator if there is an upper bound

Returns

log density

multi_normal_cholesky_truncated_rng()

vector multi_normal_cholesky_truncated_rng	(	vector	u,
		vector	mu,
		matrix	L,
		vector	lb,
		vector	ub,
		vector	lb_ind,
		vector	ub_ind
	)

Truncated Multivariate Random Number Generator

Copyright

Ben Goodrich, Ethan Alt, Sean Pinkney 2017, 2021, 2021
https://groups.google.com/g/stan-users/c/GuWUJogum1o/m/LvxjlUBnBwAJ
Accessed and modified Apr. 18, 2021

Parameters

u	Vector number on [0,1], not checked but function will return NaN
mu	Vector
L	Cholesky Factor Corr
lb	Vector of lower bounds
ub	Vector of upper bounds
lb_ind	Vector indicator if there is an lower bound
ub_ind	Vector indicator if there is an upper bound

Returns

vector of truncated random normal variates