diff --git a/src/smith/differentiable_numerics/tests/CMakeLists.txt b/src/smith/differentiable_numerics/tests/CMakeLists.txt index dc5925a793..d48fc25f76 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 ) smith_add_tests( SOURCES ${differentiable_numerics_test_source} 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 0000000000..8fe1a26551 --- /dev/null +++ b/src/smith/differentiable_numerics/tests/test_porous_heat_sink.cpp @@ -0,0 +1,259 @@ +#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/materials/thermal_material.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" + +using namespace smith; + +using ShapeDispSpace = H1<1, 2>; +using Space = H1<1>; +using GammaSpace = L2<1>; + +struct MeshFixture : public testing::Test { + double length = 1.0; + double width = 1.0; + int num_elements_x = 64; + int num_elements_y = 64; + double elem_size = length / num_elements_x; + + void SetUp() + { + 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); + } // Construct the appropriate dimension mesh and give it to the data store + + axom::sidre::DataStore datastore; + std::shared_ptr mesh; +}; + +struct HeatSinkOptions { + /// \f$ \kappa_0 \f$ : reference fluid conductivity, \f$ W/m/K \f$ + double kappa_0 = 0.5; + + /// \f$ \sigma_0 \f$ :reference solid conductivity, \f$ W/m/K \f$ + double sigma_0 = 5.0; + + /// \f$ \eta \f$ : bruggeman correlation exponent + double eta = 1.5; + + /// \f$ \epsilon_m \f$ : liquid porosity + double epsilon_m = 1.0; + + /// \f$ \epsilon_n \f$ : heat sink porosity + double epsilon_n = 0.5; + + /// \f$ a_m \f$ : liquid specific surface area, \f$ m^2/m^3 \f$ + double a_m = 0.0; + + /// \f$ a_n \f$ : heat sink specific surface area, \f$ m^2/m^3 \f$ + double a_n = 5e+6; + + /// \f$ h \f$ : heat transfer coefficient \f$ W/m^2/K \f$ \f$ + double h = 0.01; + + /// \f$ q_{app} \f$ : applied heat flux, \f$ W/m^2 \f$ \f$ + double q_app = 0.0; + + /// \f$ T_{app} \f$ : applied temperature, \f$ K \f$ + double T_app = 0.0; + + /// \f$ f_{mb} \f$ : modified Bruggeman correlation scaling coefficient + double f_mb = 1.0; +}; + +TEST_F(MeshFixture,B) +{ + + std::string physics_name = "heatsink"; + auto graph = std::make_shared(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}; + // lambda helper functions for the underlying PDEs. Note that these must be generic lambdas (i.e. using auto) + // as this function will be called both with doubles (for residual evaluation) or dual numbers (auto diff-enabled + // tangents). + 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); + }; + + // Gamma field function for Full-cell heatsink + 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_) { + // Get the value and the gradient from the input tuple + auto [T_1_val, dT_1_dX] = T_1; + auto [T_2_val, dT_2_dX] = T_2; + auto [gamma_val, dgamma_dX] = gamma_; + + // The first element is the "source" term which acts on the scalar part of the test function. + // The second element is the "flux" term which acts on the gradient part of the test function. + 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_) { + // Get the value and the gradient from the input tuple + auto [T_1_val, dT_1_dX] = T_1; + auto [T_2_val, dT_2_dX] = T_2; + auto [gamma_val, dgamma_dX] = gamma_; + + // The first element is the "source" term which acts on the scalar part of the test function. + // The second element is the "flux" term which acts on the gradient part of the test function. + 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()); + + // Apply Dirichlet BC T_2=0 at bulk + auto zero_bcs = std::make_shared([](const mfem::Vector&) { return 0.0; }); + T2_bc_manager->addEssential(std::set{1}, zero_bcs, space(T2), 0); + + // Apply Neumann BC DT_1.n = q_app + 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}; + // Create an offset representing a block vector view of the combined solution vector + 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; + // BoomerAMG solvers for the blocks as before + // solvers.push_back(std::make_unique()); + // solvers.push_back(std::make_unique()); + // LU direct solvers + // 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()); + // Iterative solver for blocks + smith::LinearSolverOptions iter_solver_options = {.linear_solver = smith::LinearSolver::GMRES, + .preconditioner = smith::Preconditioner::HypreJacobi, + .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)); + // std::unique_ptr diff_precond = std::make_unique(block_offsets_, std::move(solvers)); + std::unique_ptr diff_precond = std::make_unique(block_offsets_, std::move(solvers), smith::BlockTriangularType::Symmetric); + std::unique_ptr linear_solver = std::make_unique(mesh->getComm()); + mfem::GMRESSolver* iter_lin_solver = dynamic_cast(linear_solver.get()); + + // Set up linear solver with custom preconditioner + 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)); + // Block Solve + auto time = graph->create_state(0.0); + auto dt = graph->create_state(0.025); + int 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}, // states + {T1_params, T2_params}, // params + smith::TimeInfo(time.get(), dt.get(), cycle), d_linear_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); +} + +int main(int argc, char* argv[]) +{ + ::testing::InitGoogleTest(&argc, argv); + MPI_Init(&argc, &argv); + + axom::slic::SimpleLogger logger; + + int result = RUN_ALL_TESTS(); + MPI_Finalize(); + + return result; +} diff --git a/src/smith/numerics/CMakeLists.txt b/src/smith/numerics/CMakeLists.txt index e9243a7829..6df0f7eb12 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 0000000000..136be9c501 --- /dev/null +++ b/src/smith/numerics/block_preconditioner.cpp @@ -0,0 +1,393 @@ +#include "smith/numerics/block_preconditioner.hpp" +#include "mfem.hpp" + +namespace smith { + +/** + * @brief Construct a new nxn block diagonal preconditioner + * + * @param offsets The array of offsets describing a block vector of (b_1, ..., b_n) + */ + +BlockDiagonalPreconditioner::BlockDiagonalPreconditioner(mfem::Array& offsets, std::vector> solvers) + : block_offsets_(offsets), nblocks_(offsets.Size() - 1), block_jacobian_(nullptr), solver_diag_(block_offsets_), mfem_solvers(std::move(solvers)) +{ + MFEM_VERIFY(mfem_solvers.size() == static_cast(nblocks_), + "Number of solvers must match number of blocks"); +} + +/** +* @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) +*/ +void BlockDiagonalPreconditioner::Mult(const mfem::Vector& in, mfem::Vector& out) const +{ + solver_diag_.Mult(in, out); +} + +/** +* @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 +*/ +void BlockDiagonalPreconditioner::SetOperator(const mfem::Operator& jacobian) +{ + // 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 < nblocks_; i++) + { + const mfem::Operator &A_ii = block_jacobian_->GetBlock(i, i); + + // Attach operator to solver + mfem_solvers[i]->SetOperator(A_ii); + + // Place the solver into the diagonal block of solver_diag_ + solver_diag_.SetBlock(i, i, mfem_solvers[i].get()); + } +} + +BlockDiagonalPreconditioner::~BlockDiagonalPreconditioner() {} + +/** + * @brief Construct a new nxn block triangular preconditioner + * + * @param offsets The array of offsets describing a block vector of (b_1, ..., b_n) + */ + +BlockTriangularPreconditioner::BlockTriangularPreconditioner(mfem::Array& offsets, std::vector> solvers, BlockTriangularType type) + : block_offsets_(offsets), nblocks_(offsets.Size() - 1), block_jacobian_(nullptr), mfem_solvers(std::move(solvers)), type_(type) +{ + MFEM_VERIFY(mfem_solvers.size() == static_cast(nblocks_), + "Number of solvers must match number of blocks"); +} + +/** +* @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 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..nblocks_-1 + for (int i = 0; i < nblocks_; 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[i]->Mult(rhs_i, xi); + } +} + + +/** +* @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 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 = nblocks_-1..0 + for (int i = nblocks_ - 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 < nblocks_; 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[i]->Mult(rhs_i, xi); + } +} + +/** +* @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) +*/ +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 < nblocks_; 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; + } + } +} + +/** +* @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 +*/ +void BlockTriangularPreconditioner::SetOperator(const mfem::Operator& jacobian) +{ + // Cast the supplied jacobian to a block operator object + block_jacobian_ = dynamic_cast(&jacobian); + + // Configure all diagonal solves + for (int i = 0; i < nblocks_; i++) + { + const mfem::Operator &A_ii = block_jacobian_->GetBlock(i, i); + mfem_solvers[i]->SetOperator(A_ii); + } +} + +BlockTriangularPreconditioner::~BlockTriangularPreconditioner() {} + + /** + * @brief Construct a new 2x2 block Schur complement factorization preconditioner + * + * @param offsets The array of offsets describing a block vector of (b_1, ..., b_n) + */ +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) +{ +} + +/** +* @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 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; +} + + +/** +* @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 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; +} +/** +* @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) +*/ +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; + } + } +} + +/** +* @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 +*/ +void BlockSchurPreconditioner::SetOperator(const mfem::Operator& jacobian) +{ + 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(MPI_COMM_WORLD, + 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 = 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 0000000000..6732469047 --- /dev/null +++ b/src/smith/numerics/block_preconditioner.hpp @@ -0,0 +1,114 @@ +#pragma once + +#include +#include +#include "mfem.hpp" + +namespace smith { + +class BlockDiagonalPreconditioner : public mfem::Solver { + public: + BlockDiagonalPreconditioner(mfem::Array& offsets, std::vector> solvers); + + virtual void Mult(const mfem::Vector& in, mfem::Vector& out) const; + + 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 nblocks_; + + // 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 class BlockTriangularType { + Lower, + Upper, + Symmetric +}; + +class BlockTriangularPreconditioner : public mfem::Solver { + public: + BlockTriangularPreconditioner(mfem::Array& offsets, std::vector> solvers, BlockTriangularType type = BlockTriangularType::Lower); + + virtual void Mult(const mfem::Vector& in, mfem::Vector& out) const; + + 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 nblocks_; + + // 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_; + + // Lower and Upper sweeps to be used in Mult + void LowerSweep(const mfem::Vector& in, mfem::Vector& out) const; + void UpperSweep(const mfem::Vector& in, mfem::Vector& out) const; +}; + +enum class BlockSchurType { + Diagonal, + Lower, + Upper, + Full +}; + +class BlockSchurPreconditioner : public mfem::Solver { + public: + BlockSchurPreconditioner(mfem::Array& offsets, std::vector> solvers, BlockSchurType type=BlockSchurType::Diagonal); + + virtual void Mult(const mfem::Vector& in, mfem::Vector& out) const; + + 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_11, *A_22, *A_12, *A_21; + + mutable mfem::HypreParMatrix *S_approx; + + BlockSchurType type_; + + // Lower and Upper blocks to be used in Mult + void LowerBlock(const mfem::Vector& in, mfem::Vector& out) const; + void UpperBlock(const mfem::Vector& in, mfem::Vector& out) const; +}; +} // namespace smith diff --git a/src/smith/numerics/tests/CMakeLists.txt b/src/smith/numerics/tests/CMakeLists.txt index 3df818dd92..1994b6298f 100644 --- a/src/smith/numerics/tests/CMakeLists.txt +++ b/src/smith/numerics/tests/CMakeLists.txt @@ -10,6 +10,7 @@ set(numerics_serial_test_sources equationsolver.cpp operator.cpp odes.cpp + test_block_preconditioner.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 0000000000..415104e80f --- /dev/null +++ b/src/smith/numerics/tests/test_block_preconditioner.cpp @@ -0,0 +1,458 @@ +#include +#include + +#include "mfem.hpp" +#include "smith/numerics/block_preconditioner.hpp" + +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) +{ + Array offsets({0, 2, 4}); + std::vector> solvers; + solvers.push_back(std::make_unique()); + + EXPECT_DEATH( + { smith::BlockDiagonalPreconditioner P(offsets, std::move(solvers)); }, + "Number of solvers must match number of blocks" + ); +} + +// 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[]) +{ + MPI_Init(&argc, &argv); + ::testing::InitGoogleTest(&argc, argv); + int result = RUN_ALL_TESTS(); + MPI_Finalize(); + return result; +}