openpmd.py plugin for PICMI

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

@@ -11,7 +11,10 @@
 from .macro_particle_count import MacroParticleCount
 from .png import Png
 from .timestepspec import TimeStepSpec
+from .rangespec import RangeSpec
 from .checkpoint import Checkpoint
+from .openpmd import OpenPMD
+from .openpmd_sources.source_base import SourceBase
 
 __all__ = [
     "Auto",
@@ -20,5 +23,8 @@
     "MacroParticleCount",
     "Png",
     "TimeStepSpec",
+    "RangeSpec",
     "Checkpoint",
+    "OpenPMD",
+    "SourceBase",
 ]

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

@@ -1,7 +1,7 @@
 """
 This file is part of PIConGPU.
 Copyright 2025 PIConGPU contributors
-Authors: Pawel Ordyna
+Authors: Pawel Ordyna, Masoud Afshari
 License: GPLv3+
 """
 
@@ -42,6 +42,7 @@ def get_as_pypicongpu(
         dict_species_picmi_to_pypicongpu: dict[PICMISpecies, PyPIConGPUSpecies],
         time_step_size,
         num_steps,
+        simulation_box=None,  # Added to match OpenPMD signature, not used
     ) -> PyPIConGPUAuto:
         self.check()
         pypicongpu_auto = PyPIConGPUAuto()

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

@@ -108,6 +108,7 @@ def get_as_pypicongpu(
         pypicongpu_by_picmi_species: Dict,
         time_step_size: float,
         num_steps: int,
+        simulation_box=None,  # Added to match OpenPMD signature, not used
     ) -> PyPIConGPUCheckpoint:
         self.check()
 

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

@@ -79,6 +79,7 @@ def get_as_pypicongpu(
         dict_species_picmi_to_pypicongpu: dict[PICMISpecies, PyPIConGPUSpecies],
         time_step_size,
         num_steps,
+        simulation_box=None,  # Added to match OpenPMD signature, not used
     ) -> PyPIConGPUEnergyHistogram:
         self.check()
 

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

@@ -49,6 +49,7 @@ def get_as_pypicongpu(
         dict_species_picmi_to_pypicongpu: dict[PICMISpecies, PyPIConGPUSpecies],
         time_step_size,
         num_steps,
+        simulation_box=None,  # Added to match OpenPMD signature, not used
     ) -> PyPIConGPUMacroParticleCount:
         self.check()
 

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

