polySolver.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 "polySolver.h"
 
#include <memory>
#include <stdexcept>
#include <string>
#include <utility>
 
#include "mfem.hpp"
 
#include "4DSTARTypes.h"
#include "config.h"
#include "integrators.h"
#include "mfem.hpp"
#include "polytropeOperator.h"
#include "polyCoeff.h"
#include "probe.h"
#include "resourceManager.h"
#include "resourceManagerTypes.h"
#include "utilities.h"
#include "quill/LogMacros.h"
 
 
namespace serif::polytrope {
 
namespace laneEmden {
 
    double a (const int k, const double n) { // NOLINT(*-no-recursion)
        if ( k == 0 ) { return 1; }
        if ( k == 1 ) { return 0; }
        else { return -(c(k-2, n)/(std::pow(k, 2)+k)); }
 
    }
 
    double c(const int m, const double n) { // NOLINT(*-no-recursion)
        if ( m == 0 ) { return std::pow(a(0, n), n); }
        else {
            double termOne = 1.0/(m*a(0, n));
            double acc = 0;
            for (int k = 1; k <= m; k++) {
                acc += (k*n-m+k)*a(k, n)*c(m-k, n);
            }
            return termOne*acc;
        }
    }
 
    double thetaSeriesExpansion(const double xi, const double n, const int order) {
        double acc = 0;
        for (int k = 0; k < order; k++) {
            acc += a(k, n) * std::pow(xi, k);
        }
        return acc;
    }
} // namespace laneEmden
 
 
PolySolver::PolySolver(mfem::Mesh& mesh, const double n, const double order)
    : m_config(serif::config::Config::getInstance()), // Updated
      m_logManager(serif::probe::LogManager::getInstance()),
      m_logger(m_logManager.getLogger("log")),
      m_polytropicIndex(n),
      m_feOrder(order),
      m_mesh(mesh) {
 
    // Use feOrder - 1 for the RT space to satisfy Brezzi-Babuska condition
    // for the H1 and RT [H(div)] spaces
    m_fecH1 = std::make_unique<mfem::H1_FECollection>(m_feOrder, m_mesh.SpaceDimension());
    m_fecRT = std::make_unique<mfem::RT_FECollection>(m_feOrder - 1, m_mesh.SpaceDimension());
 
    m_feTheta = std::make_unique<mfem::FiniteElementSpace>(&m_mesh, m_fecH1.get());
    m_fePhi = std::make_unique<mfem::FiniteElementSpace>(&m_mesh, m_fecRT.get());
 
    m_theta = std::make_unique<mfem::GridFunction>(m_feTheta.get());
    m_phi = std::make_unique<mfem::GridFunction>(m_fePhi.get());
 
    assembleBlockSystem();
}
 
PolySolver::PolySolver(const double n, const double order)
    : PolySolver(prepareMesh(n), n, order){}
 
mfem::Mesh& PolySolver::prepareMesh(const double n) {
    if (n > 4.99 || n < 0.0) {
        throw std::runtime_error("The polytropic index n must be less than 5.0 and greater than 0.0. Currently it is " + std::to_string(n));
    }
    const serif::resource::ResourceManager& rm = serif::resource::ResourceManager::getInstance();
    const serif::resource::types::Resource& genericResource = rm.getResource("mesh:polySphere");
    const auto &meshIO = std::get<std::unique_ptr<serif::mesh::MeshIO>>(genericResource);
    meshIO->LinearRescale(polycoeff::x1(n)); // Assumes polycoeff is now serif::polytrope::polycoeff
    return meshIO->GetMesh();
}
 
PolySolver::~PolySolver() = default;
 
void PolySolver::assembleBlockSystem() {
    mfem::Array<int> blockOffsets = computeBlockOffsets();
 
    const std::unique_ptr<formBundle> forms = buildIndividualForms(blockOffsets);
 
    // const double penalty_param = m_config.get<double>("Poly::Solver::ZeroDerivativePenalty", 1.0);
    // mfem::Array<int> thetaCenterDofs, phiCenterDofs;
    // std::tie(thetaCenterDofs, phiCenterDofs) = findCenterElement();
    // mfem::SparseMatrix& D_mat = forms->D->SpMat();
    //
    // for (int i = 0; i < phiCenterDofs.Size(); ++i)
    // {
    //     const int dof_idx = phiCenterDofs[i];
    //     if (dof_idx >= 0 && dof_idx < D_mat.Height()) {
    //         D_mat(dof_idx, dof_idx) += penalty_param;
    //     }
    // }
 
    // --- Build the BlockOperator ---
    m_polytropOperator = std::make_unique<PolytropeOperator>(
        std::move(forms->M),
        std::move(forms->Q),
        std::move(forms->D),
        std::move(forms->S),
        std::move(forms->f),
        blockOffsets);
}
 
mfem::Array<int> PolySolver::computeBlockOffsets() const {
    mfem::Array<int> blockOffsets;
    blockOffsets.SetSize(3);
    blockOffsets[0] = 0;
    blockOffsets[1] = m_feTheta->GetVSize(); // Get actual number of dofs *before* applying BCs
    blockOffsets[2] = m_fePhi->GetVSize();
    blockOffsets.PartialSum(); // Cumulative sum to get the offsets
    return blockOffsets;
}
 
std::unique_ptr<formBundle> PolySolver::buildIndividualForms(const mfem::Array<int> &blockOffsets) {
    // --- Assemble the MixedBilinear and Bilinear forms (M, D, and Q) ---
    auto forms = std::make_unique<formBundle>(
        std::make_unique<mfem::MixedBilinearForm>(m_fePhi.get(), m_feTheta.get()),
        std::make_unique<mfem::MixedBilinearForm>(m_feTheta.get(), m_fePhi.get()),
        std::make_unique<mfem::BilinearForm>(m_fePhi.get()),
        std::make_unique<mfem::BilinearForm>(m_feTheta.get()),
        std::make_unique<mfem::NonlinearForm>(m_feTheta.get())
    );
 
    // --- Add the integrators to the forms ---
    forms->M->AddDomainIntegrator(new mfem::MixedVectorWeakDivergenceIntegrator()); // M ∫∇ψ^θ·N^φ dV
    forms->Q->AddDomainIntegrator(new mfem::MixedVectorGradientIntegrator());       // Q ∫ψ^φ·∇N^θ dV
    forms->D->AddDomainIntegrator(new mfem::VectorFEMassIntegrator());              // D ∫ψ^φ·N^φ dV
    forms->S->AddDomainIntegrator(new mfem::DiffusionIntegrator());                 // S ∫∇ψ^θ·∇N^θ dV
    forms->f->AddDomainIntegrator(new polyMFEMUtils::NonlinearPowerIntegrator(m_polytropicIndex)); // Assumes polyMFEMUtils is now serif::polytrope::polyMFEMUtils
 
    // --- Assemble and Finalize the forms ---
    assembleAndFinalizeForm(forms->M);
    assembleAndFinalizeForm(forms->Q);
    assembleAndFinalizeForm(forms->D);
    assembleAndFinalizeForm(forms->S);
 
    // Note: The NonlinearForm does not need to / cannot be finalized, as it is not a matrix form. Rather, the operator
    //       will evaluate the nonlinear form during the solve phase.
 
 
    return forms;
}
 
void PolySolver::assembleAndFinalizeForm(auto &f) {
    // This constructs / ensures the matrix representation for each form
    // Assemble => Computes the local element matrices across the domain. Adds these to the global matrix . Adds these to the global matrix.
    // Finalize => Builds the sparsity pattern and allows the SparseMatrix representation to be extracted (CSR).
    f->Assemble();
    f->Finalize();
}
 
void PolySolver::solve() const {
    // --- Set the initial guess for the solution ---
    setInitialGuess();
    setOperatorEssentialTrueDofs();
 
    // --- Cast the GridFunctions to mfem::Vector ---
    const auto thetaVec = static_cast<mfem::Vector>(*m_theta); // NOLINT(*-slicing)
    const auto phiVec = static_cast<mfem::Vector>(*m_phi); // NOLINT(*-slicing)
 
    // --- Finalize the operator ---
    // Finalize with the initial state of theta for the initial jacobian calculation
    m_polytropOperator->finalize(thetaVec);
 
    // --- Broadcast initial condition to the full state vector ---
    const mfem::Array<int>& full_block_offsets = m_polytropOperator->get_block_offsets();
    mfem::Vector x_full(full_block_offsets.Last());
    mfem::BlockVector x_full_block(x_full, full_block_offsets);
    x_full_block.GetBlock(0) = thetaVec; // NOLINT(*-slicing)
    x_full_block.GetBlock(1) = phiVec; // NOLINT(*-slicing)
 
    // --- Extract only the free DOFs from the full state vector ---
    const mfem::Array<int>& freeDofs = m_polytropOperator->get_free_dofs();
    mfem::Vector x_free(m_polytropOperator->get_reduced_system_size());
    x_full.GetSubVector(freeDofs, x_free); // Extract the free DOFs from the full vector
 
    // --- Initialize RHS ---
    mfem::Vector zero_rhs(m_polytropOperator->get_reduced_system_size());
    zero_rhs = 0.0;
 
    // --- Setup and run the Newton solver ---
    const solverBundle sb = setupNewtonSolver();
    sb.newton.Mult(zero_rhs, x_free);
 
    // --- Reconstruct the full state vector from the reduced solution ---
    mfem::BlockVector solution = m_polytropOperator->reconstruct_full_block_state_vector(x_free);
 
    // --- Save and view an approximate 1D solution ---
    saveAndViewSolution(solution);
}
 
serif::types::MFEMArrayPairSet PolySolver::getEssentialTrueDof() const {
    mfem::Array<int> theta_ess_tdof_list;
    mfem::Array<int> phi_ess_tdof_list;
 
    mfem::Array<int> thetaCenterDofs, phiCenterDofs; // phiCenterDofs are not used
    mfem::Array<double> thetaCenterVals, phiCenterVals;
    std::tie(thetaCenterDofs, phiCenterDofs) = findCenterElement();
    thetaCenterVals.SetSize(thetaCenterDofs.Size());
    // phiCenterVals.SetSize(phiCenterDofs.Size());
    //
    // phiCenterVals = 0.0;
    thetaCenterVals = 1.0;
 
    mfem::Array<int> ess_brd(m_mesh.bdr_attributes.Max());
    ess_brd = 1;
 
    mfem::Array<double> thetaSurfaceVals, phiSurfaceVals;
    m_feTheta->GetEssentialTrueDofs(ess_brd, theta_ess_tdof_list);
    m_fePhi->GetEssentialTrueDofs(ess_brd, phi_ess_tdof_list);
 
    thetaSurfaceVals.SetSize(theta_ess_tdof_list.Size());
    thetaSurfaceVals = 0.0;
    phiSurfaceVals.SetSize(phi_ess_tdof_list.Size());
    phiSurfaceVals = polycoeff::thetaSurfaceFlux(m_polytropicIndex); // Assumes polycoeff is now serif::polytrope::polycoeff
 
    // combine the essential dofs with the center dofs
    theta_ess_tdof_list.Append(thetaCenterDofs);
    thetaSurfaceVals.Append(thetaCenterVals);
 
    // phi_ess_tdof_list.Append(phiCenterDofs);
    // phiSurfaceVals.Append(phiCenterVals);
 
    serif::types::MFEMArrayPair thetaPair = std::make_pair(theta_ess_tdof_list, thetaSurfaceVals);
    serif::types::MFEMArrayPair phiPair = std::make_pair(phi_ess_tdof_list, phiSurfaceVals);
    serif::types::MFEMArrayPairSet pairSet = std::make_pair(thetaPair, phiPair);
 
    return pairSet;
}
 
std::pair<mfem::Array<int>, mfem::Array<int>> PolySolver::findCenterElement() const {
    mfem::Array<int> thetaCenterDofs;
    mfem::Array<int> phiCenterDofs;
 
    // --- 1. Find the index of the single mesh vertex at the origin ---
    int center_vertex_idx = -1;
    constexpr double tol = 1e-9; // A small tolerance for floating point comparison
 
    for (int i = 0; i < m_mesh.GetNV(); ++i) {
        const double* vertex_coords = m_mesh.GetVertex(i);
        if (std::abs(vertex_coords[0]) < tol &&
            std::abs(vertex_coords[1]) < tol &&
            std::abs(vertex_coords[2]) < tol) {
 
            center_vertex_idx = i;
            break; // Found it, assume there's only one.
            }
    }
 
    if (center_vertex_idx == -1) {
        MFEM_ABORT("Could not find the center vertex at [0,0,0]. Check mesh construction.");
    }
 
    // --- 2. Get Theta (H1) DoFs associated ONLY with that vertex ---
    // CORRECTED: Use GetVertexDofs, not GetVDofs.
    m_feTheta->GetVertexDofs(center_vertex_idx, thetaCenterDofs);
 
 
    mfem::Array<int> central_element_ids;
 
    // PERF: could probably move this to a member variable and populate during construction
    mfem::Table* vertex_to_elements_table = m_mesh.GetVertexToElementTable();
    vertex_to_elements_table->Finalize();
    mfem::Array<int> element_ids;
    vertex_to_elements_table->GetRow(center_vertex_idx, element_ids);
    delete vertex_to_elements_table;
 
    for (int i = 0; i < element_ids.Size(); ++i) {
        int element_id = element_ids[i];
        mfem::Array<int> tempDofs;
        m_fePhi->GetElementDofs(element_id, tempDofs);
 
        // decode negative dofs to their true, physical, dof indices
        for (int j = 0; j < tempDofs.Size(); ++j) {
            int dof = tempDofs[j];
            if (dof < 0) {
                dof = -dof - 1; // Convert to positive index
            }
            phiCenterDofs.Append(dof);
        }
    }
 
 
    phiCenterDofs.Sort();
    phiCenterDofs.Unique();
 
    return std::make_pair(thetaCenterDofs, phiCenterDofs);
}
 
void PolySolver::setInitialGuess() const {
    // --- Set the initial guess for the solution ---
    mfem::FunctionCoefficient thetaInitGuess (
        [this](const mfem::Vector &x) {
            const double r = x.Norml2();
            // const double radius = Probe::getMeshRadius(*m_mesh);
            // const double u = 1/radius;
 
            // return (-1.0/radius) * r + 1;
            // return -std::pow((u*r), 2)+1.0; // The series expansion is a better guess; however, this is cheaper and ensures that the value at the surface is very close to zero in a way that the series expansion does not
            return laneEmden::thetaSeriesExpansion(r, m_polytropicIndex, 10);
        }
    );
 
    mfem::VectorFunctionCoefficient phiSurfaceVectors (m_mesh.SpaceDimension(),
        [this](const mfem::Vector &x, mfem::Vector &y) {
            const double r = x.Norml2();
            mfem::Vector xh(x);
            xh /= r; // Normalize the vector
            y.SetSize(m_mesh.SpaceDimension());
            y = xh;
            y *= polycoeff::thetaSurfaceFlux(m_polytropicIndex); // Assumes polycoeff is now serif::polytrope::polycoeff
        }
    );
    // We want to apply specific boundary conditions to the surface
    mfem::Array<int> ess_brd(m_mesh.bdr_attributes.Max());
    ess_brd = 1;
 
    // θ = 0 at surface
    mfem::ConstantCoefficient surfacePotential(0);
 
    m_theta->ProjectCoefficient(thetaInitGuess);
    m_theta->ProjectBdrCoefficient(surfacePotential, ess_brd);
 
    mfem::GradientGridFunctionCoefficient phiInitGuess (m_theta.get());
    m_phi->ProjectCoefficient(phiInitGuess);
 
    // Note that this will not result in perfect boundary conditions
    // because it must maintain H(div) continuity, this is
    // why inhomogenous boundary conditions enforcement is needed for φ
    // This manifests in PolytropeOperator::Mult where we do not
    // just zero out the essential dof elements in the residuals vector
    // for φ; rather, we need to set this to something which will push the
    // solver towards a more consistent answer (x_φ - target)
    m_phi->ProjectBdrCoefficientNormal(phiSurfaceVectors, ess_brd);
 
    auto [thetaCenterDofs, phiCenterDofs] = findCenterElement();
 
    for (int i = 0; i < phiCenterDofs.Size(); ++i)
    {
        (*m_phi)(phiCenterDofs[i]) = 0.0;
    }
 
    if (m_config.get<bool>("Poly:Solver:ViewInitialGuess", false)) {
        serif::probe::glVisView(*m_theta, m_mesh, "θ init");
        serif::probe::glVisView(*m_phi, m_mesh, "φ init");
    }
 
}
 
void PolySolver::saveAndViewSolution(const mfem::BlockVector& state_vector) const {
    mfem::BlockVector x_block(const_cast<mfem::BlockVector&>(state_vector), m_polytropOperator->get_block_offsets());
    mfem::Vector& x_theta = x_block.GetBlock(0);
    mfem::Vector& x_phi = x_block.GetBlock(1);
 
    if (m_config.get<bool>("Poly:Output:View", false)) {
        serif::probe::glVisView(x_theta, *m_feTheta, "θ Solution");
        serif::probe::glVisView(x_phi, *m_fePhi, "ɸ Solution");
    }
 
    // --- Extract the Solution ---
    if (m_config.get<bool>("Poly:Output:1D:Save", true)) {
        const auto solutionPath = m_config.get<std::string>("Poly:Output:1D:Path", "polytropeSolution_1D.csv");
        auto derivSolPath = "d" + solutionPath;
 
        const auto rayCoLatitude = m_config.get<double>("Poly:Output:1D:RayCoLatitude", 0.0);
        const auto rayLongitude = m_config.get<double>("Poly:Output:1D:RayLongitude", 0.0);
        const auto raySamples = m_config.get<int>("Poly:Output:1D:RaySamples", 100);
 
        const std::vector rayDirection = {rayCoLatitude, rayLongitude};
 
        serif::probe::getRaySolution(x_theta, *m_feTheta, rayDirection, raySamples, solutionPath);
        // Probe::getRaySolution(x_phi, *m_fePhi, rayDirection, raySamples, derivSolPath);
    }
}
 
void PolySolver::setOperatorEssentialTrueDofs() const {
    const serif::types::MFEMArrayPairSet ess_tdof_pair_set = getEssentialTrueDof();
    m_polytropOperator->set_essential_true_dofs(ess_tdof_pair_set);
}
 
void PolySolver::LoadSolverUserParams(double &newtonRelTol, double &newtonAbsTol, int &newtonMaxIter, int &newtonPrintLevel,
                                      double &gmresRelTol, double &gmresAbsTol, int &gmresMaxIter, int &gmresPrintLevel) const {
    newtonRelTol = m_config.get<double>("Poly:Solver:Newton:RelTol", 1.e-4);
    newtonAbsTol = m_config.get<double>("Poly:Solver:Newton:AbsTol", 1.e-6);
    newtonMaxIter = m_config.get<int>("Poly:Solver:Newton:MaxIter", 10);
    newtonPrintLevel = m_config.get<int>("Poly:Solver:Newton:PrintLevel", 3);
 
    gmresRelTol = m_config.get<double>("Poly:Solver:GMRES:RelTol", 1.e-12);
    gmresAbsTol = m_config.get<double>("Poly:Solver:GMRES:AbsTol", 1.e-12);
    gmresMaxIter = m_config.get<int>("Poly:Solver:GMRES:MaxIter", 200);
    gmresPrintLevel = m_config.get<int>("Poly:Solver:GMRES:PrintLevel", -1);
 
    LOG_DEBUG(m_logger, "Newton Solver (relTol: {:0.2E}, absTol: {:0.2E}, maxIter: {}, printLevel: {})", newtonRelTol, newtonAbsTol, newtonMaxIter, newtonPrintLevel);
    LOG_DEBUG(m_logger, "GMRES Solver (relTol: {:0.2E}, absTol: {:0.2E}, maxIter: {}, printLevel: {})", gmresRelTol, gmresAbsTol, gmresMaxIter, gmresPrintLevel);
}
 
void PolySolver::GetDofCoordinates(const mfem::FiniteElementSpace &fes, const std::string& filename) {
    mfem::Mesh *mesh = fes.GetMesh();
    double r = serif::probe::getMeshRadius(*mesh);
    std::ofstream outputFile(filename, std::ios::out | std::ios::trunc);
    outputFile << "dof,R,r,x,y,z" << '\n';
 
    const int nElements = mesh->GetNE();
 
    mfem::Vector coords;
    mfem::IntegrationPoint ipZero;
    double p[3] = {0.0, 0.0, 0.0};
    int actual_idx;
    ipZero.Set3(p);
    for (int i = 0; i < nElements; i++) {
        mfem::Array<int> elemDofs;
        fes.GetElementDofs(i, elemDofs);
        mfem::ElementTransformation* T = mesh->GetElementTransformation(i);
        mfem::Vector physCoord(3);
        T->Transform(ipZero, physCoord);
        for (int dofID = 0; dofID < elemDofs.Size(); dofID++) {
            if (elemDofs[dofID] < 0) {
                actual_idx = -elemDofs[dofID] - 1;
            } else {
                actual_idx = elemDofs[dofID];
            }
            outputFile << actual_idx;
            if (dofID != elemDofs.Size() - 1) {
                outputFile << "|";
            } else {
                outputFile << ",";
            }
        }
        outputFile << r << "," << physCoord.Norml2() << "," << physCoord[0] << "," << physCoord[1] << "," << physCoord[2] << '\n';
    }
    outputFile.close();
}
 
solverBundle PolySolver::setupNewtonSolver() const {
    // --- Load configuration parameters ---
    double newtonRelTol, newtonAbsTol, gmresRelTol, gmresAbsTol;
    int newtonMaxIter, newtonPrintLevel, gmresMaxIter, gmresPrintLevel;
    LoadSolverUserParams(newtonRelTol, newtonAbsTol, newtonMaxIter, newtonPrintLevel, gmresRelTol, gmresAbsTol,
                         gmresMaxIter, gmresPrintLevel);
 
    solverBundle solver; // Use this solver bundle to ensure lifetime safety
    solver.solver.SetRelTol(gmresRelTol);
    solver.solver.SetAbsTol(gmresAbsTol);
    solver.solver.SetMaxIter(gmresMaxIter);
    solver.solver.SetPrintLevel(gmresPrintLevel);
 
    // solver.solver.SetPreconditioner(m_polytropOperator->GetPreconditioner());
    // --- Set up the Newton solver ---
    solver.newton.SetRelTol(newtonRelTol);
    solver.newton.SetAbsTol(newtonAbsTol);
    solver.newton.SetMaxIter(newtonMaxIter);
    solver.newton.SetPrintLevel(newtonPrintLevel);
    solver.newton.SetOperator(*m_polytropOperator);
 
    // --- Created the linear solver which is used to invert the jacobian ---
    solver.newton.SetSolver(solver.solver);
 
    return solver;
}
 
} // namespace serif::polytrope