Merge pull request #168 from AMReX-FHD/reactDiff

Multinomial

src_reactDiff/AdvanceDiffusion.cpp

@@ -33,7 +33,7 @@ void AdvanceDiffusion(MultiFab& n_old,
     }
 
     if (reactDiff_diffusion_type == 3) {
-        Abort("AdvanceDiffusion() - write multinomial case");
+        MultinomialDiffusion(n_old,n_new,diff_coef_face,geom,dt,time);
         return;
     }
 

src_reactDiff/AdvanceReactionDiffusion.cpp

@@ -26,14 +26,28 @@ void AdvanceReactionDiffusion(MultiFab& n_old,
 
     MultiFab rate1(ba,dmap,nspecies,0);
 
-    if (temporal_integrator == -3) { // multinomial diffusion
-        Abort("AdvanceReactionDiffusion() - temporal_integrator=-3 not supported yet");
-    }
-
     Vector<Real> mattingly_lin_comb_coef(2);
     mattingly_lin_comb_coef[0] = 1.;
     mattingly_lin_comb_coef[1] = 0.;
 
+    if (temporal_integrator == -3) { // multinomial diffusion
+
+        // calculate rates
+        // rates could be deterministic or stochastic depending on use_Poisson_rng
+        ChemicalRates(n_old,rate1,geom,dt,n_old,mattingly_lin_comb_coef,volume_factor);
+
+        // advance multinomial diffusion
+        MultinomialDiffusion(n_old,n_new,diff_coef_face,geom,dt,time);
+
+        // add reaction contribution and external source
+        MultiFab::Saxpy(n_new,dt,rate1,0,0,nspecies,0);
+        MultiFab::Saxpy(n_new,dt,ext_src,0,0,nspecies,0);
+        n_new.FillBoundary(geom.periodicity());
+        MultiFabPhysBC(n_new, geom, 0, nspecies, SPEC_BC_COMP, time);
+        return;
+
+    }
+
     MultiFab diff_fluxdiv (ba,dmap,nspecies,0);
     MultiFab stoch_fluxdiv(ba,dmap,nspecies,0);
 

src_reactDiff/InitN.cpp

@@ -69,7 +69,7 @@ void InitN(MultiFab& n_in,
 
             Abort("integer_populations=1 with initial_variance_mass < 0. not supported yet");
 
-        } else { // Make the number of molecules in each cell Poisson distributed with desired mean
+        } else if (initial_variance_mass > 0.) { // Make the number of molecules in each cell Poisson distributed with desired mean
 
             for ( MFIter mfi(n_in,TilingIfNotGPU()); mfi.isValid(); ++mfi ) {
 

src_reactDiff/Make.package

@@ -5,6 +5,7 @@ CEXE_sources   += AdvanceTimestep.cpp
 CEXE_sources   += DiffusiveNFluxdiv.cpp
 CEXE_sources   += ImplicitDiffusion.cpp
 CEXE_sources   += InitN.cpp
+CEXE_sources   += MultinomialDiffusion.cpp
 CEXE_sources   += reactDiff_functions.cpp
 CEXE_sources   += StochasticNFluxdiv.cpp
 CEXE_sources   += WritePlotFile.cpp

src_reactDiff/MultinomialDiffusion.cpp

@@ -0,0 +1,141 @@
+#include "reactDiff_functions.H"
+
+#include "AMReX_MLMG.H"
+#include <AMReX_MLABecLaplacian.H>
+
+#include <random>
+
+void MultinomialDiffusion(MultiFab& n_old,
+                          MultiFab& n_new,
+                          const std::array< MultiFab, AMREX_SPACEDIM >& diff_coef_face,
+                          const Geometry& geom,
+                          const Real& dt,
+                          const Real& time)
+{
+    BoxArray ba = n_old.boxArray();
+    DistributionMapping dmap = n_old.DistributionMap();
+
+    MultiFab cell_update(ba, dmap, nspecies, 1);
+    cell_update.setVal(0.);
+
+    // set new state to zero everywhere, including ghost cells
+    n_new.setVal(0.);
+
+    // copy old state into new in valid region only
+    MultiFab::Copy(n_new,n_old,0,0,nspecies,0);
+
+    const GpuArray<Real,AMREX_SPACEDIM> dx = geom.CellSizeArray();
+
+    Real dv = (AMREX_SPACEDIM==2) ? dx[0]*dx[1]*cell_depth : dx[0]*dx[1]*dx[2]*cell_depth;
+
+    for (MFIter mfi(n_new); mfi.isValid(); ++mfi)
+    {
+        const Box& bx = mfi.validbox();
+
+        const Array4<Real> & n_arr = n_new.array(mfi);
+
+        const Array4<Real> & update = cell_update.array(mfi);
+
+        AMREX_D_TERM(const Array4<const Real> & diffx = diff_coef_face[0].array(mfi);,
+                     const Array4<const Real> & diffy = diff_coef_face[1].array(mfi);,
+                     const Array4<const Real> & diffz = diff_coef_face[2].array(mfi););
+
+        amrex::ParallelFor(bx, nspecies, [=] AMREX_GPU_DEVICE (int i, int j, int k, int n) noexcept
+        {
+
+            GpuArray<Real,2*AMREX_SPACEDIM> p;
+            GpuArray<Real,2*AMREX_SPACEDIM> fluxes;
+
+            p[0] = diffx(i,j,k,n)*dt/(dx[0]*dx[0]);
+            p[1] = diffx(i+1,j,k,n)*dt/(dx[0]*dx[0]);
+            p[2] = diffy(i,j,k,n)*dt/(dx[1]*dx[1]);
+            p[3] = diffy(i,j+1,k,n)*dt/(dx[1]*dx[1]);
+#if (AMREX_SPACEDIM == 3)
+            p[4] = diffz(i,j,k,n)*dt/(dx[2]*dx[2]);
+            p[5] = diffz(i,j,k+1,n)*dt/(dx[2]*dx[2]);
+#endif
+
+            int N = std::max(0., std::round(n_arr(i,j,k,n)*dv));
+
+            multinomial_rng(fluxes, N, p);
+
+            // lo-x face
+            update(i  ,j,k,n) -= fluxes[0];
+            update(i-1,j,k,n) += fluxes[0];
+
+            // hi-x face
+            update(i  ,j,k,n) -= fluxes[1];
+            update(i+1,j,k,n) += fluxes[1];
+
+            // lo-y face
+            update(i,j,  k,n) -= fluxes[2];
+            update(i,j-1,k,n) += fluxes[2];
+
+            // hi-y face
+            update(i,j  ,k,n) -= fluxes[3];
+            update(i,j+1,k,n) += fluxes[3];
+
+#if (AMREX_SPACEDIM == 3)
+            // lo-z face
+            update(i,j,k,  n) -= fluxes[4];
+            update(i,j,k-1,n) += fluxes[4];
+
+            // hi-z face
+            update(i,j,k,  n) -= fluxes[5];
+            update(i,j,k+1,n) += fluxes[5];
+#endif
+        });
+    }
+
+    for (MFIter mfi(n_new); mfi.isValid(); ++mfi)
+    {
+        const Box& bx = mfi.growntilebox(1);
+
+        const Array4<Real> & n_arr = n_new.array(mfi);
+
+        const Array4<Real> & update = cell_update.array(mfi);
+
+        amrex::ParallelFor(bx, nspecies, [=] AMREX_GPU_DEVICE (int i, int j, int k, int n) noexcept
+        {
+            n_arr(i,j,k,n) += update(i,j,k,n) / dv;
+        });
+    }
+
+    n_new.SumBoundary(geom.periodicity());
+    n_new.FillBoundary(geom.periodicity());
+    MultiFabPhysBC(n_new, geom, 0, nspecies, SPEC_BC_COMP, time);
+}
+
+AMREX_GPU_HOST_DEVICE void multinomial_rng(GpuArray<Real,2*AMREX_SPACEDIM>& samples,
+                                           const int& N,
+                                           GpuArray<Real,2*AMREX_SPACEDIM>& p)
+{
+#if (AMREX_USE_CUDA)
+    Abort("MultinomialRNG not supported for CUDA");
+#else
+
+    Real sum_p = 0;
+    for (int sample=0; sample<2*AMREX_SPACEDIM; ++sample) {
+        sum_p += p[sample];
+    }
+    if (sum_p > 1.) {
+        Print() << "sum_p = " << sum_p << std::endl;
+        Abort("multinomial_rng: probabilities must sum to 1 or less");
+    }
+
+    std::default_random_engine generator;
+    generator.seed(std::chrono::system_clock::now().time_since_epoch().count());
+
+    sum_p = 0.;
+    int sum_n = 0;
+
+    for (int sample=0; sample<2*AMREX_SPACEDIM; ++sample) {
+        std::binomial_distribution<int> distribution(N-sum_n, p[sample]/(1.-sum_p));
+        samples[sample] = distribution(generator);
+
+        sum_n += samples[sample];
+        sum_p += p[sample];
+    }
+
+#endif
+}

src_reactDiff/reactDiff_functions.H

@@ -66,21 +66,21 @@ void AdvanceTimestep(MultiFab& n_old,
 ////////////////////////
 // In ImplicitDiffusion.cpp
 ////////////////////////
-void DiffusiveNFluxdiv(MultiFab& n_in,
-                       MultiFab& diff_fluxdiv,
+void ImplicitDiffusion(MultiFab& n_old,
+                       MultiFab& n_new,
+                       const MultiFab& rhs,
                        const std::array< MultiFab, AMREX_SPACEDIM >& diff_coef_face,
                        const Geometry& geom,
+                       const Real& dt_fac,
                        const Real& time);
 
 ////////////////////////
 // In DiffusiveNFluxdiv.cpp
 ////////////////////////
-void ImplicitDiffusion(MultiFab& n_old,
-                       MultiFab& n_new,
-                       const MultiFab& rhs,
+void DiffusiveNFluxdiv(MultiFab& n_in,
+                       MultiFab& diff_fluxdiv,
                        const std::array< MultiFab, AMREX_SPACEDIM >& diff_coef_face,
                        const Geometry& geom,
-                       const Real& dt_fac,
                        const Real& time);
 
 ////////////////////////
@@ -90,6 +90,20 @@ void InitN(MultiFab& n_in,
            const Geometry& geom,
            const Real& time);
 
+////////////////////////
+// In MultinomialDiffusion.cpp
+////////////////////////
+void MultinomialDiffusion(MultiFab& n_old,
+                          MultiFab& n_new,
+                          const std::array< MultiFab, AMREX_SPACEDIM >& diff_coef_face,
+                          const Geometry& geom,
+                          const Real& dt,
+                          const Real& time);
+
+AMREX_GPU_HOST_DEVICE void multinomial_rng(GpuArray<Real,2*AMREX_SPACEDIM>& samples,
+                                           const int& N,
+                                           GpuArray<Real,2*AMREX_SPACEDIM>& p);
+
 ////////////////////////
 // In StochasticNFluxdiv.cpp
 ////////////////////////