Implemented new particle diagnostics in picmi (#984)

* Implemented new particle diagnostics in picmi

* Cleaned up picmi adding new particle diagnostics

* In PICMI examples, use name option for diagnostics

* For travis, update ubuntu version to bionic

.travis.yml

@@ -5,7 +5,7 @@
 #
 # License: BSD-3-Clause-LBNL
 
-dist: xenial
+dist: bionic
 language: c++
 sudo: true
 cache: pip

Examples/Modules/gaussian_beam/PICMI_inputs_gaussian_beam.py

@@ -43,12 +43,14 @@
 electrons = picmi.Species(particle_type='electron', name='electrons', initial_distribution=electron_beam)
 protons = picmi.Species(particle_type='proton', name='protons', initial_distribution=proton_beam)
 
-field_diag1 = picmi.FieldDiagnostic(grid = grid,
+field_diag1 = picmi.FieldDiagnostic(name = 'diag1',
+                                    grid = grid,
                                     period = 10,
                                     data_list = ['E', 'B', 'J', 'part_per_cell'],
                                     warpx_file_prefix = 'plotfiles/plt')
 
-part_diag1 = picmi.ParticleDiagnostic(period = 10,
+part_diag1 = picmi.ParticleDiagnostic(name = 'diag1',
+                                      period = 10,
                                       species = [electrons, protons],
                                       data_list = ['weighting', 'momentum', 'fields'])
 

Examples/Physics_applications/laser_acceleration/PICMI_inputs_laser_acceleration.py

@@ -100,10 +100,12 @@
 # diagnostics
 ##########################
 
-field_diag1 = picmi.FieldDiagnostic(grid = grid,
+field_diag1 = picmi.FieldDiagnostic(name = 'diag1',
+                                    grid = grid,
                                     period = 10)
 
-part_diag1 = picmi.ParticleDiagnostic(period = 10,
+part_diag1 = picmi.ParticleDiagnostic(name = 'diag1',
+                                      period = 10,
                                       species = [electrons])
 
 ##########################

Examples/Physics_applications/plasma_acceleration/PICMI_inputs_plasma_acceleration.py

@@ -51,13 +51,15 @@
 sim.add_species(beam, layout=picmi.GriddedLayout(grid=grid, n_macroparticle_per_cell=number_per_cell_each_dim))
 sim.add_species(plasma, layout=picmi.GriddedLayout(grid=grid, n_macroparticle_per_cell=number_per_cell_each_dim))
 
-field_diag = picmi.FieldDiagnostic(grid = grid,
+field_diag = picmi.FieldDiagnostic(name = 'diag1',
+                                   grid = grid,
                                    period = max_steps,
                                    data_list = ['Ex', 'Ey', 'Ez', 'Jx', 'Jy', 'Jz', 'part_per_cell'],
                                    write_dir = 'diags',
                                    warpx_file_prefix = 'plotfiles/plt')
 
-part_diag = picmi.ParticleDiagnostic(period = max_steps,
+part_diag = picmi.ParticleDiagnostic(name = 'diag1',
+                                     period = max_steps,
                                      species = [beam, plasma],
                                      data_list = ['ux', 'uy', 'uz', 'weighting', 'Ex', 'Ey', 'Ez'],
                                      write_dir = 'diags')

Examples/Physics_applications/plasma_acceleration/PICMI_inputs_plasma_acceleration_mr.py

@@ -57,13 +57,15 @@
 sim.add_species(beam, layout=picmi.GriddedLayout(grid=grid, n_macroparticle_per_cell=number_per_cell_each_dim))
 sim.add_species(plasma, layout=picmi.GriddedLayout(grid=grid, n_macroparticle_per_cell=number_per_cell_each_dim))
 
-field_diag = picmi.FieldDiagnostic(grid = grid,
+field_diag = picmi.FieldDiagnostic(name = 'diag1',
+                                   grid = grid,
                                    period = 2,
                                    data_list = ['Ex', 'Ey', 'Ez', 'Jx', 'Jy', 'Jz', 'part_per_cell'],
                                    write_dir = 'diags',
                                    warpx_file_prefix = 'plotfiles/plt')
 
-part_diag = picmi.ParticleDiagnostic(period = 2,
+part_diag = picmi.ParticleDiagnostic(name = 'diag1',
+                                     period = 2,
                                      species = [beam, plasma],
                                      data_list = ['ux', 'uy', 'uz', 'weighting', 'Ex', 'Ey', 'Ez'],
                                      write_dir = 'diags')

Examples/Tests/Langmuir/PICMI_inputs_langmuir2d.py

@@ -60,12 +60,14 @@
 # diagnostics
 ##########################
 
-field_diag1 = picmi.FieldDiagnostic(grid = grid,
+field_diag1 = picmi.FieldDiagnostic(name = 'diag1',
+                                    grid = grid,
                                     period = diagnostic_interval,
                                     data_list = ['Ex', 'Jx'],
                                     warpx_file_prefix = 'plotfiles/plt')
 
-part_diag1 = picmi.ParticleDiagnostic(period = diagnostic_interval,
+part_diag1 = picmi.ParticleDiagnostic(name = 'diag1',
+                                      period = diagnostic_interval,
                                       species = [electrons],
                                       data_list = ['weighting', 'ux', 'Ex'])
 
@@ -90,7 +92,7 @@
 
 # write_inputs will create an inputs file that can be used to run
 # with the compiled version.
-#sim.write_input_file(file_name = 'inputs2d_from_PICMI')
+sim.write_input_file(file_name = 'inputs2d_from_PICMI')
 
 # Alternatively, sim.step will run WarpX, controlling it from Python
 sim.step()

Examples/Tests/Langmuir/PICMI_inputs_langmuir_rt.py

@@ -62,12 +62,14 @@
 # diagnostics
 ##########################
 
-field_diag1 = picmi.FieldDiagnostic(grid = grid,
+field_diag1 = picmi.FieldDiagnostic(name = 'diag1',
+                                    grid = grid,
                                     period = diagnostic_interval,
                                     data_list = ['Ex', 'Jx'],
                                     warpx_file_prefix = 'plotfiles/plt')
 
-part_diag1 = picmi.ParticleDiagnostic(period = diagnostic_interval,
+part_diag1 = picmi.ParticleDiagnostic(name = 'diag1',
+                                      period = diagnostic_interval,
                                       species = [electrons],
                                       data_list = ['weighting', 'ux', 'Ex'])
 

Examples/Tests/Langmuir/PICMI_inputs_langmuir_rz_multimode_analyze.py

@@ -95,12 +95,14 @@
 # diagnostics
 ##########################
 
-field_diag1 = picmi.FieldDiagnostic(grid = grid,
+field_diag1 = picmi.FieldDiagnostic(name = 'diag1',
+                                    grid = grid,
                                     period = diagnostic_interval,
                                     data_list = ['E', 'B', 'J', 'part_per_cell'],
                                     warpx_file_prefix = 'plotfiles/plt')
 
-part_diag1 = picmi.ParticleDiagnostic(period = diagnostic_interval,
+part_diag1 = picmi.ParticleDiagnostic(name = 'diag1',
+                                      period = diagnostic_interval,
                                       species = [electrons],
                                       data_list = ['weighting', 'momentum', 'fields'])
 

Python/pywarpx/Bucket.py

@@ -15,13 +15,23 @@ class Bucket(object):
     def __init__(self, instancename, **defaults):
         self._localsetattr('instancename', instancename)
         self._localsetattr('argvattrs', {})
-        self.argvattrs.update(defaults)
+        for name, value in defaults.items():
+            self.add_new_attr(name, value)
 
     def _localsetattr(self, name, value):
         object.__setattr__(self, name, value)
 
+    def add_new_attr(self, name, value):
+        """Names starting with "_" are make instance attributes.
+        Otherwise the attribute is added to the args list.
+        """
+        if name.startswith('_'):
+            self._localsetattr(name, value)
+        else:
+            self.argvattrs[name] = value
+
     def __setattr__(self, name, value):
-        self.argvattrs[name] = value
+        self.add_new_attr(name, value)
 
     def __getattr__(self, name):
         try:

Python/pywarpx/Diagnostics.py

@@ -6,6 +6,20 @@
 
 from .Bucket import Bucket
 
-diagnostics = Bucket('diagnostics', diags_names=[])
-diagnostics_list = []
+diagnostics = Bucket('diagnostics', _diagnostics_dict={})
+
+class Diagnostic(Bucket):
+    """
+    This is the same as a Bucket, but checks that any attributes are always given the same value.
+    """
+    def add_new_attr_with_check(self, name, value):
+        if name.startswith('_'):
+            self._localsetattr(name, value)
+        else:
+            if name in self.argvattrs:
+                assert value == self.argvattrs[name], Exception(f'Diagnostic attributes not consistent for {self.instancename}')
+            self.argvattrs[name] = value
+
+    def __setattr__(self, name, value):
+        self.add_new_attr_with_check(name, value)
 

Python/pywarpx/WarpX.py

@@ -15,7 +15,7 @@
 from .Lasers import lasers, lasers_list
 from . import Particles
 from .Particles import particles, particles_list
-from .Diagnostics import diagnostics, diagnostics_list
+from .Diagnostics import diagnostics
 
 
 class WarpX(Bucket):
@@ -54,9 +54,13 @@ def create_argv_list(self):
         for laser in lasers_list:
             argv += laser.attrlist()
 
+        diagnostics.diags_names = diagnostics._diagnostics_dict.keys()
         argv += diagnostics.attrlist()
-        for diagnostic in diagnostics_list:
+        for diagnostic in diagnostics._diagnostics_dict.values():
+            diagnostic.species = diagnostic._species_dict.keys()
             argv += diagnostic.attrlist()
+            for species_diagnostic in diagnostic._species_dict.values():
+                argv += species_diagnostic.attrlist()
 
         return argv
 

Python/pywarpx/picmi.py

@@ -77,8 +77,7 @@ def initialize_inputs(self, layout, initialize_self_fields=False):
                                              mass = self.mass,
                                              charge = self.charge,
                                              injection_style = 'python',
-                                             initialize_self_fields = int(initialize_self_fields),
-                                             plot_vars = set())
+                                             initialize_self_fields = int(initialize_self_fields))
         pywarpx.Particles.particles_list.append(self.species)
 
         if self.initial_distribution is not None:
@@ -638,16 +637,17 @@ def init(self, kw):
         self.file_prefix = kw.pop('warpx_file_prefix', None)
 
     def initialize_inputs(self):
-        self.diagnostics_number = len(pywarpx.Diagnostics.diagnostics_list) + 1
 
-        # --- This a placeholder until the name attribute is added to the picmi standard
         name = getattr(self, 'name', None)
         if name is None:
-            self.name = 'diag{}'.format(self.diagnostics_number)
+            diagnostics_number = len(pywarpx.diagnostics._diagnostics_dict) + 1
+            self.name = 'diag{}'.format(diagnostics_number)
 
-        self.diagnostic = pywarpx.Bucket.Bucket(self.name)
-        pywarpx.diagnostics.diags_names.append(self.name)
-        pywarpx.Diagnostics.diagnostics_list.append(self.diagnostic)
+        try:
+            self.diagnostic = pywarpx.diagnostics._diagnostics_dict[self.name]
+        except KeyError:
+            self.diagnostic = pywarpx.Diagnostics.Diagnostic(self.name, _species_dict={})
+            pywarpx.diagnostics._diagnostics_dict[self.name] = self.diagnostic
 
         self.diagnostic.diag_type = 'Full'
         self.diagnostic.format = self.format
@@ -712,37 +712,66 @@ def initialize_inputs(self):
 
 
 class ParticleDiagnostic(picmistandard.PICMI_ParticleDiagnostic):
+    def init(self, kw):
+
+        self.format = kw.pop('warpx_format', 'plotfile')
+        self.openpmd_backend = kw.pop('warpx_openpmd_backend', None)
+        self.file_prefix = kw.pop('warpx_file_prefix', 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)
+
     def initialize_inputs(self):
 
-        plot_vars = set()
+        name = getattr(self, 'name', None)
+        if name is None:
+            diagnostics_number = len(pywarpx.diagnostics._diagnostics_dict) + 1
+            self.name = 'diag{}'.format(diagnostics_number)
+
+        try:
+            self.diagnostic = pywarpx.diagnostics._diagnostics_dict[self.name]
+        except KeyError:
+            self.diagnostic = pywarpx.Diagnostics.Diagnostic(self.name, _species_dict={})
+            pywarpx.diagnostics._diagnostics_dict[self.name] = self.diagnostic
+
+        self.diagnostic.diag_type = 'Full'
+        self.diagnostic.format = self.format
+        self.diagnostic.openpmd_backend = self.openpmd_backend
+        self.diagnostic.period = self.period
+
+        variables = set()
         for dataname in self.data_list:
             if dataname == 'position':
                 # --- The positions are alway written out anyway
                 pass
             elif dataname == 'momentum':
-                plot_vars.add('ux')
-                plot_vars.add('uy')
-                plot_vars.add('uz')
+                variables.add('ux')
+                variables.add('uy')
+                variables.add('uz')
             elif dataname == 'weighting':
-                plot_vars.add('w')
+                variables.add('w')
             elif dataname == 'fields':
-                plot_vars.add('Ex')
-                plot_vars.add('Ey')
-                plot_vars.add('Ez')
-                plot_vars.add('Bx')
-                plot_vars.add('By')
-                plot_vars.add('Bz')
+                variables.add('Ex')
+                variables.add('Ey')
+                variables.add('Ez')
+                variables.add('Bx')
+                variables.add('By')
+                variables.add('Bz')
             elif dataname in ['ux', 'uy', 'uz', 'Ex', 'Ey', 'Ez', 'Bx', 'By', 'Bz']:
-                plot_vars.add(dataname)
-
-        if plot_vars:
-            species = self.species
-            if not np.iterable(species):
-                species = [species]
-            for specie in species:
-                for var in plot_vars:
-                    specie.species.plot_vars.add(var)
+                variables.add(dataname)
 
+        if np.iterable(self.species):
+            species_list = self.species
+        else:
+            species_list = [species]
+
+        for specie in species_list:
+            diag = pywarpx.Bucket.Bucket(self.name + '.' + specie.name,
+                                         variables = variables,
+                                         random_fraction = self.random_fraction,
+                                         uniform_stride = self.uniform_stride)
+            diag.__setattr__('plot_filter_function(t,x,y,z,ux,uy,uz)', self.plot_filter_function)
+            self.diagnostic._species_dict[specie.name] = diag
 
 # ----------------------------
 # Lab frame diagnostics