Merge pull request #1497 from martinmodrak/bugfix/1496-poisson-phi-cutoff

Fixing negative binomial phi cutoff

Author: bob-carpenter

stan/math/opencl/kernel_generator/load.hpp

@@ -49,7 +49,7 @@ class load_
    * Creates a deep copy of this expression.
    * @return copy of \c *this
    */
-  inline load_<T&> deep_copy() const & { return load_<T&>(a_); }
+  inline load_<T&> deep_copy() const& { return load_<T&>(a_); }
   inline load_<T> deep_copy() && { return load_<T>(std::forward<T>(a_)); }
 
   /**

stan/math/prim/prob/neg_binomial_2_lpmf.hpp

@@ -3,15 +3,14 @@
 
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/binomial_coefficient_log.hpp>
 #include <stan/math/prim/fun/digamma.hpp>
-#include <stan/math/prim/fun/lgamma.hpp>
 #include <stan/math/prim/fun/log.hpp>
 #include <stan/math/prim/fun/max_size.hpp>
 #include <stan/math/prim/fun/multiply_log.hpp>
 #include <stan/math/prim/fun/size.hpp>
 #include <stan/math/prim/fun/size_zero.hpp>
 #include <stan/math/prim/fun/value_of.hpp>
-#include <stan/math/prim/prob/poisson_lpmf.hpp>
 #include <cmath>
 
 namespace stan {
@@ -47,7 +46,7 @@ return_type_t<T_location, T_precision> neg_binomial_2_lpmf(
   size_t size_phi = stan::math::size(phi);
   size_t size_mu_phi = max_size(mu, phi);
   size_t size_n_phi = max_size(n, phi);
-  size_t max_size_seq_view = max_size(n, mu, phi);
+  size_t size_all = max_size(n, mu, phi);
 
   VectorBuilder<true, T_partials_return, T_location> mu_val(size_mu);
   for (size_t i = 0; i < size_mu; ++i) {
@@ -76,39 +75,30 @@ return_type_t<T_location, T_precision> neg_binomial_2_lpmf(
     n_plus_phi[i] = n_vec[i] + phi_val[i];
   }
 
-  for (size_t i = 0; i < max_size_seq_view; i++) {
-    // if phi is large we probably overflow, defer to Poisson:
-    if (phi_val[i] > 1e5) {
-      // TODO(martinmodrak) This is wrong (doesn't pass propto information),
-      // and inaccurate for n = 0, but shouldn't break most models.
-      // Also the 1e5 cutoff is too small.
-      // Will be addressed better in PR #1497
-      logp += poisson_lpmf(n_vec[i], mu_val[i]);
-    } else {
-      if (include_summand<propto>::value) {
-        logp -= lgamma(n_vec[i] + 1.0);
-      }
-      if (include_summand<propto, T_precision>::value) {
-        logp += multiply_log(phi_val[i], phi_val[i]) - lgamma(phi_val[i]);
-      }
-      if (include_summand<propto, T_location>::value) {
-        logp += multiply_log(n_vec[i], mu_val[i]);
-      }
-      if (include_summand<propto, T_precision>::value) {
-        logp += lgamma(n_plus_phi[i]);
-      }
-      logp -= n_plus_phi[i] * log_mu_plus_phi[i];
+  for (size_t i = 0; i < size_all; i++) {
+    if (include_summand<propto, T_precision>::value) {
+      logp += binomial_coefficient_log(n_plus_phi[i] - 1, n_vec[i]);
+    }
+    if (include_summand<propto, T_location>::value) {
+      logp += multiply_log(n_vec[i], mu_val[i]);
     }
+    logp += -phi_val[i] * (log1p(mu_val[i] / phi_val[i]))
+            - n_vec[i] * log_mu_plus_phi[i];
 
     if (!is_constant_all<T_location>::value) {
       ops_partials.edge1_.partials_[i]
-          += n_vec[i] / mu_val[i] - n_plus_phi[i] / mu_plus_phi[i];
+          += n_vec[i] / mu_val[i] - (n_vec[i] + phi_val[i]) / (mu_plus_phi[i]);
     }
     if (!is_constant_all<T_precision>::value) {
-      ops_partials.edge2_.partials_[i] += 1.0 - n_plus_phi[i] / mu_plus_phi[i]
-                                          + log_phi[i] - log_mu_plus_phi[i]
-                                          - digamma(phi_val[i])
-                                          + digamma(n_plus_phi[i]);
+      T_partials_return log_term;
+      if (mu_val[i] < phi_val[i]) {
+        log_term = log1p(-mu_val[i] / (mu_plus_phi[i]));
+      } else {
+        log_term = log_phi[i] - log_mu_plus_phi[i];
+      }
+      ops_partials.edge2_.partials_[i]
+          += (mu_val[i] - n_vec[i]) / (mu_plus_phi[i]) + log_term
+             - (digamma(phi_val[i]) - digamma(n_plus_phi[i]));
     }
   }
   return ops_partials.build(logp);

test/unit/math/prim/prob/neg_binomial_2_log_test.cpp

@@ -212,7 +212,10 @@ TEST(ProbNegBinomial2, log_matches_lpmf) {
 TEST(ProbDistributionsNegBinomial2Log, neg_binomial_2_log_grid_test) {
   std::vector<double> mu_log_to_test
       = {-101, -27, -3, -1, -0.132, 0, 4, 10, 87};
-  std::vector<double> phi_to_test = {2e-5, 0.36, 1, 2.3e5, 1.8e10, 6e16};
+  // TODO(martinmodrak) Reducing the span of the test, should be fixed
+  // along with #1495
+  // std::vector<double> phi_to_test = {2e-5, 0.36, 1, 10, 2.3e5, 1.8e10, 6e16};
+  std::vector<double> phi_to_test = {0.36, 1, 10};
   std::vector<int> n_to_test = {0, 1, 10, 39, 101, 3048, 150054};
 
   // TODO(martinmdorak) Only weak tolerance for this quick fix

test/unit/math/prim/prob/neg_binomial_2_test.cpp

@@ -2,6 +2,7 @@
 #include <test/unit/math/prim/prob/vector_rng_test_helper.hpp>
 #include <test/unit/math/prim/prob/NegativeBinomial2LogTestRig.hpp>
 #include <test/unit/math/prim/prob/VectorIntRNGTestRig.hpp>
+#include <test/unit/math/expect_near_rel.hpp>
 #include <gtest/gtest.h>
 #include <boost/random/mersenne_twister.hpp>
 #include <boost/math/distributions.hpp>
@@ -238,27 +239,40 @@ TEST(ProbDistributionsNegBinomial2, chiSquareGoodnessFitTest4) {
 }
 
 TEST(ProbDistributionsNegBinomial2, extreme_values) {
-  int N = 100;
-  double mu = 8;
-  double phi = 1e12;
-  for (int n = 0; n < 10; ++n) {
-    phi *= 10;
-    double logp = stan::math::neg_binomial_2_log<false>(N, mu, phi);
-    EXPECT_LT(logp, 0);
+  std::vector<int> n_to_test = {0, 1, 5, 100, 12985, 1968422};
+  std::vector<double> mu_to_test = {1e-5, 0.1, 8, 713, 28311, 19850054};
+  for (double mu : mu_to_test) {
+    for (int n : n_to_test) {
+      // Test across a range of phi
+      for (double phi = 1e12; phi < 1e22; phi *= 10) {
+        double logp = stan::math::neg_binomial_2_log<false>(n, mu, phi);
+        EXPECT_LT(logp, 0) << "n = " << n << ", mu = " << mu
+                           << ", phi = " << phi;
+      }
+    }
   }
 }
 
-TEST(ProbDistributionsNegBinomial2, vectorAroundCutoff) {
-  int y = 10;
-  double mu = 9.36;
-  std::vector<double> phi;
-  phi.push_back(1);
-  phi.push_back(1e15);
-  double vector_value = stan::math::neg_binomial_2_lpmf(y, mu, phi);
-  double scalar_value = stan::math::neg_binomial_2_lpmf(y, mu, phi[0])
-                        + stan::math::neg_binomial_2_lpmf(y, mu, phi[1]);
-
-  EXPECT_FLOAT_EQ(vector_value, scalar_value);
+TEST(ProbDistributionsNegBinomial2, zeroOne) {
+  using stan::test::expect_near_rel;
+
+  std::vector<double> mu_to_test = {2.345e-5, 0.2, 13, 150, 1621, 18432, 1e10};
+  double phi_start = 1e-8;
+  double phi_max = 1e22;
+  for (double mu : mu_to_test) {
+    for (double phi = phi_start; phi < phi_max; phi *= stan::math::pi()) {
+      std::stringstream msg;
+      msg << ", mu = " << mu << ", phi = " << phi;
+
+      double expected_value_0 = phi * (-log1p(mu / phi));
+      double value_0 = stan::math::neg_binomial_2_lpmf(0, mu, phi);
+      expect_near_rel("n = 0 " + msg.str(), value_0, expected_value_0);
+
+      double expected_value_1 = (phi + 1) * (-log1p(mu / phi)) + log(mu);
+      double value_1 = stan::math::neg_binomial_2_lpmf(1, mu, phi);
+      expect_near_rel("n = 1 " + msg.str(), value_1, expected_value_1);
+    }
+  }
 }
 
 TEST(ProbDistributionsNegativeBinomial2Log, distributionCheck) {

test/unit/math/rev/fun/lbeta_test.cpp

@@ -74,11 +74,6 @@ TEST(MathFunctions, lbeta_identities_gradient) {
   // Successors: beta(a,b) = beta(a + 1, b) + beta(a, b + 1)
   for (double x : to_test) {
     for (double y : to_test) {
-      // TODO(martinmodrak) this restriction on testing should be lifted once
-      // the log_sum_exp bug (#1679) is resolved
-      if (x > 1e10 || y > 1e10) {
-        continue;
-      }
       auto rh = [](const var& a, const var& b) {
         return stan::math::log_sum_exp(lbeta(a + 1, b), lbeta(a, b + 1));
       };