Add gaussian_parse_momentum_function

dpgrote · dpgrote · commit e62eaa8ce644 · 2023-10-31T17:49:07.000-07:00
diff --git a/Docs/source/usage/parameters.rst b/Docs/source/usage/parameters.rst
@@ -908,6 +908,17 @@ Particle initialization
       ``<species_name>.momentum_function_uy(x,y,z)`` and ``<species_name>.momentum_function_uz(x,y,z)``,
       which gives the distribution of each component of the momentum as a function of space.
 
+    * ``gaussian_parse_momentum_function``: Gaussian momentum distribution where the mean momentum :math:`u = (u_{x},u_{y},u_{z})=(\gamma v_{x}/c,\gamma v_{y}/c,\gamma v_{z}/c)` and the standard deviation are given by functions of position in the input file.
+      It requires the following arguments which specify the distribution of each component of the mean momentum
+      and standard deviation as a function of space.
+
+      * ``<species_name>.momentum_function_ux_m(x,y,z)``: mean :math:`u_{x}`
+      * ``<species_name>.momentum_function_uy_m(x,y,z)``: mean :math:`u_{y}`
+      * ``<species_name>.momentum_function_uz_m(x,y,z)``: mean :math:`u_{z}`
+      * ``<species_name>.momentum_function_ux_th(x,y,z)``: standard deviation of :math:`u_{x}`
+      * ``<species_name>.momentum_function_uy_th(x,y,z)``: standard deviation of :math:`u_{y}`
+      * ``<species_name>.momentum_function_uz_th(x,y,z)``: standard deviation of :math:`u_{z}`
+
 * ``<species_name>.theta_distribution_type`` (`string`) optional (default ``constant``)
     Only read if ``<species_name>.momentum_distribution_type`` is ``maxwell_boltzmann`` or ``maxwell_juttner``.
     See documentation for these distributions (above) for constraints on values of theta. Temperatures less than zero are not allowed.
diff --git a/Examples/Tests/initial_distribution/analysis_distribution.py b/Examples/Tests/initial_distribution/analysis_distribution.py
@@ -261,6 +261,34 @@
 assert(f7_error < tolerance)
 
 
+#============================================
+# Gaussian with parser mean and standard deviation
+#============================================
+
+# load data
+bin_value_ux, bin_data_ux = read_reduced_diags_histogram("h9x.txt")[2:]
+bin_value_uy, bin_data_uy = read_reduced_diags_histogram("h9y.txt")[2:]
+bin_value_uz, bin_data_uz = read_reduced_diags_histogram("h9z.txt")[2:]
+
+def Gaussian(mean, sigma, u):
+    V = 8.0 # volume in m^3
+    n = 1.0e21 # number density in 1/m^3
+    return (n*V/(sigma*np.sqrt(2.*np.pi)))*np.exp(-(u - mean)**2/(2.*sigma**2))
+
+du = 2./50
+f_ux = Gaussian(0.1 , 0.2 , bin_value_ux)*du
+f_uy = Gaussian(0.12, 0.21, bin_value_uy)*du
+f_uz = Gaussian(0.14, 0.22, bin_value_uz)*du
+
+f9_error = np.sum(np.abs(f_ux - bin_data_ux)/f_ux.max()
+                 +np.abs(f_uy - bin_data_uy)/f_ux.max()
+                 +np.abs(f_uz - bin_data_uz)/f_uz.max()) / bin_value_ux.size
+
+print('gaussian_parse_momentum_function velocity difference:', f9_error)
+
+assert(f9_error < tolerance)
+
+
 #============================================
 # Cuboid distribution in momentum space
 #============================================
diff --git a/Examples/Tests/initial_distribution/inputs b/Examples/Tests/initial_distribution/inputs
@@ -29,7 +29,7 @@ algo.particle_shape = 1
 #################################
 ############ PLASMA #############
 #################################
-particles.species_names = gaussian maxwell_boltzmann maxwell_juttner beam maxwell_juttner_parser velocity_constant velocity_parser uniform
+particles.species_names = gaussian maxwell_boltzmann maxwell_juttner beam maxwell_juttner_parser velocity_constant velocity_parser uniform gaussian_parser
 particles.rigid_injected_species = beam
 
 gaussian.charge                     = -q_e
@@ -119,6 +119,20 @@ velocity_parser.beta_distribution_type     = "parser"
 velocity_parser.beta_function(x,y,z)       = "-0.2 + 0.4 * heaviside(z,0)"
 velocity_parser.bulk_vel_dir               = -y
 
