/* randist/mvgauss.c

 * 

 * Copyright (C) 2016 Timothée Flutre, Patrick Alken

 * 

 * This program is free software; you can redistribute it and/or modify

 * it under the terms of the GNU General Public License as published by

 * the Free Software Foundation; either version 3 of the License, or (at

 * your option) any later version.

 * 

 * This program is distributed in the hope that it will be useful, but

 * WITHOUT ANY WARRANTY; without even the implied warranty of

 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU

 * General Public License for more details.

 * 

 * You should have received a copy of the GNU General Public License

 * along with this program; if not, write to the Free Software

 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.

 */

#include <config.h>

#include <math.h>

#include <gsl/gsl_math.h>

#include <gsl/gsl_rng.h>

#include <gsl/gsl_randist.h>

#include <gsl/gsl_vector.h>

#include <gsl/gsl_matrix.h>

#include <gsl/gsl_blas.h>

#include <gsl/gsl_linalg.h>

#include <gsl/gsl_statistics.h>

static int multivar_vcov (const double data[], size_t d, size_t tda, size_t n,

                          double vcov[], size_t tda2);

/* Generate a random vector from a multivariate Gaussian distribution using

 * the Cholesky decomposition of the variance-covariance matrix, following

 * "Computational Statistics" from Gentle (2009), section 7.4.

 *

 * mu      mean vector (dimension d)

 * L       matrix resulting from the Cholesky decomposition of

 *         variance-covariance matrix Sigma = L L^T (dimension d x d)

 * result  output vector (dimension d)

 */

int

gsl_ran_multivariate_gaussian (const gsl_rng * r,

                               const gsl_vector * mu,

                               const gsl_matrix * L,

                               gsl_vector * result)

{

  const size_t M = L->size1;

  const size_t N = L->size2;

  if (M != N)

    {

      GSL_ERROR("requires square matrix", GSL_ENOTSQR);

    }

  else if (mu->size != M)

    {

      GSL_ERROR("incompatible dimension of mean vector with variance-covariance matrix", GSL_EBADLEN);

    }

  else if (result->size != M)

    {

      GSL_ERROR("incompatible dimension of result vector", GSL_EBADLEN);

    }

  else

    {

      size_t i;

      for (i = 0; i < M; ++i)

        gsl_vector_set(result, i, gsl_ran_ugaussian(r));

      gsl_blas_dtrmv(CblasLower, CblasNoTrans, CblasNonUnit, L, result);

      gsl_vector_add(result, mu);

      return GSL_SUCCESS;

    }

}

/* Compute the log of the probability density function at a given quantile

 * vector for a multivariate Gaussian distribution using the Cholesky

 * decomposition of the variance-covariance matrix.

 *

 * x       vector of quantiles (dimension d)

 * mu      mean vector (dimension d)

 * L       matrix resulting from the Cholesky decomposition of

 *         variance-covariance matrix Sigma = L L^T (dimension d x d)

 * result  output of the density (dimension 1)

 * work    vector used for intermediate computations (dimension d)

 */

int

gsl_ran_multivariate_gaussian_log_pdf (const gsl_vector * x,

                                       const gsl_vector * mu,

                                       const gsl_matrix * L,

                                       double * result,

                                       gsl_vector * work)

{

  const size_t M = L->size1;

  const size_t N = L->size2;

  if (M != N)

    {

      GSL_ERROR("requires square matrix", GSL_ENOTSQR);

    }

  else if (mu->size != M)

    {

      GSL_ERROR("incompatible dimension of mean vector with variance-covariance matrix", GSL_EBADLEN);

    }

  else if (x->size != M)

    {

      GSL_ERROR("incompatible dimension of quantile vector", GSL_EBADLEN);

    }

  else if (work->size != M)

    {

      GSL_ERROR("incompatible dimension of work vector", GSL_EBADLEN);

    }

  else

    {

      size_t i;

      double quadForm;        /* (x - mu)' Sigma^{-1} (x - mu) */

      double logSqrtDetSigma; /* log [ sqrt(|Sigma|) ] */

      /* compute: work = x - mu */

      for (i = 0; i < M; ++i)

        {

          double xi = gsl_vector_get(x, i);

          double mui = gsl_vector_get(mu, i);

          gsl_vector_set(work, i, xi - mui);

        }

      /* compute: work = L^{-1} * (x - mu) */

      gsl_blas_dtrsv(CblasLower, CblasNoTrans, CblasNonUnit, L, work);

      /* compute: quadForm = (x - mu)' Sigma^{-1} (x - mu) */

      gsl_blas_ddot(work, work, &quadForm);

      /* compute: log [ sqrt(|Sigma|) ] = sum_i log L_{ii} */

      logSqrtDetSigma = 0.0;

      for (i = 0; i < M; ++i)

        {

          double Lii = gsl_matrix_get(L, i, i);

          logSqrtDetSigma += log(Lii);

        }

      *result = -0.5*quadForm - logSqrtDetSigma - 0.5*M*log(2.0*M_PI);

      return GSL_SUCCESS;

    }

}