@@ -0,0 +1,109 @@
+"""
+This file is part of PIConGPU.
+Copyright 2025-2025 PIConGPU contributors
+Authors: Masoud Afshari
+License: GPLv3+
+"""
+
+from ...pypicongpu.output.openpmd import OpenPMD as PyPIConGPUOpenPMD
+from ...pypicongpu.output.openpmd_sources.source_base import SourceBase as PyPIConGPUSource
+from .timestepspec import TimeStepSpec
+from .rangespec import RangeSpec
+from .openpmd_sources.source_base import SourceBase
+
+import typeguard
+from typing import Optional, Dict, Union, List, Literal, Tuple
+
+
+@typeguard.typechecked
+class OpenPMD:
+    """
+    openPMD diagnostic output
+
+    This diagnostic writes simulation data (base fields, derived fields and/or particles) to disk using the openPMD
+    standard, with configurable periods, data sources and backend settings.
+
+    @param period specification of the time steps for data output, outputs will always be written at the end of a PIC time step.
+    @param source list of data source objects to include in the dump (e.g., [ChargeDensity(filter="all")]),
+        Setting this to None will cause an empty dump
+    @param range contiguous range of cells to dump the base- and derived field for, specified as a RangeSpec object
+        or a string in the format "begin:end" (1D), "begin:end,begin:end" (2D), or "begin:end,begin:end,begin:end" (3D).
+        Example: "0:10,5:15,2:8" specifies cells 0 to 10 (x), 5 to 15 (y), 2 to 8 (z).
+        Notes: Values are clipped to the simulation box. Begin and/or end may be omitted (":") to indicate the full extent
+        of the dimension. Negative indices are supported (e.g., "-5:-1" for last 5 cells). The default ":,:,:," (3D),
+        ":,:" (2D), or ":" (1D) includes all cells in the simulation box.
+    @param file relative or absolute file path prefix for openPMD output files. Relative paths are interpreted as relative to the simulation output directory, the default value None indicates the PIC code's default.
+    @param ext file extension controlling the openPMD backend, options are "bp" (default backend ADIOS2), "h5" (HDF5), "sst" (ADIOS2/SST for streaming).
+    @param infix filename infix for the iteration layout (e.g., "_%06T"), use "NULL" for the group-based layout, ext="sst" requires infix="NULL".
+    @param json openPMD backend configuration as a JSON string, dictionary, or filename (filename must be prepended with "@").
+    @param json_restart backend-specific parameters for restarting, as a JSON string, dictionary, or filename (filenames must be prepended with "@").
+    @param data_preparation_strategy strategy for particle data preparation, options: "doubleBuffer" or "adios" (ADIOS2-based), "mappedMemory" or "hdf5" (HDF5-based), the default value None indicates the PIC code default
+    @param toml path to a TOML file for openPMD configuration. Replaces the JSON or keyword configuration.
+    @param particle_io_chunk_size size of particle data chunks used in writing (in MiB), reduces host memory footprint for certain backends, default "None" indicates the PIC code default.
+    @param file_writing file writing mode for writing, options: "create" (new files), "append" (for checkpoint-restart workflows).
+    """
+
+    def check(self):
+        """
+        Validate the provided parameters.
+        """
+        if self.particle_io_chunk_size is not None and self.particle_io_chunk_size < 1:
+            raise ValueError("particle_io_chunk_size (in MiB) must be positive")
+        if self.ext == "sst" and self.infix is not None and self.infix != "NULL":
+            raise ValueError("infix must be 'NULL' when ext is 'sst'")
+        if self.source is not None and not all(isinstance(s, SourceBase) for s in self.source):
+            raise ValueError("source must be a list of SourceBase objects")
+
+    def __init__(
+        self,
+        period: TimeStepSpec,
+        source: Optional[List[SourceBase]] = None,
+        range: Optional[Union[str, RangeSpec]] = ":,:,:",
+        file: Optional[str] = None,
+        ext: Optional[Literal["bp", "h5", "sst"]] = "bp",
+        infix: Optional[str] = "NULL",
+        json: Optional[Union[str, Dict]] = None,
+        json_restart: Optional[Union[str, Dict]] = None,
+        data_preparation_strategy: Optional[Literal["doubleBuffer", "adios", "mappedMemory", "hdf5"]] = None,
+        toml: Optional[str] = None,
+        particle_io_chunk_size: Optional[int] = None,
+        file_writing: Optional[Literal["create", "append"]] = "create",
+    ):
+        self.period = period
+        self.source = source
+        self.range = RangeSpec(range) if isinstance(range, str) else range
+        self.file = file
+        self.ext = ext
+        self.infix = infix
+        self.json = json if json is not None else {}
+        self.json_restart = json_restart if json_restart is not None else {}
+        self.data_preparation_strategy = data_preparation_strategy
+        self.toml = toml
+        self.particle_io_chunk_size = particle_io_chunk_size
+        self.file_writing = file_writing
+
+        self.check()
+
+    def get_as_pypicongpu(
+        self,
+        pypicongpu_by_picmi_species: Dict,
+        time_step_size: float,
+        num_steps: int,
+        simulation_box: Tuple[int, ...],
+    ) -> PyPIConGPUOpenPMD:
+        self.check()
+        pypicongpu_openpmd = PyPIConGPUOpenPMD(
+            period=self.period.get_as_pypicongpu(time_step_size, num_steps),
+            source=PyPIConGPUSource([s.get_as_pypicongpu() for s in self.source]) if self.source is not None else None,
+            range=self.range.get_as_pypicongpu(simulation_box),
+            file=self.file,
+            ext=self.ext,
+            infix=self.infix,
+            json=self.json,
+            json_restart=self.json_restart,
+            data_preparation_strategy=self.data_preparation_strategy,
+            toml=self.toml,
+            particle_io_chunk_size=self.particle_io_chunk_size,
+            file_writing=self.file_writing,
+        )
+        return pypicongpu_openpmd

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

