diff --git a/include/core/matrix_free_operator.h b/include/core/matrix_free_operator.h
index cd2c8a6e..7a587431 100644
--- a/include/core/matrix_free_operator.h
+++ b/include/core/matrix_free_operator.h
@@ -240,34 +240,42 @@ matrixFreeOperator<dim, degree, number>::local_compute_diagonal(
   [[maybe_unused]] const unsigned int         &dummy,
   const std::pair<unsigned int, unsigned int> &cell_range) const
 {
-  FEEvaluation<dim, degree, degree + 1, 1, number> phi(data);
+  // Field index
+  unsigned int field_index = 0;
 
-  AlignedVector<VectorizedArray<number>> diagonal(phi.dofs_per_cell);
+  // Constructor for FEEvaluation objects
+  variableContainer<dim, degree, number> variable_list(data);
+
+  variable_list.init_diagonal(field_index);
 
   for (unsigned int cell = cell_range.first; cell < cell_range.second; ++cell)
     {
-      phi.reinit(cell);
-      for (unsigned int i = 0; i < phi.dofs_per_cell; ++i)
+      variable_list.reinit_cell(cell);
+
+      for (unsigned int i = 0; i < variable_list.n_dofs_per_cell; ++i)
         {
-          for (unsigned int j = 0; j < phi.dofs_per_cell; ++j)
-            {
-              phi.submit_dof_value(VectorizedArray<number>(), j);
-            }
-          phi.submit_dof_value(make_vectorized_array<number>(1.), i);
+          variable_list.reinit_and_eval_diagonal(field_index, i);
 
-          phi.evaluate(EvaluationFlags::gradients);
-          for (const unsigned int q : phi.quadrature_point_indices())
+          // Number of quadrature points
+          unsigned int n_q_points = variable_list.get_num_q_points();
+
+          for (unsigned int q = 0; q < n_q_points; ++q)
             {
-              phi.submit_gradient(coefficient(cell, q) * phi.get_gradient(q), q);
+              // Set the quadrature point
+              variable_list.q_point = q;
+
+              // Grab the quadrature point location
+              Point<dim, VectorizedArray<number>> q_point_loc =
+                variable_list.get_q_point_location();
+
+              // Calculate the residuals
+              nonexplicit_RHS(variable_list, q_point_loc);
             }
-          phi.integrate(EvaluationFlags::gradients);
-          diagonal[i] = phi.get_dof_value(i);
-        }
-      for (unsigned int i = 0; i < phi.dofs_per_cell; ++i)
-        {
-          phi.submit_dof_value(diagonal[i], i);
+
+          variable_list.integrate_and_assemble_local_diagonal(field_index, i);
         }
-      phi.distribute_local_to_global(dst);
+
+      variable_list.distribute_diagonal(dst, field_index);
     }
 }
 
diff --git a/include/core/variableContainer.h b/include/core/variableContainer.h
index ad1efe87..491edfb7 100644
--- a/include/core/variableContainer.h
+++ b/include/core/variableContainer.h
@@ -71,21 +71,45 @@ class variableContainer
   integrate_and_distribute(dealii::LinearAlgebra::distributed::Vector<number> &dst);
 
   /**
-   * \brief Initialize the `AlignedVector` for the diagonal based on the dofs per cell.
+   * \brief Initialize the `AlignedVector` for the diagonal based on the DoFs per cell.
+   * Additionally, assign `n_dofs_per_cell` based on the current field index.
    */
   void
-  init_diagonal();
+  init_diagonal(unsigned int global_var_index);
 
   /**
-   * \brief Initialize the operation ptr to the cell .
+   * \brief Initialize the operation ptr to the cell for all field indices.
    */
   void
   reinit_cell(unsigned int cell);
 
+  /**
+   * \brief Submit the DoFs for the diagonal identity matrix and evaluate based on the
+   * evaluation flags.
+   */
+  void
+  reinit_and_eval_diagonal(unsigned int global_var_index, unsigned int i);
+
+  /**
+   * \brief Integrate the residuals and assemble the local diagonal
+   */
+  void
+  integrate_and_assemble_local_diagonal(unsigned int global_var_index, unsigned int i);
+
+  /**
+   * \brief Distribute the diagonal from local to global.
+   */
+  void
+  distribute_diagonal(dealii::LinearAlgebra::distributed::Vector<number> &dst,
+                      unsigned int global_var_index);
+
   // The quadrature point index, a method to get the number of quadrature points
   // per cell, and a method to get the xyz coordinates for the quadrature point
   unsigned int q_point = 0;
 
+  // Number of DoFs per cell
+  unsigned int n_dofs_per_cell;
+
   [[nodiscard]] unsigned int
   get_num_q_points() const;
 
@@ -115,7 +139,7 @@ class variableContainer
   // The number of variables
   unsigned int num_var;
 
-  // Diagonal
+  // Diagonal matrix that is used for preconditioning
   std::unique_ptr<dealii::AlignedVector<dealii::VectorizedArray<number>>> diagonal;
 };
 
