Actions to apply the approximate inverse of custom Schur complement approximations

src/smith/differentiable_numerics/tests/test_nonlinear_mixed_diffusion.cpp

@@ -89,14 +89,17 @@ struct ThermalConductivity {
 
 enum class PreconditionerCase
 {
-  StateDependentCustomFullSchur
+  StateDependentCustomFullSchur,
+  StateDependentCustomActionSchur
 };
 
 const char* preconditionerCaseName(PreconditionerCase preconditioner)
 {
   switch (preconditioner) {
     case PreconditionerCase::StateDependentCustomFullSchur:
       return "StateDependentCustomFullSchur";
+    case PreconditionerCase::StateDependentCustomActionSchur:
+      return "StateDependentCustomActionSchur";
   }
   return "Unknown";
 }
@@ -135,6 +138,59 @@ class MeshFixture : public testing::Test {
 class NonlinearMixedDiffusionPreconditionerTest : public MeshFixture,
                                                   public ::testing::WithParamInterface<PreconditionerCase> {};
 
+/**
+ * @brief Custom Schur action solver that applies an approximate inverse of
+ *        the temperature diffusion Schur approximation.
+ *
+ * This exercises the CustomAction path: the block Schur preconditioner does not
+ * receive a custom Schur operator. Instead, the block-1 solver owns the
+ * state-dependent operator and applies an internal AMG-preconditioned GMRES solve.
+ */
+class TemperatureDiffusionSchurActionSolver : public smith::SchurComplementActionSolver {
+ public:
+  TemperatureDiffusionSchurActionSolver(std::unique_ptr<mfem::HypreParMatrix> initial_schur_operator,
+                                        smith::StateDependentWeakFormOperator schur_operator_update)
+      : schur_operator_update_(std::move(schur_operator_update)), schur_operator_(std::move(initial_schur_operator))
+  {
+    configureLinearSolver();
+  }
+
+  void updateForState(const mfem::Vector& state, const mfem::Array<int>& block_offsets) override
+  {
+    schur_operator_ = schur_operator_update_(state, block_offsets);
+    configureLinearSolver();
+  }
+
+  void Mult(const mfem::Vector& x, mfem::Vector& y) const override
+  {
+    y.SetSize(x.Size());
+    y = 0.0;
+    linear_solver_->Mult(x, y);
+  }
+
+ private:
+  void configureLinearSolver()
+  {
+    amg_solver_ = std::make_unique<mfem::HypreBoomerAMG>();
+    amg_solver_->SetPrintLevel(0);
+    amg_solver_->SetOperator(*schur_operator_);
+
+    linear_solver_ = std::make_unique<mfem::GMRESSolver>(schur_operator_->GetComm());
+    linear_solver_->SetRelTol(1.0e-8);
+    linear_solver_->SetAbsTol(0.0);
+    linear_solver_->SetMaxIter(100);
+    linear_solver_->SetPrintLevel(0);
+    linear_solver_->SetPreconditioner(*amg_solver_);
+    linear_solver_->SetOperator(*schur_operator_);
+    linear_solver_->iterative_mode = false;
+  }
+
+  smith::StateDependentWeakFormOperator schur_operator_update_;
+  std::unique_ptr<mfem::HypreParMatrix> schur_operator_;
+  std::unique_ptr<mfem::HypreBoomerAMG> amg_solver_;
+  std::unique_ptr<mfem::GMRESSolver> linear_solver_;
+};
+
 // Exercises state-dependent Schur preconditioners on nonlinear mixed diffusion.
 TEST_P(NonlinearMixedDiffusionPreconditionerTest, BlockSolve)
 {
@@ -242,6 +298,22 @@ TEST_P(NonlinearMixedDiffusionPreconditionerTest, BlockSolve)
         std::make_unique<smith::BlockSchurPreconditioner>(std::move(sub_solvers), linear_options.block_schur_type,
                                                           linear_options.schur_approx_type, std::move(overrides));
     nonlinear_solver = smith::buildNonlinearBlockSolver(nonlin_opts, linear_options, *mesh, std::move(preconditioner));
+  } else if (preconditioner_case == PreconditionerCase::StateDependentCustomActionSchur) {
+    linear_options.schur_approx_type = smith::SchurApproxType::CustomAction;
+
+    const std::vector<smith::LinearSolverOptions> sub_solver_options{flux_solver_options};
+    auto sub_solvers = smith::buildBlockPreconditionerSubSolvers(sub_solver_options, mesh->getComm());
+
+    auto action_solver = std::make_unique<TemperatureDiffusionSchurActionSolver>(
+        smith::buildWeakFormOperator(temperature_schur_form, shape_disp, {temperature}, {1.0}, time_info,
+                                     temp_bc_manager.get()),
+        smith::makeStateDependentWeakFormOperator(temperature_schur_form, shape_disp, {temperature}, {1.0}, time_info,
+                                                  temp_bc_manager.get(), {smith::StateBlockBinding{1, 0}}));
+    sub_solvers.push_back(std::move(action_solver));
+
+    auto preconditioner = std::make_unique<smith::BlockSchurPreconditioner>(
+        std::move(sub_solvers), linear_options.block_schur_type, linear_options.schur_approx_type);
+    nonlinear_solver = smith::buildNonlinearBlockSolver(nonlin_opts, linear_options, *mesh, std::move(preconditioner));
   }
 
   ASSERT_TRUE(nonlinear_solver != nullptr);
@@ -267,7 +339,8 @@ TEST_P(NonlinearMixedDiffusionPreconditionerTest, BlockSolve)
 }
 
 INSTANTIATE_TEST_SUITE_P(StateDependentSchur, NonlinearMixedDiffusionPreconditionerTest,
-                         ::testing::Values(PreconditionerCase::StateDependentCustomFullSchur),
+                         ::testing::Values(PreconditionerCase::StateDependentCustomFullSchur,
+                                           PreconditionerCase::StateDependentCustomActionSchur),
                          preconditionerCaseNameGenerator);
 
 int main(int argc, char* argv[])

src/smith/numerics/block_preconditioner.cpp

@@ -342,6 +342,27 @@ BlockSchurPreconditioner::BlockSchurPreconditioner(std::vector<std::unique_ptr<m
   }
 }
 
