polytropeOperator.cpp

Go to the documentation of this file.

/* ***********************************************************************
//
//   Copyright (C) 2025 -- The 4D-STAR Collaboration
//   File Author: Emily Boudreaux
//   Last Modified: April 21, 2025
//
//   4DSSE is free software; you can use it and/or modify
//   it under the terms and restrictions the GNU General Library Public
//   License version 3 (GPLv3) as published by the Free Software Foundation.
//
//   4DSSE is distributed in the hope that it will be useful,
//   but WITHOUT ANY WARRANTY; without even the implied warranty of
//   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
//   See the GNU Library General Public License for more details.
//
//   You should have received a copy of the GNU Library General Public License
//   along with this software; if not, write to the Free Software
//   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
//
// *********************************************************************** */
#include "polytropeOperator.h"
#include "4DSTARTypes.h"
#include "utilities.h"
 
#include "mfem.hpp"
#include "mfem_smout.h"
#include <memory>
 
#include "config.h"
 
namespace serif {
namespace polytrope {
    // --- SchurCompliment Class Implementation ---
 
    // SchurCompliment: Constructors
SchurCompliment::SchurCompliment(
    const mfem::SparseMatrix &QOp,
    const mfem::SparseMatrix &DOp,
    const mfem::SparseMatrix &MOp,
    const mfem::Solver &GradInvOp ) :
mfem::Operator(DOp.Height(), DOp.Width()) {
    // Initialize sizes
    SetOperator(QOp, DOp, MOp, GradInvOp);
    m_nPhi = m_DOp->Height();
    m_nTheta = m_MOp->Height();
}
 
SchurCompliment::SchurCompliment(
    const mfem::SparseMatrix &QOp,
    const mfem::SparseMatrix &DOp,
    const mfem::SparseMatrix &MOp) :
mfem::Operator(DOp.Height(), DOp.Width())
{
    updateConstantTerms(QOp, DOp, MOp);
    m_nPhi = m_DOp->Height();
    m_nTheta = m_MOp->Height();
}
 
// SchurCompliment: Public Interface Methods
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
}
 
void SchurCompliment::SetOperator(const mfem::SparseMatrix &QOp, const mfem::SparseMatrix &DOp, const mfem::SparseMatrix &MOp, const mfem::Solver &GradInvOp) {
    updateConstantTerms(QOp, DOp, MOp);
    updateInverseNonlinearJacobian(GradInvOp);
}
 
void SchurCompliment::updateInverseNonlinearJacobian(const mfem::Solver &gradInv) {
    m_GradInvOp = &gradInv;
}
 
// SchurCompliment: Private Helper Methods
void SchurCompliment::updateConstantTerms(const mfem::SparseMatrix &QOp, const mfem::SparseMatrix &DOp, const mfem::SparseMatrix &MOp) {
    m_QOp = &QOp;
    m_DOp = &DOp;
    m_MOp = &MOp;
}
 
 
// --- GMRESInverter Class Implementation ---
 
// GMRESInverter: Constructor
GMRESInverter::GMRESInverter(const SchurCompliment &op) :
mfem::Operator(op.Height(), op.Width()),
m_op(op) {
    m_solver.SetOperator(m_op);
    m_solver.SetMaxIter(100);
    m_solver.SetRelTol(1e-1);
    m_solver.SetAbsTol(1e-1);
}
 
// GMRESInverter: Public Interface Methods
void GMRESInverter::Mult(const mfem::Vector &x, mfem::Vector &y) const {
    m_solver.Mult(x, y); // Approximates m_op^-1 * x
}
 
 
// --- PolytropeOperator Class Implementation ---
    // PolytropeOperator: Constructor
    PolytropeOperator::PolytropeOperator(
      std::unique_ptr<mfem::MixedBilinearForm> M,
      std::unique_ptr<mfem::MixedBilinearForm> Q,
      std::unique_ptr<mfem::BilinearForm> D,
      std::unique_ptr<mfem::BilinearForm> S,
      std::unique_ptr<mfem::NonlinearForm> f,
      const mfem::Array<int> &blockOffsets) :
 
