Add ability to load external particle fields from file (ECP-WarpX#4997)

* Update documentation * Allocate fields * Always parse filename * Add option "FromFile" in external particles fields * Create two separate examples * Update documentation * Allocate aux array * Update aux to take into account copy * Add external fields from file to auxiliary data * Update tests * Remove debugging Print statements * If averaged PSATD is used, we copy Efield_avg_fp to Efield_aux, otherwise copy Efield_fp. * Do not access cells that have invalid data when summing external fields * Update path of checkpoint files * Update comments * Update number of components * Fix clang tidy error * Make the code compatible with momentum-conserving gather * Add Python interface and tests * Use PICMI version that defined LoadAppliedField * Update file names * Update PICMI standard version * Allocate dedicated MultiFabs for external particle fields * Some clean-up * Performance optimization: do not use external fields in gather kernel * Allow external particle fields and external grid fields to be simultaneously defined * External fields are incompatible with the moving window * Support load balancing * Update Source/Initialization/WarpXInitData.cpp * Update Source/Initialization/WarpXInitData.cpp --------- Co-authored-by: Remi Lehe <[email protected]>
EZoni · Jul 9, 2024 · ce9f4af · ce9f4af
1 parent 0c33365
commit ce9f4af
Show file tree

Hide file tree

Showing 19 changed files with 497 additions and 42 deletions.
diff --git a/Docs/source/usage/parameters.rst b/Docs/source/usage/parameters.rst
@@ -1747,8 +1747,8 @@ are applied to the grid directly. In particular, these fields can be seen in the
     One can refer to input files in ``Examples/Tests/LoadExternalField`` for more information.
     Regarding how to prepare the openPMD data file, one can refer to
     the `openPMD-example-datasets <https://github.com/openPMD/openPMD-example-datasets>`__.
-    Note that if both `B_ext_grid_init_style` and `E_ext_grid_init_style` are set to
-    `read_from_file`, the openPMD file specified by `warpx.read_fields_from_path`
+    Note that if both ``B_ext_grid_init_style`` and ``E_ext_grid_init_style`` are set to
+    ``read_from_file``, the openPMD file specified by ``warpx.read_fields_from_path``
     should contain both B and E external fields data.
 
 * ``warpx.E_external_grid`` & ``warpx.B_external_grid`` (list of `3 floats`)
@@ -1804,6 +1804,21 @@ are applied to the particles directly, at each timestep. As a results, these fie
 
       Note that the position is defined in Cartesian coordinates, as a function of (x,y,z), even for RZ.
 
+    * ``read_from_file``: load the external field from an openPMD file.
+        An additional parameter, indicating the path of an openPMD data file, ``particles.read_fields_from_path``
+        must be specified, from which the external E field data can be loaded into WarpX.
+        One can refer to input files in ``Examples/Tests/LoadExternalField`` for more information.
+        Regarding how to prepare the openPMD data file, one can refer to
+        the `openPMD-example-datasets <https://github.com/openPMD/openPMD-example-datasets>`__.
+        Note that if both ``B_ext_particle_init_style`` and ``E_ext_particle_init_style`` are set to
+        ``read_from_file``, the openPMD file specified by ``particles.read_fields_from_path``
+        should contain both B and E external fields data.
+
+        .. note::
+
+            When using ``read_from_file``, the fields loaded from the file will be interpolated
+            to the resolution of the grid used for the simulation.
+
     * ``repeated_plasma_lens``: apply a series of plasma lenses.
       The properties of the lenses are defined in the lab frame by the input parameters:
 

diff --git a/...ests/LoadExternalField/PICMI_inputs_3d.py → ...ernalField/PICMI_inputs_3d_grid_fields.py b/...ests/LoadExternalField/PICMI_inputs_3d.py → ...ernalField/PICMI_inputs_3d_grid_fields.py
@@ -93,15 +93,15 @@
         species=[ions],
         data_list = ['ux', 'uy', 'uz', 'x', 'y', 'z', 'weighting'],
         write_dir='.',
-        warpx_file_prefix='Python_LoadExternalField3D_plt'
+        warpx_file_prefix='Python_LoadExternalGridField3D_plt'
         )
 field_diag = picmi.FieldDiagnostic(
         name='diag1',
         grid=grid,
         period=300,
         data_list = ['Bx', 'By', 'Bz', 'Ex', 'Ey', 'Ez', 'Jx', 'Jy', 'Jz'],
         write_dir='.',
-        warpx_file_prefix='Python_LoadExternalField3D_plt'
+        warpx_file_prefix='Python_LoadExternalGridField3D_plt'
         )
 
 #################################

