Gerchberg-Saxton optimization: 
spectral phase shaping of pulse to optimize time-domain amplitude profile

optimize spectral phase to obtain desired amplitude
Fourier transform strategy: 
start with input spectral amplitude shape (will not change this), either flat or random spectral phase
inverse Fourier transform to time domain
discard temporal amplitude profile, put in target amplitude profile, keep temporal phase
FT to spectral domain
replace spectral amplitude with input amplitude, keep spectral phase
iterate, hope for convergence

constants and unit definitions

s=1;ms=(10^(-3))s;us=(10^(-6))s;ns=(10^(-9))s;ps=(10^(-12))s; fs=(10^(-15))s;
Hz=1;kHz=(10^3)Hz;GHz=(10^9)Hz;THz=(10^12)Hz;
m=1;cm=(10^(-2))m;mm=(10^(-3))m;um=(10^(-6))m;nm=(10^(-9))m;angs=(10^(-10))m;
kg=1;g=(10^(-3));
J=1;mJ=(10^(-3))J;uJ=(10^(-6))J;nJ=(10^(-9))J;
eV=e J;keV=(10^3)eV;MeV=(10^6)eV;meV=(10^(-3))eV;
W=1;GW=(10^9)W;MW=(10^6)W;kW=(10^3)W;mW=(10^(-3))W;
Torr=1/760;

c=299792458 m/s;      
eps0=8.854187817 (10^(-12));   (*permittivity of free space*)
mu0=4 π (10^(-7));
re=2.81794092 (10^(-15))m;
e=1.60217733 (10^(-19));		(*Coul electron charge*)
hbar=1.05457266 (10^(-34))J s;

Fourier Gerchberg-Saxton algorithm

Define FFT functions for temporal domain

FFT[fLst_,npts_,dt_]:=(Sqrt[npts (dt^2)])RotateLeft[Fourier[RotateRight[fLst,npts/2]],npts/2];
InverseFFT[FLst_,npts_,dt_]:=(Sqrt[(1/(npts (dt^2)))])RotateLeft[InverseFourier[RotateRight[FLst,npts/2]],npts/2];

unwrap phase function

this version allows you to choose where the starting point is
phase is often noisy at edges of array, so start from center: dir = 0
dir = -1 start from left
dir = +1 start from right

unwrap[phs_,npts_,dir_]:=
Block[{phstmp},
phstmp=Table[1,{npts}];
Which[
dir==-1,
(Subscript[phstmp, [[1]]])=(Subscript[phs, [[1]]]);
shift=0;
For[i=2,i≤npts,i++,
If[(Abs[(Subscript[phs, [[i-1]]])-(Subscript[phs, [[i]]])]>2.5),
shift=shift+Sign[(Subscript[phs, [[i-1]]])-(Subscript[phs, [[i]]])] 2 π//N,
Null
];
(Subscript[phstmp, [[i]]])=(Subscript[phs, [[i]]])+shift;
],

dir==+1,
(Subscript[phstmp, [[npts]]])=(Subscript[phs, [[npts]]]);
shift=0;
For[i=npts-1,i≥1,i--,
If[(Abs[(Subscript[phs, [[i+1]]])-(Subscript[phs, [[i]]])]>2.5),
shift=shift+Sign[(Subscript[phs, [[i+1]]])-(Subscript[phs, [[i]]])] 2 π//N,
Null
];
(Subscript[phstmp, [[i]]])=(Subscript[phs, [[i]]])+shift;
],

dir==0,
(Subscript[phstmp, [[npts/2]]])=(Subscript[phs, [[npts/2]]]);
shift=0;
For[i=npts/2-1,i≥1,i--,
If[(Abs[(Subscript[phs, [[i+1]]])-(Subscript[phs, [[i]]])]>2.5),
shift=shift+Sign[(Subscript[phs, [[i+1]]])-(Subscript[phs, [[i]]])] 2 π//N,
Null
];
(Subscript[phstmp, [[i]]])=(Subscript[phs, [[i]]])+shift;
];
shift=0;
For[i=npts/2+1,i≤npts,i++,
If[(Abs[(Subscript[phs, [[i-1]]])-(Subscript[phs, [[i]]])]>2.5),
shift=shift+Sign[(Subscript[phs, [[i-1]]])-(Subscript[phs, [[i]]])] 2 π//N,
Null
];
(Subscript[phstmp, [[i]]])=(Subscript[phs, [[i]]])+shift;
]
];
    Return[phstmp]
];

set up grids

frequency variable: ω

ntpts=(2^12);
dt = 3.fs;
λ0=800.nm;
ω0=2π c/λ0;

tSpan=ntpts dt;
tLst=Table[i dt,{i,-ntpts/2,ntpts/2-1}];
dω=2π/(ntpts dt);
ωLst=Table[ω0+iω dω,{iω,-ntpts/2,ntpts/2-1}];
ωSpan=ntpts dω; 

λLst=2π c/ωLst;
dλ=λ0 dω/ω0;
Print["tSpan: ",tSpan/fs," fs, dt: ",dt/fs," fs"];
Print["ωSpan: ",ωSpan fs," rad/fs, dω: ",dω fs," /fs"];
Print["λMin: ",λLst[[ntpts]]/nm," nm. λMax: ",λLst[[1]]/nm," nm, dλ: ",dλ/nm," nm"];

define pulse shape

τIn and Δω are 1/e2 half widths in intensity

τIn=50fs;
Δω=2/τIn;

AtInLst=Exp[-(tLst^2)/((τIn)^2)];
AωInLst=FFT[AtInLst,ntpts,dt];

Print["time domain"];
tRange=3 τIn;
ListPlot[Transpose[{tLst/fs,((Abs[AtInLst])^2)}],Joined->True,PlotRange->{{-tRange/fs,tRange/fs},All},AxesLabel->{"t (fs)","intensity"}]

Print["spectral domain"];
ωRange=3Δω;
ListPlot[Transpose[{ωLst fs,((Abs[AωInLst])^2)}],Joined->True,PlotRange->{{(ω0-ωRange)fs,(ω0+ωRange)fs},All},AxesLabel->{"ω (rad/fs)","intensity"}]

test sample phase profiles

a2=.5/((Δω)^2);
a4 =-0.1/((Δω)^4);
phaseTstLst=2π(a2 (((ωLst-ω0))^2)+a4 (((ωLst-ω0))^4));
ϕRange=4 2π;

ListPlot[Transpose[{ωLst fs,phaseTstLst}],Joined->True,PlotRange->{{(ω0-ωRange)fs,(ω0+ωRange)fs},{-ϕRange,ϕRange}}]

AtLst=InverseFFT[AωInLst Exp[I phaseTstLst],ntpts,dt];
Print["time domain"];
ListPlot[Transpose[{tLst/fs,((Abs[AtLst])^2)}],Joined->True,PlotRange->{{-tRange/fs,tRange/fs},All}]

define target function: superGaussian

for nSG>4, underflow in wings of shape, so zero out these points

nSG=4;
τTarget=170fs;
aTst=(10^(-1));
targetFcn[t_,τ_,a_]:=a Exp[-(t^nSG)/(τ^nSG)];
BtTargetLst=Range[ntpts];
For[it=1,it≤ntpts,it++,
(Subscript[BtTargetLst, [[it]]])=If[((((Subscript[tLst, [[it]]])/τTarget))^nSG)>(10^7),0,targetFcn[(Subscript[tLst, [[it]]]),τTarget,aTst]
]
];
Print["time domain: target function"];
tRange=2 τTarget;
ListPlot[Transpose[{tLst/fs,((Abs[BtTargetLst])^2)}],Joined->True,PlotRange->{{-tRange/fs,tRange/fs},All}]

BωTargetLst=FFT[BtTargetLst,ntpts,dt];

Print["frequency domain: amplitude"];
ωRange=2Δω;
ListPlot[{Transpose[{ωLst fs,(((Abs[AωInLst])^2)/(Max[((Abs[AωInLst])^2)]))}],Transpose[{ωLst fs,(((Abs[BωTargetLst])^2)/(Max[((Abs[BωTargetLst])^2)]))}]},Joined->True,PlotRange->{{(ω0-ωRange)fs,(ω0+ωRange)fs},All}]
Print["frequency domain: wrapped phase"];
ListPlot[Transpose[{ωLst fs,Arg[BωTargetLst]}],Joined->True,PlotRange->{{(ω0-ωRange)fs,(ω0+ωRange)fs},All}]