diff --git a/src/core/variableContainer.cc b/src/core/variableContainer.cc
index 40fd09d5..21af241c 100644
--- a/src/core/variableContainer.cc
+++ b/src/core/variableContainer.cc
@@ -145,7 +145,26 @@ variableContainer<dim, degree, number>::integrate_and_distribute(
 
 template <int dim, int degree, typename number>
 void
-variableContainer<dim, degree, number>::init_diagonal()
+variableContainer<dim, degree, number>::init_diagonal(unsigned int global_var_index)
+{
+  const auto &var_info = varInfoList[global_var_index];
+
+  if (var_info.var_type == fieldType::SCALAR)
+    {
+      n_dofs_per_cell = scalar_vars_map.at(global_var_index)->dofs_per_cell;
+    }
+  else if (var_info.var_type == fieldType::VECTOR)
+    {
+      n_dofs_per_cell = vector_vars_map.at(global_var_index)->dofs_per_cell;
+    }
+
+  diagonal = std::make_unique<dealii::AlignedVector<dealii::VectorizedArray<number>>>(
+    n_dofs_per_cell);
+}
+
+template <int dim, int degree, typename number>
+void
+variableContainer<dim, degree, number>::reinit_cell(unsigned int cell)
 {
   for (unsigned int i = 0; i < num_var; i++)
     {
@@ -155,40 +174,96 @@ variableContainer<dim, degree, number>::init_diagonal()
 
       if (var_info.var_type == fieldType::SCALAR)
         {
-          diagonal =
-            std::make_unique<dealii::AlignedVector<dealii::VectorizedArray<number>>>(
-              scalar_vars_map.at(var_index)->dofs_per_cell);
+          scalar_vars_map.at(var_index).get()->reinit(cell);
         }
       else if (var_info.var_type == fieldType::VECTOR)
         {
-          diagonal =
-            std::make_unique<dealii::AlignedVector<dealii::VectorizedArray<number>>>(
-              vector_vars_map.at(var_index)->dofs_per_cell);
+          vector_vars_map.at(var_index).get()->reinit(cell);
         }
     }
 }
 
 template <int dim, int degree, typename number>
 void
-variableContainer<dim, degree, number>::reinit_cell(unsigned int cell)
+variableContainer<dim, degree, number>::reinit_and_eval_diagonal(
+  unsigned int global_var_index,
+  unsigned int i)
 {
-  for (unsigned int i = 0; i < num_var; i++)
+  const auto &var_info = varInfoList[global_var_index];
+
+  if (var_info.var_type == fieldType::SCALAR)
     {
-      const auto &var_info = varInfoList[i];
+      auto *scalar_FEEval_ptr = scalar_vars_map.at(global_var_index).get();
+      for (unsigned int j = 0; j < n_dofs_per_cell; ++j)
+        {
+          scalar_FEEval_ptr->submit_dof_value(dealii::VectorizedArray<number>(), j);
+        }
+      scalar_FEEval_ptr->submit_dof_value(dealii::make_vectorized_array<number>(1.0), i);
+    }
+  else if (var_info.var_type == fieldType::VECTOR)
+    {
+      AssertThrow(false, FeatureNotImplemented("Vector multigrid"));
+    }
 
+  for (unsigned int i = 0; i < num_var; i++)
+    {
+      const auto        &var_info  = varInfoList[i];
       const unsigned int var_index = var_info.global_var_index;
 
       if (var_info.var_type == fieldType::SCALAR)
         {
-          scalar_vars_map.at(var_index).get()->reinit(cell);
+          scalar_vars_map.at(var_index)->evaluate(var_info.evaluation_flags);
         }
       else if (var_info.var_type == fieldType::VECTOR)
         {
-          vector_vars_map.at(var_index).get()->reinit(cell);
+          vector_vars_map.at(var_index)->evaluate(var_info.evaluation_flags);
         }
     }
 }
 
+template <int dim, int degree, typename number>
+void
+variableContainer<dim, degree, number>::integrate_and_assemble_local_diagonal(
+  unsigned int global_var_index,
+  unsigned int i)
+{
+  const auto &var_info = varInfoList[global_var_index];
+
+  if (var_info.var_type == fieldType::SCALAR)
+    {
+      auto *scalar_FEEval_ptr = scalar_vars_map.at(global_var_index).get();
+      scalar_FEEval_ptr->integrate(var_info.residual_flags);
+      (*diagonal)[i] = scalar_FEEval_ptr->get_dof_value(i);
+    }
+  else if (var_info.var_type == fieldType::VECTOR)
+    {
+      AssertThrow(false, FeatureNotImplemented("Vector multigrid"));
+    }
+}
+
+template <int dim, int degree, typename number>
+void
+variableContainer<dim, degree, number>::distribute_diagonal(
+  dealii::LinearAlgebra::distributed::Vector<number> &dst,
+  unsigned int                                        global_var_index)
+{
+  const auto &var_info = varInfoList[global_var_index];
+
+  if (var_info.var_type == fieldType::SCALAR)
+    {
+      auto *scalar_FEEval_ptr = scalar_vars_map.at(global_var_index).get();
+      for (unsigned int i = 0; i < n_dofs_per_cell; ++i)
+        {
+          scalar_FEEval_ptr->submit_dof_value((*diagonal)[i], i);
+        }
+      scalar_FEEval_ptr->distribute_local_to_global(dst);
+    }
+  else if (var_info.var_type == fieldType::VECTOR)
+    {
+      AssertThrow(false, FeatureNotImplemented("Vector multigrid"));
+    }
+}
+
 template <int dim, int degree, typename number>
 dealii::VectorizedArray<number>
 variableContainer<dim, degree, number>::get_scalar_value(