(***********************************************************************
This file was generated automatically by the Mathematica front end.
It contains Initialization cells from a Notebook file, which typically
will have the same name as this file except ending in ".nb" instead of
".m".

This file is intended to be loaded into the Mathematica kernel using
the package loading commands Get or Needs.  Doing so is equivalent to
using the Evaluate Initialization Cells menu command in the front end.

DO NOT EDIT THIS FILE.  This entire file is regenerated automatically 
each time the parent Notebook file is saved in the Mathematica front end.
Any changes you make to this file will be overwritten.
***********************************************************************)



(* 
    Package for using the split-step method for solving the cubic
    nonlinear Schrodinger equation

    D[u[x,t],t] + D[u[x,t],x,x] + 2 Abs[u[x,t]]^2 u[x,t] == v(x) u[x,t]
*)


BeginPackage["SplitStep`"];



SplitStepSolve::usage = "SplitStepSolve[initialcondition, {x, xmin, xmax, n}, {{tmin, tmax, toutput}, deltat}] computes the solution of the cubic nonlinear Schrodinger equation on the spatial interval xmin to xmax from tmin to tmax using time steps of deltat and producing output at intervals toutput.";

Options[SplitStepSolve] = {Potential->0, Order->4,Output->AbsValue, WorkingPrecision->$MachinePrecision, AccuracyGoal->Automatic, PrecisionGoal->Automatic,StoppingCondition->None, MovingWindow->False, AbsorbingBoundary->False};

Discretize::usage = "Discretize[f,{x,xmin,xmax}, n] will produce a list of n values of f[x] from f[xmin] up to, but not including f[xmax].";

Mass::usage = "Mass[Abs[u], L] gives the mass analogue for the solution represented by the vector u over the length L.";

Momentum::usage = "Momentum[u,L] gives the momentum analogue for the solution represented by the vector u over the length L.";

Centroid::usage = "Centroid[Abs[u], grid] gives the position of the centroid of the discretized solution u on the given grid.";

Variance::usage = "Variance[Abs[u],grid] gives the variance of the solution represented by u on the given grid.";

Value::usage = "Value is an output option for SplitStepSolve which gives the solution vector.";

AbsValue::usage = "AbsValue is an output option for SplitStepSolve which gives the modulus of the solution vector.";

MaxValue::usage = "MaxValue is an output option for SplitStepSolve which gives the maximum modulus of the solution vector.";

MaxPosition::usage = "MaxPosition is an output option for SplitStepSolve which gives the position of maximum modulus of the solution vector.";

Grid::usage = "MaxPosition is an output option for SplitStepSolve which gives the current grid.  This is useful to have when MovingWindow->True.";

Time::usage = "Refers to the output time at which values are given.";

SplitStepSolve::precg = NIntegrate::precg;
SplitStepSolve::accg = NIntegrate::accg;

SplitStepSolve::timest = "The time step `1` should be a positive real number.";

SplitStepSolve::fixts = "The AccuracyGoal and PrecisionGoal options have no effect with fixed time steps.";

SplitStepSolve::output = "`1` is not a valid output option for SplitStepSolve.";

SplitStepSolve::ctrmov = "You have chosen a moving window. Without either the centroid or the grid in the output, you will not be able to know the real position of the solution.";

SplitStepSolve::absorb = "The value of the option AbsorbingBoundary->`1` should be False, True, or a list of 2 numbers specifying the size and extent of the absorbing boundary condition.";

SplitStepSolve::order = "The value of the option Order->`1` should be either 1 or and even integer giving the temporal convergence order.";

SplitStepSolve::stopt = "The computation was not completed because the value of the option StoppingCondition->`1` evaluated to neither True nor False at `2`.";



Begin["`Private`"]





Linear[dt_, L_][u_List] := InverseFourier[Fourier[u] CachedD2Vector[dt, Length[u], L]]



ExpFourierD2Vector[dt_, n_, L_] := 
Block[{t, s = -I dt (2 Pi/L)^2},
	t = Table[Exp[s j^2], {j, 1, n/2 - 1}];
	Join[{1}, t, {Exp[s (n/2)^2]}, Reverse[t]]
]



