TMB implementations of LKJ, Wishart, and inverse Wishart distributions

[jaganmn]

Having dlkj, dwishart, and dinvwishart in one of the TMB namespaces would enable users to assign likelihoods to correlation or covariance matrices. I have implemented these here. Documentation and tests against independent R implementations are there, too.

Following UNSTRUCTURED_CORR_t, users would pass a vector theta or c(log_sd, theta) instead of the full correlation or covariance matrix. This simplifies computation quite a bit.

Anyway, you are welcome to use or modify the code. At your discretion, I could also work on a pull request.

[kaskr]

Nice work @jaganmn - thanks for sharing! I'll have it mind for future addition to TMB.

[jaganmn]

I wonder if this should happen soon-ish, so that glmmTMB can get the density functions directly from TMB instead of copying the code into its own sources. Probably it is not the only TMB-using package that would want dlkj, etc. @bbolker

[kaskr]

@jaganmn I was about to copy your code, but then got confused over a couple of details...

Given that your version of dwishart is formulated in transformed parameter space, did you remember to add the Jacobian adjustment of the transformation? A quick test of your R-code suggests perhaps not:

integrate(Vectorize(function(x)dwishart(x,df=2,scale=1)),-100,100)
## 6.838354 with absolute error < 7.9e-05

[jaganmn]

The d* functions in the header take as argument the transformed parameter and compute the probability density in the space of the untransformed parameter. The distribution of the transformed parameter would not be called an LKJ/Wishart/inverse Wishart distribution AFAIK.

I think that your integral approximation exceeds unity due to numerical instability for very small nonzero standard deviations. For example, in the one-dimensional case with scale = 0 (indicating an identity matrix of size one), the Wishart distribution reduces to a chi-squared distribution and so we can compare with dchisq:

set.seed(0)

x1 <- rnorm(10L)
dwishart1 <- Vectorize(function(x) dwishart(x, df = 2, scale = 0))
dchisq1 <- function(x) dchisq(exp(2 * x), df = 2)

max(abs(dwishart1(x1) - dchisq1(x1)))
## [1] 8.673617e-19
integrate(dwishart1, -100, 100)
## 50.02898 with absolute error < 0.0037
integrate(dchisq1, -100, 100)
## 50.02898 with absolute error < 0.0037

x2 <- rlnorm(10L)
dwishart2 <- Vectorize(function(x) dwishart(0.5 * log(x), df = 2, scale = 0))
dchisq2 <- function(x) dchisq(x, df = 2)

max(abs(dwishart2(x2) - dchisq2(x2)))
## [1] 5.551115e-17
integrate(dwishart2, exp(-10), exp(10))
## 0.9999773 with absolute error < 6.1e-07
integrate(dchisq2, exp(-10), exp(10))
## 0.9999773 with absolute error < 6.1e-07

One can think about ways to improve the speed/stability, e.g., constant additive terms in the log densities can typically be dropped (as in dlkj) and expressions like atomic::matinv(L_S) and invL_S.transpose() * invL_S can probably be optimized ...

[kaskr]

The point of my example was only to demonstrate that the density doesn't integrate to one.

[jaganmn]

Er, right, yes, I managed to confuse myself, not having read carefully enough (your comment and my own documentation) ...

As you say, dwishart does not integrate to one because it does not attempt to compute the Jacobian of the inverse transformation. My remark about numerical instability is nonsense. Indeed:

dwishart1.j <- Vectorize(function(x) dwishart(x, df = 2, scale = 0) * 2 * exp(2 * x))
integrate(dwishart1.j, -100, 100)
## 1 with absolute error < 3.6e-07

AFAICT the current behaviour is "correct" as long as it is clearly documented what is being computed, namely a probability density in the space of the inverse transformed parameter.

[kaskr]

Conventionally, all densities in TMB integrate to one, and I'm quite determined not to break that pattern.

Beyond Bayesian priors the two typical use cases are:

1. You have Wishart data and want to estimate the parameters.
2. You want to assign a Wishart distribution to some random effects.

I claim that your version is not directly applicable in either of these cases.
First case requires you to transform the data in order to feed it into your dwishart.
Same for the second case if you work in constrained parameter space. In unconstrained parameter space one would need to add the jacobian adjustment and apply a post transformation.

So, to follow the TMB convention there should probably be two versions:

// Constrained parameter space
// (MJ version with added parameter transform)
dwishart(...);
// Unconstrained parameter space.
// (MJ version with added jacobian adjustment)
dwishart_unconstrained(...);

[jaganmn]

OK. I'll work on those alternate implementations "soon" and update here.

[bbolker]

FWIW I continue to be interested in this ...

[jaganmn]

Not forgotten ... maybe I'll work on this during the glmmTMB hack.

[bbolker]

This came up in conversation recently ... FWIW I think the dwishart_unconstrained (or dlkj_unconstrained) would be more immediately useful, as the current use case is applying priors to covariance matrices (where we would typically expect users to specify parameters in the untransformed/unconstrained space ...)

I'm thinking a bit about "parameter-expanded" priors as well (which might only be useful for direct comparisons with MCMCglmm results) — these are defined in the MCMCglmm course notes, and I think also in the Gelman and Hill Applied Regression Modeling book (they reduce to half-$t$ in the central, one-dimensional case, but I don't know offhand if there's a useful closed form for the multidimensional case ...)

