Mini-PR: Improve Python particle attribute access CI test (#2581)

* updated the Python particle attribute access CI test

* additional fixes needed for tests to pass

* changes requested by Axel during code review

* Apply suggestions from code review

Co-authored-by: Axel Huebl <axel.huebl@plasma.ninja>

Author: roelof-groenewald

Examples/Tests/ParticleDataPython/PICMI_inputs_2d.py

@@ -1,19 +1,19 @@
-from mpi4py import MPI
 from pywarpx import picmi
 import numpy as np
 
-##########################
-# MPI communicator setup
-##########################
+import argparse
+import sys
 
-# split processor 0 into separate communicator from others
-comm_world = MPI.COMM_WORLD
-rank = comm_world.Get_rank()
-if rank == 0:
-    color = 0
-else:
-    color = 1
-new_comm = comm_world.Split(color)
+# Create the parser and add the argument
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    '-u', '--unique', action='store_true',
+    help="Whether injected particles should be treated as unique"
+)
+
+# Parse the input
+args, left = parser.parse_known_args()
+sys.argv = sys.argv[:1] + left
 
 ##########################
 # numerics parameters
@@ -75,7 +75,7 @@
     period = 10,
     data_list = ['phi'],
     write_dir = '.',
-    warpx_file_prefix = f'Python_particle_attr_access_plt_{color}'
+    warpx_file_prefix = f"Python_particle_attr_access_{'unique_' if args.unique else ''}plt"
 )
 
 ##########################
@@ -98,7 +98,7 @@
 sim.add_diagnostic(field_diag)
 
 sim.initialize_inputs()
-sim.initialize_warpx(mpi_comm=new_comm)
+sim.initialize_warpx()
 
 ##########################
 # python particle data access
@@ -108,9 +108,11 @@
 
 _libwarpx.add_real_comp('electrons', 'newPid')
 
+my_id = _libwarpx.libwarpx.warpx_getMyProc()
+
 def add_particles():
 
-    nps = 10
+    nps = 10 * (my_id + 1)
     x = np.random.rand(nps) * 0.03
     y = np.zeros(nps)
     z = np.random.random(nps) * 0.03
@@ -122,7 +124,7 @@ def add_particles():
 
     _libwarpx.add_particles(
         species_name='electrons', x=x, y=y, z=z, ux=ux, uy=uy, uz=uz,
-        w=w, newPid=newPid, unique_particles=(not color)
+        w=w, newPid=newPid, unique_particles=args.unique
     )
 
 callbacks.installbeforestep(add_particles)
@@ -131,16 +133,14 @@ def add_particles():
 # simulation run
 ##########################
 
-sim.step(max_steps - 1, mpi_comm=new_comm)
+sim.step(max_steps - 1)
 
 ##########################
-# check that the new PIDs are properly set
+# check that the new PIDs
+# are properly set
 ##########################
 
-if color == 0:
-    assert (_libwarpx.get_particle_count('electrons') == 90)
-else:
-    assert (_libwarpx.get_particle_count('electrons') == 90)
+assert (_libwarpx.get_particle_count('electrons') == 270 / (2 - args.unique))
 assert (_libwarpx.get_particle_comp_index('electrons', 'w') == 0)
 assert (_libwarpx.get_particle_comp_index('electrons', 'newPid') == 4)
 

Examples/Tests/ParticleDataPython/analysis.py

@@ -7,12 +7,8 @@
 # This script just checks that the PICMI file executed successfully.
 # If it did there will be a plotfile for the final step.
 
-import yt
+import sys
 
-plotfile = 'Python_particle_attr_access_plt_000010'
-ds = yt.load( plotfile )  # noqa
+step = int(sys.argv[1][-5:])
 
-plotfile = 'Python_particle_attr_access_plt_100010'
-ds = yt.load( plotfile )  # noqa
-
-assert True
+assert step == 10

Regression/WarpX-tests.ini

@@ -2619,6 +2619,22 @@ compileTest = 0
 doVis = 0
 analysisRoutine = Examples/Tests/ParticleDataPython/analysis.py
 
+[Python_particle_attr_access_unique]
+buildDir = .
+inputFile = Examples/Tests/ParticleDataPython/PICMI_inputs_2d.py
+runtime_params =
+customRunCmd = python PICMI_inputs_2d.py --unique
+dim = 2
+addToCompileString = USE_PYTHON_MAIN=TRUE PYINSTALLOPTIONS="--user --prefix="
+restartTest = 0
+useMPI = 1
+numprocs = 2
+useOMP = 1
+numthreads = 1
+compileTest = 0
+doVis = 0
+analysisRoutine = Examples/Tests/ParticleDataPython/analysis.py
+
 [Python_prev_positions]
 buildDir = .
 inputFile = Examples/Tests/ParticleDataPython/PICMI_inputs_prev_pos_2d.py