Branch data 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
|