      // TODO: Need to update this so that the size is that of the reduced system operator
      mfem::Operator(blockOffsets.Last()), // Initialize the base class with the total size of the block offset vector
      m_blockOffsets(blockOffsets),
      m_jacobian(nullptr) {
 
        m_M = std::move(M);
        m_Q = std::move(Q);
        m_D = std::move(D);
        m_S = std::move(S);
        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);
    }
 
// PolytropeOperator: Core Operator Overrides
void PolytropeOperator::Mult(const mfem::Vector &xFree, mfem::Vector &yFree) const {
    if (!m_isFinalized) {
        MFEM_ABORT("PolytropeOperator::Mult called before finalize");
    }
 
    // TODO: confirm that the vectors xFree and m_freeDofs are always parallel
    m_state.SetSubVector(m_freeDofs, xFree); // Scatter the free dofs from the input vector xFree into the state vector
    mfem::Vector y;
    y.SetSize(m_blockOffsets.Last());
 
    // -- Create BlockVector views for input x and output y
    mfem::BlockVector x_block(const_cast<mfem::Vector&>(m_state), m_blockOffsets);
    mfem::BlockVector y_block(y, m_blockOffsets);
 
    // -- Get Vector views for individual blocks
    const mfem::Vector &x_theta = x_block.GetBlock(0);
    const mfem::Vector &x_phi = x_block.GetBlock(1);
    mfem::Vector &y_R0 = y_block.GetBlock(0); // Residual Block 0 (theta)
    mfem::Vector &y_R1 = y_block.GetBlock(1); // Residual Block 1 (phi)
 
    int theta_size = m_blockOffsets[1] - m_blockOffsets[0];
    int phi_size = m_blockOffsets[2] - m_blockOffsets[1];
 
    mfem::Vector f_term(theta_size);
    mfem::Vector Mphi_term(theta_size);
    mfem::Vector Dphi_term(phi_size);
    mfem::Vector Qtheta_term(phi_size);
    mfem::Vector Stheta_term(theta_size);
 
 
    // Calculate R0 and R1 terms
    // R0 = f(θ) + (1+c)Mɸ + cSθ
    // R1 = (1+c)Dɸ - (1+c)Qθ
 
    MFEM_ASSERT(m_f.get() != nullptr, "NonlinearForm m_f is null in PolytropeOperator::Mult");
 
    m_f->Mult(x_theta, f_term);      // f(θ)
    m_M->Mult(x_phi, Mphi_term);     // Mɸ
    m_D->Mult(x_phi, Dphi_term);     // Dɸ
    m_Q->Mult(x_theta, Qtheta_term); // Qθ
    m_S->Mult(x_theta, Stheta_term); // Sθ
 
    // PERF: these multiplications may be able to be refactored into the matrices so they only need to be done once.
    Stheta_term *= m_stabilizationCoefficient;            // cSθ
    Qtheta_term *= m_IncrementedStabilizationCoefficient; // (1+c)Qθ
    Mphi_term *= m_IncrementedStabilizationCoefficient;   // (1+c)Mɸ
    Dphi_term *= m_IncrementedStabilizationCoefficient;   // (1+c)Dɸ
 
    mfem::Vector y_R0_temp(theta_size);
    add(f_term, Mphi_term, y_R0_temp);      // R0 = f(θ) + (1+c)Mɸ
    add(y_R0_temp, Stheta_term, y_R0);      // R0 = f(θ) + (1+c)Mɸ + cSθ
    subtract(Dphi_term, Qtheta_term, y_R1); // R1 = (1+c)Dɸ - (1+c)Qθ
 
    // --- Scatter the residual vector y into the free dofs ---
    yFree.SetSize(m_reducedBlockOffsets.Last());
    MFEM_ASSERT(m_freeDofs.Size() == m_reducedBlockOffsets.Last(), "PolytropeOperator::Mult: Size of free dofs does not match reduced block offsets size.");
    for (int i = 0, j = 0; i < y.Size(); ++i) {
        if (m_freeDofs.Find(i) != -1) {
            yFree[j] = y[i];
            j++;
        }
    }
}
 