@@ -0,0 +1,35 @@
+from .auto import Auto
+from .bound_electron_density import BoundElectronDensity
+from .charge_density import ChargeDensity
+from .counter import Counter
+from .density import Density
+from .derived_attributes import DerivedAttributes
+from .energy import Energy
+from .energy_density import EnergyDensity
+from .energy_density_cutoff import EnergyDensityCutoff
+from .larmor_power import LarmorPower
+from .macro_counter import MacroCounter
+from .mid_current_density_component import MidCurrentDensityComponent
+from .momentum import Momentum
+from .momentum_density import MomentumDensity
+from .weighted_velocity import WeightedVelocity
+from .source_base import SourceBase
+
+__all__ = [
+    "Auto",
+    "BoundElectronDensity",
+    "ChargeDensity",
+    "Counter",
+    "Density",
+    "DerivedAttributes",
+    "Energy",
+    "EnergyDensity",
+    "EnergyDensityCutoff",
+    "LarmorPower",
+    "MacroCounter",
+    "MidCurrentDensityComponent",
+    "Momentum",
+    "MomentumDensity",
+    "WeightedVelocity",
+    "SourceBase",
+]

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

@@ -0,0 +1,49 @@
+"""
+This file is part of PIConGPU.
+Copyright 2025-2025 PIConGPU contributors
+Authors: Masoud Afshari
+License: GPLv3+
+"""
+
+from .source_base import SourceBase
+from ....pypicongpu.output.openpmd_sources import Auto as PyPIConGPUAuto
+import typeguard
+import typing
+
+
+@typeguard.typechecked
+class Auto(SourceBase):
+    """
+    Default data source for openPMD output
+
+    This class provides a convenient way to dump default simulation data (e.g., all
+    particle species and fields) using the openPMD standard, with defaults determined
+    by the PIC code in particle-in-cell simulations.
+
+    @param filter Name of a filter to select data contributing to the source.
+        Default: None (PIC code-dependent).
+    """
+
+    # filter = util.build_typesafe_property(typing.Optional[str])
+
+    def __init__(self, filter: typing.Optional[str] = None):
+        self.filter = filter
+        self.check()
+
+    def check(self) -> None:
+        """
+        Validate the filter parameter.
+
+        @throw ValueError If the filter is not a string or None.
+        """
+        if self.filter is not None and not isinstance(self.filter, str):
+            raise ValueError(f"Filter must be a string or None, got {type(self.filter)}")
+
+    def get_as_pypicongpu(self) -> PyPIConGPUAuto:
+        """
+        Convert this Auto source to a PyPIConGPU Auto source.
+
+        @return A PyPIConGPU Auto instance with the same filter.
+        """
+        self.check()
+        return PyPIConGPUAuto(filter=self.filter)

lib/python/picongpu/picmi/diagnostics/openpmd_sources/bound_electron_density.py

