Ci test for picmi pypicongpu

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

@@ -12,7 +12,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 import SourceBase
 
 __all__ = [
     "Auto",
@@ -22,5 +25,8 @@
     "MacroParticleCount",
     "Png",
     "TimeStepSpec",
+    "RangeSpec",
     "Checkpoint",
+    "OpenPMD",
+    "SourceBase",
 ]

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

@@ -1,18 +1,16 @@
 """
 This file is part of PIConGPU.
 Copyright 2025 PIConGPU contributors
-Authors: Pawel Ordyna
+Authors: Pawel Ordyna, Masoud Afshari
 License: GPLv3+
 """
 
-
 from ...pypicongpu.output.auto import Auto as PyPIConGPUAuto
 from ...pypicongpu.species.species import Species as PyPIConGPUSpecies
-
 from ..species import Species as PICMISpecies
 from .timestepspec import TimeStepSpec
-
 import typeguard
+import warnings
 
 
 @typeguard.typechecked
@@ -22,28 +20,39 @@ class Auto:
 
     Parameters
     ----------
-    period: int
-        Number of simulation steps between consecutive outputs.
+    period: int or TimeStepSpec
+        Number of simulation steps between consecutive outputs (e.g., 10 for every 10 steps).
+        Use 0 to disable output.
+        Alternatively, a TimeStepSpec can be provided for PIConGPU-specific step selection
+        (e.g., TimeStepSpec[5, 10], TimeStepSpec[-10:]).
         Unit: steps (simulation time steps).
     """
 
-    period: TimeStepSpec
-    """Number of simulation steps between consecutive outputs. Unit: steps (simulation time steps)."""
-
-    def __init__(self, period: TimeStepSpec) -> None:
-        self.period = period
+    def __init__(self, period: int | TimeStepSpec) -> None:
+        if not isinstance(period, (int, TimeStepSpec)):
+            raise TypeError("period must be an integer or TimeStepSpec")
+        if isinstance(period, int):
+            if period < 0:
+                raise ValueError("period must be non-negative")
+            self.period = TimeStepSpec[::period]("steps") if period > 0 else TimeStepSpec()("steps")
+        else:
+            self.period = period
+            if self.period.unit_system is None:
+                self.period = self.period("steps")
 
     def check(self):
-        pass
+        if not self.period.get_as_pypicongpu(1.0, 100).get_rendering_context().get("specs", []):
+            warnings.warn("Auto output is disabled because period is set to 0 or an empty TimeStepSpec")
 
     def get_as_pypicongpu(
         self,
-        # not used here, but needed for the interface
         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()
         pypicongpu_auto.period = self.period.get_as_pypicongpu(time_step_size, num_steps)
+
         return pypicongpu_auto

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

@@ -7,9 +7,9 @@
 
 from ...pypicongpu.output.checkpoint import Checkpoint as PyPIConGPUCheckpoint
 from .timestepspec import TimeStepSpec
-
 import typeguard
-from typing import Optional, Dict
+import warnings
+from typing import Optional, Union, Dict
 
 
 @typeguard.typechecked
@@ -24,62 +24,42 @@ class Checkpoint:
 
     Parameters
     ----------
-    period: TimeStepSpec, optional
-        Specify on which time steps to create checkpoints.
-        Unit: steps (simulation time steps). Required if timePeriod is not provided.
-
+    period: int or TimeStepSpec, optional
+        Number of simulation steps between consecutive checkpoints (e.g., 10 for every 10 steps).
+        Use 0 to disable checkpointing.
+        Alternatively, a TimeStepSpec can be provided for specific step selection
+        (e.g., TimeStepSpec([5, 10]), TimeStepSpec([slice(-10, None, 1)])).
+        Unit: steps or seconds (via TimeStepSpec unit).
     timePeriod: int, optional
-        Specify the interval in minutes for creating checkpoints.
-        Unit: minutes (must be a non-negative integer). Required if period is not provided.
-
+        Time interval between checkpoints in simulation time steps.
+        Use 0 or None to disable time-based checkpointing.
     directory: str, optional
-        Directory inside simOutput for writing checkpoints (default: "checkpoints").
-
+         Directory inside simOutput for writing checkpoints (default: "checkpoints").
     file: str, optional
-        Relative or absolute fileset prefix for checkpoint files.
-
+        Relative or absolute fileset prefix for checkpoint files. Default: None.
     restart: bool, optional
-        If True, restart simulation from the latest checkpoint.
-
+        Enable restarting from checkpoints. Default: None.
     tryRestart: bool, optional
-        If True, restart from the latest checkpoint if available, else start from scratch.
-
+        If True, restart from the latest checkpoint if available, else start from scratch. Default: None.
     restartStep: int, optional
-        Specific checkpoint step to restart from.
-
+        Specific step to restart from. Default: None.
     restartDirectory: str, optional
-        Directory inside simOutput containing checkpoints for restart (default: "checkpoints").
-
+        Directory inside simOutput containing checkpoints for restart.
     restartFile: str, optional
-        Relative or absolute fileset prefix for reading checkpoints.
-
+        Specific file to restart from. Default: None.
     restartChunkSize: int, optional
-        Number of particles processed in one kernel call during restart.
-
+        Chunk size for reading restart data. Default: None.
     restartLoop: int, optional
-        Number of times to restart the simulation after it finishes.
-
-    openPMD: Dict, optional
-        Dictionary of openPMD-specific settings (e.g., ext, json, infix).
+        Number of restart loops. Default: None.
+    openPMD: dict, optional
+        Configuration for openPMD output (e.g., {"ext": "h5"}). Default: None.
     """
 
