block_preconditioner.hpp

#pragma once
 
#include <memory>
#include <functional>
#include "mfem.hpp"
 
namespace smith {
 
using BlockOverride = std::pair<int, std::unique_ptr<const mfem::Operator>>;
 
class BlockDiagonalPreconditioner : public mfem::Solver {
 public:
  BlockDiagonalPreconditioner(mfem::Array<int>& offsets, std::vector<std::unique_ptr<mfem::Solver>> solvers,
                              std::vector<BlockOverride> overrides = {});
 
  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<int>& 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<std::unique_ptr<mfem::Solver>> mfem_solvers_;
 
  // size num_blocks_, nullptr means "use Jacobian diagonal block"
  std::vector<std::unique_ptr<const mfem::Operator>> block_op_overrides_;
};
 
enum class BlockTriangularType
{
  Lower,    
  Upper,    
  Symmetric 
};
 
class BlockTriangularPreconditioner : public mfem::Solver {
 public:
  BlockTriangularPreconditioner(mfem::Array<int>& offsets, std::vector<std::unique_ptr<mfem::Solver>> solvers,
                                BlockTriangularType type = BlockTriangularType::Lower,
                                std::vector<BlockOverride> overrides = {});
 
  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<int>& 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<std::unique_ptr<mfem::Solver>> mfem_solvers_;
 
  // Block Triangular type
  BlockTriangularType type_;
 
  void LowerSweep(const mfem::Vector& in, mfem::Vector& out) const;
 
  void UpperSweep(const mfem::Vector& in, mfem::Vector& out) const;
 
  // size num_blocks_, nullptr means "use Jacobian diagonal block"
  std::vector<std::unique_ptr<const mfem::Operator>> block_op_overrides_;
};
 
enum class BlockSchurType
{
  Diagonal, 
  Lower,    
  Upper,    
  Full      
};
 
enum class SchurApproxType
{
  DiagInv, 
  A22Only, 
  Custom,  
};
 
class BlockSchurPreconditioner : public mfem::Solver {
 public:
  BlockSchurPreconditioner(mfem::Array<int>& offsets, std::vector<std::unique_ptr<mfem::Solver>> solvers,
                           BlockSchurType type = BlockSchurType::Diagonal,
                           SchurApproxType approxType = SchurApproxType::DiagInv,
                           std::vector<BlockOverride> overrides = {});
 
  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<int>& 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<std::unique_ptr<mfem::Solver>> mfem_solvers_;
 
  // Views of the linearized Jacobian blocks
  const mfem::Operator* A_12_ = nullptr;
  const mfem::Operator* A_21_ = nullptr;
 
  // Schur complement approximation operator used by solver for block (1,1).
  //
  // For DiagInv and A22Only, the approximation is rebuilt on each SetOperator call and stored in
  // S_approx_owned_. For Custom, the approximation is provided via block_op_overrides_[1] and referenced
  // non-owningly via S_approx_view_.
  mutable std::unique_ptr<const mfem::Operator> S_approx_owned_;
  const mfem::Operator* S_approx_view_ = nullptr;
 
  BlockSchurType type_;
 
  SchurApproxType approxType_;
 
  void LowerBlock(const mfem::Vector& in, mfem::Vector& out) const;
 
  void UpperBlock(const mfem::Vector& in, mfem::Vector& out) const;
 
  // size num_blocks_, nullptr means "use Jacobian diagonal block"
  std::vector<std::unique_ptr<const mfem::Operator>> block_op_overrides_;
 
  mfem::HypreParMatrix* BuildSchurDiagApprox_(const mfem::HypreParMatrix& A11, const mfem::HypreParMatrix& A12,
                                              const mfem::HypreParMatrix& A21, const mfem::HypreParMatrix& A22) const;
};
}  // namespace smith