Fix load balancing with FEM

alpine/LoadBalancer.hpp

@@ -63,6 +63,18 @@ class LoadBalancer {
             fs_m->getStype() == "FEM_PRECON") {
             phi_m->updateLayout(*fl);
             phi_m->setFieldBC(phi_m->getFieldBC());
+
+            if (fs_m->getStype() == "FEM") {
+                // also update the layout in the FEM space
+                auto& space = std::get<FEMSolver_t<T, Dim>>(fs_m->getSolver()).getSpace();
+                space.updateLayout(*fl);
+            } 
+
+            if (fs_m->getStype() == "FEM_PRECON") {
+                // also update the layout in the FEM space
+                auto& space = std::get<FEMPreconSolver_t<T, Dim>>(fs_m->getSolver()).getSpace();
+                space.updateLayout(*fl);
+            } 
         }
 
         // Update layout with new FieldLayout

src/FEM/LagrangeSpace.h

@@ -110,6 +110,14 @@ namespace ippl {
          */
         void initializeElementIndices(Layout_t& layout);
 
+        ///////////////////////////////////////////////////////////////////////
+        /**
+         * @brief Function to update the element partition and the layout of 
+         * fields in the LagrangeSpace if the layout has been changed during
+         * the simulation (for example by the load balancer).
+         */
+        void updateLayout(Layout_t& layout);
+
         /// Degree of Freedom operations //////////////////////////////////////
         ///////////////////////////////////////////////////////////////////////
 

src/FEM/LagrangeSpace.hpp

@@ -106,6 +106,17 @@ namespace ippl {
             });
     }
 
+    // Update resultField and elementIndices according to changed domain decomposition.
+    template <typename T, unsigned Dim, unsigned Order, typename ElementType,
+              typename QuadratureType, typename FieldLHS, typename FieldRHS>
+    void LagrangeSpace<T, Dim, Order, ElementType, QuadratureType, FieldLHS,
+                       FieldRHS>::updateLayout(Layout_t& layout) {
+        // repartition elements
+        initializeElementIndices(layout);
+        // update layout of resultField member variable
+        resultField.updateLayout(layout);
+    }
+
     ///////////////////////////////////////////////////////////////////////
     /// Degree of Freedom operations //////////////////////////////////////
     ///////////////////////////////////////////////////////////////////////