HMM_StochTrans_StefanPME.hpp

Go to the documentation of this file.

// Implementation of the HMM in 2D for the stochastic Stefan and PME models through the Doss-Sussman transform
//
//   { d_t u - div(K \grad(zeta(u))) = f,       inside Omega
//   { K \grad(zeta(u)) . nTF = g,       on GammaN
//   {                              u = g,           on GammaD
//   { u(0) = u_ini
//
//  At the moment, only pure Neumann or pure Dirichlet
//
// Author: Jerome Droniou (jerome.droniou@monash.edu)
//
 
#ifndef _HMM_STOCHTRANS_STEFANPME_HPP
#define _HMM_STOCHTRANS_STEFANPME_HPP
 
#include <functional>
#include <utility>
#include <iostream>
 
#include <boost/timer/timer.hpp>
 
// Matrices and linear solvers
#include <Eigen/Sparse>
#include <Eigen/Dense>
#include <boost/math/constants/constants.hpp>
#include <random>
 
#include "mesh.hpp"
#include "hybridcore.hpp"
#include "quad2d.hpp"
#include "TestCase/TestCaseNonLinearity.hpp"
#include "TestCase/TestCaseStefanPME.hpp"
#include "BoundaryConditions/BoundaryConditions.hpp"
 
#ifdef WITH_MKL
#include <Eigen/PardisoSupport>
#include <mkl.h>
#endif
 
namespace HArDCore2D {
 
// ----------------------------------------------------------------------------
//                Types and functions to define Wiener process
// ----------------------------------------------------------------------------
// Generic types
template <typename T>
using SpatialFunctionType = std::function<T(const VectorRd&)>; 
 
template <typename T>
using PiecewiseSpatialFunctionType = std::function<T(const VectorRd&, const Cell*)>; 
 
template <typename T>
using TemporalSpatialFunctionType = std::function<T(const double&, const VectorRd&)>; 
 
template <typename T>
using PiecewiseTemporalSpatialFunctionType = std::function<T(const double&, const VectorRd&, const Cell*)>; 
 
static SpatialFunctionType<double> zero_scalar_function = [](const VectorRd & x)->double { return 0.;}; 
 
template<typename T>
using EigFunctionType = std::function<T(const size_t&, const size_t&, const VectorRd&)>; 
 
template<typename T>
using DoubleVector = std::vector<std::vector<T>>;   
 
// Definition of eigenfunctions (depending on 2 indices)
static const double PI = boost::math::constants::pi<double>();
 
static std::function<double(const size_t &, const size_t &)> mui = [](const size_t &k, const size_t &l)->double
        {
          return 1./(std::pow(k, 2) + std::pow(l, 2));
        };
 
static EigFunctionType<double> mui_ei = [](const size_t &k, const size_t &l, const VectorRd & p)->double
        {
         return mui(k,l) * sin(k * PI * p.x()) * sin(l * PI * p.y()); 
        };
 
static EigFunctionType<VectorRd> grad_mui_ei = [](const size_t &k, const size_t &l, const VectorRd & p)->VectorRd
        {
         return mui(k,l) * VectorRd(
                                  k * PI * cos(k * PI * p.x()) * sin(l * PI * p.y()),
                                  l * PI * sin(k * PI * p.x()) * cos(l * PI * p.y())
                                  ); 
        };
 
static EigFunctionType<double> Delta_mui_ei = [](const size_t &k, const size_t &l, const VectorRd & p)->double
        {
         return - mui(k,l) * (std::pow(k, 2) + std::pow(l, 2)) * std::pow(PI,2) * sin(k * PI * p.x()) * sin(l * PI * p.y()); 
        };
 
 
// Compute standard deviation of a vector
static inline double StDev(const Eigen::VectorXd & v){
  return sqrt( (v.array()*v.array()).mean() - std::pow(v.mean(), 2) );
};
 
// ----------------------------------------------------------------------------
//                            Class definition
// ----------------------------------------------------------------------------
/* The StochStefanPME class provides an implementation of the HMM method for the stochastic Stefan and PME problems with Doss-Sussmann transformation.
*/
 
class StochStefanPME { 
 
public:
 
  // Types for solution
  using solution_function_type = TemporalSpatialFunctionType<double>;       
  using source_function_type = PiecewiseSpatialFunctionType<double>;        
  using grad_function_type = PiecewiseTemporalSpatialFunctionType<VectorRd>;        
  using lapl_function_type = PiecewiseTemporalSpatialFunctionType<double>;      
  using tensor_function_type = PiecewiseSpatialFunctionType<Eigen::Matrix2d>;       
 
  // types for Wiener process
  using WienerType = SpatialFunctionType<double>; 
  using GradWienerType = SpatialFunctionType<VectorRd>; 
  using ReactionType = SpatialFunctionType<double>; 
 
  StochStefanPME(
        HybridCore &hmm,                                        
        tensor_function_type kappa,     
        TestCaseNonLinearity::nonlinearity_function_type zeta,  
        BoundaryConditions BC,                  
        solution_function_type exact_solution,  
        grad_function_type grad_exact_solution,     
        solution_function_type time_der_exact_solution, 
        lapl_function_type minus_Lapl_exact_solution, 
        double weight,              
        std::string solver_type,        
    std::ostream & output = std::cout    
        );
  
