implementing phase space plugin in picmi

lib/python/picongpu/picmi/__init__.py

@@ -8,6 +8,7 @@
 from .species import Species
 from .layout import PseudoRandomLayout
 from . import constants
+from .phase_space import PhaseSpace
 
 from .distribution import FoilDistribution, UniformDistribution, GaussianDistribution
 from .interaction import Interaction
@@ -38,6 +39,7 @@
     "Keldysh",
     "ThomasFermi",
     "Interaction",
+    "PhaseSpace",
 ]
 
 

lib/python/picongpu/picmi/debug.py

@@ -0,0 +1,5 @@
+import picmistandard
+# print(dir(picmistandard))  # Check if PICMI_PhaseSpace is listed
+
+print(picmistandard.__file__)
+# /home/afshar87/.local/lib/python3.10/site-packages/picmistandard

lib/python/picongpu/picmi/phase_space.py

@@ -0,0 +1,87 @@
+"""
+This file is part of PIConGPU.
+Copyright 2021-2024 PIConGPU contributors
+Authors: Masoud Afshari
+License: GPLv3+
+"""
+
+from ..pypicongpu import util
+from ..pypicongpu.output.phase_space import PhaseSpace
+from ..pypicongpu.species.species import Species as PyPIConGPUSpecies
+
+
+from .species import Species as PICMISpecies
+
+import picmistandard
+import typeguard
+
+
+@typeguard.typechecked
+class PhaseSpace(picmistandard.PICMI_PhaseSpace):
+    """PICMI object for Phase Space diagnostics"""
+
+    def __init__(
+        self,
+        species: PICMISpecies,
+        period: int,
+        spatial_coordinate: str,
+        momentum: str,
+        min_momentum: float,
+        max_momentum: float,
+        **kw,
+    ):
+        if period <= 0:
+            raise ValueError("Period must be > 0")
+        if min_momentum >= max_momentum:
+            raise ValueError("min_momentum must be less than max_momentum")
+
+        self.species = species
+        self.period = period
+        self.spatial_coordinate = spatial_coordinate
+        self.momentum = momentum
+        self.min_momentum = min_momentum
+        self.max_momentum = max_momentum
+
+        super().__init__(
+            species,
+            period,
+            spatial_coordinate,
+            momentum,
+            min_momentum,
+            max_momentum,
+            **kw,
+        )
+
+    def get_as_pypicongpu(
+        # to get the corresponding PyPIConGPUSpecies instance for the given PICMISpecies.
+        self,
+        dict_species_picmi_to_pypicongpu: dict[PICMISpecies, PyPIConGPUSpecies],
+    ) -> PhaseSpace:
+        # print(f"dict_species_picmi_to_pypicongpu keys: {list(dict_species_picmi_to_pypicongpu.keys())}")
+        # print(f"self.species: {self.species}")
+
+        util.unsupported("extra attributes", self.__dict__.keys())
+
+        if self.species not in dict_species_picmi_to_pypicongpu:
+            raise ValueError(f"Species {self.species} is not mapped in dict_species_picmi_to_pypicongpu!")
+
+        # 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).
+        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.")
+
+        # Print type before passing to PhaseSpace
+        print(f"DEBUG: Mapped species: {pypicongpu_species}, Type: {type(pypicongpu_species)}")
+
+        pypicongpu_phase_space = PhaseSpace(
+            species=pypicongpu_species,
+            period=self.period,
+            spatial_coordinate=self.spatial_coordinate,
+            momentum=self.momentum,
+            min_momentum=self.min_momentum,
+            max_momentum=self.max_momentum,
+        )
+
+        return pypicongpu_phase_space

lib/python/picongpu/picmi/simulation.py

@@ -10,6 +10,8 @@
 from .species import Species
 from .interaction.ionization import IonizationModel
 
+from picongpu.pypicongpu.species.initmanager import InitManager
+
 from . import constants
 from .grid import Cartesian3DGrid
 from .interaction import Interaction
@@ -340,7 +342,8 @@ def __fill_in_ionization_electrons(
                 pypicongpu_by_picmi_species, ionization_model_conversion_by_species
             )
 
