correctness fixes in neg_binomial_* functions

stan/math/prim/prob/neg_binomial_2_cdf.hpp

@@ -3,11 +3,13 @@
 
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/err.hpp>
-#include <stan/math/prim/fun/size_zero.hpp>
 #include <stan/math/prim/fun/digamma.hpp>
 #include <stan/math/prim/fun/inc_beta.hpp>
 #include <stan/math/prim/fun/inc_beta_dda.hpp>
 #include <stan/math/prim/fun/inc_beta_ddz.hpp>
+#include <stan/math/prim/fun/inv.hpp>
+#include <stan/math/prim/fun/size_zero.hpp>
+#include <stan/math/prim/fun/square.hpp>
 #include <stan/math/prim/fun/value_of.hpp>
 #include <limits>
 
@@ -34,6 +36,8 @@ return_type_t<T_location, T_precision> neg_binomial_2_cdf(
   scalar_seq_view<T_n> n_vec(n);
   scalar_seq_view<T_location> mu_vec(mu);
   scalar_seq_view<T_precision> phi_vec(phi);
+  size_t size_phi = size(phi);
+  size_t size_n_phi = max_size(n, phi);
   size_t max_size_seq_view = max_size(n, mu, phi);
 
   operands_and_partials<T_location, T_precision> ops_partials(mu, phi);
@@ -48,18 +52,18 @@ return_type_t<T_location, T_precision> neg_binomial_2_cdf(
 
   VectorBuilder<!is_constant_all<T_precision>::value, T_partials_return,
                 T_precision>
-      digamma_phi_vec(size(phi));
-
-  VectorBuilder<!is_constant_all<T_precision>::value, T_partials_return,
+      digamma_phi_vec(size_phi);
+  VectorBuilder<!is_constant_all<T_precision>::value, T_partials_return, T_n,
                 T_precision>
-      digamma_sum_vec(size(phi));
+      digamma_sum_vec(size_n_phi);
 
   if (!is_constant_all<T_precision>::value) {
-    for (size_t i = 0; i < size(phi); i++) {
+    for (size_t i = 0; i < size_phi; i++) {
+      digamma_phi_vec[i] = digamma(value_of(phi_vec[i]));
+    }
+    for (size_t i = 0; i < size_n_phi; i++) {
       const T_partials_return n_dbl = value_of(n_vec[i]);
       const T_partials_return phi_dbl = value_of(phi_vec[i]);
-
-      digamma_phi_vec[i] = digamma(phi_dbl);
       digamma_sum_vec[i] = digamma(n_dbl + phi_dbl + 1);
     }
   }
@@ -71,29 +75,30 @@ return_type_t<T_location, T_precision> neg_binomial_2_cdf(
       return ops_partials.build(1.0);
     }
 
-    const T_partials_return n_dbl = value_of(n_vec[i]);
+    const T_partials_return n_dbl_p1 = value_of(n_vec[i]) + 1;
     const T_partials_return mu_dbl = value_of(mu_vec[i]);
     const T_partials_return phi_dbl = value_of(phi_vec[i]);
-
-    const T_partials_return p_dbl = phi_dbl / (mu_dbl + phi_dbl);
-    const T_partials_return d_dbl
-        = 1.0 / ((mu_dbl + phi_dbl) * (mu_dbl + phi_dbl));
-
-    const T_partials_return P_i = inc_beta(phi_dbl, n_dbl + 1.0, p_dbl);
+    const T_partials_return inv_mu_plus_phi = inv(mu_dbl + phi_dbl);
+    const T_partials_return p_dbl = phi_dbl * inv_mu_plus_phi;
+    const T_partials_return d_dbl = square(inv_mu_plus_phi);
+    const T_partials_return P_i = inc_beta(phi_dbl, n_dbl_p1, p_dbl);
+    const T_partials_return inc_beta_ddz_i
+        = is_constant_all<T_location, T_precision>::value
+              ? 0
+              : inc_beta_ddz(phi_dbl, n_dbl_p1, p_dbl) * d_dbl / P_i;
 
     P *= P_i;
 
     if (!is_constant_all<T_location>::value) {
-      ops_partials.edge1_.partials_[i]
-          += -inc_beta_ddz(phi_dbl, n_dbl + 1.0, p_dbl) * phi_dbl * d_dbl / P_i;
+      ops_partials.edge1_.partials_[i] -= inc_beta_ddz_i * phi_dbl;
     }
 
     if (!is_constant_all<T_precision>::value) {
       ops_partials.edge2_.partials_[i]
-          += inc_beta_dda(phi_dbl, n_dbl + 1, p_dbl, digamma_phi_vec[i],
+          += inc_beta_dda(phi_dbl, n_dbl_p1, p_dbl, digamma_phi_vec[i],
                           digamma_sum_vec[i])
                  / P_i
-             + inc_beta_ddz(phi_dbl, n_dbl + 1.0, p_dbl) * mu_dbl * d_dbl / P_i;
+             + inc_beta_ddz_i * mu_dbl;
     }
   }
 
@@ -104,7 +109,7 @@ return_type_t<T_location, T_precision> neg_binomial_2_cdf(
   }
 
   if (!is_constant_all<T_precision>::value) {
-    for (size_t i = 0; i < size(phi); ++i) {
+    for (size_t i = 0; i < size_phi; ++i) {
       ops_partials.edge2_.partials_[i] *= P;
     }
   }

stan/math/prim/prob/neg_binomial_2_lccdf.hpp

@@ -27,7 +27,6 @@ return_type_t<T_location, T_precision> neg_binomial_2_lccdf(
 
   scalar_seq_view<T_location> mu_vec(mu);
   scalar_seq_view<T_precision> phi_vec(phi);
-
   size_t size_beta = max_size(mu, phi);
 
   VectorBuilder<true, return_type_t<T_location, T_precision>, T_location,
@@ -37,7 +36,7 @@ return_type_t<T_location, T_precision> neg_binomial_2_lccdf(
     beta_vec[i] = phi_vec[i] / mu_vec[i];
   }
 
-  return neg_binomial_ccdf_log(n, phi, beta_vec.data());
+  return neg_binomial_lccdf(n, phi, beta_vec.data());
 }
 
 }  // namespace math

stan/math/prim/prob/neg_binomial_2_lcdf.hpp

@@ -3,6 +3,7 @@
 
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/err.hpp>
+#include <stan/math/prim/fun/constants.hpp>
 #include <stan/math/prim/fun/size_zero.hpp>
 #include <stan/math/prim/prob/beta_cdf_log.hpp>
 #include <cmath>
@@ -29,23 +30,25 @@ return_type_t<T_location, T_precision> neg_binomial_2_lcdf(
   scalar_seq_view<T_n> n_vec(n);
   scalar_seq_view<T_location> mu_vec(mu);
   scalar_seq_view<T_precision> phi_vec(phi);
-
+  size_t size_n = size(n);
   size_t size_phi_mu = max_size(mu, phi);
+
+  for (size_t i = 0; i < size_n; i++) {
+    if (n_vec[i] < 0) {
+      return LOG_ZERO;
+    }
+  }
+
   VectorBuilder<true, return_type_t<T_location, T_precision>, T_location,
                 T_precision>
       phi_mu(size_phi_mu);
   for (size_t i = 0; i < size_phi_mu; i++) {
     phi_mu[i] = phi_vec[i] / (phi_vec[i] + mu_vec[i]);
   }
 
-  size_t size_n = size(n);
   VectorBuilder<true, return_type_t<T_n>, T_n> np1(size_n);
   for (size_t i = 0; i < size_n; i++) {
-    if (n_vec[i] < 0) {
-      return log(0.0);
-    } else {
-      np1[i] = n_vec[i] + 1.0;
-    }
+    np1[i] = n_vec[i] + 1.0;
   }
 
   return beta_cdf_log(phi_mu.data(), phi, np1.data());

stan/math/prim/prob/neg_binomial_2_log_lpmf.hpp

@@ -48,58 +48,63 @@ return_type_t<T_log_location, T_precision> neg_binomial_2_log_lpmf(
   scalar_seq_view<T_n> n_vec(n);
   scalar_seq_view<T_log_location> eta_vec(eta);
   scalar_seq_view<T_precision> phi_vec(phi);
+  size_t size_eta = size(eta);
+  size_t size_phi = size(phi);
+  size_t size_eta_phi = max_size(eta, phi);
+  size_t size_n_phi = max_size(n, phi);
   size_t max_size_seq_view = max_size(n, eta, phi);
 
   operands_and_partials<T_log_location, T_precision> ops_partials(eta, phi);
 
-  size_t len_ep = max_size(eta, phi);
-  size_t len_np = max_size(n, phi);
-
-  VectorBuilder<true, T_partials_return, T_log_location> eta__(size(eta));
-  for (size_t i = 0, max_size_seq_view = size(eta); i < max_size_seq_view;
-       ++i) {
-    eta__[i] = value_of(eta_vec[i]);
+  VectorBuilder<true, T_partials_return, T_log_location> eta_val(size_eta);
+  for (size_t i = 0; i < size_eta; ++i) {
+    eta_val[i] = value_of(eta_vec[i]);
   }
 
-  VectorBuilder<true, T_partials_return, T_precision> phi__(size(phi));
-  for (size_t i = 0, max_size_seq_view = size(phi); i < max_size_seq_view;
-       ++i) {
-    phi__[i] = value_of(phi_vec[i]);
+  VectorBuilder<!is_constant_all<T_log_location, T_precision>::value,
+                T_partials_return, T_log_location>
+      exp_eta(size_eta);
+  if (!is_constant_all<T_log_location, T_precision>::value) {
+    for (size_t i = 0; i < size_eta; ++i) {
+      exp_eta[i] = exp(eta_val[i]);
+    }
   }
 
-  VectorBuilder<true, T_partials_return, T_precision> log_phi(size(phi));
-  for (size_t i = 0, max_size_seq_view = size(phi); i < max_size_seq_view;
-       ++i) {
-    log_phi[i] = log(phi__[i]);
+  VectorBuilder<true, T_partials_return, T_precision> phi_val(size_phi);
+  VectorBuilder<true, T_partials_return, T_precision> log_phi(size_phi);
+  for (size_t i = 0; i < size_phi; ++i) {
+    phi_val[i] = value_of(phi_vec[i]);
+    log_phi[i] = log(phi_val[i]);
   }
 
   VectorBuilder<true, T_partials_return, T_log_location, T_precision>
-      logsumexp_eta_logphi(len_ep);
-  for (size_t i = 0; i < len_ep; ++i) {
-    logsumexp_eta_logphi[i] = log_sum_exp(eta__[i], log_phi[i]);
+      logsumexp_eta_logphi(size_eta_phi);
+  for (size_t i = 0; i < size_eta_phi; ++i) {
+    logsumexp_eta_logphi[i] = log_sum_exp(eta_val[i], log_phi[i]);
   }
 
-  VectorBuilder<true, T_partials_return, T_n, T_precision> n_plus_phi(len_np);
-  for (size_t i = 0; i < len_np; ++i) {
-    n_plus_phi[i] = n_vec[i] + phi__[i];
+  VectorBuilder<true, T_partials_return, T_n, T_precision> n_plus_phi(
+      size_n_phi);
+  for (size_t i = 0; i < size_n_phi; ++i) {
+    n_plus_phi[i] = n_vec[i] + phi_val[i];
   }
 
   for (size_t i = 0; i < max_size_seq_view; i++) {
-    if (phi__[i] > 1e5) {
+    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 way too low.
       // Will be adressed better once PR #1497 is merged
-      logp += poisson_log_lpmf(n_vec[i], eta__[i]);
+      logp += poisson_log_lpmf(n_vec[i], eta_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__[i], phi__[i]) - lgamma(phi__[i]);
+        logp += multiply_log(phi_val[i], phi_val[i]) - lgamma(phi_val[i]);
       }
       if (include_summand<propto, T_log_location>::value) {
-        logp += n_vec[i] * eta__[i];
+        logp += n_vec[i] * eta_val[i];
       }
       if (include_summand<propto, T_precision>::value) {
         logp += lgamma(n_plus_phi[i]);
@@ -109,12 +114,12 @@ return_type_t<T_log_location, T_precision> neg_binomial_2_log_lpmf(
 
     if (!is_constant_all<T_log_location>::value) {
       ops_partials.edge1_.partials_[i]
-          += n_vec[i] - n_plus_phi[i] / (phi__[i] / exp(eta__[i]) + 1.0);
+          += n_vec[i] - n_plus_phi[i] / (phi_val[i] / exp_eta[i] + 1.0);
     }
     if (!is_constant_all<T_precision>::value) {
       ops_partials.edge2_.partials_[i]
-          += 1.0 - n_plus_phi[i] / (exp(eta__[i]) + phi__[i]) + log_phi[i]
-             - logsumexp_eta_logphi[i] - digamma(phi__[i])
+          += 1.0 - n_plus_phi[i] / (exp_eta[i] + phi_val[i]) + log_phi[i]
+             - logsumexp_eta_logphi[i] - digamma(phi_val[i])
              + digamma(n_plus_phi[i]);
     }
   }

stan/math/prim/prob/neg_binomial_2_lpmf.hpp

@@ -3,10 +3,10 @@
 
 #include <stan/math/prim/meta.hpp>
 #include <stan/math/prim/err.hpp>
-#include <stan/math/prim/fun/size_zero.hpp>
-#include <stan/math/prim/fun/multiply_log.hpp>
 #include <stan/math/prim/fun/digamma.hpp>
 #include <stan/math/prim/fun/lgamma.hpp>
+#include <stan/math/prim/fun/multiply_log.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>
@@ -42,74 +42,74 @@ return_type_t<T_location, T_precision> neg_binomial_2_lpmf(
   scalar_seq_view<T_n> n_vec(n);
   scalar_seq_view<T_location> mu_vec(mu);
   scalar_seq_view<T_precision> phi_vec(phi);
+  size_t size_mu = size(mu);
+  size_t size_phi = 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);
 
   operands_and_partials<T_location, T_precision> ops_partials(mu, phi);
 
-  size_t len_ep = max_size(mu, phi);
-  size_t len_np = max_size(n, phi);
-
-  VectorBuilder<true, T_partials_return, T_location> mu__(size(mu));
-
-  for (size_t i = 0, max_size_seq_view = size(mu); i < max_size_seq_view; ++i) {
-    mu__[i] = value_of(mu_vec[i]);
-  }
-
-  VectorBuilder<true, T_partials_return, T_precision> phi__(size(phi));
-  for (size_t i = 0, max_size_seq_view = size(phi); i < max_size_seq_view;
-       ++i) {
-    phi__[i] = value_of(phi_vec[i]);
+  VectorBuilder<true, T_partials_return, T_location> mu_val(size_mu);
+  for (size_t i = 0; i < size_mu; ++i) {
+    mu_val[i] = value_of(mu_vec[i]);
   }
 
-  VectorBuilder<true, T_partials_return, T_precision> log_phi(size(phi));
-  for (size_t i = 0, max_size_seq_view = size(phi); i < max_size_seq_view;
-       ++i) {
-    log_phi[i] = log(phi__[i]);
+  VectorBuilder<true, T_partials_return, T_precision> phi_val(size_phi);
+  VectorBuilder<true, T_partials_return, T_precision> log_phi(size_phi);
+  for (size_t i = 0; i < size_phi; ++i) {
+    phi_val[i] = value_of(phi_vec[i]);
+    log_phi[i] = log(phi_val[i]);
   }
 
+  VectorBuilder<true, T_partials_return, T_location, T_precision> mu_plus_phi(
+      size_mu_phi);
   VectorBuilder<true, T_partials_return, T_location, T_precision>
-      log_mu_plus_phi(len_ep);
-  for (size_t i = 0; i < len_ep; ++i) {
-    log_mu_plus_phi[i] = log(mu__[i] + phi__[i]);
+      log_mu_plus_phi(size_mu_phi);
+  for (size_t i = 0; i < size_mu_phi; ++i) {
+    mu_plus_phi[i] = mu_val[i] + phi_val[i];
+    log_mu_plus_phi[i] = log(mu_plus_phi[i]);
   }
 
-  VectorBuilder<true, T_partials_return, T_n, T_precision> n_plus_phi(len_np);
-  for (size_t i = 0; i < len_np; ++i) {
-    n_plus_phi[i] = n_vec[i] + phi__[i];
+  VectorBuilder<true, T_partials_return, T_n, T_precision> n_plus_phi(
+      size_n_phi);
+  for (size_t i = 0; i < size_n_phi; ++i) {
+    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__[i] > 1e5) {
+    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 adressed better in PR #1497
-      logp += poisson_lpmf(n_vec[i], mu__[i]);
+      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__[i], phi__[i]) - lgamma(phi__[i]);
+        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__[i]);
+        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];
+      logp -= n_plus_phi[i] * log_mu_plus_phi[i];
     }
 
     if (!is_constant_all<T_location>::value) {
       ops_partials.edge1_.partials_[i]
-          += n_vec[i] / mu__[i] - (n_vec[i] + phi__[i]) / (mu__[i] + phi__[i]);
+          += n_vec[i] / mu_val[i] - n_plus_phi[i] / mu_plus_phi[i];
     }
     if (!is_constant_all<T_precision>::value) {
-      ops_partials.edge2_.partials_[i]
-          += 1.0 - n_plus_phi[i] / (mu__[i] + phi__[i]) + log_phi[i]
-             - log_mu_plus_phi[i] - digamma(phi__[i]) + digamma(n_plus_phi[i]);
+      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]);
     }
   }
   return ops_partials.build(logp);