  StochStefanPME::source_function_type compute_source(
            const bool & use_exact_source,    
            const double & t,       
            const WienerType & W,   
            const GradWienerType & grad_W, 
            const WienerType & Delta_exp_W, 
            const ReactionType & mu_squared  
            );
 
  std::pair<UVector, size_t> iterate(
                        const double tps,     
                        const double dt,    
                        const UVector& Xn,    
                        const UVector& Ih_W,    
                        const WienerType & W,   
                        const ReactionType & mu_squared, 
                        const source_function_type & source  
                        );
 
  void SetWienerProcess(
          const int & caseW,    
          const double & t,     
          WienerType & W,     
          UVector & I_W,      
          GradWienerType & grad_W,    
          WienerType & Delta_exp_W,   
          ReactionType & mu_squared,  
          bool verbose = false    
          );
                          
  void SimulateWienerProcess(
          const double & dt,     
          DoubleVector<double> & Wtime,    
          WienerType & W,     
          const DoubleVector<Eigen::VectorXd> & Ih_eigs,  
          UVector & I_W,      
          GradWienerType & grad_W,    
          WienerType & Delta_exp_W,   
          ReactionType & mu_squared  
          );
 
  Eigen::VectorXd apply_nonlinearity_eW(const std::string type, const Eigen::VectorXd & Y, const Eigen::VectorXd & I_W) const;
  UVector apply_nonlinearity_eW(const std::string type, const UVector& Y, const Eigen::VectorXd & I_W) const;
 
  double L2_MassLumped(const UVector& Xh) const; 
 
  double Lp_MassLumped(const UVector& Xh, double p) const; 
 
  double EnergyNorm(const UVector& Xh) const; 
 
  double Integral(const UVector& Xh) const; 
 
  inline double get_assembly_time() const {
    return double(_assembly_time) * pow(10, -9);
  }
  
  inline double get_solution_time() const {
    return double(_solution_time) * pow(10, -9);
  }
 
  inline double get_itime(size_t idx) const {
    return double(_itime[idx]) * pow(10, -9);
  }
  
  inline double get_solving_error() const {
    return _solving_error; 
  }
  
  inline Eigen::MatrixXd get_MassT(size_t iT) const {
    return MassT[iT];
  }
  
private:
  Eigen::MatrixXd diffusion_operator(const size_t iT) const;
 
  Eigen::VectorXd load_operator(const double tps, const size_t iT, const source_function_type & source) const;
 
    Eigen::VectorXd residual_scheme(const double dt, const UVector& Xh, const std::vector<Eigen::VectorXd> & RHS_nl, const UVector& Ih_W, const ReactionType & mu_squared) const;
 
    const HybridCore& m_hmm;
  const tensor_function_type kappa;
    const TestCaseNonLinearity::nonlinearity_function_type zeta;
    const BoundaryConditions m_BC;
  const solution_function_type exact_solution;
  const grad_function_type grad_exact_solution;
  const solution_function_type time_der_exact_solution;
  const lapl_function_type minus_Lapl_exact_solution;
    const double _weight;
    const std::string solver_type;
  std::ostream & m_output;
 
  // For random numbers
  std::random_device m_rd{};
  std::mt19937 m_gen{m_rd()};
 
    // To store local diffusion, mass and source terms
    std::vector<Eigen::MatrixXd> DiffT;
    std::vector<Eigen::MatrixXd> MassT;
 
