Move shared functionality between `picmi.UniformDistribution` and `picmi.AnalyticDistribution` into a parent class (#3476)

* allow density function in picmi

* created a new parent class to handle common functionality between `UniformDistribution` and `AnalyticDistribution`

* fix bug caught by the `Python_Langmuir_rz_multimode` CI test

* Place picmi inherited parent classes first to properly render documentation.

* fix issue due to secondary parent class `init` not being called

1 files changed, 43 insertions, 63 deletions

diff --git a/Python/pywarpx/picmi.py b/Python/pywarpx/picmi.py
index 30b9f4c7c..a85cc3784 100644
--- a/Python/pywarpx/picmi.py
+++ b/Python/pywarpx/picmi.py
@@ -321,9 +321,20 @@ class GaussianBunchDistribution(picmistandard.PICMI_GaussianBunchDistribution):
             species.uz = self.centroid_velocity[2]/constants.c
 
 
-class UniformDistribution(picmistandard.PICMI_UniformDistribution):
-    def initialize_inputs(self, species_number, layout, species, density_scale):
+class DensityDistributionBase(object):
+    """This is a base class for several predefined density distributions. It
+    captures universal initialization logic."""
+
+    def set_mangle_dict(self):
+        if not hasattr(self, 'mangle_dict'):
+            self.mangle_dict = None
+
+        if hasattr(self, "user_defined_kw") and self.mangle_dict is None:
+            # Only do this once so that the same variables can be used multiple
+            # times
+            self.mangle_dict = pywarpx.my_constants.add_keywords(self.user_defined_kw)
 
+    def set_species_attributes(self, species, layout):
         if isinstance(layout, GriddedLayout):
             # --- Note that the grid attribute of GriddedLayout is ignored
             species.injection_style = "nuniformpercell"
@@ -342,14 +353,16 @@ class UniformDistribution(picmistandard.PICMI_UniformDistribution):
         species.zmin = self.lower_bound[2]
         species.zmax = self.upper_bound[2]
 
-        # --- Only constant density is supported at this time
-        species.profile = "constant"
-        species.density = self.density
-        if density_scale is not None:
-            species.density *= density_scale
+        if self.fill_in:
+            species.do_continuous_injection = 1
 
         # --- Note that WarpX takes gamma*beta as input
-        if np.any(np.not_equal(self.rms_velocity, 0.)):
+        if (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)
+        elif np.any(np.not_equal(self.rms_velocity, 0.)):
             species.momentum_distribution_type = "gaussian"
             species.ux_m = self.directed_velocity[0]/constants.c
             species.uy_m = self.directed_velocity[1]/constants.c
@@ -363,74 +376,41 @@ class UniformDistribution(picmistandard.PICMI_UniformDistribution):
             species.uy = self.directed_velocity[1]/constants.c
             species.uz = self.directed_velocity[2]/constants.c
 
-        if self.fill_in:
-            species.do_continuous_injection = 1
-
+    def setup_parse_momentum_functions(self, species):
+        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)
+            else:
+                expression = f'{self.directed_velocity[idir]}'
+            species.__setattr__(f'momentum_function_u{sdir}(x,y,z)', f'({expression})/{constants.c}')
 
-class AnalyticDistribution(picmistandard.PICMI_AnalyticDistribution):
-    def init(self, kw):
-        self.mangle_dict = None
 
+class UniformDistribution(picmistandard.PICMI_UniformDistribution, DensityDistributionBase):
     def initialize_inputs(self, species_number, layout, species, density_scale):
 
-        if isinstance(layout, GriddedLayout):
-            # --- Note that the grid attribute of GriddedLayout is ignored
-            species.injection_style = "nuniformpercell"
-            species.num_particles_per_cell_each_dim = layout.n_macroparticle_per_cell
-        elif isinstance(layout, PseudoRandomLayout):
-            assert (layout.n_macroparticles_per_cell is not None), Exception('WarpX only supports n_macroparticles_per_cell for the PseudoRandomLayout with UniformDistribution')
-            species.injection_style = "nrandompercell"
-            species.num_particles_per_cell = layout.n_macroparticles_per_cell
-        else:
-            raise Exception('WarpX does not support the specified layout for UniformDistribution')
+        self.set_mangle_dict()
+        self.set_species_attributes(species, layout)
 
-        species.xmin = self.lower_bound[0]
-        species.xmax = self.upper_bound[0]
-        species.ymin = self.lower_bound[1]
-        species.ymax = self.upper_bound[1]
-        species.zmin = self.lower_bound[2]
-        species.zmax = self.upper_bound[2]
+        # --- Only constant density is supported by this class
+        species.profile = "constant"
+        species.density = self.density
+        if density_scale is not None:
+            species.density *= density_scale
 
-        if self.mangle_dict is None:
-            # Only do this once so that the same variables are used in this distribution
-            # is used multiple times
-            self.mangle_dict = pywarpx.my_constants.add_keywords(self.user_defined_kw)
-        expression = pywarpx.my_constants.mangle_expression(self.density_expression, self.mangle_dict)
+
+class AnalyticDistribution(picmistandard.PICMI_AnalyticDistribution, DensityDistributionBase):
+    def initialize_inputs(self, species_number, layout, species, density_scale):
+
+        self.set_mangle_dict()
+        self.set_species_attributes(species, layout)
 
         species.profile = "parse_density_function"
+        expression = pywarpx.my_constants.mangle_expression(self.density_expression, self.mangle_dict)
         if density_scale is None:
             species.__setattr__('density_function(x,y,z)', expression)
         else:
             species.__setattr__('density_function(x,y,z)', "{}*({})".format(density_scale, expression))
 
-        # --- Note that WarpX takes gamma*beta as input
-        if np.any(np.not_equal(self.momentum_expressions, None)):
-            species.momentum_distribution_type = 'parse_momentum_function'
-            self.setup_parse_momentum_functions(species)
-        elif np.any(np.not_equal(self.rms_velocity, 0.)):
-            species.momentum_distribution_type = "gaussian"
-            species.ux_m = self.directed_velocity[0]/constants.c
-            species.uy_m = self.directed_velocity[1]/constants.c
-            species.uz_m = self.directed_velocity[2]/constants.c
-            species.ux_th = self.rms_velocity[0]/constants.c
-            species.uy_th = self.rms_velocity[1]/constants.c
-            species.uz_th = self.rms_velocity[2]/constants.c
-        else:
-            species.momentum_distribution_type = "constant"
-            species.ux = self.directed_velocity[0]/constants.c
-            species.uy = self.directed_velocity[1]/constants.c
-            species.uz = self.directed_velocity[2]/constants.c
-
-        if self.fill_in:
-            species.do_continuous_injection = 1
-
-    def setup_parse_momentum_functions(self, species):
-        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)
-            else:
-                expression = f'{self.directed_velocity[idir]}'
-            species.__setattr__(f'momentum_function_u{sdir}(x,y,z)', f'({expression})/{constants.c}')
 
 class ParticleListDistribution(picmistandard.PICMI_ParticleListDistribution):
     def init(self, kw):