(* This file looks at diffusion algorithm projection to a sphere and cylinder after 3D diffusion.  The two options are to project directly to the sphere or cylinder or to first project to a plane that is normal to the sphere or cylinder and then project from the plane to the sphere or cylinder.  I was thinking that the latter method would be better.  However this file shows clearly that the former method is better and in fact is surprisingly accurate. *)

(*************************)
(* Define some functions *)
(*************************)

(* Gaussian *)
gauss[x_,sigma_]:=1/(sigma*Sqrt[2*Pi])*Exp[-x^2/(2*sigma^2)]

(* Produces a random gaussian deviate centered at 0 with std. dev. sigma *)
randomvect[sigma_]:=RandomVariate[MultinormalDistribution[{0,0,0},{{sigma^2,0,0},{0,sigma^2,0},{0,0,sigma^2}}]]

(* Projects vector to x-y plane which has given z value *)
project2plane[xvect_,z_]:={xvect[[1]],xvect[[2]],z}

(* Projects vector to sphere of given radius that is centered about origin. *)
project2sphere[xvect_,radius_]:=Module[{ratio},
ratio=radius/Sqrt[xvect[[1]]^2+xvect[[2]]^2+xvect[[3]]^2];
ratio*xvect]

(* Projects vector to cylinder of given radius that is along x-axis. *)
project2cylinder[xvect_,radius_]:=Module[{ratio},
ratio=radius/Sqrt[xvect[[2]]^2+xvect[[3]]^2];
Return[{xvect[[1]],xvect[[2]]*ratio,xvect[[3]]*ratio}]]

(* Compute polar angle in spherical coordinates for given vector.  This angle is angle away from z axis. *)
computetheta[xvect_]:=ArcTan[xvect[[3]],Sqrt[xvect[[1]]^2+xvect[[2]]^2]]

(* Compute cylinder angle value for given vector.  This angle is angle away from z axis for a cylinder on x axis. *)
computecylangle[xvect_]:=ArcTan[xvect[[3]],xvect[[2]]]

(* Diffuse on sphere surface starting from xvect.  Runs for total standard deviation sigma but performed in given number of steps. *)
diffusemultistep[xvect_,sigma_,radius_,steps_]:=Module[{x,s},
x=xvect;
s=sigma/Sqrt[steps];
For[i=1,i≤steps,i++,{
x=x+randomvect[s];
x=project2sphere[x,radius]; }];
Return[x]]

(* Diffuse on cylinder starting from xvect.  Runs for total standard deviation sigma but performed in given number of steps. *)
diffusecylmultistep[xvect_,sigma_,radius_,steps_]:=Module[{x,s},
x=xvect;
s=sigma/Sqrt[steps];
For[i=1,i≤steps,i++,{
x=x+randomvect[s];
x=project2cylinder[x,radius]; }];
Return[x]]

(* Takes a list of values and converts this to a histogram that can be graphed using the given bin width.  Result should integrate to 1. *)
histlist[dist_,binwidth_]:=Module[{hist},
hist=HistogramList[dist,{binwidth}];
hist[[1]]=hist[[1]]-binwidth/2;
hist[[1]]=Drop[hist[[1]],1];
hist[[2]]=hist[[2]]/Total[hist[[2]]]/binwidth;
Return[Transpose[hist]]];

(* Theory for cylinder distribution.  Assumes that sigma is less than about Pi *)
cyltheory[theta_,sigma_]:=gauss[theta,sigma]+gauss[2*Pi+theta,sigma]+gauss[-2*Pi+theta,sigma]

(* Theory for sphere distribution.  Result is eq. 7 from Ghosh_Sinha_2012. *)
spheretheory[theta_,sigma_]:=Sin[theta]/2*Sum[(2*l+1)*Exp[-l*(l+1)*sigma^2/2]*LegendreP[l,Cos[theta]],{l,0,30}]

(*****************)
(* Test fuctions *)
(*****************)

x={0,0,1}+randomvect[0.1]

project2plane[x,5]

x2=project2sphere[x,1]

Norm[x2]

computetheta[x2]

diffusemultistep[{0,0,1},1,1,10]

(************************************)
(* Compute distributions for sphere *)
(************************************)

(* Sphere has radius 1.  s is rms step length.  binwidth is histogram bin width in radians.  samples is number of random samples. *)

s=0.5
binwidth=0.1
samples=100000

distdirect2sphere=Table[computetheta[project2sphere[{0,0,1}+randomvect[s],1]],{samples}];

disttwostep=Table[computetheta[project2sphere[project2plane[{0,0,1}+randomvect[s],1],1]],{samples}];

histdirect2sphere=histlist[distdirect2sphere,binwidth];
histtwostep=histlist[disttwostep,binwidth];

Show[{Plot[spheretheory[theta,s],{theta,0,Pi},PlotStyle->Black],ListLinePlot[{histdirect2sphere,histtwostep},PlotLegends->{"direct","2 step"},PlotStyle->{Red,Blue}],
ListPlot[{histdirect2sphere,histtwostep},PlotStyle->{Directive[Red,PointSize[Medium]],Directive[Blue,PointSize[Medium]]}]},PlotRange->All,AxesStyle->Directive[Bold,12]]

(* Look at statistics for different methods.  This is mean value. *)
NIntegrate[theta*spheretheory[theta,s],{theta,0,Pi}]
Mean[distdirect2sphere]
Mean[disttwostep]

