Add new option to PICMI

Python/pywarpx/picmi.py

@@ -392,11 +392,17 @@ def set_species_attributes(self, species, layout):
             species.do_continuous_injection = 1
 
         # --- Note that WarpX takes gamma*beta as input
-        if (hasattr(self, "momentum_expressions")
+        if (hasattr(self, "momentum_spread_expressions")
+            and np.any(np.not_equal(self.momentum_spread_expressions, None))
+        ):
+            species.momentum_distribution_type = 'gaussian_parse_momentum_function'
+            self.setup_parse_momentum_functions(species, self.momentum_expressions, '_m', self.directed_velocity)
+            self.setup_parse_momentum_functions(species, self.momentum_spread_expressions, '_th', [0.,0.,0.])
+        elif (hasattr(self, "momentum_expressions")
             and np.any(np.not_equal(self.momentum_expressions, None))
         ):
             species.momentum_distribution_type = 'parse_momentum_function'
-            self.setup_parse_momentum_functions(species)
+            self.setup_parse_momentum_functions(species, self.momentum_expressions, '', self.directed_velocity)
         elif np.any(np.not_equal(self.rms_velocity, 0.)):
             species.momentum_distribution_type = "gaussian"
             species.ux_m = self.directed_velocity[0]/constants.c
@@ -411,13 +417,13 @@ def set_species_attributes(self, species, layout):
             species.uy = self.directed_velocity[1]/constants.c
             species.uz = self.directed_velocity[2]/constants.c
 
-    def setup_parse_momentum_functions(self, species):
+    def setup_parse_momentum_functions(self, species, expressions, suffix, defaults):
         for sdir, idir in zip(['x', 'y', 'z'], [0, 1, 2]):
-            if self.momentum_expressions[idir] is not None:
-                expression = pywarpx.my_constants.mangle_expression(self.momentum_expressions[idir], self.mangle_dict)
+            if expressions[idir] is not None:
+                expression = pywarpx.my_constants.mangle_expression(expressions[idir], self.mangle_dict)
             else:
-                expression = f'{self.directed_velocity[idir]}'
-            species.__setattr__(f'momentum_function_u{sdir}(x,y,z)', f'({expression})/{constants.c}')
+                expression = f'{defaults[idir]}'
+            species.__setattr__(f'momentum_function_u{sdir}{suffix}(x,y,z)', f'({expression})/{constants.c}')
 
 
 class UniformFluxDistribution(picmistandard.PICMI_UniformFluxDistribution, DensityDistributionBase):
@@ -458,6 +464,9 @@ def initialize_inputs(self, species_number, layout, species, density_scale):
 
 
 class AnalyticDistribution(picmistandard.PICMI_AnalyticDistribution, DensityDistributionBase):
+    def init(self, kw):
+        self.momentum_spread_expressions = kw.pop('warpx_momentum_spread_expressions', [None, None, None])
+
     def initialize_inputs(self, species_number, layout, species, density_scale):
 
         self.set_mangle_dict()