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..d5787b3b94 --- /dev/null +++ b/src/smith/differentiable_numerics/tests/test_porous_heat_sink.cpp @@ -0,0 +1,263 @@ +#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 = 6; + int parallel_refinement = 2; + + 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; + + //else if (x[0] <= 2.0 / 16.0 && x[1] <= 11.0 / 16.0) return 1.0; + //else if (x[0] >= 9.0 / 16.0 && x[0] <= 12.0 / 16.0 && x[1] <= 11.0 / 16.0) return 1.0; + //else if (x[1] <= 3.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..67148bd250 --- /dev/null +++ b/src/smith/numerics/block_preconditioner.cpp @@ -0,0 +1,388 @@ +#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++) + { + const mfem::Operator &A_ij = block_jacobian_->GetBlock(i, j); + + // Is there a way to detect if coupling is zero? + + 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++) + { + const mfem::Operator &A_ij = block_jacobian_->GetBlock(i, j); + + // Is there a way to detect if coupling is zero? + 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