def plot_frames(iom, blockid=0, view=None, plotphase=True, plotcomponents=False, plotabssqr=False, imgsize=(12,9)):
"""Plot the wave function for a series of timesteps.
:param iom: An ``IOManager`` instance providing the simulation data.
:param view: The aspect ratio.
:param plotphase: Whether to plot the complex phase. (slow)
:param plotcomponents: Whether to plot the real/imaginary parts..
:param plotabssqr: Whether to plot the absolute value squared.
"""
parameters = iom.load_parameters()
grid = iom.load_grid(blockid="global")
timegrid = iom.load_wavefunction_timegrid(blockid=blockid)
# Precompute eigenvectors for efficiency
Potential = PotentialFactory().create_potential(parameters)
eigenvectors = Potential.evaluate_eigenvectors_at(grid)
for step in timegrid:
print(" Plotting frame of timestep # " + str(step))
wave = iom.load_wavefunction(blockid=blockid, timestep=step)
values = [ wave[j,...] for j in xrange(parameters["ncomponents"]) ]
# Transform the values to the eigenbasis
# TODO: improve this:
if parameters["algorithm"] == "fourier":
ve = Potential.project_to_eigen(grid, values, eigenvectors)
else:
ve = values
# Plot the probability densities projected to the eigenbasis
fig = figure(figsize=imgsize)
for index, component in enumerate(ve):
ax = fig.add_subplot(parameters["ncomponents"],1,index+1)
ax.ticklabel_format(style="sci", scilimits=(0,0), axis="y")
if plotcomponents is True:
ax.plot(grid, real(component))
ax.plot(grid, imag(component))
ax.set_ylabel(r"$\Re \varphi_"+str(index)+r", \Im \varphi_"+str(index)+r"$")
if plotabssqr is True:
ax.plot(grid, component*conj(component))
ax.set_ylabel(r"$\langle \varphi_"+str(index)+r"| \varphi_"+str(index)+r"\rangle$")
if plotphase is True:
plotcf(grid, angle(component), component*conj(component))
ax.set_ylabel(r"$\langle \varphi_"+str(index)+r"| \varphi_"+str(index)+r"\rangle$")
ax.set_xlabel(r"$x$")
# Set the aspect window
if view is not None:
ax.set_xlim(view[:2])
ax.set_ylim(view[2:])
fig.suptitle(r"$\Psi$ at time $"+str(step*parameters["dt"])+r"$")
fig.savefig("wavefunction_block"+str(blockid)+"_"+ (7-len(str(step)))*"0"+str(step) +GD.output_format)
close(fig)
print(" Plotting frames finished")
def compute_energy(iom, blockid=0):
"""
:param iom: An ``IOManager`` instance providing the simulation data.
:param blockid: The data block from which the values are read.
"""
parameters = iom.load_parameters()
# Number of time steps we saved
timesteps = iom.load_wavefunction_timegrid(blockid=blockid)
nrtimesteps = timesteps.shape[0]
# Retrieve simulation data
if iom.has_grid(blockid=blockid):
grid = iom.load_grid(blockid=blockid)
else:
grid = iom.load_grid(blockid="global")
opT, opV = iom.load_fourieroperators(blockid=blockid)
# We want to save norms, thus add a data slot to the data file
iom.add_energy(parameters, timeslots=nrtimesteps, blockid=blockid)
# Precalculate eigenvectors for efficiency
Potential = PotentialFactory().create_potential(parameters)
eigenvectors = Potential.evaluate_eigenvectors_at(grid)
nst = Potential.get_number_components()
WF = WaveFunction(parameters)
# Iterate over all timesteps
for i, step in enumerate(timesteps):
print(" Computing energies of timestep # " + str(step))
values = iom.load_wavefunction(timestep=step, blockid=blockid)
values = [ values[j,...] for j in xrange(parameters["ncomponents"]) ]
# Project wavefunction values to eigenbasis
values = Potential.project_to_eigen(grid, values, eigenvectors)
WF.set_values(values)
ekinlist = []
epotlist = []
# For each component of |\Psi>
values = WF.get_values()
for index, item in enumerate(values):
# tmp is the Vector (0, 0, 0, \psi_i, 0, 0, ...)
tmp = [ zeros(item.shape) for z in xrange(nst) ]
tmp[index] = item
# Project this vector to the canonical basis
tmp = Potential.project_to_canonical(grid, tmp, eigenvectors)
WF.set_values(tmp)
# And calculate the energies of these components
ekinlist.append(WF.kinetic_energy(opT, summed=True))
epotlist.append(WF.potential_energy(opV, summed=True))
iom.save_energy((ekinlist, epotlist), timestep=step, blockid=blockid)
def plot_frames(iom, blockid=0, view=None, imgsize=(12,9)):
"""Plot the phase of a wavefunction for a series of timesteps.
:param iom: An ``IOManager`` instance providing the simulation data.
:param view: The aspect ratio.
"""
parameters = iom.load_parameters()
grid = iom.load_grid(blockid="global")
timegrid = iom.load_wavefunction_timegrid(blockid=blockid)
# Precompute eigenvectors for efficiency
Potential = PotentialFactory().create_potential(parameters)
eigenvectors = Potential.evaluate_eigenvectors_at(grid)
for step in timegrid:
print(" Plotting frame of timestep # " + str(step))
wave = iom.load_wavefunction(blockid=blockid, timestep=step)
values = [ wave[j,...] for j in xrange(parameters["ncomponents"]) ]
# Transform the values to the eigenbasis
# TODO: improve this:
if parameters["algorithm"] == "fourier":
ve = Potential.project_to_eigen(grid, values, eigenvectors)
else:
ve = values
fig = figure(figsize=imgsize)
for index, component in enumerate(ve):
ax = fig.add_subplot(parameters["ncomponents"],1,index+1)
ax.ticklabel_format(style="sci", scilimits=(0,0), axis="y")
# Plot the wavefunction
ax.plot(grid, component*conj(component), color="gray")
ax.set_ylabel(r"$\langle \varphi_"+str(index)+r"| \varphi_"+str(index)+r"\rangle$")
ax.set_xlabel(r"$x$")
# Compute the phase from the wavefunction restricted to "important" regions
restr_grid = grid[component*conj(component) > 10e-8]
restr_comp = component[component*conj(component) > 10e-8]
# Plot the phase
ax.plot(restr_grid, angle(restr_comp), "-", color="green")
ax.plot(restr_grid, ComplexMath.continuate(angle(restr_comp)), ".", color="green")
# Set the aspect window
if view is not None:
ax.set_xlim(view[:2])
#ax.set_ylim(view[2:])
fig.suptitle(r"$\arg \Psi$ at time $"+str(step*parameters["dt"])+r"$")
fig.savefig("wavefunction_phase_block"+str(blockid)+"_"+ (7-len(str(step)))*"0"+str(step) +GD.output_format)
close(fig)
print(" Plotting frames finished")
def compute_norm(iom, blockid=0):
"""Compute the norm of a wavepacket timeseries.
:param iom: An ``IOManager`` instance providing the simulation data.
:param blockid: The data block from which the values are read.
"""
parameters = iom.load_parameters()
if iom.has_grid(blockid=blockid):
grid = iom.load_grid(blockid=blockid)
else:
grid = iom.load_grid(blockid="global")
# Number of time steps we saved
timesteps = iom.load_wavefunction_timegrid(blockid=blockid)
nrtimesteps = timesteps.shape[0]
# We want to save norms, thus add a data slot to the data file
iom.add_norm(parameters, timeslots=nrtimesteps, blockid=blockid)
# Precalculate eigenvectors for efficiency
Potential = PotentialFactory().create_potential(parameters)
eigenvectors = Potential.evaluate_eigenvectors_at(grid)
WF = WaveFunction(parameters)
# Iterate over all timesteps
for i, step in enumerate(timesteps):
print(" Computing norms of timestep "+str(step))
values = iom.load_wavefunction(timestep=step, blockid=blockid)
values = [ values[j,...] for j in xrange(parameters["ncomponents"]) ]
# Calculate the norm of the wave functions projected into the eigenbasis
values_e = Potential.project_to_eigen(grid, values, eigenvectors)
WF.set_values(values_e)
norms = WF.get_norm()
iom.save_norm(norms, timestep=step, blockid=blockid)
def load_data(resultspath, which_norm="wf"):
# Group the data from different simulations
ids = FT.get_result_dirs(resultspath)
ids = FT.group_by(ids, "eps")
nsims = FT.get_number_simulations(resultspath)
groupdata = []
axisdata = [ [] for i in xrange(nsims) ]
normdata = [ [] for i in xrange(nsims) ]
iom_f = IOManager()
iom_h = IOManager()
for index, sims in enumerate(ids):
# Sorting based on file names
dirs_f = FT.gather_all(sims, "fourier")
dirs_h = FT.gather_all(sims, "hagedorn")
if len(dirs_f) != len(dirs_h):
raise ValueError("Found different number of fourier and hagedorn simulations!")
dirs_f = FT.sort_by(dirs_f, "eps", as_string=True)
dirs_h = FT.sort_by(dirs_h, "eps", as_string=True)
# Loop over all simulations
for dir_f, dir_h in zip(dirs_f, dirs_h):
print("Comparing simulation " + dir_h + " with " + dir_f)
resultsfile_f = FT.get_results_file(dir_f)
iom_f.open_file(filename=resultsfile_f)
resultsfile_h = FT.get_results_file(dir_h)
iom_h.open_file(filename=resultsfile_h)
# Read the parameters
parameters_f = iom_f.load_parameters()
parameters_h = iom_h.load_parameters()
grid = iom_f.load_grid(blockid="global")
# Precalculate eigenvectors for efficiency
Potential = PotentialFactory().create_potential(parameters_f)
eigenvectors = Potential.evaluate_eigenvectors_at(grid)
# Get the data
# Number of time steps we saved
timesteps = iom_f.load_wavefunction_timegrid()
# Scalar parameter that discriminates the simulations
axisdata[index].append((parameters_f, timesteps))
WF = WaveFunction(parameters_f)
WF.set_grid(grid)
norms = []
for i, step in enumerate(timesteps):
# Load the data that belong to the current timestep
data_f = iom_f.load_wavefunction(timestep=step)
data_h = iom_h.load_wavefunction(timestep=step)
data_f = Potential.project_to_eigen(grid, data_f, eigenvectors)
data_f = array(data_f)
data_diff = data_f - data_h
# Compute the norm || u_f - u_h ||
if which_norm == "wf":
# Rearrange data to fit the input of WF and handle over
WF.set_values([ data_diff[n,:] for n in xrange(parameters_f.ncomponents) ])
curnorm = WF.get_norm()
# More than one component? If yes, compute also the overall norm
if parameters_f.ncomponents > 1:
nosum = WF.get_norm(summed=True)
curnorm = list(curnorm) + [nosum]
elif which_norm == "max":
curnorm = [ max( abs(data_diff[n,:]) ) for n in xrange(parameters_f.ncomponents) ]
# More than one component? If yes, compute also the overall norm
if parameters_f.ncomponents > 1:
nosum = max(curnorm)
curnorm = list(curnorm) + [nosum]
print(" at time " + str(step*parameters_f.dt) + " the error norm is " + str(curnorm))
norms.append(curnorm)
# Append norm values to global data structure
norms = array(norms)
normdata[index].append(norms)
# Scalar parameter of the different curves
# We add this here because the simulation parameters are
# already loaded but not overwritten yet be the next iteration
# Remember: we need only a single epsilon out of each eps_group.
groupdata.append(parameters_f.dt)
iom_f.finalize()
iom_h.finalize()
return (groupdata, axisdata, normdata)
potname = pot.replace("_","\_")
if type(potdef["potential"]) == str:
potformula = latex(sympify(potdef["potential"]))
else:
potformula = latex(Matrix(sympify(potdef["potential"])))
if potdef.has_key("defaults"):
potdefaults = potdef["defaults"]
else:
potdefaults = {}
# Create the potential "the right way"
params["potential"] = pot
P = PotentialFactory().create_potential(params)
y = P.evaluate_eigenvalues_at(x)
# Plot the potential
figure()
for yvals in y:
plot(x, yvals)
xlabel(r"$x$")
ylabel(r"$\lambda_i\left(x\right)$")
xlim(min(x), max(x))
savefig(pot + ".pdf")
# The latex code
ls = """
\\begin{minipage}{0.5\\linewidth}
Name: \\texttt{""" + potname + """}
def plot_frames(iom, gid=0, view=None, plotphase=True, plotcomponents=False, plotabssqr=False, imgsize=(12,9)):
"""Plot the wave function for a series of timesteps.
:param iom: An ``IOManager`` instance providing the simulation data.
:param gid: The group ID of the group where the two packets are stored.
:param view: The aspect ratio.
:param plotphase: Whether to plot the complex phase. (slow)
:param plotcomponents: Whether to plot the real/imaginary parts..
:param plotabssqr: Whether to plot the absolute value squared.
"""
parameters = iom.load_parameters()
# Block IDs for mother and child wavepacket
bidm, bidc = iom.get_block_ids(groupid=gid)
grid = iom.load_grid(blockid="global")
# Precompute eigenvectors for efficiency
Potential = PotentialFactory().create_potential(parameters)
eigenvectors = Potential.evaluate_eigenvectors_at(grid)
timegrid = iom.load_wavefunction_timegrid(blockid=bidm)
for step in timegrid:
print(" Timestep # " + str(step))
# Retrieve spawn data for both packets
try:
wave_m = iom.load_wavefunction(timestep=step, blockid=bidm)
values_m = [ wave_m[j,...] for j in xrange(parameters["ncomponents"]) ]
have_mother_data = True
except ValueError:
have_mother_data = False
# Retrieve spawn data
try:
wave_s = iom.load_wavefunction(timestep=step, blockid=bidc)
values_s = [ wave_s[j,...] for j in xrange(parameters["ncomponents"]) ]
have_spawn_data = True
except ValueError:
have_spawn_data = False
# Plot the probability densities projected to the eigenbasis
fig = figure(figsize=imgsize)
# Create a bunch of subplots
axes = []
for index, component in enumerate(values_m):
ax = fig.add_subplot(parameters["ncomponents"],1,index+1)
ax.ticklabel_format(style="sci", scilimits=(0,0), axis="y")
axes.append(ax)
# Plot original Wavefunction
if have_mother_data is True:
for index, component in enumerate(values_m):
if plotcomponents is True:
axes[index].plot(grid, real(component))
axes[index].plot(grid, imag(component))
axes[index].set_ylabel(r"$\Re \varphi_"+str(index)+r", \Im \varphi_"+str(index)+r"$")
if plotabssqr is True:
axes[index].plot(grid, component*conj(component))
axes[index].set_ylabel(r"$\langle \varphi_"+str(index)+r"| \varphi_"+str(index)+r"\rangle$")
if plotphase is True:
plotcf(grid, angle(component), component*conj(component))
axes[index].set_ylabel(r"$\langle \varphi_"+str(index)+r"| \varphi_"+str(index)+r"\rangle$")
# Overlay spawned parts
if have_spawn_data is True:
for index, component in enumerate(values_s):
axes[index].plot(grid, component*conj(component), "-r")
axes[index].set_ylabel(r"$\langle \varphi_"+str(index)+r"| \varphi_"+str(index)+r"\rangle$")
# Set the axis properties
for index in xrange(len(values_m)):
axes[index].set_xlabel(r"$x$")
# Set the aspect window
if view is not None:
axes[index].set_xlim(view[:2])
axes[index].set_ylim(view[2:])
fig.suptitle(r"$\Psi$ at time $"+str(step*parameters["dt"])+r"$")
fig.savefig("wavefunction_"+ (5-len(str(step)))*"0"+str(step) +GD.output_format)
close(fig)
print(" Plotting frames finished")
def plot_frames_split(iom, gid=0, view=None, plotphase=True, plotcomponents=False, plotabssqr=False, imgsize=(12,9)):
"""Plot the wave function for a series of timesteps.
:param iom: An ``IOManager`` instance providing the simulation data.
:param gid: The group ID of the group where the two packets are stored.
:param view: The aspect ratio.
:param plotphase: Whether to plot the complex phase. (slow)
:param plotcomponents: Whether to plot the real/imaginary parts..
:param plotabssqr: Whether to plot the absolute value squared.
"""
parameters = iom.load_parameters()
n = parameters["ncomponents"]
# Block IDs for mother and child wavepacket
bidm, bidc = iom.get_block_ids(groupid=gid)
grid = iom.load_grid(blockid="global")
# Precompute eigenvectors for efficiency
Potential = PotentialFactory().create_potential(parameters)
eigenvectors = Potential.evaluate_eigenvectors_at(grid)
timegrid = iom.load_wavefunction_timegrid(blockid=bidm)
for step in timegrid:
print(" Timestep # " + str(step))
# Split grid
gl = grid[grid<=X0]
gr = grid[grid>X0]
# Retrieve spawn data for both packets and split the data as necessary
try:
wave_m = iom.load_wavefunction(timestep=step, blockid=bidm)
values_m = [ wave_m[j,...] for j in xrange(parameters["ncomponents"]) ]
yl = values_m[0][grid<=X0]
yr = values_m[0][grid>X0]
have_mother_data = True
except ValueError:
have_mother_data = False
# Retrieve spawn data
try:
wave_s = iom.load_wavefunction(timestep=step, blockid=bidc)
values_s = [ wave_s[j,...] for j in xrange(parameters["ncomponents"]) ]
ysl = values_s[0][grid<=X0]
ysr = values_s[0][grid>X0]
have_spawn_data = True
except ValueError:
have_spawn_data = False
# Plot the probability densities projected to the eigenbasis
fig = figure(figsize=imgsize)
# Plot the probability density, left to X0
ax1 = fig.add_subplot(1,2,1)
ax1.ticklabel_format(style="sci", scilimits=(0,0), axis="y")
# mother
if have_mother_data is True:
plotcf(gl, angle(yl), conj(yl)*yl)
# spawned
if have_spawn_data is True:
plot(gl, conj(ysl)*ysl, "-r")
if view is not None:
ax1.set_xlim(view[0],0)
ax1.set_ylim(view[2:4])
ax1.set_xlabel(r"$x \le 0$")
ax1.set_ylabel(r"$\langle\varphi |\varphi \rangle$")
# Plot the probability density, right to X0
ax2 = fig.add_subplot(1,2,2)
ax2.ticklabel_format(style="sci", scilimits=(0,0), axis="y")
# mother
if have_mother_data is True:
plotcf(gr, angle(yr), conj(yr)*yr)
# spawned
if have_spawn_data is True:
plot(gr, conj(ysr)*ysr, "-r")
if view is not None:
ax2.set_xlim(0, view[1])
ax2.set_ylim(view[4:])
ax2.set_xlabel(r"$x > 0$")
ax2.set_ylabel(r"$\langle\varphi |\varphi \rangle$")
fig.suptitle(r"Time $"+str(step*parameters["dt"])+r"$")
fig.savefig("wavepackets_"+ (5-len(str(step)))*"0"+str(step) +GD.output_format)
close(fig)
print(" Plotting frames finished")