Python: Remove PC Wrapper in Usage (BLAST-WarpX#6357)

Modernize examples to use the new `sim.particles` access to
`MultiParticleContainer` directly.

Author: ax3l

Examples/Physics_applications/capacitive_discharge/inputs_base_1d_picmi.py

@@ -11,7 +11,7 @@
 from scipy.sparse import csc_matrix
 from scipy.sparse import linalg as sla
 
-from pywarpx import callbacks, libwarpx, particle_containers, picmi
+from pywarpx import callbacks, libwarpx, picmi
 from pywarpx.LoadThirdParty import load_cupy
 
 constants = picmi.constants
@@ -409,36 +409,32 @@ def rethermalize_neutrals(self):
         if step % 1000 != 10:
             return
 
-        if not hasattr(self, "neutral_cont"):
-            self.neutral_cont = particle_containers.ParticleContainerWrapper(
-                self.neutrals.name
-            )
-
-        ux_arrays = self.neutral_cont.uxp
-        uy_arrays = self.neutral_cont.uyp
-        uz_arrays = self.neutral_cont.uzp
+        if not hasattr(self, "neutral_particles"):
+            self.neutral_particles = self.sim.particles.get(self.neutrals.name)
 
         xp, _ = load_cupy()
 
         vel_std = np.sqrt(constants.kb * self.gas_temp / self.m_ion)
-        for ii in range(len(ux_arrays)):
-            nps = len(ux_arrays[ii])
-            ux_arrays[ii][:] = xp.array(vel_std * self.rng.normal(size=nps))
-            uy_arrays[ii][:] = xp.array(vel_std * self.rng.normal(size=nps))
-            uz_arrays[ii][:] = xp.array(vel_std * self.rng.normal(size=nps))
+        for pti in self.neutral_particles.iterator(level=0):
+            nps = pti.size
+            pti["ux"][:] = xp.array(vel_std * self.rng.normal(size=nps))
+            pti["uy"][:] = xp.array(vel_std * self.rng.normal(size=nps))
+            pti["uz"][:] = xp.array(vel_std * self.rng.normal(size=nps))
 
     def _get_rho_ions(self):
         # deposit the ion density in rho_fp
-        he_ions_wrapper = particle_containers.ParticleContainerWrapper("he_ions")
-        he_ions_wrapper.deposit_charge_density(level=0)
+        self.rho.set_val(0.0)
+        he_ions = self.sim.particles.get("he_ions")
+        he_ions.deposit_charge(self.rho, lev=0)
+        libwarpx.warpx.sync_rho()
 
-        rho_data = self.rho_wrapper[...]
+        rho_data = self.rho[...]
         self.ion_density_array += rho_data / constants.q_e / self.diag_steps
 
     def run_sim(self):
         self.sim.step(self.max_steps - self.diag_steps)
 
-        self.rho_wrapper = self.sim.fields.get("rho_fp", level=0)
+        self.rho = self.sim.fields.get("rho_fp", level=0)
         callbacks.installafterstep(self._get_rho_ions)
 
         self.sim.step(self.diag_steps)
@@ -453,9 +449,11 @@ def run_sim(self):
         # query the particle z-coordinates if this is run during CI testing
         # to cover that functionality
         if self.test:
-            he_ions_wrapper = particle_containers.ParticleContainerWrapper("he_ions")
-            nparts = he_ions_wrapper.get_particle_count(local=True)
-            z_coords = np.concatenate(he_ions_wrapper.zp)
+            he_ions = self.sim.particles.get("he_ions")
+            nparts = he_ions.number_of_particles(only_local=True)
+            z_coords = np.concatenate(
+                list(pti["z"] for pti in he_ions.iterator(level=0))
+            )
             assert len(z_coords) == nparts
             assert np.all(z_coords >= 0.0) and np.all(z_coords <= self.gap)
 

Examples/Physics_applications/laser_acceleration/inputs_test_3d_laser_acceleration_python.py

@@ -18,12 +18,10 @@ def my_simple_callback():
     and
     https://warpx.readthedocs.io/en/latest/usage/workflows/python_extend.html#fields
     """
-    from pywarpx import particle_containers
-
     print("  my_simple_callback")
 
     # electrons: access (and potentially manipulate)
-    electrons = particle_containers.ParticleContainerWrapper("electrons")
+    electrons = sim.particles.get("electrons")
     print(f"    {electrons}")
 
     # electric field: access (and potentially manipulate)
@@ -42,14 +40,9 @@ def my_advanced_callback():
     amr = sim.extension.amr
     amr.Print(f"    {amr.ParallelDescriptor.NProcs()} MPI process(es) active")
 
-    # the active simulation
-    warpx = sim.extension.warpx
-
     # electrons: access (and potentially manipulate)
-    electrons_pc = warpx.multi_particle_container().get_particle_container_from_name(
-        "electrons"
-    )
-    print(f"    {electrons_pc}")
+    electrons = sim.particles.get("electrons")
+    print(f"    {electrons}")
 
     # electric field: access (and potentially manipulate)
     Ex_mf = sim.fields.get("Efield_fp", dir="x", level=0)

Examples/Tests/ohm_solver_ion_Landau_damping/inputs_test_2d_ohm_solver_landau_damping_picmi.py

@@ -14,7 +14,7 @@
 import numpy as np
 from mpi4py import MPI as mpi
 
-from pywarpx import callbacks, libwarpx, particle_containers, picmi
+from pywarpx import callbacks, libwarpx, picmi
 
 constants = picmi.constants
 
@@ -254,7 +254,7 @@ def setup_run(self):
         simulation.initialize_warpx()
 
         # get ion particle container wrapper
-        self.ion_part_container = particle_containers.ParticleContainerWrapper("ions")
+        self.ions = simulation.particles.get("ions")
 
     def text_diag(self):
         """Diagnostic function to print out timing data and particle numbers."""
@@ -269,7 +269,7 @@ def text_diag(self):
 
         status_dict = {
             "step": step,
-            "nplive ions": self.ion_part_container.nps,
+            "nplive ions": self.ions.size,
             "wall_time": wall_time,
             "step_rate": step_rate,
             "diag_steps": self.diag_steps,

Examples/Tests/ohm_solver_ion_beam_instability/inputs_test_1d_ohm_solver_ion_beam_picmi.py

@@ -15,7 +15,7 @@
 import numpy as np
 from mpi4py import MPI as mpi
 
-from pywarpx import callbacks, libwarpx, particle_containers, picmi
+from pywarpx import callbacks, libwarpx, picmi
 
 constants = picmi.constants
 
@@ -315,12 +315,8 @@ def setup_run(self):
 
         # create particle container wrapper for the ion species to access
         # particle data
-        self.ion_container_wrapper = particle_containers.ParticleContainerWrapper(
-            self.ions.name
-        )
-        self.beam_ion_container_wrapper = particle_containers.ParticleContainerWrapper(
-            self.beam_ions.name
-        )
+        self.ion_container = simulation.particles.get(self.ions.name)
+        self.beam_ion_container = simulation.particles.get(self.beam_ions.name)
 
     def _create_data_arrays(self):
         self.prev_time = time.time()
@@ -347,8 +343,8 @@ def text_diag(self):
 
         status_dict = {
             "step": step,
-            "nplive beam ions": self.ion_container_wrapper.nps,
-            "nplive ions": self.beam_ion_container_wrapper.nps,
+            "nplive ions": self.ion_container.size,
+            "nplive beam ions": self.beam_ion_container.size,
             "wall_time": wall_time,
             "step_rate": step_rate,
             "diag_steps": self.diag_steps,
@@ -383,8 +379,8 @@ def energy_diagnostic(self):
             self._create_data_arrays()
 
         # get the simulation energies
-        Ec_par, Ec_perp = self._get_kinetic_energy(self.ion_container_wrapper)
-        Eb_par, Eb_perp = self._get_kinetic_energy(self.beam_ion_container_wrapper)
+        Ec_par, Ec_perp = self._get_kinetic_energy(self.ion_container)
+        Eb_par, Eb_perp = self._get_kinetic_energy(self.beam_ion_container)
 
         if libwarpx.amr.ParallelDescriptor.MyProc() != 0:
             return
@@ -400,18 +396,17 @@ def energy_diagnostic(self):
     def _get_kinetic_energy(self, container_wrapper):
         """Utility function to retrieve the total kinetic energy in the
         simulation."""
-        try:
-            ux = np.concatenate(container_wrapper.get_particle_ux())
-            uy = np.concatenate(container_wrapper.get_particle_uy())
-            uz = np.concatenate(container_wrapper.get_particle_uz())
-            w = np.concatenate(container_wrapper.get_particle_weight())
-        except ValueError:
-            return 0.0, 0.0
-
-        my_E_perp = 0.5 * self.M * np.sum(w * (ux**2 + uy**2))
-        E_perp = comm.allreduce(my_E_perp, op=mpi.SUM)
+        my_E_perp = 0
+        my_E_par = 0
+        for pti in container_wrapper.iterator(level=0):
+            ux = pti["ux"]
+            uy = pti["uy"]
+            uz = pti["uz"]
+            w = pti["w"]
+            my_E_perp += 0.5 * self.M * np.sum(w * (ux**2 + uy**2))
+            my_E_par += 0.5 * self.M * np.sum(w * uz**2)
 
-        my_E_par = 0.5 * self.M * np.sum(w * uz**2)
+        E_perp = comm.allreduce(my_E_perp, op=mpi.SUM)
         E_par = comm.allreduce(my_E_par, op=mpi.SUM)
 
         return E_par, E_perp

Examples/Tests/particle_boundary_interaction/inputs_test_rz_particle_boundary_interaction_picmi.py

@@ -148,16 +148,14 @@ def mirror_reflection():
     nz = concat(buffer.get_particle_scraped_this_step("electrons", "eb", "nz", lev))
 
     # STEP 2: use these parameters to inject particle from the same position in the plasma
-    elect_pc = particle_containers.ParticleContainerWrapper(
-        "electrons"
-    )  # general particle container
+    electrons = sim.particles.get("electrons")  # general particle container
 
     ####this part is specific to the case of simple reflection.
     un = ux * nx + uy * ny + uz * nz
     ux_reflect = -2 * un * nx + ux  # for a "mirror reflection" u(sym)=-2(u.n)n+u
     uy_reflect = -2 * un * ny + uy
     uz_reflect = -2 * un * nz + uz
-    elect_pc.add_particles(
+    electrons.add_particles(
         x=x + (dt - delta_t) * ux_reflect,
         y=y + (dt - delta_t) * uy_reflect,
         z=z + (dt - delta_t) * uz_reflect,

Examples/Tests/particle_data_python/inputs_test_2d_particle_attr_access_picmi.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 
-from pywarpx import callbacks, libwarpx, particle_containers, picmi
+from pywarpx import callbacks, libwarpx, picmi
 
 # Create the parser and add the argument
 parser = argparse.ArgumentParser()
@@ -106,8 +106,8 @@
 # below will be reproducible from run to run
 np.random.seed(30025025)
 
-elec_wrapper = particle_containers.ParticleContainerWrapper("electrons")
-elec_wrapper.add_real_comp("newPid")
+electrons = sim.particles.get("electrons")
+electrons.add_real_comp("newPid")
 
 my_id = libwarpx.amr.ParallelDescriptor.MyProc()
 
@@ -123,7 +123,7 @@ def add_particles():
     w = np.ones(nps) * 2.0
     newPid = 5.0
 
-    elec_wrapper.add_particles(
+    electrons.add_particles(
         x=x,
         y=y,
         z=z,
@@ -149,12 +149,12 @@ def add_particles():
 # are properly set
 ##########################
 
-assert elec_wrapper.nps == 270 / (2 - args.unique)
-assert elec_wrapper.particle_container.get_real_comp_index("w") == 2
-assert elec_wrapper.particle_container.get_real_comp_index("newPid") == 6
+assert electrons.size == 270 / (2 - args.unique)
+assert electrons.get_real_comp_index("w") == 2
+assert electrons.get_real_comp_index("newPid") == 6
 
-new_pid_vals = elec_wrapper.get_particle_real_arrays("newPid", 0)
-for vals in new_pid_vals:
+for pti in electrons.iterator(level=0):
+    vals = pti["newPid"]
     assert np.allclose(vals, 5)
 
 ##########################

Examples/Tests/particle_data_python/inputs_test_2d_prev_positions_picmi.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 
-from pywarpx import particle_containers, picmi
+from pywarpx import picmi
 
 constants = picmi.constants
 
@@ -107,23 +107,21 @@
 # exist
 ##########################
 
-elec_wrapper = particle_containers.ParticleContainerWrapper("electrons")
-elec_count = elec_wrapper.nps
+electrons = sim.particles.get("electrons")
 
 # check that the runtime attributes have the right indices
-assert elec_wrapper.particle_container.get_real_comp_index("prev_x") == 6
-assert elec_wrapper.particle_container.get_real_comp_index("prev_z") == 7
+assert electrons.get_real_comp_index("prev_x") == 6
+assert electrons.get_real_comp_index("prev_z") == 7
 
 # sanity check that the prev_z values are reasonable and
 # that the correct number of values are returned
-prev_z_vals = elec_wrapper.get_particle_real_arrays("prev_z", 0)
 running_count = 0
-
-for z_vals in prev_z_vals:
+for pti in electrons.iterator(level=0):
+    z_vals = pti["prev_z"]
     running_count += len(z_vals)
     assert np.all(z_vals < zmax)
 
-assert running_count == elec_wrapper.get_particle_count(True)
+assert running_count == electrons.number_of_particles(only_local=True)
 
 ##########################
 # take the final sim step

Examples/Tests/restart/inputs_test_2d_id_cpu_read_picmi.py

@@ -7,7 +7,7 @@
 
 import numpy as np
 
-from pywarpx import callbacks, particle_containers, picmi
+from pywarpx import callbacks, libwarpx, picmi
 
 ##########################
 # physics parameters
@@ -112,8 +112,8 @@
 np.random.seed(30025025)
 
 # wrap the electrons particle container
-electron_wrapper = particle_containers.ParticleContainerWrapper("electrons")
-electron_wrapper.add_real_comp("newPid")
+electrons = sim.particles.get("electrons")
+electrons.add_real_comp("newPid")
 
 
 def add_particles():
@@ -127,9 +127,7 @@ def add_particles():
     w = np.ones(nps) * 2.0
     newPid = 5.0
 
-    electron_wrapper.add_particles(
-        x=x, y=y, z=z, ux=ux, uy=uy, uz=uz, w=w, newPid=newPid
-    )
+    electrons.add_particles(x=x, y=y, z=z, ux=ux, uy=uy, uz=uz, w=w, newPid=newPid)
 
 
 callbacks.installbeforestep(add_particles)
@@ -145,5 +143,15 @@ def add_particles():
 # check that the ids and cpus are read properly
 ###############################################
 
-assert np.sum(np.concatenate(electron_wrapper.get_particle_id())) == 5050
-assert np.sum(np.concatenate(electron_wrapper.get_particle_cpu())) == 0
+ids_sum = 0
+cpu_sum = 0
+for pti in electrons.iterator(level=0):
+    idcpu = pti["idcpu"]
+    ids = libwarpx.amr.unpack_ids(idcpu)
+    cpu = libwarpx.amr.unpack_cpus(idcpu)
+
+    ids_sum += np.sum(ids)
+    cpu_sum += np.sum(cpu)
+
+assert ids_sum == 5050
+assert cpu_sum == 0