InverseTransformSampling.h

Go to the documentation of this file.

// Class InverseTransformSampling
//   This class can be used for generating samples of a given distribution function class using
//   inverse of its cumulative distribution on host or device.
//
 
#ifndef IPPL_INVERSE_TRANSFORM_SAMPLING_H
#define IPPL_INVERSE_TRANSFORM_SAMPLING_H
 
#include "Random/Utility.h"
 
namespace ippl {
    namespace random {
        template <typename T, unsigned Dim, class DeviceType, class Distribution>
        class InverseTransformSampling {
        public:
            using view_type = typename ippl::detail::ViewType<Vector<T, Dim>, 1>::view_type;
            using size_type = ippl::detail::size_type;
 
            Distribution dist_m;
            size_type ntotal_m;
            Vector<T, Dim> umin_m, umax_m;

            template <class RegionLayout>
            InverseTransformSampling(Distribution& dist_r, Vector<T, Dim>& rmax_r,
                                     Vector<T, Dim>& rmin_r, RegionLayout& rlayout_r,
                                     size_type& ntotal_r)
                : dist_m(dist_r)
                , ntotal_m(ntotal_r) {
                const typename RegionLayout::host_mirror_type regions =
                    rlayout_r.gethLocalRegions();
                int rank = ippl::Comm->rank();
 
                Vector<T, Dim> locrmax, locrmin;
                for (unsigned d = 0; d < Dim; ++d) {
                    locrmax[d] = regions(rank)[d].max();
                    locrmin[d] = regions(rank)[d].min();
                }
 
                updateBounds(rmax_r, rmin_r, locrmax, locrmin);
            }

            InverseTransformSampling(Distribution& dist_r, Vector<T, Dim>& rmax_r,
                                     Vector<T, Dim>& rmin_r, Vector<T, Dim>& locrmax_r,
                                     Vector<T, Dim>& locrmin_r, size_type& ntotal_r)
                : dist_m(dist_r)
                , ntotal_m(ntotal_r) {
                updateBounds(rmax_r, rmin_r, locrmax_r, locrmin_r);
            }

            InverseTransformSampling(Distribution& dist_r, Vector<T, Dim>& rmax_r,
                                     Vector<T, Dim>& rmin_r, size_type& ntotal_r)
                : dist_m(dist_r)
                , ntotal_m(ntotal_r) {
                updateBounds(rmax_r, rmin_r);
 
                nlocal_m = ntotal_m;
            }

            void updateBounds(Vector<T, Dim>& rmax, Vector<T, Dim>& rmin, Vector<T, Dim>& locrmax,
                              Vector<T, Dim>& locrmin) {
                int rank = ippl::Comm->rank();
                Vector<T, Dim> nr_m, dr_m;
                for (unsigned d = 0; d < Dim; ++d) {
                    nr_m[d]   = dist_m.getCdf(locrmax[d], d) - dist_m.getCdf(locrmin[d], d);
                    dr_m[d]   = dist_m.getCdf(rmax[d], d) - dist_m.getCdf(rmin[d], d);
                    umin_m[d] = dist_m.getCdf(locrmin[d], d);
                    umax_m[d] = dist_m.getCdf(locrmax[d], d);
                }
                T pnr = std::accumulate(nr_m.begin(), nr_m.end(), 1.0, std::multiplies<T>());
                T pdr = std::accumulate(dr_m.begin(), dr_m.end(), 1.0, std::multiplies<T>());
 
                T factor = pnr / pdr;
                nlocal_m = (size_type)(factor * ntotal_m);
 
                size_type nglobal = 0;
                ippl::Comm->allreduce(&nlocal_m, &nglobal, 1, std::plus<size_type>());
 
                int rest = (int)(ntotal_m - nglobal);
                if (rank < rest) {
                    ++nlocal_m;
                }
            }

            void updateBounds(Vector<T, Dim>& rmax, Vector<T, Dim>& rmin) {
                Vector<T, Dim> nr_m, dr_m;
                for (unsigned d = 0; d < Dim; ++d) {
                    umin_m[d] = dist_m.getCdf(rmin[d], d);
                    umax_m[d] = dist_m.getCdf(rmax[d], d);
                }
            }

            ~InverseTransformSampling() {}

            template <class GeneratorPool>
            struct fill_random {
                using value_type = T;
                Distribution targetdist_m;
                view_type sample_m;
                GeneratorPool pool_m;
                Vector<T, Dim> minbound_m;
                Vector<T, Dim> maxbound_m;
                unsigned int dim_m;

                KOKKOS_FUNCTION
                fill_random(Distribution& dist_r, view_type& x_r, GeneratorPool& rand_pool_r,
                            Vector<T, Dim>& umin_r, Vector<T, Dim>& umax_r, unsigned int& d_r)
                    : targetdist_m(dist_r)
                    , sample_m(x_r)
                    , pool_m(rand_pool_r)
                    , minbound_m(umin_r)
                    , maxbound_m(umax_r)
                    , dim_m(d_r) {}

                KOKKOS_INLINE_FUNCTION void operator()(const size_t i) const {
                    typename GeneratorPool::generator_type rand_gen = pool_m.get_state();
 
                    value_type u = 0.0;
 
                    u = rand_gen.drand(minbound_m[dim_m], maxbound_m[dim_m]);
 
                    // first guess for Newton-Raphson
                    sample_m(i)[dim_m] = targetdist_m.getEstimate(u, dim_m);
 
                    // solve
                    ippl::random::detail::NewtonRaphson<value_type, Distribution> solver(
                        targetdist_m);
 
                    solver.solve(dim_m, sample_m(i)[dim_m], u);
 
                    pool_m.free_state(rand_gen);
                }
            };

            KOKKOS_INLINE_FUNCTION size_type getLocalSamplesNum() const { return nlocal_m; }

            KOKKOS_INLINE_FUNCTION void setLocalSamplesNum(size_type nlocal) { nlocal_m = nlocal; }

            void generate(view_type view, Kokkos::Random_XorShift64_Pool<> rand_pool64) {
                Vector<T, Dim> minbound_m = umin_m;
                Vector<T, Dim> maxbound_m = umax_m;
                Distribution targetdist_m = dist_m;
                size_type numlocal_m      = nlocal_m;
                for (unsigned d = 0; d < Dim; ++d) {
                    Kokkos::parallel_for(numlocal_m, fill_random<Kokkos::Random_XorShift64_Pool<>>(
                                                         targetdist_m, view, rand_pool64,
                                                         minbound_m, maxbound_m, d));
                    Kokkos::fence();
                }
            }

            void generate(view_type view, size_type startIndex, size_type endIndex,
                          Kokkos::Random_XorShift64_Pool<> rand_pool64) {
                Vector<T, Dim> minbound_m = umin_m;
                Vector<T, Dim> maxbound_m = umax_m;
                Distribution targetdist_m = dist_m;
                for (unsigned d = 0; d < Dim; ++d) {
                    Kokkos::parallel_for(
                        Kokkos::RangePolicy<>(startIndex, endIndex),
                        fill_random<Kokkos::Random_XorShift64_Pool<>>(
                            targetdist_m, view, rand_pool64, minbound_m, maxbound_m, d));
                    Kokkos::fence();
                }
            }
 
        private:
            size_type nlocal_m;
        };
 
    }  // namespace random
}  // namespace ippl
 
#endif