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