diff --git a/Examples/Tests/LoadExternalField/PICMI_inputs_3d_particle_fields.py b/Examples/Tests/LoadExternalField/PICMI_inputs_3d_particle_fields.py
@@ -0,0 +1,140 @@
+#!/usr/bin/env python3
+#
+# --- Input file for loading initial field from openPMD file.
+
+from pywarpx import picmi
+
+constants = picmi.constants
+
+#################################
+####### GENERAL PARAMETERS ######
+#################################
+
+max_steps = 300
+
+max_grid_size = 40
+nx = max_grid_size
+ny = max_grid_size
+nz = max_grid_size
+
+xmin = -1
+xmax = 1
+ymin = xmin
+ymax = xmax
+zmin = 0
+zmax = 5
+
+number_per_cell = 200
+
+#################################
+############ NUMERICS ###########
+#################################
+
+verbose = 1
+dt = 4.4e-7
+use_filter = 0
+
+# Order of particle shape factors
+particle_shape = 1
+
+#################################
+############ PLASMA #############
+#################################
+
+ion_dist = picmi.ParticleListDistribution(
+        x=0.0,
+        y=0.2,
+        z=2.5,
+        ux=9.5e-05*constants.c,
+        uy=0.0*constants.c,
+        uz=1.34e-4*constants.c,
+        weight=1.0
+        )
+
+ions = picmi.Species(
+        particle_type='H',
+        name='proton', charge='q_e',mass="m_p",
+        warpx_do_not_deposit=1,
+        initial_distribution=ion_dist
+        )
+
+#################################
+######## INITIAL FIELD ##########
+#################################
+
+applied_field = picmi.LoadAppliedField(
+        read_fields_from_path="../../../../openPMD-example-datasets/example-femm-3d.h5",
+        load_E=False
+        )
+
+#################################
+###### GRID AND SOLVER ##########
+#################################
+
+grid = picmi.Cartesian3DGrid(
+        number_of_cells=[nx, ny, nz],
+        warpx_max_grid_size=max_grid_size,
+        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']
+        )
+solver = picmi.ElectrostaticSolver(grid=grid)
+
+#################################
+######### DIAGNOSTICS ###########
+#################################
+
+particle_diag = picmi.ParticleDiagnostic(
+        name='diag1',
+        period=300,
+        species=[ions],
+        data_list = ['ux', 'uy', 'uz', 'x', 'y', 'z', 'weighting'],
+        write_dir='.',
+        warpx_file_prefix='Python_LoadExternalParticleField3D_plt'
+        )
+field_diag = picmi.FieldDiagnostic(
+        name='diag1',
+        grid=grid,
+        period=300,
+        data_list = ['Bx', 'By', 'Bz', 'Ex', 'Ey', 'Ez', 'Jx', 'Jy', 'Jz'],
+        write_dir='.',
+        warpx_file_prefix='Python_LoadExternalParticleField3D_plt'
+        )
+
+#################################
+####### SIMULATION SETUP ########
+#################################
+
+sim = picmi.Simulation(
+        solver=solver,
+        max_steps=max_steps,
+        verbose=verbose,
+        warpx_serialize_initial_conditions=False,
+        warpx_grid_type='collocated',
+        warpx_do_dynamic_scheduling=False,
+        warpx_use_filter=use_filter,
+        time_step_size=dt,
+        particle_shape=particle_shape
+        )
+
+sim.add_applied_field(applied_field)
+
+sim.add_species(
+        ions,
+        layout=picmi.PseudoRandomLayout(
+            n_macroparticles_per_cell=number_per_cell, grid=grid
+            )
+        )
+
+sim.add_diagnostic(field_diag)
+sim.add_diagnostic(particle_diag)
+
+#################################
+##### SIMULATION EXECUTION ######
+#################################
+
+#sim.write_input_file('PICMI_inputs_3d')
+sim.step(max_steps)
diff --git a/Examples/Tests/LoadExternalField/inputs_rz → ...s/LoadExternalField/inputs_rz_grid_fields b/Examples/Tests/LoadExternalField/inputs_rz → ...s/LoadExternalField/inputs_rz_grid_fields
diff --git a/Examples/Tests/LoadExternalField/inputs_rz_particle_fields b/Examples/Tests/LoadExternalField/inputs_rz_particle_fields
@@ -0,0 +1,65 @@
+warpx.serialize_initial_conditions = 0
+warpx.do_dynamic_scheduling = 0
+particles.do_tiling = 0
+
+particles.B_ext_particle_init_style = "read_from_file"
+particles.read_fields_from_path = "../../../../openPMD-example-datasets/example-femm-thetaMode.h5"
+
+warpx.grid_type = collocated
+warpx.do_electrostatic = labframe
+
+#################################
+####### GENERAL PARAMETERS ######
+#################################
+max_step = 300
+amr.n_cell = 40 40
+warpx.numprocs = 1 1
+amr.max_level = 0
+geometry.dims = RZ
+
+geometry.prob_lo =  0.0  0.0
+geometry.prob_hi =  1.0  5.0
+
+#################################
+###### Boundary Condition #######
+#################################
+boundary.field_lo = none pec
+boundary.field_hi = pec  pec
+boundary.potential_lo_x = 0
+boundary.potential_hi_x = 0
+boundary.potential_lo_y = 0
+boundary.potential_hi_y = 0
+boundary.potential_lo_z = 0
+boundary.potential_hi_z = 0
+
+#################################
+############ NUMERICS ###########
+#################################
+warpx.serialize_initial_conditions = 1
+warpx.verbose = 1
+warpx.const_dt = 4.40917904849092e-7
+warpx.use_filter = 0
+
+# Order of particle shape factors
+algo.particle_shape = 1
+
+#################################
+############ PLASMA #############
+#################################
+particles.species_names = proton
+proton.injection_style = "SingleParticle"
+proton.single_particle_pos = 0.0  0.2  2.5
+proton.single_particle_u = 9.506735958279367e-05  0.0  0.00013435537232359165
+proton.single_particle_weight = 1.0
+proton.do_not_deposit = 1
+proton.mass = m_p
+proton.charge = q_e
+
+# Diagnostics
+diagnostics.diags_names = diag1 chk
+diag1.intervals = 300
+diag1.diag_type = Full
+
+chk.intervals = 150
+chk.diag_type = Full
+chk.format = checkpoint
diff --git a/Python/pywarpx/picmi.py b/Python/pywarpx/picmi.py
@@ -1486,6 +1486,13 @@ def applied_field_initialize_inputs(self):
                 expression = pywarpx.my_constants.mangle_expression(expression, self.mangle_dict)
                 pywarpx.warpx.__setattr__(f'B{sdir}_external_grid_function(x,y,z)', expression)
 
+class LoadAppliedField(picmistandard.PICMI_LoadAppliedField):
+    def applied_field_initialize_inputs(self):
+        pywarpx.particles.read_fields_from_path = self.read_fields_from_path
+        if self.load_E:
+            pywarpx.particles.E_ext_particle_init_style = 'read_from_file'
+        if self.load_B:
+            pywarpx.particles.B_ext_particle_init_style = 'read_from_file'
 
 class ConstantAppliedField(picmistandard.PICMI_ConstantAppliedField):
     def applied_field_initialize_inputs(self):

diff --git a/.../benchmarks_json/LoadExternalFieldRZ.json → ...chmarks_json/LoadExternalFieldRZGrid.json b/.../benchmarks_json/LoadExternalFieldRZ.json → ...chmarks_json/LoadExternalFieldRZGrid.json
diff --git a/Regression/Checksum/benchmarks_json/LoadExternalFieldRZParticles.json b/Regression/Checksum/benchmarks_json/LoadExternalFieldRZParticles.json
@@ -0,0 +1,22 @@
+{
+  "lev=0": {
+    "Br": 0.6380326770459374,
+    "Bt": 0.0,
+    "Bz": 4.850698699951099,
+    "Er": 0.0,
+    "Et": 0.0,
+    "Ez": 0.0,
+    "jr": 0.0,
+    "jt": 0.0,
+    "jz": 0.0
+  },
+  "proton": {
+    "particle_momentum_x": 1.2945379007125463e-23,
+    "particle_momentum_y": 7.901629565307941e-23,
+    "particle_momentum_z": 2.0004592432172922e-23,
+    "particle_position_x": 0.12402004585603355,
+    "particle_position_y": 4.363249204658946,
+    "particle_theta": 0.22200801011337173,
+    "particle_weight": 1.0
+  }
+}
diff --git a/...arks_json/Python_LoadExternalField3D.json → ..._json/Python_LoadExternalGridField3D.json b/...arks_json/Python_LoadExternalField3D.json → ..._json/Python_LoadExternalGridField3D.json
diff --git a/Regression/Checksum/benchmarks_json/Python_LoadExternalParticleField3D.json b/Regression/Checksum/benchmarks_json/Python_LoadExternalParticleField3D.json
@@ -0,0 +1,22 @@
+{
+  "lev=0": {
+    "Bx": 29.8448958859583,
+    "By": 29.844895885958305,
+    "Bz": 197.55124310523018,
+    "Ex": 0.0,
+    "Ey": 0.0,
+    "Ez": 0.0,
+    "jx": 0.0,
+    "jy": 0.0,
+    "jz": 0.0
+  },
+  "proton": {
+    "particle_momentum_x": 1.2118953253556959e-23,
+    "particle_momentum_y": 7.822512598009088e-23,
+    "particle_momentum_z": 2.2761898274347412e-23,
+    "particle_position_x": 0.12238072371983327,
+    "particle_position_y": 0.009653941154598592,
+    "particle_position_z": 4.317544769595657,
+    "particle_weight": 1.0
+  }
+}
diff --git a/Regression/WarpX-tests.ini b/Regression/WarpX-tests.ini
@@ -2266,10 +2266,29 @@ doVis = 0
 compareParticles = 0
 analysisRoutine = Examples/Tests/resampling/analysis_leveling_thinning.py
 
