Add scraping boundary diagnostics with PICMI (ECP-WarpX#5109)

* Removal of asserttion which prevented from usung the averaged PSATD algorithms with PML BC

* Added ParticleBoundaryScrapingDiagnostic class.

* Temp commit: testing an unreleased PICMI version.

* Temp commit: updated PICMI version.

* Temp commit: updated PICMI version in setup.py

* Activated scraping boundary diagnostics in the CI test: spacecraft_charging

* Use latest picmistandard version

* Revert "Removal of asserttion which prevented from usung the averaged PSATD algorithms with PML BC"

This reverts commit 9ee80aa.

* Update Examples/Physics_applications/spacecraft_charging/PICMI_inputs_rz.py

* Set writing interval to -1 for particle boundary scraping diagnostics in spacecraft_charging CI test.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

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

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.29.0
+picmistandard==0.30.0
 # for development against an unreleased PICMI version, use:
 # picmistandard @ git+https://github.com/picmi-standard/picmi.git#subdirectory=PICMI_Python
 

Examples/Physics_applications/spacecraft_charging/PICMI_inputs_rz.py

@@ -275,6 +275,13 @@ def compute_actual_charge_on_spacecraft():
     warpx_file_prefix="spacecraft_charging_plt",
 )
 
+part_scraping_boundary_diag = picmi.ParticleBoundaryScrapingDiagnostic(
+    name="diag2",
+    period=-1,  # only at the end, because we also use the buffers in this test
+    species=[electrons, protons],
+    warpx_format="openpmd",
+)
+
 ##########################
 # simulation setup
 ##########################
@@ -298,6 +305,7 @@ def compute_actual_charge_on_spacecraft():
 
 sim.add_diagnostic(field_diag)
 sim.add_diagnostic(part_diag)
+sim.add_diagnostic(part_scraping_boundary_diag)
 
 ##########################
 # simulation run

Python/pywarpx/picmi.py

@@ -3969,3 +3969,178 @@ def diagnostic_initialize_inputs(self):
                     self.diagnostic.__setattr__(key, expression)
                 else:
                     self.diagnostic.__setattr__(key, value)
