Handle degenerate negative-binomial parameters before sampling

include/boost/random/negative_binomial_distribution.hpp

@@ -16,6 +16,7 @@
 #include <iosfwd>
 #include <limits>
 
+#include <boost/assert.hpp>
 #include <boost/random/detail/config.hpp>
 #include <boost/random/gamma_distribution.hpp>
 #include <boost/random/poisson_distribution.hpp>
@@ -32,6 +33,8 @@ namespace random {
  * \f$\displaystyle P(i) = {k+i-1\choose i}p^k(1-p)^i\f$.
  *
  * This implementation uses a gamma-poisson mixture.
+ * When @c k is zero or @c p is one, the result is always zero.
+ * For positive @c k, @c p must be strictly positive.
  */
 template<class IntType = int, class RealType = double>
 class negative_binomial_distribution {
@@ -46,7 +49,7 @@ class negative_binomial_distribution {
          * Construct a param_type object.  @c k and @c p
          * are the parameters of the distribution.
          *
-         * Requires: k >=0 && 0 <= p <= 1
+         * Requires: k >=0 && 0 <= p <= 1 && (k == 0 || p > 0)
          */
         explicit param_type(IntType k_arg = 1, RealType p_arg = RealType (0.5))
           : _k(k_arg), _p(p_arg)
@@ -94,7 +97,7 @@ class negative_binomial_distribution {
      * Construct a @c negative_binomial_distribution object. @c k and @c p
      * are the parameters of the distribution.
      *
-     * Requires: k >=0 && 0 <= p <= 1
+     * Requires: k >=0 && 0 <= p <= 1 && (k == 0 || p > 0)
      */
     explicit negative_binomial_distribution(IntType k_arg = 1,
                                             RealType p_arg = RealType(0.5))
@@ -116,6 +119,9 @@ class negative_binomial_distribution {
     template<class URNG>
     IntType operator()(URNG& urng) const
     {
+        if(_k == IntType(0) || _p == RealType(1))
+            return IntType(0);
+        BOOST_ASSERT(_p > RealType(0));
         gamma_distribution<RealType> gamma(_k, (1-_p)/_p);
         poisson_distribution<IntType, RealType> poisson(gamma(urng));
         return poisson(urng);

test/test_negative_binomial_distribution.cpp

@@ -10,7 +10,10 @@
  */
 
 #include <boost/random/negative_binomial_distribution.hpp>
+#include <boost/random/linear_congruential.hpp>
+#include <iomanip>
 #include <limits>
+#include <sstream>
 
 #define BOOST_RANDOM_DISTRIBUTION boost::random::negative_binomial_distribution<>
 #define BOOST_RANDOM_ARG1 k
@@ -35,3 +38,107 @@
 #define BOOST_RANDOM_TEST2_MIN 50
 
 #include "test_distribution.ipp"
+
+namespace {
+
+// Forwards to a real engine so non-degenerate draws terminate, and counts every call
+struct counting_engine
+{
+    typedef boost::random::minstd_rand::result_type result_type;
+
+    counting_engine() : calls {0} {}
+    static result_type min() { return (boost::random::minstd_rand::min)(); }
+    static result_type max() { return (boost::random::minstd_rand::max)(); }
+    result_type operator()() { ++calls; return base(); }
+
+    boost::random::minstd_rand base;
+    unsigned calls;
+};
+
+template<class IntType, class RealType>
+void check_degenerate(IntType k, RealType p)
+{
+    typedef boost::random::negative_binomial_distribution<IntType, RealType> distribution;
+    typedef typename distribution::param_type param_type;
+    counting_engine engine;
+    const param_type param(k, p);
+    distribution direct(k, p);
+    distribution from_param(param);
+    distribution assigned;
+    assigned.param(param);
+    distribution explicit_param(5, RealType(0.5));
+    distribution streamed;
+    std::stringstream stream;
+    // At default precision min() prints as a subnormal, which libc++ fails to read back
+    stream << std::setprecision(std::numeric_limits<RealType>::max_digits10) << direct;
+    stream >> streamed;
+    BOOST_REQUIRE(!stream.fail());
+
+    for(unsigned i = 0; i < 4; ++i) {
+        BOOST_TEST(direct(engine) == IntType(0));
+        BOOST_TEST(from_param(engine) == IntType(0));
+        BOOST_TEST(assigned(engine) == IntType(0));
+        BOOST_TEST(explicit_param(engine, param) == IntType(0));
+        BOOST_TEST(streamed(engine) == IntType(0));
+    }
+    BOOST_TEST(engine.calls == 0u);
+    BOOST_TEST(explicit_param.k() == IntType(5));
+    BOOST_TEST(explicit_param.p() == RealType(0.5));
+}
+
+template<class IntType, class RealType>
+void check_zero_target()
+{
+    const RealType probabilities[] = {
+        RealType(0), (std::numeric_limits<RealType>::min)(),
+        RealType(0.25), RealType(0.5), RealType(0.75), RealType(1)
+    };
+    for(unsigned i = 0; i < sizeof(probabilities) / sizeof(probabilities[0]); ++i)
+        check_degenerate(IntType(0), probabilities[i]);
+}
+
+template<class IntType, class RealType>
+void check_unit_probability()
+{
+    const IntType targets[] = {
+        IntType(0), IntType(1), IntType(10), (std::numeric_limits<IntType>::max)()
+    };
+    for(unsigned i = 0; i < sizeof(targets) / sizeof(targets[0]); ++i)
+        check_degenerate(targets[i], RealType(1));
+}
+
+}
+
+BOOST_AUTO_TEST_CASE(test_zero_target)
+{
+    check_zero_target<int, float>();
+    check_zero_target<int, double>();
+    check_zero_target<int, long double>();
+    check_zero_target<boost::int64_t, double>();
+    check_zero_target<unsigned, float>();
+    check_zero_target<boost::uint64_t, long double>();
+}
+
+BOOST_AUTO_TEST_CASE(test_unit_probability)
+{
+    check_unit_probability<int, float>();
+    check_unit_probability<int, double>();
+    check_unit_probability<int, long double>();
+    check_unit_probability<boost::int64_t, double>();
+    check_unit_probability<unsigned, float>();
+    check_unit_probability<boost::uint64_t, long double>();
+}
+
+// Positive control: a non-degenerate draw must consume the engine,
+// so the engine.calls == 0 checks above would catch a sampling call
+BOOST_AUTO_TEST_CASE(test_non_degenerate_uses_engine)
+{
+    counting_engine engine;
+    boost::random::negative_binomial_distribution<> dist {1, 0.5};
+    for (unsigned i {0}; i < 4; ++i)
+    {
+        const int value {dist(engine)};
+        BOOST_TEST(value >= 0);
+    }
+    BOOST_TEST(engine.calls > 0u);
+}