Add ability to load particles from openPMD, through the PICMI interface (#6305)

Do no merge until picmi-standard/picmi#127 is
merged and a new PICMI version is released.

It also adds a corresponding automated test

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: RemiLehe

Docs/requirements.txt

@@ -11,7 +11,7 @@ docutils>=0.17.1
 
 # PICMI API docs
 # note: keep in sync with version in ../requirements.txt
-picmistandard==0.33.0
+picmistandard==0.34.0
 # for development against an unreleased PICMI version, use:
 # picmistandard @ git+https://github.com/picmi-standard/picmi.git#subdirectory=PICMI_Python
 

Docs/source/usage/python.rst

@@ -152,6 +152,8 @@ Particle distributions can be used for to initialize particles in a particle spe
 
 .. autoclass:: pywarpx.picmi.ParticleListDistribution
 
+.. autoclass:: pywarpx.picmi.FromFileDistribution
+
 Particle layouts determine how to microscopically place macro particles in a grid cell.
 
 .. autoclass:: pywarpx.picmi.GriddedLayout

Examples/Tests/gaussian_beam/CMakeLists.txt

@@ -21,6 +21,16 @@ add_warpx_test(
     test_3d_focusing_gaussian_beam_from_openpmd_prepare  # dependency
 )
 
+add_warpx_test(
+    test_3d_focusing_gaussian_beam_from_openpmd_picmi  # name
+    3  # dims
+    2  # nprocs
+    inputs_test_3d_focusing_gaussian_beam_from_openpmd_picmi.py  # inputs
+    "analysis_focusing_beam.py"  # analysis
+    "analysis_default_regression.py --path diags/diag1000000"  # checksum
+    test_3d_focusing_gaussian_beam_from_openpmd_prepare  # dependency
+)
+
 add_warpx_test(
     test_3d_focusing_gaussian_beam_from_openpmd_prepare  # name
     3  # dims

Examples/Tests/gaussian_beam/inputs_test_3d_focusing_gaussian_beam_from_openpmd_picmi.py

@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+
+from pywarpx import picmi
+
+constants = picmi.constants
+
+# Constants from the focusing beam example
+nano = 1.0e-9
+micro = 1.0e-6
+
+sigmax = 516.0 * nano
+sigmay = 7.7 * nano
+sigmaz = 300.0 * micro
+
+# Box dimensions
+Lx = 20 * sigmax
+Ly = 20 * sigmay
+Lz = 20 * sigmaz
+
+nx = 256
+ny = 256
+nz = 256
+
+xmin = -0.5 * Lx
+xmax = 0.5 * Lx
+ymin = -0.5 * Ly
+ymax = 0.5 * Ly
+zmin = -0.5 * Lz
+zmax = 0.5 * Lz
+
+em_order = 3
+
+grid = picmi.Cartesian3DGrid(
+    number_of_cells=[nx, ny, nz],
+    lower_bound=[xmin, ymin, zmin],
+    upper_bound=[xmax, ymax, zmax],
+    lower_boundary_conditions=["dirichlet", "dirichlet", "dirichlet"],
+    upper_boundary_conditions=["dirichlet", "dirichlet", "dirichlet"],
+    lower_boundary_conditions_particles=["absorbing", "absorbing", "absorbing"],
+    upper_boundary_conditions_particles=["absorbing", "absorbing", "absorbing"],
+    warpx_max_grid_size=256,
+)
+
+solver = picmi.ElectromagneticSolver(grid=grid)
+
+# Create species with external file injection
+
+beam1_distribution = picmi.FromFileDistribution(
+    file_path="../test_3d_focusing_gaussian_beam_from_openpmd_prepare/openpmd_generated_particles.h5",
+)
+
+beam1 = picmi.Species(
+    particle_type="electron",
+    name="beam1",
+    initial_distribution=beam1_distribution,
+)
+
+diag1 = picmi.ParticleDiagnostic(
+    name="diag1",
+    period=1,
+    warpx_dump_last_timestep=1,
+)
+
+openpmd = picmi.ParticleDiagnostic(
+    name="openpmd",
+    period=1,
+    species=[beam1],
+    data_list=["weighting", "x", "y", "z"],
+    warpx_format="openpmd",
+    warpx_dump_last_timestep=1,
+)
+
+sim = picmi.Simulation(
+    solver=solver,
+    max_steps=0,
+    verbose=1,
+)
+
+sim.add_species(beam1, layout=None)
+
+sim.add_diagnostic(diag1)
+sim.add_diagnostic(openpmd)
+
+# write_inputs will create an inputs file that can be used to run
+# with the compiled version.
+# sim.write_input_file(file_name = 'inputs_from_PICMI')
+
+# Alternatively, sim.step will run WarpX, controlling it from Python
+sim.step()

Python/pywarpx/picmi.py

@@ -905,6 +905,17 @@ def distribution_initialize_inputs(
             )
 
 
+class FromFileDistribution(picmistandard.PICMI_FromFileDistribution):
+    def init(self, kw):
+        pass
+
+    def distribution_initialize_inputs(
+        self, species_number, layout, species, density_scale, source_name
+    ):
+        species.add_new_group_attr(source_name, "injection_style", "external_file")
+        species.add_new_group_attr(source_name, "injection_file", self.file_path)
+
+
 class ParticleDistributionPlanarInjector(
     picmistandard.PICMI_ParticleDistributionPlanarInjector
 ):

Python/setup.py

@@ -77,7 +77,7 @@
     package_dir={"pywarpx": "pywarpx"},
     description="""Wrapper of WarpX""",
     package_data=package_data,
-    install_requires=["numpy", "picmistandard==0.33.0", "periodictable"],
+    install_requires=["numpy", "picmistandard==0.34.0", "periodictable"],
     python_requires=">=3.8",  # left for CI, truly ">=3.9"
     zip_safe=False,
 )

Regression/Checksum/benchmarks_json/test_3d_focusing_gaussian_beam_from_openpmd_picmi.json

@@ -0,0 +1,12 @@
+{
+  "beam1": {
+    "particle_momentum_x": 4.222627241785717e-14,
+    "particle_momentum_y": 1.131583197131103e-15,
+    "particle_momentum_z": 1.3360487849541085e-10,
+    "particle_position_x": 1.1288391053970308,
+    "particle_position_y": 0.02391393806881765,
+    "particle_position_z": 478.7122905262462,
+    "particle_weight": 19999660000.0
+  },
+  "lev=0": {}
+}

Tools/machines/karolina-it4i/install_dependencies.sh

@@ -53,7 +53,7 @@ python -m pip install --user --upgrade matplotlib
 #python -m pip install --user --upgrade yt
 
 # install or update WarpX dependencies
-python -m pip install --user --upgrade picmistandard==0.33.0
+python -m pip install --user --upgrade picmistandard==0.34.0
 python -m pip install --user --upgrade lasy
 
 # optional: for optimas (based on libEnsemble & ax->botorch->gpytorch->pytorch)

requirements.txt

@@ -4,7 +4,7 @@ periodictable~=1.5
 
 # PICMI
 # note: don't forget to update the version in Docs/requirements.txt, too
-picmistandard==0.33.0
+picmistandard==0.34.0
 #   for development against an unreleased PICMI version, use:
 #picmistandard @ git+https://github.com/picmi-standard/picmi.git#subdirectory=PICMI_Python