[michaelspence]

Thanks a lot for this and this was a great help as I wanted to use the LKJ distribution in TMB. I did notice a couple of things and thought it might be helpful to flag them.

1. I think the normalising constant in the LKJ distribution is incorrect. The LKJ paper and Wikipedia say the LKJ distribution is

f(X,\eta)=C(\eta)det(X)^{\eta-1} (eq 1),

with C(\eta) being a function that contains beta functions and other things. I think the distribution should be:

f(X,\eta)=1/C(\eta)det(X)^{\eta-1} (eq 2).

See the attached document (LKJ_distribution.pdf) that shows this is true for 2 dimensional correlation matrix with eta=1 and eta=2 in Section 2. I also believe that in stan they use the same normalising constant as in eq 2.

2. I think the log determinant of Jacobian for the transformation should be:

-1/2(n+1)\sum_{i=1}^{n}log(1+\sum_{j=1}^{n}y_{i,j}^2)

The attached document shows in the n=3 case this is correct (Section 3) this is correct. In Section 4 tries to find the Jacobian in the general case for any n. We wrote the determinant of the Jacobian as a recursive formula, as shown in equation 42, however, we struggled to calculate the determinant of an n by n matrix in general (denoted N_n in the document).

The attached cpp file (lkj_stuff.cpp) shows the proposed changes to dlkj, specifically:

template<class Type>
Type dlkj_unconstrained(const vector<Type> &x, Type shape, int give_log = 0)
{
	Eigen::Index len = x.size();
	if (len == 0)
	{
		Rf_error("%s: length of '%s' is zero",
		         __func__, "x");
	}
	if (len > CHOOSE_MAX)
	{
		Rf_error("%s: length of '%s' (%lld) exceeds maximum %llu",
		         __func__, "x", static_cast<long long int>(len), CHOOSE_MAX);
	}
	Eigen::Index n = static_cast<Eigen::Index>(0.5 * (1.0 + sqrt(1.0 + 8.0 * len)));
	Type eta_minus_one = exp(shape) - Type(1.0);

	matrix<Type> R(n, n);
	for (Eigen::Index j = 0, k = 0; j < n; ++j)
	{
		for (Eigen::Index i = 0; i < j; ++i, ++k)
		{
			R(i, j) = x(k);
		}
		R(j, j) = Type(1.0);
		for (Eigen::Index i = j + 1; i < n; ++i)
		{
			R(i, j) = Type(0.0);
		}
	}

	Type log_det_X = -(R.array() * R.array()).colwise().sum().log().sum();

	Type log_res = Type(0.0);
	/* Terms from conventional density function : */
	log_res += eta_minus_one * log_det_X -// changed to minus
		Type(M_LN2 * n * (n - 1)) * (eta_minus_one + Type((n + n - 1) / 6.0));
	for (Eigen::Index j = 1; j < n; ++j)
	{
		Type a = Type(2.0) * eta_minus_one + Type(j + 1);
		log_res -= Type(2 * j) * lgamma(Type(0.5) * a) - Type(j) * lgamma(a); // 
	}
	/* Terms from Jacobian determinant : */
	log_res += Type(0.5 * (n + 1)) * log_det_X;
	return (give_log) ? log_res : exp(log_res);
}

The we compared this version with the old version.

> dlkj(0,log(2),F)
[1] 1.333333
> dlkj2(0,log(2),F)
[1] 0.75

We test the integration in 2-d case

> integrate(Vectorize(function(x){dlkj2(x,log(2),F)}),-100,100)
1 with absolute error < 6.7e-06

and use (quasi) Monte Carlo integration to integrate over possible x values. In the 3-d and 4-d case

> #### test the integration using Monte Carlo sampling
> N <- 3e7
> p_dim <- 4
> d <- p_dim*(p_dim-1)/2
> width <- 10 ### can't make it too wide as the MC variance will be large for d>1
> us <- randtoolbox::sobol(N,dim=d) ## sobol ### space filling
> #us <- matrix(runif(N*d),N,d)
> xs <-us * width - width/2
> 
> mean(apply(xs[,1:3],1,dlkj2,shape=log(2),give_log=F))*width^3 ## check 3d case
[1] 0.9988517
> 
> mean(apply(xs,1,dlkj2,shape=log(2),give_log=F))*width^d ## check 4d case
[1] 0.9989281

As the number of dimensions increases the variance of the Monte Carlo estimation becomes large, due to the large space of xs.

In lkj_distribution.cpp there are two functions: dlkj_unconstrained and dlkj, which both integrate to one and follow the TMB style. The code follows jaganmn's parameterisation of the distribution with log(eta) being an input.

We show that in the 2d case the dlkj function in "lkj_distribution.cpp" integrates to 1:

> integrate(Vectorize(function(x){
+   X <- matrix(c(1,x,x,1),2,2)
+   dlkj_test(X,log(2),F)
+ }),-1,1)
1 with absolute error < 1.1e-14

[jaganmn]

@michaelspence Thanks for picking up the thread, which I have clearly repeatedly dropped ... I did spend some time last year thinking about the Jacobian determinant, using Theorems 1.28 and 1.29 in this reference. Maybe those will help you here? In the mean time, I will try to find my notes. I ought to have some bandwidth in December, but I'll be happy if people decide to jump ahead without me ...