Line data Source code
1 : // Class InverseTransformSampling
2 : // This class can be used for generating samples of a given distribution function class using
3 : // inverse of its cumulative distribution on host or device.
4 : //
5 :
6 : #ifndef IPPL_INVERSE_TRANSFORM_SAMPLING_H
7 : #define IPPL_INVERSE_TRANSFORM_SAMPLING_H
8 :
9 : #include "Random/Utility.h"
10 :
11 : namespace ippl {
12 : namespace random {
13 : /*!
14 : * @file InverseTransformSampling.h
15 : * @class InverseTransformSampling
16 : * @brief A class for inverse transform sampling.
17 : *
18 : * This class performs inverse transform sampling for a given distribution.
19 : *
20 : * @tparam T Data type.
21 : * @tparam Dim Dimensionality of the sample space.
22 : * @tparam DeviceType The device type for Kokkos.
23 : * @tparam Distribution Type of the distribution to sample from.
24 : */
25 : template <typename T, unsigned Dim, class DeviceType, class Distribution>
26 : class InverseTransformSampling {
27 : public:
28 : using view_type = typename ippl::detail::ViewType<Vector<T, Dim>, 1>::view_type;
29 : using size_type = ippl::detail::size_type;
30 :
31 : Distribution dist_m;
32 : size_type ntotal_m;
33 : Vector<T, Dim> umin_m, umax_m;
34 :
35 : /*!
36 : * @brief Constructor for InverseTransformSampling class with domain decomposition.
37 : *
38 : * @param dist_r The distribution to sample from.
39 : * @param rmax_r Maximum global range.
40 : * @param rmin_r Minimum global range.
41 : * @param rlayout_r The region layout.
42 : * @param ntotal_r Total number of samples to generate.
43 : */
44 : template <class RegionLayout>
45 0 : InverseTransformSampling(Distribution& dist_r, Vector<T, Dim>& rmax_r,
46 : Vector<T, Dim>& rmin_r, RegionLayout& rlayout_r,
47 : size_type& ntotal_r)
48 0 : : dist_m(dist_r)
49 0 : , ntotal_m(ntotal_r) {
50 0 : const typename RegionLayout::host_mirror_type regions =
51 : rlayout_r.gethLocalRegions();
52 0 : int rank = ippl::Comm->rank();
53 :
54 0 : Vector<T, Dim> locrmax, locrmin;
55 0 : for (unsigned d = 0; d < Dim; ++d) {
56 0 : locrmax[d] = regions(rank)[d].max();
57 0 : locrmin[d] = regions(rank)[d].min();
58 : }
59 :
60 0 : updateBounds(rmax_r, rmin_r, locrmax, locrmin);
61 0 : }
62 :
63 : /*!
64 : * @brief Constructor for InverseTransformSampling class without applying domain
65 : * decomposition..
66 : *
67 : * @param dist_r The distribution to sample from.
68 : * @param rmax_r Maximum global range.
69 : * @param rmin_r Minimum global range.
70 : * @param locrmax_r Maximum local (per rank) range.
71 : * @param locrmin_r Minimum local (per rank) range.
72 : * @param ntotal_ Total number of samples to generate.
73 : */
74 0 : InverseTransformSampling(Distribution& dist_r, Vector<T, Dim>& rmax_r,
75 : Vector<T, Dim>& rmin_r, Vector<T, Dim>& locrmax_r,
76 : Vector<T, Dim>& locrmin_r, size_type& ntotal_r)
77 0 : : dist_m(dist_r)
78 0 : , ntotal_m(ntotal_r) {
79 0 : updateBounds(rmax_r, rmin_r, locrmax_r, locrmin_r);
80 0 : }
81 :
82 : /*!
83 : * @brief Constructor for InverseTransformSampling class.
84 : * In this method, we do not consider any domain decomposition.
85 : *
86 : * @param dist_r The distribution to sample from.
87 : * @param rmax_r Maximum global range for sampling.
88 : * @param rmin_r Minimum global range for sampling.
89 : * @param ntotal_r Total number of samples to generate.
90 : */
91 : InverseTransformSampling(Distribution& dist_r, Vector<T, Dim>& rmax_r,
92 : Vector<T, Dim>& rmin_r, size_type& ntotal_r)
93 : : dist_m(dist_r)
94 : , ntotal_m(ntotal_r) {
95 : updateBounds(rmax_r, rmin_r);
96 :
97 : nlocal_m = ntotal_m;
98 : }
99 :
100 : /*!
101 : * @brief Updates the sampling bounds and reinitializes internal variables.
102 : *
103 : * This method allows the user to update the minimum and maximum bounds
104 : * for the sampling process It recalculates
105 : * the cumulative distribution function (CDF) values for the new bounds and
106 : * updates the internal variables to reflect these changes.
107 : *
108 : * @param rmax The new maximum range for sampling. This vector defines
109 : * the upper bounds for each dimension.
110 : * @param rmin The new minimum range for sampling. This vector defines
111 : * the lower bounds for each dimension.
112 : * @param locrmax The new local maximum range for sampling. This vector defines
113 : * the upper bounds for each dimension for a given rank.
114 : * @param locrmin The new minimum range for sampling. This vector defines
115 : * the lower bounds for each dimension for a given rank.
116 : */
117 0 : void updateBounds(Vector<T, Dim>& rmax, Vector<T, Dim>& rmin, Vector<T, Dim>& locrmax,
118 : Vector<T, Dim>& locrmin) {
119 0 : int rank = ippl::Comm->rank();
120 0 : Vector<T, Dim> nr_m, dr_m;
121 0 : for (unsigned d = 0; d < Dim; ++d) {
122 0 : nr_m[d] = dist_m.getCdf(locrmax[d], d) - dist_m.getCdf(locrmin[d], d);
123 0 : dr_m[d] = dist_m.getCdf(rmax[d], d) - dist_m.getCdf(rmin[d], d);
124 0 : umin_m[d] = dist_m.getCdf(locrmin[d], d);
125 0 : umax_m[d] = dist_m.getCdf(locrmax[d], d);
126 : }
127 0 : T pnr = std::accumulate(nr_m.begin(), nr_m.end(), 1.0, std::multiplies<T>());
128 0 : T pdr = std::accumulate(dr_m.begin(), dr_m.end(), 1.0, std::multiplies<T>());
129 :
130 0 : T factor = pnr / pdr;
131 0 : nlocal_m = (size_type)(factor * ntotal_m);
132 :
133 0 : size_type nglobal = 0;
134 0 : ippl::Comm->allreduce(&nlocal_m, &nglobal, 1, std::plus<size_type>());
135 :
136 0 : int rest = (int)(ntotal_m - nglobal);
137 0 : if (rank < rest) {
138 0 : ++nlocal_m;
139 : }
140 0 : }
141 :
142 : /*!
143 : * @brief Updates the sampling bounds using the CDF without any domain decomposition.
144 : *
145 : * This method allows the user to update the minimum and maximum bounds
146 : * for the inverse transform sampling method. It recalculates
147 : * the cumulative distribution function (CDF) values for the new bounds and
148 : * updates the internal variables to reflect these changes.
149 : *
150 : * @param new_rmax The new maximum range for sampling. This vector defines
151 : * the upper bounds for each dimension.
152 : * @param new_rmin The new minimum range for sampling. This vector defines
153 : * the lower bounds for each dimension.
154 : */
155 0 : void updateBounds(Vector<T, Dim>& rmax, Vector<T, Dim>& rmin) {
156 0 : Vector<T, Dim> nr_m, dr_m;
157 0 : for (unsigned d = 0; d < Dim; ++d) {
158 0 : umin_m[d] = dist_m.getCdf(rmin[d], d);
159 0 : umax_m[d] = dist_m.getCdf(rmax[d], d);
160 : }
161 0 : }
162 :
163 : /*!
164 : * @brief Deconstructor for InverseTransformSampling class.
165 : */
166 0 : ~InverseTransformSampling() {}
167 :
168 : /*!
169 : * @brief Functor that is used for generating samples.
170 : */
171 : template <class GeneratorPool>
172 : struct fill_random {
173 : using value_type = T;
174 : Distribution targetdist_m;
175 : view_type sample_m;
176 : GeneratorPool pool_m;
177 : Vector<T, Dim> minbound_m;
178 : Vector<T, Dim> maxbound_m;
179 : unsigned int dim_m;
180 :
181 : /*!
182 : * @brief Constructor for the fill_random functor.
183 : *
184 : * @param dist_ The distribution to sample from.
185 : * @param x_ The view to generate samples in.
186 : * @param rand_pool_ The random number generator pool.
187 : * @param umin_ Minimum cumulative distribution values.
188 : * @param umax_ Maximum cumulative distribution values.
189 : */
190 : KOKKOS_FUNCTION
191 0 : fill_random(Distribution& dist_r, view_type& x_r, GeneratorPool& rand_pool_r,
192 : Vector<T, Dim>& umin_r, Vector<T, Dim>& umax_r, unsigned int& d_r)
193 0 : : targetdist_m(dist_r)
194 0 : , sample_m(x_r)
195 0 : , pool_m(rand_pool_r)
196 0 : , minbound_m(umin_r)
197 0 : , maxbound_m(umax_r)
198 0 : , dim_m(d_r) {}
199 :
200 : /*!
201 : * @brief Operator to fill random values.
202 : *
203 : * @param i Index for the random values.
204 : */
205 0 : KOKKOS_INLINE_FUNCTION void operator()(const size_t i) const {
206 0 : typename GeneratorPool::generator_type rand_gen = pool_m.get_state();
207 :
208 0 : value_type u = 0.0;
209 :
210 0 : u = rand_gen.drand(minbound_m[dim_m], maxbound_m[dim_m]);
211 :
212 : // first guess for Newton-Raphson
213 0 : sample_m(i)[dim_m] = targetdist_m.getEstimate(u, dim_m);
214 :
215 : // solve
216 0 : ippl::random::detail::NewtonRaphson<value_type, Distribution> solver(
217 0 : targetdist_m);
218 :
219 0 : solver.solve(dim_m, sample_m(i)[dim_m], u);
220 :
221 0 : pool_m.free_state(rand_gen);
222 0 : }
223 : };
224 :
225 : /*!
226 : * @brief Get the local number of samples.
227 : *
228 : * @returns The local number of samples.
229 : */
230 0 : KOKKOS_INLINE_FUNCTION size_type getLocalSamplesNum() const { return nlocal_m; }
231 :
232 : /*!
233 : * @brief Set the local number of particles.
234 : *
235 : * @param nlocal The new number of local particles.
236 : */
237 0 : KOKKOS_INLINE_FUNCTION void setLocalSamplesNum(size_type nlocal) { nlocal_m = nlocal; }
238 :
239 : /*!
240 : * @brief Generate random samples using inverse transform sampling.
241 : *
242 : * @param view The view to fill with random samples.
243 : * @param rand_pool64 The random number generator pool.
244 : */
245 0 : void generate(view_type view, Kokkos::Random_XorShift64_Pool<> rand_pool64) {
246 0 : Vector<T, Dim> minbound_m = umin_m;
247 0 : Vector<T, Dim> maxbound_m = umax_m;
248 0 : Distribution targetdist_m = dist_m;
249 0 : size_type numlocal_m = nlocal_m;
250 0 : for (unsigned d = 0; d < Dim; ++d) {
251 0 : Kokkos::parallel_for(numlocal_m, fill_random<Kokkos::Random_XorShift64_Pool<>>(
252 : targetdist_m, view, rand_pool64,
253 : minbound_m, maxbound_m, d));
254 0 : Kokkos::fence();
255 : }
256 0 : }
257 :
258 : /*!
259 : * @brief Generate random samples using inverse transform sampling for a specific range
260 : * of particles
261 : *
262 : * @param view The view to fill with random samples.
263 : * @param startIndex The starting index of view.
264 : * @param endIndex The ending index of view.
265 : * @param rand_pool64 The random number generator pool.
266 : */
267 0 : void generate(view_type view, size_type startIndex, size_type endIndex,
268 : Kokkos::Random_XorShift64_Pool<> rand_pool64) {
269 0 : Vector<T, Dim> minbound_m = umin_m;
270 0 : Vector<T, Dim> maxbound_m = umax_m;
271 0 : Distribution targetdist_m = dist_m;
272 0 : for (unsigned d = 0; d < Dim; ++d) {
273 0 : Kokkos::parallel_for(
274 0 : Kokkos::RangePolicy<>(startIndex, endIndex),
275 0 : fill_random<Kokkos::Random_XorShift64_Pool<>>(
276 : targetdist_m, view, rand_pool64, minbound_m, maxbound_m, d));
277 0 : Kokkos::fence();
278 : }
279 0 : }
280 :
281 : private:
282 : size_type nlocal_m;
283 : };
284 :
285 : } // namespace random
286 : } // namespace ippl
287 :
288 : #endif
|