+gaussian_parser.charge                     = -q_e
+gaussian_parser.mass                       = m_e
+gaussian_parser.injection_style            = "NRandomPerCell"
+gaussian_parser.num_particles_per_cell     = 1000
+gaussian_parser.profile                    = constant
+gaussian_parser.density                    = 1.0e21
+gaussian_parser.momentum_distribution_type = "gaussian_parse_momentum_function"
+gaussian_parser.momentum_function_ux_m(x,y,z) = 0.1*z
+gaussian_parser.momentum_function_uy_m(x,y,z) = 0.12*z
+gaussian_parser.momentum_function_uz_m(x,y,z) = 0.14*z
+gaussian_parser.momentum_function_ux_th(x,y,z) = 0.2*z
+gaussian_parser.momentum_function_uy_th(x,y,z) = 0.21*z
+gaussian_parser.momentum_function_uz_th(x,y,z) = 0.22*z
+
 uniform.charge                     = q_e
 uniform.mass                       = m_e
 uniform.injection_style            = "NRandomPerCell"
@@ -144,7 +158,7 @@ uniform.uz_max                     = 11.2
 # 6 for maxwell-boltzmann with constant velocity
 # 7 for maxwell-boltzmann with parser velocity
 # 8 for cuboid in momentum space
-warpx.reduced_diags_names              = h1x h1y h1z h2x h2y h2z h3 h3_filtered h4x h4y h4z bmmntr h5_neg h5_pos h6 h6uy h7 h7uy_pos h7uy_neg h8x h8y h8z
+warpx.reduced_diags_names              = h1x h1y h1z h2x h2y h2z h3 h3_filtered h4x h4y h4z bmmntr h5_neg h5_pos h6 h6uy h7 h7uy_pos h7uy_neg h8x h8y h8z h9x h9y h9z
 
 h1x.type                                 = ParticleHistogram
 h1x.intervals                            = 1
@@ -350,6 +364,33 @@ h8z.bin_min                              = 0
 h8z.bin_max                              = 15
 h8z.histogram_function(t,x,y,z,ux,uy,uz) = "uz"
 
+h9x.type                                 = ParticleHistogram
+h9x.intervals                            = 1
+h9x.path                                 = "./"
+h9x.species                              = gaussian_parser
+h9x.bin_number                           = 50
+h9x.bin_min                              = -1
+h9x.bin_max                              = 1
+h9x.histogram_function(t,x,y,z,ux,uy,uz) = "ux/z"
+
+h9y.type                                 = ParticleHistogram
+h9y.intervals                            = 1
+h9y.path                                 = "./"
+h9y.species                              = gaussian_parser
+h9y.bin_number                           = 50
+h9y.bin_min                              = -1
+h9y.bin_max                              = 1
+h9y.histogram_function(t,x,y,z,ux,uy,uz) = "uy/z"
+
+h9z.type                                 = ParticleHistogram
+h9z.intervals                            = 1
+h9z.path                                 = "./"
+h9z.species                              = gaussian_parser
+h9z.bin_number                           = 50
+h9z.bin_min                              = -1
+h9z.bin_max                              = 1
+h9z.histogram_function(t,x,y,z,ux,uy,uz) = "uz/z"
+
 # our little beam monitor
 bmmntr.type                              = BeamRelevant
 bmmntr.intervals                         = 1
diff --git a/Source/Fluids/WarpXFluidContainer.H b/Source/Fluids/WarpXFluidContainer.H
@@ -174,6 +174,9 @@ protected:
     std::unique_ptr<amrex::Parser> ux_parser;
     std::unique_ptr<amrex::Parser> uy_parser;
     std::unique_ptr<amrex::Parser> uz_parser;
+    std::unique_ptr<amrex::Parser> ux_th_parser;
+    std::unique_ptr<amrex::Parser> uy_th_parser;
+    std::unique_ptr<amrex::Parser> uz_th_parser;
 
     // Keep a pointer to TemperatureProperties to ensure the lifetime of the
     // contained Parser
