Feature/2569 Analysis API: compute effective sample size

src/stan/analyze/mcmc/compute_effective_sample_size.hpp

@@ -0,0 +1,96 @@
+#ifndef STAN_ANALYZE_MCMC_COMPUTE_EFFECTIVE_SAMPLE_SIZE_HPP
+#define STAN_ANALYZE_MCMC_COMPUTE_EFFECTIVE_SAMPLE_SIZE_HPP
+
+#include <stan/math/prim/mat.hpp>
+#include <algorithm>
+#include <vector>
+
+namespace stan {
+  namespace analyze {
+
+    /**
+     * Returns the effective sample size for the specified parameter
+     * across all kept samples.
+     *
+     * See more details in Stan reference manual section "Effective
+     * Sample Size". http://mc-stan.org/users/documentation
+     *
+     * Current implementation assumes chains are all of equal size and
+     * draws are stored in contiguous blocks of memory.
+     *
+     * @param std::vector stores pointers to arrays of chains
+     * @param std::vector stores sizes of chains
+     * @return effective sample size for the specified parameter
+     */
+    double compute_effective_sample_size(std::vector<const double*> draws,
+                                         std::vector<size_t> sizes) {
+      int num_chains = sizes.size();
+
+      // need to generalize to each jagged draws per chain
+      size_t num_draws = sizes[0];
+      for (int chain = 1; chain < num_chains; ++chain) {
+        num_draws = std::min(num_draws, sizes[chain]);
+      }
+
+      Eigen::Matrix<Eigen::VectorXd, Eigen::Dynamic, 1> acov(num_chains);
+      Eigen::VectorXd chain_mean(num_chains);
+      Eigen::VectorXd chain_var(num_chains);
+      for (int chain = 0; chain < num_chains; ++chain) {
+        Eigen::Map<const Eigen::Matrix<double, Eigen::Dynamic, 1>>
+          draw(draws[chain], sizes[chain]);
+        math::autocovariance<double>(draw, acov(chain));
+        chain_mean(chain) = draw.mean();
+        chain_var(chain) = acov(chain)(0) * num_draws / (num_draws - 1);
+      }
+
+      double mean_var = chain_var.mean();
+      double var_plus = mean_var * (num_draws - 1) / num_draws;
+      if (num_chains > 1)
+        var_plus += math::variance(chain_mean);
+      Eigen::VectorXd rho_hat_s(num_draws);
+      rho_hat_s.setZero();
+      Eigen::VectorXd acov_s(num_chains);
+      for (int chain = 0; chain < num_chains; ++chain)
+        acov_s(chain) = acov(chain)(1);
+      double rho_hat_even = 1;
+      double rho_hat_odd = 1 - (mean_var - acov_s.mean()) / var_plus;
+      rho_hat_s(1) = rho_hat_odd;
+      // Geyer's initial positive sequence
+      int max_s = 1;
+      for (size_t s = 1;
+           (s < (num_draws - 2) && (rho_hat_even + rho_hat_odd) >= 0);
+           s += 2) {
+        for (int chain = 0; chain < num_chains; ++chain)
+          acov_s(chain) = acov(chain)(s + 1);
+        rho_hat_even = 1 - (mean_var - acov_s.mean()) / var_plus;
+        for (int chain = 0; chain < num_chains; ++chain)
+          acov_s(chain) = acov(chain)(s + 2);
+        rho_hat_odd = 1 - (mean_var - acov_s.mean()) / var_plus;
+        if ((rho_hat_even + rho_hat_odd) >= 0) {
+          rho_hat_s(s + 1) = rho_hat_even;
+          rho_hat_s(s + 2) = rho_hat_odd;
+        }
+        max_s = s + 2;
+      }
+      // Geyer's initial monotone sequence
+      for (int s = 3; s <= max_s - 2; s += 2) {
+        if (rho_hat_s(s + 1) + rho_hat_s(s + 2) >
+            rho_hat_s(s - 1) + rho_hat_s(s)) {
+          rho_hat_s(s + 1) = (rho_hat_s(s - 1) + rho_hat_s(s)) / 2;
+          rho_hat_s(s + 2) = rho_hat_s(s + 1);
+        }
+      }
+
+      return num_chains * num_draws / (1 + 2 * rho_hat_s.sum());
+    }
+
+    double compute_effective_sample_size(std::vector<const double*> draws,
+                                         size_t size) {
+      int num_chains = draws.size();
+      std::vector<size_t> sizes(num_chains, size);
+      return compute_effective_sample_size(draws, sizes);
+    }
+  }
+}
+
+#endif

