Added AnalyticFluxDistribution class (#5422)

Added an AnalyticFluxDistribution class with a parsed `flux_expression`.
Depends on picmi-standard/picmi#121

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: David Grote <[email protected]>

Docs/requirements.txt

@@ -13,7 +13,7 @@ openpmd-viewer  # for checksumAPI
 
 # PICMI API docs
 # note: keep in sync with version in ../requirements.txt
-picmistandard==0.31.0
+picmistandard==0.33.0
 # for development against an unreleased PICMI version, use:
 # picmistandard @ git+https://github.com/picmi-standard/picmi.git#subdirectory=PICMI_Python
 

Docs/source/usage/python.rst

@@ -146,6 +146,10 @@ Particle distributions can be used for to initialize particles in a particle spe
 
 .. autoclass:: pywarpx.picmi.AnalyticDistribution
 
+.. autoclass:: pywarpx.picmi.UniformFluxDistribution
+
+.. autoclass:: pywarpx.picmi.AnalyticFluxDistribution
+
 .. autoclass:: pywarpx.picmi.ParticleListDistribution
 
 Particle layouts determine how to microscopically place macro particles in a grid cell.

Python/pywarpx/picmi.py

@@ -688,25 +688,42 @@ def setup_parse_momentum_functions(
             )
 
 
-class UniformFluxDistribution(
-    picmistandard.PICMI_UniformFluxDistribution, DensityDistributionBase
+class UniformDistribution(
+    picmistandard.PICMI_UniformDistribution, DensityDistributionBase
 ):
+    def distribution_initialize_inputs(
+        self, species_number, layout, species, density_scale, source_name
+    ):
+        self.set_mangle_dict()
+        self.set_species_attributes(species, layout, source_name)
+
+        # --- Only constant density is supported by this class
+        species.add_new_group_attr(source_name, "profile", "constant")
+        species.add_new_group_attr(source_name, "density", self.density)
+        if density_scale is not None:
+            species.add_new_group_attr(source_name, "density", density_scale)
+
+
+class FluxDistributionBase(object):
+    """This is a base class for both uniform and analytic flux distributions."""
+
     def init(self, kw):
         self.inject_from_embedded_boundary = kw.pop(
             "warpx_inject_from_embedded_boundary", False
         )
 
+    def initialize_flux_profile_func(self, species, density_scale, source_name):
+        """Initialize the flux profile and flux function."""
+        pass
+
     def distribution_initialize_inputs(
         self, species_number, layout, species, density_scale, source_name
     ):
         self.fill_in = False
         self.set_mangle_dict()
         self.set_species_attributes(species, layout, source_name)
 
-        species.add_new_group_attr(source_name, "flux_profile", "constant")
-        species.add_new_group_attr(source_name, "flux", self.flux)
-        if density_scale is not None:
-            species.add_new_group_attr(source_name, "flux", density_scale)
+        self.initialize_flux_profile_func(species, density_scale, source_name)
 
         if not self.inject_from_embedded_boundary:
             species.add_new_group_attr(
@@ -737,20 +754,62 @@ def distribution_initialize_inputs(
             )
 
 
-class UniformDistribution(
-    picmistandard.PICMI_UniformDistribution, DensityDistributionBase
+class AnalyticFluxDistribution(
+    picmistandard.PICMI_AnalyticFluxDistribution,
+    FluxDistributionBase,
+    DensityDistributionBase,
 ):
-    def distribution_initialize_inputs(
-        self, species_number, layout, species, density_scale, source_name
-    ):
-        self.set_mangle_dict()
-        self.set_species_attributes(species, layout, source_name)
+    """
+    Parameters
+    ----------
 
-        # --- Only constant density is supported by this class
-        species.add_new_group_attr(source_name, "profile", "constant")
-        species.add_new_group_attr(source_name, "density", self.density)
+    warpx_inject_from_embedded_boundary: bool
+        When true, the flux is injected from the embedded boundaries instead
+        of a plane.
+    """
+
+    def init(self, kw):
+        FluxDistributionBase.init(self, kw)
+
+    def initialize_flux_profile_func(self, species, density_scale, source_name):
+        species.add_new_group_attr(source_name, "flux_profile", "parse_flux_function")
         if density_scale is not None:
-            species.add_new_group_attr(source_name, "density", density_scale)
+            species.add_new_group_attr(source_name, "flux", density_scale)
+        expression = pywarpx.my_constants.mangle_expression(self.flux, self.mangle_dict)
+        if density_scale is None:
+            species.add_new_group_attr(
+                source_name, "flux_function(x,y,z,t)", expression
+            )
+        else:
+            species.add_new_group_attr(
+                source_name,
+                "flux_function(x,y,z,t)",
+                "{}*({})".format(density_scale, expression),
+            )
+
+
+class UniformFluxDistribution(
+    picmistandard.PICMI_UniformFluxDistribution,
+    FluxDistributionBase,
+    DensityDistributionBase,
+):
+    """
+    Parameters
+    ----------
+
+    warpx_inject_from_embedded_boundary: bool
+        When true, the flux is injected from the embedded boundaries instead
+        of a plane.
+    """
+
+    def init(self, kw):
+        FluxDistributionBase.init(self, kw)
+
+    def initialize_flux_profile_func(self, species, density_scale, source_name):
+        species.add_new_group_attr(source_name, "flux_profile", "constant")
+        species.add_new_group_attr(source_name, "flux", self.flux)
+        if density_scale is not None:
+            species.add_new_group_attr(source_name, "flux", density_scale)
 
 
 class AnalyticDistribution(

Python/setup.py

@@ -70,7 +70,7 @@
     package_dir={"pywarpx": "pywarpx"},
     description="""Wrapper of WarpX""",
     package_data=package_data,
-    install_requires=["numpy", "picmistandard==0.31.0", "periodictable"],
+    install_requires=["numpy", "picmistandard==0.33.0", "periodictable"],
     python_requires=">=3.8",
     zip_safe=False,
 )

Tools/machines/karolina-it4i/install_dependencies.sh

@@ -53,7 +53,7 @@ python -m pip install --user --upgrade matplotlib
 #python -m pip install --user --upgrade yt
 
 # install or update WarpX dependencies
-python -m pip install --user --upgrade picmistandard==0.31.0
+python -m pip install --user --upgrade picmistandard==0.33.0
 python -m pip install --user --upgrade lasy
 
 # optional: for optimas (based on libEnsemble & ax->botorch->gpytorch->pytorch)

requirements.txt

@@ -4,7 +4,7 @@ periodictable~=1.5
 
 # PICMI
 # note: don't forget to update the version in Docs/requirements.txt, too
-picmistandard==0.31.0
+picmistandard==0.33.0
 #   for development against an unreleased PICMI version, use:
 #picmistandard @ git+https://github.com/picmi-standard/picmi.git#subdirectory=PICMI_Python