(*      Finite Difference Example : 2D diffusion problem, based on       *)
(*      "Mathematica for Scientists and Engineers", T.M.Bahder           *)
(*      This case tries to set up ADI scheme and Tridiagonal solutions   *)
(*      In this version (4) :                                            *)
(*              set up tridiagonal matrices, one row at a time and solve *)

BeginPackage["Diffusion2D`","LinearAlgebra`Tridiagonal`"]

Diffusion2D::usage="
Diffusion2D[x1_,x2_,y1_,y2_,U_,K_,A_,dT_,niterations_,ndisp_] 
Performs niterations of two-dimensional diffusion on an array U defined 
in rectangle (x1 < x < x2, y1 < y < y2), in steps of dT, \n
diffusion coefficient K[x_,y_], absorption A[x_,y_] are functions of
two dimensional plane. Solution is displayed every ndisp iterations. Output 
is a set {U,grid} where U is final iteration and grid is sampling array"

Begin["`Private`"]
Diffusion2D[x1_,x2_,y1_,y2_,Uin_,K_Symbol,A_Symbol,dT_,niterations_Integer,
        ndisp_Integer] :=
Module[{xmin=N[x1],xmax=N[x2],ymin=N[y1],ymax=N[y2],
        hx,hy,u1,u1T,u2,kappa,mua},
 npx= Length[Uin[[1]]];
 npy = Length[Uin];
 hx = (xmax - xmin)/(npx -1);
 hy = (ymax - ymin)/(npy -1);
 (* Tables in Mathematica are ordered by rows, then columns *)
 u1  = Uin; (* soln at time step n *)
 u1T = Table[0,{npx},{npy}]; (* Transpose of soln at time step n *)
 u2  = Table[0,{npy},{npx}]; (* soln at time step n+1/2 *)
 Uout = Table[0,{niterations},{npy},{npx}]; (* output everything *)
 Uout[[1]] = Uin;
 grid = Table[{xmin + (i-1)hx, ymin +(j-1)hy},{j,npy},{i,npx}];
 mua = Table[A[grid[[j,i]][[1]],grid[[j,i]][[2]]],{j,npy},{i,npx}]; (* absorption *)
 kappa = Table[K[grid[[j,i]][[1]],grid[[j,i]][[2]]],{j,npy},{i,npx}]; (* diffusion *)
 (*     I can't figure this one out *)
 (* mua   = (A[#[[1]],#[[2]]])& /@ grid; *) (* absorption *)
 (* kappa = (K[#[[1]],#[[2]]])& /@ grid; *) (* diffusion *)
 Print["hx ",hx," hy ",hy," dT ",dT];
 Do [ {
    Print["Iteration ",k];
    (*  Print["u1 : ",MatrixForm[u1]];*)
        (*      solve time step n +1/2 in x     *)
    Do [{
        c = -dT/hx^2 Table[(kappa[[j,i+1]]-kappa[[j,i-1]])/8 +kappa[[j,i]]/2,
                                        {i,2,npx-1}] ;
        a =  dT/hx^2 Table[(kappa[[j,i+1]]-kappa[[j,i-1]])/8 -kappa[[j,i]]/2,
                                        {i,2,npx-1}];
        b = Table[1 + dT kappa[[j,i]]/hx^2 + dT mua[[j,i]]/2,
                                        {i,2,npx-1}];
        s = Table[ u1[[j,i]] + 
         dT/(8 hy^2)(kappa[[j+1,i]]-kappa[[j-1,i]])(u1[[j+1,i]]-u1[[j-1,i]])-
         dT/(2 hy^2)kappa[[j,i]](2 u1[[j,i]]-u1[[j-1,i]]-u1[[j+1,i]]) -
         dT mua[[j,i]] u1[[j,i]]/2,
        {i,2,npx-1} ];
        soln = TridiagonalSolve[Drop[a,1],b,Drop[c,-1],s];
        u2[[j]] = Flatten[Join[{0,soln,0}]];
    },{j,2,npy-1}];
    (*Print[MatrixForm[u2]];*)
    Print["Done time step 1/2"];
        (*      solve time step n  in y   *)
    Do [{
        c = -dT/hy^2 Table[(kappa[[j+1,i]]-kappa[[j-1,i]])/8 +kappa[[j,i]]/2,
                                        {j,2,npy-1}] ;
        a =  dT/hy^2 Table[(kappa[[j+1,i]]-kappa[[j-1,i]])/8 -kappa[[j,i]]/2,
                                        {j,2,npy-1}];
        b = Table[1 + dT kappa[[j,i]]/hy^2 + dT mua[[j,i]]/2,
                                        {j,2,npy-1}];
        s = Table[ u2[[j,i]] + 
         dT/(8 hx^2)(kappa[[j,i+1]]-kappa[[j,i-1]])(u2[[j,i+1]]-u2[[j,i-1]])-
         dT/(2 hx^2)kappa[[j,i]](2 u2[[j,i]]-u2[[j,i-1]]-u2[[j,i+1]]) -
         dT mua[[j,i]] u2[[j,i]]/2,
         {j,2,npy-1} ];
        soln = TridiagonalSolve[Drop[a,1],b,Drop[c,-1],s];
        u1T[[i]] = Flatten[Join[{0,soln,0}]];
    },{i,2,npx-1}];
    u1=Transpose[u1T];
    Uout[[k]] = u1;
    (*Print[MatrixForm[u1]];*)
    Print["Done time step 1"];
    If[Mod[k,ndisp]==0,ListPlot3D[u1,PlotRange->All],0];
 },{k,niterations}];
 {Uout,grid,mua,kappa} (*final output *)
]      (* EndModule *)
End[]
EndPackage[]

ReadImage/: ReadImage[filename_, n_] := 
   ReadList[filename, Table[Number, {n}]]
DispIt/: DispIt[t_] := ListPlot3D[t, Lighting -> False, 
    Mesh -> False, AspectRatio -> 3, 
    ViewPoint -> {0, 0, -1000}, Axes -> False]
DispIt/: DispIt[t_, a_] := ListPlot3D[t, Lighting -> False, 
    Mesh -> False, AspectRatio -> a, 
    ViewPoint -> {0, 0, -1000}, Axes -> False]
EpsOut/: EpsOut[which_] := Display["!psfix -epsf > out.ps", 
    which]
SetOptions[ListPlot, PlotJoined -> True];

nptsx = 64;
tmp=ReadImage["arridge.aim",nptsx];
nptsy=Length[tmp];
imr=tmp;
For[i=1,i≤nptsy,i++,imr[[nptsy+1-i]]=tmp[[i]]];

(* synthesise a simple image *)
nptsy = nptsx;
ims=Table[Table[150 Random[],{nptsx}],{nptsy}];
ystart=Floor[(nptsy/4)];
xstart=Floor[(nptsx/4)];
yw=Floor[(nptsy/4)];
xw=Floor[(nptsx/4)];
ycen=Floor[((3 nptsy)/4)];
xcen=Floor[((3 nptsx)/4)];
For[m=1,m≤nptsy,m++,For[n=1,n≤nptsx,n++,ims[[m,n]]+=If[(((m-ycen))^2)+(((n-xcen))^2)<0.8 (xw^2),200,0];]];
For[m=1,m≤yw,m++,For[n=1,n≤xw,n++,ims[[ystart+m,xstart+n]]+=100]];

(* use synthetic image *)
im = ims;

ListDensityPlot[im,Mesh->False];

K[x_,y_]:=1; (*these functions are constant here but could be*)
A[x_,y_]:=0.001;(*spatially varying*)
{imdiff,grid,mua,kappa} = Diffusion2D[0,63,0,63,im,K,A,1,5,12];

For[i = 1, i < Length[imdiff],i++,
ListDensityPlot[imdiff[[i]],Mesh->False]
];

K[x_,y_]:=0.5 (1 + Cos[ 2 Pi x/64] Sin[Pi y/64]); 
A[x_,y_]:=0.00001;
{imdiffa,grid,mua,kappa} = Diffusion2D[0,63,0,63,im,K,A,1,5,12];

ListDensityPlot[kappa,Mesh->False];

For[i = 1, i < Length[imdiffa],i++,
ListDensityPlot[imdiffa[[i]],Mesh->False]
];

g1 =ListPlot[imdiff[[5,32]],PlotJoined->True,DisplayFunction->Identity];
g2= ListPlot[imdiffa[[5,32]],PlotJoined->True,PlotStyle->{RGBColor[1,0,0]},DisplayFunction->Identity];
Show[g1,g2,DisplayFunction->$DisplayFunction];

(* Form diffusion through Perona-Malik scheme  *)

G[x_,s_]:=((E^(-((x^2)/(2 (s^2)))))/(Sqrt[2 π]))
Gx[x_,s_]:=(∂_xx)G[xx,s]/.xx->x
Gxx[x_,s_]:=(∂_({xx,2}))G[xx,s]/.xx->x

(******************************************************
Now the Fourier Versions 
*******************************************************)(* this function rotates quadrants of an image *)
RotImage[im_]:=Module[{nptsx,nptsy},nptsy=Length[im];nptsx=Length[im[[1]]];RotateLeft[Transpose[RotateLeft[Transpose[im],Floor[(nptsx/2)+1]]],Floor[(nptsy/2)+1]]]
(* inverse of the above *)
IRotImage[im_]:=Module[{nptsx,nptsy},nptsy=Length[im];nptsx=Length[im[[1]]];RotateRight[Transpose[RotateRight[Transpose[im],Floor[(nptsx/2)+1]]],Floor[(nptsy/2)+1]]]

fim=Fourier[im];
ycen = nptsy/2;
xcen = nptsx/2;
(* Repeat at larger scale *)
sx = 2.5;
sy = 2.5;
fGh= Table[0,{nptsx}];
fGv= Table[0,{nptsy}];
fGhx=fGh; fGhxx = fGh;
fGvy=fGv; fGvyy = fGv;
For[n=1,n<=nptsx,n++,
fGh[[n]] = G[(((n-xcen) 2 π)/nptsx),(1/sx)];
fGhx[[n]] = (I (n-xcen) 2 π /nptsx )fGh[[n]];
fGhxx[[n]] = (I (n-xcen) 2 π /nptsx )fGhx[[n]];
];
For[m=1,m<=nptsy,m++,
fGv[[m]]=G[(((m-ycen) 2 π)/nptsy),(1/sy)]  ;
fGvy[[m]] = I (m-ycen) 2 π /nptsy fGv[[m]];
fGvyy[[m]] = I (m-ycen) 2 π /nptsy fGvy[[m]];
];

fhmaskim=IRotImage[Outer[Times,fGv,fGhx]];
fvmaskim=IRotImage[Outer[Times,fGvy,fGh]];
(*fhhmaskim=IRotImage[Outer[Times,fGv,fGhxx]];
fvvmaskim=IRotImage[Outer[Times,fGvyy,fGh]];
fhvmaskim=IRotImage[Outer[Times,fGvy,fGhx]];
*)

hfiltimf=Chop[InverseFourier[N[fim fhmaskim]]];
vfiltimf=Chop[InverseFourier[N[fim fvmaskim]]];
(*hhfiltimf=Chop[InverseFourier[N[fim fhhmaskim]]];
vvfiltimf=Chop[InverseFourier[N[fim fvvmaskim]]];
hvfiltimf=Chop[InverseFourier[N[fim fhvmaskim]]];
*)


gradmag= Sqrt[Abs[hfiltimf]^2 + Abs[vfiltimf]^2];
ListDensityPlot[gradmag,Mesh->False,PlotRange->All];

Gmax = Max[gradmag]

K[x_,y_] := Exp[-gradmag[[Floor[1+y],Floor[1+x]]]/(0.5 Gmax)]
A[x_,y_] := 0.000001
{imdiffpm,grid,mua,kappa} = Diffusion2D[0,63,0,63,im,K,A,1,5,12];

ListDensityPlot[kappa,Mesh->False];

For[i = 1, i < Length[imdiffpm],i++,
ListDensityPlot[imdiffpm[[i]],Mesh->False]
];

g3= ListPlot[imdiffpm[[5,32]],PlotJoined->True,PlotStyle->{RGBColor[0,1,0]},DisplayFunction->Identity];
Show[g1,g2,g3,DisplayFunction->$DisplayFunction];