CompiledExpFourierD2Vector = 
Compile[{{dt, _Real}, {n, _Integer}, {L, _Real}}, 
	Block[{t, s = -I dt (2 N[Pi]/L)^2},
		t = Table[Exp[s j^2], {j, 1, n/2 - 1}];
		Join[{1.}, t, {Exp[s (n/2)^2]}, Reverse[t]]
	]
];



CachedD2Vector[dt_, n_, L_] := 
(CachedD2Vector[dt, n, L] = 
	If[MachineNumberQ[dt] || MachineNumberQ[L],
		CompiledExpFourierD2Vector[dt,n,L],
		ExpFourierD2Vector[dt, n, L]
	]
)



FourierDVector[ord_, n_, L_] := 
Block[{t, s},
	If[ord === 0, Return[Table[1,{n}]]];
	s  = (2 Pi I/L)^ord;
	t = Table[s j^ord, {j, 1, n/2 - 1}];
	Join[{0}, t (-1)^ord, {s (n/2)^ord}, Reverse[t]]
]



(*
	The potential is assumed to be a nonlinear function of v
	for the basic cubic NLSE, use, e.g.
	v = Function[2 # Conjugate[#]]
*)
Nonlinear[dt_, v_][u_List] := u Exp[I dt v[u]]



SplitStep[1, dt_, v_, L_][u_] := Linear[dt, L][Nonlinear[dt, v][u]]



SplitStep[2, dt_, v_, L_][u_] :=
	Linear[dt/2, L][Nonlinear[dt, v][Linear[dt/2,L][u]]]



(* 
	This is the method of Yoshida.
	Depends on the time-reversability of the NLSE 
*)
SplitStep[n_?EvenQ, dt_, p___][u_] := 
Block[{z0, z1, alpha, beta, k = 1/(n - 1)},
	z1 = 1/(2 - 2^k); 
	z0 = - z1 2^k; 
	SplitStep[n - 2, z1 dt, p][SplitStep[n - 2, z0 dt, p][SplitStep[n - 2, z1 dt, p][u]]]]



(* 
	Take the SplitStep method for order ord and simplify it for 
	actual numerical use.
*)
SimplifySplitStep[ord_?EvenQ, {dt_, p___}, prec___] := 
Block[{Linear, Nonlinear, u}, 
	(* This MUST be Block here to work.  Since the methods are defined
		in terms of Linear (and Nonlinear), using Block allows you to 
		give a temporary definition for it without having that affect
		the main one you want to be used in numerical evaluation (or 
		having the definition for numerical evaluation evaluate here.)
	*)
	Module[{new}, (* Use Module for local variables *)
		(* We make a (temporary) definition for Linear which will combine
		   two Linear steps into a single one.  *)
		Linear[dt1_, pl___][Linear[dt2_, pl___][x__]] := 
			Linear[Simplify[dt1 + dt2], pl][x];
		new = SplitStep[ord, dt, p][Slot[1]];
		new = N[new, prec];
		(* Now we want to make a Function that we can call to invoke a step 
		   of the simplified composite method.  Using Function here is fast
		   and it prevents evaluation until time of use. *)
		args = {dt, p};
		Function[Evaluate[args], Evaluate[Function[Evaluate[new]]]]
	]
]



Discretize[f_, {x_,xmin_, xmax_}, n_] := 
Block[{fx = Function[x,f]},
  Table[fx[xmin + i(xmax - xmin)/n],{i,0,n-1}]
]



SplitStepSolve[InitialCondition_, {x_, xmin_, xmax_, n_}, {{tbegin_, tend_, T_:1},dt_}, opts___] := SplitStepSolve[Discretize[InitialCondition, {x,xmin,xmax}, n], {x, xmin, xmax}, {{tbegin, tend, T}, dt}, opts];

SplitStepSolve[InitialCondition_List, {x_, xmin_, xmax_}, {{tbegin_, tend_, T_:1},dtin_},  opts___]  := 
Block[{n = Length[InitialCondition],
		L = xmax - xmin,
		pot, rv, v, ndt, order, out, wp, alpha, beta, stoptest, absorb,
		moving,grid,deltax,position, last,
		$MaxPrecision, $MinPrecision},
	wp = WorkingPrecision /. {opts} /. Options[SplitStepSolve];
	If[Or[dtin === Automatic, And[ListQ[dtin], Length[dtin] > 0, dtin[[1]] === Automatic]],
		dt = Automatic;
		dts = N[If[And[ListQ[dtin], Length[dtin] > 1],dtin[[2]], T/1000], wp];
		If[Not[NumberQ[dts]],
			Message[SplitStepSolve::timest, dts];
			Return[$Failed]
		],
	(* else *)
		dt = N[dtin, wp];
		If[Not[NumberQ[dts]],
			Message[SplitStepSolve::timest, dts];
			Return[$Failed]
		]
	];
	(* 
		Process the Accuracy and Precision goal for Automatic time step control
	*)
	If[dt === Automatic,
		atol = AutoAGPG[wp,AccuracyGoal /. {opts} /. Options[SplitStepSolve], Accuracy];
		If[atol === $Failed, Return[$Failed]];
		rtol = AutoAGPG[wp, PrecisionGoal /. {opts} /. Options[SplitStepSolve], Precision];
		If[rtol === $Failed, Return[$Failed]],
	(* else *)
		If[MemberQ[{opts}, AccuracyGoal | PrecisionGoal,Infinity],
			Message[SplitStepSolve::fixts]]
	];			
	(* 
		The Potential.  Unless it is already a list, we discretize it
	*)
	pot = Potential /. {opts} /. Options[SplitStepSolve];
	If[!ListQ[pot],
	    v = Discretize[pot,{x,xmin,xmax},n],
	    v = pot];
	(*
		The output to keep track of 
	*)
	out = Output /. {opts} /. Options[SplitStepSolve];
	If[!ListQ[out],out = {out}];
	If[Catch[Scan[
			If[!MatchQ[#,Value | AbsValue | MaxValue | MaxPosition | Mass | Momentum | Centroid | Variance | Grid | Plot | {Plot, ___}],Message[SplitStepSolve::output, #]; Throw[True]]&,
			out];False],
		Return[$Failed]];
	moving = TrueQ[MovingWindow /. {opts} /. Options[SplitStepSolve]];
	If[moving,
	    If[Not[MemberQ[out, Centroid | Grid]],
		Message[SplitStepSolve::ctrmov]];
	    If[ListQ[pot],
		ppot = True;
		Message[SplitStepSolve::periodicv],
		pder = pot;
		ppot = Catch[Do[
		    If[(pder /. x->xmax) != (pder /. x->xmin),Throw[False]];
		    pder = D[pder, x],
		    {5}]; True]
	    ];
	];
	(* 
		The is a simple default for AbsorbingBoundary->True which 
		works pretty well most of the time.
	*)
	absorb = AbsorbingBoundary /. {opts} /. Options[SplitStepSolve];
	If[absorb =!= False,
	    If[TrueQ[absorb], absorb = {1,1/4}];
	    If[Not[ListQ[absorb] && (Length[absorb] == 2)],
			Message[SplitStepSolve::absorb, absorb];
			Return[$Failed];
	    ];
	    absorb = Discretize[AbsorbingBoundaryConditionPotential[absorb[[1]], {xmin,xmax}, {xmin, xmax} + L absorb[[2]] {1,-1}][x],{x,xmin,xmax},n];
	];
	(* 
		The order should be 1 or even 
	*)
	order = Order /. {opts} /. Options[SplitStepSolve];
	If[Not[Or[order === 1, EvenQ[order]]],
		Message[SplitStepSolve::order];
		Return[$Failed];
	];
	If[wp === $MachinePrecision, np = Sequence[], np = wp];
	(*
		Derive the integration method
	*)
	splits = SimplifySplitStep[order, {hsym, vsym, Lsym},np];
	grid = Discretize[x,{x,xmin,xmax},n];
	origgrid = grid;
	deltax = (xmax - xmin)/n;
	position = Round[(xmax + xmin)/(2*deltax)];
	If[wp === $MachinePrecision,
		(*
			With machine precision, we set up complex PackedArrays
			this ensures top speed.
		*)
		grid = Developer`ToPackedArray[grid, Real];
		ulast = Developer`ToPackedArray[#, Complex]& /@ InitialCondition;
		v = Developer`ToPackedArray[v, Complex];
		If[absorb =!= False, 
			absorb = Developer`ToPackedArray[absorb, Complex];
			av = v + absorb,
		(* else *)
			av = v],
	(* else 
			otherwise, we numericize to the desired precision 
		*)
		grid = N[grid, wp];
		ulast = N[InitialCondition, wp];
		v = N[v, wp];
		If[absorb =!= False, 
			absorb = N[absorb, wp];
			av = N[v + absorb + v, wp],
		(*else*)
			av = v];
		(*
			Fix the precision.  These have been localized inside Block, so 
			this will not affect your session.  For iterative processes
			it is often a good idea to fix the precision.
		*)
		$MaxPrecision = $MinPrecision = wp
	];
	stoptest = StoppingCondition /. {opts} /. Options[SplitStepSolve];
	If[stoptest === None, stoptest = Function[True]];
	glast = grid;
	vext = v;
	tt = tbegin;
	If[dt === Automatic,
		h = dts,
	(* else *)
		ndt = Ceiling[T/dt];
		h = N[T/ndt, np]];
	Catch[
		databag = Internal`Bag[];
		stepbag = Internal`Bag[];
		Internal`StuffBag[databag,StepOutput[tbegin,ulast,out, grid]];
		Do[
		    last = Internal`BagPart[databag, -1];
			stopd = stoptest[last];
			If[stopd =!= False,
				If[stopd =!= True,
					Message[SplitStepSolve::stopt, stoptest,Rule[Time, Time /. last]]];
				Break[]];
			If[moving,
				cent = Centroid /. last;
			    shift = Round[(cent/deltax) - position];
				If[shift != 0, 
					position += shift;
					grid = origgrid + position*deltax;
					ulast = RotateLeft[ulast, shift];
					If[ppot, 
						v = RotateLeft[v, shift],
					(* else *)
						v = Discretize[pot,{x,xmin + position*deltax, xmin+position*deltax}]
					];
					av = If[absorb === False,v, v+absorb];
				];
			];
			If[Developer`ZeroQ[Max[Abs[av]]],
				vfun = Function[2 # Conjugate[#]],
				vfun = Function[2 # Conjugate[#] - av]
			];
			If[dt === Automatic, 
				While[Not[Or[tt == t, tt > t]],
					If[tt + 2 h >= t,
					If[Or[tt + h == t, tt + h > t],
						h = t - tt,
						(* Avoid a very small step *)
						h = (t - tt)/2
						];
					];
					{h,hused,ulast} = DoubleStep[Function[{hh},Evaluate[splits[hh,vfun, L]]], h, ulast, order, {atol, rtol}];
					tt += hused;				
					Internal`StuffBag[stepbag,{tt,hused}];
				],
			(* else: fixed step size *)
				ulast = Nest[splits[h, vfun, L], ulast, ndt];
				tt += h*ndt;
				Internal`StuffBag[stepbag,{tt,h}];
			];
			glast = grid;
			Internal`StuffBag[databag, StepOutput[t,ulast,out, grid]],
			{t, tbegin + T, tend, T}];
	];
	DownValues[CachedD2Vector] = Last[DownValues[CachedD2Vector]];
	Internal`BagPart[databag, All]
]





DoubleStep[split_, htry_, u_, ord_, {atol_, rtol_}] :=  
Block[{h = 2 htry, half, one, two, errest = Infinity, count = 0, fac = 1},
    half = split[h/2][u];
    While[errest > 1,
		count++;
		h /= 2;
		one = half;
		half = split[h/2][u];
		two = split[h/2][half];
		w = atol + rtol*Map[Max, Transpose[{Abs[u], Abs[two]}]];
		diff = Max[Abs[(two - one)/w]];
		errest = diff/(2^ord - 1);
    ];
    (* Success: increase h if we didn't have to reduce *)
    If[count < 2,
	If[Developer`ZeroQ[errest],
	    fac = 2,
	    fac = Min[2,(9/10)*(1/errest)^(1/(ord + 1))]
	]
    ];
    {h*fac, h, two}];



AutoAGPG[wprec_, apgin_, pora_] :=
Module[{apg = apgin},
    If[apg === Automatic,
		apg = If[wprec <= $MachinePrecision, 6, wprec - 10],
		If[Not[NumberQ[apg] && Positive[apg]],
			If[pora === Precision,
				Message[SplitStepSolve::precg, apg],
				Message[SplitStepSolve::accg, apg]];
			Return[$Failed];
		]
    ];
    SetPrecision[10.^(-apg), wprec]
];



AbsorbingBoundaryConditionPotential[g_, {xmin_, xmax_}, {xL_, xR_}][x_] := 
    - I (g/2) If[x < xL, 
	1 - Cos[Pi (x - xL)/(xmin - xL)],
	If[x > xR, 1 - Cos[Pi (x - xR)/(xmax - xR)], 0]]



StepOutput[t_, u_, things_, grid_] := 
Block[{datum, L, ua, plot = False},
	L = n*(grid[[2]] - grid[[1]]);
	ua = Abs[u];
	datum = Join[{Time->t},things];
	Do[
		datum[[j]] = Switch[datum[[j]],
			Value, Value->u,
			AbsValue, AbsValue->Abs[u],
			MaxValue, MaxValue->Max[Abs[u]],
			MaxPosition, tmp = Abs[u]; MaxPosition->Position[tmp, Max[tmp]][[1,1]],
			Grid, Grid->grid, 
			Mass, Mass->Mass[ua, L],
			Centroid, Centroid->Centroid[ua, grid],
			Variance, Variance->Variance[ua, grid],
			Momentum, Momentum->Momentum[u, L, FourierDVector[1,Length[u], L]],
			Plot | {Plot,popts___}, 
				popts = Sequence[];
				If[ListQ[datum[[j]]] && Length[datum[[j]]] > 1,
					popts = Apply[Sequence, Drop[datum[[j]],1]]];
				plot = True;Show[Block[{$DisplayFunction = Identity}, {
					ListPlot[Transpose[{grid, Abs[u]}], PlotJoined->True],
					ListPlot[Transpose[{grid, Re[u]}], PlotJoined->True, PlotStyle->RGBColor[1,0,0]],
					ListPlot[Transpose[{grid, Im[u]}], PlotJoined->True, PlotStyle->RGBColor[0,1,0]],
					ListPlot[Transpose[{grid, Abs[u]}], PlotStyle->{PointSize[0.025], RGBColor[0,0,1]}]}], popts]; Plot
		],
		{j,2,Length[datum]}];
	If[TrueQ[plot],datum = DeleteCases[datum, Plot]];
	datum];



Mass[ua_, {xmin_, xmax_}] := Mass[ua, xmax-xmin]

Mass[ua_, L_] := Tr[ua^2] L/Length[ua]

Centroid[ua_, grid_] := Tr[ua^2 grid]/Tr[ua^2]

Variance[ua_, grid_] := With[{c = Centroid[ua, grid]},
	Tr[(ua (c - grid))^2]/Tr[ua^2]]

Momentum[u_, L_, fder_:Automatic] :=
Block[{fd = fder, ux},
	If[fder === Automatic,
		fd = FourierDVector[1,Length[u], L]];
	ux = InverseFourier[Fourier[u]*fd];
	(L Im[Conjugate[u].ux]/Length[u])
]

End[]; (* Private *)

EndPackage[];