Merge pull request #5578 from chillenzer/species-refactoring

PICMI Species handling refactoring

Author: PrometheusPi

docs/source/usage/picmi/custom_template.rst

@@ -74,7 +74,7 @@ The output of the picongpu simulation in the default templates is configured in
           $USED_CHARGE_CONSERVATION_FLAGS
 
 
-          {{#species_initmanager.species}}
+          {{#species}}
               --{{{name}}}_macroParticlesCount.period {{{period}}}
 
               --{{{name}}}_energy.period {{{period}}}
@@ -92,7 +92,7 @@ The output of the picongpu simulation in the default templates is configured in
                   --{{{name}}}_png.slicePoint 0.5
                   --{{{name}}}_png.folder png_{{{name}}}_{{{axis}}}
               {{/png_axis}}
-          {{/species_initmanager.species}}
+          {{/species}}
 
       {{/output.auto}}
 
@@ -148,9 +148,9 @@ Instead of hard coding the output we might want to automatically generate one in
          --Cu_energyHistogram.minEnergy 0
          --Cu_energyHistogram.maxEnergy 256000
 
-          {{#species_initmanager.species}}
+          {{#species}}
               --{{{name}}}_macroParticlesCount.period 1
-          {{/species_initmanager.species}}
+          {{/species}}
       {{/output.auto}}
 
 Let's go in detail through the above example.
@@ -218,9 +218,9 @@ And of course it will be available as such in the rendering of the template whic
          --Cu_energyHistogram.minEnergy 0
          --Cu_energyHistogram.maxEnergy 256000
 
-          {{#species_initmanager.species}}
+          {{#species}}
               --{{{name}}}_macroParticlesCount.period 1
-          {{/species_initmanager.species}}
+          {{/species}}
       {{/output.auto}}
 
 .. warning::

lib/python/picongpu/picmi/copy_attributes.py

@@ -89,7 +89,11 @@ def copy_attributes(
     """
     assignments = {
         to_name: _value_generator(from_name)
-        for from_name, _ in inspect.getmembers(from_instance)
+        for from_name, _ in (
+            type(from_instance).model_fields.items()
+            if isinstance(from_instance, BaseModel)
+            else inspect.getmembers(from_instance)
+        )
         if from_name not in ignore
         and not from_name.startswith("_")
         and has_attribute(to, to_name := from_name.removeprefix(remove_prefix))

lib/python/picongpu/picmi/diagnostics/__init__.py

@@ -5,7 +5,6 @@
 License: GPLv3+
 """
 
-from .auto import Auto
 from .binning import Binning, BinningAxis, BinSpec
 from .phase_space import PhaseSpace
 from .energy_histogram import EnergyHistogram
@@ -19,7 +18,6 @@
 from .unit_dimension import UnitDimension
 
 __all__ = [
-    "Auto",
     "BackendConfig",
     "OpenPMDConfig",
     "Binning",

lib/python/picongpu/picmi/diagnostics/auto.py

@@ -8,15 +8,13 @@
 from pathlib import Path
 from ..copy_attributes import default_converts_to
 
-import typeguard
 
 from picongpu.picmi.diagnostics.backend_config import OpenPMDConfig
 
 from ...pypicongpu.output.binning import Binning as PyPIConGPUBinning
 from ...pypicongpu.output.binning import BinningAxis as PyPIConGPUBinningAxis
 from ...pypicongpu.output.binning import BinSpec as PyPIConGPUBinSpec
-from ...pypicongpu.species.species import Species as PyPIConGPUSpecies
-from ..species import Species as PICMISpecies
+from ..species import Species as Species
 from .particle_functor import ParticleFunctor as BinningFunctor
 from .timestepspec import TimeStepSpec
 
@@ -30,7 +28,6 @@ def __init__(self, kind, start, stop, nsteps):
         self.nsteps = nsteps
 
 
-@typeguard.typechecked
 class BinningAxis:
     def __init__(
         self,
@@ -53,14 +50,13 @@ def get_as_pypicongpu(self) -> PyPIConGPUBinningAxis:
         )
 
 
-@typeguard.typechecked
 class Binning:
     def __init__(
         self,
         name: str,
         deposition_functor: BinningFunctor,
         axes: list[BinningAxis],
-        species: PICMISpecies | list[PICMISpecies],
+        species: Species | list[Species],
         period: TimeStepSpec | None = None,
         openPMD: dict | None = None,
         openPMDExt: str | None = None,
@@ -70,7 +66,7 @@ def __init__(
         self.name = name
         self.deposition_functor = deposition_functor
         self.axes = axes
-        if isinstance(species, PICMISpecies):
+        if isinstance(species, Species):
             species = [species]
         self.species = species
         self.period = period or TimeStepSpec[:]
@@ -86,19 +82,14 @@ def result_path(self, prefix_path):
 
     def get_as_pypicongpu(
         self,
-        dict_species_picmi_to_pypicongpu: dict[PICMISpecies, PyPIConGPUSpecies],
         time_step_size,
         num_steps,
     ) -> PyPIConGPUBinning:
-        if len(not_found := [s for s in self.species if s not in dict_species_picmi_to_pypicongpu.keys()]) > 0:
-            raise ValueError(f"Species {not_found} are not known to Simulation")
-        pypic_species = list(map(dict_species_picmi_to_pypicongpu.get, self.species))
-
         return PyPIConGPUBinning(
             name=self.name,
             deposition_functor=self.deposition_functor.get_as_pypicongpu(),
             axes=list(map(BinningAxis.get_as_pypicongpu, self.axes)),
-            species=pypic_species,
+            species=[s.get_as_pypicongpu() for s in self.species],
             period=self.period.get_as_pypicongpu(time_step_size, num_steps),
             openPMD=self.openPMD,
             openPMDExt=self.openPMDExt,

lib/python/picongpu/picmi/diagnostics/binning.py

@@ -5,20 +5,20 @@
 License: GPLv3+
 """
 
-import typeguard
+from pydantic import BaseModel
+
+from picongpu.picmi.copy_attributes import default_converts_to
 
-from picongpu.picmi.diagnostics.util import diagnostic_converts_to
 
 from ...pypicongpu.output.energy_histogram import (
     EnergyHistogram as PyPIConGPUEnergyHistogram,
 )
-from ..species import Species as PICMISpecies
+from ..species import Species as Species
 from .timestepspec import TimeStepSpec
 
 
-@diagnostic_converts_to(PyPIConGPUEnergyHistogram)
-@typeguard.typechecked
-class EnergyHistogram:
+@default_converts_to(PyPIConGPUEnergyHistogram)
+class EnergyHistogram(BaseModel):
     """
     Specifies the parameters for the output of Energy Histogram of species such as electrons.
 
@@ -53,24 +53,17 @@ class EnergyHistogram:
         Optional name for the energy histogram plugin.
     """
 
-    def check(self, dict_species_picmi_to_pypicongpu, *args, **kwargs):
+    def check(self, *args, **kwargs):
         if self.min_energy >= self.max_energy:
             raise ValueError("min_energy must be less than max_energy")
         if self.bin_count <= 0:
             raise ValueError("bin_count must be > 0")
-        if self.species not in dict_species_picmi_to_pypicongpu.keys():
-            raise ValueError(f"Species {self.species} is not known to Simulation")
-
-    def __init__(
-        self,
-        species: PICMISpecies,
-        period: TimeStepSpec,
-        bin_count: int,
-        min_energy: float,
-        max_energy: float,
-    ):
-        self.species = species
-        self.period = period
-        self.bin_count = bin_count
-        self.min_energy = min_energy
-        self.max_energy = max_energy
+
+    species: Species
+    period: TimeStepSpec
+    bin_count: int
+    min_energy: float
+    max_energy: float
+
+    class Config:
+        arbitrary_types_allowed = True

lib/python/picongpu/picmi/diagnostics/energy_histogram.py

@@ -5,21 +5,19 @@
 License: GPLv3+
 """
 
-import typeguard
+from pydantic import BaseModel
 
-from picongpu.picmi.diagnostics.util import diagnostic_converts_to
+from picongpu.picmi.copy_attributes import default_converts_to
 
 from ...pypicongpu.output.macro_particle_count import (
     MacroParticleCount as PyPIConGPUMacroParticleCount,
 )
-from ...pypicongpu.species.species import Species as PyPIConGPUSpecies
-from ..species import Species as PICMISpecies
+from ..species import Species as Species
 from .timestepspec import TimeStepSpec
 
 
-@diagnostic_converts_to(PyPIConGPUMacroParticleCount)
-@typeguard.typechecked
-class MacroParticleCount:
+@default_converts_to(PyPIConGPUMacroParticleCount)
+class MacroParticleCount(BaseModel):
     """
     Specifies the parameters for counting the total number of macro particles of a given species.
 
@@ -39,13 +37,8 @@ class MacroParticleCount:
         Optional name for the macro particle count plugin.
     """
 
-    def __init__(self, species: PICMISpecies, period: TimeStepSpec):
-        self.species = species
-        self.period = period
+    species: Species
+    period: TimeStepSpec
 
-    def check(self, dict_species_picmi_to_pypicongpu: dict[PICMISpecies, PyPIConGPUSpecies], *args, **kwargs):
-        if self.species not in dict_species_picmi_to_pypicongpu.keys():
-            raise ValueError(f"Species {self.species} is not known to Simulation")
-        pypicongpu_species = dict_species_picmi_to_pypicongpu.get(self.species)
-        if pypicongpu_species is None:
-            raise ValueError(f"Species {self.species} is not mapped to a PyPIConGPUSpecies.")
+    class Config:
+        arbitrary_types_allowed = True

lib/python/picongpu/picmi/diagnostics/macro_particle_count.py

@@ -10,7 +10,7 @@
 
 from pydantic import BaseModel
 
-from picongpu.picmi.species import Species
+from picongpu.picmi.species import Species as Species
 
 from .backend_config import BackendConfig, OpenPMDConfig
 from .timestepspec import TimeStepSpec

lib/python/picongpu/picmi/diagnostics/particle_dump.py

@@ -7,18 +7,17 @@
 
 from typing import Literal
 
-import typeguard
+from pydantic import BaseModel
 
-from picongpu.picmi.diagnostics.util import diagnostic_converts_to
+from picongpu.picmi.copy_attributes import default_converts_to
 
 from ...pypicongpu.output.phase_space import PhaseSpace as PyPIConGPUPhaseSpace
-from ..species import Species as PICMISpecies
+from ..species import Species as Species
 from .timestepspec import TimeStepSpec
 
 
-@diagnostic_converts_to(PyPIConGPUPhaseSpace)
-@typeguard.typechecked
-class PhaseSpace:
+@default_converts_to(PyPIConGPUPhaseSpace)
+class PhaseSpace(BaseModel):
     """
     Specifies the parameters for the output of Phase Space of species such as electrons.
 
@@ -52,30 +51,16 @@ class PhaseSpace:
         Optional name for the phase-space plugin.
     """
 
-    def check(self, dict_species_picmi_to_pypicongpu, *args, **kwargs):
+    species: Species
+    period: TimeStepSpec
+    spatial_coordinate: Literal["x", "y", "z"]
+    momentum_coordinate: Literal["px", "py", "pz"]
+    min_momentum: float
+    max_momentum: float
+
+    def check(self, *args, **kwargs):
         if self.min_momentum >= self.max_momentum:
             raise ValueError("min_momentum must be less than max_momentum")
 
-        if self.species not in dict_species_picmi_to_pypicongpu.keys():
-            raise ValueError(f"Species {self.species} is not known to Simulation")
-
-        # checks if PICMISpecies instance exists in the dictionary. If yes, it returns the corresponding PyPIConGPUSpecies instance.
-        # self.species refers to the species attribute of the class  PhaseSpace(picmistandard.PICMI_PhaseSpace).
-        if dict_species_picmi_to_pypicongpu.get(self.species) is None:
-            raise ValueError(f"Species {self.species} is not mapped to a PyPIConGPUSpecies.")
-
-    def __init__(
-        self,
-        species: PICMISpecies,
-        period: TimeStepSpec,
-        spatial_coordinate: Literal["x", "y", "z"],
-        momentum_coordinate: Literal["px", "py", "pz"],
-        min_momentum: float,
-        max_momentum: float,
-    ):
-        self.species = species
-        self.period = period
-        self.spatial_coordinate = spatial_coordinate
-        self.momentum_coordinate = momentum_coordinate
-        self.min_momentum = min_momentum
-        self.max_momentum = max_momentum
+    class Config:
+        arbitrary_types_allowed = True