@@ -0,0 +1,65 @@
+"""
+This file is part of PIConGPU.
+Copyright 2025-2025 PIConGPU contributors
+Authors: Masoud Afshari
+License: GPLv3+
+"""
+
+from .source_base import SourceBase
+from ....pypicongpu.output.openpmd_sources import BoundElectronDensity as PyPIConGPUBoundElectronDensity
+from ...species import Species as PICMISpecies
+import typeguard
+import typing
+
+
+@typeguard.typechecked
+class BoundElectronDensity(SourceBase):
+    """
+    Bound electron density data source for openPMD output
+
+    This source calculates the density of bound electrons from a specified particle species,
+    optionally filtered by a selection criterion, for particle-in-cell simulations.
+
+    @param species Particle species contributing to the bound electron density (e.g., ions).
+    @param filter Name of a filter to select particles contributing to the source.
+        Default: "all" (includes all particles of the specified species).
+    """
+
+    def __init__(self, species: PICMISpecies, filter: str = "all"):
+        self.species = species
+        self.filter = filter
+        self.check()
+
+    def check(self) -> None:
+        """
+        Validate the parameters.
+
+        @throw ValueError If filter is not a string or species is not a PICMISpecies.
+        """
+        if not isinstance(self.filter, str):
+            raise ValueError(f"Filter must be a string, got {type(self.filter)}")
+        if not isinstance(self.species, PICMISpecies):
+            raise ValueError(f"Species must be a PICMISpecies, got {type(self.species)}")
+
+    def get_as_pypicongpu(
+        self,
+        dict_species_picmi_to_pypicongpu: dict[PICMISpecies, typing.Any],
+    ) -> PyPIConGPUBoundElectronDensity:
+        """
+        Convert this BoundElectronDensity source to a PyPIConGPU BoundElectronDensity source.
+
+        @param dict_species_picmi_to_pypicongpu Mapping of PICMI species to PyPIConGPU species.
+        @return A PyPIConGPU BoundElectronDensity instance with the same filter and species.
+        @throw ValueError If the species is not known to the simulation or not mapped to a PyPIConGPUSpecies.
+        """
+        self.check()
+
+        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.")
+
+        return PyPIConGPUBoundElectronDensity(filter=self.filter, species=pypicongpu_species)

lib/python/picongpu/picmi/diagnostics/openpmd_sources/charge_density.py

@@ -0,0 +1,65 @@
+"""
+This file is part of PIConGPU.
+Copyright 2025-2025 PIConGPU contributors
+Authors: Masoud Afshari
+License: GPLv3+
+"""
+
+from .source_base import SourceBase
+from ....pypicongpu.output.openpmd_sources import ChargeDensity as PyPIConGPUChargeDensity
+from ...species import Species as PICMISpecies
+import typeguard
+import typing
+
+
+@typeguard.typechecked
+class ChargeDensity(SourceBase):
+    """
+    Charge density data source for openPMD output
+
+    This source calculates the charge density from a specified particle species, optionally
+    filtered by a selection criterion, for particle-in-cell simulations.
+
+    @param species Particle species contributing to the charge density (e.g., electrons, protons).
+    @param filter Name of a filter to select particles contributing to the source.
+        Default: "all" (includes all particles of the specified species).
+    """
+
+    def __init__(self, species: PICMISpecies, filter: str = "all"):
+        self.species = species
+        self.filter = filter
+        self.check()
+
+    def check(self) -> None:
+        """
+        Validate the parameters.
+
+        @throw ValueError If filter is not a string or species is not a PICMISpecies.
+        """
+        if not isinstance(self.filter, str):
+            raise ValueError(f"Filter must be a string, got {type(self.filter)}")
+        if not isinstance(self.species, PICMISpecies):
+            raise ValueError(f"Species must be a PICMISpecies, got {type(self.species)}")
+
+    def get_as_pypicongpu(
+        self,
+        dict_species_picmi_to_pypicongpu: dict[PICMISpecies, typing.Any],
+    ) -> PyPIConGPUChargeDensity:
+        """
+        Convert this ChargeDensity source to a PyPIConGPU ChargeDensity source.
+
+        @param dict_species_picmi_to_pypicongpu Mapping of PICMI species to PyPIConGPU species.
+        @return A PyPIConGPU ChargeDensity instance with the same filter and species.
+        @throw ValueError If the species is not known to the simulation or not mapped to a PyPIConGPUSpecies.
+        """
+        self.check()
+
+        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.")
+
+        return PyPIConGPUChargeDensity(filter=self.filter, species=pypicongpu_species)