-[LoadExternalFieldRZ]
+[LoadExternalFieldRZGrid]
 buildDir = .
-inputFile = Examples/Tests/LoadExternalField/inputs_rz
-runtime_params = warpx.abort_on_warning_threshold=medium chk.file_prefix=LoadExternalFieldRZ_chk chk.file_min_digits=5
+inputFile = Examples/Tests/LoadExternalField/inputs_rz_grid_fields
+runtime_params = warpx.abort_on_warning_threshold=medium chk.file_prefix=LoadExternalFieldRZGrid_chk chk.file_min_digits=5
+dim = 2
+addToCompileString = USE_RZ=TRUE
+cmakeSetupOpts = -DWarpX_DIMS=RZ -DWarpX_OPENPMD=ON
+restartTest = 1
+restartFileNum = 150
+useMPI = 1
+numprocs = 1
+useOMP = 1
+numthreads = 1
+compileTest = 0
+doVis = 0
+compareParticles = 1
+particleTypes = proton
+analysisRoutine = Examples/Tests/LoadExternalField/analysis_rz.py
+
+[LoadExternalFieldRZParticles]
+buildDir = .
+inputFile = Examples/Tests/LoadExternalField/inputs_rz_particle_fields
+runtime_params = warpx.abort_on_warning_threshold=medium chk.file_prefix=LoadExternalFieldRZParticles_chk chk.file_min_digits=5
 dim = 2
 addToCompileString = USE_RZ=TRUE
 cmakeSetupOpts = -DWarpX_DIMS=RZ -DWarpX_OPENPMD=ON
@@ -3448,11 +3467,31 @@ compileTest = 0
 doVis = 0
 analysisRoutine = Examples/analysis_default_openpmd_regression.py
 
-[Python_LoadExternalField3D]
+[Python_LoadExternalGridField3D]
 buildDir = .
-inputFile = Examples/Tests/LoadExternalField/PICMI_inputs_3d.py
+inputFile = Examples/Tests/LoadExternalField/PICMI_inputs_3d_grid_fields.py
 runtime_params =
-customRunCmd = python3 PICMI_inputs_3d.py
+customRunCmd = python3 PICMI_inputs_3d_grid_fields.py
+dim = 3
+addToCompileString = USE_PYTHON_MAIN=TRUE
+cmakeSetupOpts = -DWarpX_DIMS=3 -DWarpX_APP=OFF -DWarpX_PYTHON=ON -DWarpX_OPENPMD=ON
+target = pip_install
+restartTest = 0
+useMPI = 1
+numprocs = 1
+useOMP = 1
+numthreads = 1
+compileTest = 0
+doVis = 0
+compareParticles = 1
+particleTypes = proton
+analysisRoutine = Examples/Tests/LoadExternalField/analysis_3d.py
+
+[Python_LoadExternalParticleField3D]
+buildDir = .
+inputFile = Examples/Tests/LoadExternalField/PICMI_inputs_3d_particle_fields.py
+runtime_params =
+customRunCmd = python3 PICMI_inputs_3d_particle_fields.py
 dim = 3
 addToCompileString = USE_PYTHON_MAIN=TRUE
 cmakeSetupOpts = -DWarpX_DIMS=3 -DWarpX_APP=OFF -DWarpX_PYTHON=ON -DWarpX_OPENPMD=ON

diff --git a/Source/Initialization/ExternalField.cpp b/Source/Initialization/ExternalField.cpp
@@ -172,7 +172,6 @@ ExternalFieldParams::ExternalFieldParams(const amrex::ParmParse& pp_warpx)
     //
     //  External fields from file
     //
-
     if (E_ext_grid_type == ExternalFieldType::read_from_file ||
         B_ext_grid_type == ExternalFieldType::read_from_file){
             const std::string read_fields_from_path="./";