1 #ifndef STAN_MATH_PRIM_FUN_WELFORD_VAR_ESTIMATOR_HPP 2 #define STAN_MATH_PRIM_FUN_WELFORD_VAR_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_var_estimator { 11 public: welford_var_estimator(int n)12 explicit welford_var_estimator(int n) 13 : m_(Eigen::VectorXd::Zero(n)), m2_(Eigen::VectorXd::Zero(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_ += delta.cwiseProduct(q - m_); 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_variance(Eigen::VectorXd & var)35 void sample_variance(Eigen::VectorXd& var) { 36 if (num_samples_ > 1) { 37 var = m2_ / (num_samples_ - 1.0); 38 } 39 } 40 41 protected: 42 double num_samples_; 43 Eigen::VectorXd m_; 44 Eigen::VectorXd m2_; 45 }; 46 47 } // namespace math 48 } // namespace stan 49 #endif 50