Add weighted choice without replacement

Author: AntoinePrv

include/xtensor/xrandom.hpp

@@ -14,12 +14,17 @@
 #ifndef XTENSOR_RANDOM_HPP
 #define XTENSOR_RANDOM_HPP
 
+#include <algorithm>
 #include <functional>
 #include <random>
 #include <utility>
+#include <type_traits>
+
+#include <xtl/xspan.hpp>
 
 #include "xbuilder.hpp"
 #include "xgenerator.hpp"
+#include "xindex_view.hpp"
 #include "xtensor.hpp"
 #include "xtensor_config.hpp"
 #include "xview.hpp"
@@ -180,6 +185,12 @@ namespace xt
         template <class T, class E = random::default_engine_type>
         xtensor<typename T::value_type, 1> choice(const xexpression<T>& e, std::size_t n, bool replace = true,
                                                   E& engine = random::get_default_random_engine());
+
+        template <class T, class W, class E = random::default_engine_type>
+        xtensor<typename T::value_type, 1> choice(const xexpression<T>& e, std::size_t n,
+                                                  const xexpression<W>& weights,
+                                                  bool replace = true,
+                                                  E& engine = random::get_default_random_engine());
     }
 
     namespace detail
@@ -792,6 +803,83 @@ namespace xt
             }
             return result;
         }
+
+        /**
+         * Randomly select n unique elements from xexpression e using the weights distribution w.
+         *
+         * Note: this function makes a copy of your data, and only 1D data is accepted.
+         *
+         * @param e expression to sample from
+         * @param n number of elements to sample
+         * @param e expression for the weight distribution.
+         *          Weights must be positive and real-valued but need not sum to 1.
+         * @param replace whether to sample with or without replacement
+         * @param engine random number engine
+         *
+         * @return xtensor containing 1D container of sampled elements
+         */
+        template <class T, class W, class E>
+        xtensor<typename T::value_type, 1>
+        choice(const xexpression<T>& e, std::size_t n, const xexpression<W>& weights, bool replace, E& engine)
+        {
+            const auto& de = e.derived_cast();
+            const auto& dweights = weights.derived_cast();
+            if (de.dimension() != 1)
+            {
+                XTENSOR_THROW(std::runtime_error, "Sample and weight expression must be 1 dimensional");
+            }
+            if (de.size() < n && !replace)
+            {
+                XTENSOR_THROW(std::runtime_error, "If replace is false, then the sample expression's size must be > n");
+            }
+            if (de.size() != dweights.size() || de.dimension() != dweights.dimension())
+            {
+                XTENSOR_THROW(std::runtime_error, "Sample and weight expression must have the same size");
+            }
+            static_assert(std::is_floating_point<typename W::value_type>::value,
+                          "Weight expression must be of floating point type");
+            using result_type = xtensor<typename T::value_type, 1>;
+            using size_type = typename result_type::size_type;
+            using weight_type = typename W::value_type;
+            result_type result;
+            result.resize({n});
+
+            if (replace)
+            {
+                XTENSOR_THROW(std::runtime_error, "Not implemented");
+            }
+            else
+            {
+                // Algorithm from
+                // Efraimidis PS, Spirakis PG (2006). "Weighted random sampling with a reservoir."
+                // Information Processing  Letters, 97 (5), 181-185. ISSN 0020-0190.
+                // doi:10.1016/j.ipl.2005.11.003.
+                // https://www.sciencedirect.com/science/article/pii/S002001900500298X
+                //
+                // The keys computed are replaced with weight/randexp(1) instead rand()^(1/weight) as done in wrlmlR:
+                // https://web.archive.org/web/20201021162211/https://krlmlr.github.io/wrswoR/
+                // https://web.archive.org/web/20201021162520/https://github.com/krlmlr/wrswoR/blob/master/src/sample_int_crank.cpp
+                // As well as in JuliaStats:
+                // https://web.archive.org/web/20201021162949/https://github.com/JuliaStats/StatsBase.jl/blob/master/src/sampling.jl
+
+                // Compute (modified) keys as weight/randexp(1).
+                xtensor<weight_type, 1> keys;
+                keys.resize({dweights.size()});
+                std::exponential_distribution<weight_type> randexp{1.};
+                std::transform(dweights.storage().begin(), dweights.storage().end(), keys.begin(),
+                               [&randexp, &engine](auto w){ return w / randexp(engine); });
+
+                // Find indexes for the n biggest key
+                xtensor<size_type, 1> indices = arange<size_type>(0, dweights.size());
+                std::partial_sort(indices.storage().begin(), indices.storage().begin() + n, indices.storage().end(),
+                                  [&keys](auto i, auto j) { return keys[i] > keys[j]; });
+
+                // Return samples with the n biggest keys
+                result = index_view(de, xtl::span<size_type>{indices.data(), n});
+            }
+            return result;
+        }
+
     }
 }
 

test/test_xrandom.cpp

@@ -7,6 +7,8 @@
 * The full license is in the file LICENSE, distributed with this software. *
 ****************************************************************************/
 
+#include <type_traits>
+
 #include "gtest/gtest.h"
 #include "test_common_macros.hpp"
 #if (defined(__GNUC__) && !defined(__clang__))
@@ -19,6 +21,7 @@
 #endif
 #include "xtensor/xarray.hpp"
 #include "xtensor/xview.hpp"
+#include "xtensor/xset_operation.hpp"
 
 namespace xt
 {
@@ -167,6 +170,33 @@ namespace xt
         XT_ASSERT_THROW(xt::random::choice(multidim_input, 5, true), std::runtime_error);
     }
 
+    TEST(xrandom, weighted_choice)
+    {
+        xarray<int> a = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12};
+        xarray<double> w = {1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0};
+        xt::random::seed(42);
+        auto ac1 = xt::random::choice(a, 6, w, false);
+        auto ac2 = xt::random::choice(a, 6, w, false);
+        xt::random::seed(42);
+        auto ac3 = xt::random::choice(a, 6, w, false);
+
+        static_assert(std::is_same<decltype(a)::value_type, decltype(ac1)::value_type>::value, "Elements must be same type");
+        ASSERT_EQ(ac1, ac3);
+        ASSERT_NE(ac1, ac2);
+        ASSERT_TRUE(all(isin(ac1, a)));
+        ASSERT_TRUE(all(equal(ac1 % 2, 1)));
+
+        xarray<double> b = {-1, 1};
+        xarray<double> v = {1, 1};
+        xt::random::seed(42);
+        XT_ASSERT_THROW(xt::random::choice(b, 5, v, false), std::runtime_error);
+        XT_ASSERT_NO_THROW(xt::random::choice(b, 5, v, true));
+        xarray<double> multidim_input = { {1,2,3}, {3,4,5} };
+        XT_ASSERT_THROW(xt::random::choice(multidim_input, 5, v, true), std::runtime_error);
+        xarray<double> bad_count_weights = {1, 1, 4};
+        XT_ASSERT_THROW(xt::random::choice(b, 5, bad_count_weights, true), std::runtime_error);
+    }
+
     TEST(xrandom, shuffle)
     {
         xarray<double> a = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12};