diff --git a/src/smith/differentiable_numerics/tests/CMakeLists.txt b/src/smith/differentiable_numerics/tests/CMakeLists.txt index 06681ec85..ba088e0fb 100644 --- a/src/smith/differentiable_numerics/tests/CMakeLists.txt +++ b/src/smith/differentiable_numerics/tests/CMakeLists.txt @@ -10,6 +10,7 @@ set(differentiable_numerics_test_source test_field_state.cpp test_explicit_dynamics.cpp test_solid_mechanics_state_advancer.cpp + test_porous_heat_sink.cpp test_thermo_mechanics.cpp ) diff --git a/src/smith/differentiable_numerics/tests/test_porous_heat_sink.cpp b/src/smith/differentiable_numerics/tests/test_porous_heat_sink.cpp new file mode 100644 index 000000000..abf2c4dfa --- /dev/null +++ b/src/smith/differentiable_numerics/tests/test_porous_heat_sink.cpp @@ -0,0 +1,362 @@ +#include + +#include "smith/infrastructure/application_manager.hpp" +#include "smith/numerics/equation_solver.hpp" +#include "smith/numerics/solver_config.hpp" +#include "smith/mesh_utils/mesh_utils.hpp" +#include "smith/physics/state/state_manager.hpp" +#include "smith/physics/boundary_conditions/boundary_condition_manager.hpp" +#include "smith/physics/functional_weak_form.hpp" +#include "smith/differentiable_numerics/field_state.hpp" +#include "smith/differentiable_numerics/state_advancer.hpp" +#include "smith/smith_config.hpp" +#include "smith/differentiable_numerics/differentiable_solver.hpp" +#include "smith/differentiable_numerics/nonlinear_solve.hpp" +#include "smith/differentiable_numerics/paraview_writer.hpp" +#include "smith/numerics/block_preconditioner.hpp" + +#include "gretl/data_store.hpp" +#include "gretl/wang_checkpoint_strategy.hpp" + +using namespace smith; + +using ShapeDispSpace = H1<1, 2>; +using Space = H1<1>; +using GammaSpace = L2<1>; + +enum class BlockSolverType +{ + Direct, + Iterative, + BoomerAMG +}; +enum class BlockPrecondType +{ + Diagonal, + TriLower, + TriUpper, + TriSym, + SchurLower, + SchurUpper, + SchurDiag, + SchurFull +}; + +struct BlockTestParams { + BlockSolverType solver_type; + BlockPrecondType precond_type; +}; + +std::string BlockParamNameGenerator(const ::testing::TestParamInfo& info) +{ + auto solver_to_str = [](BlockSolverType t) { + switch (t) { + case BlockSolverType::Direct: + return "Direct"; + case BlockSolverType::Iterative: + return "Iterative"; + case BlockSolverType::BoomerAMG: + return "BoomerAMG"; + } + return "Unknown"; + }; + auto precond_to_str = [](BlockPrecondType t) { + switch (t) { + case BlockPrecondType::Diagonal: + return "Diag"; + case BlockPrecondType::TriLower: + return "TriLower"; + case BlockPrecondType::TriUpper: + return "TriUpper"; + case BlockPrecondType::TriSym: + return "TriSym"; + case BlockPrecondType::SchurLower: + return "SchurLower"; + case BlockPrecondType::SchurUpper: + return "SchurUpper"; + case BlockPrecondType::SchurDiag: + return "SchurDiag"; + case BlockPrecondType::SchurFull: + return "SchurFull"; + } + return "Unknown"; + }; + return std::string(solver_to_str(info.param.solver_type)) + "_" + precond_to_str(info.param.precond_type); +} + +struct HeatSinkOptions { + double kappa_0 = 0.5; + double sigma_0 = 5.0; + double eta = 1.5; + double epsilon_m = 1.0; + double epsilon_n = 0.5; + double a_m = 0.0; + double a_n = 5e+6; + double h = 0.01; + double q_app = 0.0; + double T_app = 0.0; + double f_mb = 1.0; +}; + +class MeshFixture : public testing::Test { + protected: + double length = 1.0; + double width = 1.0; + int num_elements_x = 32; + int num_elements_y = 32; + double elem_size = length / num_elements_x; + + axom::sidre::DataStore datastore; + std::shared_ptr mesh; + + void SetUp() override + { + smith::StateManager::initialize(datastore, "porous_heat"); + + MPI_Barrier(MPI_COMM_WORLD); + int serial_refinement = 0; + 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); + } +}; + +class BlockPreconditionerTest : public MeshFixture, public ::testing::WithParamInterface {}; + +TEST_P(BlockPreconditionerTest, BlockSolve) +{ + const auto& test_params = GetParam(); + + std::string physics_name = "heatsink"; + auto graph = std::make_shared(std::make_unique(100)); + auto shape_disp = createFieldState(*graph, ShapeDispSpace{}, physics_name + "_shape_displacement", mesh->tag()); + auto T1 = createFieldState(*graph, Space{}, physics_name + "_T1", mesh->tag()); + auto T2 = createFieldState(*graph, Space{}, physics_name + "_T2", mesh->tag()); + auto gamma = createFieldState(*graph, GammaSpace{}, physics_name + "_gamma", mesh->tag()); + smith::FunctionalWeakForm<2, Space, smith::Parameters> T1_form("T1_eqn", mesh, space(T1), + spaces({T1, T2, gamma})); + smith::FunctionalWeakForm<2, Space, smith::Parameters> T2_form("T2_eqn", mesh, space(T2), + spaces({T1, T2, gamma})); + + HeatSinkOptions heatsink_options{.kappa_0 = 0.5, .sigma_0 = 5., .a_n = 1000.0, .h = 0.01, .q_app = -10.0}; + + auto epsilon = [heatsink_options = heatsink_options](auto gamma_) { + return (1.0 - gamma_) * heatsink_options.epsilon_m + heatsink_options.f_mb * gamma_ * heatsink_options.epsilon_n; + }; + + auto a = [heatsink_options = heatsink_options](auto gamma_) { + return (1.0 - gamma_) * heatsink_options.a_m + gamma_ * heatsink_options.a_n; + }; + + auto sigma = [heatsink_options = heatsink_options, epsilon](auto gamma_) { + using std::pow; + return pow(1.0 - epsilon(gamma_), heatsink_options.eta) * heatsink_options.sigma_0; + }; + + auto kappa = [heatsink_options = heatsink_options, epsilon](auto gamma_) { + using std::pow; + return pow(epsilon(gamma_), heatsink_options.eta) * heatsink_options.kappa_0; + }; + + auto q_n = [heatsink_options = heatsink_options](auto T_1, auto T_2) { return heatsink_options.h * (T_1 - T_2); }; + + auto gamma_fun = [](const mfem::Vector& x) -> double { + if (x[0] >= 4.0 / 16.0 && x[0] <= 7.0 / 16.0 && x[1] >= 5.0 / 16.0) + return 1.0; + else if (x[0] >= 14.0 / 16.0 && x[1] >= 5.0 / 16.0) + return 1.0; + else if (x[1] >= 13.0 / 16.0) + return 1.0; + + return 1e-8; + }; + + auto gamma_coef = std::make_shared(gamma_fun); + gamma.get()->project(gamma_coef); + + T1_form.addBodyIntegral(DependsOn<0, 1, 2>{}, mesh->entireBodyName(), + [sigma, a, q_n](double /* t */, auto /* x */, auto T_1, auto T_2, auto gamma_) { + auto [T_1_val, dT_1_dX] = T_1; + auto [T_2_val, dT_2_dX] = T_2; + auto [gamma_val, dgamma_dX] = gamma_; + return smith::tuple{a(gamma_val) * q_n(T_1_val, T_2_val), sigma(gamma_val) * dT_1_dX}; + }); + T2_form.addBodyIntegral(DependsOn<0, 1, 2>{}, mesh->entireBodyName(), + [kappa, a, q_n](double /* t */, auto /* x */, auto T_1, auto T_2, auto gamma_) { + auto [T_1_val, dT_1_dX] = T_1; + auto [T_2_val, dT_2_dX] = T_2; + auto [gamma_val, dgamma_dX] = gamma_; + return smith::tuple{-1.0 * a(gamma_val) * q_n(T_1_val, T_2_val), + kappa(gamma_val) * dT_2_dX}; + }); + + auto T1_bc_manager = std::make_shared(mesh->mfemParMesh()); + auto T2_bc_manager = std::make_shared(mesh->mfemParMesh()); + + auto zero_bcs = std::make_shared([](const mfem::Vector&) { return 0.0; }); + T2_bc_manager->addEssential(std::set{1}, zero_bcs, space(T2), 0); + + mesh->addDomainOfBoundaryElements("heat_spreader", by_attr<2>(2)); + T1_form.addBoundaryIntegral(DependsOn<>{}, "heat_spreader", + [heatsink_options = heatsink_options](double, auto) { return heatsink_options.q_app; }); + + // Block Diagonal Preconditioner + smith::LinearSolverOptions default_linear_options = {.linear_solver = smith::LinearSolver::GMRES, + .preconditioner = smith::Preconditioner::HypreAMG, + .relative_tol = 1.0e-8, + .absolute_tol = 1.0e-12, + .max_iterations = 200, + .print_level = 1}; + + mfem::Array block_offsets_; + block_offsets_.SetSize(3); + block_offsets_[0] = 0; + block_offsets_[1] = T1.get()->space().TrueVSize(); + block_offsets_[2] = T2.get()->space().TrueVSize(); + block_offsets_.PartialSum(); + + std::vector> solvers; + std::vector> preconds; + + // Parameter sweep: construct solvers according to test parameters + if (test_params.solver_type == BlockSolverType::Direct) { + smith::LinearSolverOptions direct_solver_options{.linear_solver = smith::LinearSolver::Strumpack}; + auto [solver1, precond1] = smith::buildLinearSolverAndPreconditioner(direct_solver_options, mesh->getComm()); + auto [solver2, precond2] = smith::buildLinearSolverAndPreconditioner(direct_solver_options, mesh->getComm()); + solvers.push_back(std::move(solver1)); + solvers.push_back(std::move(solver2)); + } else if (test_params.solver_type == BlockSolverType::Iterative) { + smith::LinearSolverOptions iter_solver_options = {.linear_solver = smith::LinearSolver::GMRES, + .preconditioner = smith::Preconditioner::HypreAMG, + .relative_tol = 1.0e-3, + .absolute_tol = 1.0e-6, + .max_iterations = 100, + .print_level = 1}; + auto [solver1, precond1] = smith::buildLinearSolverAndPreconditioner(iter_solver_options, mesh->getComm()); + auto [solver2, precond2] = smith::buildLinearSolverAndPreconditioner(iter_solver_options, mesh->getComm()); + solvers.push_back(std::move(solver1)); + solvers.push_back(std::move(solver2)); + // So that preconds don't go out of scope + preconds.push_back(std::move(precond1)); + preconds.push_back(std::move(precond2)); + } else if (test_params.solver_type == BlockSolverType::BoomerAMG) { + auto solver1 = std::make_unique(); + auto solver2 = std::make_unique(); + solvers.push_back(std::move(solver1)); + solvers.push_back(std::move(solver2)); + } + + std::unique_ptr diff_precond; + + switch (test_params.precond_type) { + case BlockPrecondType::Diagonal: + diff_precond = std::make_unique(block_offsets_, std::move(solvers)); + break; + case BlockPrecondType::TriLower: + diff_precond = std::make_unique(block_offsets_, std::move(solvers), + smith::BlockTriangularType::Lower); + break; + case BlockPrecondType::TriUpper: + diff_precond = std::make_unique(block_offsets_, std::move(solvers), + smith::BlockTriangularType::Upper); + break; + case BlockPrecondType::TriSym: + diff_precond = std::make_unique(block_offsets_, std::move(solvers), + smith::BlockTriangularType::Symmetric); + break; + case BlockPrecondType::SchurLower: + diff_precond = std::make_unique(block_offsets_, std::move(solvers), + smith::BlockSchurType::Lower); + break; + case BlockPrecondType::SchurUpper: + diff_precond = std::make_unique(block_offsets_, std::move(solvers), + smith::BlockSchurType::Upper); + break; + case BlockPrecondType::SchurDiag: + diff_precond = std::make_unique(block_offsets_, std::move(solvers), + smith::BlockSchurType::Diagonal); + break; + case BlockPrecondType::SchurFull: + diff_precond = std::make_unique(block_offsets_, std::move(solvers), + smith::BlockSchurType::Full); + break; + } + + std::unique_ptr linear_solver = std::make_unique(mesh->getComm()); + mfem::GMRESSolver* iter_lin_solver = dynamic_cast(linear_solver.get()); + + iter_lin_solver->iterative_mode = false; + iter_lin_solver->SetRelTol(default_linear_options.relative_tol); + iter_lin_solver->SetAbsTol(default_linear_options.absolute_tol); + iter_lin_solver->SetMaxIter(default_linear_options.max_iterations); + iter_lin_solver->SetPrintLevel(default_linear_options.print_level); + iter_lin_solver->SetPreconditioner(*diff_precond); + + std::shared_ptr d_linear_solver = + std::make_shared(std::move(linear_solver), std::move(diff_precond)); + + auto time = graph->create_state(0.0); + auto dt = graph->create_state(0.025); + size_t cycle = 0; + std::vector params; + auto& T1_params = params; + auto& T2_params = params; + std::vector T1_arguments{T1, T2, gamma}; + std::vector T2_arguments{T1, T2, gamma}; + 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), d_linear_solver.get(), + {T1_bc_manager.get(), T2_bc_manager.get()}); + + // Convergence check + const double rel = iter_lin_solver->GetFinalRelNorm(); + EXPECT_LT(rel, default_linear_options.relative_tol) << "GMRES final relative norm too large"; + + auto pv_writer = smith::createParaviewWriter(*mesh, sols, physics_name); + pv_writer.write(0, 0.0, sols); + + SUCCEED(); +} + +INSTANTIATE_TEST_SUITE_P(BlockPrecondSweep, BlockPreconditionerTest, + ::testing::Values( + // Direct solvers + BlockTestParams{BlockSolverType::Direct, BlockPrecondType::Diagonal}, + BlockTestParams{BlockSolverType::Direct, BlockPrecondType::TriLower}, + BlockTestParams{BlockSolverType::Direct, BlockPrecondType::TriUpper}, + BlockTestParams{BlockSolverType::Direct, BlockPrecondType::TriSym}, + BlockTestParams{BlockSolverType::Direct, BlockPrecondType::SchurLower}, + BlockTestParams{BlockSolverType::Direct, BlockPrecondType::SchurUpper}, + BlockTestParams{BlockSolverType::Direct, BlockPrecondType::SchurDiag}, + BlockTestParams{BlockSolverType::Direct, BlockPrecondType::SchurFull}, + + // Iterative solvers + BlockTestParams{BlockSolverType::Iterative, BlockPrecondType::Diagonal}, + BlockTestParams{BlockSolverType::Iterative, BlockPrecondType::TriLower}, + BlockTestParams{BlockSolverType::Iterative, BlockPrecondType::TriUpper}, + BlockTestParams{BlockSolverType::Iterative, BlockPrecondType::TriSym}, + BlockTestParams{BlockSolverType::Iterative, BlockPrecondType::SchurLower}, + BlockTestParams{BlockSolverType::Iterative, BlockPrecondType::SchurUpper}, + BlockTestParams{BlockSolverType::Iterative, BlockPrecondType::SchurDiag}, + BlockTestParams{BlockSolverType::Iterative, BlockPrecondType::SchurFull}, + + // BoomerAMG solvers + BlockTestParams{BlockSolverType::BoomerAMG, BlockPrecondType::Diagonal}, + BlockTestParams{BlockSolverType::BoomerAMG, BlockPrecondType::TriLower}, + BlockTestParams{BlockSolverType::BoomerAMG, BlockPrecondType::TriUpper}, + BlockTestParams{BlockSolverType::BoomerAMG, BlockPrecondType::TriSym}, + BlockTestParams{BlockSolverType::BoomerAMG, BlockPrecondType::SchurLower}, + BlockTestParams{BlockSolverType::BoomerAMG, BlockPrecondType::SchurUpper}, + BlockTestParams{BlockSolverType::BoomerAMG, BlockPrecondType::SchurDiag}, + BlockTestParams{BlockSolverType::BoomerAMG, BlockPrecondType::SchurFull}), + BlockParamNameGenerator); + +int main(int argc, char* argv[]) +{ + ::testing::InitGoogleTest(&argc, argv); + smith::ApplicationManager applicationManager(argc, argv); + return RUN_ALL_TESTS(); +} \ No newline at end of file diff --git a/src/smith/numerics/CMakeLists.txt b/src/smith/numerics/CMakeLists.txt index e9243a782..6df0f7eb1 100644 --- a/src/smith/numerics/CMakeLists.txt +++ b/src/smith/numerics/CMakeLists.txt @@ -14,6 +14,7 @@ set(numerics_headers petsc_solvers.hpp trust_region_solver.hpp dense_petsc.hpp + block_preconditioner.hpp ) set(numerics_sources @@ -21,6 +22,7 @@ set(numerics_sources trust_region_solver.cpp odes.cpp petsc_solvers.cpp + block_preconditioner.cpp ) set(numerics_depends smith_infrastructure smith_functional smith_boundary_conditions) diff --git a/src/smith/numerics/block_preconditioner.cpp b/src/smith/numerics/block_preconditioner.cpp new file mode 100644 index 000000000..ad4635a6c --- /dev/null +++ b/src/smith/numerics/block_preconditioner.cpp @@ -0,0 +1,326 @@ +#include "smith/numerics/block_preconditioner.hpp" +#include "mfem.hpp" +#include "axom/slic/core/SimpleLogger.hpp" + +namespace smith { + +BlockDiagonalPreconditioner::BlockDiagonalPreconditioner(mfem::Array& offsets, + std::vector> solvers) + : block_offsets_(offsets), + num_blocks_(offsets.Size() - 1), + block_jacobian_(nullptr), + solver_diag_(block_offsets_), + mfem_solvers_(std::move(solvers)) +{ + SLIC_ERROR_IF(mfem_solvers_.size() != static_cast(num_blocks_), + "Number of solvers must match number of blocks"); +} + +void BlockDiagonalPreconditioner::Mult(const mfem::Vector& in, mfem::Vector& out) const { solver_diag_.Mult(in, out); } + +void BlockDiagonalPreconditioner::SetOperator(const mfem::Operator& jacobian) +{ + height = jacobian.Height(); + width = jacobian.Width(); + // Cast the supplied jacobian to a block operator object + block_jacobian_ = dynamic_cast(&jacobian); + + // For each diagonal block A_ii, configure the corresponding solver + for (int i = 0; i < num_blocks_; i++) { + const mfem::Operator& A_ii = block_jacobian_->GetBlock(i, i); + + // Attach operator to solver + mfem_solvers_[static_cast(i)]->SetOperator(A_ii); + + // Place the solver into the diagonal block of solver_diag_ + solver_diag_.SetBlock(i, i, mfem_solvers_[static_cast(i)].get()); + } +} + +BlockDiagonalPreconditioner::~BlockDiagonalPreconditioner() {} + +BlockTriangularPreconditioner::BlockTriangularPreconditioner(mfem::Array& offsets, + std::vector> solvers, + BlockTriangularType type) + : block_offsets_(offsets), + num_blocks_(offsets.Size() - 1), + block_jacobian_(nullptr), + mfem_solvers_(std::move(solvers)), + type_(type) +{ + SLIC_ERROR_IF(mfem_solvers_.size() != static_cast(num_blocks_), + "Number of solvers must match number of blocks"); +} + +void BlockTriangularPreconditioner::LowerSweep(const mfem::Vector& in, mfem::Vector& out) const +{ + mfem::BlockVector b(const_cast(in), block_offsets_); + mfem::BlockVector x(out, block_offsets_); + + // Forward sweep: i = 0 .. num_blocks_ - 1 + for (int i = 0; i < num_blocks_; i++) { + mfem::Vector& bi = b.GetBlock(i); + mfem::Vector& xi = x.GetBlock(i); + + // rhs_i = b_i + mfem::Vector rhs_i(bi.Size()); + rhs_i = bi; + + // Subtract sum_{j < i} A_ij x_j + for (int j = 0; j < i; j++) { + if (block_jacobian_->IsZeroBlock(i, j)) { + continue; // no coupling + } + const mfem::Operator& A_ij = block_jacobian_->GetBlock(i, j); + + mfem::Vector tmp(rhs_i.Size()); + const mfem::Vector& xj = x.GetBlock(j); + + A_ij.Mult(xj, tmp); // tmp = A_ij x_j + rhs_i.Add(-1.0, tmp); // rhs_i -= A_ij x_j + } + + // Solve A_ii x_i = rhs_i with the i-th block solver + mfem_solvers_[static_cast(i)]->Mult(rhs_i, xi); + } +} + +void BlockTriangularPreconditioner::UpperSweep(const mfem::Vector& in, mfem::Vector& out) const +{ + mfem::BlockVector b(const_cast(in), block_offsets_); + mfem::BlockVector x(out, block_offsets_); + + // Backward sweep: i = num_blocks_ - 1 .. 0 + for (int i = num_blocks_ - 1; i >= 0; i--) { + mfem::Vector& bi = b.GetBlock(i); + mfem::Vector& xi = x.GetBlock(i); + + // rhs_i = b_i + mfem::Vector rhs_i(bi.Size()); + rhs_i = bi; + + // Subtract sum_{j > i} A_ij x_j + for (int j = i + 1; j < num_blocks_; j++) { + if (block_jacobian_->IsZeroBlock(i, j)) { + continue; // no coupling + } + const mfem::Operator& A_ij = block_jacobian_->GetBlock(i, j); + + mfem::Vector tmp(rhs_i.Size()); + const mfem::Vector& xj = x.GetBlock(j); + + A_ij.Mult(xj, tmp); // tmp = A_ij x_j + rhs_i.Add(-1.0, tmp); // rhs_i -= A_ij x_j + } + + // Solve A_ii x_i = rhs_i + mfem_solvers_[static_cast(i)]->Mult(rhs_i, xi); + } +} + +void BlockTriangularPreconditioner::Mult(const mfem::Vector& in, mfem::Vector& out) const +{ + switch (type_) { + case BlockTriangularType::Lower: + // x = P_lower^{-1} b + LowerSweep(in, out); + break; + + case BlockTriangularType::Upper: + // x = P_upper^{-1} b + UpperSweep(in, out); + break; + + case BlockTriangularType::Symmetric: { + // Symmetric: x = P_upper^{-1} D P_lower^{-1} b + // 1) tmp = P_lower^{-1} b + mfem::Vector tmp(out.Size()); + LowerSweep(in, tmp); + + // 2) tmp = D * tmp where D = diag(A_ii) + { + mfem::BlockVector tmp_block(tmp, block_offsets_); + + for (int i = 0; i < num_blocks_; i++) { + mfem::Vector& tmp_i = tmp_block.GetBlock(i); + mfem::Vector tmp_i_scaled(tmp_i.Size()); + + const mfem::Operator& A_ii = block_jacobian_->GetBlock(i, i); + A_ii.Mult(tmp_i, tmp_i_scaled); // tmp_i_scaled = A_ii * tmp_i + + tmp_i = tmp_i_scaled; // write back into block vector + } + } + + // 3) out = P_upper^{-1} tmp + UpperSweep(tmp, out); + break; + } + } +} + +void BlockTriangularPreconditioner::SetOperator(const mfem::Operator& jacobian) +{ + height = jacobian.Height(); + width = jacobian.Width(); + // Cast the supplied jacobian to a block operator object + block_jacobian_ = dynamic_cast(&jacobian); + + // Configure all diagonal solves + for (int i = 0; i < num_blocks_; i++) { + const mfem::Operator& A_ii = block_jacobian_->GetBlock(i, i); + mfem_solvers_[static_cast(i)]->SetOperator(A_ii); + } +} + +BlockTriangularPreconditioner::~BlockTriangularPreconditioner() {} + +BlockSchurPreconditioner::BlockSchurPreconditioner(mfem::Array& offsets, + std::vector> solvers, + BlockSchurType type) + : block_offsets_(offsets), + block_jacobian_(nullptr), + solver_diag_(block_offsets_), + mfem_solvers_(std::move(solvers)), + type_(type) +{ + SLIC_ERROR_IF(block_offsets_.Size() - 1 != 2, "This precondition is specifically for 2X2 block systems"); +} + +void BlockSchurPreconditioner::LowerBlock(const mfem::Vector& in, mfem::Vector& out) const +{ + // Interpret in, out as block vectors: in = [b1; b2], out = [x1; x2] + mfem::BlockVector b(const_cast(in), block_offsets_); + mfem::BlockVector x(out, block_offsets_); + + mfem::Vector& b1 = b.GetBlock(0); + mfem::Vector& b2 = b.GetBlock(1); + mfem::Vector& x1 = x.GetBlock(0); + mfem::Vector& x2 = x.GetBlock(1); + + // 1) Solve A11 x1 = b1 + mfem_solvers_[0]->Mult(b1, x1); + + // 2) Build x2 = b2 - A21 x1 + A_21_->Mult(x1, x2); // x2 = A21 x1 + x2.Neg(); // x2 = -A21 x1 + x2 += b2; // x2 = b2 - A21 x1 + + // 3) Reassign x1. + x1 = b1; +} + +void BlockSchurPreconditioner::UpperBlock(const mfem::Vector& in, mfem::Vector& out) const +{ + // Interpret in, out as block vectors: in = [b1; b2], out = [x1; x2] + mfem::BlockVector b(const_cast(in), block_offsets_); + mfem::BlockVector x(out, block_offsets_); + + mfem::Vector& b1 = b.GetBlock(0); + mfem::Vector& b2 = b.GetBlock(1); + mfem::Vector& x1 = x.GetBlock(0); + mfem::Vector& x2 = x.GetBlock(1); + + // 1) Build x1 = A12 b2 + mfem::Vector rhs1(b1.Size()); + A_12_->Mult(b2, rhs1); // rhs1 = A12 b2 + + // 2) Solve A11 x1 = rhs1 + mfem_solvers_[0]->Mult(rhs1, x1); + + // 3) Build b1 - A11^-1 A12 b2 + x1.Neg(); // x1 = -x1 + x1 += b1; // = b1 - A12 x2 + + // 4) Assign x2 + x2 = b2; +} + +void BlockSchurPreconditioner::Mult(const mfem::Vector& in, mfem::Vector& out) const +{ + switch (type_) { + case BlockSchurType::Diagonal: { + // x = [A11^-1, 0; 0, S^-1] b + solver_diag_.Mult(in, out); + break; + } + + case BlockSchurType::Lower: { + // x = [A11^-1, 0; 0, S^-1][I, 0; -A21 A11^-1, I] b + mfem::Vector tmp(out.Size()); + LowerBlock(in, tmp); + solver_diag_.Mult(tmp, out); + break; + } + + case BlockSchurType::Upper: { + // x = [I, -A11^-1 A12; 0, I][A11^-1, 0; 0, S^-1] b + mfem::Vector tmp(out.Size()); + solver_diag_.Mult(in, tmp); + UpperBlock(tmp, out); + break; + } + + case BlockSchurType::Full: { + // x = [I, -A11^-1 A12; 0, I][A11^-1, 0; 0, S^-1][I, 0; -A21 A11^-1, I] b + mfem::Vector tmp(out.Size()); + mfem::Vector tmp2(out.Size()); + LowerBlock(in, tmp); + solver_diag_.Mult(tmp, tmp2); + UpperBlock(tmp2, out); + break; + } + } +} + +void BlockSchurPreconditioner::SetOperator(const mfem::Operator& jacobian) +{ + height = jacobian.Height(); + width = jacobian.Width(); + block_jacobian_ = dynamic_cast(&jacobian); + MFEM_VERIFY(block_jacobian_, "Jacobian must be a BlockOperator"); + + 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."); + + if (type_ == BlockSchurType::Lower || type_ == BlockSchurType::Full) { + A_21_ = A21; + } + if (type_ == BlockSchurType::Upper || type_ == BlockSchurType::Full) { + A_12_ = A12; + } + // Diagonal preconditioner for block (0,0) + mfem_solvers_[0]->SetOperator(*A11); + + // Extract the diagonal of A11 (no inversion!) + mfem::HypreParVector* Md = new mfem::HypreParVector(A11->GetComm(), A11->GetGlobalNumRows(), A11->GetRowStarts()); + A11->GetDiag(*Md); + + // Scale ROWS of A12 by Md^{-1} + mfem::HypreParMatrix* A12_scaled = new mfem::HypreParMatrix(*A12); + A12_scaled->InvScaleRows(*Md); + delete Md; + + // Now compute A21 * (diag(A11)^{-1} * A12) + mfem::HypreParMatrix* A21DinvA12 = mfem::ParMult(A21, A12_scaled); + delete A12_scaled; + + // S_approx = A22 - A21 * diag(A11)^{-1} * A12 + S_approx_.reset(mfem::Add(1.0, *A22, -1.0, *A21DinvA12)); + delete A21DinvA12; + + // Set the Schur complement preconditioner for block (1,1) + mfem_solvers_[1]->SetOperator(*S_approx_); + + // Set up block diagonal operator + solver_diag_.SetBlock(0, 0, mfem_solvers_[0].get()); + solver_diag_.SetBlock(1, 1, mfem_solvers_[1].get()); +} + +BlockSchurPreconditioner::~BlockSchurPreconditioner() {} +} // namespace smith diff --git a/src/smith/numerics/block_preconditioner.hpp b/src/smith/numerics/block_preconditioner.hpp new file mode 100644 index 000000000..b0d0bb09a --- /dev/null +++ b/src/smith/numerics/block_preconditioner.hpp @@ -0,0 +1,233 @@ +#pragma once + +#include +#include +#include "mfem.hpp" + +namespace smith { + +/** + * @class BlockDiagonalPreconditioner + * @brief Simple block diagonal preconditioner for block systems. + * + * Stores one solver per block and applies them to the diagonal blocks of a + * block Jacobian. + * + * Call SetOperator() with an mfem::BlockOperator, then use Mult() to apply the + * preconditioner. + */ +class BlockDiagonalPreconditioner : public mfem::Solver { + public: + /** + * @brief Construct a new N by N block diagonal preconditioner. + * + * @param offsets Offsets describing the block layout. + * @param solvers One solver per block (size must match number of blocks). + */ + BlockDiagonalPreconditioner(mfem::Array& offsets, std::vector> solvers); + + /** + * @brief The action of the precondition on the block vector (b_1, ..., b_n) + * + * @param in The block input vector (b_1, ..., b_n) + * @param out The block output vector P^-1(b_1, ..., b_n) + */ + virtual void Mult(const mfem::Vector& in, mfem::Vector& out) const; + + /** + * @brief Set the preconditioner to use the supplied linearized block Jacobian + * + * @param jacobian The supplied linearized Jacobian. Note that it is always a block operator + */ + virtual void SetOperator(const mfem::Operator& jacobian); + + virtual ~BlockDiagonalPreconditioner(); + + private: + // Offsets for extracting block vector segments + mfem::Array& block_offsets_; + + // Number of blocks + const int num_blocks_; + + // Jacobian view for block access + const mfem::BlockOperator* block_jacobian_; + + // The diagonal part of the preconditioner containing BoomerAMG applications + mfem::BlockOperator solver_diag_; + + // mfem solvers for each block + mutable std::vector> mfem_solvers_; +}; + +/** + * @enum BlockTriangularType + * @brief Selects the block triangular sweep used by BlockTriangularPreconditioner. + */ +enum class BlockTriangularType +{ + Lower, /**< Forward (lower triangular) sweep. */ + Upper, /**< Backward (upper triangular) sweep. */ + Symmetric /**< Apply a symmetric combination of lower and upper sweeps. */ +}; + +/** + * @class BlockTriangularPreconditioner + * @brief Simple block triangular preconditioner for block systems. + * + * Stores one solver per diagonal block and applies a block sweep using the + * supplied block Jacobian. + * + * Call SetOperator() with an mfem::BlockOperator, then use Mult() to apply the + * preconditioner. + */ +class BlockTriangularPreconditioner : public mfem::Solver { + public: + /** + * @brief Construct a new nxn block triangular preconditioner. + * + * @param offsets Offsets describing the block layout. + * @param solvers One solver per diagonal block (size must match number of blocks). + * @param type Sweep type (lower, upper, or symmetric). + */ + BlockTriangularPreconditioner(mfem::Array& offsets, std::vector> solvers, + BlockTriangularType type = BlockTriangularType::Lower); + + /** + * @brief The action of the precondition on the block vector (b_1, ..., b_n) + * + * @param in The block input vector (b_1, ..., b_n) + * @param out The block output vector P^-1(b_1, ..., b_n) + */ + virtual void Mult(const mfem::Vector& in, mfem::Vector& out) const; + + /** + * @brief Set the preconditioner to use the supplied linearized block Jacobian + * + * @param jacobian The supplied linearized Jacobian. Note that it is always a block operator + */ + virtual void SetOperator(const mfem::Operator& jacobian); + + virtual ~BlockTriangularPreconditioner(); + + private: + // Offsets for extracting block vector segments + mfem::Array& block_offsets_; + + // Number of blocks + const int num_blocks_; + + // Jacobian view for block access + const mfem::BlockOperator* block_jacobian_; + + // mfem solvers for each block + mutable std::vector> mfem_solvers_; + + // Block Triangular type + BlockTriangularType type_; + + /** + * @brief The action of the lower sweep on the block vector (b_1, ..., b_n) + * + * @param in The block input vector (b_1, ..., b_n) + * @param out The block output vector P_lower^-1(b_1, ..., b_n) + */ + void LowerSweep(const mfem::Vector& in, mfem::Vector& out) const; + + /** + * @brief The action of the upper sweep on the block vector (b_1, ..., b_n) + * + * @param in The block input vector (b_1, ..., b_n) + * @param out The block output vector P_upper^-1(b_1, ..., b_n) + */ + void UpperSweep(const mfem::Vector& in, mfem::Vector& out) const; +}; + +/** + * @enum BlockSchurType + * @brief Selects the block Schur preconditioner variant. + */ +enum class BlockSchurType +{ + Diagonal, /**< Block diagonal: apply $ A_{11}^{-1} $ and $ S^{-1} $ only. */ + Lower, /**< Lower factor form. */ + Upper, /**< Upper factor form. */ + Full /**< Full factor form (lower, diagonal, upper). */ +}; + +/** + * @class BlockSchurPreconditioner + * @brief Simple 2x2 block Schur complement preconditioner for block systems. + * + * Uses two solvers, one for $ A_{11} $ and one for an approximate Schur complement $ S $. + * Call SetOperator() with an mfem::BlockOperator, then use Mult() to apply the + * selected Schur preconditioner type. + */ +class BlockSchurPreconditioner : public mfem::Solver { + public: + /** + * @brief Construct a new 2x2 block Schur complement preconditioner. + * + * @param offsets Offsets describing the 2-block layout. + * @param solvers Two solvers, for $ A_{11} $ and the Schur complement approximation. + * @param type Preconditioner variant (diagonal, lower, upper, or full). + */ + BlockSchurPreconditioner(mfem::Array& offsets, std::vector> solvers, + BlockSchurType type = BlockSchurType::Diagonal); + + /** + * @brief The action of the precondition on the block vector (b_1, b_2) + * + * @param in The block input vector (b_1, b_2) + * @param out The block output vector P^-1(b_1, b_2) + */ + virtual void Mult(const mfem::Vector& in, mfem::Vector& out) const; + + /** + * @brief Set the preconditioner to use the supplied linearized block Jacobian. + * + * The Schur complement approximation is given by S_approx = A22 - A21 * diag(A11)^{-1} * A12 + * + * @param jacobian The supplied linearized Jacobian. Note that it is always a block operator + */ + virtual void SetOperator(const mfem::Operator& jacobian); + + virtual ~BlockSchurPreconditioner(); + + private: + // Offsets for extracting block vector segments + mfem::Array& block_offsets_; + + // Jacobian view for block access + const mfem::BlockOperator* block_jacobian_; + + // The diagonal part of the preconditioner containing BoomerAMG applications + mfem::BlockOperator solver_diag_; + + // mfem solvers for each block + mutable std::vector> mfem_solvers_; + + // Views of the linearized Jacobian blocks + const mfem::Operator *A_12_, *A_21_; + + mutable std::unique_ptr S_approx_; + + BlockSchurType type_; + + /** + * @brief The action of the lower sweep on the block vector (b_1, b_2) + * + * @param in The block input vector (b_1, b_2) + * @param out The block output vector [I, 0; -A21 A11^-1, I] (b_1, b_2) + */ + void LowerBlock(const mfem::Vector& in, mfem::Vector& out) const; + + /** + * @brief The action of the upper block on the block vector (b_1, b_2) + * + * @param in The block input vector (b_1, b_2) + * @param out The block output vector [I - A11^-1 A12; 0, I](b_1, b_2) + */ + void UpperBlock(const mfem::Vector& in, mfem::Vector& out) const; +}; +} // namespace smith diff --git a/src/smith/numerics/equation_solver.cpp b/src/smith/numerics/equation_solver.cpp index fcdd7af12..c53878466 100644 --- a/src/smith/numerics/equation_solver.cpp +++ b/src/smith/numerics/equation_solver.cpp @@ -1013,7 +1013,8 @@ void StrumpackSolver::SetOperator(const mfem::Operator& op) strumpack_mat_ = std::make_unique(*matrix); } - + height = op.Height(); + width = op.Width(); strumpack_solver_.SetOperator(*strumpack_mat_); } diff --git a/src/smith/numerics/tests/CMakeLists.txt b/src/smith/numerics/tests/CMakeLists.txt index 3df818dd9..afce05689 100644 --- a/src/smith/numerics/tests/CMakeLists.txt +++ b/src/smith/numerics/tests/CMakeLists.txt @@ -10,6 +10,8 @@ set(numerics_serial_test_sources equationsolver.cpp operator.cpp odes.cpp + test_block_preconditioner.cpp + test_block_preconditioner_backend.cpp ) smith_add_tests( SOURCES ${numerics_serial_test_sources} diff --git a/src/smith/numerics/tests/test_block_preconditioner.cpp b/src/smith/numerics/tests/test_block_preconditioner.cpp new file mode 100644 index 000000000..8f7e0a176 --- /dev/null +++ b/src/smith/numerics/tests/test_block_preconditioner.cpp @@ -0,0 +1,452 @@ +#include +#include + +#include "mfem.hpp" +#include "smith/numerics/block_preconditioner.hpp" +#include "smith/infrastructure/application_manager.hpp" + +#include "axom/slic.hpp" + +bool abort_called = false; + +void testAbortHandler() +{ + abort_called = true; // record abort instead of exiting +} + +using namespace mfem; + +/* ============================================================ + Helper utilities + ============================================================ */ + +// Build c * I +std::unique_ptr makeScaledIdentity(int n, double c) +{ + auto A = std::make_unique(n); + for (int i = 0; i < n; i++) { + A->Add(i, i, c); + } + A->Finalize(); + return A; +} + +// Simple SPD matrix: tridiagonal +std::unique_ptr makeSPDMatrix(int n) +{ + auto A = std::make_unique(n); + for (int i = 0; i < n; i++) { + A->Add(i, i, 2.0); + if (i > 0) A->Add(i, i - 1, -1.0); + if (i < n - 1) A->Add(i, i + 1, -1.0); + } + A->Finalize(); + return A; +} + +class IdentitySolver : public mfem::Solver { + public: + IdentitySolver() = default; + + 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; } +}; + +// Exact diagonal inverse solver +class ExactDiagonalSolver : public mfem::Solver { + public: + ExactDiagonalSolver() = default; + + void SetOperator(const mfem::Operator& op) override + { + const auto* A = dynamic_cast(&op); + MFEM_VERIFY(A, "ExactDiagonalSolver requires SparseMatrix"); + + A_ = A; + height = A_->Height(); + width = A_->Width(); + + diag_.SetSize(height); + A_->GetDiag(diag_); + } + + void Mult(const mfem::Vector& x, mfem::Vector& y) const override + { + MFEM_ASSERT(A_, "Operator not set"); + + y.SetSize(x.Size()); + for (int i = 0; i < x.Size(); i++) { + y[i] = x[i] / diag_[i]; + } + } + + private: + const mfem::SparseMatrix* A_ = nullptr; + mfem::Vector diag_; +}; + +std::vector> makeExactDiagonalSolvers() +{ + std::vector> solvers; + solvers.push_back(std::make_unique()); + solvers.push_back(std::make_unique()); + return solvers; +} + +// Tridiagonal rectangular sparse matrix with constant main diagonal 2.0, off-diagonals -1.0 +std::unique_ptr makeRectTridiagonal(int rows, int cols) +{ + auto A = std::make_unique(rows, cols); + + for (int i = 0; i < rows; i++) { + if (i < cols) A->Add(i, i, 2.0); // main diagonal + if (i > 0 && i - 1 < cols) A->Add(i, i - 1, -1.0); // lower diagonal + if (i + 1 < cols) A->Add(i, i + 1, -1.0); // upper diagonal + } + + A->Finalize(); + return A; +} + +/* ============================================================ + Tests + ============================================================ */ + +// Makes sure an error is thrown when length of solvers does not match the +// number of blocks +TEST(BlockDiagonal, ThrowsOnWrongNumberOfSolvers) +{ + // Replace abort handler + axom::slic::setAbortFunction(testAbortHandler); + + abort_called = false; + + Array offsets({0, 2, 4}); + std::vector> solvers; + solvers.push_back(std::make_unique()); + + std::cout << abort_called << std::endl; + smith::BlockDiagonalPreconditioner P(offsets, std::move(solvers)); + std::cout << abort_called << std::endl; + + EXPECT_TRUE(abort_called); +} + +// If the solver for each block is identity, the block solver is identity +TEST(BlockTriangular, IdentityActsAsIdentity) +{ + Array offsets({0, 2, 5}); + + auto A11 = makeScaledIdentity(2, 1.0); + auto A22 = makeScaledIdentity(3, 1.0); + + BlockOperator A(offsets); + A.SetBlock(0, 0, A11.get()); + A.SetBlock(1, 1, A22.get()); + + std::vector> solvers; + solvers.push_back(std::make_unique()); + solvers.push_back(std::make_unique()); + + smith::BlockTriangularPreconditioner P(offsets, std::move(solvers), smith::BlockTriangularType::Symmetric); + P.SetOperator(A); + + Vector x(5), y(5); + x.Randomize(); + + P.Mult(x, y); + + mfem::Vector diff(x); + diff -= y; + + EXPECT_NEAR(diff.Norml2(), 0.0, 1e-14); +} + +// BlockDiagonalPreconditioner ignores off-diagonal blocks +TEST(BlockDiagonal, IgnoresOffDiagonalBlocks) +{ + Array offsets({0, 2, 4}); + + auto A11 = makeScaledIdentity(2, 2.0); + auto A22 = makeScaledIdentity(2, 4.0); + auto A12 = makeScaledIdentity(2, 1.0); + auto A21 = makeScaledIdentity(2, 1.0); + + BlockOperator A(offsets); + A.SetBlock(0, 0, A11.get()); + A.SetBlock(0, 1, A12.get()); + A.SetBlock(1, 0, A21.get()); + A.SetBlock(1, 1, A22.get()); + + auto solvers = makeExactDiagonalSolvers(); + smith::BlockDiagonalPreconditioner P(offsets, std::move(solvers)); + P.SetOperator(A); + + Vector b(4), x(4); + 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] / 4.0, 1e-12); + EXPECT_NEAR(x[3], b[3] / 4.0, 1e-12); +} + +// LowerTriangularPreconditioner is exact with exact solvers for a lower +// triangular matrix +TEST(BlockTriangular, LowerTriangularExactSolve) +{ + Array offsets({0, 2, 4}); + + auto A11 = makeScaledIdentity(2, 2.0); + auto A22 = makeScaledIdentity(2, 3.0); + auto A21 = makeScaledIdentity(2, 1.0); + + BlockOperator A(offsets); + A.SetBlock(0, 0, A11.get()); + A.SetBlock(1, 0, A21.get()); + A.SetBlock(1, 1, A22.get()); + + auto solvers = makeExactDiagonalSolvers(); + smith::BlockTriangularPreconditioner P(offsets, std::move(solvers), smith::BlockTriangularType::Lower); + P.SetOperator(A); + + Vector b(4), x(4), Ax(4); + b.Randomize(); + + P.Mult(b, x); + A.Mult(x, Ax); + + mfem::Vector res(Ax); + res -= b; + EXPECT_NEAR(res.Norml2(), 0.0, 1e-12); +} + +// Symmetric BlockTriangularPreconditioner is actually symmetric +TEST(BlockTriangular, SymmetricGSIsSelfAdjoint) +{ + Array offsets({0, 2, 4}); + + auto A11 = makeSPDMatrix(2); + auto A22 = makeSPDMatrix(2); + auto A12 = makeScaledIdentity(2, 0.5); + auto A21 = makeScaledIdentity(2, 0.5); + + BlockOperator A(offsets); + A.SetBlock(0, 0, A11.get()); + A.SetBlock(0, 1, A12.get()); + A.SetBlock(1, 0, A21.get()); + A.SetBlock(1, 1, A22.get()); + + auto solvers = makeExactDiagonalSolvers(); + smith::BlockTriangularPreconditioner P(offsets, std::move(solvers), smith::BlockTriangularType::Symmetric); + P.SetOperator(A); + + Vector x(4), y(4), Px(4), Py(4); + x.Randomize(); + y.Randomize(); + + P.Mult(x, Px); + P.Mult(y, Py); + + double lhs = InnerProduct(Px, y); + double rhs = InnerProduct(x, Py); + + EXPECT_NEAR(lhs, rhs, 1e-12); +} + +// 3x3 block triangular system, exact solve +TEST(BlockTriangular, LowerTriangularExactSolve_3Blocks) +{ + Array offsets({0, 2, 5, 7}); // block sizes: 2,3,2 + + auto sz = [&](int i) { return offsets[i + 1] - offsets[i]; }; + + // Diagonal blocks (square, exact solves) + auto A11 = makeScaledIdentity(sz(0), 2.0); + auto A22 = makeScaledIdentity(sz(1), 3.0); + auto A33 = makeScaledIdentity(sz(2), 4.0); + + // Off-diagonal blocks (rectangular, nonzero) + auto A21 = makeRectTridiagonal(sz(1), sz(0)); // 3x2 + auto A31 = makeRectTridiagonal(sz(2), sz(0)); // 2x2 + auto A32 = makeRectTridiagonal(sz(2), sz(1)); // 2x3 + + mfem::BlockOperator A(offsets); + + // Set blocks + A.SetBlock(0, 0, A11.get()); + A.SetBlock(1, 0, A21.get()); + A.SetBlock(1, 1, A22.get()); + A.SetBlock(2, 0, A31.get()); + A.SetBlock(2, 1, A32.get()); + A.SetBlock(2, 2, A33.get()); + + // Exact diagonal solvers for each block + std::vector> solvers; + solvers.push_back(std::make_unique()); + solvers.push_back(std::make_unique()); + solvers.push_back(std::make_unique()); + + smith::BlockTriangularPreconditioner P(offsets, std::move(solvers), smith::BlockTriangularType::Lower); + + P.SetOperator(A); + + mfem::Vector b(sz(0) + sz(1) + sz(2)), x(sz(0) + sz(1) + sz(2)), Ax(sz(0) + sz(1) + sz(2)); + b.Randomize(); + + P.Mult(b, x); + A.Mult(x, Ax); + + mfem::Vector r(Ax); + r -= b; + EXPECT_NEAR(r.Norml2(), 0.0, 1e-12); +} + +// UpperTriangularPreconditioner is exact with exact solvers for an upper +// triangular matrix +TEST(BlockTriangular, UpperTriangularExactSolve) +{ + Array offsets({0, 2, 4}); + + auto A11 = makeScaledIdentity(2, 2.0); + auto A22 = makeScaledIdentity(2, 3.0); + auto A12 = makeScaledIdentity(2, 1.0); + + BlockOperator A(offsets); + A.SetBlock(0, 0, A11.get()); + A.SetBlock(0, 1, A12.get()); + A.SetBlock(1, 1, A22.get()); + + auto solvers = makeExactDiagonalSolvers(); + smith::BlockTriangularPreconditioner P(offsets, std::move(solvers), smith::BlockTriangularType::Upper); + + P.SetOperator(A); + + Vector b(4), x(4), Ax(4); + b.Randomize(); + + P.Mult(b, x); + A.Mult(x, Ax); + + mfem::Vector r(Ax); + r -= b; + EXPECT_NEAR(r.Norml2(), 0.0, 1e-12); +} + +// Ensures that P^-1 0 = 0 +TEST(BlockTriangular, ZeroInputGivesZeroOutput) +{ + Array offsets({0, 2, 5}); + + auto A11 = makeScaledIdentity(2, 2.0); + auto A22 = makeScaledIdentity(3, 3.0); + + BlockOperator A(offsets); + A.SetBlock(0, 0, A11.get()); + A.SetBlock(1, 1, A22.get()); + + std::vector> solvers; + solvers.push_back(std::make_unique()); + solvers.push_back(std::make_unique()); + + smith::BlockTriangularPreconditioner P(offsets, std::move(solvers), smith::BlockTriangularType::Symmetric); + + P.SetOperator(A); + + Vector x(5), y(5); + x = 0.0; + + P.Mult(x, y); + EXPECT_NEAR(y.Norml2(), 0.0, 0.0); +} + +// Block Triangular still works if some off-diagonal blocks are zero/missing +TEST(BlockTriangular, HandlesMissingOffDiagonalBlocks) +{ + Array offsets({0, 2, 4}); + + auto A11 = makeScaledIdentity(2, 2.0); + auto A22 = makeScaledIdentity(2, 3.0); + + BlockOperator A(offsets); + A.SetBlock(0, 0, A11.get()); + A.SetBlock(1, 1, A22.get()); + + auto solvers = makeExactDiagonalSolvers(); + smith::BlockTriangularPreconditioner P(offsets, std::move(solvers), smith::BlockTriangularType::Symmetric); + + EXPECT_NO_THROW(P.SetOperator(A)); +} + +TEST(BlockDiagonal, WorksForSingleBlock) +{ + // Single block (non-block system) + auto A0 = makeScaledIdentity(3, 2.0); // 3x3 diagonal + + // Wrap as a block operator with offsets 0,3 + Array offsets({0, 3}); + BlockOperator A(offsets); + A.SetBlock(0, 0, A0.get()); + + // One solver for the single block + std::vector> solvers; + solvers.push_back(std::make_unique()); + + smith::BlockDiagonalPreconditioner P(offsets, std::move(solvers)); + EXPECT_NO_THROW(P.SetOperator(A)); + + // Apply preconditioner + Vector b(3), x(3), Ax(3); + b.Randomize(); + + P.Mult(b, x); + A.Mult(x, Ax); + + // Check that result is exact + Vector r(Ax); + r -= b; + EXPECT_NEAR(r.Norml2(), 0.0, 1e-12); +} + +TEST(BlockTriangular, WorksForSingleBlockAllTypes) +{ + auto A0 = makeScaledIdentity(3, 2.0); + Array offsets({0, 3}); + BlockOperator A(offsets); + A.SetBlock(0, 0, A0.get()); + + for (auto type : + {smith::BlockTriangularType::Lower, smith::BlockTriangularType::Upper, smith::BlockTriangularType::Symmetric}) { + std::vector> solvers; + solvers.push_back(std::make_unique()); + + smith::BlockTriangularPreconditioner P(offsets, std::move(solvers), type); + EXPECT_NO_THROW(P.SetOperator(A)); + + Vector b(3), x(3), Ax(3); + b.Randomize(); + P.Mult(b, x); + A.Mult(x, Ax); + + Vector r(Ax); + r -= b; + EXPECT_NEAR(r.Norml2(), 0.0, 1e-12); + } +} + +/* ============================================================ + main + ============================================================ */ + +int main(int argc, char* argv[]) +{ + ::testing::InitGoogleTest(&argc, argv); + smith::ApplicationManager applicationManager(argc, argv); + return RUN_ALL_TESTS(); +} diff --git a/src/smith/numerics/tests/test_block_preconditioner_backend.cpp b/src/smith/numerics/tests/test_block_preconditioner_backend.cpp new file mode 100644 index 000000000..a8182368e --- /dev/null +++ b/src/smith/numerics/tests/test_block_preconditioner_backend.cpp @@ -0,0 +1,192 @@ +#include +#include +#include "mfem.hpp" +#include "smith/numerics/block_preconditioner.hpp" +#include "smith/numerics/equation_solver.hpp" +#include "smith/infrastructure/application_manager.hpp" + +using namespace mfem; +using namespace smith; + +// ----------- Parameter Structs ------------ + +struct BlockPrecTestParams { + enum class BlockPattern + { + Diagonal2x2, + Lower2x2 + }; + enum class PrecKind + { + Diagonal, + Triangular + }; + enum class SolverBackend + { + HypreBoomerAMG, + Strumpack + }; + + BlockPattern pattern; + PrecKind prec_kind; + smith::BlockTriangularType tri_type; // Only for triangular + std::string name; + double rel_tol; + SolverBackend backend; +}; + +// ----------- Test Fixture ------------ + +class BlockPreconditionerParamTest : public ::testing::TestWithParam { + protected: + MPI_Comm comm_; + int rank_; + void SetUp() override + { + comm_ = MPI_COMM_WORLD; + MPI_Comm_rank(comm_, &rank_); + } +}; + +// ----------- Test Body ------------ + +TEST_P(BlockPreconditionerParamTest, SolvesBlockSystemApproximately) +{ + const auto& params = GetParam(); + + // Build FE problem + const int dim = 2, ne = 2, order = 1; + Mesh mesh = Mesh::MakeCartesian2D(ne, ne, Element::QUADRILATERAL, 1, 1.0, 1.0); + ParMesh pmesh(comm_, mesh); + mesh.Clear(); + H1_FECollection fec(order, dim); + ParFiniteElementSpace fes(&pmesh, &fec); + + Array ess_bdr(pmesh.bdr_attributes.Max()); + ess_bdr[0] = 1; + Array ess_tdof_list; + fes.GetEssentialTrueDofs(ess_bdr, ess_tdof_list); + + ConstantCoefficient one(1.0); + ParBilinearForm a(&fes); + a.AddDomainIntegrator(new DiffusionIntegrator(one)); + a.Assemble(); + a.Finalize(); + + ParLinearForm b_(&fes); + b_ = 0.0; + ParGridFunction x_(&fes); + x_ = 0.0; + + OperatorPtr A; + Vector X, B; + a.FormLinearSystem(ess_tdof_list, x_, b_, A, X, B); + + // Block setup + HypreParMatrix* A_hypre = A.As(); + ASSERT_TRUE(A_hypre != nullptr); + + int N = A_hypre->NumRows(); + Array block_offsets(3); + block_offsets[0] = 0; + block_offsets[1] = N; + block_offsets[2] = 2 * N; + + BlockOperator J(block_offsets); + + J.SetBlock(0, 0, A_hypre); + auto A11_copy = std::make_unique(*A_hypre); + J.SetBlock(1, 1, A11_copy.get()); + + HypreParMatrix* C = nullptr; + if (params.pattern == BlockPrecTestParams::BlockPattern::Lower2x2) { + // Build mass matrix for off-diagonal + ParBilinearForm c(&fes); + c.AddDomainIntegrator(new MassIntegrator(one)); + c.Assemble(); + c.Finalize(); + C = c.ParallelAssemble(); + J.SetBlock(1, 0, C); + } + + // Build solver array + std::vector> block_solvers; + block_solvers.reserve(2); + for (int i = 0; i < 2; ++i) { + if (params.backend == BlockPrecTestParams::SolverBackend::HypreBoomerAMG) { + auto solver = std::make_unique(); + solver->SetPrintLevel(0); + block_solvers.push_back(std::move(solver)); + } else { +#ifdef MFEM_USE_STRUMPACK + block_solvers.push_back(std::make_unique(0, comm_)); +#else + MFEM_ABORT("Strumpack backend requested but MFEM_USE_STRUMPACK not enabled."); +#endif + } + } + + // Instantiate preconditioner + std::unique_ptr P; + if (params.prec_kind == BlockPrecTestParams::PrecKind::Diagonal) { + P = std::make_unique(block_offsets, std::move(block_solvers)); + } else { + P = std::make_unique(block_offsets, std::move(block_solvers), + params.tri_type); + } + P->SetOperator(J); + + BlockVector r(block_offsets), x(block_offsets), b(block_offsets); + b.Randomize(1); + + P->Mult(b, x); + J.Mult(x, r); + r -= b; + double resid_err = r.Norml2(); + double rel_err = resid_err / b.Norml2(); + + if (rank_ == 0) { + std::cout << "Test " << params.name << ", rel_err = " << rel_err << std::endl; + } + ASSERT_LT(rel_err, params.rel_tol); + + if (C) delete C; +} + +// ----------- Instantiate Test Cases ------------ + +INSTANTIATE_TEST_SUITE_P( + BlockPreconditionerTests, BlockPreconditionerParamTest, + ::testing::Values( + // BlockDiagonalPreconditioner + HypreBoomerAMG on [A 0; 0 A] + BlockPrecTestParams{BlockPrecTestParams::BlockPattern::Diagonal2x2, BlockPrecTestParams::PrecKind::Diagonal, + smith::BlockTriangularType::Lower, // unused for diagonal + "Diag_HypreBoomerAMG", 1e-1, BlockPrecTestParams::SolverBackend::HypreBoomerAMG}, + // BlockTriangularPreconditioner + HypreBoomerAMG on [A 0; C A] + BlockPrecTestParams{BlockPrecTestParams::BlockPattern::Lower2x2, BlockPrecTestParams::PrecKind::Triangular, + smith::BlockTriangularType::Lower, "TriLower_HypreBoomerAMG", 1e-1, + BlockPrecTestParams::SolverBackend::HypreBoomerAMG}), + [](const ::testing::TestParamInfo& param_info) { return param_info.param.name; }); + +#ifdef MFEM_USE_STRUMPACK +INSTANTIATE_TEST_SUITE_P( + BlockPreconditionerStrumpackTests, BlockPreconditionerParamTest, + ::testing::Values( + // BlockDiagonalPreconditioner + Strumpack on [A 0; 0 A] + BlockPrecTestParams{BlockPrecTestParams::BlockPattern::Diagonal2x2, BlockPrecTestParams::PrecKind::Diagonal, + smith::BlockTriangularType::Lower, "Diag_Strumpack", 1e-10, + BlockPrecTestParams::SolverBackend::Strumpack}, + // BlockTriangularPreconditioner + Strumpack on [A 0; C A] + BlockPrecTestParams{BlockPrecTestParams::BlockPattern::Lower2x2, BlockPrecTestParams::PrecKind::Triangular, + smith::BlockTriangularType::Lower, "TriLower_Strumpack", 1e-10, + BlockPrecTestParams::SolverBackend::Strumpack}), + [](const ::testing::TestParamInfo& param_info) { return param_info.param.name; }); +#endif +// ----------- Google Test Main ------------ + +int main(int argc, char** argv) +{ + ::testing::InitGoogleTest(&argc, argv); + smith::ApplicationManager applicationManager(argc, argv); + return RUN_ALL_TESTS(); +} \ No newline at end of file