+void SchurComplementActionSolver::setBlockContext(const mfem::BlockOperator& jacobian,
+                                                  const mfem::Array<int>& block_offsets)
+{
+  MFEM_VERIFY(block_offsets.Size() == 3, "SchurComplementActionSolver requires 2x2 block offsets");
+  block_jacobian_ = &jacobian;
+  block_offsets_ = block_offsets;
+  height = block_offsets[2] - block_offsets[1];
+  width = height;
+}
+
+void SchurComplementActionSolver::SetOperator(const mfem::Operator& op)
+{
+  height = op.Height();
+  width = op.Width();
+}
+
+void SchurComplementActionSolver::updateForState([[maybe_unused]] const mfem::Vector& state,
+                                                 [[maybe_unused]] const mfem::Array<int>& block_offsets)
+{
+}
+
 void BlockSchurPreconditioner::LowerBlock(const mfem::Vector& in, mfem::Vector& out) const
 {
   // Interpret in, out as block vectors: in = [b1; b2], out = [x1; x2]
@@ -540,12 +561,19 @@ void BlockSchurPreconditioner::SetOperator(const mfem::Operator& jacobian)
   } else if (approxType_ == SchurApproxType::Custom) {
     S_approx_owned_.reset();
     S_approx_view_ = &block_op_providers_[1]->currentOperator();
+  } else if (approxType_ == SchurApproxType::CustomAction) {
+    S_approx_owned_.reset();
+    auto* action_solver = dynamic_cast<SchurComplementActionSolver*>(mfem_solvers_[1].get());
+    MFEM_VERIFY(action_solver, "SchurApproxType::CustomAction requires a SchurComplementActionSolver for block 1");
+    action_solver->setBlockContext(*block_jacobian_, block_offsets_);
   }
 
-  MFEM_VERIFY(S_approx_view_, "Schur complement approximation operator must be set");
+  if (approxType_ != SchurApproxType::CustomAction) {
+    MFEM_VERIFY(S_approx_view_, "Schur complement approximation operator must be set");
 
-  // Set the Schur complement preconditioner for block (1,1)
-  mfem_solvers_[1]->SetOperator(*S_approx_view_);
+    // Set the Schur complement preconditioner for block (1,1)
+    mfem_solvers_[1]->SetOperator(*S_approx_view_);
+  }
   mfem_solvers_[1]->iterative_mode = false;
 
   // Set up block diagonal operator

src/smith/numerics/block_preconditioner.hpp