+
+
+class ParticleBoundaryScrapingDiagnostic(
+    picmistandard.PICMI_ParticleBoundaryScrapingDiagnostic, WarpXDiagnosticBase
+):
+    """
+    See `Input Parameters <https://warpx.readthedocs.io/en/latest/usage/parameters.html>`__ for more information.
+
+    Parameters
+    ----------
+    warpx_format: openpmd
+        Diagnostic file format
+
+    warpx_openpmd_backend: {bp, h5, json}, optional
+        Openpmd backend file format
+
+    warpx_openpmd_encoding: 'v' (variable based), 'f' (file based) or 'g' (group based), optional
+        Only read if ``<diag_name>.format = openpmd``. openPMD file output encoding.
+        File based: one file per timestep (slower), group/variable based: one file for all steps (faster)).
+        Variable based is an experimental feature with ADIOS2. Default: `'f'`.
+
+    warpx_file_prefix: string, optional
+        Prefix on the diagnostic file name
+
+    warpx_file_min_digits: integer, optional
+        Minimum number of digits for the time step number in the file name
+
+    warpx_random_fraction: float or dict, optional
+        Random fraction of particles to include in the diagnostic. If a float
+        is given the same fraction will be used for all species, if a dictionary
+        is given the keys should be species with the value specifying the random
+        fraction for that species.
+
+    warpx_uniform_stride: integer or dict, optional
+        Stride to down select to the particles to include in the diagnostic.
+        If an integer is given the same stride will be used for all species, if
+        a dictionary is given the keys should be species with the value
+        specifying the stride for that species.
+
+    warpx_dump_last_timestep: bool, optional
+        If true, the last timestep is dumped regardless of the diagnostic period/intervals.
+
+    warpx_plot_filter_function: string, optional
+        Analytic expression to down select the particles to in the diagnostic
+    """
+
+    def init(self, kw):
+        self.format = kw.pop("warpx_format", "openpmd")
+        self.openpmd_backend = kw.pop("warpx_openpmd_backend", None)
+        self.openpmd_encoding = kw.pop("warpx_openpmd_encoding", None)
+        self.file_prefix = kw.pop("warpx_file_prefix", None)
+        self.file_min_digits = kw.pop("warpx_file_min_digits", None)
+        self.random_fraction = kw.pop("warpx_random_fraction", None)
+        self.uniform_stride = kw.pop("warpx_uniform_stride", None)
+        self.plot_filter_function = kw.pop("warpx_plot_filter_function", None)
+        self.dump_last_timestep = kw.pop("warpx_dump_last_timestep", None)
+
+        self.user_defined_kw = {}
+        if self.plot_filter_function is not None:
+            # This allows variables to be used in the plot_filter_function, but
+            # in order not to break other codes, the variables must begin with "warpx_"
+            for k in list(kw.keys()):
+                if k.startswith("warpx_") and re.search(
+                    r"\b%s\b" % k, self.plot_filter_function
+                ):
+                    self.user_defined_kw[k] = kw[k]
+                    del kw[k]
+
+        self.mangle_dict = None
+
+    def diagnostic_initialize_inputs(self):
+        self.add_diagnostic()
+
+        self.diagnostic.diag_type = "BoundaryScraping"
+        self.diagnostic.format = self.format
+        self.diagnostic.openpmd_backend = self.openpmd_backend
+        self.diagnostic.openpmd_encoding = self.openpmd_encoding
+        self.diagnostic.file_min_digits = self.file_min_digits
+        self.diagnostic.dump_last_timestep = self.dump_last_timestep
+        self.diagnostic.intervals = self.period
+        self.diagnostic.set_or_replace_attr("write_species", True)
+        if "fields_to_plot" not in self.diagnostic.argvattrs:
+            self.diagnostic.fields_to_plot = "none"
+
+        self.set_write_dir()
+
+        # --- Use a set to ensure that fields don't get repeated.
+        variables = set()
+
+        if self.data_list is not None:
+            for dataname in self.data_list:
+                if dataname == "position":
+                    if pywarpx.geometry.dims != "1":  # because then it's WARPX_DIM_1D_Z
+                        variables.add("x")
+                    if pywarpx.geometry.dims == "3":
+                        variables.add("y")
+                    variables.add("z")
+                    if pywarpx.geometry.dims == "RZ":
+                        variables.add("theta")
+                elif dataname == "momentum":
+                    variables.add("ux")
+                    variables.add("uy")
+                    variables.add("uz")
+                elif dataname == "weighting":
+                    variables.add("w")
+                elif dataname in ["x", "y", "z", "theta", "ux", "uy", "uz"]:
+                    if pywarpx.geometry.dims == "1" and (
+                        dataname == "x" or dataname == "y"
+                    ):
+                        raise RuntimeError(
+                            f"The attribute {dataname} is not available in mode WARPX_DIM_1D_Z"
+                            f"chosen by dim={pywarpx.geometry.dims} in pywarpx."
+                        )
+                    elif pywarpx.geometry.dims != "3" and dataname == "y":
+                        raise RuntimeError(
+                            f"The attribute {dataname} is not available outside of mode WARPX_DIM_3D"
+                            f"The chosen value was dim={pywarpx.geometry.dims} in pywarpx."
+                        )
+                    elif pywarpx.geometry.dims != "RZ" and dataname == "theta":
+                        raise RuntimeError(
+                            f"The attribute {dataname} is not available outside of mode WARPX_DIM_RZ."
+                            f"The chosen value was dim={pywarpx.geometry.dims} in pywarpx."
+                        )
+                    else:
+                        variables.add(dataname)
+                else:
+                    # possibly add user defined attributes
+                    variables.add(dataname)
+
+            # --- Convert the set to a sorted list so that the order
+            # --- is the same on all processors.
+            variables = list(variables)
+            variables.sort()
+
+        # species list
+        if self.species is None:
+            species_names = pywarpx.particles.species_names
+        elif np.iterable(self.species):
+            species_names = [species.name for species in self.species]
+        else:
+            species_names = [self.species.name]
+
+        # check if random fraction is specified and whether a value is given per species
+        random_fraction = {}
+        random_fraction_default = self.random_fraction
+        if isinstance(self.random_fraction, dict):
+            random_fraction_default = 1.0
+            for key, val in self.random_fraction.items():
+                random_fraction[key.name] = val
+
+        # check if uniform stride is specified and whether a value is given per species
+        uniform_stride = {}
+        uniform_stride_default = self.uniform_stride
+        if isinstance(self.uniform_stride, dict):
+            uniform_stride_default = 1
+            for key, val in self.uniform_stride.items():
+                uniform_stride[key.name] = val
+
+        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)
+
+        for name in species_names:
+            diag = pywarpx.Bucket.Bucket(
+                self.name + "." + name,
+                variables=variables,
+                random_fraction=random_fraction.get(name, random_fraction_default),
+                uniform_stride=uniform_stride.get(name, uniform_stride_default),
+            )
+            expression = pywarpx.my_constants.mangle_expression(
+                self.plot_filter_function, self.mangle_dict
+            )
+            diag.__setattr__("plot_filter_function(t,x,y,z,ux,uy,uz)", expression)
+            self.diagnostic._species_dict[name] = diag

Python/setup.py

@@ -70,7 +70,7 @@
     package_dir={"pywarpx": "pywarpx"},
     description="""Wrapper of WarpX""",
     package_data=package_data,
-    install_requires=["numpy", "picmistandard==0.29.0", "periodictable"],
+    install_requires=["numpy", "picmistandard==0.30.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.29.0
+python -m pip install --user --upgrade picmistandard==0.30.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.29.0
+picmistandard==0.30.0
 #   for development against an unreleased PICMI version, use:
 #picmistandard @ git+https://github.com/picmi-standard/picmi.git#subdirectory=PICMI_Python