def propagate(psi_0,time,E_rot,E_0_squared_max,sigma,eig,vec):
''' Propagate the wave function psi_0 (given as coefficients)
through a gaussian laser pulse centered at time t = 0.
psi_0: initial wave function
time: array of timesteps to integrate. psi_0 is given at time[0].
The time step must be constant!
E_rot: array of rotational energies for each J.
E_0_squared_max: Max value of the square of the E field during pulse
sigma: temporal variance of the gaussian
eig, vec: diagonalized interaction matrix V = vec * diag(eig) * vec.T
'''
if (numpy.any(numpy.abs(numpy.diff(numpy.diff(time)))>1e-20)):
raise RuntimeError("Pulse time steps must be equidistant.");
# Within the maximum time step, a Gaussian is approximately constant.
if (len(E_rot) > 3):
dt_max = min(2.4*sigma/150,1/(E_rot[-1]-E_rot[-3]));
else:
dt_max = min(2.4*sigma/150,1/E_rot[-1]);
dt = time[1]-time[0];
# If our given time step is larger than this, use a smaller time step
if (dt > dt_max):
scale = int(numpy.ceil(dt/dt_max));
dt = dt/scale;
else:
scale = 1;
expRot = numpy.exp(-1j*dt*E_rot);
expRot2 = numpy.exp(-1j*dt*E_rot/2);
psi_t = numpy.empty((len(time),len(psi_0)), dtype=numpy.complex)
psi_t[0,:] = psi_0;
vecT = numpy.ascontiguousarray(vec.T);
if (libpropagation):
res = libpropagation.propagate_field(len(time),scale,len(psi_0),time[0],dt,E_0_squared_max,sigma,psi_t,eig,vec,vecT,expRot,expRot2);
if (res != 0):
raise RuntimeError("Basis size too small");
else:
i = 1;
for t in time[:-1]: # use split step
psi_t[i,:] = expRot2*psi_t[i-1,:];
for k in range(scale):
# I(t), gaussian beam
if (k > 0): # Double half step:
psi_t[i,:] = expRot*psi_t[i,:];
tp = t+k*dt;
E_0_squared = E_0_squared_max * numpy.exp(-(tp**2)/(2*sigma**2));
psi_t[i,:] = ((numpy.exp(-dt*E_0_squared*1j*eig))*vec).dot(vecT.dot(psi_t[i,:]));
if (numpy.max(numpy.abs(psi_t[i,-2:])**2) > 1e-5):
raise RuntimeError("Basis size too small");
psi_t[i,:] = expRot2*psi_t[i,:];
i = i + 1;
return psi_t;