@@ -274,9 +274,37 @@ enum class BlockSchurType
  */
 enum class SchurApproxType
 {
-  DiagInv, /**< Use assembled \f$ S \approx A_{22} - A_{21} \\mathrm{diag}(A_{11})^{-1} A_{12} \f$. */
-  A22Only, /**< Use \f$ S \approx A_{22} \f$. */
-  Custom   /**< Use a custom operator provider for block index 1. */
+  DiagInv,     /**< Use assembled \f$ S \approx A_{22} - A_{21} \\mathrm{diag}(A_{11})^{-1} A_{12} \f$. */
+  A22Only,     /**< Use \f$ S \approx A_{22} \f$. */
+  Custom,      /**< Use a custom operator provider for block index 1. */
+  CustomAction /**< Use a custom block-1 solver that applies the Schur inverse/action directly. */
+};
+
+/**
+ * @class SchurComplementActionSolver
+ * @brief Base class for custom solvers that apply an approximate Schur inverse/action directly.
+ *
+ * Subclasses normally override updateForState() to refresh any state-dependent
+ * approximation and Mult() to apply the approximate Schur inverse/action. Override
+ * setBlockContext() only when the current 2x2 block Jacobian is needed. In a
+ * nonlinear solve, updateForState() is called before Jacobian assembly and
+ * setBlockContext() is called later from BlockSchurPreconditioner::SetOperator().
+ */
+class SchurComplementActionSolver : public mfem::Solver, public StateDependentSolver {
+ public:
+  /// @brief Configure this action solver with the current 2x2 block context.
+  virtual void setBlockContext(const mfem::BlockOperator& jacobian, const mfem::Array<int>& block_offsets);
+
+  /// @overload
+  void SetOperator(const mfem::Operator& op) override;
+
+  /// @overload
+  void updateForState([[maybe_unused]] const mfem::Vector& state,
+                      [[maybe_unused]] const mfem::Array<int>& block_offsets) override;
+
+ protected:
+  const mfem::BlockOperator* block_jacobian_ = nullptr;  ///< Current 2x2 block Jacobian context.
+  mfem::Array<int> block_offsets_;                       ///< Current block offsets.
 };
 
 /**
@@ -334,7 +362,7 @@ class BlockSchurPreconditioner : public BlockPreconditioner {
   //
   // For DiagInv and A22Only, the approximation is rebuilt on each SetOperator call and stored in
   // S_approx_owned_. For Custom, the approximation is provided via block_op_providers_[1] and referenced
-  // non-owningly via S_approx_view_.
+  // non-owningly via S_approx_view_. For CustomAction, the block-1 solver is configured directly.
   mutable std::unique_ptr<const mfem::Operator> S_approx_owned_;
   const mfem::Operator* S_approx_view_ = nullptr;
 

src/smith/numerics/tests/test_block_preconditioner_custom_operators.cpp

@@ -287,6 +287,38 @@ class OperatorDiagonalSolver : public mfem::Solver {
   mfem::Vector diag_;
 };
 
+class MockSchurActionSolver : public smith::SchurComplementActionSolver {
+ public:
+  MockSchurActionSolver(double scale, std::shared_ptr<int> context_calls, std::shared_ptr<int> update_calls)
+      : scale_(scale), context_calls_(std::move(context_calls)), update_calls_(std::move(update_calls))
+  {
+  }
+
+  void setBlockContext(const mfem::BlockOperator& jacobian, const mfem::Array<int>& block_offsets) override
+  {
+    smith::SchurComplementActionSolver::setBlockContext(jacobian, block_offsets);
+    ++(*context_calls_);
+  }
+
+  void updateForState([[maybe_unused]] const mfem::Vector& state,
+                      [[maybe_unused]] const mfem::Array<int>& block_offsets) override
+  {
+    ++(*update_calls_);
+  }
+
+  void Mult(const mfem::Vector& x, mfem::Vector& y) const override
+  {
+    y.SetSize(x.Size());
+    y = x;
+    y *= scale_;
+  }
+
+ private:
+  double scale_;
+  std::shared_ptr<int> context_calls_;
+  std::shared_ptr<int> update_calls_;
+};
+
 }  // namespace
 /* ============================================================
    Tests
@@ -763,6 +795,50 @@ TEST(BlockSchurPreconditionerCustom, StateDependentProviderUpdatesSchurSolve)
   EXPECT_NEAR(x[3], b[3] / 7.0, 1e-12);
 }
 
+// Verifies custom Schur action solvers can replace explicit Schur operators.
+TEST(BlockSchurPreconditionerCustom, CustomActionUsesActionSolverWithoutSchurOperatorProvider)
+{
+  constexpr int n = 2;
+  Array<int> offsets({0, n, 2 * n});
+
+  auto A11o = makeHypreScaledIdentity(n, 2.0);
+  auto A12o = makeHypreScaledIdentity(n, 0.0);
+  auto A21o = makeHypreScaledIdentity(n, 0.0);
+  auto A22o = makeHypreScaledIdentity(n, 3.0);
+
+  BlockOperator A(offsets);
+  A.SetBlock(0, 0, A11o.A.get());
+  A.SetBlock(0, 1, A12o.A.get());
+  A.SetBlock(1, 0, A21o.A.get());
+  A.SetBlock(1, 1, A22o.A.get());
+
+  auto context_calls = std::make_shared<int>(0);
+  auto update_calls = std::make_shared<int>(0);
+
+  std::vector<std::unique_ptr<Solver>> solvers;
+  solvers.push_back(std::make_unique<IdentitySolver>());
+  solvers.push_back(std::make_unique<MockSchurActionSolver>(0.25, context_calls, update_calls));
+
+  smith::BlockSchurPreconditioner P(std::move(solvers), smith::BlockSchurType::Diagonal,
+                                    smith::SchurApproxType::CustomAction);
+
+  Vector state(2 * n);
+  state = 0.0;
+  P.updateForState(state, offsets);
+  P.SetOperator(A);
+
+  Vector b(2 * n), x(2 * n);
+  b.Randomize();
+  P.Mult(b, x);
+
+  EXPECT_EQ(*update_calls, 1);
+  EXPECT_EQ(*context_calls, 1);
+  EXPECT_NEAR(x[0], b[0], 1e-12);
+  EXPECT_NEAR(x[1], b[1], 1e-12);
+  EXPECT_NEAR(x[2], 0.25 * b[2], 1e-12);
+  EXPECT_NEAR(x[3], 0.25 * b[3], 1e-12);
+}
+
 TEST(BlockDiagonalPreconditionerCustom, ThrowsOnOutOfRangeOverrideIndex)
 {
   Array<int> offsets({0, 2, 4});