//ff-mpirun -np 4 elasticity_saddlepoint.edp -wg -ns -ffddm_schwarz_coarse_correction BNN -ffddm_geneo_threshold_A 0.1 -ffddm_geneo_threshold_S 0.3 -ffddm_geneo_nu 40 -Dmdim=2 -global 10 -nlayers 10

/* 3D heterogenous elasticity problem with nearly incompressible
   material solved with the GenEO saddle point solver from
   Nataf, F., & Tournier, P. H. (2021). A GenEO Domain Decomposition
   method for Saddle Point problems.
   arXiv preprint available at https://arxiv.org/abs/1911.01858 */

// the command line parameter '-Dmdim' sets the value of the macro 'mdim' in the script;
// '-Dmdim=2' or '-Dmdim=3' switches between the 2D and 3D test case

include "getARGV.idp"
int global = getARGV("-global", 10); // discretization parameter
int nlayers = getARGV("-nlayers", global); // number of alternating material layers in the y direction

// material coefficients:
real nu1 = 0.4999;
real nu2 = 0.35;
real E1 = 1e+7;
real E2 = 200e+9;

// heterogeneous material distribution:
func fE = int(y*nlayers)%2 == 0 ?  E1 :  E2;
func fnu = int(y*nlayers)%2 == 0 ? nu1 : nu2;

real gtol = 1e-5;

/********* 2D *********/
IFMACRO(mdim,2)
macro dimension 2// EOM
include "ffddm.idp"

