SERiF/src/poly/solver/private/polySolver.cpp

#include "mfem.hpp"

#include <string>
#include <iostream>
#include <memory>

#include "meshIO.h"
#include "polySolver.h"
#include "polyMFEMUtils.h"
#include "polyCoeff.h"


// TODO: Come back to this and think of a better way to get the mesh file
const std::string SPHERICAL_MESH = std::string(getenv("MESON_SOURCE_ROOT")) + "/src/resources/mesh/sphere.msh";

PolySolver::PolySolver(double n, double order)
: n(n),
  order(order),
  meshIO(SPHERICAL_MESH),
  mesh(meshIO.GetMesh()),
  feCollection(std::make_unique<mfem::H1_FECollection>(order, mesh.SpaceDimension())),
  feSpace(std::make_unique<mfem::FiniteElementSpace>(&mesh, feCollection.get())),
  lambdaFeSpace(std::make_unique<mfem::FiniteElementSpace>(&mesh, feCollection.get(), 1)), // Scalar space for lambda
  compositeIntegrator(std::make_unique<polyMFEMUtils::CompositeNonlinearIntegrator>()),
  nonlinearForm(std::make_unique<mfem::NonlinearForm>(feSpace.get())),
  C(std::make_unique<mfem::LinearForm>(feSpace.get())),
  u(std::make_unique<mfem::GridFunction>(feSpace.get())),
  diffusionCoeff(std::make_unique<mfem::VectorConstantCoefficient>([&](){
        mfem::Vector diffusionCoeffVec(mesh.SpaceDimension());
        diffusionCoeffVec = 1.0;
        return diffusionCoeffVec;
    }())),
  nonLinearSourceCoeff(std::make_unique<mfem::ConstantCoefficient>(1.0)),
  gaussianCoeff(std::make_unique<polyMFEMUtils::GaussianCoefficient>(0.1)) {
    assembleNonlinearForm();
    assembleConstraintForm();
}

PolySolver::~PolySolver() {}

void PolySolver::assembleNonlinearForm() {
    // Add the \int_{\Omega}\nabla v\cdot\nabla\theta d\Omegaterm
    auto wrappedDiffusionIntegrator = std::make_unique<polyMFEMUtils::BilinearIntegratorWrapper>(
        new mfem::DiffusionIntegrator(*diffusionCoeff)
    );
    compositeIntegrator->add_integrator(wrappedDiffusionIntegrator.release());

    // Add the \int_{\Omega}v\theta^{n} d\Omega term
    auto nonLinearIntegrator = std::make_unique<polyMFEMUtils::NonlinearPowerIntegrator>(*nonLinearSourceCoeff, n);
    compositeIntegrator->add_integrator(nonLinearIntegrator.release());

    // Add the \int_{\Omega}v\eta(r) d\Omega term
    auto constraintIntegrator = std::make_unique<polyMFEMUtils::ConstraintIntegrator>(*gaussianCoeff);
    compositeIntegrator->add_integrator(constraintIntegrator.release());

    nonlinearForm->AddDomainIntegrator(compositeIntegrator.release());
}

void PolySolver::assembleConstraintForm() {
    auto constraintIntegrator = std::make_unique<mfem::DomainLFIntegrator>(*gaussianCoeff);
    C->AddDomainIntegrator(constraintIntegrator.release());
    C->Assemble();
}

void PolySolver::solve(){
    // --- Set the initial guess for the solution ---
    mfem::FunctionCoefficient initCoeff (
        [this](const mfem::Vector &x) {
            return 1.0; // Update this to be a better init guess
        }
    );
    u->ProjectCoefficient(initCoeff);

    // --- Combine DOFs (u and λ) into a single vector ---
    int lambdaDofOffset  = feSpace->GetTrueVSize(); // Get the size of θ space
    int totalTrueDofs = lambdaDofOffset + lambdaFeSpace->GetTrueVSize();

    std::cout << "Total True Dofs: " << totalTrueDofs << std::endl;
    std::cout << "Lambda Dof Offset: " << lambdaDofOffset << std::endl;
    std::cout << "Lambda True Dofs: " << lambdaFeSpace->GetTrueVSize() << std::endl;
    std::cout << "U True Dofs: " << feSpace->GetTrueVSize() << std::endl;
    if (totalTrueDofs != lambdaDofOffset + 1) {
        throw std::runtime_error("The total number of true dofs is not equal to the sum of the lambda offset and the lambda space size");
    }

    mfem::Vector U(totalTrueDofs);
    U = 0.0;

    mfem::Vector u_view(U.GetData(), lambdaDofOffset);
    u->GetTrueDofs(u_view);

    // --- Setup the Newton Solver ---
    mfem::NewtonSolver newtonSolver;
    newtonSolver.SetRelTol(1e-8);
    newtonSolver.SetAbsTol(1e-10);
    newtonSolver.SetMaxIter(200);
    newtonSolver.SetPrintLevel(1);

    // --- Setup the GMRES Solver ---
    // --- GMRES is good for indefinite systems ---
    mfem::GMRESSolver gmresSolver;
    gmresSolver.SetRelTol(1e-10);
    gmresSolver.SetAbsTol(1e-12);
    gmresSolver.SetMaxIter(2000);
    gmresSolver.SetPrintLevel(0);
    newtonSolver.SetSolver(gmresSolver);
    // TODO: Change numeric tolerance to grab from the tol module

    // --- Setup the Augmented Operator ---
    polyMFEMUtils::AugmentedOperator aug_op(*nonlinearForm, *C, lambdaDofOffset);
    newtonSolver.SetOperator(aug_op);

    // --- Create the RHS of the augmented system ---
    mfem::Vector B(totalTrueDofs);
    B = 0.0;

    // Set the constraint value (∫η(r) dΩ) in the last entry of B
    // This sets the last entry to 1.0, this mighht be a problem later on...
    mfem::ConstantCoefficient one(1.0);
    mfem::LinearForm constraint_rhs(lambdaFeSpace.get());
    constraint_rhs.AddDomainIntegrator(
        new mfem::DomainLFIntegrator(*gaussianCoeff)
    );

    constraint_rhs.Assemble();
    B[lambdaDofOffset] = constraint_rhs(0); // Get that single value for the rhs. Only one value because it's a scalar space


    // --- Solve the augmented system ---
    newtonSolver.Mult(B, U);

    // --- Extract the Solution ---
    mfem::Vector u_sol_view(U.GetData(), lambdaDofOffset);
    u->SetData(u_sol_view);
    double lambda = U[lambdaDofOffset];

    std::cout << "λ = " << lambda << std::endl;
    // TODO : Add a way to get the solution out of the solver
}