diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md index c6efdb4f98..2f67dbc2fb 100644 --- a/RELEASE-NOTES.md +++ b/RELEASE-NOTES.md @@ -30,6 +30,8 @@ The Smith project release numbers follow [Semantic Versioning](http://semver.org - Added composable solid-mechanics and thermo-mechanics examples, tutorials, and regression tests covering coupled sensitivities, finite-difference checks, field parameters, and solves. - Added axisymmetric solid mechanics materials and loads for 2D `(r, z)` meshes. +- Added diagonal, triangular, and Schur block preconditioners for linearized block systems. +- Added support for custom block operators that are rebuilt from the current nonlinear state at each Newton iteration. ### Removed diff --git a/src/smith/differentiable_numerics/CMakeLists.txt b/src/smith/differentiable_numerics/CMakeLists.txt index 71ade68cc9..4984f0d5e3 100644 --- a/src/smith/differentiable_numerics/CMakeLists.txt +++ b/src/smith/differentiable_numerics/CMakeLists.txt @@ -10,6 +10,7 @@ set(differentiable_numerics_sources differentiable_physics.cpp lumped_mass_explicit_newmark_state_advancer.cpp nonlinear_block_solver.cpp + weak_form_block_operator.cpp system_solver.cpp system_base.cpp nonlinear_solve.cpp @@ -24,6 +25,7 @@ set(differentiable_numerics_headers state_advancer.hpp reaction.hpp nonlinear_block_solver.hpp + weak_form_block_operator.hpp system_solver.hpp differentiable_physics.hpp timestep_estimator.hpp diff --git a/src/smith/differentiable_numerics/nonlinear_block_solver.cpp b/src/smith/differentiable_numerics/nonlinear_block_solver.cpp index 5c657e30ff..47816cf839 100644 --- a/src/smith/differentiable_numerics/nonlinear_block_solver.cpp +++ b/src/smith/differentiable_numerics/nonlinear_block_solver.cpp @@ -58,7 +58,7 @@ NonlinearBlockSolver::NonlinearBlockSolver(std::unique_ptr s, MP std::shared_ptr NonlinearBlockSolver::cloneFresh() const { - if (!retained_nonlinear_options_ || !retained_linear_options_) { + if (state_dependent_solver_ || !retained_nonlinear_options_ || !retained_linear_options_) { return nullptr; } @@ -199,6 +199,9 @@ std::vector NonlinearBlockSolver::solve( for (int row_i = 0; row_i < num_rows; ++row_i) { *u_guesses[static_cast(row_i)] = u->GetBlock(row_i); } + if (state_dependent_solver_) { + state_dependent_solver_->updateForState(*u, block_offsets); + } matrix_of_jacs_ = jacobian_funcs(u_guesses); if (num_rows == 1) { auto& J = matrix_of_jacs_[0][0]; @@ -296,4 +299,21 @@ std::shared_ptr buildNonlinearBlockSolver(NonlinearSolverO nonlinear_opts.relative_tol, nonlinear_opts, linear_opts); } +std::shared_ptr buildNonlinearBlockSolver(NonlinearSolverOptions nonlinear_opts, + LinearSolverOptions linear_opts, + const smith::Mesh& mesh, + std::unique_ptr preconditioner) +{ + auto solid_solver = + std::make_unique(nonlinear_opts, linear_opts, std::move(preconditioner), mesh.getComm()); + auto* state_dependent_preconditioner = dynamic_cast(&solid_solver->preconditioner()); + auto nonlinear_block_solver = + std::make_shared(std::move(solid_solver), mesh.getComm(), nonlinear_opts.absolute_tol, + nonlinear_opts.relative_tol, std::nullopt, linear_opts); + if (state_dependent_preconditioner) { + nonlinear_block_solver->setStateDependentSolver(state_dependent_preconditioner); + } + return nonlinear_block_solver; +} + } // namespace smith diff --git a/src/smith/differentiable_numerics/nonlinear_block_solver.hpp b/src/smith/differentiable_numerics/nonlinear_block_solver.hpp index d75b2b1f0a..0af70cf039 100644 --- a/src/smith/differentiable_numerics/nonlinear_block_solver.hpp +++ b/src/smith/differentiable_numerics/nonlinear_block_solver.hpp @@ -15,6 +15,7 @@ #include #include #include +#include #include #include @@ -22,6 +23,8 @@ #include "smith/numerics/nonlinear_convergence.hpp" namespace mfem { +template +class Array; class Solver; class Vector; class HypreParMatrix; @@ -35,6 +38,7 @@ class BoundaryConditionManager; class FiniteElementState; class FiniteElementDual; class Mesh; +class StateDependentSolver; struct NonlinearSolverOptions; struct LinearSolverOptions; @@ -140,7 +144,10 @@ class NonlinearBlockSolver : public NonlinearBlockSolverBase { /// @brief Set the inner tolerance multiplier. void setInnerToleranceMultiplier(double multiplier) override { inner_tol_multiplier_ = multiplier; } - /// @brief Build a fresh solver instance from retained config. + /// @brief Register a solver refreshed with the current nonlinear state before Jacobian assembly. + void setStateDependentSolver(StateDependentSolver* solver) { state_dependent_solver_ = solver; } + + /// @brief Build a fresh solver instance from retained config, or nullptr when this solver has injected state. std::shared_ptr cloneFresh() const; mutable std::unique_ptr @@ -158,6 +165,7 @@ class NonlinearBlockSolver : public NonlinearBlockSolverBase { double inner_tol_multiplier_ = 1.0; ///< multiplier for tolerances during inner solves std::optional retained_nonlinear_options_ = std::nullopt; ///< retained nonlinear config std::optional retained_linear_options_ = std::nullopt; ///< retained linear config + StateDependentSolver* state_dependent_solver_ = nullptr; ///< optional solver refreshed during Jacobian evaluation }; /// @brief Create an equation-backed nonlinear block solver. @@ -168,4 +176,14 @@ std::shared_ptr buildNonlinearBlockSolver(NonlinearSolverO LinearSolverOptions linear_opts, const smith::Mesh& mesh); +/// @brief Create an equation-backed nonlinear block solver with a custom preconditioner. +/// @param nonlinear_opts nonlinear options struct +/// @param linear_opts linear options struct +/// @param mesh mesh +/// @param preconditioner custom preconditioner attached to the linear solver +std::shared_ptr buildNonlinearBlockSolver(NonlinearSolverOptions nonlinear_opts, + LinearSolverOptions linear_opts, + const smith::Mesh& mesh, + std::unique_ptr preconditioner); + } // namespace smith diff --git a/src/smith/differentiable_numerics/tests/CMakeLists.txt b/src/smith/differentiable_numerics/tests/CMakeLists.txt index d9c676c94d..a0533851cb 100644 --- a/src/smith/differentiable_numerics/tests/CMakeLists.txt +++ b/src/smith/differentiable_numerics/tests/CMakeLists.txt @@ -19,7 +19,10 @@ set(differentiable_numerics_test_source test_multiphysics_time_integrator.cpp test_thermo_mechanics_with_internal_vars.cpp test_mixed_poisson.cpp + test_nonlinear_mixed_diffusion.cpp test_amg_vdim_block_preconditioner.cpp + test_state_dependent_preconditioner.cpp + test_weak_form_block_operator.cpp ) smith_add_tests( SOURCES ${differentiable_numerics_test_source} diff --git a/src/smith/differentiable_numerics/tests/test_mixed_poisson.cpp b/src/smith/differentiable_numerics/tests/test_mixed_poisson.cpp index 9c21313fd3..206aa35b36 100644 --- a/src/smith/differentiable_numerics/tests/test_mixed_poisson.cpp +++ b/src/smith/differentiable_numerics/tests/test_mixed_poisson.cpp @@ -1,5 +1,12 @@ #include +#include +#include +#include +#include +#include +#include + #include "smith/infrastructure/application_manager.hpp" #include "smith/numerics/equation_solver.hpp" #include "smith/numerics/solver_config.hpp" @@ -13,6 +20,7 @@ #include "smith/differentiable_numerics/nonlinear_block_solver.hpp" #include "smith/differentiable_numerics/nonlinear_solve.hpp" #include "smith/differentiable_numerics/paraview_writer.hpp" +#include "smith/differentiable_numerics/weak_form_block_operator.hpp" #include "smith/numerics/block_preconditioner.hpp" #include "gretl/data_store.hpp" @@ -112,6 +120,8 @@ TEST_P(BlockPreconditionerTest, BlockSolve) "constitutive_eqn", mesh, space(flux), spaces({flux, potential})); smith::FunctionalWeakForm<2, Space, smith::Parameters> bal_form( "balance_eqn", mesh, space(potential), spaces({flux, potential})); + smith::FunctionalWeakForm<2, Space, smith::Parameters> potential_diffusion_form( + "potential_diffusion", mesh, space(potential), spaces({potential})); con_form.addBodyIntegral(DependsOn<0, 1>{}, mesh->entireBodyName(), [](auto /* t */, auto /* x */, auto SIGMA, auto U) { @@ -132,6 +142,11 @@ TEST_P(BlockPreconditionerTest, BlockSolve) auto f = 2.0 * pi * pi * sin(pi * x[0]) * sin(pi * x[1]); return smith::tuple{-f + div_sigma, smith::zero{}}; }); + potential_diffusion_form.addBodyIntegral(DependsOn<0>{}, mesh->entireBodyName(), + [](auto /* time_info */, auto /* x */, auto U) { + auto grad_u = get(U); + return smith::tuple{smith::zero{}, grad_u}; + }); // u_exact = sin(M_PI * x(0)) * sin(M_PI * x(1)); // sigma_exact // pi * cos(pi * x(0)) * sin(pi * x(1)) @@ -177,6 +192,8 @@ TEST_P(BlockPreconditionerTest, BlockSolve) auto time = graph->create_state(0.0); auto dt = graph->create_state(0.025); size_t cycle = 0; + const auto time_info = smith::TimeInfo(time.get(), dt.get(), cycle); + std::unique_ptr diffusion_precond; std::vector params; auto& flux_params = params; auto& potential_params = params; @@ -203,7 +220,16 @@ TEST_P(BlockPreconditionerTest, BlockSolve) case BlockPrecondType::SchurFullCustom: linear_options.preconditioner = smith::Preconditioner::BlockSchur; linear_options.block_schur_type = smith::BlockSchurType::Full; - /// linear_options.schur_approx_type = smith::BlockSchurType::Custom; + linear_options.schur_approx_type = smith::SchurApproxType::Custom; + std::vector overrides; + overrides.push_back(smith::makeWeakFormBlockProviderOverride(1, potential_diffusion_form, shape_disp, {potential}, + {1.0}, time_info, potential_bc_manager.get())); + + auto solvers = + smith::buildBlockPreconditionerSubSolvers(linear_options.sub_block_linear_solver_options, mesh->getComm()); + + diffusion_precond = std::make_unique( + std::move(solvers), linear_options.block_schur_type, linear_options.schur_approx_type, std::move(overrides)); break; } @@ -214,11 +240,14 @@ TEST_P(BlockPreconditionerTest, BlockSolve) nonlin_opts.max_iterations = 1; nonlin_opts.print_level = linear_options.print_level; - auto nonlinear_block_solver = smith::buildNonlinearBlockSolver(nonlin_opts, linear_options, *mesh); + auto nonlinear_block_solver = + diffusion_precond + ? smith::buildNonlinearBlockSolver(nonlin_opts, linear_options, *mesh, std::move(diffusion_precond)) + : smith::buildNonlinearBlockSolver(nonlin_opts, linear_options, *mesh); auto sols = block_solve({&con_form, &bal_form}, {{0, 1}, {0, 1}}, shape_disp, {con_arguments, bal_arguments}, - {flux_params, potential_params}, smith::TimeInfo(time.get(), dt.get(), cycle), - nonlinear_block_solver.get(), {flux_bc_manager.get(), potential_bc_manager.get()}); + {flux_params, potential_params}, time_info, nonlinear_block_solver.get(), + {flux_bc_manager.get(), potential_bc_manager.get()}); auto pv_writer = smith::createParaviewWriter(*mesh, sols, physics_name); pv_writer.write(0, 0.0, sols); diff --git a/src/smith/differentiable_numerics/tests/test_nonlinear_mixed_diffusion.cpp b/src/smith/differentiable_numerics/tests/test_nonlinear_mixed_diffusion.cpp new file mode 100644 index 0000000000..c5f3ef348b --- /dev/null +++ b/src/smith/differentiable_numerics/tests/test_nonlinear_mixed_diffusion.cpp @@ -0,0 +1,278 @@ +// Copyright (c) Lawrence Livermore National Security, LLC and +// other Smith Project Developers. See the top-level LICENSE file for +// details. +// +// SPDX-License-Identifier: (BSD-3-Clause) + +/** + * @file test_nonlinear_mixed_diffusion.cpp + * + * @brief Test field-dependent Schur complement preconditioning. + * + * Problem: Nonlinear thermal diffusion with temperature-dependent conductivity + * + * -∇·(k(T)∇T) = f in Ω + * T = 0 on ∂Ω + * + * where k(T) = k₀(1 + α·T) is temperature-dependent conductivity. + * + * Mixed formulation (introduce flux q = -k(T)∇T): + * + * q + k(T)∇T = 0 (constitutive equation) + * ∇·q = f (balance equation) + * + * This gives a 2×2 saddle-point system: + * + * [K(T) B^T] [q] [0] + * [B 0 ] [T] = [f] + * + * where K(T) = k(T)⁻¹ I and B = ∇· operator. + * + * The Schur complement S = -B K(T)⁻¹ B^T ≈ -∇·(k(T)∇) depends on the + * temperature field T. At each Newton iteration, we need to update the + * Schur complement approximation with the current temperature iterate. + * + * This example demonstrates: + * 1. Mixed formulation with field-dependent coefficients + * 2. Custom Schur complement operator that depends on solution field + * 3. State update plumbing in the Newton loop + */ + +#include + +#include +#include +#include +#include +#include +#include + +#include "smith/smith_config.hpp" +#include "smith/differentiable_numerics/field_state.hpp" +#include "smith/differentiable_numerics/nonlinear_block_solver.hpp" +#include "smith/differentiable_numerics/nonlinear_solve.hpp" +#include "smith/differentiable_numerics/weak_form_block_operator.hpp" +#include "smith/infrastructure/application_manager.hpp" +#include "smith/mesh_utils/mesh_utils.hpp" +#include "smith/numerics/block_preconditioner.hpp" +#include "smith/numerics/equation_solver.hpp" +#include "smith/numerics/solver_config.hpp" +#include "smith/physics/boundary_conditions/boundary_condition_manager.hpp" +#include "smith/physics/functional_weak_form.hpp" +#include "smith/physics/state/state_manager.hpp" + +#include "gretl/data_store.hpp" +#include "gretl/wang_checkpoint_strategy.hpp" + +using namespace smith; + +// Taylor-Hood-style elements for stability +using ShapeDispSpace = H1<1, 2>; +using FluxSpace = H1<2, 2>; // Vector flux q +using TemperatureSpace = H1<1>; // Scalar temperature T +using TemperatureSchurForm = smith::FunctionalWeakForm<2, TemperatureSpace, smith::Parameters>; + +/** + * @brief Temperature-dependent conductivity k(T) = k₀(1 + α·T) + */ +struct ThermalConductivity { + double k0 = 1.0; // Base conductivity + double alpha = 0.5; // Temperature dependence parameter + + // Support both double and dual numbers for automatic differentiation + template + auto operator()(const T& temp) const + { + return k0 * (1.0 + alpha * temp); + } +}; + +enum class PreconditionerCase +{ + StateDependentCustomFullSchur +}; + +const char* preconditionerCaseName(PreconditionerCase preconditioner) +{ + switch (preconditioner) { + case PreconditionerCase::StateDependentCustomFullSchur: + return "StateDependentCustomFullSchur"; + } + return "Unknown"; +} + +std::string preconditionerCaseNameGenerator(const ::testing::TestParamInfo& info) +{ + return preconditionerCaseName(info.param); +} + +class MeshFixture : public testing::Test { + protected: + axom::sidre::DataStore datastore; + std::shared_ptr mesh; + ThermalConductivity k_func; + + void SetUp() override + { + smith::StateManager::initialize(datastore, "nonlinear_mixed_diffusion"); + + MPI_Barrier(MPI_COMM_WORLD); + int serial_refinement = 3; + int parallel_refinement = 0; + + std::string filename = SMITH_REPO_DIR "/data/meshes/square_attribute.mesh"; + + const std::string meshtag = "mesh"; + mesh = std::make_shared(smith::buildMeshFromFile(filename), meshtag, serial_refinement, + parallel_refinement); + + // Set up nonlinear conductivity + k_func.k0 = 1.0; + k_func.alpha = 0.5; // Moderate nonlinearity + } +}; + +class NonlinearMixedDiffusionPreconditionerTest : public MeshFixture, + public ::testing::WithParamInterface {}; + +// Exercises state-dependent Schur preconditioners on nonlinear mixed diffusion. +TEST_P(NonlinearMixedDiffusionPreconditionerTest, BlockSolve) +{ + const auto preconditioner_case = GetParam(); + + std::string physics_name = std::string("nonlinear_thermal_") + preconditionerCaseName(preconditioner_case); + auto graph = std::make_shared(std::make_unique(100)); + + // Create field states + auto shape_disp = createFieldState(*graph, ShapeDispSpace{}, physics_name + "_shape_displacement", mesh->tag()); + auto flux = createFieldState(*graph, FluxSpace{}, physics_name + "_flux", mesh->tag()); + auto temperature = createFieldState(*graph, TemperatureSpace{}, physics_name + "_temperature", mesh->tag()); + + // Set up weak forms + smith::FunctionalWeakForm<2, FluxSpace, smith::Parameters> constitutive_form( + "constitutive", mesh, space(flux), spaces({flux, temperature})); + + smith::FunctionalWeakForm<2, TemperatureSpace, smith::Parameters> balance_form( + "balance", mesh, space(temperature), spaces({flux, temperature})); + + TemperatureSchurForm temperature_schur_form("temperature_schur", mesh, space(temperature), spaces({temperature})); + + // Constitutive equation: q + k(T)∇T = 0 + constitutive_form.addBodyIntegral(DependsOn<0, 1>{}, mesh->entireBodyName(), + [this](auto /* time_info */, auto /* x */, auto Q, auto T) { + auto q = get(Q); + auto temp = get(T); + auto grad_temp = get(T); + + // k(T) = k₀(1 + α·T) + auto k_val = k_func(temp); + + // Residual: q + k(T)∇T + auto residual = q + k_val * grad_temp; + + return smith::tuple{residual, smith::zero{}}; + }); + + // Balance equation: ∇·q = f + balance_form.addBodyIntegral(DependsOn<0>{}, mesh->entireBodyName(), [](auto /* time_info */, auto /* x */, auto Q) { + auto div_q = smith::tr(get(Q)); + + double f = 1.0; + + return smith::tuple{-f + div_q, smith::zero{}}; + }); + + temperature_schur_form.addBodyIntegral(DependsOn<0>{}, mesh->entireBodyName(), + [this](auto /* time_info */, auto /* x */, auto T) { + auto temp = get(T); + auto grad_temp = get(T); + + return smith::tuple{smith::zero{}, k_func(temp) * grad_temp}; + }); + + // Boundary conditions: T = 0 on all boundaries + auto flux_bc_manager = std::make_shared(mesh->mfemParMesh()); + auto temp_bc_manager = std::make_shared(mesh->mfemParMesh()); + + auto zero_bcs = std::make_shared([](const mfem::Vector&) { return 0.0; }); + temp_bc_manager->addEssential(std::set{1, 2, 3, 4}, zero_bcs, space(temperature), 0); + + // Linear solver options + smith::LinearSolverOptions linear_options; + linear_options.linear_solver = smith::LinearSolver::GMRES; + linear_options.relative_tol = 1.0e-10; + linear_options.absolute_tol = 1.0e-14; + linear_options.max_iterations = 500; + linear_options.print_level = 0; + + smith::LinearSolverOptions amg_options; + amg_options.linear_solver = smith::LinearSolver::PrecondOnly; + amg_options.preconditioner = smith::Preconditioner::HypreAMG; + const smith::LinearSolverOptions flux_solver_options = amg_options; + const smith::LinearSolverOptions schur_solver_options = amg_options; + + linear_options.preconditioner = smith::Preconditioner::BlockSchur; + linear_options.block_schur_type = smith::BlockSchurType::Full; + + // Nonlinear solver options + smith::NonlinearSolverOptions nonlin_opts; + nonlin_opts.nonlin_solver = smith::NonlinearSolver::Newton; + nonlin_opts.relative_tol = 1.0e-10; + nonlin_opts.absolute_tol = 1.0e-12; + nonlin_opts.max_iterations = 12; + nonlin_opts.print_level = 0; + + std::shared_ptr nonlinear_solver; + auto time = graph->create_state(0.0); + auto dt = graph->create_state(1.0); + size_t cycle = 0; + const auto time_info = smith::TimeInfo(time.get(), dt.get(), cycle); + + if (preconditioner_case == PreconditionerCase::StateDependentCustomFullSchur) { + linear_options.schur_approx_type = smith::SchurApproxType::Custom; + const std::vector sub_solver_options{flux_solver_options, schur_solver_options}; + auto sub_solvers = smith::buildBlockPreconditionerSubSolvers(sub_solver_options, mesh->getComm()); + + std::vector overrides; + overrides.push_back(smith::makeStateDependentWeakFormBlockProviderOverride( + 1, temperature_schur_form, shape_disp, {temperature}, {1.0}, time_info, temp_bc_manager.get(), + {smith::StateBlockBinding{1, 0}})); + + auto preconditioner = + 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)); + } + + ASSERT_TRUE(nonlinear_solver != nullptr); + + // Solve + std::vector params; + + auto sols = block_solve({&constitutive_form, &balance_form}, {{0, 1}, {0, 1}}, shape_disp, + {{flux, temperature}, {flux, temperature}}, {params, params}, time_info, + nonlinear_solver.get(), {flux_bc_manager.get(), temp_bc_manager.get()}); + + // Verify convergence + EXPECT_EQ(sols.size(), 2); + + double temp_norm = sols[1].get()->Norml2(); + double flux_norm = sols[0].get()->Norml2(); + auto& newton_solver = nonlinear_solver->nonlinear_solver_->nonlinearSolver(); + EXPECT_TRUE(newton_solver.GetConverged()); + EXPECT_TRUE(std::isfinite(temp_norm)); + EXPECT_TRUE(std::isfinite(flux_norm)); + EXPECT_GT(temp_norm, 1e-6); + EXPECT_GT(flux_norm, 1e-6); +} + +INSTANTIATE_TEST_SUITE_P(StateDependentSchur, NonlinearMixedDiffusionPreconditionerTest, + ::testing::Values(PreconditionerCase::StateDependentCustomFullSchur), + preconditionerCaseNameGenerator); + +int main(int argc, char* argv[]) +{ + ::testing::InitGoogleTest(&argc, argv); + smith::ApplicationManager applicationManager(argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/src/smith/differentiable_numerics/tests/test_porous_heat_sink.cpp b/src/smith/differentiable_numerics/tests/test_porous_heat_sink.cpp index af2b9429d2..899d788401 100644 --- a/src/smith/differentiable_numerics/tests/test_porous_heat_sink.cpp +++ b/src/smith/differentiable_numerics/tests/test_porous_heat_sink.cpp @@ -1,5 +1,11 @@ #include +#include +#include +#include +#include +#include + #include "smith/infrastructure/application_manager.hpp" #include "smith/numerics/equation_solver.hpp" #include "smith/numerics/solver_config.hpp" @@ -13,6 +19,7 @@ #include "smith/differentiable_numerics/nonlinear_block_solver.hpp" #include "smith/differentiable_numerics/nonlinear_solve.hpp" #include "smith/differentiable_numerics/paraview_writer.hpp" +#include "smith/differentiable_numerics/weak_form_block_operator.hpp" #include "smith/numerics/block_preconditioner.hpp" #include "gretl/data_store.hpp" @@ -244,6 +251,8 @@ TEST_P(BlockPreconditionerTest, BlockSolve) auto time = graph->create_state(0.0); auto dt = graph->create_state(0.025); size_t cycle = 0; + const auto time_info = smith::TimeInfo(time.get(), dt.get(), cycle); + std::unique_ptr diff_precond; std::vector params; auto& T1_params = params; auto& T2_params = params; @@ -290,7 +299,18 @@ TEST_P(BlockPreconditionerTest, BlockSolve) case BlockPrecondType::SchurFullCustom: linear_options.preconditioner = smith::Preconditioner::BlockSchur; linear_options.block_schur_type = smith::BlockSchurType::Full; - /// linear_options.schur_approx_type = smith::BlockSchurType::Custom; + linear_options.schur_approx_type = smith::SchurApproxType::Custom; + std::vector jacobian_weights{0.0, 1.0, 0.0}; + + std::vector overrides; + overrides.push_back(smith::makeWeakFormBlockProviderOverride(1, T2_form, shape_disp, T2_arguments, + jacobian_weights, time_info, T2_bc_manager.get())); + + auto solvers = + smith::buildBlockPreconditionerSubSolvers(linear_options.sub_block_linear_solver_options, mesh->getComm()); + + diff_precond = std::make_unique( + std::move(solvers), linear_options.block_schur_type, linear_options.schur_approx_type, std::move(overrides)); break; } @@ -301,11 +321,13 @@ TEST_P(BlockPreconditionerTest, BlockSolve) nonlin_opts.max_iterations = 1; nonlin_opts.print_level = linear_options.print_level; - auto nonlinear_block_solver = smith::buildNonlinearBlockSolver(nonlin_opts, linear_options, *mesh); + auto nonlinear_block_solver = + diff_precond ? smith::buildNonlinearBlockSolver(nonlin_opts, linear_options, *mesh, std::move(diff_precond)) + : smith::buildNonlinearBlockSolver(nonlin_opts, linear_options, *mesh); auto sols = block_solve({&T1_form, &T2_form}, {{0, 1}, {0, 1}}, shape_disp, {T1_arguments, T2_arguments}, - {T1_params, T2_params}, smith::TimeInfo(time.get(), dt.get(), cycle), - nonlinear_block_solver.get(), {T1_bc_manager.get(), T2_bc_manager.get()}); + {T1_params, T2_params}, time_info, nonlinear_block_solver.get(), + {T1_bc_manager.get(), T2_bc_manager.get()}); auto pv_writer = smith::createParaviewWriter(*mesh, sols, physics_name); pv_writer.write(0, 0.0, sols); diff --git a/src/smith/differentiable_numerics/tests/test_state_dependent_preconditioner.cpp b/src/smith/differentiable_numerics/tests/test_state_dependent_preconditioner.cpp new file mode 100644 index 0000000000..2c16d232e4 --- /dev/null +++ b/src/smith/differentiable_numerics/tests/test_state_dependent_preconditioner.cpp @@ -0,0 +1,183 @@ +// Copyright (c) Lawrence Livermore National Security, LLC and +// other Smith Project Developers. See the top-level LICENSE file for +// details. +// +// SPDX-License-Identifier: (BSD-3-Clause) + +#include + +#include +#include +#include + +#include +#include "mfem.hpp" +#include "axom/sidre.hpp" + +#include "smith/differentiable_numerics/nonlinear_block_solver.hpp" +#include "smith/infrastructure/application_manager.hpp" +#include "smith/numerics/block_preconditioner.hpp" +#include "smith/numerics/solver_config.hpp" +#include "smith/physics/mesh.hpp" +#include "smith/physics/state/finite_element_state.hpp" +#include "smith/physics/state/state_manager.hpp" + +namespace smith { +namespace { + +class StateScaledIdentityProvider : public BlockOperatorProvider { + public: + explicit StateScaledIdentityProvider(int size) : size_(size) { rebuild(1.0); } + + void updateForState(const mfem::Vector& state, const mfem::Array& block_offsets) override + { + double scale = 0.0; + for (int i = block_offsets[1]; i < block_offsets[2]; ++i) { + scale += state[i]; + } + + last_scale_ = scale; + ++update_count_; + rebuild(scale); + } + + const mfem::Operator& currentOperator() const override + { + MFEM_VERIFY(current_operator_, "StateScaledIdentityProvider has no current operator"); + return *current_operator_; + } + + double lastScale() const { return last_scale_; } + + int updateCount() const { return update_count_; } + + private: + void rebuild(double scale) + { + auto op = std::make_unique(size_); + for (int i = 0; i < size_; ++i) { + op->Add(i, i, scale); + } + op->Finalize(); + current_operator_ = std::move(op); + } + + int size_ = 0; + double last_scale_ = 1.0; + int update_count_ = 0; + std::unique_ptr current_operator_; +}; + +class IdentitySolver : public mfem::Solver { + public: + void SetOperator(const mfem::Operator& op) override + { + height = op.Height(); + width = op.Width(); + } + + void Mult(const mfem::Vector& x, mfem::Vector& y) const override { y = x; } +}; + +std::unique_ptr makeScaledMassMatrix(mfem::ParFiniteElementSpace& space, double scale) +{ + mfem::ParBilinearForm mass(&space); + mfem::ConstantCoefficient scale_coef(scale); + mass.AddDomainIntegrator(new mfem::MassIntegrator(scale_coef)); + mass.Assemble(); + mass.Finalize(); + return std::unique_ptr(mass.ParallelAssemble()); +} + +// Verifies the public custom-preconditioner path refreshes state-dependent block operators. +TEST(StateDependentPreconditioner, UpdatesProviderCurrentOperatorFromNonlinearState) +{ + MPI_Comm comm = MPI_COMM_WORLD; + axom::sidre::DataStore datastore; + StateManager::reset(); + StateManager::initialize(datastore, "state_dependent_preconditioner"); + + mfem::Mesh serial_mesh = mfem::Mesh::MakeCartesian2D(1, 1, mfem::Element::QUADRILATERAL, 1, 1.0, 1.0); + Mesh mesh(std::move(serial_mesh), "state_dependent_preconditioner_mesh", 0, 0, comm); + + mfem::H1_FECollection fec(1, mesh.mfemParMesh().Dimension()); + mfem::ParFiniteElementSpace space0(&mesh.mfemParMesh(), &fec, 1, smith::ordering); + mfem::ParFiniteElementSpace space1(&mesh.mfemParMesh(), &fec, 1, smith::ordering); + + auto u0 = std::make_shared(space0, "u0"); + auto u1 = std::make_shared(space1, "u1"); + *u0 = 1.0; + *u1 = 2.0; + + const double expected_initial_scale = u1->Sum(); + + auto provider = std::make_unique(u1->Size()); + auto* provider_ptr = provider.get(); + std::vector overrides; + overrides.emplace_back(1, std::move(provider)); + + std::vector> sub_solvers; + sub_solvers.push_back(std::make_unique()); + sub_solvers.push_back(std::make_unique()); + + auto preconditioner = std::make_unique(std::move(sub_solvers), BlockSchurType::Diagonal, + SchurApproxType::Custom, std::move(overrides)); + + NonlinearSolverOptions nonlin_opts; + nonlin_opts.nonlin_solver = NonlinearSolver::Newton; + nonlin_opts.max_iterations = 1; + nonlin_opts.print_level = 0; + nonlin_opts.absolute_tol = 1.0e-14; + nonlin_opts.relative_tol = 1.0e-14; + + LinearSolverOptions linear_opts; + linear_opts.linear_solver = LinearSolver::PrecondOnly; + linear_opts.preconditioner = Preconditioner::None; + linear_opts.print_level = 0; + + auto block_solver = buildNonlinearBlockSolver(nonlin_opts, linear_opts, mesh, std::move(preconditioner)); + + auto residual = [](const std::vector& states) { + std::vector residuals; + residuals.reserve(states.size()); + for (const auto& state : states) { + residuals.emplace_back(*state); + } + return residuals; + }; + + auto jacobian = [&space0, &space1](const std::vector&) { + std::vector> jac(2); + jac[0].resize(2); + jac[1].resize(2); + jac[0][0] = makeScaledMassMatrix(space0, 1.0); + jac[0][1] = makeScaledMassMatrix(space0, 0.0); + jac[1][0] = makeScaledMassMatrix(space1, 0.0); + jac[1][1] = makeScaledMassMatrix(space1, 1.0); + return jac; + }; + + [[maybe_unused]] auto solution = block_solver->solve({u0, u1}, residual, jacobian); + + ASSERT_GT(provider_ptr->updateCount(), 0); + EXPECT_NEAR(provider_ptr->lastScale(), expected_initial_scale, 1.0e-12); + + mfem::Vector x(u1->Size()); + mfem::Vector y(u1->Size()); + x = 1.0; + provider_ptr->currentOperator().Mult(x, y); + + for (int i = 0; i < y.Size(); ++i) { + EXPECT_NEAR(y[i], expected_initial_scale, 1.0e-12); + } +} + +} // namespace +} // namespace smith + +int main(int argc, char** argv) +{ + ::testing::InitGoogleTest(&argc, argv); + smith::ApplicationManager applicationManager(argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/src/smith/differentiable_numerics/tests/test_weak_form_block_operator.cpp b/src/smith/differentiable_numerics/tests/test_weak_form_block_operator.cpp new file mode 100644 index 0000000000..c78707a5bd --- /dev/null +++ b/src/smith/differentiable_numerics/tests/test_weak_form_block_operator.cpp @@ -0,0 +1,146 @@ +// Copyright (c) Lawrence Livermore National Security, LLC and +// other Smith Project Developers. See the top-level LICENSE file for +// details. +// +// SPDX-License-Identifier: (BSD-3-Clause) + +#include + +#include +#include +#include + +#include "mfem.hpp" + +#include "smith/differentiable_numerics/field_state.hpp" +#include "smith/differentiable_numerics/weak_form_block_operator.hpp" +#include "smith/infrastructure/application_manager.hpp" +#include "smith/physics/functional_weak_form.hpp" +#include "smith/physics/mesh.hpp" +#include "smith/physics/state/state_manager.hpp" + +#include "gretl/data_store.hpp" +#include "gretl/wang_checkpoint_strategy.hpp" + +namespace smith { +namespace { + +using ShapeDispSpace = H1<1, 2>; +using ScalarSpace = H1<1>; +using ScalarWeakForm = FunctionalWeakForm<2, ScalarSpace, Parameters>; + +class WeakFormBlockOperatorTest : public testing::Test { + protected: + void initialize(const std::string& physics_name) + { + StateManager::initialize(datastore, physics_name); + + auto serial_mesh = mfem::Mesh::MakeCartesian2D(1, 1, mfem::Element::QUADRILATERAL, 1, 1.0, 1.0); + mesh = std::make_shared(std::move(serial_mesh), physics_name + "_mesh"); + graph = std::make_shared(std::make_unique(100)); + } + + axom::sidre::DataStore datastore; + std::shared_ptr graph; + std::shared_ptr mesh; +}; + +void addMassIntegral(ScalarWeakForm& weak_form, const std::shared_ptr& mesh) +{ + weak_form.addBodyIntegral(DependsOn<0>{}, mesh->entireBodyName(), [](auto /* time_info */, auto /* x */, auto U) { + auto u = get(U); + return tuple{u, zero{}}; + }); +} + +void addQuadraticMassIntegral(ScalarWeakForm& weak_form, const std::shared_ptr& mesh) +{ + weak_form.addBodyIntegral(DependsOn<0>{}, mesh->entireBodyName(), [](auto /* time_info */, auto /* x */, auto U) { + auto u = get(U); + return tuple{u * u, zero{}}; + }); +} + +// Verifies the flat utility assembles an operator directly from weak-form inputs. +TEST_F(WeakFormBlockOperatorTest, BuildsOperatorFromWeakForm) +{ + const std::string physics_name = "weak_form_block_operator_build"; + initialize(physics_name); + auto shape_disp = createFieldState(*graph, ShapeDispSpace{}, physics_name + "_shape_displacement", mesh->tag()); + auto field = createFieldState(*graph, ScalarSpace{}, physics_name + "_field", mesh->tag()); + ScalarWeakForm weak_form("mass", mesh, space(field), spaces({field})); + addMassIntegral(weak_form, mesh); + + auto op = buildWeakFormOperator(weak_form, shape_disp, {field}, {1.0}, TimeInfo(0.0, 1.0)); + + ASSERT_NE(op, nullptr); + EXPECT_EQ(op->Height(), field.get()->Size()); + EXPECT_EQ(op->Width(), field.get()->Size()); +} + +// Verifies the public fixed provider factory hides the implementation builder. +TEST_F(WeakFormBlockOperatorTest, FixedOverrideProvidesWeakFormOperator) +{ + const std::string physics_name = "weak_form_block_operator_fixed"; + initialize(physics_name); + auto shape_disp = createFieldState(*graph, ShapeDispSpace{}, physics_name + "_shape_displacement", mesh->tag()); + auto field = createFieldState(*graph, ScalarSpace{}, physics_name + "_field", mesh->tag()); + ScalarWeakForm weak_form("mass", mesh, space(field), spaces({field})); + addMassIntegral(weak_form, mesh); + + auto override = makeWeakFormBlockProviderOverride(0, weak_form, shape_disp, {field}, {1.0}, TimeInfo(0.0, 1.0)); + const auto& op = override.provider->currentOperator(); + + EXPECT_EQ(op.Height(), field.get()->Size()); + EXPECT_EQ(op.Width(), field.get()->Size()); +} + +// Verifies state-dependent provider overrides refresh scratch fields without mutating graph-owned fields. +TEST_F(WeakFormBlockOperatorTest, StateDependentOverrideUpdatesFromStateBlock) +{ + const std::string physics_name = "weak_form_block_operator_state_dependent"; + initialize(physics_name); + auto shape_disp = createFieldState(*graph, ShapeDispSpace{}, physics_name + "_shape_displacement", mesh->tag()); + auto field = createFieldState(*graph, ScalarSpace{}, physics_name + "_field", mesh->tag()); + *field.get() = 2.0; + ScalarWeakForm weak_form("quadratic_mass", mesh, space(field), spaces({field})); + addQuadraticMassIntegral(weak_form, mesh); + + auto override = makeStateDependentWeakFormBlockProviderOverride( + 0, weak_form, shape_disp, {field}, {1.0}, TimeInfo(0.0, 1.0), mfem::Array(), {StateBlockBinding{0, 0}}); + + mfem::Array block_offsets(2); + block_offsets[0] = 0; + block_offsets[1] = field.get()->Size(); + + mfem::Vector x(block_offsets[1]); + mfem::Vector y(block_offsets[1]); + x = 1.0; + + mfem::Vector state(block_offsets[1]); + state = 3.0; + override.provider->updateForState(state, block_offsets); + override.provider->currentOperator().Mult(x, y); + const double first_norm = y.Norml2(); + + state = 5.0; + override.provider->updateForState(state, block_offsets); + override.provider->currentOperator().Mult(x, y); + const double second_norm = y.Norml2(); + + EXPECT_GT(first_norm, 0.0); + EXPECT_GT(second_norm, first_norm); + for (int i = 0; i < field.get()->Size(); ++i) { + EXPECT_DOUBLE_EQ((*field.get())[i], 2.0); + } +} + +} // namespace +} // namespace smith + +int main(int argc, char* argv[]) +{ + ::testing::InitGoogleTest(&argc, argv); + smith::ApplicationManager applicationManager(argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/src/smith/differentiable_numerics/weak_form_block_operator.cpp b/src/smith/differentiable_numerics/weak_form_block_operator.cpp new file mode 100644 index 0000000000..c4eb6aa982 --- /dev/null +++ b/src/smith/differentiable_numerics/weak_form_block_operator.cpp @@ -0,0 +1,232 @@ +// Copyright (c) Lawrence Livermore National Security, LLC and +// other Smith Project Developers. See the top-level LICENSE file for +// details. +// +// SPDX-License-Identifier: (BSD-3-Clause) + +#include "smith/differentiable_numerics/weak_form_block_operator.hpp" + +#include +#include +#include + +#include "smith/physics/boundary_conditions/boundary_condition_manager.hpp" +#include "smith/physics/state/finite_element_state.hpp" +#include "smith/physics/weak_form.hpp" + +namespace smith { + +namespace { + +mfem::Array copyEssentialTrueDofs(const BoundaryConditionManager* bc_manager) +{ + if (!bc_manager) { + return mfem::Array(); + } + return bc_manager->allEssentialTrueDofs(); +} + +std::vector copyFieldValues(const std::vector& fields) +{ + std::vector field_values; + field_values.reserve(fields.size()); + for (const auto& field : fields) { + field_values.emplace_back(*field.get()); + } + return field_values; +} + +std::vector getConstOperatorFieldPointers(const std::vector& fields) +{ + std::vector pointers; + pointers.reserve(fields.size()); + for (const auto& field : fields) { + pointers.push_back(&field); + } + return pointers; +} + +class WeakFormBlockOperatorBuilder { + public: + WeakFormBlockOperatorBuilder(const WeakForm& weak_form, FieldState shape_disp, std::vector fields, + std::vector jacobian_weights, TimeInfo time_info, mfem::Array ess_tdofs, + std::vector state_block_bindings) + : weak_form_(weak_form), + shape_disp_(std::move(shape_disp)), + fields_(std::move(fields)), + jacobian_weights_(std::move(jacobian_weights)), + time_info_(time_info), + ess_tdofs_(std::move(ess_tdofs)), + state_block_bindings_(std::move(state_block_bindings)) + { + validate(); + } + + std::unique_ptr build() const { return build(smith::getConstFieldPointers(fields_)); } + + std::unique_ptr updateAndBuild(const mfem::Vector& state, + const mfem::Array& block_offsets) const + { + auto operator_fields = copyFieldValues(fields_); + updateFieldsFromState(operator_fields, state, block_offsets); + return build(getConstOperatorFieldPointers(operator_fields)); + } + + private: + std::unique_ptr build(const std::vector& fields) const + { + auto op = weak_form_.jacobian(time_info_, shape_disp_.get().get(), fields, jacobian_weights_); + if (!op) { + throw std::invalid_argument("Weak-form operator builder received a null weak-form Jacobian"); + } + eliminateEssentialDofs(*op); + return op; + } + + void validate() const + { + if (jacobian_weights_.size() != fields_.size()) { + throw std::invalid_argument("Weak-form operator jacobian_weights size must match fields size"); + } + for (const auto& binding : state_block_bindings_) { + if (binding.block_index < 0) { + throw std::invalid_argument("Weak-form operator state block index must be non-negative"); + } + if (binding.field_index < 0 || binding.field_index >= static_cast(fields_.size())) { + throw std::invalid_argument("Weak-form operator field index is out of range"); + } + } + } + + void updateFieldsFromState(std::vector& fields, const mfem::Vector& state, + const mfem::Array& block_offsets) const + { + for (const auto& binding : state_block_bindings_) { + MFEM_VERIFY(binding.block_index + 1 < block_offsets.Size(), "Weak-form operator state block is out of range"); + + const int block_begin = block_offsets[binding.block_index]; + const int block_size = block_offsets[binding.block_index + 1] - block_begin; + MFEM_VERIFY(block_begin >= 0 && block_size >= 0 && block_begin + block_size <= state.Size(), + "Weak-form operator block offsets are inconsistent with the state size"); + + FiniteElementState& field = fields[static_cast(binding.field_index)]; + MFEM_VERIFY(field.Size() == block_size, + "Weak-form operator cannot update a field from a block with incompatible size"); + + mfem::Vector block_view; + block_view.MakeRef(const_cast(state), block_begin, block_size); + field = block_view; + } + } + + void eliminateEssentialDofs(mfem::HypreParMatrix& op) const + { + if (ess_tdofs_.Size() == 0) { + return; + } + mfem::HypreParMatrix* eliminated_entries = op.EliminateRowsCols(ess_tdofs_); + delete eliminated_entries; + } + + const WeakForm& weak_form_; + FieldState shape_disp_; + std::vector fields_; + std::vector jacobian_weights_; + TimeInfo time_info_; + mfem::Array ess_tdofs_; + std::vector state_block_bindings_; +}; + +} // namespace + +std::unique_ptr buildWeakFormOperator(const WeakForm& weak_form, FieldState shape_disp, + std::vector fields, + std::vector jacobian_weights, TimeInfo time_info, + mfem::Array ess_tdofs) +{ + WeakFormBlockOperatorBuilder builder(weak_form, std::move(shape_disp), std::move(fields), std::move(jacobian_weights), + time_info, std::move(ess_tdofs), {}); + return builder.build(); +} + +std::unique_ptr buildWeakFormOperator(const WeakForm& weak_form, FieldState shape_disp, + std::vector fields, + std::vector jacobian_weights, TimeInfo time_info, + const BoundaryConditionManager* bc_manager) +{ + return buildWeakFormOperator(weak_form, std::move(shape_disp), std::move(fields), std::move(jacobian_weights), + time_info, copyEssentialTrueDofs(bc_manager)); +} + +StateDependentWeakFormOperator makeStateDependentWeakFormOperator(const WeakForm& weak_form, FieldState shape_disp, + std::vector fields, + std::vector jacobian_weights, + TimeInfo time_info, mfem::Array ess_tdofs, + std::vector state_block_bindings) +{ + WeakFormBlockOperatorBuilder builder(weak_form, std::move(shape_disp), std::move(fields), std::move(jacobian_weights), + time_info, std::move(ess_tdofs), std::move(state_block_bindings)); + return [builder = std::move(builder)](const mfem::Vector& state, const mfem::Array& block_offsets) { + return builder.updateAndBuild(state, block_offsets); + }; +} + +StateDependentWeakFormOperator makeStateDependentWeakFormOperator(const WeakForm& weak_form, FieldState shape_disp, + std::vector fields, + std::vector jacobian_weights, + TimeInfo time_info, + const BoundaryConditionManager* bc_manager, + std::vector state_block_bindings) +{ + return makeStateDependentWeakFormOperator(weak_form, std::move(shape_disp), std::move(fields), + std::move(jacobian_weights), time_info, copyEssentialTrueDofs(bc_manager), + std::move(state_block_bindings)); +} + +BlockProviderOverride makeWeakFormBlockProviderOverride(int block_index, const WeakForm& weak_form, + FieldState shape_disp, std::vector fields, + std::vector jacobian_weights, TimeInfo time_info, + mfem::Array ess_tdofs) +{ + return makeFixedBlockProviderOverride( + block_index, buildWeakFormOperator(weak_form, std::move(shape_disp), std::move(fields), + std::move(jacobian_weights), time_info, std::move(ess_tdofs))); +} + +BlockProviderOverride makeWeakFormBlockProviderOverride(int block_index, const WeakForm& weak_form, + FieldState shape_disp, std::vector fields, + std::vector jacobian_weights, TimeInfo time_info, + const BoundaryConditionManager* bc_manager) +{ + return makeWeakFormBlockProviderOverride(block_index, weak_form, std::move(shape_disp), std::move(fields), + std::move(jacobian_weights), time_info, copyEssentialTrueDofs(bc_manager)); +} + +BlockProviderOverride makeStateDependentWeakFormBlockProviderOverride( + int block_index, const WeakForm& weak_form, FieldState shape_disp, std::vector fields, + std::vector jacobian_weights, TimeInfo time_info, mfem::Array ess_tdofs, + std::vector state_block_bindings) +{ + auto initial_operator = buildWeakFormOperator(weak_form, shape_disp, fields, jacobian_weights, time_info, ess_tdofs); + auto weak_form_operator_update = makeStateDependentWeakFormOperator( + weak_form, std::move(shape_disp), std::move(fields), std::move(jacobian_weights), time_info, std::move(ess_tdofs), + std::move(state_block_bindings)); + auto block_builder = [weak_form_operator_update = std::move(weak_form_operator_update)]( + const mfem::Vector& state, + const mfem::Array& block_offsets) mutable -> std::unique_ptr { + return weak_form_operator_update(state, block_offsets); + }; + return makeStateDependentBlockProviderOverride(block_index, std::move(block_builder), std::move(initial_operator)); +} + +BlockProviderOverride makeStateDependentWeakFormBlockProviderOverride( + int block_index, const WeakForm& weak_form, FieldState shape_disp, std::vector fields, + std::vector jacobian_weights, TimeInfo time_info, const BoundaryConditionManager* bc_manager, + std::vector state_block_bindings) +{ + return makeStateDependentWeakFormBlockProviderOverride( + block_index, weak_form, std::move(shape_disp), std::move(fields), std::move(jacobian_weights), time_info, + copyEssentialTrueDofs(bc_manager), std::move(state_block_bindings)); +} + +} // namespace smith diff --git a/src/smith/differentiable_numerics/weak_form_block_operator.hpp b/src/smith/differentiable_numerics/weak_form_block_operator.hpp new file mode 100644 index 0000000000..5dfee564b1 --- /dev/null +++ b/src/smith/differentiable_numerics/weak_form_block_operator.hpp @@ -0,0 +1,111 @@ +// Copyright (c) Lawrence Livermore National Security, LLC and +// other Smith Project Developers. See the top-level LICENSE file for +// details. +// +// SPDX-License-Identifier: (BSD-3-Clause) + +/** + * @file weak_form_block_operator.hpp + * + * @brief Helpers for using Smith weak forms as block preconditioner operators. + */ + +#pragma once + +#include +#include +#include + +#include "mfem.hpp" + +#include "smith/differentiable_numerics/field_state.hpp" +#include "smith/numerics/block_preconditioner.hpp" +#include "smith/physics/common.hpp" + +namespace smith { + +class BoundaryConditionManager; +class WeakForm; + +/** + * @brief Mapping from a nonlinear solve block to an input field of a weak form. + */ +struct StateBlockBinding { + int block_index; ///< Nonlinear solve block index in the monolithic state vector. + int field_index; ///< Weak-form input field index to update from the block. +}; + +/** + * @brief Callable that rebuilds a weak-form operator from a nonlinear state. + */ +using StateDependentWeakFormOperator = + std::function(const mfem::Vector&, const mfem::Array&)>; + +/** + * @brief Assemble a weak-form Jacobian operator for use in a block preconditioner. + */ +std::unique_ptr buildWeakFormOperator(const WeakForm& weak_form, FieldState shape_disp, + std::vector fields, + std::vector jacobian_weights, TimeInfo time_info, + mfem::Array ess_tdofs = mfem::Array()); + +/** + * @brief Assemble a weak-form Jacobian operator, eliminating essential dofs from a boundary-condition manager. + */ +std::unique_ptr buildWeakFormOperator(const WeakForm& weak_form, FieldState shape_disp, + std::vector fields, + std::vector jacobian_weights, TimeInfo time_info, + const BoundaryConditionManager* bc_manager); + +/** + * @brief Build a callable that updates bound weak-form fields from state and assembles the weak-form operator. + */ +StateDependentWeakFormOperator makeStateDependentWeakFormOperator(const WeakForm& weak_form, FieldState shape_disp, + std::vector fields, + std::vector jacobian_weights, + TimeInfo time_info, mfem::Array ess_tdofs, + std::vector state_block_bindings); + +/** + * @brief Build a callable using essential dofs copied from a boundary-condition manager. + */ +StateDependentWeakFormOperator makeStateDependentWeakFormOperator(const WeakForm& weak_form, FieldState shape_disp, + std::vector fields, + std::vector jacobian_weights, + TimeInfo time_info, + const BoundaryConditionManager* bc_manager, + std::vector state_block_bindings); + +/** + * @brief Build a fixed block override from a weak-form Jacobian operator. + */ +BlockProviderOverride makeWeakFormBlockProviderOverride(int block_index, const WeakForm& weak_form, + FieldState shape_disp, std::vector fields, + std::vector jacobian_weights, TimeInfo time_info, + mfem::Array ess_tdofs = mfem::Array()); + +/** + * @brief Build a fixed block override using essential dofs copied from a boundary-condition manager. + */ +BlockProviderOverride makeWeakFormBlockProviderOverride(int block_index, const WeakForm& weak_form, + FieldState shape_disp, std::vector fields, + std::vector jacobian_weights, TimeInfo time_info, + const BoundaryConditionManager* bc_manager); + +/** + * @brief Build a state-dependent block override from a weak-form Jacobian operator. + */ +BlockProviderOverride makeStateDependentWeakFormBlockProviderOverride( + int block_index, const WeakForm& weak_form, FieldState shape_disp, std::vector fields, + std::vector jacobian_weights, TimeInfo time_info, mfem::Array ess_tdofs, + std::vector state_block_bindings); + +/** + * @brief Build a state-dependent block override using essential dofs copied from a boundary-condition manager. + */ +BlockProviderOverride makeStateDependentWeakFormBlockProviderOverride( + int block_index, const WeakForm& weak_form, FieldState shape_disp, std::vector fields, + std::vector jacobian_weights, TimeInfo time_info, const BoundaryConditionManager* bc_manager, + std::vector state_block_bindings); + +} // namespace smith diff --git a/src/smith/numerics/CMakeLists.txt b/src/smith/numerics/CMakeLists.txt index cda8cb0bd6..8a9fd9c705 100644 --- a/src/smith/numerics/CMakeLists.txt +++ b/src/smith/numerics/CMakeLists.txt @@ -10,6 +10,7 @@ set(numerics_headers nonlinear_convergence.hpp odes.hpp solver_config.hpp + state_dependent_solver.hpp solver_with_preconditioner.hpp stdfunction_operator.hpp petsc_solvers.hpp diff --git a/src/smith/numerics/block_preconditioner.cpp b/src/smith/numerics/block_preconditioner.cpp index 38478a85cf..9f5a864feb 100644 --- a/src/smith/numerics/block_preconditioner.cpp +++ b/src/smith/numerics/block_preconditioner.cpp @@ -13,45 +13,128 @@ namespace smith { namespace { -void applyOverrides(int num_blocks, std::vector>& block_op_overrides, - std::vector overrides) +class FixedBlockOperatorProvider : public BlockOperatorProvider { + public: + explicit FixedBlockOperatorProvider(std::unique_ptr op) : op_(std::move(op)) + { + if (!op_) { + throw std::invalid_argument("Fixed block provider override requires a non-null operator"); + } + } + + const mfem::Operator& currentOperator() const override + { + MFEM_VERIFY(op_, "Fixed block operator provider has no operator"); + return *op_; + } + + private: + std::unique_ptr op_; +}; + +class StateDependentBlockOperatorProvider : public BlockOperatorProvider { + public: + StateDependentBlockOperatorProvider(StateDependentBlockOperatorBuilder builder, + std::unique_ptr initial_operator) + : builder_(std::move(builder)), op_(std::move(initial_operator)) + { + if (!builder_) { + throw std::invalid_argument("State-dependent block operator builder must be non-null"); + } + } + + void updateForState(const mfem::Vector& state, const mfem::Array& block_offsets) override + { + op_ = builder_(state, block_offsets); + if (!op_) { + throw std::invalid_argument("State-dependent block operator builder returned a null operator"); + } + } + + const mfem::Operator& currentOperator() const override + { + MFEM_VERIFY(op_, + "State-dependent block operator builder has no current operator; call updateForState first or " + "provide an initial operator"); + return *op_; + } + + private: + StateDependentBlockOperatorBuilder builder_; + std::unique_ptr op_; +}; + +void applyOverrides(int num_blocks, std::vector>& block_op_providers, + std::vector overrides) { for (auto& ov : overrides) { - const int i = ov.first; - auto& op = ov.second; + const int i = ov.block_index; + auto& provider = ov.provider; if (i < 0 || i >= num_blocks) { throw std::out_of_range("Override block index out of range"); } - if (!op) { - throw std::invalid_argument("Override operator must be non-null"); + if (!provider) { + throw std::invalid_argument("Override provider must be non-null"); } - if (block_op_overrides[static_cast(i)]) { + if (block_op_providers[static_cast(i)]) { throw std::invalid_argument("Duplicate override for same block index"); } - block_op_overrides[static_cast(i)] = std::move(op); + block_op_providers[static_cast(i)] = std::move(provider); + } +} + +void updateSolverForState(mfem::Solver* solver, const mfem::Vector& state, const mfem::Array& block_offsets) +{ + if (auto* state_dependent_solver = dynamic_cast(solver)) { + state_dependent_solver->updateForState(state, block_offsets); } } } // namespace +BlockProviderOverride makeFixedBlockProviderOverride(int block_index, std::unique_ptr op) +{ + return {block_index, std::make_unique(std::move(op))}; +} + +BlockProviderOverride makeStateDependentBlockProviderOverride(int block_index, + StateDependentBlockOperatorBuilder builder, + std::unique_ptr initial_operator) +{ + return {block_index, + std::make_unique(std::move(builder), std::move(initial_operator))}; +} + BlockPreconditioner::BlockPreconditioner(std::vector> solvers) : block_offsets_(), num_blocks_(static_cast(solvers.size())), block_jacobian_(nullptr), mfem_solvers_(std::move(solvers)), - block_op_overrides_(static_cast(num_blocks_)) + block_op_providers_(static_cast(num_blocks_)) +{ +} + +void BlockPreconditioner::updateForState(const mfem::Vector& state, const mfem::Array& block_offsets) { + for (auto& provider : block_op_providers_) { + if (provider) { + provider->updateForState(state, block_offsets); + } + } + for (auto& solver : mfem_solvers_) { + updateSolverForState(solver.get(), state, block_offsets); + } } BlockPreconditioner::~BlockPreconditioner() {} BlockDiagonalPreconditioner::BlockDiagonalPreconditioner(std::vector> solvers, - std::vector overrides) + std::vector overrides) : BlockPreconditioner(std::move(solvers)), solver_diag_(nullptr) { - applyOverrides(num_blocks_, block_op_overrides_, std::move(overrides)); + applyOverrides(num_blocks_, block_op_providers_, std::move(overrides)); } void BlockDiagonalPreconditioner::Mult(const mfem::Vector& in, mfem::Vector& out) const { solver_diag_->Mult(in, out); } @@ -78,8 +161,8 @@ void BlockDiagonalPreconditioner::SetOperator(const mfem::Operator& jacobian) const mfem::Operator* op = nullptr; const size_t si = static_cast(i); - if (block_op_overrides_[si]) { - op = block_op_overrides_[si].get(); // use override + if (block_op_providers_[si]) { + op = &block_op_providers_[si]->currentOperator(); // use override } else { op = &block_jacobian_->GetBlock(i, i); // use Jacobian diagonal block } @@ -96,10 +179,10 @@ BlockDiagonalPreconditioner::~BlockDiagonalPreconditioner() {} BlockTriangularPreconditioner::BlockTriangularPreconditioner(std::vector> solvers, BlockTriangularType type, - std::vector overrides) + std::vector overrides) : BlockPreconditioner(std::move(solvers)), type_(type) { - applyOverrides(num_blocks_, block_op_overrides_, std::move(overrides)); + applyOverrides(num_blocks_, block_op_providers_, std::move(overrides)); } void BlockTriangularPreconditioner::LowerSweep(const mfem::Vector& in, mfem::Vector& out) const @@ -230,8 +313,8 @@ void BlockTriangularPreconditioner::SetOperator(const mfem::Operator& jacobian) const mfem::Operator* op = nullptr; const size_t si = static_cast(i); - if (block_op_overrides_[si]) { - op = block_op_overrides_[si].get(); // use override + if (block_op_providers_[si]) { + op = &block_op_providers_[si]->currentOperator(); // use override } else { op = &block_jacobian_->GetBlock(i, i); // use Jacobian diagonal block } @@ -245,17 +328,17 @@ BlockTriangularPreconditioner::~BlockTriangularPreconditioner() {} BlockSchurPreconditioner::BlockSchurPreconditioner(std::vector> solvers, BlockSchurType type, SchurApproxType approxType, - std::vector overrides) + std::vector overrides) : BlockPreconditioner(std::move(solvers)), solver_diag_(nullptr), type_(type), approxType_(approxType) { - block_op_overrides_.resize(2); + block_op_providers_.resize(2); SLIC_ERROR_IF(mfem_solvers_.size() != 2, "This precondition is specifically for 2X2 block systems"); - applyOverrides(2, block_op_overrides_, std::move(overrides)); + applyOverrides(2, block_op_providers_, std::move(overrides)); - if (approxType_ == SchurApproxType::Custom && !block_op_overrides_[1]) { + if (approxType_ == SchurApproxType::Custom && !block_op_providers_[1]) { throw std::invalid_argument( - "SchurApproxType::Custom requires an override operator for block index 1 (custom Schur operator)"); + "SchurApproxType::Custom requires an override provider for block index 1 (custom Schur operator)"); } } @@ -415,39 +498,48 @@ void BlockSchurPreconditioner::SetOperator(const mfem::Operator& jacobian) block_offsets_.MakeRef(const_cast&>(block_jacobian_->RowOffsets())); solver_diag_ = std::make_unique(block_offsets_); - auto* A11 = dynamic_cast(&block_jacobian_->GetBlock(0, 0)); - auto* A12 = dynamic_cast(&block_jacobian_->GetBlock(0, 1)); - auto* A21 = dynamic_cast(&block_jacobian_->GetBlock(1, 0)); - auto* A22 = dynamic_cast(&block_jacobian_->GetBlock(1, 1)); - - MFEM_VERIFY(A11 && A12 && A21 && A22, - "All blocks must be HypreParMatrix for assembled Schur complement preconditioner."); + const mfem::Operator& A11_op = block_jacobian_->GetBlock(0, 0); + const mfem::Operator& A12_op = block_jacobian_->GetBlock(0, 1); + const mfem::Operator& A21_op = block_jacobian_->GetBlock(1, 0); + const mfem::Operator& A22_op = block_jacobian_->GetBlock(1, 1); if (type_ == BlockSchurType::Lower || type_ == BlockSchurType::Full) { - A_21_ = A21; + A_21_ = &A21_op; } if (type_ == BlockSchurType::Upper || type_ == BlockSchurType::Full) { - A_12_ = A12; + A_12_ = &A12_op; } // Diagonal preconditioner for block (0,0) const mfem::Operator* op = nullptr; - if (block_op_overrides_[0]) { - op = block_op_overrides_[0].get(); // use override + if (block_op_providers_[0]) { + op = &block_op_providers_[0]->currentOperator(); // use override } else { - op = A11; // use Jacobian diagonal block + op = &A11_op; // use Jacobian diagonal block } mfem_solvers_[0]->SetOperator(*op); mfem_solvers_[0]->iterative_mode = false; // Build Schur complement approximation if (approxType_ == SchurApproxType::DiagInv) { + auto* A11 = dynamic_cast(&A11_op); + auto* A12 = dynamic_cast(&A12_op); + auto* A21 = dynamic_cast(&A21_op); + auto* A22 = dynamic_cast(&A22_op); + + MFEM_VERIFY(A11 && A12 && A21 && A22, + "All blocks must be HypreParMatrix for assembled Schur complement preconditioner."); + S_approx_owned_.reset(BuildSchurDiagApprox_(*A11, *A12, *A21, *A22)); S_approx_view_ = S_approx_owned_.get(); } else if (approxType_ == SchurApproxType::A22Only) { + auto* A22 = dynamic_cast(&A22_op); + + MFEM_VERIFY(A22, "A22 block must be a HypreParMatrix for A22Only Schur complement preconditioner."); + S_approx_owned_.reset(new mfem::HypreParMatrix(*A22)); S_approx_view_ = S_approx_owned_.get(); - } else { + } else if (approxType_ == SchurApproxType::Custom) { S_approx_owned_.reset(); - S_approx_view_ = block_op_overrides_[1].get(); + S_approx_view_ = &block_op_providers_[1]->currentOperator(); } MFEM_VERIFY(S_approx_view_, "Schur complement approximation operator must be set"); diff --git a/src/smith/numerics/block_preconditioner.hpp b/src/smith/numerics/block_preconditioner.hpp index a088eaa8e6..a155c931bb 100644 --- a/src/smith/numerics/block_preconditioner.hpp +++ b/src/smith/numerics/block_preconditioner.hpp @@ -1,29 +1,101 @@ #pragma once -#include #include +#include #include #include + #include "mfem.hpp" +#include "smith/numerics/state_dependent_solver.hpp" + namespace smith { +/** + * @brief Supplies the current concrete operator for a block solver override. + */ +class BlockOperatorProvider { + public: + virtual ~BlockOperatorProvider() = default; + + /** + * @brief Refresh the owned operator for the current nonlinear state. + * @param state Monolithic state vector at the current nonlinear iterate. + * @param block_offsets Offsets describing the block layout of @a state. + */ + virtual void updateForState([[maybe_unused]] const mfem::Vector& state, + [[maybe_unused]] const mfem::Array& block_offsets) + { + } + + /** + * @brief Return the current concrete operator. + */ + virtual const mfem::Operator& currentOperator() const = 0; +}; + +/** + * @brief Builder that rebuilds an operator from the current nonlinear state. + * + * Builders are invoked by BlockPreconditioner::updateForState(). In nonlinear + * block solves created with a custom state-dependent preconditioner, that update + * is wired into the Newton loop before the preconditioner is configured with the + * current Jacobian. + */ +using StateDependentBlockOperatorBuilder = + std::function(const mfem::Vector&, const mfem::Array&)>; /** - * @brief Optional override for a diagonal block operator. + * @brief Optional provider override for a diagonal block operator. * - * The integer is the block index i and the operator replaces the Jacobian block - * A_ii (or, for 2x2 Schur systems, the block used to build/approximate the - * (1,1) Schur operator). + * The block index i identifies the provider that supplies the operator used + * in place of the Jacobian block A_ii. For 2x2 Schur systems, index 1 supplies + * the custom Schur operator when approxType is SchurApproxType::Custom. + */ +struct BlockProviderOverride { + /** + * @brief Construct from a block index and owned provider. + * @param block_index_in Block index to override. + * @param provider_in Provider supplying the override operator. + */ + BlockProviderOverride(int block_index_in, std::unique_ptr provider_in) + : block_index(block_index_in), provider(std::move(provider_in)) + { + } + + int block_index; ///< Block index to override. + std::unique_ptr provider; ///< Provider supplying the override operator. +}; + +/** + * @brief Build an override from a fixed concrete operator. + * @param block_index Block index to override. + * @param op Fixed concrete operator. + */ +BlockProviderOverride makeFixedBlockProviderOverride(int block_index, std::unique_ptr op); + +/** + * @brief Build an override from a state-dependent operator builder. + * @param block_index Block index to override. + * @param builder Callable that returns a new concrete operator for a state. + * @param initial_operator Optional operator to use before the first update. * - * Ownership of the operator is transferred to the preconditioner. + * The builder is called from BlockPreconditioner::updateForState() with the + * current monolithic nonlinear state and its block offsets. The rebuilt operator + * is then used the next time SetOperator() configures the block solvers. + * + * Provide @a initial_operator when SetOperator() may be called before the first + * state update. Otherwise, currentOperator() will fail until updateForState() + * has produced an operator. */ -using BlockOverride = std::pair>; +BlockProviderOverride makeStateDependentBlockProviderOverride( + int block_index, StateDependentBlockOperatorBuilder builder, + std::unique_ptr initial_operator = nullptr); /** * @class BlockPreconditioner * @brief Base class for block preconditioners that own one sub-solver per block. */ -class BlockPreconditioner : public mfem::Solver { +class BlockPreconditioner : public mfem::Solver, public StateDependentSolver { public: /** @brief Return the number of sub-solvers owned by this preconditioner. */ int numSubSolvers() const { return num_blocks_; } @@ -39,6 +111,9 @@ class BlockPreconditioner : public mfem::Solver { return mfem_solvers_[static_cast(i)].get(); } + /// @overload + void updateForState(const mfem::Vector& state, const mfem::Array& block_offsets) override; + virtual ~BlockPreconditioner(); protected: @@ -60,8 +135,8 @@ class BlockPreconditioner : public mfem::Solver { /// @brief Owned MFEM solver for each block. mutable std::vector> mfem_solvers_; - /// @brief Optional per-block operators; null entries use the corresponding Jacobian diagonal block. - std::vector> block_op_overrides_; + /// @brief Per-block operator providers; null entries use the corresponding Jacobian diagonal block. + std::vector> block_op_providers_; }; /** @@ -80,11 +155,11 @@ class BlockDiagonalPreconditioner : public BlockPreconditioner { * @brief Construct a new N by N block diagonal preconditioner. * * @param solvers One solver per block (size must match number of blocks). - * @param overrides Optional list of (block index, operator) pairs used in place - * of the corresponding Jacobian diagonal block. + * @param overrides Optional provider overrides used in place of the + * corresponding Jacobian diagonal blocks. */ BlockDiagonalPreconditioner(std::vector> solvers, - std::vector overrides = {}); + std::vector overrides = {}); /** * @brief The action of the precondition on the block vector (b_1, ..., b_n) @@ -136,12 +211,12 @@ class BlockTriangularPreconditioner : public BlockPreconditioner { * * @param solvers One solver per diagonal block (size must match number of blocks). * @param type Sweep type (lower, upper, or symmetric). - * @param overrides Optional list of (block index, operator) pairs used in place - * of the corresponding Jacobian diagonal block. + * @param overrides Optional provider overrides used in place of the + * corresponding Jacobian diagonal blocks. */ BlockTriangularPreconditioner(std::vector> solvers, BlockTriangularType type = BlockTriangularType::Lower, - std::vector overrides = {}); + std::vector overrides = {}); /** * @brief The action of the precondition on the block vector (b_1, ..., b_n) @@ -201,7 +276,7 @@ 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 provided via the overrides list for block index 1. */ + Custom /**< Use a custom operator provider for block index 1. */ }; /** @@ -220,15 +295,13 @@ class BlockSchurPreconditioner : public BlockPreconditioner { * @param solvers Two solvers, for $ A_{11} $ and the Schur complement approximation. * @param type Preconditioner variant (diagonal, lower, upper, or full). * @param approxType Schur complement approximation strategy for the (1,1) block. - * @param overrides Optional list of (block index, operator) pairs used in place - * of the corresponding Jacobian diagonal block. For Schur systems, index - * 0 overrides $A_{11}$ and index 1 provides a custom Schur operator when - * approxType is SchurApproxType::Custom. + * @param overrides Optional provider overrides. Index 0 overrides $A_{11}$ and + * index 1 provides a custom Schur operator when approxType is SchurApproxType::Custom. */ BlockSchurPreconditioner(std::vector> solvers, BlockSchurType type = BlockSchurType::Diagonal, SchurApproxType approxType = SchurApproxType::DiagInv, - std::vector overrides = {}); + std::vector overrides = {}); /** * @brief The action of the precondition on the block vector (b_1, b_2) @@ -260,7 +333,7 @@ class BlockSchurPreconditioner : public BlockPreconditioner { // Schur complement approximation operator used by solver for block (1,1). // // 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_overrides_[1] and referenced + // S_approx_owned_. For Custom, the approximation is provided via block_op_providers_[1] and referenced // non-owningly via S_approx_view_. mutable std::unique_ptr S_approx_owned_; const mfem::Operator* S_approx_view_ = nullptr; diff --git a/src/smith/numerics/equation_solver.cpp b/src/smith/numerics/equation_solver.cpp index d7adc577e6..8ee23bd5a2 100644 --- a/src/smith/numerics/equation_solver.cpp +++ b/src/smith/numerics/equation_solver.cpp @@ -1017,6 +1017,22 @@ EquationSolver::EquationSolver(NonlinearSolverOptions nonlinear_opts, LinearSolv attachConvergenceManager(); } +EquationSolver::EquationSolver(NonlinearSolverOptions nonlinear_opts, LinearSolverOptions lin_opts, + std::unique_ptr preconditioner, MPI_Comm comm) +{ + SLIC_ERROR_ROOT_IF(!preconditioner, "Custom EquationSolver preconditioner must be non-null"); + + auto [lin_solver, attached_preconditioner] = + buildLinearSolverAndPreconditioner(lin_opts, std::move(preconditioner), comm); + + lin_solver_ = std::move(lin_solver); + preconditioner_ = std::move(attached_preconditioner); + nonlin_solver_ = buildNonlinearSolver(nonlinear_opts, lin_opts, *preconditioner_, comm); + convergence_manager_ = std::make_shared(comm, nonlinear_opts.absolute_tol, + nonlinear_opts.relative_tol); + attachConvergenceManager(); +} + EquationSolver::EquationSolver(std::unique_ptr nonlinear_solver, std::unique_ptr linear_solver, std::unique_ptr preconditioner) @@ -1027,6 +1043,11 @@ EquationSolver::EquationSolver(std::unique_ptr nonlinear_sol nonlin_solver_ = std::move(nonlinear_solver); lin_solver_ = std::move(linear_solver); preconditioner_ = std::move(preconditioner); + if (preconditioner_) { + if (auto* iterative_solver = dynamic_cast(lin_solver_.get())) { + iterative_solver->SetPreconditioner(*preconditioner_); + } + } } void EquationSolver::attachConvergenceManager() const @@ -1265,21 +1286,19 @@ std::unique_ptr buildNonlinearSolver(NonlinearSolverOptions return nonlinear_solver; } -std::pair, std::unique_ptr> buildLinearSolverAndPreconditioner( - LinearSolverOptions linear_opts, MPI_Comm comm) -{ - auto preconditioner = buildPreconditioner(linear_opts, comm); +namespace { +std::unique_ptr buildLinearSolver(LinearSolverOptions linear_opts, MPI_Comm comm, + mfem::Solver* preconditioner) +{ if (linear_opts.linear_solver == LinearSolver::SuperLU) { - auto lin_solver = std::make_unique(linear_opts.print_level, comm); - return {std::move(lin_solver), std::move(preconditioner)}; + return std::make_unique(linear_opts.print_level, comm); } #ifdef MFEM_USE_STRUMPACK if (linear_opts.linear_solver == LinearSolver::Strumpack) { - auto lin_solver = std::make_unique(linear_opts.print_level, comm); - return {std::move(lin_solver), std::move(preconditioner)}; + return std::make_unique(linear_opts.print_level, comm); } #endif @@ -1324,7 +1343,37 @@ std::pair, std::unique_ptr> buildLin iter_lin_solver->SetPreconditioner(*preconditioner); } - return {std::move(iter_lin_solver), std::move(preconditioner)}; + return iter_lin_solver; +} + +} // namespace + +std::pair, std::unique_ptr> buildLinearSolverAndPreconditioner( + LinearSolverOptions linear_opts, MPI_Comm comm) +{ + auto preconditioner = buildPreconditioner(linear_opts, comm); + auto lin_solver = buildLinearSolver(linear_opts, comm, preconditioner.get()); + return {std::move(lin_solver), std::move(preconditioner)}; +} + +std::pair, std::unique_ptr> buildLinearSolverAndPreconditioner( + LinearSolverOptions linear_opts, std::unique_ptr preconditioner, MPI_Comm comm) +{ + auto lin_solver = buildLinearSolver(linear_opts, comm, preconditioner.get()); + return {std::move(lin_solver), std::move(preconditioner)}; +} + +std::vector> buildBlockPreconditionerSubSolvers( + const std::vector& sub_block_options, MPI_Comm comm) +{ + std::vector> sub_solvers; + sub_solvers.reserve(sub_block_options.size()); + for (const auto& opt : sub_block_options) { + auto [lin_solver, preconditioner] = buildLinearSolverAndPreconditioner(opt, comm); + sub_solvers.push_back( + std::make_unique(std::move(lin_solver), std::move(preconditioner))); + } + return sub_solvers; } bool requiresMonolithicOperator(const LinearSolverOptions& linear_opts) @@ -1437,11 +1486,7 @@ std::unique_ptr buildPreconditioner(LinearSolverOptions linear_opt preconditioner_solver = std::move(amgfcontact_preconditioner); } else if (preconditioner == Preconditioner::BlockDiagonal || preconditioner == Preconditioner::BlockTriangular || preconditioner == Preconditioner::BlockSchur) { - std::vector> inner_solvers; - for (const auto& opt : linear_opts.sub_block_linear_solver_options) { - auto [lin, prec] = buildLinearSolverAndPreconditioner(opt, comm); - inner_solvers.push_back(std::make_unique(std::move(lin), std::move(prec))); - } + auto inner_solvers = buildBlockPreconditionerSubSolvers(linear_opts.sub_block_linear_solver_options, comm); if (preconditioner == Preconditioner::BlockDiagonal) { preconditioner_solver = std::make_unique(std::move(inner_solvers)); diff --git a/src/smith/numerics/equation_solver.hpp b/src/smith/numerics/equation_solver.hpp index baac9936e1..8e239c42b8 100644 --- a/src/smith/numerics/equation_solver.hpp +++ b/src/smith/numerics/equation_solver.hpp @@ -16,6 +16,7 @@ #include #include #include +#include #include "mpi.h" #include "mfem.hpp" @@ -69,6 +70,17 @@ class EquationSolver { MPI_Comm comm = MPI_COMM_WORLD); // _build_equationsolver_end + /** + * @brief Construct an equation solver object with a custom owned preconditioner. + * + * @param nonlinear_opts The options to configure the nonlinear solution scheme + * @param lin_opts The options to configure the underlying linear solution scheme + * @param preconditioner Custom preconditioner to attach to the linear solver + * @param comm The MPI communicator for the supplied nonlinear operators and HypreParVectors + */ + EquationSolver(NonlinearSolverOptions nonlinear_opts, LinearSolverOptions lin_opts, + std::unique_ptr preconditioner, MPI_Comm comm = MPI_COMM_WORLD); + /** * Updates the solver with the provided operator * @param[in] op The operator (nonlinear system of equations) to use, "F" in F(x) = 0 @@ -114,15 +126,22 @@ class EquationSolver { /** * Returns the underlying preconditioner - * @return A pointer to the underlying preconditioner - * @note This may be null if a preconditioner is not given + * @return A reference to the underlying preconditioner */ - mfem::Solver& preconditioner() { return *preconditioner_; } + mfem::Solver& preconditioner() + { + MFEM_VERIFY(preconditioner_, "EquationSolver has no preconditioner"); + return *preconditioner_; + } /** * @overload */ - const mfem::Solver& preconditioner() const { return *preconditioner_; } + const mfem::Solver& preconditioner() const + { + MFEM_VERIFY(preconditioner_, "EquationSolver has no preconditioner"); + return *preconditioner_; + } /** * Input file parameters specific to this class @@ -305,6 +324,31 @@ std::unique_ptr buildNonlinearSolver(NonlinearSolverOptions std::pair, std::unique_ptr> buildLinearSolverAndPreconditioner( LinearSolverOptions linear_opts = {}, MPI_Comm comm = MPI_COMM_WORLD); +/** + * @brief Build the linear solver and attach a supplied preconditioner. + * + * @param linear_opts The options to configure the linear solver + * @param preconditioner Custom preconditioner to attach to the linear solver + * @param comm The MPI communicator for the supplied HypreParMatrix and HypreParVectors + * @return A pair containing the constructed linear solver and supplied preconditioner + */ +std::pair, std::unique_ptr> buildLinearSolverAndPreconditioner( + LinearSolverOptions linear_opts, std::unique_ptr preconditioner, MPI_Comm comm = MPI_COMM_WORLD); + +/** + * @brief Build wrapped sub-solvers for block preconditioners. + * + * Each sub-block option creates a linear solver and its preconditioner, then + * wraps them in SolverWithPreconditioner so the preconditioner lifetime is tied + * to the sub-solver. + * + * @param sub_block_options One linear solver option set per block. + * @param comm The MPI communicator for the supplied HypreParMatrix and HypreParVectors. + * @return One owned solver per block, ready to pass to a block preconditioner. + */ +std::vector> buildBlockPreconditionerSubSolvers( + const std::vector& sub_block_options, MPI_Comm comm = MPI_COMM_WORLD); + /** * @brief Return true if the configured linear solve stack requires block operators to be merged. * @param linear_opts Linear solver and preconditioner options. diff --git a/src/smith/numerics/solver_with_preconditioner.hpp b/src/smith/numerics/solver_with_preconditioner.hpp index 27385763fa..e9959319f3 100644 --- a/src/smith/numerics/solver_with_preconditioner.hpp +++ b/src/smith/numerics/solver_with_preconditioner.hpp @@ -10,13 +10,15 @@ #include "mfem.hpp" +#include "smith/numerics/state_dependent_solver.hpp" + namespace smith { /// @brief Simple wrapper that owns a linear solver and its preconditioner. /// /// This is used to keep a preconditioner alive when it is referenced by an /// iterative solver (e.g. GMRES) via SetPreconditioner(). -class SolverWithPreconditioner : public mfem::Solver { +class SolverWithPreconditioner : public mfem::Solver, public StateDependentSolver { public: /// @brief Construct from an owned linear solver and (optional) preconditioner. /// @param linear_solver Owned linear solver (must be non-null). @@ -45,6 +47,17 @@ class SolverWithPreconditioner : public mfem::Solver { linear_solver_->Mult(x, y); } + /// @brief Refresh state-dependent owned solver data. + void updateForState(const mfem::Vector& state, const mfem::Array& block_offsets) override + { + if (auto* state_dependent_solver = dynamic_cast(linear_solver_.get())) { + state_dependent_solver->updateForState(state, block_offsets); + } + if (auto* state_dependent_solver = dynamic_cast(preconditioner_.get())) { + state_dependent_solver->updateForState(state, block_offsets); + } + } + /// @brief Non-owning access to the underlying linear solver. mfem::Solver* linearSolver() const { return linear_solver_.get(); } diff --git a/src/smith/numerics/state_dependent_solver.hpp b/src/smith/numerics/state_dependent_solver.hpp new file mode 100644 index 0000000000..ab4e474b13 --- /dev/null +++ b/src/smith/numerics/state_dependent_solver.hpp @@ -0,0 +1,34 @@ +// Copyright (c) Lawrence Livermore National Security, LLC and +// other Smith Project Developers. See the top-level LICENSE file for +// details. +// +// SPDX-License-Identifier: (BSD-3-Clause) + +/** + * @file state_dependent_solver.hpp + * + * @brief Interface for solvers that refresh data from a nonlinear state. + */ + +#pragma once + +#include "mfem.hpp" + +namespace smith { + +/** + * @brief Interface for solvers that can refresh state-dependent internals. + */ +class StateDependentSolver { + public: + virtual ~StateDependentSolver() = default; + + /** + * @brief Refresh solver-owned data for the current nonlinear state. + * @param state Monolithic state vector at the current nonlinear iterate. + * @param block_offsets Offsets describing the block layout of @a state. + */ + virtual void updateForState(const mfem::Vector& state, const mfem::Array& block_offsets) = 0; +}; + +} // namespace smith 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 21f4aa8039..29810407f0 100644 --- a/src/smith/numerics/tests/test_block_preconditioner_custom_operators.cpp +++ b/src/smith/numerics/tests/test_block_preconditioner_custom_operators.cpp @@ -246,6 +246,47 @@ std::unique_ptr makeLocalScaledIdentityOp(int n, double c) return std::unique_ptr(mat); } +std::unique_ptr makeMutableLocalScaledIdentityOp(int n, double c) +{ + auto mat = std::make_unique(n); + for (int i = 0; i < n; i++) { + mat->Add(i, i, c); + } + mat->Finalize(); + return mat; +} + +class OperatorDiagonalSolver : public mfem::Solver { + public: + void SetOperator(const mfem::Operator& op) override + { + MFEM_VERIFY(op.Height() == op.Width(), "OperatorDiagonalSolver requires a square operator"); + height = op.Height(); + width = op.Width(); + diag_.SetSize(height); + + mfem::Vector e(width); + mfem::Vector y(height); + for (int i = 0; i < width; ++i) { + e = 0.0; + e[i] = 1.0; + op.Mult(e, y); + diag_[i] = y[i]; + } + } + + void Mult(const mfem::Vector& x, mfem::Vector& y) const override + { + y.SetSize(x.Size()); + for (int i = 0; i < x.Size(); ++i) { + y[i] = x[i] / diag_[i]; + } + } + + private: + mfem::Vector diag_; +}; + } // namespace /* ============================================================ Tests @@ -274,9 +315,11 @@ TEST(BlockDiagonalPreconditionerCustom, IdentityActsAsIdentity) override_mats.push_back(makeHypreScaledIdentity(2, 1.0)); // M1 override_mats.push_back(makeHypreScaledIdentity(3, 1.0)); // M2 - std::vector>> overrides; - overrides.emplace_back(0, std::unique_ptr(std::move(override_mats[0].A))); - overrides.emplace_back(1, std::unique_ptr(std::move(override_mats[1].A))); + std::vector overrides; + overrides.push_back( + smith::makeFixedBlockProviderOverride(0, std::unique_ptr(std::move(override_mats[0].A)))); + overrides.push_back( + smith::makeFixedBlockProviderOverride(1, std::unique_ptr(std::move(override_mats[1].A)))); smith::BlockDiagonalPreconditioner P(std::move(solvers), std::move(overrides)); @@ -309,8 +352,9 @@ TEST(BlockDiagonalPreconditionerCustom, PartialOverrideUsesJacobianForOthers) std::vector override_mats; override_mats.push_back(makeHypreScaledIdentity(2, 1.0)); // override block 1 - std::vector overrides; - overrides.emplace_back(1, std::unique_ptr(std::move(override_mats[0].A))); + std::vector overrides; + overrides.push_back( + smith::makeFixedBlockProviderOverride(1, std::unique_ptr(std::move(override_mats[0].A)))); smith::BlockDiagonalPreconditioner P(std::move(solvers), std::move(overrides)); P.SetOperator(A); @@ -344,8 +388,9 @@ TEST(BlockDiagonalPreconditionerCustom, OverrideBeatsBadJacobianBlock) std::vector override_mats; override_mats.push_back(makeHypreScaledIdentity(2, 1.0)); // override A11 with I - std::vector overrides; - overrides.emplace_back(0, std::unique_ptr(std::move(override_mats[0].A))); + std::vector overrides; + overrides.push_back( + smith::makeFixedBlockProviderOverride(0, std::unique_ptr(std::move(override_mats[0].A)))); smith::BlockDiagonalPreconditioner P(std::move(solvers), std::move(overrides)); P.SetOperator(A); @@ -381,8 +426,9 @@ TEST(BlockTriangularPreconditionerCustom, LowerSweepUsesOverrideDiagonal) std::vector override_mats; override_mats.push_back(makeHypreScaledIdentity(2, 6.0)); // override A22_used - std::vector overrides; - overrides.emplace_back(1, std::unique_ptr(std::move(override_mats[0].A))); + std::vector overrides; + overrides.push_back( + smith::makeFixedBlockProviderOverride(1, std::unique_ptr(std::move(override_mats[0].A)))); smith::BlockTriangularPreconditioner P(std::move(solvers), smith::BlockTriangularType::Lower, std::move(overrides)); P.SetOperator(A); @@ -427,8 +473,9 @@ TEST(BlockTriangularPreconditionerCustom, UpperSweepUsesOverrideDiagonal) std::vector override_mats; override_mats.push_back(makeHypreScaledIdentity(2, 4.0)); // override A11_used - std::vector overrides; - overrides.emplace_back(0, std::unique_ptr(std::move(override_mats[0].A))); + std::vector overrides; + overrides.push_back( + smith::makeFixedBlockProviderOverride(0, std::unique_ptr(std::move(override_mats[0].A)))); smith::BlockTriangularPreconditioner P(std::move(solvers), smith::BlockTriangularType::Upper, std::move(overrides)); P.SetOperator(A); @@ -524,8 +571,9 @@ TEST(BlockSchurPreconditionerCustom, FullWithExactSchurOverrideIsExactInverse) solvers.push_back(std::make_unique()); solvers.push_back(std::make_unique()); - std::vector overrides; - overrides.emplace_back(1, std::unique_ptr(std::move(override_mats[0].A))); + std::vector overrides; + overrides.push_back( + smith::makeFixedBlockProviderOverride(1, std::unique_ptr(std::move(override_mats[0].A)))); smith::BlockSchurPreconditioner P(std::move(solvers), smith::BlockSchurType::Full, smith::SchurApproxType::Custom, std::move(overrides)); @@ -598,8 +646,9 @@ TEST(BlockSchurPreconditionerCustom, Block0OverrideIsUsed) std::vector override_mats; override_mats.push_back(makeHypreScaledIdentity(n, 4.0)); - std::vector overrides; - overrides.emplace_back(0, std::unique_ptr(std::move(override_mats[0].A))); + std::vector overrides; + overrides.push_back( + smith::makeFixedBlockProviderOverride(0, std::unique_ptr(std::move(override_mats[0].A)))); smith::BlockSchurPreconditioner P(std::move(solvers), smith::BlockSchurType::Diagonal, smith::SchurApproxType::A22Only, std::move(overrides)); @@ -636,8 +685,9 @@ TEST(BlockSchurPreconditionerCustom, CustomOverrideNotConsumedOnRepeatedSetOpera std::vector override_mats; override_mats.push_back(makeHypreScaledIdentity(n, 7.0)); - std::vector overrides; - overrides.emplace_back(1, std::unique_ptr(std::move(override_mats[0].A))); + std::vector overrides; + overrides.push_back( + smith::makeFixedBlockProviderOverride(1, std::unique_ptr(std::move(override_mats[0].A)))); smith::BlockSchurPreconditioner P(std::move(solvers), smith::BlockSchurType::Diagonal, smith::SchurApproxType::Custom, std::move(overrides)); @@ -653,30 +703,77 @@ TEST(BlockSchurPreconditionerCustom, CustomOverrideNotConsumedOnRepeatedSetOpera EXPECT_NEAR(x[3], b[3] / 7.0, 1e-12); } -TEST(BlockDiagonalPreconditionerCustom, ThrowsOnOutOfRangeOverrideIndex) +// Verifies state-dependent Schur providers update the Schur solve. +TEST(BlockSchurPreconditionerCustom, StateDependentProviderUpdatesSchurSolve) { - Array offsets({0, 2, 4}); - EXPECT_THROW( - { - auto solvers = makeIdentitySolvers(2); - std::vector overrides; - overrides.emplace_back(2, makeLocalScaledIdentityOp(2, 1.0)); - smith::BlockDiagonalPreconditioner P(std::move(solvers), std::move(overrides)); - }, - std::out_of_range); + 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()); + + std::vector> solvers; + solvers.push_back(std::make_unique()); + solvers.push_back(std::make_unique()); + + auto update_count = std::make_shared(0); + std::vector overrides; + overrides.push_back(smith::makeStateDependentBlockProviderOverride( + 1, [update_count](const mfem::Vector& state, const mfem::Array& block_offsets) { + ++(*update_count); + const int block_size = block_offsets[2] - block_offsets[1]; + return makeMutableLocalScaledIdentityOp(block_size, state[block_offsets[1]]); + })); + + smith::BlockSchurPreconditioner P(std::move(solvers), smith::BlockSchurType::Diagonal, smith::SchurApproxType::Custom, + std::move(overrides)); + + Vector state(2 * n); + state = 0.0; + state[offsets[1]] = 5.0; + P.updateForState(state, offsets); + P.SetOperator(A); + + Vector b(2 * n), x(2 * n); + b.Randomize(); + P.Mult(b, x); + + EXPECT_NEAR(x[0], b[0] / 2.0, 1e-12); + EXPECT_NEAR(x[1], b[1] / 2.0, 1e-12); + EXPECT_NEAR(x[2], b[2] / 5.0, 1e-12); + EXPECT_NEAR(x[3], b[3] / 5.0, 1e-12); + + state[offsets[1]] = 7.0; + P.updateForState(state, offsets); + P.SetOperator(A); + P.Mult(b, x); + + EXPECT_EQ(*update_count, 2); + EXPECT_NEAR(x[0], b[0] / 2.0, 1e-12); + EXPECT_NEAR(x[1], b[1] / 2.0, 1e-12); + EXPECT_NEAR(x[2], b[2] / 7.0, 1e-12); + EXPECT_NEAR(x[3], b[3] / 7.0, 1e-12); } -TEST(BlockDiagonalPreconditionerCustom, ThrowsOnNullOverrideOperator) +TEST(BlockDiagonalPreconditionerCustom, ThrowsOnOutOfRangeOverrideIndex) { Array offsets({0, 2, 4}); EXPECT_THROW( { auto solvers = makeIdentitySolvers(2); - std::vector overrides; - overrides.emplace_back(0, std::unique_ptr()); + std::vector overrides; + overrides.push_back(smith::makeFixedBlockProviderOverride(2, makeLocalScaledIdentityOp(2, 1.0))); smith::BlockDiagonalPreconditioner P(std::move(solvers), std::move(overrides)); }, - std::invalid_argument); + std::out_of_range); } TEST(BlockDiagonalPreconditionerCustom, ThrowsOnDuplicateOverrideIndex) @@ -685,9 +782,9 @@ TEST(BlockDiagonalPreconditionerCustom, ThrowsOnDuplicateOverrideIndex) EXPECT_THROW( { auto solvers = makeIdentitySolvers(2); - std::vector overrides; - overrides.emplace_back(0, makeLocalScaledIdentityOp(2, 1.0)); - overrides.emplace_back(0, makeLocalScaledIdentityOp(2, 2.0)); + std::vector overrides; + overrides.push_back(smith::makeFixedBlockProviderOverride(0, makeLocalScaledIdentityOp(2, 1.0))); + overrides.push_back(smith::makeFixedBlockProviderOverride(0, makeLocalScaledIdentityOp(2, 2.0))); smith::BlockDiagonalPreconditioner P(std::move(solvers), std::move(overrides)); }, std::invalid_argument); @@ -699,37 +796,23 @@ TEST(BlockTriangularPreconditionerCustom, ThrowsOnOutOfRangeOverrideIndex) EXPECT_THROW( { auto solvers = makeIdentitySolvers(2); - std::vector overrides; - overrides.emplace_back(2, makeLocalScaledIdentityOp(2, 1.0)); + std::vector overrides; + overrides.push_back(smith::makeFixedBlockProviderOverride(2, makeLocalScaledIdentityOp(2, 1.0))); smith::BlockTriangularPreconditioner P(std::move(solvers), smith::BlockTriangularType::Lower, std::move(overrides)); }, std::out_of_range); } -TEST(BlockTriangularPreconditionerCustom, ThrowsOnNullOverrideOperator) -{ - Array offsets({0, 2, 4}); - EXPECT_THROW( - { - auto solvers = makeIdentitySolvers(2); - std::vector overrides; - overrides.emplace_back(0, std::unique_ptr()); - smith::BlockTriangularPreconditioner P(std::move(solvers), smith::BlockTriangularType::Lower, - std::move(overrides)); - }, - std::invalid_argument); -} - TEST(BlockTriangularPreconditionerCustom, ThrowsOnDuplicateOverrideIndex) { Array offsets({0, 2, 4}); EXPECT_THROW( { auto solvers = makeIdentitySolvers(2); - std::vector overrides; - overrides.emplace_back(1, makeLocalScaledIdentityOp(2, 1.0)); - overrides.emplace_back(1, makeLocalScaledIdentityOp(2, 2.0)); + std::vector overrides; + overrides.push_back(smith::makeFixedBlockProviderOverride(1, makeLocalScaledIdentityOp(2, 1.0))); + overrides.push_back(smith::makeFixedBlockProviderOverride(1, makeLocalScaledIdentityOp(2, 2.0))); smith::BlockTriangularPreconditioner P(std::move(solvers), smith::BlockTriangularType::Lower, std::move(overrides)); }, @@ -742,22 +825,23 @@ TEST(BlockSchurPreconditionerCustom, ThrowsOnOutOfRangeOverrideIndex) EXPECT_THROW( { auto solvers = makeIdentitySolvers(2); - std::vector overrides; - overrides.emplace_back(2, makeLocalScaledIdentityOp(2, 1.0)); + std::vector overrides; + overrides.push_back(smith::makeFixedBlockProviderOverride(2, makeLocalScaledIdentityOp(2, 1.0))); smith::BlockSchurPreconditioner P(std::move(solvers), smith::BlockSchurType::Diagonal, smith::SchurApproxType::A22Only, std::move(overrides)); }, std::out_of_range); } -TEST(BlockSchurPreconditionerCustom, ThrowsOnNullOverrideOperator) +// Verifies Schur preconditioners reject null override providers. +TEST(BlockSchurPreconditionerCustom, ThrowsOnNullOverrideProvider) { Array offsets({0, 2, 4}); EXPECT_THROW( { auto solvers = makeIdentitySolvers(2); - std::vector overrides; - overrides.emplace_back(0, std::unique_ptr()); + std::vector overrides; + overrides.emplace_back(0, std::unique_ptr()); smith::BlockSchurPreconditioner P(std::move(solvers), smith::BlockSchurType::Diagonal, smith::SchurApproxType::A22Only, std::move(overrides)); }, @@ -770,9 +854,9 @@ TEST(BlockSchurPreconditionerCustom, ThrowsOnDuplicateOverrideIndex) EXPECT_THROW( { auto solvers = makeIdentitySolvers(2); - std::vector overrides; - overrides.emplace_back(1, makeLocalScaledIdentityOp(2, 1.0)); - overrides.emplace_back(1, makeLocalScaledIdentityOp(2, 2.0)); + std::vector overrides; + overrides.push_back(smith::makeFixedBlockProviderOverride(1, makeLocalScaledIdentityOp(2, 1.0))); + overrides.push_back(smith::makeFixedBlockProviderOverride(1, makeLocalScaledIdentityOp(2, 2.0))); smith::BlockSchurPreconditioner P(std::move(solvers), smith::BlockSchurType::Diagonal, smith::SchurApproxType::A22Only, std::move(overrides)); }, diff --git a/src/smith/numerics/tests/test_equationsolver.cpp b/src/smith/numerics/tests/test_equationsolver.cpp index c6b895860f..ea3272489e 100644 --- a/src/smith/numerics/tests/test_equationsolver.cpp +++ b/src/smith/numerics/tests/test_equationsolver.cpp @@ -110,6 +110,31 @@ class ManagedHalvingSolver : public mfem::NewtonSolver, public smith::Convergenc std::shared_ptr convergence_manager_ = nullptr; }; +class ScalingPreconditioner : public mfem::Solver { + public: + explicit ScalingPreconditioner(double scale) : scale_(scale) {} + + void SetOperator(const mfem::Operator& op) override + { + height = op.Height(); + width = op.Width(); + ++set_operator_count_; + } + + void Mult(const mfem::Vector& x, mfem::Vector& y) const override + { + y.SetSize(x.Size()); + y = x; + y *= scale_; + } + + int setOperatorCount() const { return set_operator_count_; } + + private: + double scale_; + int set_operator_count_ = 0; +}; + } // namespace class EquationSolverSuite : public testing::TestWithParam { @@ -232,6 +257,40 @@ TEST(EquationSolverManualConvergence, InjectedManagedSolverSupportsScalarConverg EXPECT_LE(residual.Norml2(), 1.0e-2 * eq_solver.nonlinearSolver().GetInitialNorm()); } +// Verifies EquationSolver owns and attaches a caller-provided preconditioner. +TEST(EquationSolverCustomPreconditioner, OptionsConstructorOwnsAndAttachesPreconditioner) +{ + LinearSolverOptions lin_opts; + lin_opts.linear_solver = LinearSolver::PrecondOnly; + lin_opts.preconditioner = Preconditioner::None; + + NonlinearSolverOptions nonlin_opts; + nonlin_opts.nonlin_solver = NonlinearSolver::Newton; + nonlin_opts.print_level = 0; + + auto preconditioner = std::make_unique(3.0); + auto* preconditioner_ptr = preconditioner.get(); + EquationSolver eq_solver(nonlin_opts, lin_opts, std::move(preconditioner), MPI_COMM_WORLD); + + EXPECT_EQ(&eq_solver.preconditioner(), preconditioner_ptr); + + mfem::SparseMatrix op(2); + op.Add(0, 0, 1.0); + op.Add(1, 1, 1.0); + op.Finalize(); + + eq_solver.linearSolver().SetOperator(op); + + mfem::Vector x(2), y(2); + x[0] = 2.0; + x[1] = -1.0; + eq_solver.linearSolver().Mult(x, y); + + EXPECT_EQ(preconditioner_ptr->setOperatorCount(), 1); + EXPECT_NEAR(y[0], 6.0, 1e-12); + EXPECT_NEAR(y[1], -3.0, 1e-12); +} + /** * @brief Nonlinear solvers to test. Always includes NonlinearSolver::Newton and NonlinearSolver::LBFGS * If SMITH_USE_SUNDIALS is set, adds: NonlinearSolver::KINFullStep, NonlinearSolver::KINBacktrackingLineSearch, and