src/stan/mcmc/chains.hpp

@@ -3,6 +3,7 @@
 
 #include <stan/io/stan_csv_reader.hpp>
 #include <stan/math/prim/mat.hpp>
+#include <stan/analyze/mcmc/compute_effective_sample_size.hpp>
 #include <boost/accumulators/accumulators.hpp>
 #include <boost/accumulators/statistics/stats.hpp>
 #include <boost/accumulators/statistics/mean.hpp>
@@ -209,89 +210,6 @@ namespace stan {
         return ac2;
       }
 
-      /**
-       * Returns the effective sample size for the specified parameter
-       * across all kept samples.
-       *
-       * The implementation is close to the effective sample size
-       * description in BDA3 (p. 286-287).  See more details in Stan
-       * reference manual section "Effective Sample Size".
-       *
-       * Current implementation takes the minimum number of samples
-       * across chains as the number of samples per chain.
-       *
-       * @param VectorXd
-       * @param Dynamic
-       * @param samples
-       *
-       * @return
-       */
-      double effective_sample_size(const Eigen::Matrix<Eigen::VectorXd,
-                                   Dynamic, 1> &samples) const {
-        int chains = samples.size();
-
-        // need to generalize to each jagged samples per chain
-        int n_samples = samples(0).size();
-        for (int chain = 1; chain < chains; chain++) {
-          n_samples = std::min(n_samples,
-                               static_cast<int>(samples(chain).size()));
-        }
-
-        Eigen::Matrix<Eigen::VectorXd, Dynamic, 1> acov(chains);
-        for (int chain = 0; chain < chains; chain++) {
-          acov(chain) = autocovariance(samples(chain));
-        }
-
-        Eigen::VectorXd chain_mean(chains);
-        Eigen::VectorXd chain_var(chains);
-        for (int chain = 0; chain < chains; chain++) {
-          double n_kept_samples = num_kept_samples(chain);
-          chain_mean(chain) = mean(samples(chain));
-          chain_var(chain) = acov(chain)(0)*n_kept_samples/(n_kept_samples-1);
-        }
-
-        double mean_var = mean(chain_var);
-        double var_plus = mean_var*(n_samples-1)/n_samples;
-        if (chains > 1)
-          var_plus += variance(chain_mean);
-        Eigen::VectorXd rho_hat_t(n_samples);
-        rho_hat_t.setZero();
-        Eigen::VectorXd acov_t(chains);
-        for (int chain = 0; chain < chains; chain++)
-          acov_t(chain) = acov(chain)(1);
-        double rho_hat_even = 1;
-        double rho_hat_odd = 1 - (mean_var - mean(acov_t)) / var_plus;
-        rho_hat_t(1) = rho_hat_odd;
-        // Geyer's initial positive sequence
-        int max_t = 1;
-        for (int t = 1;
-             (t < (n_samples - 2) && (rho_hat_even + rho_hat_odd) >= 0);
-             t += 2) {
-          for (int chain = 0; chain < chains; chain++)
-            acov_t(chain) = acov(chain)(t + 1);
-          rho_hat_even = 1 - (mean_var - mean(acov_t)) / var_plus;
-          for (int chain = 0; chain < chains; chain++)
-            acov_t(chain) = acov(chain)(t + 2);
-          rho_hat_odd = 1 - (mean_var - mean(acov_t)) / var_plus;
-          if ((rho_hat_even + rho_hat_odd) >= 0) {
-            rho_hat_t(t + 1) = rho_hat_even;
-            rho_hat_t(t + 2) = rho_hat_odd;
-          }
-          max_t = t + 2;
-        }
-        // Geyer's initial monotone sequence
-        for (int t = 3; t <= max_t - 2; t += 2) {
-          if (rho_hat_t(t + 1) + rho_hat_t(t + 2) >
-              rho_hat_t(t - 1) + rho_hat_t(t)) {
-            rho_hat_t(t + 1) = (rho_hat_t(t - 1) + rho_hat_t(t)) / 2;
-            rho_hat_t(t + 2) = rho_hat_t(t + 1);
-          }
-        }
-        double ess = chains * n_samples;
-        ess /= (1 + 2 * rho_hat_t.sum());
-        return ess;
-      }
-
       /**
        * Return the split potential scale reduction (split R hat)
        * for the specified parameter.
@@ -689,12 +607,17 @@ namespace stan {
 
       // FIXME: reimplement using autocorrelation.
       double effective_sample_size(const int index) const {
-        Eigen::Matrix<Eigen::VectorXd, Dynamic, 1>
-          samples(num_chains());
-        for (int chain = 0; chain < num_chains(); chain++) {
-          samples(chain) = this->samples(chain, index);
+        int n_chains = num_chains();
+        std::vector<const double*> draws(n_chains);
+        std::vector<size_t> sizes(n_chains);
+        int n_kept_samples = 0;
+        for (int chain = 0; chain < n_chains; ++chain) {
+          n_kept_samples = num_kept_samples(chain);
+          draws[chain]
+            = samples_(chain).col(index).bottomRows(n_kept_samples).data();
+          sizes[chain] = n_kept_samples;
         }
-        return effective_sample_size(samples);
+        return analyze::compute_effective_sample_size(draws, sizes);
       }
 
       double effective_sample_size(const std::string& name) const {

src/test/unit/analyze/mcmc/compute_effective_sample_size_test.cpp

@@ -0,0 +1,141 @@
+#include <stan/mcmc/chains.hpp>
+#include <stan/io/stan_csv_reader.hpp>
+#include <gtest/gtest.h>
+#include <fstream>
+#include <sstream>
+
+class ComputeEss : public testing::Test {
+public:
+    void SetUp() {
+    blocker1_stream.open("src/test/unit/mcmc/test_csv_files/blocker.1.csv");
+    blocker2_stream.open("src/test/unit/mcmc/test_csv_files/blocker.2.csv");
+  }
+
+  void TearDown() {
+    blocker1_stream.close();
+    blocker2_stream.close();
+  }
+  std::ifstream blocker1_stream, blocker2_stream;
+};
+
+TEST_F(ComputeEss,compute_effective_sample_size) {
+  std::stringstream out;
+  stan::io::stan_csv blocker1 = stan::io::stan_csv_reader::parse(blocker1_stream, &out);
+  stan::io::stan_csv blocker2 = stan::io::stan_csv_reader::parse(blocker2_stream, &out);
+  EXPECT_EQ("", out.str());
+
+  stan::mcmc::chains<> chains(blocker1);
+  chains.add(blocker2);
+
+  Eigen::VectorXd n_eff(48);
+  n_eff << 466.099,136.953,1170.390,541.256,
+    518.051,589.244,764.813,688.294,
+    323.777,502.892,353.823,588.142,
+    654.336,480.914,176.978,182.649,
+    642.389,470.949,561.947,581.187,
+    446.389,397.641,338.511,678.772,
+    1442.250,837.956,869.865,951.124,
+    619.336,875.805,233.260,786.568,
+    910.144,231.582,907.666,747.347,
+    720.660,195.195,944.547,767.271,
+    723.665,1077.030,470.903,954.924,
+    497.338,583.539,697.204,98.421;
+
+  // replicates calls to stan::anlayze::compute_effective_sample_size
+  // for any interface *without* access to chains class
+  Eigen::Matrix<Eigen::VectorXd, Eigen::Dynamic, 1>
+    samples(chains.num_chains());
+  std::vector<const double*> draws(chains.num_chains());
+  std::vector<size_t> sizes(chains.num_chains());
+  for (int index = 4; index < chains.num_params(); index++) {
+    for (int chain = 0; chain < chains.num_chains(); ++chain) {
+      samples(chain) = chains.samples(chain, index);
+      draws[chain] = &samples(chain)(0);
+      sizes[chain] = samples(chain).size();
+    }
+    ASSERT_NEAR(n_eff(index - 4),
+                stan::analyze::compute_effective_sample_size(draws, sizes), 1.0)
+      << "n_effective for index: " << index << ", parameter: "
+      << chains.param_name(index);
+  }
+}
+
+TEST_F(ComputeEss,compute_effective_sample_size_convenience) {
+  std::stringstream out;
+  stan::io::stan_csv blocker1 = stan::io::stan_csv_reader::parse(blocker1_stream, &out);
+  stan::io::stan_csv blocker2 = stan::io::stan_csv_reader::parse(blocker2_stream, &out);
+  EXPECT_EQ("", out.str());
+
+  stan::mcmc::chains<> chains(blocker1);
+  chains.add(blocker2);
+
+  Eigen::VectorXd n_eff(48);
+  n_eff << 466.099,136.953,1170.390,541.256,
+    518.051,589.244,764.813,688.294,
+    323.777,502.892,353.823,588.142,
+    654.336,480.914,176.978,182.649,
+    642.389,470.949,561.947,581.187,
+    446.389,397.641,338.511,678.772,
+    1442.250,837.956,869.865,951.124,
+    619.336,875.805,233.260,786.568,
+    910.144,231.582,907.666,747.347,
+    720.660,195.195,944.547,767.271,
+    723.665,1077.030,470.903,954.924,
+    497.338,583.539,697.204,98.421;
+
+  Eigen::Matrix<Eigen::VectorXd, Eigen::Dynamic, 1>
+    samples(chains.num_chains());
+  std::vector<const double*> draws(chains.num_chains());
+  std::vector<size_t> sizes(chains.num_chains());
+  for (int index = 4; index < chains.num_params(); index++) {
+    for (int chain = 0; chain < chains.num_chains(); ++chain) {
+      samples(chain) = chains.samples(chain, index);
+      draws[chain] = &samples(chain)(0);
+
+    }
+    size_t size  = samples(0).size();
+    ASSERT_NEAR(n_eff(index - 4),
+                stan::analyze::compute_effective_sample_size(draws, size), 1.0)
+      << "n_effective for index: " << index << ", parameter: "
+      << chains.param_name(index);
+  }
+}
+
+TEST_F(ComputeEss,chains_compute_effective_sample_size) {
+  std::stringstream out;
+  stan::io::stan_csv blocker1 = stan::io::stan_csv_reader::parse(blocker1_stream, &out);
+  stan::io::stan_csv blocker2 = stan::io::stan_csv_reader::parse(blocker2_stream, &out);
+  EXPECT_EQ("", out.str());
+
+  stan::mcmc::chains<> chains(blocker1);
+  chains.add(blocker2);
+
+  Eigen::VectorXd n_eff(48);
+  n_eff << 466.099,136.953,1170.390,541.256,
+    518.051,589.244,764.813,688.294,
+    323.777,502.892,353.823,588.142,
+    654.336,480.914,176.978,182.649,
+    642.389,470.949,561.947,581.187,
+    446.389,397.641,338.511,678.772,
+    1442.250,837.956,869.865,951.124,
+    619.336,875.805,233.260,786.568,
+    910.144,231.582,907.666,747.347,
+    720.660,195.195,944.547,767.271,
+    723.665,1077.030,470.903,954.924,
+    497.338,583.539,697.204,98.421;
+
+  // replicates calls to stan::anlayze::compute_effective_sample_size
+  // for any interface with access to chains class
+  for (int index = 4; index < chains.num_params(); index++) {
+    ASSERT_NEAR(n_eff(index - 4),
+                chains.effective_sample_size(index), 1.0)
+      << "n_effective for index: " << index << ", parameter: "
+      << chains.param_name(index);
+  }
+
+  for (int index = 0; index < chains.num_params(); index++) {
+    std::string name = chains.param_name(index);
+    ASSERT_EQ(chains.effective_sample_size(index),
+              chains.effective_sample_size(name));
+  }
+}