(* More statistics.  This is rms step length on sphere surface.  The factor of 2 is here because this is 2D diffusion where the mean squared step length is 4Dt rather than 2Dt. *)
sphstddev=Sqrt[NIntegrate[theta^2*spheretheory[theta,s]/2,{theta,0,Pi}]]
{x=Sqrt[Total[distdirect2sphere^2]/Length[distdirect2sphere]/2],x/sphstddev}
{x=Sqrt[Total[disttwostep^2]/Length[disttwostep]/2],x/sphstddev}

(* Sphere calculations, combined to single functions. *)
stddevsphereexact[s_]:=Sqrt[NIntegrate[theta^2*spheretheory[theta,s]/2,{theta,0,Pi}]];
stddevspheredirect[s_,samples_]:=Module[{dist,stddev},
dist=Table[computetheta[project2sphere[{0,0,1}+randomvect[s],1]],{samples}];
stddev=Sqrt[Total[dist^2]/Length[dist]/2];
Return[stddev]];
stddevspheretwostep[s_,samples_]:=Module[{dist,stddev},
dist=Table[computetheta[project2sphere[project2plane[{0,0,1}+randomvect[s],1],1]],{samples}];
stddev=Sqrt[Total[dist^2]/Length[dist]/2];
Return[stddev]];

stddevsph=Transpose[Table[{{s,stddevsphereexact[s]},{s,stddevspheredirect[s,100000]},{s,stddevspheretwostep[s,100000]}},{s,0.1,1.5,0.1}]];

stddevsph={Prepend[stddevsph[[1]],{0,0}],Prepend[stddevsph[[2]],{0,0}],Prepend[stddevsph[[3]],{0,0}]};

Show[{ListLinePlot[stddevsph,PlotStyle->{Black,Red,Blue}],
ListPlot[stddevsph,PlotStyle->{Directive[Black,PointSize[0]],Directive[Red,PointSize[Medium]],Directive[Blue,PointSize[Medium]]}]},AxesStyle->Directive[Bold,12]]

stddevsph[[1,2;;11,1]]

(* Maximum fractional error for s between 0 and 1 for direct and 2 step *)
Max[Abs[stddevsph[[2,2;;11,2]]/stddevsph[[1,2;;11,2]]-1]]
Max[Abs[stddevsph[[3,2;;11,2]]/stddevsph[[1,2;;11,2]]-1]]

(*************************************)
(* Compute distributions for cylinder *)
(*************************************)

s=0.5
binwidth=0.1
samples=100000

distdirect2cyl=Table[computecylangle[project2cylinder[{0,0,1}+randomvect[s],1]],{samples}];

distcyltwostep=Table[computecylangle[project2cylinder[project2plane[{0,0,1}+randomvect[s],1],1]],{samples}];

histdirect2cyl=histlist[Abs[distdirect2cyl],binwidth];
histcyltwostep=histlist[Abs[distcyltwostep],binwidth];

Show[{Plot[2*cyltheory[theta,s],{theta,0,Pi},PlotStyle->Black],ListLinePlot[{histdirect2cyl,histcyltwostep},PlotLegends->{"direct","2 step"},PlotStyle->{Red,Blue}],
ListPlot[{histdirect2cyl,histcyltwostep},PlotStyle->{Directive[Red,PointSize[Medium]],Directive[Blue,PointSize[Medium]]}]},PlotRange->All,AxesOrigin->{0,0},AxesStyle->Directive[Bold,12]]

(* Look at statistics for different methods *)
cylstddev=Sqrt[NIntegrate[theta^2*cyltheory[theta,s],{theta,-Pi,Pi}]]
{x=StandardDeviation[distdirect2cyl],x/cylstddev-1}
{x=StandardDeviation[distcyltwostep],x/cylstddev-1}

(* Cylinder calculations, combined to single functions. *)
stddevcylexact[s_]:=Sqrt[NIntegrate[theta^2*cyltheory[theta,s],{theta,-Pi,Pi}]];
stddevcyldirect[s_,samples_]:=Module[{dist,stddev},
dist=Table[computecylangle[project2cylinder[{0,0,1}+randomvect[s],1]],{samples}];
stddev=StandardDeviation[dist];
Return[stddev]];
stddevcyltwostep[s_,samples_]:=Module[{dist,stddev},
dist=Table[computecylangle[project2cylinder[project2plane[{0,0,1}+randomvect[s],1],1]],{samples}];
stddev=StandardDeviation[dist];
Return[stddev]];

stddevcyl=Transpose[Table[{{s,stddevcylexact[s]},{s,stddevcyldirect[s,100000]},{s,stddevcyltwostep[s,100000]}},{s,0.1,1.5,0.1}]];

stddevcyl={Prepend[stddevcyl[[1]],{0,0}],Prepend[stddevcyl[[2]],{0,0}],Prepend[stddevcyl[[3]],{0,0}]};

Show[{ListLinePlot[stddevcyl,PlotStyle->{Black,Red,Blue}],
ListPlot[stddevcyl,PlotStyle->{Directive[Black,PointSize[0]],Directive[Red,PointSize[Medium]],Directive[Blue,PointSize[Medium]]}]},AxesStyle->Directive[Bold,12]]

(* Maximum fractional error for s between 0 and 1 *)
Max[Abs[stddevcyl[[2,2;;11,2]]/stddevcyl[[1,2;;11,2]]-1]]
Max[Abs[stddevcyl[[3,2;;11,2]]/stddevcyl[[1,2;;11,2]]-1]]