// model for mutation-selection, Magal & Webb 2000
// theta-IMEX scheme

// domain + triangulation
int n =50; // size
mesh Th = square(n,n,flags=1);
// FE space, Lagrange-P2
fespace Vh(Th,P2);
Vh u,v, uold,uold1;

// define operators: Gradient
macro Grad(u) [ dx(u),dy(u) ] //

// model parameters
// drug dose
real m = 0.02;
// function beta(x,y)
func bb = .03 + .25*y*(1-y^2)+.05*(1-x)*x;
func dd = m*(.15 + 1.3*(1-x)) + .02;

// parameters
real alpha = 1e-5;
real I0;
// integration
real dt = 5.;

// computing stiffness matrix & load vector

// variational formulations
varf a1(u,v) = int2d(Th)(u*v/dt);
varf a2(u,v)= int2d(Th)(alpha*( Grad(u)'*Grad(v) ));
varf ll1(u,v) = int2d(Th) (bb *u*v) ;
varf ll2(u,v) = int2d(Th) (dd *u*v) ;
// bilinear part
matrix A1 =a1(Vh,Vh),A2=a2(Vh,Vh);
// linear part
matrix L1 = ll1(Vh,Vh), L2 = ll2(Vh,Vh);


Vh rhs,rhs0 ,rhs1, err;
real theta =2./3., tol = 1e-4, dtol=1;
// initial data
uold = .5*sin(5.*pi*x)*sin(5.*pi*y); uold=abs(uold); // u{1}

matrix A = A1+theta*A2;
set (A,init=1,solver=sparsesolver);

// loop for convergence to stationary solution
while (dtol > tol)
{
	I0 = int2d(Th)(uold); // compute integral
//	cout << "Int(u) = "<<I0<< endl;
	rhs[] = A1*uold[];
	rhs1[] = A2*uold[]; rhs1 = -(1-theta)*rhs1;
	rhs[] += rhs1[];
	rhs1[] = L1*uold[];
	rhs1 = (1-I0)*rhs1;
	rhs[] += rhs1[];
	rhs1[]= L2*uold[]; 
	rhs[] -=rhs1[];
	u[]=A^-1*rhs[];
	err[] = u[] - uold[];
	uold =u;
	dtol = err[].linfty;
	cout << "dtol = "<<dtol<< endl;
}
	cout << "covergence!"<<endl;
	plot(u,value=1,nbiso=25,fill=1);