diff --git a/src/smith/differentiable_numerics/tests/test_nonlinear_mixed_diffusion.cpp b/src/smith/differentiable_numerics/tests/test_nonlinear_mixed_diffusion.cpp index c5f3ef348..9bc941d02 100644 --- a/src/smith/differentiable_numerics/tests/test_nonlinear_mixed_diffusion.cpp +++ b/src/smith/differentiable_numerics/tests/test_nonlinear_mixed_diffusion.cpp @@ -89,7 +89,8 @@ struct ThermalConductivity { enum class PreconditionerCase { - StateDependentCustomFullSchur + StateDependentCustomFullSchur, + StateDependentCustomActionSchur }; const char* preconditionerCaseName(PreconditionerCase preconditioner) @@ -97,6 +98,8 @@ 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 {}; +/** + * @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 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& 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(); + amg_solver_->SetPrintLevel(0); + amg_solver_->SetOperator(*schur_operator_); + + linear_solver_ = std::make_unique(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 schur_operator_; + std::unique_ptr amg_solver_; + std::unique_ptr 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(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 sub_solver_options{flux_solver_options}; + auto sub_solvers = smith::buildBlockPreconditionerSubSolvers(sub_solver_options, mesh->getComm()); + + auto action_solver = std::make_unique( + 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( + 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[]) diff --git a/src/smith/numerics/block_preconditioner.cpp b/src/smith/numerics/block_preconditioner.cpp index 9f5a864fe..c810ba323 100644 --- a/src/smith/numerics/block_preconditioner.cpp +++ b/src/smith/numerics/block_preconditioner.cpp @@ -342,6 +342,27 @@ BlockSchurPreconditioner::BlockSchurPreconditioner(std::vector& 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& 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(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 diff --git a/src/smith/numerics/block_preconditioner.hpp b/src/smith/numerics/block_preconditioner.hpp index a155c931b..c3664e99c 100644 --- a/src/smith/numerics/block_preconditioner.hpp +++ b/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& block_offsets); + + /// @overload + void SetOperator(const mfem::Operator& op) override; + + /// @overload + void updateForState([[maybe_unused]] const mfem::Vector& state, + [[maybe_unused]] const mfem::Array& block_offsets) override; + + protected: + const mfem::BlockOperator* block_jacobian_ = nullptr; ///< Current 2x2 block Jacobian context. + mfem::Array 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 S_approx_owned_; const mfem::Operator* S_approx_view_ = nullptr; diff --git a/src/smith/numerics/tests/test_block_preconditioner_custom_operators.cpp b/src/smith/numerics/tests/test_block_preconditioner_custom_operators.cpp index 29810407f..042fa3d25 100644 --- a/src/smith/numerics/tests/test_block_preconditioner_custom_operators.cpp +++ b/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 context_calls, std::shared_ptr 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& block_offsets) override + { + smith::SchurComplementActionSolver::setBlockContext(jacobian, block_offsets); + ++(*context_calls_); + } + + void updateForState([[maybe_unused]] const mfem::Vector& state, + [[maybe_unused]] const mfem::Array& 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 context_calls_; + std::shared_ptr 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 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(0); + auto update_calls = std::make_shared(0); + + std::vector> solvers; + solvers.push_back(std::make_unique()); + solvers.push_back(std::make_unique(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 offsets({0, 2, 4});