diff options
Diffstat (limited to 'Python/pywarpx')
| -rw-r--r-- | Python/pywarpx/Diagnostics.py | 1 | ||||
| -rw-r--r-- | Python/pywarpx/WarpX.py | 10 | ||||
| -rw-r--r-- | Python/pywarpx/__init__.py | 2 | ||||
| -rw-r--r-- | Python/pywarpx/picmi.py | 228 |
4 files changed, 184 insertions, 57 deletions
diff --git a/Python/pywarpx/Diagnostics.py b/Python/pywarpx/Diagnostics.py index 408ce6d14..06562c28d 100644 --- a/Python/pywarpx/Diagnostics.py +++ b/Python/pywarpx/Diagnostics.py @@ -7,6 +7,7 @@ from .Bucket import Bucket diagnostics = Bucket('diagnostics', _diagnostics_dict={}) +reduced_diagnostics = Bucket('warpx', _diagnostics_dict={}) class Diagnostic(Bucket): """ diff --git a/Python/pywarpx/WarpX.py b/Python/pywarpx/WarpX.py index 175917ab6..517c74d42 100644 --- a/Python/pywarpx/WarpX.py +++ b/Python/pywarpx/WarpX.py @@ -14,7 +14,7 @@ from .Boundary import boundary from .Bucket import Bucket from .Collisions import collisions, collisions_list from .Constants import my_constants -from .Diagnostics import diagnostics +from .Diagnostics import diagnostics, reduced_diagnostics from .EB2 import eb2 from .Geometry import geometry from .Interpolation import interpolation @@ -76,6 +76,14 @@ class WarpX(Bucket): for species_diagnostic in diagnostic._species_dict.values(): argv += species_diagnostic.attrlist() + reduced_diagnostics.reduced_diags_names = reduced_diagnostics._diagnostics_dict.keys() + argv += reduced_diagnostics.attrlist() + for diagnostic in reduced_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 def init(self, mpi_comm=None): diff --git a/Python/pywarpx/__init__.py b/Python/pywarpx/__init__.py index 8f4960456..58d214398 100644 --- a/Python/pywarpx/__init__.py +++ b/Python/pywarpx/__init__.py @@ -9,7 +9,7 @@ from .Amr import amr from .Boundary import boundary from .Collisions import collisions from .Constants import my_constants -from .Diagnostics import diagnostics +from .Diagnostics import diagnostics, reduced_diagnostics from .EB2 import eb2 from .Geometry import geometry from .Interpolation import interpolation diff --git a/Python/pywarpx/picmi.py b/Python/pywarpx/picmi.py index 4640a2e2e..3005b06bd 100644 --- a/Python/pywarpx/picmi.py +++ b/Python/pywarpx/picmi.py @@ -1595,8 +1595,43 @@ class Simulation(picmistandard.PICMI_Simulation): # Simulation frame diagnostics # ---------------------------- +class WarpXDiagnosticBase(object): + """ + Base class for all WarpX diagnostic containing functionality shared by + all WarpX diagnostic installations. + """ + def add_diagnostic(self): + # reduced diagnostics go in a different bucket than regular diagnostics + if isinstance(self, ReducedDiagnostic): + bucket = pywarpx.reduced_diagnostics + name_template = 'reduced_diag' + else: + bucket = pywarpx.diagnostics + name_template = 'diag' + + name = getattr(self, 'name', None) + if name is None: + diagnostics_number = (len(bucket._diagnostics_dict) + 1) + self.name = f'{name_template}{diagnostics_number}' + + try: + self.diagnostic = bucket._diagnostics_dict[self.name] + except KeyError: + self.diagnostic = pywarpx.Diagnostics.Diagnostic( + self.name, _species_dict={} + ) + bucket._diagnostics_dict[self.name] = self.diagnostic + + -class FieldDiagnostic(picmistandard.PICMI_FieldDiagnostic): + def set_write_dir(self): + if self.write_dir is not None or self.file_prefix is not None: + write_dir = (self.write_dir or 'diags') + file_prefix = (self.file_prefix or self.name) + self.diagnostic.file_prefix = os.path.join(write_dir, file_prefix) + + +class FieldDiagnostic(picmistandard.PICMI_FieldDiagnostic, WarpXDiagnosticBase): """ See `Input Parameters <https://warpx.readthedocs.io/en/latest/usage/parameters.html>`_ for more information. @@ -1637,16 +1672,7 @@ class FieldDiagnostic(picmistandard.PICMI_FieldDiagnostic): def initialize_inputs(self): - 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.add_diagnostic() self.diagnostic.diag_type = 'Full' self.diagnostic.format = self.format @@ -1709,16 +1735,13 @@ class FieldDiagnostic(picmistandard.PICMI_FieldDiagnostic): self.diagnostic.plot_finepatch = self.plot_finepatch self.diagnostic.plot_crsepatch = self.plot_crsepatch - if self.write_dir is not None or self.file_prefix is not None: - write_dir = (self.write_dir or 'diags') - file_prefix = (self.file_prefix or self.name) - self.diagnostic.file_prefix = os.path.join(write_dir, file_prefix) + self.set_write_dir() ElectrostaticFieldDiagnostic = FieldDiagnostic -class Checkpoint(picmistandard.base._ClassWithInit): +class Checkpoint(picmistandard.base._ClassWithInit, WarpXDiagnosticBase): """ Sets up checkpointing of the simulation, allowing for later restarts @@ -1749,23 +1772,17 @@ class Checkpoint(picmistandard.base._ClassWithInit): def initialize_inputs(self): - 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.add_diagnostic() self.diagnostic.intervals = self.period self.diagnostic.diag_type = 'Full' self.diagnostic.format = 'checkpoint' self.diagnostic.file_min_digits = self.file_min_digits - if self.write_dir is not None or self.file_prefix is not None: - write_dir = (self.write_dir or 'diags') - file_prefix = (self.file_prefix or self.name) - self.diagnostic.file_prefix = os.path.join(write_dir, file_prefix) + self.set_write_dir() -class ParticleDiagnostic(picmistandard.PICMI_ParticleDiagnostic): + +class ParticleDiagnostic(picmistandard.PICMI_ParticleDiagnostic, WarpXDiagnosticBase): """ See `Input Parameters <https://warpx.readthedocs.io/en/latest/usage/parameters.html>`_ for more information. @@ -1815,16 +1832,7 @@ class ParticleDiagnostic(picmistandard.PICMI_ParticleDiagnostic): def initialize_inputs(self): - 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.add_diagnostic() self.diagnostic.diag_type = 'Full' self.diagnostic.format = self.format @@ -1832,10 +1840,7 @@ class ParticleDiagnostic(picmistandard.PICMI_ParticleDiagnostic): self.diagnostic.file_min_digits = self.file_min_digits self.diagnostic.intervals = self.period - if self.write_dir is not None or self.file_prefix is not None: - write_dir = (self.write_dir or 'diags') - file_prefix = (self.file_prefix or self.name) - self.diagnostic.file_prefix = os.path.join(write_dir, file_prefix) + self.set_write_dir() # --- Use a set to ensure that fields don't get repeated. variables = set() @@ -1890,7 +1895,8 @@ class ParticleDiagnostic(picmistandard.PICMI_ParticleDiagnostic): # ---------------------------- -class LabFrameFieldDiagnostic(picmistandard.PICMI_LabFrameFieldDiagnostic): +class LabFrameFieldDiagnostic(picmistandard.PICMI_LabFrameFieldDiagnostic, + WarpXDiagnosticBase): """ See `Input Parameters <https://warpx.readthedocs.io/en/latest/usage/parameters.html>`_ for more information. @@ -1952,16 +1958,7 @@ class LabFrameFieldDiagnostic(picmistandard.PICMI_LabFrameFieldDiagnostic): def initialize_inputs_new(self): - 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.add_diagnostic() self.diagnostic.diag_type = 'BackTransformed' self.diagnostic.format = self.format @@ -2006,10 +2003,7 @@ class LabFrameFieldDiagnostic(picmistandard.PICMI_LabFrameFieldDiagnostic): fields_to_plot.sort() self.diagnostic.fields_to_plot = fields_to_plot - if self.write_dir is not None or self.file_prefix is not None: - write_dir = (self.write_dir or 'diags') - file_prefix = (self.file_prefix or self.name) - self.diagnostic.file_prefix = os.path.join(write_dir, file_prefix) + self.set_write_dir() class LabFrameParticleDiagnostic(picmistandard.PICMI_LabFrameParticleDiagnostic): @@ -2033,3 +2027,127 @@ class LabFrameParticleDiagnostic(picmistandard.PICMI_LabFrameParticleDiagnostic) pywarpx.warpx.num_snapshots_lab = self.num_snapshots pywarpx.warpx.dt_snapshots_lab = self.dt_snapshots pywarpx.warpx.lab_data_directory = self.write_dir + + +class ReducedDiagnostic(picmistandard.base._ClassWithInit, WarpXDiagnosticBase): + """ + Sets up a reduced diagnostic in the simulation. + + See `Input Parameters <https://warpx.readthedocs.io/en/latest/usage/parameters.html#reduced-diagnostics>`_ + for more information. + + Parameters + ---------- + diag_type: string + The type of reduced diagnostic. See the link above for all the different + types of reduced diagnostics available. + + name: string + The name of this diagnostic which will also be the name of the data + file written to disk. + + period: integer + The simulation step interval at which to output this diagnostic. + + path: string + The file path in which the diagnostic file should be written. + + extension: string + The file extension used for the diagnostic output. + + separator: string + The separator between row values in the output file. + """ + + def __init__(self, diag_type, name=None, period=1, path=None, + extension=None, separator=None, **kw): + + self.name = name + self.type = diag_type + self.intervals = period + self.path = path + self.extension = extension + self.separator = separator + + self._species = kw.pop('species', None) + + # Now we need to handle all the specific inputs required for the + # different reduced diagnostic types. + # Note: only a limited number are presently supported. + + # The simple diagnostics do not require any additional arguments + self._simple_reduced_diagnostics = [ + 'ParticleEnergy', 'ParticleMomentum', 'FieldEnergy', + 'FieldMomentum', 'FieldMaximum', 'RhoMaximum', 'ParticleNumber', + 'LoadBalanceCosts', 'LoadBalanceEfficiency', + ] + # The species diagnostics require a species to be provided + self._species_reduced_diagnostics = [ + 'BeamRelevant', 'ParticleHistogram', 'ParticleExtrema', + ] + + if self.type in self._simple_reduced_diagnostics: + pass + elif self.type in self._species_reduced_diagnostics: + if self._species is None: + raise AttributeError( + f"{self.type} reduced diagnostic requires a species." + ) + if self.type == 'ParticleHistogram': + raise NotImplementedError( + f"{self.type} reduced diagnostic is not yet supported " + "in pywarpx." + ) + elif self.type == "FieldProbe": + kw = self._handle_field_probe(**kw) + else: + raise RuntimeError( + f"{self.type} reduced diagnostic is not yet supported " + "in pywarpx." + ) + + self.handle_init(kw) + + def _handle_field_probe(self, **kw): + """Utility function to grab required inputs for a field probe from kw""" + self.probe_geometry = kw.pop("probe_geometry") + self.x_probe = kw.pop("x_probe") + self.y_probe = kw.pop("y_probe") + self.z_probe = kw.pop("z_probe") + + self.interp_order = kw.pop("interp_order", None) + self.integrate = kw.pop("integrate", None) + self.do_moving_window_FP = kw.pop("do_moving_window_FP", None) + + if self.probe_geometry.lower() != 'point': + self.resolution = kw.pop("resolution") + + if self.probe_geometry.lower() == 'line': + self.x1_probe = kw.pop("x1_probe") + self.y1_probe = kw.pop("y1_probe") + self.z1_probe = kw.pop("z1_probe") + + if self.probe_geometry.lower() == 'plane': + self.detector_radius = kw.pop("detector_radius") + + self.target_normal_x = kw.pop("target_normal_x") + self.target_normal_y = kw.pop("target_normal_y") + self.target_normal_z = kw.pop("target_normal_z") + + self.target_up_x = kw.pop("target_up_x") + self.target_up_y = kw.pop("target_up_y") + self.target_up_z = kw.pop("target_up_z") + + return kw + + def initialize_inputs(self): + + self.add_diagnostic() + + for key in self.__dict__.keys(): + if not key.startswith('_') and key not in ['name', 'diagnostic']: + self.diagnostic.__setattr__(key, self.__dict__[key]) + + if self._species is not None: + diag = pywarpx.Bucket.Bucket(self.name + '.' + self._species.name) + self.diagnostic._species_dict[self._species.name] = diag |