tboudreaux/SERiF
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(operator): added matrix free inverter and SchurComplement operator

Browse Source 
In order to maintain memory efficienty I have implimented a matrix free SchurComplement operator as well as an operator which uses a few iterations of GMRES to approxinate the inverse of any general operator.
This commit is contained in:
Emily Boudreaux
2025-05-18 15:32:08 -04:00
parent ddab27b833
commit b16ba8a7b6
2 changed files with 166 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files

237
src/poly/utils/private/operator.cpp
View File

@@ -24,40 +24,7 @@
#include "mfem_smout.h"
#include <memory>
#include "../../../../utils/debugUtils/MFEMAnalysisUtils/MFEMAnalysis-cpp/src/include/mfem_smout.h"
void approxJacobiInvert(const mfem::SparseMatrix& mat, std::unique_ptr<mfem::SparseMatrix>& invMat, const std::string& name="matrix") {
// PERF: This likely can be made much more efficient and will probably be called in tight loops, a good
// PERF: place for some easy optimization might be here.
// Confirm that mat is a square matrix
MFEM_ASSERT(mat.Height() == mat.Width(), "Matrix " + name + " is not square, cannot invert.");
mfem::Vector diag;
mat.GetDiag(diag);
// Invert the diagonal
for (int i = 0; i < diag.Size(); i++) {
MFEM_ASSERT(diag(i) != 0, "Diagonal element (" + std::to_string(i) +") in " + name + " is zero, cannot invert.");
diag(i) = 1.0 / diag(i);
}
// If the matrix is already inverted, just set the diagonal to avoid reallocation
if (invMat != nullptr) {
MFEM_ASSERT(invMat->Height() == invMat->Width(), "invMat (result matrix) is not square, cannot invert " + name + " into it.");
MFEM_ASSERT(invMat->Height() == mat.Height(), "Incompatible matrix sizes for inversion of " + name + ", expected " + std::to_string(mat.Height()) + " but got " + std::to_string(invMat->Height()));
for (int i = 0; i < diag.Size(); i++) {
MFEM_ASSERT(diag(i) != 0, "Diagonal element (" + std::to_string(i) +") in " + name + " is zero, resulting matrix would be singular.");
invMat->Elem(i, i) = diag(i);
}
} else { // The matrix has not been allocated yet so that needs to be done. Sparse Matrix has a constructor that can build from the diagonals
invMat = std::make_unique<mfem::SparseMatrix>(diag);
}
}
static int s_newtonStepGrad = 0;
static int s_newtonStepMult = 0;
//
PolytropeOperator::PolytropeOperator(
std::unique_ptr<mfem::MixedBilinearForm> M,
@@ -69,54 +36,51 @@ PolytropeOperator::PolytropeOperator(
mfem::Operator(blockOffsets.Last()), // Initialize the base class with the total size of the block offset vector
m_blockOffsets(blockOffsets),
m_jacobian(nullptr),
m_polytropicIndex(index){
m_jacobian(nullptr) {
m_M = std::move(M);
m_Q = std::move(Q);
m_D = std::move(D);
m_f = std::move(f);
// Use Gauss-Seidel smoother to approximate the inverse of the matrix
// t = 0 (symmetric Gauss-Seidel), 1 (forward Gauss-Seidel), 2 (backward Gauss-Seidel)
// iterations = 3
m_invNonlinearJacobian = std::make_unique<mfem::GSSmoother>(0, 3);
}
void PolytropeOperator::finalize(const mfem::Vector &initTheta) {
if (m_isFinalized) {
return; // do nothing if already finalized
return;
}
m_Mmat = std::make_unique<mfem::SparseMatrix>(m_M->SpMat());
m_Qmat = std::make_unique<mfem::SparseMatrix>(m_Q->SpMat());
m_Dmat = std::make_unique<mfem::SparseMatrix>(m_D->SpMat());
m_M -> EliminateTestEssentialBC(m_theta_ess_tdofs.first);
m_M -> EliminateTrialEssentialBC(m_phi_ess_tdofs.first);
// Remove the essential dofs from the constant matrices
for (const auto& dof: m_theta_ess_tdofs.first) {
m_Mmat->EliminateRow(dof);
m_Qmat->EliminateCol(dof);
}
m_Q -> EliminateTestEssentialBC(m_phi_ess_tdofs.first);
m_Q -> EliminateTrialEssentialBC(m_theta_ess_tdofs.first);
for (const auto& dof: m_phi_ess_tdofs.first) {
m_Mmat->EliminateCol(dof);
m_Qmat->EliminateRow(dof);
m_Dmat->EliminateRowCol(dof);
}
m_D -> EliminateEssentialBC(m_phi_ess_tdofs.first);
// m_negM_op = std::make_unique<mfem::ScaledOperator>(m_Mmat.get(), -1.0);
// m_negQ_op = std::make_unique<mfem::ScaledOperator>(m_Qmat.get(), -1.0);
m_negM_op = std::make_unique<mfem::ScaledOperator>(m_M.get(), -1.0);
m_negQ_op = std::make_unique<mfem::ScaledOperator>(m_Q.get(), -1.0);
// TODO Replace this with a scaled operator when I am done writing these out to disk
m_Mmat->operator*=(-1.0); // Sparse matrix only defines an inplace operator*= not an operator* so we need to do this
m_Qmat->operator*=(-1.0);
m_schurCompliment = std::make_unique<SchurCompliment>(*m_Q, *m_D, *m_M);
// Set up the constant parts of the jacobian now
m_jacobian = std::make_unique<mfem::BlockOperator>(m_blockOffsets);
m_jacobian->SetBlock(0, 1, m_Mmat.get()); //<- -M (constant)
m_jacobian->SetBlock(1, 0, m_Qmat.get()); //<- -Q (constant)
m_jacobian->SetBlock(1, 1, m_Dmat.get()); //<- D (constant)
m_jacobian->SetBlock(0, 1, m_negM_op.get()); //<- -M (constant)
m_jacobian->SetBlock(1, 0, m_negQ_op.get()); //<- -Q (constant)
m_jacobian->SetBlock(1, 1, m_D.get()); //<- D (constant)
m_invSchurCompliment = std::make_unique<GMRESInverter>(*m_schurCompliment);
m_isFinalized = true;
// Build the initial preconditioner based on some initial guess
const auto &grad = m_f->GetGradient(initTheta);
updatePreconditioner(grad);
m_isFinalized = true;
}
const mfem::BlockOperator &PolytropeOperator::GetJacobianOperator() const {
@@ -168,8 +132,8 @@ void PolytropeOperator::Mult(const mfem::Vector &x, mfem::Vector &y) const {
MFEM_ASSERT(m_f.get() != nullptr, "NonlinearForm m_f is null in PolytropeOperator::Mult");
m_f->Mult(x_theta, f_term); // fixme: this may be the wrong way to assemble m_f?
m_M->Mult(x_phi, Mphi_term); // Does the order of operations work out between this and the next subtract to make Mφ negative properly? I think so but double check
m_f->Mult(x_theta, f_term);
m_M->Mult(x_phi, Mphi_term);
m_D->Mult(x_phi, Dphi_term);
m_Q->Mult(x_theta, Qtheta_term);
@@ -193,59 +157,31 @@ void PolytropeOperator::Mult(const mfem::Vector &x, mfem::Vector &y) const {
std::cout << "||r_θ|| = " << y_block.GetBlock(0).Norml2();
std::cout << ", ||r_φ|| = " << y_block.GetBlock(1).Norml2() << std::endl;
//
// std::cout << "Writing out residuals to file..." << std::endl;
// std::ofstream outfile("Residuals_n" + std::to_string(m_polytropicIndex) + "_" + std::to_string(s_newtonStepMult) + ".csv", std::ios::trunc);
// if (!outfile.is_open()) {
// MFEM_ABORT("Could not open file for writing residuals");
// }
// outfile << "# Residuals for Newton Step: " << s_newtonStepMult << '\n';
// outfile << "# theta size: " << theta_size << '\n';
// outfile << "# phi size: " << phi_size << '\n';
// outfile << "dof,R\n";
// for (int i = 0; i < y_block.GetBlock(0).Size(); i++) {
// outfile << i << "," << y_block.GetBlock(0)[i] << '\n';
// }
// for (int i = 0; i < y_block.GetBlock(1).Size(); i++) {
// outfile << y_block.GetBlock(0).Size() + i << "," << y_block.GetBlock(1)[i] << '\n';
// }
// outfile.close();
// s_newtonStepMult++;
// std::cout << "Done writing out residuals to file." << std::endl;
//
}
void PolytropeOperator::updateInverseNonlinearJacobian(const mfem::Operator &grad) const {
if (const auto *sparse_mat = dynamic_cast<const mfem::SparseMatrix*>(&grad); sparse_mat != nullptr) {
approxJacobiInvert(*sparse_mat, m_invNonlinearJacobian, "Nonlinear Jacobian");
} else {
MFEM_ABORT("PolytropeOperator::GetGradient called on nonlinear jacobian where nonlinear jacobian is not dynamically castable to a sparse matrix");
}
m_invNonlinearJacobian->SetOperator(grad);
}
void PolytropeOperator::updateInverseSchurCompliment() const {
// TODO Add a flag in to make sure this tracks in parallel (i.e. every time the non linear jacobian inverse is updated set the flag to true and then check if the flag is true here and if so do work (if not throw error). then at the end of this function set it to false.
// TODO: This entire function could probably be refactored out
if (!m_isFinalized) {
MFEM_ABORT("PolytropeOperator::updateInverseSchurCompliment called before finalize");
}
if (m_invNonlinearJacobian == nullptr) {
MFEM_ABORT("PolytropeOperator::updateInverseSchurCompliment called before updateInverseNonlinearJacobian");
}
mfem::SparseMatrix* schurCompliment(&m_D->SpMat()); // Represents S in the preconditioner, starts as a copy of D before being modified in place
if (m_schurCompliment == nullptr) {
MFEM_ABORT("PolytropeOperator::updateInverseSchurCompliment called before updateInverseSchurCompliment");
}
// Calculate S = D - Q df^{-1} M
mfem::SparseMatrix* temp_QxdfInv = mfem::Mult(*m_Qmat, *m_invNonlinearJacobian); // Q * df^{-1}
const mfem::SparseMatrix* temp_QxdfInvxM = mfem::Mult(*temp_QxdfInv, *m_Mmat); // Q * df^{-1} * M
// PERF: Could potentially add some caching in here to only update the preconditioner when some condition has been met
schurCompliment->Add(1, *temp_QxdfInvxM); // D - Q * df^{-1} * M
saveSparseMatrixBinary(*schurCompliment, "schurCompliment.bin");
approxJacobiInvert(*schurCompliment, m_invSchurCompliment, "Schur Compliment");
m_schurCompliment->updateInverseNonlinearJacobian(*m_invNonlinearJacobian);
if (m_schurPreconditioner == nullptr) {
m_schurPreconditioner = std::make_unique<mfem::BlockDiagonalPreconditioner>(m_blockOffsets);
}
// ⎡ḟ(θ)^-1 0⎤
// ⎣0 S^-1⎦
m_schurPreconditioner->SetDiagonalBlock(0, m_invNonlinearJacobian.get());
m_schurPreconditioner->SetDiagonalBlock(1, m_invSchurCompliment.get());
@@ -256,7 +192,6 @@ void PolytropeOperator::updatePreconditioner(const mfem::Operator &grad) const {
updateInverseSchurCompliment();
}
mfem::Operator& PolytropeOperator::GetGradient(const mfem::Vector &x) const {
if (!m_isFinalized) {
MFEM_ABORT("PolytropeOperator::GetGradient called before finalize");
@@ -266,24 +201,10 @@ mfem::Operator& PolytropeOperator::GetGradient(const mfem::Vector &x) const {
const mfem::Vector& x_theta = x_block.GetBlock(0);
auto &grad = m_f->GetGradient(x_theta);
// auto *gradPtr = &grad;
// updatePreconditioner(grad);
updatePreconditioner(grad);
m_jacobian->SetBlock(0, 0, &grad);
std::cout << "Writing out jacobian to file..." << std::endl;
std::vector<mfem::Operator*> ops;
ops.push_back(&m_jacobian->GetBlock(0, 0));
ops.push_back(&m_jacobian->GetBlock(0, 1));
ops.push_back(&m_jacobian->GetBlock(1, 0));
ops.push_back(&m_jacobian->GetBlock(1, 1));
saveBlockFormToBinary(ops, {{0, 0}, {0, 1}, {1, 0}, {1, 1}}, {false, false, false, false}, "jacobian_" + std::to_string(m_polytropicIndex) + "_" + std::to_string(s_newtonStepGrad) + ".bin");
s_newtonStepGrad++;
std::cout << "Done writing out jacobian to file." << std::endl;
// // Exit the code here
// // throw std::runtime_error("Done writing out jacobian to file");
return *m_jacobian;
}
void PolytropeOperator::SetEssentialTrueDofs(const SSE::MFEMArrayPair& theta_ess_tdofs, const SSE::MFEMArrayPair& phi_ess_tdofs) {
@@ -304,4 +225,92 @@ void PolytropeOperator::SetEssentialTrueDofs(const SSE::MFEMArrayPairSet& ess_td
SSE::MFEMArrayPairSet PolytropeOperator::GetEssentialTrueDofs() const {
return std::make_pair(m_theta_ess_tdofs, m_phi_ess_tdofs);
}
}
GMRESInverter::GMRESInverter(const SchurCompliment &op) :
mfem::Operator(op.Height(), op.Width()),
m_op(op) {
m_solver.SetOperator(m_op);
// PERF: It might be a good idea to turn down the total number of iterations and the tolerances
// PERF: since we only need an approximation of the inverse
}
void GMRESInverter::Mult(const mfem::Vector &x, mfem::Vector &y) const {
m_solver.Mult(x, y); // Approximates m_op^-1 * x
}
SchurCompliment::SchurCompliment(
const mfem::MixedBilinearForm &QOp,
const mfem::BilinearForm &DOp,
const mfem::MixedBilinearForm &MOp,
const mfem::Solver &GradInvOp) :
mfem::Operator(DOp.Height(), DOp.Width())
{
SetOperator(QOp, DOp, MOp, GradInvOp);
m_nPhi = m_DOp->Height();
m_nTheta = m_MOp->Height();
}
SchurCompliment::SchurCompliment(
const mfem::MixedBilinearForm &QOp,
const mfem::BilinearForm &DOp,
const mfem::MixedBilinearForm &MOp) :
mfem::Operator(DOp.Height(), DOp.Width())
{
updateConstantTerms(QOp, DOp, MOp);
m_nPhi = m_DOp->Height();
m_nTheta = m_MOp->Height();
}
void SchurCompliment::SetOperator(const mfem::MixedBilinearForm &QOp, const mfem::BilinearForm &DOp, const mfem::MixedBilinearForm &MOp, const mfem::Solver &GradInvOp) {
updateConstantTerms(QOp, DOp, MOp);
updateInverseNonlinearJacobian(GradInvOp);
}
void SchurCompliment::updateInverseNonlinearJacobian(const mfem::Solver &gradInv) {
m_GradInvOp = &gradInv;
}
void SchurCompliment::updateConstantTerms(const mfem::MixedBilinearForm &QOp, const mfem::BilinearForm &DOp, const mfem::MixedBilinearForm &MOp) {
m_QOp = &QOp;
m_DOp = &DOp;
m_MOp = &MOp;
}
void SchurCompliment::Mult(const mfem::Vector &x, mfem::Vector &y) const {
// Check that the input vector is the correct size
if (x.Size() != m_nPhi) {
MFEM_ABORT("Input vector x has size " + std::to_string(x.Size()) + ", expected " + std::to_string(m_nPhi));
}
if (y.Size() != m_nPhi) {
MFEM_ABORT("Output vector y has size " + std::to_string(y.Size()) + ", expected " + std::to_string(m_nPhi));
}
// Check that the operators are set
if (m_QOp == nullptr) {
MFEM_ABORT("QOp is null in SchurCompliment::Mult");
}
if (m_DOp == nullptr) {
MFEM_ABORT("DOp is null in SchurCompliment::Mult");
}
if (m_MOp == nullptr) {
MFEM_ABORT("MOp is null in SchurCompliment::Mult");
}
if (m_GradInvOp == nullptr) {
MFEM_ABORT("GradInvOp is null in SchurCompliment::Mult");
}
mfem::Vector v1(m_nTheta); // M * x
m_MOp -> Mult(x, v1); // M * x
mfem::Vector v2(m_nTheta); // GradInv * M * x
m_GradInvOp -> Mult(v1, v2); // GradInv * M * x
mfem::Vector v3(m_nPhi); // Q * GradInv * M * x
m_QOp -> Mult(v2, v3); // Q * GradInv * M * x
mfem::Vector v4(m_nPhi); // D * x
m_DOp -> Mult(x, v4); // D * x
subtract(v4, v3, y); // (D - Q * GradInv * M) * x
}

47
src/poly/utils/public/operator.h
View File

@@ -26,6 +26,42 @@
#include "probe.h"
class SchurCompliment final : public mfem::Operator {
public:
SchurCompliment(const mfem::MixedBilinearForm &QOp, const mfem::BilinearForm &DOp, const mfem::MixedBilinearForm &MOp, const mfem::Solver &GradInvOp);
SchurCompliment(const mfem::MixedBilinearForm &QOp, const mfem::BilinearForm &DOp, const mfem::MixedBilinearForm &MOp);
~SchurCompliment() override = default;
void Mult(const mfem::Vector &x, mfem::Vector &y) const override;
void SetOperator(const mfem::MixedBilinearForm &QOp, const mfem::BilinearForm &DOp, const mfem::MixedBilinearForm &MOp, const mfem::Solver &GradInvOp);
void updateInverseNonlinearJacobian(const mfem::Solver &gradInv);
private:
void updateConstantTerms(const mfem::MixedBilinearForm &QOp, const mfem::BilinearForm &DOp, const mfem::MixedBilinearForm &MOp);
private:
// Note that these are not owned by this class
const mfem::MixedBilinearForm* m_QOp = nullptr;
const mfem::BilinearForm* m_DOp = nullptr;
const mfem::MixedBilinearForm* m_MOp = nullptr;
const mfem::Solver* m_GradInvOp = nullptr;
int m_nPhi = 0;
int m_nTheta = 0;
mutable std::unique_ptr<mfem::SparseMatrix> m_matrixForm;
};
class GMRESInverter final : public mfem::Operator {
public:
explicit GMRESInverter(const SchurCompliment& op);
~GMRESInverter() override = default;
void Mult(const mfem::Vector &x, mfem::Vector &y) const override;
private:
const SchurCompliment& m_op;
mfem::GMRESSolver m_solver;
};
class PolytropeOperator final : public mfem::Operator {
public:
PolytropeOperator(
@@ -70,13 +106,16 @@ private:
SSE::MFEMArrayPair m_theta_ess_tdofs;
SSE::MFEMArrayPair m_phi_ess_tdofs;
std::unique_ptr<mfem::SparseMatrix> m_Mmat, m_Qmat, m_Dmat;
// std::unique_ptr<mfem::SparseMatrix> m_Mmat, m_Qmat, m_Dmat;
std::unique_ptr<mfem::ScaledOperator> m_negM_op;
std::unique_ptr<mfem::ScaledOperator> m_negQ_op;
mutable std::unique_ptr<mfem::BlockOperator> m_jacobian;
mutable std::unique_ptr<mfem::SparseMatrix> m_invSchurCompliment;
mutable std::unique_ptr<mfem::SparseMatrix> m_invNonlinearJacobian;
// mutable std::unique_ptr<mfem::SparseMatrix> m_invSchurCompliment;
// mutable std::unique_ptr<mfem::SparseMatrix> m_invNonlinearJacobian;
mutable std::unique_ptr<SchurCompliment> m_schurCompliment;
mutable std::unique_ptr<GMRESInverter> m_invSchurCompliment;
mutable std::unique_ptr<mfem::Solver> m_invNonlinearJacobian;
/*
* The schur preconditioner has the form
@@ -91,10 +130,10 @@ private:
mutable std::unique_ptr<mfem::BlockDiagonalPreconditioner> m_schurPreconditioner;
bool m_isFinalized = false;
double m_polytropicIndex;
private:
void updateInverseNonlinearJacobian(const mfem::Operator &grad) const;
void updateInverseSchurCompliment() const;
void updatePreconditioner(const mfem::Operator &grad) const;
};