  // Computation statistics
  size_t _assembly_time;
  size_t _solution_time;
  double _solving_error;
    std::vector<double> _itime = std::vector<double>(10, 0.);
 
};
 
// Constructor
StochStefanPME::StochStefanPME(HybridCore &hmm, tensor_function_type kappa, TestCaseNonLinearity::nonlinearity_function_type zeta, BoundaryConditions BC, solution_function_type exact_solution, grad_function_type grad_exact_solution, solution_function_type time_der_exact_solution, lapl_function_type minus_Lapl_exact_solution, double weight, std::string solver_type, std::ostream & output)
  : m_hmm(hmm),
        kappa(kappa),
        zeta(zeta),
        m_BC(BC),
        exact_solution(exact_solution),
        grad_exact_solution(grad_exact_solution),
        time_der_exact_solution(time_der_exact_solution),
        minus_Lapl_exact_solution(minus_Lapl_exact_solution),
        _weight(weight),
        solver_type(solver_type),
    m_output(output) {
 
        //---- Compute diffusion and mass matrix (do not change during time stepping or Newton) ------//
        const Mesh* mesh = m_hmm.get_mesh();
        DiffT.resize(mesh->n_cells());
        MassT.resize(mesh->n_cells());
        for (size_t iT = 0; iT < mesh->n_cells(); iT++) {
            size_t n_edgesT = mesh->cell(iT)->n_edges();
            size_t n_local_dofs = 1 + n_edgesT;
            double measT = mesh->cell(iT)->measure();
            DiffT[iT] = diffusion_operator(iT);
            MassT[iT] = Eigen::MatrixXd::Zero(n_local_dofs, n_local_dofs);
            MassT[iT](0,0) = (1-_weight) * measT;
            MassT[iT].bottomRightCorner(n_edgesT, n_edgesT) = _weight * (measT / n_edgesT) * Eigen::MatrixXd::Identity(n_edgesT, n_edgesT);
        }
}
 
 
//****************************************
//  Set the source
//****************************************
StochStefanPME::source_function_type StochStefanPME::compute_source(
              const bool & use_exact_source,
              const double & t,
              const WienerType & W,
              const GradWienerType & grad_W,
              const WienerType & Delta_exp_W,
              const ReactionType & mu_squared
              ) {
 
   StochStefanPME::source_function_type source_term;
  
   if (use_exact_source){
    // Exact source given by exact solution (latter is only valid for diffusion=Id)
    source_term = [&](const VectorRd & p, const Cell* cell)->double {
          double exp_W = exp(W(p));
          double exp_minus_W = 1./exp_W;
          VectorRd grad_exp_W = exp(W(p)) * grad_W(p);
          
            double val = time_der_exact_solution(t, p) 
                    - exp_minus_W * zeta(exp_W * exact_solution(t, p), "hess") * 
                          (
                          grad_exp_W * exact_solution(t, p) + exp_W * grad_exact_solution(t, p, cell)
                          ).squaredNorm() 
                    - exp_minus_W * zeta(exp_W * exact_solution(t, p), "der") * 
                          (
                          Delta_exp_W(p) * exact_solution(t, p) + 2 * grad_exp_W.dot(grad_exact_solution(t,p,cell))
                          - exp_W * minus_Lapl_exact_solution(t, p, cell)
                          )
                    + 0.5 * mu_squared(p) * exact_solution(t, p);
    
        return val;
    };
 
  }else{
    source_term = [&](const VectorRd & p, const Cell* cell)->double { return 0.; };                  
  }
  
  return source_term;
};
 
 
//****************************************
//  Perform one time iteration
//****************************************
 
std::pair<UVector, size_t> StochStefanPME::iterate(const double tps, const double dt, const UVector& Xn, const UVector& Ih_W, const WienerType & W, const ReactionType & mu_squared, const source_function_type & source) {
 
    boost::timer::cpu_timer timer; 
    boost::timer::cpu_timer timerint; 
  const auto mesh = m_hmm.get_mesh();
    size_t n_edges_dofs = mesh->n_edges();
    size_t n_cell_dofs = mesh->n_cells();
 
  //----- COMPUTE SOURCE: load + previous time -----//
  std::vector<Eigen::VectorXd> RHS_nl;
  for (size_t iT = 0; iT < mesh->n_cells(); iT++) {
    RHS_nl.push_back(dt * load_operator(tps, iT, source) + MassT[iT]*Xn.restr(iT));
  }
 
    //----- PARAMETERS FOR NEWTON ------------//
    constexpr double tol = 1e-10;
    constexpr size_t maxiter = 400;
  // relaxation parameter
  double relax = 1;     
    size_t iter = 0;
    UVector Xhprev = Xn;
    UVector Xh = UVector(Eigen::VectorXd::Zero(n_cell_dofs + n_edges_dofs), *m_hmm.get_mesh(), 0, 0);
 
    // Vector of Dirichlet BC, either on u (if weight>0) or zeta(e^W u) (if weight=0)
    Eigen::VectorXd DirBC = Eigen::VectorXd::Zero(n_cell_dofs + n_edges_dofs);
    for (size_t ibF = 0; ibF < mesh->n_b_edges(); ibF++){
    Edge* edge = mesh->b_edge(ibF);
    if (m_BC.type(*edge)=="dir"){
          size_t iF = edge->global_index();
          QuadratureRule quadF = generate_quadrature_rule(*edge, 5);
          for (QuadratureNode quadrule : quadF){
              if (_weight == 0){              
                  DirBC(n_cell_dofs + iF) += quadrule.w * zeta(exp(W(quadrule.vector())) * exact_solution(tps, quadrule.vector()), "fct");
              }else{
                  DirBC(n_cell_dofs + iF) += quadrule.w * exact_solution(tps, quadrule.vector());
              }
          }
          // Take average
          DirBC(n_cell_dofs + iF) /= mesh->b_edge(ibF)->measure();
    }
    }
    // Adjust initial iteration of Newton to match these BCs
    Xhprev.asVectorXd().tail(mesh->n_b_edges()) = DirBC.tail(mesh->n_b_edges());
 
    // ---- NEWTON ITERATONS ------ //
    Eigen::VectorXd RES = residual_scheme(dt, Xhprev, RHS_nl, Ih_W, mu_squared);        // Scheme is F(X)=C, then RES = C - F(u_prev)
  std::vector<double> residual(maxiter, 1.0);
    residual[iter] = RES.norm();
 
    while ( (iter==0) || ( (iter < maxiter) && (residual[iter] > tol) ) ){
        iter++;
 
    timerint.start();
        // System matrix and initialisation of RHS (only on the edge dofs)
        Eigen::SparseMatrix<double> GlobMat(n_edges_dofs, n_edges_dofs);
        std::vector<Eigen::Triplet<double>> triplets_GlobMat;
        Eigen::VectorXd GlobRHS = RES.tail(n_edges_dofs);
 
        // Static condensation: matrix and source to recover cell unknowns
        Eigen::SparseMatrix<double> ScMat(n_cell_dofs, n_edges_dofs);
        std::vector<Eigen::Triplet<double>> triplets_ScMat;
        Eigen::VectorXd ScRHS = Eigen::VectorXd::Zero(n_cell_dofs);
 
        // Assemble local contributions
        for (size_t iT = 0; iT < mesh->n_cells(); iT++) {
            Cell* T = mesh->cell(iT);
            size_t n_edgesT = T->n_edges();
            
            // exp(W) and exp(-W) in the cell
            Eigen::VectorXd WT = Ih_W.restr(iT);
            Eigen::MatrixXd exp_minus_WT = exp(-WT.array()).matrix().asDiagonal();
                    
            // Local static condensation of element unknowns
            Eigen::VectorXd XTprev = Xhprev.restr(iT);
            Eigen::VectorXd Zetaprime_expWT_XTprev = apply_nonlinearity_eW("der", XTprev, WT);
            Eigen::MatrixXd MatZetaprime_expWT_XTprev = Zetaprime_expWT_XTprev.asDiagonal();
            
            // Reaction term
            Eigen::MatrixXd ReacT = MassT[iT];
      ReacT(0,0) *= 0.5 * mu_squared(T->center_mass());
            for (size_t i=0; i<n_edgesT; i++){
              ReacT(i+1,i+1) *= 0.5 * mu_squared(T->edge(i)->center_mass());
            }
 
            Eigen::MatrixXd MatT = dt * exp_minus_WT * DiffT[iT] * MatZetaprime_expWT_XTprev + MassT[iT] + dt * ReacT;
            Eigen::MatrixXd ATT = MatT.topLeftCorner(1, 1);
            Eigen::MatrixXd ATF = MatT.topRightCorner(1, n_edgesT);
            Eigen::MatrixXd AFT = MatT.bottomLeftCorner(n_edgesT, 1);
            Eigen::MatrixXd AFF = MatT.bottomRightCorner(n_edgesT, n_edgesT);
 
            Eigen::PartialPivLU<Eigen::MatrixXd> invATT;
            invATT.compute(ATT);
                
            Eigen::MatrixXd invATT_ATF = invATT.solve(ATF);
            Eigen::VectorXd RES_cell = RES.segment(iT, 1);
            Eigen::VectorXd invATT_RES_cell = invATT.solve(RES_cell);
 
            // Local matrix and right-hand side on the face unknowns
            Eigen::MatrixXd MatF = Eigen::MatrixXd::Zero(n_edgesT, n_edgesT);
            Eigen::VectorXd bF = Eigen::VectorXd::Zero(n_edgesT);
            MatF = AFF - AFT * invATT_ATF;
            bF = - AFT * invATT_RES_cell;       // only additional RHS coming from static condensation, beyond RES
                            
            // Assemble static condensation operator
            ScRHS.segment(iT, 1) = invATT_RES_cell;
 
            for (size_t i = 0; i < 1; i++){
                size_t iGlobal = iT + i;
                for (size_t j = 0; j < n_edgesT; j++){
                    size_t jGlobal = T->edge(j)->global_index();
              triplets_ScMat.emplace_back(iGlobal, jGlobal, invATT_ATF(i, j));
                }
            }
 
          // GLOBAL MATRIX and RHS on the edge unknowns, using the matrices MatF obtained after static condensation
            for (size_t i = 0; i < n_edgesT; i++){
                size_t iGlobal = T->edge(i)->global_index();
                for (size_t j = 0; j < n_edgesT; j++){
                    size_t jGlobal = T->edge(j)->global_index();
                    triplets_GlobMat.emplace_back(iGlobal, jGlobal, MatF(i, j));
                }
            GlobRHS(iGlobal) += bF(i);
            }
        }
 
 
        // Assemble the global linear system (without BC), and matrix to recover statically-condensed cell dofs
        GlobMat.setFromTriplets(std::begin(triplets_GlobMat), std::end(triplets_GlobMat));
        ScMat.setFromTriplets(std::begin(triplets_ScMat), std::end(triplets_ScMat));
 
        //  Dirichlet boundary conditions are trivial. Newton requires to solve
        //          F'(X^k) (X^{k+1} - X^k) = B - F(X^k)
        //  Both X^k and X^{k+1} have fixed Dirichlet values on boundary edges, so the system we solve
    //  is Dirichlet homogeneous, and we just select the non-dirichlet edges (first ones in the list)
    size_t n_edge_unknowns = mesh->n_edges() - m_BC.n_dir_edges();
        Eigen::SparseMatrix<double> A = GlobMat.topLeftCorner(n_edge_unknowns, n_edge_unknowns);
//std::cout << "Number non-zero terms in matrix: " << A.nonZeros() << "\n";
        Eigen::VectorXd B = GlobRHS.head(n_edge_unknowns);
 
    timerint.stop();
    _assembly_time += double(timerint.elapsed().wall);
 
    timerint.start();
        // Solve condensed system and recover cell unknowns. dX = X^{k+1}-X^k//
        Eigen::VectorXd dX_edge_unknowns;
    if(solver_type == "pardiso"){
    #ifdef WITH_MKL
      Eigen::PardisoLU<Eigen::SparseMatrix<double>> solver;
      solver.compute(A);
      if (solver.info() != Eigen::Success) {
        std::cerr << "[main] ERROR: Could not factorize matrix" << std::endl;
      }
      dX_edge_unknowns = solver.solve(B);
      if (solver.info() != Eigen::Success) {
        std::cerr << "[main] ERROR: Could not solve linear system" << std::endl;
      }
  #endif
    }else{
      Eigen::BiCGSTAB<Eigen::SparseMatrix<double>, Eigen::IncompleteLUT<double> > solver;
      solver.compute(A);
      if (solver.info() != Eigen::Success) {
        std::cerr << "[main] ERROR: Could not factorize matrix" << std::endl;
        exit(1);
      }
      dX_edge_unknowns = solver.solve(B);
      if (solver.info() != Eigen::Success) {
        std::cerr << "[main] ERROR: Could not solve direct system" << std::endl;
        exit(1);
      }
    }
 
    // Recover all unknowns
        Eigen::VectorXd dX = Eigen::VectorXd::Zero(n_cell_dofs + n_edges_dofs);
        dX.segment(n_cell_dofs, n_edge_unknowns) = dX_edge_unknowns;
        dX.head(n_cell_dofs) = ScRHS - ScMat * dX.tail(n_edges_dofs);
 
        // Recover the fixed boundary values  and cell unknowns (from static condensation and Newton).
        // We start by putting the Dirichlet BC at the end.
        Xh.asVectorXd() = DirBC;
    Xh.asVectorXd().head(n_cell_dofs + n_edge_unknowns) = relax*dX.head(n_cell_dofs + n_edge_unknowns) + Xhprev.asVectorXd().head(n_cell_dofs + n_edge_unknowns);
 
    // Compute new residual. We might not accept new Xh but start over reducing relax. Otherwise, accept Xh and increase relax
    // to a max of 1
    Eigen::VectorXd RES_temp = residual_scheme(dt, Xh, RHS_nl, Ih_W, mu_squared);
    double residual_temp = RES_temp.norm();
 
    if (iter > 1 && residual_temp > 10*residual[iter-1]){
      // We don't update the residual and Xh, but we will do another iteration with relax reduced
      relax = relax/2.0;
      iter--;
    } else {
      // Increase relax
      relax = std::min(1.0, 1.2*relax);
        // Store Xh in Xhprev, compute new residual
        Xhprev.asVectorXd() = Xh.asVectorXd();
        RES = RES_temp;
        residual[iter] = residual_temp;
    }
    timerint.stop();
    _solution_time += double(timerint.elapsed().wall);
 
    // For a verbose output of Newton iterations:
//      m_output << "     ...Iteration: " << iter << "; residual = " << residual[iter] << "; relax = "<< relax << std::endl;
 
    } //-------- END NEWTON ITERATIONS -----------//
 
//  m_output << " (Newton it.: " << iter << ")" << std::endl;
//  m_output << "(" << iter << "), ";
 
  return std::make_pair(Xh, iter);
}
 
 
//*****************************
//  Set the Wiener process
//*****************************
 
void StochStefanPME::SetWienerProcess(const int & caseW, const double & t, WienerType & W, UVector & I_W, GradWienerType & grad_W, WienerType & Delta_exp_W, ReactionType & mu_squared, bool verbose) {
 
  switch(caseW){
  case 0:
    if(verbose) { std::cout << "[main] Wiener process = zero" << std::endl; }
      W = [&](const VectorRd & p)->double { return 0.; };
    grad_W = [&](const VectorRd & p)->VectorRd { return VectorRd::Zero(); };
      Delta_exp_W = [&](const VectorRd & p)->double { return 0.; };
      mu_squared = [](const VectorRd & p)->double { return 0.; };
      break;
    
    case 1:
      if(verbose) { std::cout << "[main] Wiener process = x + 2y" << std::endl; }
      W = [&](const VectorRd & p)->double { return p.x() + 2*p.y(); };
    grad_W = [&](const VectorRd & p)->VectorRd { return VectorRd(1., 2.); };
      Delta_exp_W = [&](const VectorRd & p)->double { return 5*exp(W(p)); };
      mu_squared = [](const VectorRd & p)->double { return 0.; };
    break;
  
  case 2:
      if(verbose) { std::cout << "[main] Wiener process = sin(pi x)sin(pi y)" << std::endl; }
      W = [&](const VectorRd & p)->double { return mui_ei(1,1,p); };
    grad_W = [&](const VectorRd & p)->VectorRd { 
                        return grad_mui_ei(1,1,p); 
                     };
      Delta_exp_W = [&](const VectorRd & p)->double { 
                          return exp(W(p)) * (Delta_mui_ei(1,1,p) + grad_W(p).squaredNorm()); 
                        };
      mu_squared = [](const VectorRd & p)->double { return std::pow(mui_ei(1,1,p), 2); };
    break;
  
  case 3:
      if(verbose) { std::cout << "[main] Wiener process = cos(t) * sin(pi x)sin(pi y)" << std::endl; }
      W = [&](const VectorRd & p)->double { return cos(t) * mui_ei(1,1,p); };
    grad_W = [&](const VectorRd & p)->VectorRd { 
                      return cos(t) * grad_mui_ei(1,1,p); 
                     };
      Delta_exp_W = [&](const VectorRd & p)->double { 
                          return exp(W(p)) * ( cos(t) * Delta_mui_ei(1,1,p) + grad_W(p).squaredNorm() ); 
                        };
      mu_squared = [](const VectorRd & p)->double { return std::pow(mui_ei(1,1,p), 2); };
    break;
 
  default:
    std::cout << "[main] Wiener process unknown" << std::endl;
    exit(1);
  }
 
  I_W.asVectorXd() = m_hmm.interpolate(W, 0, 0, 3).asVectorXd();
};
 
 
void StochStefanPME::SimulateWienerProcess(const double & dt, DoubleVector<double> & Wtime, WienerType & W, const DoubleVector<Eigen::VectorXd> & Ih_eigs, UVector & I_W, GradWienerType & grad_W, WienerType & Delta_exp_W, ReactionType & mu_squared) {
 
  size_t nb_eig = Wtime.size();
  
  // Simulate increment of Wtime
  std::normal_distribution<double> Nd{0, sqrt(dt)};
  for (size_t k=0; k<nb_eig; ++k){
    for (size_t l=0; l<nb_eig; ++l){
      Wtime[k][l] += Nd(m_gen);
    }
  }
  
  // Compute Wiener process
  W = [&, nb_eig](const VectorRd & p)->double { 
          double val = 0;
          for (size_t k=0; k<nb_eig; ++k){
            for (size_t l=0; l<nb_eig; ++l){
              val += Wtime[k][l] * mui_ei(k+1, l+1, p);
            } 
          } 
          return val; 
        };
  grad_W = [&, nb_eig](const VectorRd & p)->VectorRd { 
          VectorRd val = VectorRd::Zero();
          for (size_t k=0; k<nb_eig; ++k){
            for (size_t l=0; l<nb_eig; ++l){
              val += Wtime[k][l] * grad_mui_ei(k+1, l+1, p);
            } 
          } 
          return val; 
        };
  Delta_exp_W = [&, nb_eig](const VectorRd & p)->double { 
          double Delta_Wp = 0.;
          VectorRd grad_Wp = VectorRd::Zero();
          for (size_t k=0; k<nb_eig; ++k){
            for (size_t l=0; l<nb_eig; ++l){
                Delta_Wp +=  Wtime[k][l] * Delta_mui_ei(k+1, l+1, p);
                grad_Wp += Wtime[k][l] * grad_mui_ei(k+1, l+1, p);
              }          
            }            
          return exp(W(p)) * (Delta_Wp + grad_Wp.squaredNorm());
          };
  mu_squared = [&, nb_eig](const VectorRd & p)->double { 
          double val = 0.;
          for (size_t k=0; k<nb_eig; ++k){
            for (size_t l=0; l<nb_eig; ++l){
              val += std::pow(mui_ei(k+1, l+1, p), 2);
            }
          }
          return val;
        };
 
  // Compute interpolate 
  I_W.asVectorXd() = Eigen::VectorXd::Zero(I_W.asVectorXd().size());
  for (size_t k=0; k<nb_eig; ++k){
    for (size_t l=0; l<nb_eig; ++l){
      I_W.asVectorXd() += Wtime[k][l] * Ih_eigs[k][l];
    } 
  } 
 
};
 
 
//****************************************
//  Apply nonlinearity zeta to a vector
//****************************************
 
Eigen::VectorXd StochStefanPME::apply_nonlinearity_eW(const std::string type, const Eigen::VectorXd & Y, const Eigen::VectorXd & W) const {
    // Type="fct" or "der" as per the nonlinear function zeta
    // The function applied to each coefficient of Y depends on its nature. If it's a coefficient appearing in
    // the mass-lumping (which depends on _weight), the unknown represents u and we apply the full nonlinear function zeta after multiplying by e^W.
    // Otherwise, the unknown is zeta(e^W u) itself and the function we apply is the identity s->s (with nlin/der: s-> 1)
  // The derivative also takes into account the potential factor e^W that comes out when the unknown is u itself
  
    Eigen::VectorXd ZetaY;
    // Initialisation depends on "type". We initialise as if the nonlinearity was the identity, we'll change the
    // appropriate coefficients afterwards
    if (type == "fct"){
        ZetaY = Y;
    }else if (type == "der"){
        ZetaY = Eigen::VectorXd::Ones(Y.size());
    }else{
        m_output << "type unknown in apply_nonlinearity: " << type << "\n";
        exit(EXIT_FAILURE);
    }
 
    const Mesh* mesh = m_hmm.get_mesh();
    // We check if Y is a local vector or a global one, this determines the number of cell and edge unknowns
    size_t n_cell_dofs = 0;
    if (size_t(Y.size()) == mesh->n_cells() + mesh->n_edges()){
        n_cell_dofs = mesh->n_cells();
    }else{
        n_cell_dofs = 1;
    }
    size_t n_edges_dofs = Y.size() - n_cell_dofs;
 
  // Determine range of application of function (cell, all, or edge unknowns)
  size_t start = 0;
  size_t end = n_cell_dofs + n_edges_dofs;  
    if (_weight == 0){
      start = 0;
      end = n_cell_dofs;
    }else if (_weight == 1){
        start = n_cell_dofs;
        end = n_cell_dofs + n_edges_dofs;
    }
 
  // Apply non-linearity and correct for derivative
    for (size_t i = start; i < end; i++){
      ZetaY(i) = zeta(exp(W(i)) * Y(i), type);
      if (type == "der"){
        ZetaY(i) *= exp(W(i));
      }
    }
 
    return ZetaY;
}
 
// Overloaded version: apply to the values of UVector, don't change the other parameters
UVector StochStefanPME::apply_nonlinearity_eW(const std::string type, const UVector & Y, const Eigen::VectorXd & W) const {
 
  UVector ZetaY = Y;
  ZetaY.asVectorXd() = apply_nonlinearity_eW(type, Y.asVectorXd(), W);
  return ZetaY;
}
 
 
//*****************************************
//  Local diffusion matrix and load term
//*****************************************
 
Eigen::MatrixXd StochStefanPME::diffusion_operator(const size_t iT) const {
 
  const auto mesh = m_hmm.get_mesh();
  size_t dim = mesh->dim();
    Cell* T = mesh->cell(iT);
  double mT = T->measure();
    const size_t n_edgesT = T->n_edges();
 
  size_t local_dofs = 1 + n_edgesT;
  Eigen::MatrixXd StiffT = Eigen::MatrixXd::Zero(local_dofs, local_dofs);
    
    // Diffusion in the cell.
  Eigen::Vector2d xT = T->center_mass();
  Eigen::Matrix2d kappaT = kappa(xT, T);
 
  // Matrix of \nabla_K
  Eigen::MatrixXd nablaK = Eigen::MatrixXd::Zero(dim, local_dofs);
  for (size_t iE=0; iE < n_edgesT; iE++){
    Edge* E = T->edge(iE);
    double mE = E->measure();
    nablaK.block(0, iE+1, dim, 1) = mE * T->edge_normal(iE);
  }
  nablaK /= mT;
 
  // Contribution per diamond
  for (size_t iE=0; iE < n_edgesT; iE++){
    Edge* E = T->edge(iE);
    Eigen::Vector2d xE = E->center_mass();
    Eigen::Vector2d nTE = T->edge_normal(iE);
    double mE = E->measure();
    // Distance center to edge and measure of diamond
    double dTE =  (xE-xT).dot(nTE);
    double mTE = mE * dTE / dim;
 
    // stabilisation (without scaling or normal: u_E-u_K-nablaK u . (xE-xK))
    Eigen::MatrixXd stabE = Eigen::MatrixXd::Zero(1, local_dofs);
    stabE(0) = -1.0;
    stabE(1+iE) = 1.0;
    stabE -= (xE-xT).transpose() * nablaK;
 
    // Complete gradient with stabilisation
    Eigen::MatrixXd GRAD = nablaK + (pow(dim, 0.5)/dTE) * nTE * stabE;
 
    // Contribution to stiffness
    StiffT += mTE * GRAD.transpose() * kappaT * GRAD;
  }
 
    return StiffT; 
}
 
 
Eigen::VectorXd StochStefanPME::load_operator(const double tps, const size_t iT, const source_function_type & source) const {
 
  const auto mesh = m_hmm.get_mesh();
    Cell* T = mesh->cell(iT);
    size_t n_edgesT = T->n_edges();
    size_t local_dofs = 1 + n_edgesT;
 
    // Beware, MassT[iT] must have been created
    if (MassT.size() < iT || size_t(MassT[iT].rows()) != local_dofs){
        m_output << "Called load_operator without creating MassT[iT]\n";
        exit(EXIT_FAILURE);
    }
    Eigen::VectorXd fvec = Eigen::VectorXd::Zero(1 + n_edgesT);
    for (size_t i = 0; i < local_dofs; i++){
        VectorRd node;
        if (i == 0){
            node = T->center_mass();
        }else{
            node = T->edge(i-1)->center_mass();
        }
        fvec(i) = source(node, T);
    }
  // Contribution of source to loading term
  Eigen::VectorXd load = MassT[iT] * fvec;
 
  // Adding Neumann boundary conditions
  // The boundary conditions coming from Neumann edges are computed based on the piecewise constant
  // discrete trace: if g is the value of the outer flux,
  //      \int g v -> \int g Tv
  //  where (Tv)_e = average of v on e, for e Neumann edge.
  // Only edge-based basis functions therefore have a contribution, which will simply be \int_e g.
  if (T->is_boundary()){
    for (size_t ilF = 0; ilF < T->n_edges(); ilF++) {
      Edge* edge = T->edge(ilF);
      if (m_BC.type(*edge)=="neu"){
                if (edge->is_boundary()){
                  const auto& nTF = T->edge_normal(ilF);
                    QuadratureRule quadF = generate_quadrature_rule(*edge, 5);
                    std::function<double(VectorRd)> Kgrad_n = [&](VectorRd p){
                                return zeta(exact_solution(tps, p), "der") * nTF.dot(kappa(p,T) * grad_exact_solution(tps, p, T));
                        };
                    for (QuadratureNode qr : quadF){
                        load(DimPoly<Cell>(1) + ilF) += qr.w * Kgrad_n(qr.vector());
                    }
                }
            }
        }
    }
 
  return load;
}
 
//***********************************************************************
// Residual non-linear scheme: if scheme is F(X)=C, residual is C-F(X)
//***********************************************************************
 
Eigen::VectorXd StochStefanPME::residual_scheme(const double dt, const UVector& Xh, const std::vector<Eigen::VectorXd> & RHS_nl, const UVector& Ih_W, const ReactionType & mu_squared) const{
 
    const Mesh* mesh = m_hmm.get_mesh();    
    Eigen::VectorXd RES = Eigen::VectorXd::Zero(mesh->n_cells() + mesh->n_edges());
 
    // Compute the residual: C - (dt * e^(-W) * DIFF * Zeta(e^W Xh) + MASS * Xh), where DIFF is the diffusion matrix, MASS is the mass matrix and C is the source term
    size_t n_cell_dofs = mesh->n_cells();
    for (size_t iT = 0; iT < mesh->n_cells(); iT++){
        Cell* T =  mesh->cell(iT);
        size_t n_edgesT = T->n_edges();
 
        // exp(W) and exp(-W) in the cell
        Eigen::VectorXd WT = Ih_W.restr(iT);
        Eigen::MatrixXd exp_minus_WT = exp(-WT.array()).matrix().asDiagonal();
 
        // Local unknown in the cell, and zeta of these unknowns
        Eigen::VectorXd XT = Xh.restr(iT);
        Eigen::VectorXd Zeta_expWT_XT = apply_nonlinearity_eW("fct", XT, WT);
 
        // Reaction term
        Eigen::MatrixXd ReacT = MassT[iT];
    ReacT(0,0) *= 0.5 * mu_squared(T->center_mass());
        for (size_t i=0; i<n_edgesT; i++){
          ReacT(i+1,i+1) *= 0.5 * mu_squared(T->edge(i)->center_mass());
        }
 
        Eigen::VectorXd REST = RHS_nl[iT] - (dt * exp_minus_WT * DiffT[iT] * Zeta_expWT_XT + (MassT[iT] + dt * ReacT) * XT);
        RES(iT) += REST(0);
        for (size_t ilE = 0; ilE < n_edgesT; ilE++){
            size_t iE = T->edge(ilE)->global_index();
            RES(n_cell_dofs + iE) += REST(1+ilE);
        }
    }
 
    // No residual on Dirichlet edges (do not correspond to equations of the system)
    RES.tail(m_BC.n_dir_edges()) = Eigen::VectorXd::Zero(m_BC.n_dir_edges());
 
    return RES;
}
 
//*******************************************************
// Norms: L2 mass lumped, Energy (pure diffusion)
//*******************************************************
 
double StochStefanPME::L2_MassLumped(const UVector& Xh) const {
  const Mesh* mesh = m_hmm.get_mesh();
    double value = 0.0;
 
  for (size_t iT = 0; iT < mesh-> n_cells(); iT++) {
        Eigen::VectorXd XTF = Xh.restr(iT);
        value += XTF.transpose() * MassT[iT] * XTF;
  }
 
  return sqrt(value);
}
 
 
double StochStefanPME::Lp_MassLumped(const UVector& Xh, const double p) const {
  const Mesh* mesh = m_hmm.get_mesh();
    double value = 0.0;
 
  for (size_t iT = 0; iT < mesh-> n_cells(); iT++) {
        Eigen::VectorXd XTF = Xh.restr(iT);
    // Coefficient-wise power p
    Eigen::VectorXd XTF_powerp = (XTF.array().abs()).pow(p);
    // Sum local contributions; this assumes that MassT is diagonal
        value += (MassT[iT] * XTF_powerp.matrix() ).sum();
  }
 
  return std::pow(value, 1.0/p);
}
 
double StochStefanPME::EnergyNorm(const UVector& Xh) const {
  const Mesh* mesh = m_hmm.get_mesh();
    double value = 0.0;
 
  for (size_t iT = 0; iT < mesh-> n_cells(); iT++) {
        Eigen::VectorXd XTF = Xh.restr(iT);
        value += XTF.transpose() * DiffT[iT] * XTF;
  }
 
  return sqrt(value);
}
 
double StochStefanPME::Integral(const UVector& Xh) const {
  const Mesh* mesh = m_hmm.get_mesh();
    double value = 0.0;
 
  for (size_t iT = 0; iT < mesh-> n_cells(); iT++) {
        value += (MassT[iT] * Xh.restr(iT)).sum();
  }
 
  return value;
}
 
 
 
 
} // end of namespace HArDCore2D
 
#endif //_HMM_STOCHTRANS_STEFANPME_HPP