Update spacecraft charging example

Examples/Physics_applications/spacecraft_charging/inputs_test_rz_spacecraft_charging_picmi.py

@@ -35,16 +35,18 @@ def __init__(self):
         self.spacecraft_capacitance = None
 
         # shortcuts
-        self.Direction = libwarpx.libwarpx_so.Direction
-        self.dir_r, self.dir_z = self.Direction(0), self.Direction(2)
+        self.dir_r, self.dir_z = (
+            libwarpx.libwarpx_so.Direction(0),
+            libwarpx.libwarpx_so.Direction(2),
+        )
+
 
     def correct_space_charge_fields(self, q=None):
         """
         Function that will be called at each iteration,
         after each electrostatic solve in WarpX
         """
         assert self.saved_first_iteration_fields
-        warpx = sim.extension.warpx
 
         # Compute the charge that WarpX thinks there is on the spacecraft
         # from phi and rho after the Poisson solver
@@ -53,12 +55,17 @@ def correct_space_charge_fields(self, q=None):
             q = compute_actual_charge_on_spacecraft()
 
         # Correct fields so as to recover the actual charge
-        Er = warpx.multifab("Efield_fp", dir=self.dir_r, level=0)
-        Er.saxpy(q - q_v, self.normalized_Er, 0, 0, 1, 0)
-        Ez = warpx.multifab("Efield_fp", dir=self.dir_z, level=0)
-        Ez.saxpy(q - q_v, self.normalized_Ez, 0, 0, 1, 0)
-        phi = warpx.multifab("phi_fp", level=0)
-        phi.saxpy(q - q_v, self.normalized_phi, 0, 0, 1, 0)
+        warpx = sim.extension.warpx
+        fields = warpx.multifab_register()
+        Er = fields.get("Efield_fp", self.dir_r, 0)
+        normalized_Er = fields.get("normalized_Er", 0)
+        Er.saxpy(q - q_v, normalized_Er, 0, 0, 1, 0)
+        Ez = fields.get("Efield_fp", self.dir_z, 0)
+        normalized_Ez = fields.get("normalized_Ez", 0)
+        Ez.saxpy(q - q_v, normalized_Ez, 0, 0, 1, 0)
+        phi = fields.get("phi_fp", 0)
+        normalized_phi = fields.get("normalized_phi", 0)
+        phi.saxpy(q - q_v, normalized_phi, 0, 0, 1, 0)
 
         self.spacecraft_potential += (q - q_v) * self.spacecraft_capacitance
         warpx.set_potential_on_eb("%f" % self.spacecraft_potential)
@@ -76,54 +83,60 @@ def save_normalized_vacuum_Efields(
         self.spacecraft_capacitance = 1.0 / q_v  # the potential was set to 1V
 
         warpx = sim.extension.warpx
-        amrex = sim.extension.amr
-        reg = warpx.multifab_register()
-
-        self.phi = warpx.multifab("phi_fp", 0)
-        self.Er = warpx.multifab("Efield_fp", dir=self.dir_r, level=0)
-        self.Ez = warpx.multifab("Efield_fp", dir=self.dir_z, level=0)
-        # allocate and initialize normalized fields
-        self.normalized_Er = reg.alloc_init(
+        fields = warpx.multifab_register()
+
+        phi = fields.get("phi_fp", 0)
+        Er = fields.get("Efield_fp", self.dir_r, 0)
+        Ez = fields.get("Efield_fp", self.dir_z, 0)
+        # Allocate the fields `normalized_Er`, `normalized_Ez`, and `normalized_phi
+        # in WarpX's multifab register. This allows to get these fields at later
+        # iterations with fields.get( ... ).
+        # These new fields are automatically redistributed when doing load balancing.
+        normalized_Er = fields.alloc_init(
             "normalized_Er",
             0,
-            warpx.boxArray(0).convert(amrex.IntVect(0, 1)),
+            Er.box_array(),
             warpx.DistributionMap(0),
             1,
-            self.Er.n_grow_vect,
+            Er.n_grow_vect,
             0.0,
             True,
             True,
         )
-        self.normalized_Ez = reg.alloc_init(
+
+        normalized_Ez = fields.alloc_init(
             "normalized_Ez",
             0,
-            warpx.boxArray(0).convert(amrex.IntVect(1, 0)),
+            Ez.box_array(),
             warpx.DistributionMap(0),
             1,
-            self.Ez.n_grow_vect,
+            Ez.n_grow_vect,
             0.0,
             True,
             True,
         )
-        self.normalized_phi = reg.alloc_init(
+
+        normalized_phi = fields.alloc_init(
             "normalized_phi",
             0,
-            warpx.boxArray(0).convert(amrex.IntVect(1, 1)),
+            phi.box_array(),
             warpx.DistributionMap(0),
             1,
-            self.phi.n_grow_vect,
+            phi.n_grow_vect,
             0.0,
             True,
             True,
         )
 
         # Record fields
-        self.normalized_Er.copymf(self.Er, 0, 0, 1, self.Er.n_grow_vect)
-        self.normalized_Er.mult(1 / q_v, 0)
-        self.normalized_Ez.copymf(self.Ez, 0, 0, 1, self.Ez.n_grow_vect)
-        self.normalized_Ez.mult(1 / q_v, 0)
-        self.normalized_phi.copymf(self.phi, 0, 0, 1, self.phi.n_grow_vect)
-        self.normalized_phi.mult(1 / q_v, 0)
+
+        normalized_Er.copymf(Er, 0, 0, 1, Er.n_grow_vect)
+        normalized_Er.mult(1 / q_v, 0)
+        normalized_Ez.copymf(Ez, 0, 0, 1, Ez.n_grow_vect)
+        normalized_Ez.mult(1 / q_v, 0)
+        normalized_phi.copymf(phi, 0, 0, 1, phi.n_grow_vect)
+        normalized_phi.mult(1 / q_v, 0)
+
 
         self.saved_first_iteration_fields = True
         self.correct_space_charge_fields(q=0)
@@ -137,8 +150,9 @@ def compute_virtual_charge_on_spacecraft():
     that WarpX thinks there should be on the spacecraft.
     """
     warpx = sim.extension.warpx
-    rho = warpx.multifab("rho_fp", level=0)
-    phi = warpx.multifab("phi_fp", level=0)
+    fields = warpx.multifab_register()
+    rho = fields.get("rho_fp", 0)
+    phi = fields.get("phi_fp", 0)
 
     dr, dz = warpx.Geom(lev=0).data().CellSize()