CalcMeanCov.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 <cmath>
 
namespace Belle2 {
  template<int N = 1, class RealType = double>
  class CalcMeanCov {
    static_assert(N > 0, "Number of parameters, N, must be positive");
  public:
    CalcMeanCov() { clear(); }
 
    typedef RealType value_type;
 
    void clear()
    {
      m_entries = 0;
      for (int i = 0; i < N; i++)
        m_mean[i] = 0;
      for (int i = 0; i < N * (N + 1) / 2; i++)
        m_covariance[i] = 0;
    }
 
    template<class... T> void addWeighted(value_type weight, T... values)
    {
      static_assert(sizeof...(values) == N,
                    "Number of arguments must be equal N");
      m_entries += weight;
      addValue<0>(weight, values...);
    }
 
    template<class... T> void add(T... values)
    {
      addWeighted(1.0, values...);
    }
 
    void add(const CalcMeanCov<N, RealType>& other)
    {
      const value_type n = m_entries + other.m_entries;
      if (n == 0)
        return;
      for (int i = 0; i < N; ++i) {
        const value_type delta = (other.m_mean[i] - m_mean[i]);
        //Update covariance matrix
        for (int j = i; j < N; ++j) {
          const value_type delta2 = (other.m_mean[j] - m_mean[j]);
          m_covariance[getIndex(i, j)] += other.m_covariance[getIndex(i, j)]  + m_entries * other.m_entries
                                          / n * delta * delta2;
        }
        //Update mean value
        m_mean[i] += other.m_entries * delta / n;
      }
      m_entries = n;
    }
 
    void addWeightedArray(value_type weight, value_type* values)
    {
      m_entries += weight;
      addArrayValues(weight, values);
    }
 
    void addArray(value_type* values)
    {
      addWeightedArray(1.0, values);
    }
 
    value_type getEntries() const { return m_entries; }
    value_type getMean(int i) const { return m_mean[i]; }
    value_type getCovariance(int i, int j) const
    {
      if (m_entries == 0.0) return 0.0;
      return m_covariance[getIndex(i, j)] / m_entries;
    }
    value_type getCorrelation(int i, int j) const
    {
      if (i == j) return 1.0;
      if (m_entries == 0.0) return 0.0;
      return getCovariance(i, j) / (getStddev(i) * getStddev(j));
    }
    value_type getVariance(int i) const { return getCovariance(i, i); }
    value_type getStddev(int i) const { return std::sqrt(getVariance(i)); }
    value_type getSum(int i) const { return getMean(i) * m_entries; }
 
    template <int i = 0> value_type getMean() const
    {
      static_assert(i >= 0 && i < N, "index i out of range");
      return m_mean[i];
    }
    template <int i, int j> value_type getCovariance() const
    {
      static_assert(i >= 0 && i < N, "index i out of range");
      static_assert(j >= 0 && j < N, "index j out of range");
      if (m_entries == 0.0) return 0.0;
      return m_covariance[getIndex(i, j)] / m_entries;
    }
    template <int i, int j> value_type getCorrelation() const
    {
      if (i == j) return 1.0;
      if (m_entries == 0.0) return 0.0;
      return getCovariance<i, j>() / (getStddev<i>() * getStddev<j>());
    }
    template <int i = 0> value_type getVariance() const
    {
      return getCovariance<i, i>();
    }
    template <int i = 0> value_type getStddev() const
    {
      return std::sqrt(getVariance<i>());
    }
    template <int i = 0> value_type getSum() const
    {
      return getMean<i>() * m_entries;
    }
 
  private:
    template<int i, class... T>
    void addValue(value_type weight, value_type x, T... values)
    {
      const value_type delta = (x - m_mean[i]);
      //Update covariance elements
      updateCov<i, i>(weight, delta, x, values...);
      //Update mean value
      m_mean[i] += weight * delta / m_entries;
      //recursively call again for remaining values
      addValue < i + 1 > (weight, values...);
    }
    template<int i> void addValue(value_type) {}
 
    template<int i, int j, class... T>
    void updateCov(value_type weight, value_type delta, value_type x,
                   T... values)
    {
      const value_type delta2 = (x - m_mean[j]);
      m_covariance[getIndex(i, j)] += (m_entries - weight) * weight
                                      / m_entries * delta * delta2;
      updateCov < i, j + 1 > (weight, delta, values...);
    }
 
    template<int i, int j> void updateCov(value_type, value_type) {}
 
    void addArrayValues(value_type weight, value_type* x)
    {
      for (int i = 0; i < N; ++i) {
        const value_type delta = (x[i] - m_mean[i]);
        //Update covariance matrix
        for (int j = i; j < N; ++j) {
          const value_type delta2 = (x[j] - m_mean[j]);
          m_covariance[getIndex(i, j)] += (m_entries - weight) * weight
                                          / m_entries * delta * delta2;
        }
        //Update mean value
        m_mean[i] += weight * delta / m_entries;
      }
    }
 
    constexpr int getIndex(unsigned int i, unsigned int j) const
    {
      //swap indices if i >= j
      return (i < j) ? ((j + 1) * j / 2 + i) : ((i + 1) * i / 2 + j);
    }
 
    value_type m_entries;
    value_type m_mean[N];
    value_type m_covariance[N * (N + 1) / 2];
  };
 
} //Belle2 namespace