macro grad(u) [dx(u),dy(u)] // EOM
real Sqrt = sqrt(2.);
macro epsilon(u)[dx(u), dy(u#B), (dy(u) + dx(u#B)) / Sqrt]// EOM
macro div(u)(dx(u) + dy(u#B))// EOM

mesh ThGlobal = square(5*global, global, [5*x,y]);

func Pku = [P2,P2];
func Pkp = [P0];

macro VarfA(varfName, meshName, PhName)
fespace Ph0(meshName,P0);
Ph0  E = fE, nu = fnu;
func mu = E/(2*(1+nu));
func lambda = E*nu/((1+nu)*(1-2*nu));
varf varfName([u, uB], [v, vB]) = int2d(meshName)(2.0 * mu * (epsilon(u)' * epsilon(v)))
                                 +on(4, u = 0, uB = 0)
                                 -int2d(meshName)(100*vB);
 // EOM

macro VarfB(varfName, meshName, PhName)
varf varfName([u, uB], [q]) = int2d(meshName)(div(u) * q);
// EOM

macro VarfC(varfName, meshName, PhName)
fespace Ph0(meshName,P0);
Ph0  E = fE, nu = fnu;
func mu = E/(2*(1+nu));
func lambda = E*nu/((1+nu)*(1-2*nu));

varf varfName([p], [q]) = int2d(meshName)(1./lambda * p * q);
// EOM

macro defA(i)[i, i#B] // EOM
macro initA(i)[i, i] // EOM

macro defB(i)[i] // EOM
macro initB(i)[i] // EOM
ENDIFMACRO

/********* 3D *********/
IFMACRO(mdim,3)
macro dimension 3// EOM
include "ffddm.idp"

real Sqrt = sqrt(2.);
macro epsilon(u)[dx(u), dy(u#B), dz(u#C), (dz(u#B) + dy(u#C)) / Sqrt, (dz(u) + dx(u#C)) / Sqrt, (dy(u) + dx(u#B)) / Sqrt]// EOM
macro div(u)(dx(u) + dy(u#B) + dz(u#C))// EOM

int[int] LL = [2,3, 2,1, 2,2];
mesh3 ThGlobal = cube(5*global, global, global, [5*x,y,z], label = LL);

func Pku = [P2,P2,P2];
func Pkp = [P1];

macro VarfA(varfName, meshName, PhName)
fespace Ph0(meshName,P0);
Ph0  E = fE, nu = fnu;
func mu = E/(2*(1+nu));
func lambda = E*nu/((1+nu)*(1-2*nu));
varf varfName([u, uB, uC], [v, vB, vC]) = int3d(meshName)(2.0 * mu * (epsilon(u)' * epsilon(v)))
                                 +on(1, u = 0, uB = 0, uC = 0)
                                 -int3d(meshName)(100*vB);
// EOM

macro VarfB(varfName, meshName, PhName)
varf varfName([u, uB, uC], [q]) = int3d(meshName)(div(u) * q);
// EOM

macro VarfC(varfName, meshName, PhName)
fespace Ph0(meshName,P0);
Ph0  E = fE, nu = fnu;
func mu = E/(2*(1+nu));
func lambda = E*nu/((1+nu)*(1-2*nu));

varf varfName([p], [q]) = int3d(meshName)(1./lambda * p * q);
// EOM

macro defA(i)[i, i#B, i#C] // EOM
macro initA(i)[i, i, i] // EOM

macro defB(i)[i] // EOM
macro initB(i)[i] // EOM
ENDIFMACRO

/* build mesh decomposition with minimal overlap size 1*2 (prefix M) and also the Augmented decomposition
   with overlap size 2*2 (prefix MAug) */
ffddmbuildDmeshAug(M,ThGlobal,mpiCommWorld)

/* build Dfespace for A on distributed mesh M (prefix AFE) ; this also builds Dfespace tilde with larger overlap
   on distributed mesh MAug (prefix AFEAug) needed for the rectangular operator B on tilde decomposition.
   This also builds 'AFEAuginterp' matrices (restriction matrices from AFEAugVhi to AFEVhi local FE spaces). */
ffddmbuildDfespace(AFE,M,real,defA,initA,Pku)

ffddmbuildDfespace(CFEAug,MAug,real,defB,initB,Pkp) // Dfespace for C

vsym=1;
vtgvelim=1e+30;
vtgv=1e+30;

ffddmsetupOperator(A,AFE,VarfA); /* Operator A */
ffddmsetupPrecond(A,null); /* one level preconditioner for A */

ffddmtau = getARGV("-ffddm_geneo_threshold_A", 0.3);
ffddmgeneosetup(A,VarfA); /* second level for A */

Acorr == "ADEF1";

ffddmsetupOperatorRect(Btilde,AFEAug,CFEAug,VarfB) /* Operator B on tilde decomposition */

ffddmsetupOperator(Ctilde,CFEAug,VarfC) /* Operator C on tilde decomposition */

matrix Ctildei;
{
  VarfC(vC, MAugThi, null)
  Ctildei = vC(CFEAugVhi,CFEAugVhi,sym=1); /* Ctilde_i from Remark 3.2 */
}

/* define and build all necessary operators for the application of N_S^-1 */
ffddmtau = getARGV("-ffddm_geneo_threshold_S", 0.1);
ffddmSPsetup(SP,A,Btilde,Ctilde,Ctildei)

AFEVhi defA(Fx);
ffddmbuildrhs(A,VarfA,Fx[]) /* Build F_u */

CFEAugVhi [Fp] = [0]; /* Build F_p */

AFEVhi defA(Gx);

Averbosity = 1;
/* Solve A G_u = F_u with M_A^-1 as a preconditioner, step 1 of Algorithm 2 */
Gx[] = AfGMRES(Gx[], Fx[], gtol, 100, "right");

/* Compute G_p = F_p - B G_u, step 2 of Algorithm 2 */
CFEAugVhi [Gp];
Gp[] = Fp[] - SPBmatvec(Gx[]);

/* Step 3 of Algorithm 2: solve (C + B A^-1 B^T) P = -G_p with N_S^-1 as a preconditioner */
Gp[] *= -1;

CFEAugVhi [P];
P[] = 0;
P[] = SPfGMRES(P[], Gp[], gtol, 100, "right");

Awritesummary
SPwritesummary

ffddmplot(CFEAug,P,"P");

// Step 4 of Algorithm 2: compute G_u = F_u - B^T P
Gx[] = Fx[] - SPBTmatvec(P[]);

// Step 5 of Algorithm 2: solve A U = G_u with M_A^-1 as a preconditioner
AFEVhi defA(Ux);
Ux[] = 0;
Ux[] = AfGMRES(Ux[], Gx[], gtol, 100, "right");

ffddmplot(CFEAug,UxB,"Uy");

IFMACRO(BLOCK)
/* Block solve with preconditioner P of eq. (28) */
real[int] rhs(AFEVhi.ndof+CFEAugVhi.ndof);
rhs(0:AFEVhi.ndof-1) = Fx[];
rhs(AFEVhi.ndof:rhs.n-1) = Fp[];

real[int] x0(rhs.n), u(rhs.n);

u = SPBlockfGMRES(x0, rhs, gtol, 100, "right");

SPBlockwritesummary

real[int] Us = Ux[];

Ux[] = u(0:AFEVhi.ndof-1);

ffddmplot(CFEAug,UxB,"UxB block");

Ux[] -= Us;

ffddmplot(CFEAug,UxB,"difference");
ENDIFMACRO