mfem::Operator& PolytropeOperator::GetGradient(const mfem::Vector &xFree) const {
    if (!m_isFinalized) {
        MFEM_ABORT("PolytropeOperator::GetGradient called before finalize");
    }
    m_state.SetSubVector(m_freeDofs, xFree); // Scatter the free dofs from the input vector xFree into the state vector
    // --- Get the gradient of f ---
    mfem::BlockVector x_block(const_cast<mfem::Vector&>(m_state), m_blockOffsets);
    const mfem::Vector& x_theta = x_block.GetBlock(0);
 
    // PERF: There are a lot of copies and loops here, probably performance could be gained by flattering some of these.
    auto &grad = m_f->GetGradient(x_theta);
    // updatePreconditioner(grad);
    const auto gradMatrix = dynamic_cast<mfem::SparseMatrix*>(&grad); // ∂f(θ)/∂θ
 
    if (gradMatrix == nullptr) {
        MFEM_ABORT("PolytropeOperator::GetGradient: Gradient is not a SparseMatrix.");
    }
 
    m_gradReduced = std::make_unique<mfem::SparseMatrix> (
        serif::utilities::build_reduced_matrix(
            *gradMatrix,
            m_theta_ess_tdofs.first,
            m_theta_ess_tdofs.first
            )
        ); // ∂f(θ)/∂θ (now reduced to only the free DOFs)
 
 
    // TODO: Confirm this actually does what I want it to do
    *m_gradReduced += *m_ScaledSReduced; // ∂f(θ)/∂θ + cS (reduced to only the free DOFs)
 
    m_jacobian->SetBlock(0, 0, m_gradReduced.get());
 
    return *m_jacobian;
}
 
 
// PolytropeOperator: Setup and Finalization
void PolytropeOperator::finalize(const mfem::Vector &initTheta) {
    using serif::utilities::build_reduced_matrix;
 
    if (m_isFinalized) {
        return;
    }
 
    // These functions must be called in this order since they depend on each others post state
    // TODO: Refactor this so that either there are explicit checks to make sure the order is correct or make
    //       them pure functions
    construct_matrix_representations();
    construct_reduced_block_offsets();
    construct_jacobian_constant_terms();
    scatter_boundary_conditions();
    populate_free_dof_array();
 
    // Override the size based on the reduced system
    height = m_reducedBlockOffsets.Last();
    width = m_reducedBlockOffsets.Last();
 
    m_isFinalized = true;
 
}
 
// PolytropeOperator: Essential True DOF Management
void PolytropeOperator::set_essential_true_dofs(const serif::types::MFEMArrayPair& theta_ess_tdofs, const serif::types::MFEMArrayPair& phi_ess_tdofs) {
    m_isFinalized = false;
    m_theta_ess_tdofs = theta_ess_tdofs;
    m_phi_ess_tdofs = phi_ess_tdofs;
 
    if (m_f) {
        m_f->SetEssentialTrueDofs(theta_ess_tdofs.first);
    } else {
        MFEM_ABORT("m_f is null in PolytropeOperator::SetEssentialTrueDofs");
    }
}
 
void PolytropeOperator::set_essential_true_dofs(const serif::types::MFEMArrayPairSet& ess_tdof_pair_set) {
    set_essential_true_dofs(ess_tdof_pair_set.first, ess_tdof_pair_set.second);
}
 
serif::types::MFEMArrayPairSet PolytropeOperator::get_essential_true_dofs() const {
     return std::make_pair(m_theta_ess_tdofs, m_phi_ess_tdofs);
}
 
// PolytropeOperator: Getter Methods
const mfem::BlockOperator &PolytropeOperator::get_jacobian_operator() const {
    if (m_jacobian == nullptr) {
        MFEM_ABORT("Jacobian has not been initialized before GetJacobianOperator() call.");
    }
    if (!m_isFinalized) {
        MFEM_ABORT("PolytropeOperator not finalized prior to call to GetJacobianOperator().");
    }
 
    return *m_jacobian;
}
 
mfem::BlockDiagonalPreconditioner& PolytropeOperator::get_preconditioner() const {
    if (m_schurPreconditioner == nullptr) {
        MFEM_ABORT("Schur preconditioner has not been initialized before GetPreconditioner() call.");
    }
    if (!m_isFinalized) {
        MFEM_ABORT("PolytropeOperator not finalized prior to call to GetPreconditioner().");
    }
    return *m_schurPreconditioner;
}
 