-    def __get_init_manager(self) -> pypicongpu.species.InitManager:
+    # def __get_init_manager(self) -> pypicongpu.species.InitManager:
+    def __get_init_manager(self) -> tuple[InitManager, typing.Dict[Species, pypicongpu.species.Species]]:
         """
         create & fill an Initmanager
 
@@ -361,7 +364,8 @@ def __get_init_manager(self) -> pypicongpu.species.InitManager:
         self.__fill_in_ionization_electrons(pypicongpu_by_picmi_species, ionization_model_conversion_by_species)
 
         # init PyPIConGPU init manager
-        initmgr = pypicongpu.species.InitManager()
+        # initmgr = pypicongpu.species.InitManager()
+        initmgr = InitManager()  # This works because InitManager is imported
 
         for pypicongpu_species in pypicongpu_by_picmi_species.values():
             initmgr.all_species.append(pypicongpu_species)
@@ -373,7 +377,7 @@ def __get_init_manager(self) -> pypicongpu.species.InitManager:
         initmgr.all_operations += self.__get_operations_not_placed(pypicongpu_by_picmi_species)
         initmgr.all_operations += self.__get_operations_from_individual_species(pypicongpu_by_picmi_species)
 
-        return initmgr
+        return initmgr, pypicongpu_by_picmi_species
 
     def write_input_file(
         self, file_name: str, pypicongpu_simulation: typing.Optional[pypicongpu.simulation.Simulation] = None
@@ -455,7 +459,12 @@ def get_as_pypicongpu(self) -> pypicongpu.simulation.Simulation:
             # explictly disable laser (as required by pypicongpu)
             s.laser = None
 
-        s.init_manager = self.__get_init_manager()
+        s.init_manager, pypicongpu_by_picmi_species = self.__get_init_manager()
+
+        pypicongpu_particle_plugins = []
+        for entry in self.diagnostics:
+            pypicongpu_particle_plugins.append(entry.get_as_pypicongpu(pypicongpu_by_picmi_species))
+        s.plugins = pypicongpu_particle_plugins
 
         # set typical ppc if not set explicitly by user
         if self.picongpu_typical_ppc is None:

lib/python/picongpu/pypicongpu/__init__.py

@@ -4,6 +4,7 @@
 
 from .simulation import Simulation
 from .runner import Runner
+from .output.phase_space import PhaseSpace
 
 from . import laser
 from . import grid
@@ -25,6 +26,7 @@
     "util",
     "grid",
     "customuserinput",
+    "PhaseSpace",
 ]
 
 # note: put down here b/c linter complains if imports are not at top

lib/python/picongpu/pypicongpu/output/phase_space.py

@@ -0,0 +1,36 @@
+"""
+This file is part of PIConGPU.
+Copyright 2021-2024 PIConGPU contributors
+Authors: Masoud Afshari
+License: GPLv3+
+"""
+
+from .. import util
+from ..species import Species
+
+from .plugin import Plugin
+
+import typeguard
+import typing
+from typing import Literal
+
+
+@typeguard.typechecked
+class PhaseSpace(Plugin):
+    species = util.build_typesafe_property(Species)
+    period = util.build_typesafe_property(int)
+    spatial_coordinate = util.build_typesafe_property(Literal["x", "y", "z"])
+    momentum = util.build_typesafe_property(Literal["px", "py", "pz"])
+    min_momentum = util.build_typesafe_property(float)
+    max_momentum = util.build_typesafe_property(float)
+
+    def _get_serialized(self) -> typing.Dict:
+        """Return the serialized representation of the object."""
+        return {
+            "species": self.species.get_cxx_typename(),
+            "period": self.period,
+            "spatial_coordinate": self.spatial_coordinate,
+            "momentum": self.momentum,
+            "min_momentum": self.min_momentum,
+            "max_momentum": self.max_momentum,
+        }

lib/python/picongpu/pypicongpu/output/plugin.py

@@ -0,0 +1,18 @@
+"""
+This file is part of PIConGPU.
+Copyright 2021-2024 PIConGPU contributors
+Authors: Brian Edward Marre, Masoud Afshari
+License: GPLv3+
+"""
+
+from ..rendering import RenderedObject
+
+import typeguard
+
+
+@typeguard.typechecked
+class Plugin(RenderedObject):
+    """general interface for all plugins"""
+
+    def __init__(self):
+        raise NotImplementedError("abstract base class only")

lib/python/picongpu/pypicongpu/simulation.py

@@ -14,6 +14,7 @@
 from . import output
 from .rendering import RenderedObject
 from .customuserinput import InterfaceCustomUserInput
+from .output.plugin import Plugin
 
 import typing
 import typeguard
@@ -66,6 +67,8 @@ class Simulation(RenderedObject):
     moving_window = util.build_typesafe_property(typing.Optional[MovingWindow])
     """used moving Window, set to None to disable"""
 
+    plugins = util.build_typesafe_property(typing.Optional[Plugin])
+
     def __get_output_context(self) -> dict:
         """retrieve all output objects"""
         auto = output.Auto()