-    def check(self):
-        if self.period is None and self.timePeriod is None:
-            raise ValueError("At least one of period or timePeriod must be provided")
-        if self.timePeriod is not None and self.timePeriod < 0:
-            raise ValueError("timePeriod must be a non-negative integer")
-        if self.restartStep is not None and self.restartStep < 0:
-            raise ValueError("restartStep must be non-negative")
-        if self.restartChunkSize is not None and self.restartChunkSize < 1:
-            raise ValueError("restartChunkSize must be positive")
-        if self.restartLoop is not None and self.restartLoop < 0:
-            raise ValueError("restartLoop must be non-negative")
-
     def __init__(
         self,
-        period: Optional[TimeStepSpec] = None,
+        period: Optional[Union[int, TimeStepSpec]] = None,
         timePeriod: Optional[int] = None,
-        directory: Optional[str] = None,
+        directory: Optional[str] = "checkpoints",
         file: Optional[str] = None,
         restart: Optional[bool] = None,
         tryRestart: Optional[bool] = None,
@@ -90,7 +70,16 @@ def __init__(
         restartLoop: Optional[int] = None,
         openPMD: Optional[Dict] = None,
     ):
-        self.period = period
+        if period is None and timePeriod is None:
+            raise ValueError("At least one of period or timePeriod must be provided to enable checkpointing")
+        if period is not None and not isinstance(period, (int, TimeStepSpec)):
+            raise TypeError("period must be an integer or TimeStepSpec")
+        if isinstance(period, int):
+            if period < 0:
+                raise ValueError("period must be non-negative")
+            self.period = TimeStepSpec([slice(None, None, period)] if period > 0 else [])("steps")
+        else:
+            self.period = period if period is not None else TimeStepSpec([])("steps")
         self.timePeriod = timePeriod
         self.directory = directory
         self.file = file
@@ -102,19 +91,32 @@ def __init__(
         self.restartChunkSize = restartChunkSize
         self.restartLoop = restartLoop
         self.openPMD = openPMD
+        self.check()
+
+    def check(self):
+        if self.timePeriod is not None and (not isinstance(self.timePeriod, int) or self.timePeriod < 0):
+            raise ValueError("timePeriod must be a non-negative integer")
+        if self.restartStep is not None and self.restartStep < 0:
+            raise ValueError("restartStep must be non-negative")
+        if self.restartChunkSize is not None and self.restartChunkSize <= 0:
+            raise ValueError("restartChunkSize must be positive")
+        if self.restartLoop is not None and self.restartLoop < 0:
+            raise ValueError("restartLoop must be non-negative")
+        if not self.period.specs and (self.timePeriod is None or self.timePeriod == 0):
+            warnings.warn(
+                "Checkpoint is disabled because period is set to 0 or an empty TimeStepSpec and timePeriod is None or 0"
+            )
 
     def get_as_pypicongpu(
         self,
         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()
-
         pypicongpu_checkpoint = PyPIConGPUCheckpoint()
-        pypicongpu_checkpoint.period = (
-            self.period.get_as_pypicongpu(time_step_size, num_steps) if self.period is not None else None
-        )
+        pypicongpu_checkpoint.period = self.period.get_as_pypicongpu(time_step_size, num_steps) if self.period else None
         pypicongpu_checkpoint.timePeriod = self.timePeriod
         pypicongpu_checkpoint.directory = self.directory
         pypicongpu_checkpoint.file = self.file
@@ -126,5 +128,4 @@ def get_as_pypicongpu(
         pypicongpu_checkpoint.restartChunkSize = self.restartChunkSize
         pypicongpu_checkpoint.restartLoop = self.restartLoop
         pypicongpu_checkpoint.openPMD = self.openPMD
-
         return pypicongpu_checkpoint

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

@@ -1,6 +1,6 @@
 """
 This file is part of PIConGPU.
-Copyright 2021-2024 PIConGPU contributors
+Copyright 2021-2025 PIConGPU contributors
 Authors: Masoud Afshari
 License: GPLv3+
 """
@@ -10,11 +10,9 @@
     EnergyHistogram as PyPIConGPUEnergyHistogram,
 )
 from ...pypicongpu.species.species import Species as PyPIConGPUSpecies
-
-
 from ..species import Species as PICMISpecies
-
 import typeguard
+from typing import Union
 
 
 @typeguard.typechecked
@@ -27,75 +25,71 @@ class EnergyHistogram:
 
     Parameters
     ----------
-    species: string
-        Name of the particle species to track (e.g., "electron", "proton").
-
-    period: int
-        Number of simulation steps between consecutive outputs.
-        If set to a non-zero value, the energy histogram of all electrons is computed.
-        By default, the value is 0 and no histogram for the electrons is computed.
-        Unit: steps (simulation time steps).
-
+    species: PICMISpecies
+        Particle species to count (e.g., an instance with name="electron" or "proton").
+    period: int or TimeStepSpec
+        Number of simulation steps between consecutive outputs (e.g., 10 for every 10 steps).
+        Use 0 to disable output. Alternatively, a TimeStepSpec can be provided for
+        PyPIConGPU-specific step selection (e.g., TimeStepSpec([slice(0, None, 10)])).
     bin_count: int
-        Number of bins for the energy histogram.
-
+        Number of bins for the energy histogram. Must be positive.
     min_energy: float
         Minimum value for the energy histogram range.
         Unit: keV
-        Default is 0, meaning 0 keV.
-
     max_energy: float
-        Maximum value for the energy histogram range.
+        Maximum value for the energy histogram range. Must be greater than min_energy.
         Unit: keV
-        There is no default value.
-
-    name: string, optional
-        Optional name for the energy histogram plugin.
     """
 
-    def check(self):
-        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")
-
     def __init__(
         self,
         species: PICMISpecies,
-        period: TimeStepSpec,
+        period: Union[int, TimeStepSpec],
         bin_count: int,
         min_energy: float,
         max_energy: float,
     ):
+        if isinstance(period, int):
+            if period < 0:
+                raise ValueError("period must be non-negative")
+            self.period = TimeStepSpec([slice(None, None, period)]) if period > 0 else TimeStepSpec()
+        else:
+            self.period = period
         self.species = species
-        self.period = period
         self.bin_count = bin_count
         self.min_energy = min_energy
         self.max_energy = max_energy
 
+    def check(self):
+        if not isinstance(self.species, PICMISpecies):
+            raise TypeError("species must be a PICMISpecies")
+        if not isinstance(self.species.name, str) or not self.species.name:
+            raise TypeError("species must have a non-empty name")
+        if not isinstance(self.period, TimeStepSpec):
+            raise TypeError("period must be a TimeStepSpec")
+        if self.bin_count <= 0:
+            raise ValueError("bin_count must be > 0")
+        if self.min_energy >= self.max_energy:
+            raise ValueError("min_energy must be less than max_energy")
+
     def get_as_pypicongpu(
-        # to get the corresponding PyPIConGPUSpecies instance for the given PICMISpecies.
         self,
         dict_species_picmi_to_pypicongpu: dict[PICMISpecies, PyPIConGPUSpecies],
-        time_step_size,
-        num_steps,
+        time_step_size: float,
+        num_steps: int,
+        simulation_box=None,  # Added to match OpenPMD signature, not used
     ) -> PyPIConGPUEnergyHistogram:
         self.check()
-
-        if self.species not in dict_species_picmi_to_pypicongpu.keys():
+        if self.species not in dict_species_picmi_to_pypicongpu:
             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.
-        pypicongpu_species = dict_species_picmi_to_pypicongpu.get(self.species)
-
+        pypicongpu_species = dict_species_picmi_to_pypicongpu[self.species]
         if pypicongpu_species is None:
             raise ValueError(f"Species {self.species} is not mapped to a PyPIConGPUSpecies.")
-
-        pypicongpu_energy_histogram = PyPIConGPUEnergyHistogram()
-        pypicongpu_energy_histogram.species = pypicongpu_species
-        pypicongpu_energy_histogram.period = self.period.get_as_pypicongpu(time_step_size, num_steps)
-        pypicongpu_energy_histogram.bin_count = self.bin_count
-        pypicongpu_energy_histogram.min_energy = self.min_energy
-        pypicongpu_energy_histogram.max_energy = self.max_energy
-
+        pypicongpu_energy_histogram = PyPIConGPUEnergyHistogram(
+            species=pypicongpu_species,
+            period=self.period.get_as_pypicongpu(time_step_size, num_steps),
+            bin_count=self.bin_count,
+            min_energy=self.min_energy,
+            max_energy=self.max_energy,
+        )
         return pypicongpu_energy_histogram