diff --git a/Source/Fluids/WarpXFluidContainer.cpp b/Source/Fluids/WarpXFluidContainer.cpp
@@ -29,7 +29,7 @@ WarpXFluidContainer::WarpXFluidContainer(int nlevs_max, int ispecies, const std:
     const ParmParse pp_species_name(species_name);
     SpeciesUtils::parseDensity(species_name, h_inj_rho, density_parser);
     SpeciesUtils::parseMomentum(species_name, "none", h_inj_mom,
-        ux_parser, uy_parser, uz_parser, h_mom_temp, h_mom_vel);
+        ux_parser, uy_parser, uz_parser, ux_th_parser, uy_th_parser, uz_th_parser, h_mom_temp, h_mom_vel);
     if (h_inj_rho) {
 #ifdef AMREX_USE_GPU
         d_inj_rho = static_cast<InjectorDensity*>
diff --git a/Source/Initialization/InjectorMomentum.H b/Source/Initialization/InjectorMomentum.H
@@ -504,6 +504,45 @@ struct InjectorMomentumParser
     amrex::ParserExecutor<3> m_ux_parser, m_uy_parser, m_uz_parser;
 };
 
+// struct whose getMomentumm returns local momentum and thermal spread computed from parser.
+struct InjectorMomentumGaussianParser
+{
+    InjectorMomentumGaussianParser (amrex::ParserExecutor<3> const& a_ux_m_parser,
+                                    amrex::ParserExecutor<3> const& a_uy_m_parser,
+                                    amrex::ParserExecutor<3> const& a_uz_m_parser,
+                                    amrex::ParserExecutor<3> const& a_ux_th_parser,
+                                    amrex::ParserExecutor<3> const& a_uy_th_parser,
+                                    amrex::ParserExecutor<3> const& a_uz_th_parser) noexcept
+        : m_ux_m_parser(a_ux_m_parser), m_uy_m_parser(a_uy_m_parser), m_uz_m_parser(a_uz_m_parser),
+          m_ux_th_parser(a_ux_th_parser), m_uy_th_parser(a_uy_th_parser), m_uz_th_parser(a_uz_th_parser) {}
+
+    AMREX_GPU_HOST_DEVICE
+    amrex::XDim3
+    getMomentum (amrex::Real x, amrex::Real y, amrex::Real z,
+                 amrex::RandomEngine const& engine) const noexcept
+    {
+        amrex::Real const ux_m = m_ux_m_parser(x,y,z);
+        amrex::Real const uy_m = m_uy_m_parser(x,y,z);
+        amrex::Real const uz_m = m_uz_m_parser(x,y,z);
+        amrex::Real const ux_th = m_ux_th_parser(x,y,z);
+        amrex::Real const uy_th = m_uy_th_parser(x,y,z);
+        amrex::Real const uz_th = m_uz_th_parser(x,y,z);
+        return amrex::XDim3{amrex::RandomNormal(ux_m, ux_th, engine),
+                            amrex::RandomNormal(uy_m, uy_th, engine),
+                            amrex::RandomNormal(uz_m, uz_th, engine)};
+    }
+
+    AMREX_GPU_HOST_DEVICE
+    amrex::XDim3
+    getBulkMomentum (amrex::Real x, amrex::Real y, amrex::Real z) const noexcept
+    {
+        return amrex::XDim3{m_ux_m_parser(x,y,z), m_uy_m_parser(x,y,z), m_uz_m_parser(x,y,z)};
+    }
+
+    amrex::ParserExecutor<3> m_ux_m_parser, m_uy_m_parser, m_uz_m_parser;
+    amrex::ParserExecutor<3> m_ux_th_parser, m_uy_th_parser, m_uz_th_parser;
+};
+
 // Base struct for momentum injector.
 // InjectorMomentum contains a union (called Object) that holds any one
 // instance of:
@@ -512,6 +551,7 @@ struct InjectorMomentumParser
 // - InjectorMomentumGaussianFlux   : to generate v*gaussian distribution;
 // - InjectorMomentumRadialExpansion: to generate radial expansion;
 // - InjectorMomentumParser         : to generate momentum from parser;
+// - InjectorMomentumGaussianParser : to generate momentum and thermal spread from parser;
 // The choice is made at runtime, depending in the constructor called.
 // This mimics virtual functions.
 struct InjectorMomentum
@@ -532,6 +572,19 @@ struct InjectorMomentum
           object(t, a_ux_parser, a_uy_parser, a_uz_parser)
     { }
 