ϕRange=1;
ListPlot[Transpose[{ωLst fs,unwrap[Arg[BωTargetLst],ntpts,0]}],Joined->True,PlotRange->{{(ω0-ωRange)fs,(ω0+ωRange)fs},{-ϕRange,ϕRange}}]

try fitting pulse to target duration by guessing initial 2nd order phase (not fully implemented)

The strategy here is to start with an initial guess of 2nd order phase to get close to the right duration, then implement from there. The version below was implemented for spatial propagation, but the idea is the same. The target superGaussian must be longer in duration than the transform limit, and this can be achieved by stretching in either direction. Starting the optimization with the a choice of chirp direction generally leads to faster convergence toward the one of the two equivalent local optimum phase profiles.

ans=FindFit[Transpose[{XLst,Abs[AXLst]}],targetFcn[x,w,a],{{w,400um},{a,(10^(-3))}},x];
wFit=w/.ans;
aFit=a/.ans;
Print["wFit: ",wFit/um," um, aFit: ",aFit];
ListPlot[
{Transpose[{XLst/um,((Abs[AXLst])^2)}],
Transpose[{XLst/um,((targetFcn[XLst,wFit,aFit])^2)}]},
Joined->True,PlotRange->{{-XRange/um,XRange/um},All}]

GS loop to optimize phase

“A” represents what we start with, “B” represents the shaped profiles. 

Gerchberg-Saxton loop: 
AωInLst: input spectrum
ω-domain: BωLst (initially = |AωInLst| exp[ i ϕInput ] ) 

In this version ϕStart is random phase. Other possibilities could be flat phase, or quadratic phase that does the initial stretch to the approximately the right duration (see above). 

invFFT to t-domain: BtLst
set amplitude of field in t-domain to target function, keep phase from invFFT:
	BtLst = Abs[BtTargetLst] Exp[i Arg[BtLst] ]
invFFT to ω-space: new BωLst
set amplitude back to input, keep phase:
	BωLst = Abs[AωInLst] Exp[i Arg[BωLst] ]

At end: 
	Arg[BωLst] is phase for shaper
	Arg[BtLst] is residual phase in the time domain

these plotting definitions are designed to be called within the loop, but this hasn’t been implemented yet.

timePlt[BtLst_,BtTargetLst_]:=ListLinePlot[{Transpose[{tLst/fs,((Abs[BtLst])^2)}],Transpose[{tLst/fs,((Abs[BtTargetLst])^2)}]},PlotRange->{{-tRange/fs,tRange/fs},All},PlotStyle->{{Red,Thick},{Blue,Dashed,Thick}},AxesOrigin->{0,0},Frame->True,FrameLabel->{{None,None},{None,"t space" }}];
lensPlt[BωLst_,AωLst_]:=ListLinePlot[{Transpose[{ωLst fs,((Abs[BωLst])^2)}],Transpose[{ωLst fs,((Abs[AωLst])^2)}]},PlotRange->{{(ω0-ωRange)fs,(ω0+ωRange)fs},All},PlotStyle->{{Red,Thick},{Blue,Dashed,Thick}},AxesOrigin->{0,0},Frame->True,FrameLabel->{{None,None},{None,"ω space"}}];

errPlt[errorLst_]:=ListLinePlot[Log[10,errorLst],PlotRange->{{0,nSteps},Automatic},PlotStyle->{Red,Thick},AxesStyle->Thick,GridLines->{None,{Log[10,minαError]}},Frame->True,FrameLabel->{{None,None},{None,"Log[error]"}}];

initialize starting phase

ϕStart=Table[2π RandomReal[],{iω,1,ntpts}];

run loop

