SERiF/src/poly/utils/public/polyMFEMUtils.h

#ifndef POLYMFEMUTILS_H
#define POLYMFEMUTILS_H

#include "mfem.hpp"
#include <string>


/**
 * @file polyMFEMUtils.h
 * @brief A collection of utilities for working with MFEM and solving the lane-emden equation.
 */


/**
 * @namespace polyMFEMUtils
 * @brief A namespace for utilities for working with MFEM and solving the lane-emden equation.
 */
namespace polyMFEMUtils {
  /**
   * @brief A class for nonlinear power integrator.
   */
  class NonlinearPowerIntegrator: public mfem::NonlinearFormIntegrator {
  private: 
    mfem::Coefficient &coeff_;
    double polytropicIndex;
  public:
    /**
     * @brief Constructor for NonlinearPowerIntegrator.
     * 
     * @param coeff The function coefficient.
     * @param n The polytropic index.
     */
    NonlinearPowerIntegrator(mfem::Coefficient &coeff, double n);

    /**
     * @brief Assembles the element vector.
     * 
     * @param el The finite element.
     * @param Trans The element transformation.
     * @param elfun The element function.
     * @param elvect The element vector to be assembled.
     */
    virtual void AssembleElementVector(const mfem::FiniteElement &el, mfem::ElementTransformation &Trans, const mfem::Vector &elfun, mfem::Vector &elvect) override;

    /**
     * @brief Assembles the element gradient.
     * 
     * @param el The finite element.
     * @param Trans The element transformation.
     * @param elfun The element function.
     * @param elmat The element matrix to be assembled.
     */
    virtual void AssembleElementGrad (const mfem::FiniteElement &el, mfem::ElementTransformation &Trans, const mfem::Vector &elfun, mfem::DenseMatrix &elmat) override;
  };

  /**
   * @brief A wrapper class for bilinear integrator.
   */
  class BilinearIntegratorWrapper : public mfem::NonlinearFormIntegrator
  {
  private:
    mfem::BilinearFormIntegrator *integrator;
  public:
    /**
      * @brief Constructor for BilinearIntegratorWrapper.
      * 
      * @param integratorInput The bilinear form integrator input.
      */
    BilinearIntegratorWrapper(mfem::BilinearFormIntegrator *integratorInput);

    /**
      * @brief Destructor for BilinearIntegratorWrapper.
      */
    virtual ~BilinearIntegratorWrapper();

    /**
      * @brief Assembles the element vector.
      * 
      * @param el The finite element.
      * @param Trans The element transformation.
      * @param elfun The element function.
      * @param elvect The element vector to be assembled.
      */
    virtual void AssembleElementVector(const mfem::FiniteElement &el, mfem::ElementTransformation &Trans, const mfem::Vector &elfun, mfem::Vector &elvect) override;

    /**
      * @brief Assembles the element gradient.
      * 
      * @param el The finite element.
      * @param Trans The element transformation.
      * @param elfun The element function.
      * @param elmat The element matrix to be assembled.
      */
    virtual void AssembleElementGrad(const mfem::FiniteElement &el, mfem::ElementTransformation &Trans, const mfem::Vector &elfun, mfem::DenseMatrix &elmat) override;
  };

  /**
   * @brief A class for composite nonlinear integrator.
   */
  class CompositeNonlinearIntegrator: public mfem::NonlinearFormIntegrator {
    private:
      std::vector<mfem::NonlinearFormIntegrator*> integrators;
    public:
      /**
       * @brief Constructor for CompositeNonlinearIntegrator.
       */
      CompositeNonlinearIntegrator();

      /**
       * @brief Destructor for CompositeNonlinearIntegrator.
       */
      virtual ~CompositeNonlinearIntegrator();

      /**
       * @brief Adds an integrator to the composite integrator.
       * 
       * @param integrator The nonlinear form integrator to add.
       */
      void add_integrator(mfem::NonlinearFormIntegrator *integrator);

      /**
       * @brief Assembles the element vector.
       * 
       * @param el The finite element.
       * @param Trans The element transformation.
       * @param elfun The element function.
       * @param elvect The element vector to be assembled.
       */
      virtual void AssembleElementVector(const mfem::FiniteElement &el, mfem::ElementTransformation &Trans, const mfem::Vector &elfun, mfem::Vector &elvect) override;

      /**
       * @brief Assembles the element gradient.
       * 
       * @param el The finite element.
       * @param Trans The element transformation.
       * @param elfun The element function.
       * @param elmat The element matrix to be assembled.
       */
      virtual void AssembleElementGrad(const mfem::FiniteElement &el, mfem::ElementTransformation &Trans, const mfem::Vector &elfun, mfem::DenseMatrix &elmat) override;
  };

  /**
   * @brief A class for constraint integrator.
   */
  class ConstraintIntegrator: public mfem::NonlinearFormIntegrator {
    private:
      mfem::Coefficient &eta;

    public:
      /**
       * @brief Constructor for ConstraintIntegrator.
       * 
       * @param eta The coefficient.
       */
      ConstraintIntegrator(mfem::Coefficient &eta_);

      /**
       * @brief Assembles the element vector.
       * 
       * @param el The finite element.
       * @param Trans The element transformation.
       * @param elfun The element function.
       * @param elvect The element vector to be assembled.
       */
      virtual void AssembleElementVector(const mfem::FiniteElement &el, mfem::ElementTransformation &Trans, const mfem::Vector &elfun, mfem::Vector &elvect) override;

      /**
       * @brief Assembles the element gradient.
       * 
       * @param el The finite element.
       * @param Trans The element transformation.
       * @param elfun The element function.
       * @param elmat The element matrix to be assembled.
       */
      virtual void AssembleElementGrad(const mfem::FiniteElement &el, mfem::ElementTransformation &Trans, const mfem::Vector &elfun, mfem::DenseMatrix &elmat) override;
  };

  class GaussianCoefficient : public mfem::Coefficient {
    private:
      double stdDev;
      double norm_coeff;
    public:
      GaussianCoefficient(double stdDev);
      virtual double Eval(mfem::ElementTransformation &T, const mfem::IntegrationPoint &ip) override;
  };

  class AugmentedOperator : public mfem::Operator {
    private:
      mfem::NonlinearForm &nfl;
      mfem::LinearForm &C;
      int lambdaDofOffset;
      mutable mfem::SparseMatrix *lastJacobian = nullptr;

    public:
      AugmentedOperator(mfem::NonlinearForm &nfl_, mfem::LinearForm &C_, int lambdaDofOffset_);
      ~AugmentedOperator();

      virtual void Mult(const mfem::Vector &x, mfem::Vector &y) const override;

      virtual mfem::Operator &GetGradient(const mfem::Vector &x) const override;
  };
} // namespace polyMFEMUtils

#endif // POLYMFEMUTILS_H