Allow integrating hypre AMS/ADS in FSPs

Author: nmnobre

src/numerics/petsc_preconditioner.C

@@ -183,10 +183,6 @@ void PetscPreconditioner<T>::set_petsc_aux_data(PC & pc, System & sys, const uns
     // If not ams/ads, we quit with nothing to do
     if (std::string(hypre_type) == "ams")
     {
-      // If multiple variables, we error out as senseless
-      libmesh_error_msg_if(sys.n_vars() > 1,
-                           "Error applying hypre AMS to a system with multiple "
-                           "variables");
       // If not a 1st order Nédélec or a 2d 1st order Raviart-Thomas system, we
       // error out as we do not support anything else at the moment
       libmesh_error_msg_if(sys.variable(v).type() != FEType(1, NEDELEC_ONE) &&
@@ -199,10 +195,6 @@ void PetscPreconditioner<T>::set_petsc_aux_data(PC & pc, System & sys, const uns
     }
     else if (std::string(hypre_type) == "ads")
     {
-      // If multiple variables, we error out as senseless
-      libmesh_error_msg_if(sys.n_vars() > 1,
-                           "Error applying hypre ADS to a system with multiple "
-                           "variables");
       // If not a 3d 1st order Raviart-Thomas system, we error out as we do not
       // support anything else at the moment
       libmesh_error_msg_if(sys.variable(v).type() != FEType(1, RAVIART_THOMAS) ||
@@ -213,6 +205,31 @@ void PetscPreconditioner<T>::set_petsc_aux_data(PC & pc, System & sys, const uns
       set_hypre_ads_data(pc, sys, v);
     }
   }
+  else if (pc_type && std::string(pc_type) == PCFIELDSPLIT)
+  {
+    // Annoyingly, if using Schur complement preconditioning, we need to call PCSetUp()
+    auto pmatrix = cast_ptr<PetscMatrixBase<T> *>(sys.request_matrix("System Matrix"));
+    pmatrix->close();
+    Mat mat = pmatrix->mat();
+    LibmeshPetscCallA(comm, PCSetOperators(pc, mat, mat));
+    LibmeshPetscCallA(comm, PCSetUp(pc));
+
+    // Get sub KSP contexts for all splits
+    PetscInt n_splits;
+    KSP * subksps;
+    LibmeshPetscCallA(comm, PCFieldSplitGetSubKSP(pc, &n_splits, &subksps));
+
+    // Get the sub PC context for each split and recursively call this function
+    for (auto s : make_range(n_splits))
+    {
+      PC subpc;
+      LibmeshPetscCallA(comm, KSPGetPC(subksps[s], &subpc));
+      set_petsc_aux_data(subpc, sys, s);
+    }
+
+    // Free the array of sub KSP contexts
+    LibmeshPetscCallA(comm, PetscFree(subksps));
+  }
 }
 
 

src/solvers/petsc_nonlinear_solver.C

@@ -514,6 +514,12 @@ extern "C"
         Jac->close();
       }
 
+    KSP kspc;
+    PC pcc;
+    LibmeshPetscCall2(sys.comm(), SNESGetKSP(snes, &kspc));
+    LibmeshPetscCall2(sys.comm(), KSPGetPC(kspc, &pcc));
+    PetscPreconditioner<Number>::set_petsc_aux_data(pcc, sys);
+
     PetscFunctionReturn(LIBMESH_PETSC_SUCCESS);
   }
 
@@ -985,10 +991,6 @@ PetscNonlinearSolver<T>::solve (SparseMatrix<T> &  pre_in,  // System Preconditi
 #endif
   LibmeshPetscCall(SNESSetFromOptions(_snes));
 
-  PC pc;
-  LibmeshPetscCall(KSPGetPC(ksp, &pc));
-  PetscPreconditioner<T>::set_petsc_aux_data(pc, this->system());
-
 #if defined(LIBMESH_HAVE_PETSC_HYPRE) && PETSC_VERSION_LESS_THAN(3, 23, 0) &&                      \
     !PETSC_VERSION_LESS_THAN(3, 12, 0) && defined(PETSC_HAVE_HYPRE_DEVICE)
   {