+    // This constructor stores a InjectorMomentumGaussianParser in union object.
+    InjectorMomentum (InjectorMomentumGaussianParser* t,
+                      amrex::ParserExecutor<3> const& a_ux_m_parser,
+                      amrex::ParserExecutor<3> const& a_uy_m_parser,
+                      amrex::ParserExecutor<3> const& a_uz_m_parser,
+                      amrex::ParserExecutor<3> const& a_ux_th_parser,
+                      amrex::ParserExecutor<3> const& a_uy_th_parser,
+                      amrex::ParserExecutor<3> const& a_uz_th_parser)
+        : type(Type::gaussianparser),
+          object(t, a_ux_m_parser, a_uy_m_parser, a_uz_m_parser,
+                    a_ux_th_parser, a_uy_th_parser, a_uz_th_parser)
+    { }
+
     // This constructor stores a InjectorMomentumGaussian in union object.
     InjectorMomentum (InjectorMomentumGaussian* t,
                       amrex::Real a_ux_m, amrex::Real a_uy_m, amrex::Real a_uz_m,
@@ -606,6 +659,10 @@ struct InjectorMomentum
         {
             return object.gaussian.getMomentum(x,y,z,engine);
         }
+        case Type::gaussianparser:
+        {
+            return object.gaussianparser.getMomentum(x,y,z,engine);
+        }
         case Type::gaussianflux:
         {
             return object.gaussianflux.getMomentum(x,y,z,engine);
@@ -654,6 +711,10 @@ struct InjectorMomentum
         {
             return object.gaussian.getBulkMomentum(x,y,z);
         }
+        case Type::gaussianparser:
+        {
+            return object.gaussianparser.getBulkMomentum(x,y,z);
+        }
         case Type::gaussianflux:
         {
             return object.gaussianflux.getBulkMomentum(x,y,z);
@@ -686,7 +747,7 @@ struct InjectorMomentum
         }
     }
 
-    enum struct Type { constant, gaussian, gaussianflux, uniform, boltzmann, juttner, radial_expansion, parser };
+    enum struct Type { constant, gaussian, gaussianflux, uniform, boltzmann, juttner, radial_expansion, parser, gaussianparser };
     Type type;
 
 private:
@@ -728,14 +789,24 @@ private:
                 amrex::ParserExecutor<3> const& a_uy_parser,
                 amrex::ParserExecutor<3> const& a_uz_parser) noexcept
             : parser(a_ux_parser, a_uy_parser, a_uz_parser) {}
+        Object (InjectorMomentumGaussianParser*,
+                amrex::ParserExecutor<3> const& a_ux_m_parser,
+                amrex::ParserExecutor<3> const& a_uy_m_parser,
+                amrex::ParserExecutor<3> const& a_uz_m_parser,
+                amrex::ParserExecutor<3> const& a_ux_th_parser,
+                amrex::ParserExecutor<3> const& a_uy_th_parser,
+                amrex::ParserExecutor<3> const& a_uz_th_parser) noexcept
+            : gaussianparser(a_ux_m_parser, a_uy_m_parser, a_uz_m_parser,
+                             a_ux_th_parser, a_uy_th_parser, a_uz_th_parser) {}
         InjectorMomentumConstant constant;
         InjectorMomentumGaussian gaussian;
         InjectorMomentumGaussianFlux gaussianflux;
         InjectorMomentumUniform uniform;
         InjectorMomentumBoltzmann boltzmann;
         InjectorMomentumJuttner juttner;
         InjectorMomentumRadialExpansion radial_expansion;
-        InjectorMomentumParser   parser;
+        InjectorMomentumParser parser;
+        InjectorMomentumGaussianParser gaussianparser;
     };
     Object object;
 };
