1 #ifndef STAN_MATH_PRIM_FUN_WELFORD_COVAR_ESTIMATOR_HPP 2 #define STAN_MATH_PRIM_FUN_WELFORD_COVAR_ESTIMATOR_HPP 3 4 #include <stan/math/prim/fun/Eigen.hpp> 5 #include <vector> 6 7 namespace stan { 8 namespace math { 9 10 class welford_covar_estimator { 11 public: welford_covar_estimator(int n)12 explicit welford_covar_estimator(int n) 13 : m_(Eigen::VectorXd::Zero(n)), m2_(Eigen::MatrixXd::Zero(n, n)) { 14 restart(); 15 } 16 restart()17 void restart() { 18 num_samples_ = 0; 19 m_.setZero(); 20 m2_.setZero(); 21 } 22 add_sample(const Eigen::VectorXd & q)23 void add_sample(const Eigen::VectorXd& q) { 24 ++num_samples_; 25 26 Eigen::VectorXd delta(q - m_); 27 m_ += delta / num_samples_; 28 m2_ += (q - m_) * delta.transpose(); 29 } 30 num_samples()31 int num_samples() { return num_samples_; } 32 sample_mean(Eigen::VectorXd & mean)33 void sample_mean(Eigen::VectorXd& mean) { mean = m_; } 34 sample_covariance(Eigen::MatrixXd & covar)35 void sample_covariance(Eigen::MatrixXd& covar) { 36 if (num_samples_ > 1) { 37 covar = m2_ / (num_samples_ - 1.0); 38 } 39 } 40 41 protected: 42 double num_samples_; 43 Eigen::VectorXd m_; 44 Eigen::MatrixXd m2_; 45 }; 46 47 } // namespace math 48 } // namespace stan 49 #endif 50