Fix ESS and R-hat when some chains are constant, and align ESS and MCSE with posterior

src/stan/analyze/mcmc/autocovariance.hpp

@@ -77,13 +77,21 @@ void autocorrelation(const Eigen::MatrixBase<DerivedA>& y,
  * <p>This method is just a light wrapper around the three-argument
  * autocovariance function
  *
+ * <p>If all values of the input sequence are equal, the autocovariances
+ * are zero at every lag.
+ *
  * @tparam T Scalar type.
  * @param y Input sequence.
  * @param acov Autocovariances.
  */
 template <typename T, typename DerivedA, typename DerivedB>
 void autocovariance(const Eigen::MatrixBase<DerivedA>& y,
                     Eigen::MatrixBase<DerivedB>& acov) {
+  // the autocorrelations of a constant sequence are 0 / 0
+  if (y.size() > 0 && (y.array() == y(0)).all()) {
+    acov = Eigen::Matrix<T, Eigen::Dynamic, 1>::Zero(y.size());
+    return;
+  }
   Eigen::FFT<T> fft;
   autocorrelation(y, acov, fft);
 

src/stan/analyze/mcmc/check_chains.hpp

@@ -12,32 +12,23 @@ namespace stan {
 namespace analyze {
 
 /**
- * Checks that values across all matrix columns finite and non-identical.
+ * Checks that all values are finite and that at least one chain is not
+ * constant.  Some chains can be constant: a chain that does not move gives
+ * a large R-hat and a small ESS, which the diagnostics must report.
  *
  * @param chains matrix of draws, one column per chain
  * @return bool true if OK, false otherwise
  */
 inline bool is_finite_and_varies(const Eigen::MatrixXd& chains) {
-  size_t num_chains = chains.cols();
-  size_t num_samples = chains.rows();
-  if (num_samples < 2) {
+  if (chains.rows() < 2 || !chains.allFinite()) {
     return false;
   }
-  Eigen::VectorXd first_draws = Eigen::VectorXd::Zero(num_chains);
-  for (std::size_t i = 0; i < num_chains; ++i) {
-    first_draws(i) = chains.col(i)(0);
-    for (int j = 0; j < num_samples; ++j) {
-      if (!std::isfinite(chains.col(i)(j)))
-        return false;
+  for (Eigen::Index i = 0; i < chains.cols(); ++i) {
+    if (!chains.col(i).isApproxToConstant(chains(0, i))) {
+      return true;
     }
-    if (chains.col(i).isApproxToConstant(first_draws(i))) {
-      return false;
-    }
-  }
-  if (num_chains > 1 && first_draws.isApproxToConstant(first_draws(0))) {
-    return false;
   }
-  return true;
+  return false;
 }
 
 }  // namespace analyze

src/stan/analyze/mcmc/ess.hpp

@@ -14,7 +14,9 @@ namespace analyze {
 /**
  * Computes the effective sample size (ESS) for the specified
  * parameter across all chains.  The number of draws per chain must be > 3,
- * and the values across all draws must be finite and not constant.
+ * and the values across all draws must be finite and not constant;
+ * otherwise the return value is NaN.  A constant chain is permitted: its
+ * autocovariances are zero.
  * See https://arxiv.org/abs/1903.08008, section 3.2 for discussion.
  *
  * Sample autocovariance is computed using the implementation in this namespace
@@ -54,6 +56,10 @@ double ess(const Eigen::MatrixXd& chains) {
   if (num_chains > 1) {
     var_plus += math::variance(chain_mean);  // B (between chain var)
   }
+  // all draws are equal, or a draw is not finite
+  if (!(var_plus > 0)) {
+    return std::numeric_limits<double>::quiet_NaN();
+  }
 
   // Geyer's initial positive sequence, eqn (11)
   Eigen::VectorXd rho_hat_t = Eigen::VectorXd::Zero(draws_per_chain);
@@ -92,15 +98,20 @@ double ess(const Eigen::MatrixXd& chains) {
     t += 2;
   }
 
-  auto max_t = t;  // max lag, used for truncation
+  // truncation lag: the even lag of the last pair computed, which is the
+  // first negative pair or the pair at the lag limit
+  Eigen::Index max_t = t - 1;
   //  see discussion p. 8, par "In extreme antithetic cases, "
-  if (rho_hat_even > 0) {
-    rho_hat_t(max_t + 1) = rho_hat_even;
-  }
+  double rho_hat_last = (rho_hat_even > 0 || rho_hat_even + rho_hat_odd >= 0)
+                            ? rho_hat_even
+                            : 0;
 
   double draws_total = num_chains * draws_per_chain;
-  //  eqn (13): Geyer's truncation rule, w/ modification
-  double tau_hat = -1 + 2 * rho_hat_t.head(max_t).sum() + rho_hat_t(max_t + 1);
+  //  eqn (13): Geyer's truncation rule, w/ modification;
+  //  tau_hat is 2 if no pair was computed (max_t = 0)
+  double tau_hat = -1
+                   + 2 * rho_hat_t.head(std::max<Eigen::Index>(max_t, 1)).sum()
+                   + rho_hat_last;
   // safety check for negative values and with max ess equal to ess*log10(ess)
   tau_hat = std::max(tau_hat, 1 / std::log10(draws_total));
   return (draws_total / tau_hat);

src/stan/analyze/mcmc/mcse.hpp

@@ -13,27 +13,28 @@ namespace stan {
 namespace analyze {
 
 /**
- * Computes the mean Monte Carlo error estimate for the central 90% interval.
+ * Computes the Monte Carlo standard error of the mean.
  * See https://arxiv.org/abs/1903.08008, section 4.4.
  * Follows implementation in the R posterior package.
  *
  * @param chains matrix of draws across all chains
  * @return mcse
  */
 inline double mcse_mean(const Eigen::MatrixXd& chains) {
-  if (chains.rows() < 4 || !is_finite_and_varies(chains))
+  if (chains.rows() < 8 || !is_finite_and_varies(chains))
     return std::numeric_limits<double>::quiet_NaN();
 
   double sample_var
       = (chains.array() - chains.mean()).square().sum() / (chains.size() - 1);
-  return std::sqrt(sample_var / ess(chains));
+  // posterior pkg fn `ess_mean` computes on split chains
+  return std::sqrt(sample_var / ess(split_chains(chains)));
 }
 
 /**
- * Computes the standard deviation of the Monte Carlo error estimate
- * https://arxiv.org/abs/1903.08008, section 4.4.
- * Follows implementation in the R posterior package:
- * https://github.com/stan-dev/posterior/blob/98bf52329d68f3307ac4ecaaea659276ee1de8df/R/convergence.R#L478-L496
+ * Computes the Monte Carlo standard error of the standard deviation.
+ * See https://arxiv.org/abs/1903.08008, section 4.4.
+ * Follows implementation in the R posterior package.  The effective
+ * sample size is that of the squared deviations from the mean.
  *
  * @param chains matrix of draws across all chains
  * @return mcse
@@ -42,18 +43,17 @@ inline double mcse_sd(const Eigen::MatrixXd& chains) {
   if (chains.rows() < 8 || !is_finite_and_varies(chains)) {
     return std::numeric_limits<double>::quiet_NaN();
   }
-  // center the data, take abs value
-  Eigen::MatrixXd draws_ctr = (chains.array() - chains.mean()).abs().matrix();
+  // squared deviations from the mean
+  Eigen::MatrixXd draws_sq = (chains.array() - chains.mean()).square().matrix();
 
   // posterior pkg fn `ess_mean` computes on split chains
-  double ess_mean = ess(split_chains(draws_ctr));
+  double ess_mean = ess(split_chains(draws_sq));
 
   // estimated variance (2nd moment)
-  double Evar = draws_ctr.array().square().mean();
+  double Evar = draws_sq.mean();
 
   // variance of variance, adjusted for ESS
-  double fourth_moment = draws_ctr.array().pow(4).mean();
-  double varvar = (fourth_moment - std::pow(Evar, 2)) / ess_mean;
+  double varvar = (draws_sq.array() - Evar).square().mean() / ess_mean;
 
   // variance of standard deviation - use Taylor series approximation
   double varsd = varvar / Evar / 4.0;

src/stan/analyze/mcmc/split_rank_normalized_ess.hpp

@@ -43,14 +43,9 @@ inline std::pair<double, double> split_rank_normalized_ess(
                             .cast<double>();
   double ess_tail_95 = ess(q95);
 
-  double ess_tail;
-  if (std::isnan(ess_tail_05)) {
-    ess_tail = ess_tail_95;
-  } else if (std::isnan(ess_tail_95)) {
-    ess_tail = ess_tail_05;
-  } else {
-    ess_tail = std::min(ess_tail_05, ess_tail_95);
-  }
+  double ess_tail = std::isnan(ess_tail_05) || std::isnan(ess_tail_95)
+                        ? std::numeric_limits<double>::quiet_NaN()
+                        : std::min(ess_tail_05, ess_tail_95);
   return std::make_pair(ess_bulk, ess_tail);
 }
 

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

@@ -58,3 +58,22 @@ TEST(ProbAutocovariance, test2) {
   EXPECT_NEAR(1.10, ac(4), 0.01);
   EXPECT_NEAR(0.89, ac(5), 0.01);
 }
+
+TEST(ProbAutocovariance, constant_sequence) {
+  size_t N = 10;
+  Eigen::VectorXd y = Eigen::VectorXd::Constant(N, 0.3);
+  Eigen::VectorXd ac(N);
+  stan::analyze::autocovariance<double>(y, ac);
+  EXPECT_EQ(N, ac.size());
+  for (size_t i = 0; i < N; ++i) {
+    EXPECT_EQ(0.0, ac(i));
+  }
+
+  std::vector<double> y_vec(N, 0.3);
+  std::vector<double> ac_vec;
+  stan::analyze::autocovariance(y_vec, ac_vec);
+  EXPECT_EQ(N, ac_vec.size());
+  for (size_t i = 0; i < N; ++i) {
+    EXPECT_EQ(0.0, ac_vec[i]);
+  }
+}

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

@@ -46,3 +46,29 @@ TEST(CheckChains, good_and_bad) {
   chain_1(0, 0) = std::numeric_limits<double>::infinity();
   EXPECT_FALSE(stan::analyze::is_finite_and_varies(chain_1));
 }
+
+TEST(CheckChains, some_chains_constant) {
+  Eigen::MatrixXd chains(10, 3);
+  for (int i = 0; i < 10; ++i) {
+    chains(i, 0) = i;
+    chains(i, 1) = 0.3;
+    chains(i, 2) = -2.5;
+  }
+  // one chain varies
+  EXPECT_TRUE(stan::analyze::is_finite_and_varies(chains));
+
+  // every chain constant, at different values
+  chains.col(0).setConstant(1.0);
+  EXPECT_FALSE(stan::analyze::is_finite_and_varies(chains));
+
+  // every chain constant, at the same value
+  chains.setConstant(0.3);
+  EXPECT_FALSE(stan::analyze::is_finite_and_varies(chains));
+
+  // one chain varies, another chain has a non-finite value
+  for (int i = 0; i < 10; ++i) {
+    chains(i, 0) = i;
+  }
+  chains(4, 1) = std::numeric_limits<double>::quiet_NaN();
+  EXPECT_FALSE(stan::analyze::is_finite_and_varies(chains));
+}

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

@@ -51,3 +51,29 @@ TEST_F(EssBasic, test_basic_ess) {
   EXPECT_NEAR(ess_lp_expect, ess_basic_lp, 1e-8);
   EXPECT_NEAR(ess_theta_expect, ess_basic_theta, 1e-8);
 }
+
+TEST_F(EssBasic, constant_draws) {
+  Eigen::MatrixXd chains = Eigen::MatrixXd::Constant(1000, 4, 0.3);
+  EXPECT_TRUE(std::isnan(stan::analyze::ess(chains)));
+}
+
+TEST_F(EssBasic, constant_chain) {
+  // computed via R pkg posterior
+  chains_theta.col(3).setConstant(0.3);
+  EXPECT_NEAR(866.619903755711, stan::analyze::ess(chains_theta), 1e-8);
+}
+
+TEST_F(EssBasic, every_chain_constant) {
+  // computed via R pkg posterior
+  Eigen::MatrixXd chains(1000, 4);
+  for (int i = 0; i < 4; ++i) {
+    chains.col(i).setConstant(i + 1);
+  }
+  EXPECT_NEAR(2.00803212851406, stan::analyze::ess(chains), 1e-8);
+}
+
+TEST_F(EssBasic, autocorrelation_positive_up_to_max_lag) {
+  // computed via R pkg posterior
+  chains_theta.col(3).array() += 1.0;
+  EXPECT_NEAR(2.13486710959329, stan::analyze::ess(chains_theta), 1e-8);
+}

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

@@ -36,17 +36,26 @@ class MonteCarloStandardError : public testing::Test {
 
 TEST_F(MonteCarloStandardError, test_mcse) {
   // computed via R pkg posterior
-  double mcse_mean_lp_expect = 0.02016;
-  double mcse_mean_theta_expect = 0.00323;
+  double mcse_mean_lp_expect = 0.0201647780004714;
+  double mcse_mean_theta_expect = 0.00323399160714608;
 
-  double mcse_sd_lp_expect = 0.03553;
-  double mcse_sd_theta_expect = 0.00216;
-  EXPECT_NEAR(mcse_mean_lp_expect, stan::analyze::mcse_mean(chains_lp), 1e-4);
+  double mcse_sd_lp_expect = 0.0308898300566685;
+  double mcse_sd_theta_expect = 0.00216421366907844;
+  EXPECT_NEAR(mcse_mean_lp_expect, stan::analyze::mcse_mean(chains_lp), 1e-10);
   EXPECT_NEAR(mcse_mean_theta_expect, stan::analyze::mcse_mean(chains_theta),
-              1e-4);
+              1e-10);
 
-  EXPECT_NEAR(mcse_sd_lp_expect, stan::analyze::mcse_sd(chains_lp), 1e-4);
-  EXPECT_NEAR(mcse_sd_theta_expect, stan::analyze::mcse_sd(chains_theta), 1e-4);
+  EXPECT_NEAR(mcse_sd_lp_expect, stan::analyze::mcse_sd(chains_lp), 1e-10);
+  EXPECT_NEAR(mcse_sd_theta_expect, stan::analyze::mcse_sd(chains_theta),
+              1e-10);
+}
+
+TEST_F(MonteCarloStandardError, constant_chain) {
+  // computed via R pkg posterior
+  chains_theta.col(3).setConstant(0.3);
+  EXPECT_NEAR(0.00355436328690979, stan::analyze::mcse_mean(chains_theta),
+              1e-10);
+  EXPECT_NEAR(0.0121180542262653, stan::analyze::mcse_sd(chains_theta), 1e-10);
 }
 
 TEST_F(MonteCarloStandardError, const_fail) {

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

@@ -65,6 +65,62 @@ TEST_F(RankNormalizedEss, inf_fail) {
   EXPECT_TRUE(std::isnan(ess.second));
 }
 
+TEST_F(RankNormalizedEss, constant_chain) {
+  // computed via R pkg posterior
+  chains_theta.col(3).setConstant(0.3);
+  auto ess = stan::analyze::split_rank_normalized_ess(chains_theta);
+  EXPECT_NEAR(794.887967121017, ess.first, 1e-6);
+  EXPECT_NEAR(1194.06701581485, ess.second, 1e-6);
+}
+
+TEST_F(RankNormalizedEss, half_constant_chain) {
+  // computed via R pkg posterior
+  chains_theta.col(3).tail(500).setConstant(0.3);
+  auto ess = stan::analyze::split_rank_normalized_ess(chains_theta);
+  EXPECT_NEAR(1108.06541725348, ess.first, 1e-6);
+  EXPECT_NEAR(1288.5051347075, ess.second, 1e-6);
+}
+
+TEST_F(RankNormalizedEss, chain_outside_both_tails) {
+  // computed via R pkg posterior
+  // chain 4 is above the 5% quantile, and the other chains are below the
+  // 95% quantile
+  chains_theta.col(3).array() += 1.0;
+  auto ess = stan::analyze::split_rank_normalized_ess(chains_theta);
+  EXPECT_NEAR(7.24636932115751, ess.first, 1e-8);
+  EXPECT_NEAR(32.1273579765829, ess.second, 1e-8);
+}
+
+TEST_F(RankNormalizedEss, every_chain_constant) {
+  Eigen::MatrixXd chains(1000, 4);
+  for (int i = 0; i < 4; ++i) {
+    chains.col(i).setConstant(i + 1);
+  }
+  auto ess = stan::analyze::split_rank_normalized_ess(chains);
+  EXPECT_TRUE(std::isnan(ess.first));
+  EXPECT_TRUE(std::isnan(ess.second));
+}
+
+TEST_F(RankNormalizedEss, discrete_tails) {
+  // computed via R pkg posterior
+  // 46% of the draws are 1, so the upper tail indicator is x == 1
+  Eigen::MatrixXd chains = (chains_theta.array() > 0.25).cast<double>();
+  auto ess = stan::analyze::split_rank_normalized_ess(chains);
+  EXPECT_NEAR(1951.94660323879, ess.first, 1e-6);
+  EXPECT_NEAR(1951.94660323879, ess.second, 1e-6);
+
+  // 96% of the draws are 1, so every draw is in the lower tail
+  chains = (chains_theta.array() > 0.07).cast<double>();
+  ess = stan::analyze::split_rank_normalized_ess(chains);
+  EXPECT_NEAR(1347.60516728569, ess.first, 1e-6);
+  EXPECT_TRUE(std::isnan(ess.second));
+
+  // the mirror image
+  ess = stan::analyze::split_rank_normalized_ess(-chains);
+  EXPECT_NEAR(1347.60516728569, ess.first, 1e-6);
+  EXPECT_TRUE(std::isnan(ess.second));
+}
+
 TEST_F(RankNormalizedEss, short_chains_fail) {
   chains_theta.resize(3, 4);
   auto ess = stan::analyze::split_rank_normalized_ess(chains_theta);