From da0a1f4917ec3cd7e481ce225780980a021180e8 Mon Sep 17 00:00:00 2001 From: David Grote Date: Thu, 14 May 2020 16:00:36 -0700 Subject: [PATCH] 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 | 2 +- .../PICMI_inputs_gaussian_beam.py | 6 +- .../PICMI_inputs_laser_acceleration.py | 6 +- .../PICMI_inputs_plasma_acceleration.py | 6 +- .../PICMI_inputs_plasma_acceleration_mr.py | 6 +- .../Tests/Langmuir/PICMI_inputs_langmuir2d.py | 8 +- .../Langmuir/PICMI_inputs_langmuir_rt.py | 6 +- ...MI_inputs_langmuir_rz_multimode_analyze.py | 6 +- Python/pywarpx/Bucket.py | 14 ++- Python/pywarpx/Diagnostics.py | 18 +++- Python/pywarpx/WarpX.py | 8 +- Python/pywarpx/picmi.py | 85 +++++++++++++------ 12 files changed, 121 insertions(+), 50 deletions(-) diff --git a/.travis.yml b/.travis.yml index 7ba184113d0..bb0958e9ed2 100644 --- a/.travis.yml +++ b/.travis.yml @@ -5,7 +5,7 @@ # # License: BSD-3-Clause-LBNL -dist: xenial +dist: bionic language: c++ sudo: true cache: pip diff --git a/Examples/Modules/gaussian_beam/PICMI_inputs_gaussian_beam.py b/Examples/Modules/gaussian_beam/PICMI_inputs_gaussian_beam.py index 8f3d345038a..d74fae29a04 100644 --- a/Examples/Modules/gaussian_beam/PICMI_inputs_gaussian_beam.py +++ b/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']) diff --git a/Examples/Physics_applications/laser_acceleration/PICMI_inputs_laser_acceleration.py b/Examples/Physics_applications/laser_acceleration/PICMI_inputs_laser_acceleration.py index 8f0b1153dd5..3adecf33043 100644 --- a/Examples/Physics_applications/laser_acceleration/PICMI_inputs_laser_acceleration.py +++ b/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]) ########################## diff --git a/Examples/Physics_applications/plasma_acceleration/PICMI_inputs_plasma_acceleration.py b/Examples/Physics_applications/plasma_acceleration/PICMI_inputs_plasma_acceleration.py index e958f99fd35..01c322508b4 100644 --- a/Examples/Physics_applications/plasma_acceleration/PICMI_inputs_plasma_acceleration.py +++ b/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') diff --git a/Examples/Physics_applications/plasma_acceleration/PICMI_inputs_plasma_acceleration_mr.py b/Examples/Physics_applications/plasma_acceleration/PICMI_inputs_plasma_acceleration_mr.py index cbbe851976c..6bcd9ac0127 100644 --- a/Examples/Physics_applications/plasma_acceleration/PICMI_inputs_plasma_acceleration_mr.py +++ b/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') diff --git a/Examples/Tests/Langmuir/PICMI_inputs_langmuir2d.py b/Examples/Tests/Langmuir/PICMI_inputs_langmuir2d.py index 99522b0d675..effc1b24439 100644 --- a/Examples/Tests/Langmuir/PICMI_inputs_langmuir2d.py +++ b/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() diff --git a/Examples/Tests/Langmuir/PICMI_inputs_langmuir_rt.py b/Examples/Tests/Langmuir/PICMI_inputs_langmuir_rt.py index 31b41088c0d..04cd05401b7 100644 --- a/Examples/Tests/Langmuir/PICMI_inputs_langmuir_rt.py +++ b/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']) diff --git a/Examples/Tests/Langmuir/PICMI_inputs_langmuir_rz_multimode_analyze.py b/Examples/Tests/Langmuir/PICMI_inputs_langmuir_rz_multimode_analyze.py index 1e2940b283b..46e5da9e4df 100644 --- a/Examples/Tests/Langmuir/PICMI_inputs_langmuir_rz_multimode_analyze.py +++ b/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']) diff --git a/Python/pywarpx/Bucket.py b/Python/pywarpx/Bucket.py index 714516fd3e1..7817f8f8647 100644 --- a/Python/pywarpx/Bucket.py +++ b/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: diff --git a/Python/pywarpx/Diagnostics.py b/Python/pywarpx/Diagnostics.py index 25d9ae9587c..4846a88c86f 100644 --- a/Python/pywarpx/Diagnostics.py +++ b/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) diff --git a/Python/pywarpx/WarpX.py b/Python/pywarpx/WarpX.py index edcf0524b2c..12afc4e15c7 100644 --- a/Python/pywarpx/WarpX.py +++ b/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 diff --git a/Python/pywarpx/picmi.py b/Python/pywarpx/picmi.py index ef5e810d754..de1bb1363af 100644 --- a/Python/pywarpx/picmi.py +++ b/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