Merge pull request #307 from IPPL-framework/306-updating-inverse-tran…

…sform-sampling-class

Updating Inverse Transform Sampling class

src/Random/InverseTransformSampling.h

@@ -29,53 +29,136 @@ namespace ippl {
         using size_type = ippl::detail::size_type;
 
         Distribution dist_m;
-        const Vector<T, Dim> rmax_m;
-        const Vector<T, Dim> rmin_m;
-        Vector<T, Dim> umin_m, umax_m;
         size_type ntotal_m;
+        Vector<T, Dim> umin_m, umax_m;
 
         /*!
-         * @brief Constructor for InverseTransformSampling class.
+         * @brief Constructor for InverseTransformSampling class with domain decomposition.
          *
-         * @param dist_ The distribution to sample from.
-         * @param rmax_ Maximum range for sampling.
-         * @param rmin_ Minimum range for sampling.
-         * @param rlayout The region layout.
-         * @param ntotal_ Total number of samples to generate.
+         * @param dist_r The distribution to sample from.
+         * @param rmax_r Maximum global range.
+         * @param rmin_r Minimum global range.
+         * @param rlayout_r The region layout.
+         * @param ntotal_r Total number of samples to generate.
         */
         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),
-        rmax_m(rmax_r),
-        rmin_m(rmin_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);
+
+        }
+
+        /*!
+         * @brief Constructor for InverseTransformSampling class without applying domain decomposition..
+         *
+         * @param dist_r The distribution to sample from.
+         * @param rmax_r Maximum global range.
+         * @param rmin_r Minimum global range.
+         * @param locrmax_r Maximum local (per rank) range.
+         * @param locrmin_r Minimum local (per rank) range.
+         * @param ntotal_ Total number of samples to generate.
+        */
+        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);
+
+        }
+
+        /*!
+         * @brief Constructor for InverseTransformSampling class.
+         *        In this method, we do not consider any domain decomposition.
+         *
+         * @param dist_r The distribution to sample from.
+         * @param rmax_r Maximum global range for sampling.
+         * @param rmin_r Minimum global range for sampling.
+         * @param ntotal_r Total number of samples to generate.
+        */
+        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;
+
+        }
+
+        /*!
+         * @brief Updates the sampling bounds and reinitializes internal variables.
+         *
+         * This method allows the user to update the minimum and maximum bounds 
+         * for the sampling process It recalculates
+         * the cumulative distribution function (CDF) values for the new bounds and 
+         * updates the internal variables to reflect these changes.
+         *
+         * @param rmax The new maximum range for sampling. This vector defines
+         *                 the upper bounds for each dimension.
+         * @param rmin The new minimum range for sampling. This vector defines
+         *                 the lower bounds for each dimension.
+         * @param locrmax  The new local maximum range for sampling. This vector defines
+         *                 the upper bounds for each dimension for a given rank.
+         * @param locrmin  The new minimum range for sampling. This vector defines
+         *                 the lower bounds for each dimension for a given rank.
+        */
+        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(regions(rank)[d].max(), d) - dist_m.getCdf(regions(rank)[d].min(),d);
-                dr_m[d]   = dist_m.getCdf(rmax_m[d],d) - dist_m.getCdf(rmin_m[d],d);
-                umin_m[d] = dist_m.getCdf(regions(rank)[d].min(),d);
-                umax_m[d] = dist_m.getCdf(regions(rank)[d].max(),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;
+            nlocal_m = (size_type)(factor * ntotal_m);
 
-            MPI_Allreduce(&nlocal_m, &nglobal, 1, MPI_UNSIGNED_LONG, MPI_SUM,
-                          ippl::Comm->getCommunicator());
+            size_type nglobal = 0;
+            MPI_Allreduce(&nlocal_m, &nglobal, 1, MPI_UNSIGNED_LONG, MPI_SUM, ippl::Comm->getCommunicator());
 
             int rest = (int)(ntotal_m - nglobal);
-
             if (rank < rest) {
                 ++nlocal_m;
             }
+         }
 