int PolytropeOperator::get_reduced_system_size() const {
    if (!m_isFinalized) {
        MFEM_ABORT("PolytropeOperator not finalized prior to call to GetReducedSystemSize().");
    }
    return m_reducedBlockOffsets.Last();
}
 
// PolytropeOperator: State Reconstruction
const mfem::Vector &PolytropeOperator::reconstruct_full_state_vector(const mfem::Vector &reducedState) const {
    m_state.SetSubVector(m_freeDofs, reducedState); // Scatter the reduced state vector into the full state vector
    return m_state;
}
 
const mfem::BlockVector PolytropeOperator::reconstruct_full_block_state_vector(const mfem::Vector &reducedState) const {
    m_state.SetSubVector(m_freeDofs, reducedState); // Scatter the reduced state vector into the full state vector
    mfem::BlockVector x_block(m_state, m_blockOffsets);
    return x_block;
}
 
// PolytropeOperator: DOF Population
void PolytropeOperator::populate_free_dof_array() {
    m_freeDofs.SetSize(0);
    for (int i = 0; i < m_blockOffsets.Last(); i++) {
        const int thetaSearchIndex = i;
        const int phiSearchIndex = i - m_blockOffsets[1];
        if (phiSearchIndex < 0){
            if (m_theta_ess_tdofs.first.Find(thetaSearchIndex) == -1) {
                m_freeDofs.Append(i);
            }
        } else {
            if (m_phi_ess_tdofs.first.Find(phiSearchIndex) == -1) {
                m_freeDofs.Append(i);
            }
        }
    }
}
 
// PolytropeOperator: Private Helper Methods - Construction and Setup
void PolytropeOperator::construct_matrix_representations() {
    // --- Construct the sparse matrix representations of M, Q, D, and S ---
    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_Smat = std::make_unique<mfem::SparseMatrix>(m_S->SpMat());
 
    // --- Reduce the matrices to only the free DOFs ---
    m_MReduced = std::make_unique<mfem::SparseMatrix>(
        serif::utilities::build_reduced_matrix(
            *m_Mmat,
            m_phi_ess_tdofs.first,
            m_theta_ess_tdofs.first)
    );
    m_QReduced = std::make_unique<mfem::SparseMatrix>(
        serif::utilities::build_reduced_matrix(
            *m_Qmat,
            m_theta_ess_tdofs.first,
            m_phi_ess_tdofs.first)
    );
    m_DReduced = std::make_unique<mfem::SparseMatrix>(
        serif::utilities::build_reduced_matrix(
            *m_Dmat,
            m_phi_ess_tdofs.first,
            m_phi_ess_tdofs.first)
    );
    m_SReduced = std::make_unique<mfem::SparseMatrix>(
        serif::utilities::build_reduced_matrix(
            *m_Smat,
            m_theta_ess_tdofs.first,
            m_theta_ess_tdofs.first)
    );
 
    // --- Scale the reduced matrices by the stabilization coefficients ---
    mfem::SparseMatrix MScaledTemp(*m_MReduced); // Create a temporary copy of the M matrix for scaling
    MScaledTemp *= m_IncrementedStabilizationCoefficient; // Scale the M matrix by the incremented stabilization coefficient
    m_MScaledReduced = std::make_unique<mfem::SparseMatrix>(MScaledTemp); // Store the scaled M matrix so that it persists
 
    mfem::SparseMatrix QScaledTemp(*m_QReduced);
    QScaledTemp *= m_IncrementedStabilizationCoefficient; // Scale the Q matrix by the incremented stabilization coefficient
    m_QScaledReduced = std::make_unique<mfem::SparseMatrix>(QScaledTemp); // Store the scaled Q matrix so that it persists
 
    mfem::SparseMatrix DRScaledTemp(*m_DReduced);
    DRScaledTemp *= m_IncrementedStabilizationCoefficient; // Scale the D matrix by the incremented stabilization coefficient
    m_DScaledReduced = std::make_unique<mfem::SparseMatrix>(DRScaledTemp); // Store the scaled D matrix so that it persists
 
    // Scale the S matrix by the stabilization coefficient (so that the allocations only need to be done once)
    mfem::SparseMatrix SScaledTemp(*m_SReduced); // Create a temporary copy of the S matrix for scaling
    SScaledTemp *= m_stabilizationCoefficient; // Scale the S matrix by the stabilization coefficient
    m_ScaledSReduced = std::make_unique<mfem::SparseMatrix>(SScaledTemp); // Store the scaled S matrix so that it persists
 
    // --- Create the scaled operator for -(1+c)Q ---
    m_negQ_mat = std::make_unique<mfem::ScaledOperator>(m_QScaledReduced.get(), -1.0);
}
 