nSteps=100;
plotStep=10;
BωLst=AωInLst Exp[I ϕStart];
Monitor[
For[is=1,is≤nSteps,is++,
BtLst=InverseFFT[BωLst,ntpts,dt];
If[Mod[is,plotStep]==0,
gr=GraphicsRow[{ListLinePlot[{Transpose[{tLst/fs,(((Abs[BtLst])^2)/(Max[((Abs[BtLst])^2)]))}],Transpose[{tLst/fs,(((Abs[BtTargetLst])^2)/(Max[((Abs[BtTargetLst])^2)]))}]},PlotRange->{{-tRange/fs,tRange/fs},All},PlotStyle->{{Red,Thick},{Blue,Dashed,Thick}},AxesOrigin->{0,0},Frame->True,FrameLabel->{{None,None},{None,"t space" }}],ListLinePlot[{Transpose[{ωLst fs,((Abs[BωLst])^2)}],Transpose[{ωLst fs,((Abs[AωInLst])^2)}]},PlotRange->{{(ω0-ωRange)fs,(ω0+ωRange)fs},All},PlotStyle->{{Red,Thick},{Blue,Dashed,Thick}},AxesOrigin->{0,0},Frame->True,FrameLabel->{{None,None},{None,"ω space"}}]},ImageSize->700]
];
If[is<nSteps,
BtLst=Abs[BtTargetLst]Exp[I Arg[BtLst]];
BωLst=FFT[BtLst,ntpts,dt];
BωLst=Abs[AωInLst]Exp[I Arg[BωLst]]
]
], gr];
gr

extract the shaper phase, plot

labelSize=12;
dmPhase=unwrap[Arg[BωLst],ntpts,0];
plotFactor=1;
ωRange=2Δω;
Print["red: starting spectral profile, blue: phase profile"];
ListPlot[
{Transpose[{ωLst fs,(((Abs[BωLst])^2)/(Max[((Abs[BωLst])^2)]))plotFactor}],
Transpose[{ωLst fs,dmPhase}]},
Joined->True,PlotRange->{{(ω0-ωRange)fs,(ω0+ωRange)fs},{-4,4}},
PlotStyle->{{Red,Thickness[0.004]},{Blue,Thickness[0.004]}},
AxesStyle->Thickness[0.003],LabelStyle->{FontFamily->Times, FontSize->labelSize}]

outPhaseShift=-2.0;
outPhase=unwrap[Arg[BtLst],ntpts,0]+outPhaseShift;
outInten=((Abs[BtLst])^2)/Max[((Abs[BtLst])^2)];
targetInten=((Abs[BtTargetLst])^2)/Max[((Abs[BtTargetLst])^2)];
plotFactor=5;
Print["red: final beam profile, blue: final phase profile"];
ListPlot[
{Transpose[{tLst/fs,outInten plotFactor}],
(* Transpose[{XLst/um,targetInten plotFactor}],*)
Transpose[{tLst/fs,outPhase}]},
Joined->True,PlotRange->{{-tRange/fs,tRange/fs},{0,5.1}},PlotStyle->{{Red,Thickness[0.004]},{Blue,Thickness[0.004]}},
AxesStyle->Thickness[0.003],LabelStyle->{FontFamily->Times, FontSize->labelSize}]

fit input wavefront to polynomial, remove quadratic dispersion term

find portion of input spectrum where we want to localize the fitting
better to define the width relative to input spectral width, then find number of points to select

winRange=200;
(ωLst[[ntpts/2+1+winRange]] -ω0)/Δω

ωLstWindow=Take[ωLst,{ntpts/2+1-winRange,ntpts/2+1+winRange}];
dmPhaseWindow=Take[dmPhase,{ntpts/2+1-winRange,ntpts/2+1+winRange}];

fitFcn[x_]=Fit[Transpose[{(ωLstWindow-ω0) fs,dmPhaseWindow}],{1,x,(x^2),(x^3),(x^4),(x^5),(x^6)},x]

pActual=ListPlot[Transpose[{(ωLstWindow-ω0) fs,dmPhaseWindow}],Joined->True,PlotStyle->Red,PlotRange->{{-ωRange fs,ωRange fs},Automatic}];
pFit=Plot[fitFcn[x],{x,-ωRange fs,ωRange fs},PlotStyle->Blue];
Show[pActual,pFit]

Print["phase profile without quadratic phase term"];
ϕ2=Coefficient[fitFcn[x],x,2]
pFit=Plot[fitFcn[x]-ϕ2 (x^2),{x,-ωRange fs,ωRange fs},PlotStyle->Blue]