diff --git a/Source/Initialization/InjectorMomentum.cpp b/Source/Initialization/InjectorMomentum.cpp
@@ -15,6 +15,7 @@ void InjectorMomentum::clear ()
     {
     case Type::parser:
     case Type::gaussian:
+    case Type::gaussianparser:
     case Type::gaussianflux:
     case Type::uniform:
     case Type::boltzmann:
diff --git a/Source/Initialization/PlasmaInjector.H b/Source/Initialization/PlasmaInjector.H
@@ -173,6 +173,9 @@ protected:
     std::unique_ptr<amrex::Parser> ux_parser;
     std::unique_ptr<amrex::Parser> uy_parser;
     std::unique_ptr<amrex::Parser> uz_parser;
+    std::unique_ptr<amrex::Parser> ux_th_parser;
+    std::unique_ptr<amrex::Parser> uy_th_parser;
+    std::unique_ptr<amrex::Parser> uz_th_parser;
 
     // Keep a pointer to TemperatureProperties to ensure the lifetime of the
     // contained Parser
diff --git a/Source/Initialization/PlasmaInjector.cpp b/Source/Initialization/PlasmaInjector.cpp
@@ -248,7 +248,9 @@ void PlasmaInjector::setupGaussianBeam (const amrex::ParmParse& pp_species_name)
         "Error: Symmetrization only supported to orders 4 or 8 ");
     gaussian_beam = true;
     SpeciesUtils::parseMomentum(species_name, "gaussian_beam", h_inj_mom,
-        ux_parser, uy_parser, uz_parser, h_mom_temp, h_mom_vel);
+                                ux_parser, uy_parser, uz_parser,
+                                ux_th_parser, uy_th_parser, uz_th_parser,
+                                h_mom_temp, h_mom_vel);
 #if defined(WARPX_DIM_XZ)
     WARPX_ALWAYS_ASSERT_WITH_MESSAGE( y_rms > 0._rt,
         "Error: Gaussian beam y_rms must be strictly greater than 0 in 2D "
@@ -289,7 +291,9 @@ void PlasmaInjector::setupNRandomPerCell (const amrex::ParmParse& pp_species_nam
 #endif
     SpeciesUtils::parseDensity(species_name, h_inj_rho, density_parser);
     SpeciesUtils::parseMomentum(species_name, "nrandompercell", h_inj_mom,
-        ux_parser, uy_parser, uz_parser, h_mom_temp, h_mom_vel);
+                                ux_parser, uy_parser, uz_parser,
+                                ux_th_parser, uy_th_parser, uz_th_parser,
+                                h_mom_temp, h_mom_vel);
 }
 
 void PlasmaInjector::setupNFluxPerCell (const amrex::ParmParse& pp_species_name)
@@ -367,7 +371,10 @@ void PlasmaInjector::setupNFluxPerCell (const amrex::ParmParse& pp_species_name)
 
     parseFlux(pp_species_name);
     SpeciesUtils::parseMomentum(species_name, "nfluxpercell", h_inj_mom,
-        ux_parser, uy_parser, uz_parser, h_mom_temp, h_mom_vel, flux_normal_axis, flux_direction);
+                                ux_parser, uy_parser, uz_parser,
+                                ux_th_parser, uy_th_parser, uz_th_parser,
+                                h_mom_temp, h_mom_vel,
+                                flux_normal_axis, flux_direction);
 }
 
 void PlasmaInjector::setupNuniformPerCell (const amrex::ParmParse& pp_species_name)
@@ -420,7 +427,9 @@ void PlasmaInjector::setupNuniformPerCell (const amrex::ParmParse& pp_species_na
                              num_particles_per_cell_each_dim[2];
     SpeciesUtils::parseDensity(species_name, h_inj_rho, density_parser);
     SpeciesUtils::parseMomentum(species_name, "nuniformpercell", h_inj_mom,
-        ux_parser, uy_parser, uz_parser, h_mom_temp, h_mom_vel);
+                                ux_parser, uy_parser, uz_parser,
+                                ux_th_parser, uy_th_parser, uz_th_parser,
+                                h_mom_temp, h_mom_vel);
 }
 
 void PlasmaInjector::setupExternalFile (const amrex::ParmParse& pp_species_name)
diff --git a/Source/Utils/SpeciesUtils.H b/Source/Utils/SpeciesUtils.H
@@ -30,6 +30,9 @@ namespace SpeciesUtils {
         std::unique_ptr<amrex::Parser>& ux_parser,
         std::unique_ptr<amrex::Parser>& uy_parser,
         std::unique_ptr<amrex::Parser>& uz_parser,
+        std::unique_ptr<amrex::Parser>& ux_th_parser,
+        std::unique_ptr<amrex::Parser>& uy_th_parser,
+        std::unique_ptr<amrex::Parser>& uz_th_parser,
         std::unique_ptr<TemperatureProperties>& h_mom_temp,
         std::unique_ptr<VelocityProperties>& h_mom_vel,
         int flux_normal_axis=0, int flux_direction=0);
diff --git a/Source/Utils/SpeciesUtils.cpp b/Source/Utils/SpeciesUtils.cpp

Original file line number	Diff line number	Diff line change
`@@ -15,6 +15,7 @@ void InjectorMomentum::clear ()`
`15`	`15`	`{`
`16`	`16`	`case Type::parser:`
`17`	`17`	`case Type::gaussian:`
	`18`	`+ case Type::gaussianparser:`
`18`	`19`	`case Type::gaussianflux:`
`19`	`20`	`case Type::uniform:`
`20`	`21`	`case Type::boltzmann:`