Added PlasmaLens class to PICMI (#3025)

* Added PlasmaLens class to PICMI

* Added CI test

* Added PICMI input file

* Fixed the output dir for the CI test

* Add `_plt` to Output File Prefix

Co-authored-by: Edoardo Zoni <59625522+EZoni@users.noreply.github.com>

7 files changed, 176 insertions, 6 deletions

diff --git a/Examples/Tests/plasma_lens/PICMI_inputs_3d.py b/Examples/Tests/plasma_lens/PICMI_inputs_3d.py
new file mode 100644
index 000000000..ed64ae3e5
--- /dev/null
+++ b/Examples/Tests/plasma_lens/PICMI_inputs_3d.py
@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+
+from pywarpx import picmi
+
+# Physical constants
+c = picmi.constants.c
+q_e = picmi.constants.q_e
+
+# Number of time steps
+max_steps = 84
+
+# Number of cells
+nx = 16
+ny = 16
+nz = 16
+
+# Physical domain
+xmin = -1.
+xmax =  1.
+ymin = -1.
+ymax =  1.
+zmin =  0.
+zmax =  2.
+
+# Create grid
+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'])
+
+# Particles
+vel_z = 0.5*c
+multiparticles_distribution = picmi.ParticleListDistribution(x = [0.05, 0.],
+                                                             y = [0., 0.04],
+                                                             z = [0.05, 0.05],
+                                                             ux = [0., 0.],
+                                                             uy = [0., 0.],
+                                                             uz = [vel_z, vel_z],
+                                                             weight = [1., 1.])
+
+electrons = picmi.Species(particle_type = 'electron',
+                          name = 'electrons',
+                          initial_distribution = multiparticles_distribution)
+
+# Plasma lenses
+plasma_lenses = picmi.PlasmaLens(period = 0.5,
+                                 starts = [0.1, 0.11, 0.12, 0.13],
+                                 lengths = [0.1, 0.11, 0.12, 0.13],
+                                 strengths_E = [600000., 800000., 600000., 200000.],
+                                 strengths_B = [0.0, 0.0, 0.0, 0.0])
+
+# Electromagnetic solver
+solver = picmi.ElectromagneticSolver(grid = grid,
+                                     method = 'Yee',
+                                     cfl = 0.7)
+
+# Diagnostics
+part_diag1 = picmi.ParticleDiagnostic(name = 'diag1',
+                                      period = max_steps,
+                                      species = [electrons],
+                                      data_list = ['ux', 'uy', 'uz'],
+                                      write_dir = '.',
+                                      warpx_file_prefix = 'Python_plasma_lens_plt')
+
+# Set up simulation
+sim = picmi.Simulation(solver = solver,
+                       max_steps = max_steps,
+                       verbose = 1,
+                       particle_shape = 'linear',
+                       warpx_serialize_initial_conditions = 1,
+                       warpx_do_dynamic_scheduling = 0)
+
+# Add plasma electrons
+sim.add_species(electrons, layout = None)
+
+# Add the plasma lenses
+sim.add_applied_field(plasma_lenses)
+
+# Add diagnostics
+sim.add_diagnostic(part_diag1)
+
+# Write input file that can be used to run with the compiled version
+#sim.write_input_file(file_name = 'inputs_3d_picmi')
+
+# Advance simulation until last time step
+sim.step(max_steps)
diff --git a/Examples/Tests/plasma_lens/analysis.py b/Examples/Tests/plasma_lens/analysis.py
index f13ecb00f..6a0af60ad 100755
--- a/Examples/Tests/plasma_lens/analysis.py
+++ b/Examples/Tests/plasma_lens/analysis.py
@@ -60,7 +60,11 @@ def applylens(x0, vx0, vz0, gamma, lens_length, lens_strength):
     return x1, vx1
 
 clight = c
-vel_z = eval(ds.parameters.get('my_constants.vel_z'))
+try:
+    vel_z = eval(ds.parameters.get('my_constants.vel_z'))
+except TypeError:
+    # vel_z is not saved in my_constants with the PICMI version
+    vel_z = 0.5*c
 
 plasma_lens_period = float(ds.parameters.get('particles.repeated_plasma_lens_period'))
 plasma_lens_starts = [float(x) for x in ds.parameters.get('particles.repeated_plasma_lens_starts').split()]
diff --git a/Examples/Tests/plasma_lens/inputs_3d b/Examples/Tests/plasma_lens/inputs_3d
index 5f557e4eb..48f2263f0 100644
--- a/Examples/Tests/plasma_lens/inputs_3d
+++ b/Examples/Tests/plasma_lens/inputs_3d
@@ -18,6 +18,7 @@ boundary.particle_hi = absorbing absorbing absorbing
 
 # Algorithms
 algo.particle_shape = 1
+warpx.cfl = 0.7
 
 my_constants.vel_z = 0.5*clight
 
diff --git a/Examples/Tests/plasma_lens/inputs_boosted_3d b/Examples/Tests/plasma_lens/inputs_boosted_3d
index 32d6fa47e..c997136d4 100644
--- a/Examples/Tests/plasma_lens/inputs_boosted_3d
+++ b/Examples/Tests/plasma_lens/inputs_boosted_3d
@@ -20,6 +20,7 @@ boundary.particle_hi = absorbing absorbing absorbing
 algo.particle_shape = 1
 warpx.gamma_boost = 2.
 warpx.boost_direction = z
+warpx.cfl = 0.7
 
 my_constants.vel_z = 0.5*clight
 
diff --git a/Python/pywarpx/picmi.py b/Python/pywarpx/picmi.py
index 71e1b7ca5..ae5fc8bd4 100644
--- a/Python/pywarpx/picmi.py
+++ b/Python/pywarpx/picmi.py
@@ -360,9 +360,9 @@ class ParticleListDistribution(picmistandard.PICMI_ParticleListDistribution):
         species.multiple_particles_pos_x = self.x
         species.multiple_particles_pos_y = self.y
         species.multiple_particles_pos_z = self.z
-        species.multiple_particles_vel_x = self.ux/constants.c
-        species.multiple_particles_vel_y = self.uy/constants.c
-        species.multiple_particles_vel_z = self.uz/constants.c
+        species.multiple_particles_vel_x = np.array(self.ux)/constants.c
+        species.multiple_particles_vel_y = np.array(self.uy)/constants.c
+        species.multiple_particles_vel_z = np.array(self.uz)/constants.c
         species.multiple_particles_weight = self.weight
         if density_scale is not None:
             species.multiple_particles_weight = self.weight*density_scale
@@ -953,6 +953,38 @@ class EmbeddedBoundary(picmistandard.base._ClassWithInit):
             pywarpx.warpx.__setattr__('eb_potential(x,y,z,t)', expression)
 
 
+class PlasmaLens(picmistandard.base._ClassWithInit):
+    """
+    Custom class to setup a plasma lens lattice.
+    The applied fields are dependent on the transverse position
+      - Ex = x*stengths_E
+      - Ey = y*stengths_E
+      - Bx = +y*stengths_B
+      - By = -x*stengths_B
+    """
+    def __init__(self, period, starts, lengths, strengths_E=None, strengths_B=None, **kw):
+        self.period = period
+        self.starts = starts
+        self.lengths = lengths
+        self.strengths_E = strengths_E
+        self.strengths_B = strengths_B
+
+        assert (self.strengths_E is not None) or (self.strengths_B is not None),\
+               Exception('One of strengths_E or strengths_B must be supplied')
+
+        self.handle_init(kw)
+
+    def initialize_inputs(self):
+
+        pywarpx.particles.E_ext_particle_init_style = 'repeated_plasma_lens'
+        pywarpx.particles.B_ext_particle_init_style = 'repeated_plasma_lens'
+        pywarpx.particles.repeated_plasma_lens_period = self.period
+        pywarpx.particles.repeated_plasma_lens_starts = self.starts
+        pywarpx.particles.repeated_plasma_lens_lengths = self.lengths
+        pywarpx.particles.repeated_plasma_lens_strengths_E = self.strengths_E
+        pywarpx.particles.repeated_plasma_lens_strengths_B = self.strengths_B
+
+
 class Simulation(picmistandard.PICMI_Simulation):
 
     # Set the C++ WarpX interface (see _libwarpx.LibWarpX) as an extension to
diff --git a/Regression/Checksum/benchmarks_json/Python_plasma_lens.json b/Regression/Checksum/benchmarks_json/Python_plasma_lens.json
new file mode 100644
index 000000000..c3abc56f3
--- /dev/null
+++ b/Regression/Checksum/benchmarks_json/Python_plasma_lens.json
@@ -0,0 +1,23 @@
+{
+  "electrons": {
+    "particle_cpu": 1.0,
+    "particle_id": 2.0,
+    "particle_momentum_x": 7.424668333879405e-24,
+    "particle_momentum_y": 5.93963893779683e-24,
+    "particle_momentum_z": 2.7309245307375727e-22,
+    "particle_position_x": 0.0360838943897417,
+    "particle_position_y": 0.028872102262743393,
+    "particle_position_z": 3.894799963324205
+  },
+  "lev=0": {
+    "Bx": 3.4518920711615626e-14,
+    "By": 3.4457590162508847e-14,
+    "Bz": 3.1975595242439324e-16,
+    "Ex": 4.460637209979829e-06,
+    "Ey": 4.4877026638440135e-06,
+    "Ez": 9.259738797694699e-06,
+    "jx": 4.136556647397788e-10,
+    "jy": 3.3072821927157533e-10,
+    "jz": 1.649309087211398e-08
+  }
+}
diff --git a/Regression/WarpX-tests.ini b/Regression/WarpX-tests.ini
index 6a2e96cc6..ea8e434b0 100644
--- a/Regression/WarpX-tests.ini
+++ b/Regression/WarpX-tests.ini
@@ -2959,7 +2959,7 @@ analysisRoutine = Examples/Tests/pass_mpi_communicator/analysis.py
 [Plasma_lens]
 buildDir = .
 inputFile = Examples/Tests/plasma_lens/inputs_3d
-runtime_params = warpx.cfl=0.7
+runtime_params =
 dim = 3
 addToCompileString =
 cmakeSetupOpts = -DWarpX_DIMS=3
@@ -2977,7 +2977,7 @@ analysisRoutine = Examples/Tests/plasma_lens/analysis.py
 [Plasma_lens_boosted]
 buildDir = .
 inputFile = Examples/Tests/plasma_lens/inputs_boosted_3d
-runtime_params = warpx.cfl=0.7
+runtime_params =
 dim = 3
 addToCompileString =
 cmakeSetupOpts = -DWarpX_DIMS=3
@@ -2992,6 +2992,26 @@ compareParticles = 1
 particleTypes = electrons
 analysisRoutine = Examples/Tests/plasma_lens/analysis.py
 
+[Python_plasma_lens]
+buildDir = .
+inputFile = Examples/Tests/plasma_lens/PICMI_inputs_3d.py
+runtime_params =
+customRunCmd = python3 PICMI_inputs_3d.py
+dim = 3
+addToCompileString = USE_PYTHON_MAIN=TRUE
+cmakeSetupOpts = -DWarpX_DIMS=3 -DWarpX_LIB=ON -DWarpX_APP=OFF
+target = pip_install
+restartTest = 0
+useMPI = 1
+numprocs = 2
+useOMP = 1
+numthreads = 1
+compileTest = 0
+doVis = 0
+compareParticles = 1
+particleTypes = electrons
+analysisRoutine = Examples/Tests/plasma_lens/analysis.py
+
 [background_mcc]
 buildDir = .
 inputFile = Examples/Physics_applications/capacitive_discharge/inputs_2d