Default constructor for discrete distribution

include/boost/compute/random/discrete_distribution.hpp

@@ -11,6 +11,9 @@
 #ifndef BOOST_COMPUTE_RANDOM_DISCRETE_DISTRIBUTION_HPP
 #define BOOST_COMPUTE_RANDOM_DISCRETE_DISTRIBUTION_HPP
 
+#include <numeric>
+
+#include <boost/config.hpp>
 #include <boost/type_traits.hpp>
 #include <boost/static_assert.hpp>
 
@@ -41,28 +44,42 @@ class discrete_distribution
 public:
     typedef IntType result_type;
 
+    /// Creates a new discrete distribution with a single weight p = { 1 }.
+    /// This distribution produces only zeroes.
+    discrete_distribution()
+        : m_probabilities(1, double(1)),
+          m_scanned_probabilities(1, double(1))
+    {
+
+    }
+
     /// Creates a new discrete distribution with weights given by
-    /// the range [\p first, \p last)
+    /// the range [\p first, \p last).
     template<class InputIterator>
     discrete_distribution(InputIterator first, InputIterator last)
-        : m_n((std::max)(size_t(1), static_cast<size_t>(std::distance(first, last)))),
-          m_probabilities((std::max)(size_t(1), static_cast<size_t>(std::distance(first, last))))
+        : m_probabilities(first, last),
+          m_scanned_probabilities(std::distance(first, last))
     {
-        double sum = 0;
-
-        for(InputIterator iter = first; iter!=last; iter++)
-        {
-            sum += *iter;
+        if(first != last) {
+            // after this m_scanned_probabilities.back() is a sum of all
+            // weights from the range [first, last)
+            std::partial_sum(first, last, m_scanned_probabilities.begin());
+
+            std::vector<double>::iterator i = m_probabilities.begin();
+            std::vector<double>::iterator j = m_scanned_probabilities.begin();
+            for(; i != m_probabilities.end(); ++i, ++j)
+            {
+                // dividing each weight by sum of all weights to
+                // get probabilities
+                *i = *i / m_scanned_probabilities.back();
+                // dividing each partial sum of weights by sum of
+                // all weights to get partial sums of probabilities
+                *j = *j / m_scanned_probabilities.back();
+            }
         }
-
-        InputIterator iter = first;
-        m_probabilities[0] = (*iter)/sum;
-        iter++;
-
-        for(size_t i = 1; i < m_n; i++)
-        {
-            m_probabilities[i] = m_probabilities[i-1] + (*iter)/sum;
-            iter++;
+        else {
+            m_probabilities.push_back(double(1));
+            m_scanned_probabilities.push_back(double(1));
         }
     }
 
@@ -71,52 +88,65 @@ class discrete_distribution
     {
     }
 
-    /// Returns the value of n
-    result_type n() const
-    {
-        return m_n;
-    }
-
     /// Returns the probabilities
     ::std::vector<double> probabilities() const
     {
         return m_probabilities;
     }
 
-    /// Generates uniformily distributed integers and stores
+    /// Returns the minimum potentially generated value.
+    result_type min BOOST_PREVENT_MACRO_SUBSTITUTION () const
+    {
+        return result_type(0);
+    }
+
+    /// Returns the maximum potentially generated value.
+    result_type max BOOST_PREVENT_MACRO_SUBSTITUTION () const
+    {
+        size_t type_max = static_cast<size_t>(
+            (std::numeric_limits<result_type>::max)()
+        );
+        if(m_probabilities.size() - 1 > type_max) {
+            return (std::numeric_limits<result_type>::max)();
+        }
+        return static_cast<result_type>(m_probabilities.size() - 1);
+    }
+
+    /// Generates uniformly distributed integers and stores
     /// them to the range [\p first, \p last).
     template<class OutputIterator, class Generator>
     void generate(OutputIterator first,
                   OutputIterator last,
                   Generator &generator,
                   command_queue &queue)
     {
-        std::string source = "inline uint scale_random(uint x)\n";
+        std::string source = "inline IntType scale_random(uint x)\n";
 
         source = source +
             "{\n" +
             "float rno = convert_float(x) / UINT_MAX;\n";
-        for(size_t i=0; i<m_n; i++)
+        for(size_t i = 0; i < m_scanned_probabilities.size() - 1; i++)
         {
             source = source +
-                "if(rno <= " + detail::make_literal<float>(m_probabilities[i]) + ")\n" +
+                "if(rno <= " + detail::make_literal<float>(m_scanned_probabilities[i]) + ")\n" +
                 "   return " + detail::make_literal(i) + ";\n";
         }
 
         source = source +
-            "return " + detail::make_literal(m_n - 1) + ";\n" +
+            "return " + detail::make_literal(m_scanned_probabilities.size() - 1) + ";\n" +
             "}\n";
 
         BOOST_COMPUTE_FUNCTION(IntType, scale_random, (const uint_ x), {});
 
         scale_random.set_source(source);
+        scale_random.define("IntType", type_name<IntType>());
 
         generator.generate(first, last, scale_random, queue);
     }
 
 private:
-    size_t m_n;
     ::std::vector<double> m_probabilities;
+    ::std::vector<double> m_scanned_probabilities;
 
     BOOST_STATIC_ASSERT_MSG(
         boost::is_integral<IntType>::value,

test/test_discrete_distribution.cpp

@@ -11,6 +11,8 @@
 #define BOOST_TEST_MODULE TestDiscreteDistribution
 #include <boost/test/unit_test.hpp>
 
+#include <vector>
+
 #include <boost/compute/system.hpp>
 #include <boost/compute/command_queue.hpp>
 #include <boost/compute/algorithm/count_if.hpp>
@@ -51,4 +53,185 @@ distribution.generate(vec.begin(), vec.end(), engine, queue);
     );
 }
 
+BOOST_AUTO_TEST_CASE(discrete_distribution)
+{
+    using boost::compute::uint_;
+    using boost::compute::lambda::_1;
+
+    size_t size = 100;
+    boost::compute::vector<uint_> vec(size, context);
+
+    // initialize the default random engine
+    boost::compute::default_random_engine engine(queue);
+
+    // initialize weights
+    int weights[] = {10, 40, 40, 10};
+
+    // setup the discrete distribution
+    boost::compute::discrete_distribution<uint_> distribution(
+        weights, weights + 4
+    );
+
+    std::vector<double> p = distribution.probabilities();
+    BOOST_CHECK_CLOSE(p[0], double(0.1), 0.001);
+    BOOST_CHECK_CLOSE(p[1], double(0.4), 0.001);
+    BOOST_CHECK_CLOSE(p[2], double(0.4), 0.001);
+    BOOST_CHECK_CLOSE(p[3], double(0.1), 0.001);
+
+    BOOST_CHECK_EQUAL((distribution.min)(), uint_(0));
+    BOOST_CHECK_EQUAL((distribution.max)(), uint_(3));
+
+    // generate the random values and store them to 'vec'
+    distribution.generate(vec.begin(), vec.end(), engine, queue);
+
+    BOOST_CHECK_EQUAL(
+        boost::compute::count_if(
+            vec.begin(), vec.end(), _1 < 4, queue
+        ),
+        size
+    );
+}
+
+BOOST_AUTO_TEST_CASE(discrete_distribution_default_ctor)
+{
+    using boost::compute::uint_;
+    using boost::compute::lambda::_1;
+
+    size_t size = 100;
+    boost::compute::vector<uint_> vec(size, context);
+
+    // initialize the default random engine
+    boost::compute::default_random_engine engine(queue);
+
+    // call default constructor
+    boost::compute::discrete_distribution<uint_> distribution;
+
+    std::vector<double> p = distribution.probabilities();
+    BOOST_CHECK_CLOSE(p[0], double(1), 0.001);
+
+    // generate the random values and store them to 'vec'
+    distribution.generate(vec.begin(), vec.end(), engine, queue);
+
+    BOOST_CHECK_EQUAL(
+        boost::compute::count_if(
+            vec.begin(), vec.end(), _1 == 0, queue
+        ),
+        size
+    );
+}
+
+BOOST_AUTO_TEST_CASE(discrete_distribution_one_weight)
+{
+    using boost::compute::uint_;
+    using boost::compute::lambda::_1;
+
+    size_t size = 100;
+    boost::compute::vector<uint_> vec(size, context);
+
+    // initialize the default random engine
+    boost::compute::default_random_engine engine(queue);
+
+    std::vector<int> weights(1, 1);
+    // call default constructor
+    boost::compute::discrete_distribution<uint_> distribution(
+        weights.begin(), weights.end()
+    );
+
+    std::vector<double> p = distribution.probabilities();
+    BOOST_CHECK_CLOSE(p[0], double(1), 0.001);
+
+    BOOST_CHECK_EQUAL((distribution.min)(), uint_(0));
+    BOOST_CHECK_EQUAL((distribution.max)(), uint_(0));
+
+    // generate the random values and store them to 'vec'
+    distribution.generate(vec.begin(), vec.end(), engine, queue);
+
+    BOOST_CHECK_EQUAL(
+        boost::compute::count_if(
+            vec.begin(), vec.end(), _1 == 0, queue
+        ),
+        size
+    );
+}
+
+BOOST_AUTO_TEST_CASE(discrete_distribution_empty_weights)
+{
+    using boost::compute::uint_;
+    using boost::compute::lambda::_1;
+
+    size_t size = 100;
+    boost::compute::vector<uint_> vec(size, context);
+
+    // initialize the default random engine
+    boost::compute::default_random_engine engine(queue);
+
+    std::vector<int> weights;
+    // weights.begin() == weights.end()
+    boost::compute::discrete_distribution<uint_> distribution(
+        weights.begin(), weights.end()
+    );
+
+    std::vector<double> p = distribution.probabilities();
+    BOOST_CHECK_CLOSE(p[0], double(1), 0.001);
+
+    BOOST_CHECK_EQUAL((distribution.min)(), uint_(0));
+    BOOST_CHECK_EQUAL((distribution.max)(), uint_(0));
+
+    // generate the random values and store them to 'vec'
+    distribution.generate(vec.begin(), vec.end(), engine, queue);
+
+    BOOST_CHECK_EQUAL(
+        boost::compute::count_if(
+            vec.begin(), vec.end(), _1 == 0, queue
+        ),
+        size
+    );
+}
+
+BOOST_AUTO_TEST_CASE(discrete_distribution_uchar)
+{
+    using boost::compute::uchar_;
+    using boost::compute::uint_;
+    using boost::compute::lambda::_1;
+
+    size_t size = 100;
+    boost::compute::vector<uchar_> uchar_vec(size, context);
+    boost::compute::vector<uint_> uint_vec(size, context);
+
+    // initialize the default random engine
+    boost::compute::default_random_engine engine(queue);
+
+    // initialize weights
+    std::vector<int> weights(258, 0);
+    weights[257] = 1;
+
+    // setup the discrete distribution
+    boost::compute::discrete_distribution<uchar_> distribution(
+        weights.begin(), weights.end()
+    );
+
+    BOOST_CHECK_EQUAL((distribution.min)(), uchar_(0));
+    BOOST_CHECK_EQUAL((distribution.max)(), uchar_(255));
+
+    // generate the random uchar_ values to the uchar_ vector
+    distribution.generate(uchar_vec.begin(), uchar_vec.end(), engine, queue);
+
+    BOOST_CHECK_EQUAL(
+        boost::compute::count_if(
+            uchar_vec.begin(), uchar_vec.end(), _1 == uchar_(1), queue
+        ),
+        size
+    );
+
+    // generate the random uchar_ values to the uint_ vector
+    distribution.generate(uint_vec.begin(), uint_vec.end(), engine, queue);
+
+    BOOST_CHECK_EQUAL(
+        boost::compute::count_if(
+            uint_vec.begin(), uint_vec.end(), _1 == uint_(1), queue
+        ),
+        size
+    );
+}
+
 BOOST_AUTO_TEST_SUITE_END()