int

gsl_ran_multivariate_gaussian_pdf (const gsl_vector * x,

                                   const gsl_vector * mu,

                                   const gsl_matrix * L,

                                   double * result,

                                   gsl_vector * work)

{

  double logpdf;

  int status = gsl_ran_multivariate_gaussian_log_pdf(x, mu, L, &logpdf, work);

  if (status == GSL_SUCCESS)

    *result = exp(logpdf);

  return status;

}

/* Compute the maximum-likelihood estimate of the mean vector of samples

 * from a multivariate Gaussian distribution.

 *

 * Example from R (GPL): http://www.r-project.org/

 * (samples <- matrix(c(4.348817, 2.995049, -3.793431, 4.711934, 1.190864, -1.357363), nrow=3, ncol=2))

 * colMeans(samples) # 1.183478 1.515145

 */

int

gsl_ran_multivariate_gaussian_mean (const gsl_matrix * X, gsl_vector * mu_hat)

{

  const size_t M = X->size1;

  const size_t N = X->size2;

  if (N != mu_hat->size)

    {

      GSL_ERROR("mu_hat vector has wrong size", GSL_EBADLEN);

    }

  else

    {

      size_t j;

      for (j = 0; j < N; ++j)

        {

          gsl_vector_const_view c = gsl_matrix_const_column(X, j);

          double mean = gsl_stats_mean(c.vector.data, c.vector.stride, M);

          gsl_vector_set(mu_hat, j, mean);

        }

      return GSL_SUCCESS;

    }

}

/* Compute the maximum-likelihood estimate of the variance-covariance matrix

 * of samples from a multivariate Gaussian distribution.

 */

int

gsl_ran_multivariate_gaussian_vcov (const gsl_matrix * X, gsl_matrix * sigma_hat)

{

  const size_t M = X->size1;

  const size_t N = X->size2;

  if (sigma_hat->size1 != sigma_hat->size2)

    {

      GSL_ERROR("sigma_hat must be a square matrix", GSL_ENOTSQR);

    }

  else if (N != sigma_hat->size1)

    {

      GSL_ERROR("sigma_hat does not match X matrix dimensions", GSL_EBADLEN);

    }

  else

    {

      return multivar_vcov (X->data, N, X->tda, M, sigma_hat->data, sigma_hat->tda);

    }

}

/* Example from R (GPL): http://www.r-project.org/

 * (samples <- matrix(c(4.348817, 2.995049, -3.793431, 4.711934, 1.190864, -1.357363), nrow=3, ncol=2))

 * cov(samples) # 19.03539 11.91384 \n 11.91384  9.28796

 */

static int

multivar_vcov (const double data[], size_t d, size_t tda, size_t n,

               double vcov[], size_t tda2)

{

  size_t j1 = 0, j2 = 0;

  for (j1 = 0; j1 < d; ++j1)

    {

      vcov[j1 * tda2 + j1] = gsl_stats_variance(&(data[j1]), tda, n);

      for (j2 = j1 + 1; j2 < d; ++j2)

        {

          vcov[j1 * tda2 + j2] = gsl_stats_covariance(&(data[j1]), tda,

                                                      &(data[j2]), tda, n);

          vcov[j2 * tda2 + j1] = vcov[j1 * tda2 + j2];

        }

      }

  return GSL_SUCCESS;

}