-        }
+	/*!
+         * @brief Updates the sampling bounds using the CDF without any domain decomposition.
+         *
+         * This method allows the user to update the minimum and maximum bounds 
+         * for the inverse transform sampling method. It recalculates
+         * the cumulative distribution function (CDF) values for the new bounds and 
+         * updates the internal variables to reflect these changes.
+         *
+         * @param new_rmax The new maximum range for sampling. This vector defines
+         *                 the upper bounds for each dimension.
+         * @param new_rmin The new minimum range for sampling. This vector defines
+         *                 the lower bounds for each dimension.
+        */
+        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);
+            }
+         }
 
         /*!
          * @brief Deconstructor for InverseTransformSampling class.
@@ -144,7 +227,16 @@ namespace ippl {
          * @returns The local number of samples.
         */
         KOKKOS_INLINE_FUNCTION size_type getLocalSamplesNum() const { return nlocal_m; }
-
+
+        /*!
+         * @brief Set the local number of particles.
+         *
+         * @param nlocal The new number of local particles.
+        */
+        KOKKOS_INLINE_FUNCTION void setLocalSamplesNum(size_type nlocal) {
+            nlocal_m = nlocal;
+        }
+
         /*!
          * @brief Generate random samples using inverse transform sampling.
          *

test/random/CMakeLists.txt

@@ -20,6 +20,13 @@ target_link_libraries (
     ${MPI_CXX_LIBRARIES}
 )
 
+add_executable (TestInverseTransformSamplingUpdateBounds TestInverseTransformSamplingUpdateBounds.cpp)
+target_link_libraries (
+    TestInverseTransformSamplingUpdateBounds
+    ${IPPL_LIBS}
+    ${MPI_CXX_LIBRARIES}
+)
+
 add_executable (TestInverseTransformSamplingCustom TestInverseTransformSamplingCustom.cpp)
 target_link_libraries (
     TestInverseTransformSamplingCustom

test/random/TestInverseTransformSamplingUpdateBounds.cpp

@@ -0,0 +1,135 @@
+// Testing the update functionality of the inverse transform sampling method for Normal Distribution
+//     Example:
+//     srun ./TestInverseTransformSamplingUpdateBounds --overallocate 2.0 --info 10
+
+#include <Kokkos_MathematicalConstants.hpp>
+#include <Kokkos_MathematicalFunctions.hpp>
+#include <Kokkos_Random.hpp>
+#include <chrono>
+#include <iostream>
+#include <random>
+#include <set>
+#include <string>
+#include <vector>
+#include "Utility/IpplTimings.h"
+#include "Ippl.h"
+#include "Random/InverseTransformSampling.h"
+#include "Random/NormalDistribution.h"
+
+const int Dim = 2;
+
+using view_type  = typename ippl::detail::ViewType<ippl::Vector<double, Dim>, 1>::view_type;
+
+using Mesh_t = ippl::UniformCartesian<double, Dim>;
+
+using size_type = ippl::detail::size_type;
+
+using GeneratorPool = typename Kokkos::Random_XorShift64_Pool<>;
+
+void get_boundaries(view_type& Rview, double *rmin, double *rmax){
+    double rmax_loc[Dim];
+    double rmin_loc[Dim];
+
+    for (unsigned d = 0; d < Dim; ++d) {
+            Kokkos::parallel_reduce(
+                "rel max", ippl::getRangePolicy(Rview),
+                KOKKOS_LAMBDA(const int i, double& mm) {
+                    double tmp_vel = Rview(i)[d];
+                    mm             = tmp_vel > mm ? tmp_vel : mm;
+                },
+                Kokkos::Max<double>(rmax_loc[d]));
+
+            Kokkos::parallel_reduce(
+                "rel min", ippl::getRangePolicy(Rview),
+                KOKKOS_LAMBDA(const int i, double& mm) {
+                    double tmp_vel = Rview(i)[d];
+                    mm             = tmp_vel < mm ? tmp_vel : mm;
+                },
+                Kokkos::Min<double>(rmin_loc[d]));
+     }
+     Kokkos::fence();
+     MPI_Allreduce(rmax_loc, rmax, Dim, MPI_DOUBLE, MPI_MAX, ippl::Comm->getCommunicator());
+     MPI_Allreduce(rmin_loc, rmin, Dim, MPI_DOUBLE, MPI_MIN, ippl::Comm->getCommunicator());
+     ippl::Comm->barrier();
+}
+
+void write_minmax(double *rmin1, double *rmax1, double *rmin2, double *rmax2){
+    Inform csvout(NULL, "data/rmin_max_normal_dist.csv", Inform::APPEND);
+    csvout.precision(10);
+    csvout.setf(std::ios::scientific, std::ios::floatfield);
+    csvout << "dim    min(pos)    max(pos)    min(vel)    max(vel)" << endl;
+    for(int i=0; i<Dim; i++){
+        csvout << i << " " << rmin1[i] << " " << rmax1[i] << " " << rmin2[i] << " " << rmax2[i] << endl;
+    }
+    ippl::Comm->barrier();
+}
+
+int main(int argc, char* argv[]) {
+    ippl::initialize(argc, argv);
+    {
+        Inform m("test ITS normal");
+        // set up ITS as other examples to sample position
+        ippl::Vector<int, 2> nr   = {100, 100};
+        size_type ntotal = 100000;
+
+        ippl::NDIndex<2> domain;
+        for (unsigned i = 0; i < Dim; i++) {
+            domain[i] = ippl::Index(nr[i]);
+        }
+
+        std::array<bool, Dim> isParallel;
+        isParallel.fill(true);
+
+        ippl::Vector<double, Dim> rmin   = -2.;
+        ippl::Vector<double, Dim> rmax   = 2.;
+        ippl::Vector<double, Dim> length = rmax - rmin;
+        ippl::Vector<double, Dim> hr     = length / nr;
+        ippl::Vector<double, Dim> origin = rmin;
+
+        const bool isAllPeriodic = true;
+
+        Mesh_t mesh(domain, hr, origin);
+
+        ippl::FieldLayout<Dim> fl(MPI_COMM_WORLD, domain, isParallel, isAllPeriodic);
+
+        ippl::detail::RegionLayout<double, Dim, Mesh_t> rlayout(fl, mesh);
+
+        int seed = 42;
+
+        GeneratorPool rand_pool64((size_type)(seed + 100 * ippl::Comm->rank()));
+
+        const double mu1 = 0.1;
+        const double sd1 = 0.5;
+        const double mu2 = -0.1;
+        const double sd2 = 1.0;
+        const double par[4] = {mu1, sd1, mu2, sd2};
+        using Dist_t = ippl::random::NormalDistribution<double, Dim>;
+        using sampling_t = ippl::random::InverseTransformSampling<double, Dim, Kokkos::DefaultExecutionSpace, Dist_t>;
+
+        Dist_t dist(par);
+        sampling_t sampling(dist, rmax, rmin, rlayout, ntotal);
+        size_type nlocal = sampling.getLocalSamplesNum();
+        view_type position("position", nlocal);
+        sampling.generate(position, rand_pool64);
+
+        // now, we want to sample velocity with the same density, but different bounds
+        // update bounds, and related parameters
+        ippl::Vector<double, Dim> vmin   = 0.;
+        ippl::Vector<double, Dim> vmax   = 2.;
+        sampling.updateBounds(vmax, vmin);
+        // set nlocal
+        sampling.setLocalSamplesNum(nlocal);
+        // create samples of updated ITS
+        view_type velocity("velocity", nlocal);
+        sampling.generate(velocity, rand_pool64);
+
+        double minPosition[3], maxPosition[3];
+        double minVelocity[3], maxVelocity[3];
+        get_boundaries(position, minPosition, maxPosition);
+        get_boundaries(velocity, minVelocity, maxVelocity);
+        write_minmax(minPosition, maxPosition, minVelocity, maxVelocity);
+    }
+    ippl::finalize();
+    return 0;
+}
+