Exposing mf register (BLAST-WarpX#5846)

This PR exposes the multifab register so that when we create a local
copy using pyamrex interface in a python script,
1. the multifab gets allocated and initialized
2. this allows the multifab to also be redistributed during load
balancing
3. when finalize is called, the multifab memory is cleared in a clean
way


This PR AMReX-Codes/pyamrex#439) adds wrappers
to ba.convert, ba.surroundingNodes, mf.copy

so now , we could add a new "test_mf" with desired nodality
```
        reg.alloc_init("test_nf",0,warpx.boxArray(0).convert(amrex.IntVect(1,1)),warpx.DistributionMap(0),
                       1,amrex.IntVect(1,1),0.0,True,True)
```


<img width="637" alt="Screenshot 2025-05-14 at 10 58 19 AM"
src="https://github.com/user-attachments/assets/ecc3a5ae-9ce1-4b32-82ad-169167d8dfdb"
/>

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Remi Lehe <[email protected]>

Author: 3 people

Examples/Physics_applications/spacecraft_charging/inputs_test_rz_spacecraft_charging_picmi.py

@@ -35,16 +35,17 @@ 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 +54,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)
@@ -70,20 +76,65 @@ def correct_space_charge_fields(self, q=None):
     def save_normalized_vacuum_Efields(
         self,
     ):
-        warpx = sim.extension.warpx
-
         # Compute the charge that WarpX thinks there is on the spacecraft
         # from phi and rho after the Poisson solver
         q_v = compute_virtual_charge_on_spacecraft()
         self.spacecraft_capacitance = 1.0 / q_v  # the potential was set to 1V
 
+        warpx = sim.extension.warpx
+        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,
+            Er.box_array(),
+            warpx.DistributionMap(0),
+            1,
+            Er.n_grow_vect,
+            0.0,
+            True,
+            True,
+        )
+
+        normalized_Ez = fields.alloc_init(
+            "normalized_Ez",
+            0,
+            Ez.box_array(),
+            warpx.DistributionMap(0),
+            1,
+            Ez.n_grow_vect,
+            0.0,
+            True,
+            True,
+        )
+
+        normalized_phi = fields.alloc_init(
+            "normalized_phi",
+            0,
+            phi.box_array(),
+            warpx.DistributionMap(0),
+            1,
+            phi.n_grow_vect,
+            0.0,
+            True,
+            True,
+        )
+
         # Record fields
-        self.normalized_Er = warpx.multifab("Efield_fp", dir=self.dir_r, level=0).copy()
-        self.normalized_Er.mult(1 / q_v, 0)
-        self.normalized_Ez = warpx.multifab("Efield_fp", dir=self.dir_z, level=0).copy()
-        self.normalized_Ez.mult(1 / q_v, 0)
-        self.normalized_phi = warpx.multifab("phi_fp", level=0).copy()
-        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)
@@ -97,8 +148,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()
 

Source/Python/MultiFabRegister.cpp

@@ -32,6 +32,7 @@ void init_MultiFabRegister (py::module & m)
                  bool,
                  bool
              >(&MultiFabRegister::alloc_init<std::string>),
+             py::return_value_policy::reference_internal,
              py::arg("name"),
              py::arg("level"),
              py::arg("ba"),
@@ -56,6 +57,7 @@ void init_MultiFabRegister (py::module & m)
                  bool,
                  bool
              >(&MultiFabRegister::alloc_init<std::string>),
+             py::return_value_policy::reference_internal,
              py::arg("name"),
              py::arg("dir"),
              py::arg("level"),
@@ -75,6 +77,7 @@ void init_MultiFabRegister (py::module & m)
                  int,
                  std::optional<const amrex::Real>
              >(&MultiFabRegister::alias_init<std::string, std::string>),
+             py::return_value_policy::reference_internal,
              py::arg("new_name"),
              py::arg("alias_name"),
              py::arg("level"),
@@ -89,6 +92,7 @@ void init_MultiFabRegister (py::module & m)
                  int,
                  std::optional<const amrex::Real>
              >(&MultiFabRegister::alias_init<std::string, std::string>),
+             py::return_value_policy::reference_internal,
              py::arg("new_name"),
              py::arg("alias_name"),
              py::arg("dir"),
@@ -121,6 +125,7 @@ void init_MultiFabRegister (py::module & m)
                  std::string,
                  int
              >(&MultiFabRegister::get<std::string>),
+             py::return_value_policy::reference_internal,
              py::arg("name"),
              py::arg("level")
         )
@@ -131,6 +136,7 @@ void init_MultiFabRegister (py::module & m)
                  ablastr::fields::Direction,
                  int
              >(&MultiFabRegister::get<std::string>),
+             py::return_value_policy::reference_internal,
              py::arg("name"),
              py::arg("dir"),
              py::arg("level")

Source/Python/WarpX.cpp

@@ -113,6 +113,8 @@ void init_WarpX (py::module& m)
             //py::overload_cast< int >(&WarpX::boxArray, py::const_),
             py::arg("lev")
         )
+        .def("multifab_register",&WarpX::GetMultiFabRegister,
+            py::return_value_policy::reference_internal)
         .def("multifab",
              [](WarpX & wx, std::string internal_name) {
                  if (wx.m_fields.internal_has(internal_name)) {

Source/WarpX.H

@@ -953,6 +953,7 @@ public:
 
     // Field container
     ablastr::fields::MultiFabRegister m_fields;
+    ablastr::fields::MultiFabRegister& GetMultiFabRegister() {return m_fields;}
 
 protected: