Particle fields diagnostic support for PICMI (ECP-WarpX#4262)

* added particle field interface in picmi FieldDiagnostic

* added documentation to doc string.

* undo line deletion

* catch unexpected keywords

* simplified logic

* fixed bug and removed unnecessary error check

* passing in particle field diagnostic information using a class defined in picmi rather than a dictionary

* Added ParticleFieldDiagnostic class to be used in FieldDiagnostics

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

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

* error checking on repeat name required if not using a dictionary

* fixed docstring, changed var name to be more explicit, removed error check by using default list, bug fixes

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

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

---------

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

Python/pywarpx/picmi.py

@@ -8,6 +8,7 @@
 
 """Classes following the PICMI standard
 """
+from dataclasses import dataclass
 import os
 import re
 
@@ -1941,6 +1942,35 @@ def set_write_dir(self):
             self.diagnostic.file_prefix = os.path.join(write_dir, file_prefix)
 
 
+@dataclass(frozen=True)
+class ParticleFieldDiagnostic:
+    """
+    Class holding particle field diagnostic information to be processed in FieldDiagnostic below.
+
+    Parameters
+    ----------
+    name: str
+        Name of particle field diagnostic. If a component of a vector field, for the openPMD viewer
+        to treat it as a vector, the coordinate (i.e x, y, z) should be the last character.
+
+    func: parser str
+        Parser function to be calculated for each particle per cell. Should be of the form
+        f(x,y,z,ux,uy,uz)
+
+    do_average: (0 or 1) optional, default 1
+        Whether the diagnostic is averaged by the sum of particle weights in each cell
+
+    filter: parser str, optional
+        Parser function returning a boolean for whether to include a particle in the diagnostic.
+        If not specified, all particles will be included. The function arguments are the same
+        as the `func` above.
+    """
+    name: str
+    func: str
+    do_average: int = 1
+    filter: str = None
+
+
 class FieldDiagnostic(picmistandard.PICMI_FieldDiagnostic, WarpXDiagnosticBase):
     """
     See `Input Parameters <https://warpx.readthedocs.io/en/latest/usage/parameters.html>`_ for more information.
@@ -1970,6 +2000,15 @@ class FieldDiagnostic(picmistandard.PICMI_FieldDiagnostic, WarpXDiagnosticBase):
 
     warpx_dump_rz_modes: bool, optional
         Flag whether to dump the data for all RZ modes
+
+    warpx_particle_fields_to_plot: list of ParticleFieldDiagnostics
+        List of ParticleFieldDiagnostic classes to install in the simulation. Error
+        checking is handled in the class itself.
+
+    warpx_particle_fields_species: list of strings, optional
+        Species for which to calculate particle_fields_to_plot functions. Fields will
+        be calculated separately for each specified species. If not passed, default is
+        all of the available particle species.
     """
     def init(self, kw):
 
@@ -1983,6 +2022,8 @@ def init(self, kw):
         self.file_prefix = kw.pop('warpx_file_prefix', None)
         self.file_min_digits = kw.pop('warpx_file_min_digits', None)
         self.dump_rz_modes = kw.pop('warpx_dump_rz_modes', None)
+        self.particle_fields_to_plot = kw.pop('warpx_particle_fields_to_plot', [])
+        self.particle_fields_species = kw.pop('warpx_particle_fields_species', None)
 
     def initialize_inputs(self):
 
@@ -2060,6 +2101,27 @@ def initialize_inputs(self):
             fields_to_plot.sort()
             self.diagnostic.fields_to_plot = fields_to_plot
 
+        particle_fields_to_plot_names = list()
+        for pfd in self.particle_fields_to_plot:
+            if pfd.name in particle_fields_to_plot_names:
+                raise Exception('A particle fields name can not be repeated.')
+            particle_fields_to_plot_names.append(pfd.name)
+            self.diagnostic.__setattr__(
+                f'particle_fields.{pfd.name}(x,y,z,ux,uy,uz)', pfd.func
+            )
+            self.diagnostic.__setattr__(
+                f'particle_fields.{pfd.name}.do_average', pfd.do_average
+            )
+            self.diagnostic.__setattr__(
+                f'particle_fields.{pfd.name}.filter(x,y,z,ux,uy,uz)', pfd.filter
+            )
+
+        # --- Convert to a sorted list so that the order
+        # --- is the same on all processors.
+        particle_fields_to_plot_names.sort()
+        self.diagnostic.particle_fields_to_plot = particle_fields_to_plot_names
+        self.diagnostic.particle_fields_species = self.particle_fields_species
+
         self.diagnostic.plot_raw_fields = self.plot_raw_fields
         self.diagnostic.plot_raw_fields_guards = self.plot_raw_fields_guards
         self.diagnostic.plot_finepatch = self.plot_finepatch