SERiF/src/poly/utils/private/operator.cpp

/* ***********************************************************************
//
//   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 "operator.h"
#include "4DSTARTypes.h"
#include "mfem.hpp"
#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,
  std::unique_ptr<mfem::MixedBilinearForm> Q,
  std::unique_ptr<mfem::BilinearForm> D,
  std::unique_ptr<mfem::NonlinearForm> f,
  const mfem::Array<int> &blockOffsets) : 

  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_f = std::move(f);
}

void PolytropeOperator::finalize(const mfem::Vector &initTheta) {
    if (m_isFinalized) {
        return; // do nothing if already finalized
    }

    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());

    // 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);
    }

    for (const auto& dof: m_phi_ess_tdofs.first) {
        m_Mmat->EliminateCol(dof);
        m_Qmat->EliminateRow(dof);
        m_Dmat->EliminateRowCol(dof);
    }

    // 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);

    // 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);

    // 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)

    // 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 {
    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::GetPreconditioner() 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;
}

void PolytropeOperator::Mult(const mfem::Vector &x, mfem::Vector &y) const {
    if (!m_isFinalized) {
        MFEM_ABORT("PolytropeOperator::Mult called before finalize");
    }
    // -- Create BlockVector views for input x and output y
    mfem::BlockVector x_block(const_cast<mfem::Vector&>(x), 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);

    // Calculate R0 and R1 terms
    // R0 = f(θ) - Mɸ
    // R1 = Dɸ - Qθ

    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_D->Mult(x_phi, Dphi_term);
    m_Q->Mult(x_theta, Qtheta_term);

    subtract(f_term, Mphi_term, y_R0);
    subtract(Dphi_term, Qtheta_term, y_R1);

    // -- Apply essential boundary conditions --
    for (int i = 0; i < m_theta_ess_tdofs.first.Size(); i++) {
        if (int idx = m_theta_ess_tdofs.first[i]; idx >= 0 && idx < y_R0.Size()) {
            const double &targetValue = m_theta_ess_tdofs.second[i];
            y_block.GetBlock(0)[idx] = x_theta(idx) - targetValue; // inhomogenous essential bc. This is commented out since seems it results in dramatic instabilities arising
        }
    }
    std::cout << "θ Norm: " << y_block.GetBlock(0).Norml2() << std::endl;

    for (int i = 0; i < m_phi_ess_tdofs.first.Size(); i++) {
        if (int idx = m_phi_ess_tdofs.first[i]; idx >= 0 && idx < y_R1.Size()) {
            const double &targetValue = m_phi_ess_tdofs.second[i];
            y_block.GetBlock(1)[idx] = x_phi(idx) - targetValue; // inhomogenous essential bc. This is commented out since seems it results in dramatic instabilities arising
        }
    }

    // std::cout << "φ Norm: " << y_block.GetBlock(1).Norml2() << std::endl;
    //
    // std::cout << "Writing out residuals to file..." << std::endl;
    // std::ofstream outfile("Residuals_" + 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");
    }
}

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.
    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

    // 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");

    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());

}

void PolytropeOperator::updatePreconditioner(const mfem::Operator &grad) const {
    updateInverseNonlinearJacobian(grad);
    updateInverseSchurCompliment();
}


mfem::Operator& PolytropeOperator::GetGradient(const mfem::Vector &x) const {
    if (!m_isFinalized) {
        MFEM_ABORT("PolytropeOperator::GetGradient called before finalize");
    }
    // --- Get the gradient of f ---
    mfem::BlockVector x_block(const_cast<mfem::Vector&>(x), m_blockOffsets);
    const mfem::Vector& x_theta = x_block.GetBlock(0);

    auto &grad = m_f->GetGradient(x_theta);
    // auto *gradPtr = &grad;
    // updatePreconditioner(grad);

    // // TODO: Save the state vector somewhere so that I can inspect how the actual solution
    // // to theta and phi are evolving
    // m_jacobian->SetBlock(0, 0, &grad);
    // // TODO: Save the jacobian here
    // 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(s_newtonStepGrad) + ".bin");
    // s_newtonStepGrad++;
    // std::cout << "Done writing out jacobian to file." << std::endl;
    return *m_jacobian;
}

void PolytropeOperator::SetEssentialTrueDofs(const SSE::MFEMArrayPair& theta_ess_tdofs, const SSE::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::SetEssentialTrueDofs(const SSE::MFEMArrayPairSet& ess_tdof_pair_set) {
    SetEssentialTrueDofs(ess_tdof_pair_set.first, ess_tdof_pair_set.second);
}

SSE::MFEMArrayPairSet PolytropeOperator::GetEssentialTrueDofs() const {
     return std::make_pair(m_theta_ess_tdofs, m_phi_ess_tdofs);
}