void PolytropeOperator::construct_reduced_block_offsets() {
    m_reducedBlockOffsets.SetSize(3);
    m_reducedBlockOffsets[0] = 0; // R0 block (theta)
    m_reducedBlockOffsets[1] = m_MReduced->Height(); // R1 block (theta)
    m_reducedBlockOffsets[2] = m_QReduced->Height() + m_reducedBlockOffsets[1]; // R2 block (phi + theta)
}
 
void PolytropeOperator::construct_jacobian_constant_terms() {
    m_jacobian = std::make_unique<mfem::BlockOperator>(m_reducedBlockOffsets);
 
    m_jacobian->SetBlock(0, 1, m_MScaledReduced.get());     //<-  (1+c)M (constant, reduced to only free DOFs)
    m_jacobian->SetBlock(1, 0, m_negQ_mat.get());           //<- -(1+c)Q (constant, reduced to only free DOFs)
    m_jacobian->SetBlock(1, 1, m_DScaledReduced.get());     //<-  (1+c)D (constant, reduced to only free DOFs)
}
 
void PolytropeOperator::scatter_boundary_conditions() {
    mfem::Vector thetaStateValues(m_theta_ess_tdofs.first.Size());
    for (int i = 0; i < m_theta_ess_tdofs.first.Size(); i++) {
        thetaStateValues[i] = m_theta_ess_tdofs.second[i];
    }
    mfem::Vector phiStateValues(m_phi_ess_tdofs.first.Size());
    for (int i = 0; i < m_phi_ess_tdofs.first.Size(); i++) {
        phiStateValues[i] = m_phi_ess_tdofs.second[i]; // TODO: figure out if this needs to be normalized
    }
 
    mfem::Array<int> phiDofIndices(m_phi_ess_tdofs.first.Size());
    for (int i = 0; i < m_phi_ess_tdofs.first.Size(); i++) {
        phiDofIndices[i] = m_phi_ess_tdofs.first[i] + m_blockOffsets[1];
    }
 
    m_state.SetSize(m_blockOffsets.Last());
    m_state = 0.0;
    m_state.SetSubVector(m_theta_ess_tdofs.first, thetaStateValues);
    m_state.SetSubVector(phiDofIndices, phiStateValues);
 
}
 
// PolytropeOperator: Private Helper Methods - Jacobian and Preconditioner Updates
void PolytropeOperator::update_inverse_nonlinear_jacobian(const mfem::Operator &grad) const {
    m_invNonlinearJacobian->SetOperator(grad);
}
 
void PolytropeOperator::update_inverse_schur_compliment() const {
    // 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");
    }
    if (m_schurCompliment == nullptr) {
        MFEM_ABORT("PolytropeOperator::updateInverseSchurCompliment called before updateInverseSchurCompliment");
    }
 
    m_schurCompliment->updateInverseNonlinearJacobian(*m_invNonlinearJacobian);
 
    if (m_schurPreconditioner == nullptr) {
        m_schurPreconditioner = std::make_unique<mfem::BlockDiagonalPreconditioner>(m_blockOffsets);
    }
 
    m_schurPreconditioner->SetDiagonalBlock(0, m_invNonlinearJacobian.get());
    m_schurPreconditioner->SetDiagonalBlock(1, m_invSchurCompliment.get());
 
}
 
void PolytropeOperator::update_preconditioner(const mfem::Operator &grad) const {
    update_inverse_nonlinear_jacobian(grad);
    update_inverse_schur_compliment();
}
 
} // namespace polytrope
} // namespace serif