MultivariateNormalGenerator.h

/**************************************************************************
 * basf2 (Belle II Analysis Software Framework)                           *
 * Author: The Belle II Collaboration                                     *
 *                                                                        *
 * See git log for contributors and copyright holders.                    *
 * This file is licensed under LGPL-3.0, see LICENSE.md.                  *
 **************************************************************************/
 
#pragma once
 
#include <Eigen/Dense>
#include <Math/Vector3D.h>
#include <TRandom.h>
#include <TVectorT.h>
#include <TMatrixTBase.h>
 
#include <array>
 
namespace Belle2 {
  class MultivariateNormalGenerator {
  public:
    MultivariateNormalGenerator() {}
    MultivariateNormalGenerator(int n, const double* mean, const double* cov)
    {
      setMeanCov(n, mean, cov);
    }
    MultivariateNormalGenerator(const Eigen::VectorXd& mean, const Eigen::MatrixXd& cov)
    {
      setMeanCov(mean, cov);
    }
    Eigen::VectorXd generate() const
    {
      //To get the correlated multivariate normal distribution, we multiply
      //standard normal distributed values (mean=0, sigma=1) with a
      //transformation matrix we obtained from an LDLT decomposition and add
      //the mean values.
      Eigen::VectorXd x(m_mean.rows());
      for (int i = 0; i < m_mean.rows(); ++i) x(i) = gRandom->Gaus();
      return m_mean + (m_transform * x);
    }
    void reset();
    size_t size() const { return m_mean.rows(); }
    void generate(double* output) const
    {
      Eigen::VectorXd x = generate();
      for (int i = 0; i < x.rows(); ++i) { output[i] = x(i); }
    }
 
    ROOT::Math::XYZVector generateVec3() const
    {
      Eigen::VectorXd x = generate();
      std::array<double, 3> tmp = {0.};
      for (unsigned int i = 0; i < std::min(3u, (unsigned int)size()); ++i) {
        tmp[i] = x(i);
      }
      return ROOT::Math::XYZVector(tmp[0], tmp[1], tmp[2]);
    }
 
    TVectorD generateVecT() const
    {
      Eigen::VectorXd x = generate();
      TVectorD output(x.rows());
      output.SetElements(x.data());
      return output;
    }
 
    bool setMeanCov(int n, const double* mean, const double* cov);
 
    bool setMeanCov(const Eigen::VectorXd& mean, const Eigen::MatrixXd& cov);
 
    template<class value_type> bool setMeanCov(const TVectorT<value_type>& mean,
                                               const TMatrixTBase<value_type>& cov);
 
    template<class value_type> bool setMeanCov(const ROOT::Math::XYZVector& mean,
                                               const TMatrixTBase<value_type>& cov);
 
  private:
    Eigen::VectorXd m_mean;
    Eigen::MatrixXd m_transform;
  };
 
  template<class value_type> bool MultivariateNormalGenerator::setMeanCov(
    const TVectorT<value_type>& mean, const TMatrixTBase<value_type>& cov)
  {
    Eigen::VectorXd emean(mean.GetNrows());
    Eigen::MatrixXd ecov(cov.GetNrows(), cov.GetNcols());
    for (int i = 0; i < mean.GetNrows(); ++i) { emean(i) = mean(i); }
    for (int i = 0; i < cov.GetNrows(); ++i) {
      for (int j = 0; j < cov.GetNcols(); ++j) {
        ecov(i, j) = cov(i, j);
      }
    }
    return setMeanCov(emean, ecov);
  }
 
  template<class value_type> bool MultivariateNormalGenerator::setMeanCov(
    const ROOT::Math::XYZVector& mean, const TMatrixTBase<value_type>& cov)
  {
    TVectorT<value_type> tmean(3);
    tmean[0] = mean.X();
    tmean[1] = mean.Y();
    tmean[2] = mean.Z();
    return setMeanCov(tmean, cov);
  }
}