def plot_frames(
iom,
blockid=0
): #, view=None, plotphase=True, plotcomponents=False, plotabssqr=False, imgsize=(12,9)):
parameters = iom.load_parameters()
if not parameters["dimension"] == 2:
print("No wavefunction of two space dimensions, silent return!")
return
G = BlockFactory().create_grid(parameters)
V = BlockFactory().create_potential(parameters)
print(G.get_extensions())
WF = WaveFunction(parameters)
WF.set_grid(G)
BT = BasisTransformationWF(V)
BT.set_grid(G)
timegrid = iom.load_wavefunction_timegrid(blockid=blockid)
u, v = G.get_nodes(split=True, flat=False)
u = real(u)
v = real(v)
N = WF.get_number_components()
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"])]
WF.set_values(values)
# Transform the values to the eigenbasis
# TODO: improve this:
if parameters["algorithm"] == "fourier":
BT.transform_to_eigen(WF)
else:
pass
Psi = WF.get_values()
fig = figure()
for level in xrange(N):
z = Psi[level]
subplot(N, 1, level + 1)
plotcm(z, darken=0.3)
savefig("wavefunction_level_" + str(level) + "_timestep_" +
(5 - len(str(step))) * "0" + str(step) + ".png")
close(fig)
print(" Plotting frames finished")
def plot_frames(iom, blockid=0):
parameters = iom.load_parameters()
if not parameters["dimension"] == 2:
print("No wavefunction of two space dimensions, silent return!")
return
G = BlockFactory().create_grid(parameters)
V = BlockFactory().create_potential(parameters)
WF = WaveFunction(parameters)
WF.set_grid(G)
BT = BasisTransformationWF(V)
BT.set_grid(G)
timegrid = iom.load_wavefunction_timegrid(blockid=blockid)
u, v = G.get_nodes(split=True, flat=False)
u = real(u)
v = real(v)
N = WF.get_number_components()
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"]) ]
WF.set_values(values)
# Transform the values to the eigenbasis
# TODO: improve this:
if parameters["algorithm"] == "fourier":
BT.transform_to_eigen(WF)
else:
pass
Psi = WF.get_values()
for level in xrange(N):
z = Psi[level]
# Plot the probability densities projected to the eigenbasis
fig = mlab.figure(size=(800,700))
surfcf(u, v, angle(z), abs(z))
#mlab.contour_surf(u, v, abs(z))
#mlab.outline()
#mlab.axes()
mlab.savefig("wavefunction_level_"+str(level)+"_timestep_"+(5-len(str(step)))*"0"+str(step)+".png")
mlab.close(fig)
print(" Plotting frames finished")
def plot_frames(iom, blockid=0):#, view=None, plotphase=True, plotcomponents=False, plotabssqr=False, imgsize=(12,9)):
parameters = iom.load_parameters()
if not parameters["dimension"] == 2:
print("No wavefunction of two space dimensions, silent return!")
return
G = BlockFactory().create_grid(parameters)
V = BlockFactory().create_potential(parameters)
print(G.get_extensions())
WF = WaveFunction(parameters)
WF.set_grid(G)
BT = BasisTransformationWF(V)
BT.set_grid(G)
timegrid = iom.load_wavefunction_timegrid(blockid=blockid)
u, v = G.get_nodes(split=True, flat=False)
u = real(u)
v = real(v)
N = WF.get_number_components()
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"]) ]
WF.set_values(values)
# Transform the values to the eigenbasis
# TODO: improve this:
if parameters["algorithm"] == "fourier":
BT.transform_to_eigen(WF)
else:
pass
Psi = WF.get_values()
fig = figure()
for level in xrange(N):
z = Psi[level]
subplot(N,1,level+1)
plotcm(z, darken=0.3)
savefig("wavefunction_level_"+str(level)+"_timestep_"+(5-len(str(step)))*"0"+str(step)+".png")
close(fig)
print(" Plotting frames finished")
def transform_wavefunction_to_eigen(iomin, iomout, blockidin=0, blockidout=0):
"""Compute the transformation to the eigenbasis for a wavefunction.
Save the result back to a file.
:param iomin: An :py:class:`IOManager: instance providing the simulation data.
:param iomout: An :py:class:`IOManager: instance for saving the transformed data.
:param blockidin: The data block from which the values are read. Default is `0`.
:param blockidout: The data block to which the values are written. Default is `0`.
"""
parameters = iomin.load_parameters()
# Number of time steps we saved
timesteps = iomin.load_wavefunction_timegrid(blockid=blockidin)
nrtimesteps = timesteps.shape[0]
iomout.add_wavefunction(parameters,
timeslots=nrtimesteps,
blockid=blockidout)
# The grid on the domain
grid = BlockFactory().create_grid(parameters)
# The potential used
Potential = BlockFactory().create_potential(parameters)
# Basis transformator
BT = BasisTransformationWF(Potential)
BT.set_grid(grid)
# And two empty wavefunctions
WF = WaveFunction(parameters)
WF.set_grid(grid)
# Iterate over all timesteps
for i, step in enumerate(timesteps):
print(" Compute eigentransform at timestep # " + str(step))
# Retrieve simulation data
values = iomin.load_wavefunction(timestep=step, blockid=blockidin)
values = [values[j, ...] for j in xrange(parameters["ncomponents"])]
WF.set_values(values)
# Project wavefunction values to eigenbasis
BT.transform_to_eigen(WF)
# Save the transformed values
iomout.save_wavefunction(WF.get_values(),
timestep=step,
blockid=blockidout)
def compute_norm(iom, blockid=0, eigentrafo=True):
"""Compute the norm of a wavefunction timeseries.
:param iom: An :py:class:`IOManager` instance providing the simulation data.
:param blockid: The data block from which the values are read. Default is `0`.
:param eigentrafo: Whether to make a transformation into the eigenbasis.
:type eigentrafo: Boolean, default is ``True``.
"""
parameters = iom.load_parameters()
# Number of time steps we saved
timesteps = iom.load_wavefunction_timegrid(blockid=blockid)
nrtimesteps = timesteps.shape[0]
# Construct the grid from the parameters
grid = BlockFactory().create_grid(parameters)
# Basis transformator
if eigentrafo is True:
# The potential used
Potential = BlockFactory().create_potential(parameters)
BT = BasisTransformationWF(Potential)
BT.set_grid(grid)
# And two empty wavefunctions
WF = WaveFunction(parameters)
WF.set_grid(grid)
# We want to save norms, thus add a data slot to the data file
iom.add_norm(parameters, timeslots=nrtimesteps, blockid=blockid)
# Iterate over all timesteps
for i, step in enumerate(timesteps):
print(" Computing norms of timestep " + str(step))
# Retrieve simulation data
values = iom.load_wavefunction(timestep=step, blockid=blockid)
values = [values[j, ...] for j in xrange(parameters["ncomponents"])]
WF.set_values(values)
# Project wavefunction values to eigenbasis
if eigentrafo is True:
BT.transform_to_eigen(WF)
# Calculate the norm of the wave functions projected into the eigenbasis
norms = WF.norm()
iom.save_norm(norms, timestep=step, blockid=blockid)
def compute_norm(iom, blockid=0, eigentrafo=True):
"""Compute the norm of a wavefunction timeseries.
:param iom: An :py:class:`IOManager` instance providing the simulation data.
:param blockid: The data block from which the values are read. Default is `0`.
:param eigentrafo: Whether to make a transformation into the eigenbasis.
:type eigentrafo: Boolean, default is ``True``.
"""
parameters = iom.load_parameters()
# Number of time steps we saved
timesteps = iom.load_wavefunction_timegrid(blockid=blockid)
nrtimesteps = timesteps.shape[0]
# Construct the grid from the parameters
grid = BlockFactory().create_grid(parameters)
# Basis transformator
if eigentrafo is True:
# The potential used
Potential = BlockFactory().create_potential(parameters)
BT = BasisTransformationWF(Potential)
BT.set_grid(grid)
# And two empty wavefunctions
WF = WaveFunction(parameters)
WF.set_grid(grid)
# We want to save norms, thus add a data slot to the data file
iom.add_norm(parameters, timeslots=nrtimesteps, blockid=blockid)
# Iterate over all timesteps
for i, step in enumerate(timesteps):
print(" Computing norms of timestep %d" % step)
# Retrieve simulation data
values = iom.load_wavefunction(timestep=step, blockid=blockid)
values = [values[j, ...] for j in range(parameters["ncomponents"])]
WF.set_values(values)
# Project wavefunction values to eigenbasis
if eigentrafo is True:
BT.transform_to_eigen(WF)
# Calculate the norm of the wave functions projected into the eigenbasis
norms = WF.norm()
iom.save_norm(norms, timestep=step, blockid=blockid)
def transform_wavefunction_to_eigen(iomin, iomout, blockidin=0, blockidout=0):
"""Compute the transformation to the eigenbasis for a wavefunction.
Save the result back to a file.
:param iomin: An :py:class:`IOManager: instance providing the simulation data.
:param iomout: An :py:class:`IOManager: instance for saving the transformed data.
:param blockidin: The data block from which the values are read. Default is `0`.
:param blockidout: The data block to which the values are written. Default is `0`.
"""
parameters = iomin.load_parameters()
# Number of time steps we saved
timesteps = iomin.load_wavefunction_timegrid(blockid=blockidin)
nrtimesteps = timesteps.shape[0]
iomout.add_wavefunction(parameters, timeslots=nrtimesteps, blockid=blockidout)
# The grid on the domain
grid = BlockFactory().create_grid(parameters)
# The potential used
Potential = BlockFactory().create_potential(parameters)
# Basis transformator
BT = BasisTransformationWF(Potential)
BT.set_grid(grid)
# And two empty wavefunctions
WF = WaveFunction(parameters)
WF.set_grid(grid)
# Iterate over all timesteps
for i, step in enumerate(timesteps):
print(" Compute eigentransform at timestep %d" % step)
# Retrieve simulation data
values = iomin.load_wavefunction(timestep=step, blockid=blockidin)
values = [values[j, ...] for j in range(parameters["ncomponents"])]
WF.set_values(values)
# Project wavefunction values to eigenbasis
BT.transform_to_eigen(WF)
# Save the transformed values
iomout.save_wavefunction(WF.get_values(), timestep=step, blockid=blockidout)
def compute_autocorrelation(iom, obsconfig=None, blockid=0, eigentrafo=True):
"""Compute the autocorrelation of a wavefunction timeseries.
:param iom: An :py:class:`IOManager` instance providing the simulation data.
:param obsconfig: Configuration parameters describing f.e. the inner product to use.
:type obsconfig: A :py:class:`ParameterProvider` instance.
Value has no effect in this class.
:param blockid: The data block from which the values are read.
:type blockid: Integer, Default is ``0``
:param eigentrafo: Whether to make a transformation into the eigenbasis.
:type eigentrafo: Boolean, default is ``True``.
"""
parameters = iom.load_parameters()
# Number of time steps we saved
timesteps = iom.load_wavefunction_timegrid(blockid=blockid)
nrtimesteps = timesteps.shape[0]
# Construct the grid from the parameters
grid = BlockFactory().create_grid(parameters)
# Basis transformator
if eigentrafo is True:
# The potential used
Potential = BlockFactory().create_potential(parameters)
BT = BasisTransformationWF(Potential)
BT.set_grid(grid)
# And two empty wavefunctions
WFo = WaveFunction(parameters)
WFo.set_grid(grid)
WFt = WaveFunction(parameters)
WFt.set_grid(grid)
# We want to save norms, thus add a data slot to the data file
iom.add_autocorrelation(parameters, timeslots=nrtimesteps, blockid=blockid)
# Preconfigure the
values = iom.load_wavefunction(timestep=0, blockid=blockid)
values = [values[j, ...] for j in range(parameters["ncomponents"])]
WFo.set_values(values)
# Project wavefunction values to eigenbasis
if eigentrafo is True:
BT.transform_to_eigen(WFo)
# Fourier transform the values
WFo.set_values([fftn(value) for value in WFo.get_values()])
# Iterate over all timesteps
for i, step in enumerate(timesteps):
print(" Computing autocorrelations of timestep %d" % step)
# Retrieve simulation data
values = iom.load_wavefunction(timestep=step, blockid=blockid)
values = [values[j, ...] for j in range(parameters["ncomponents"])]
WFt.set_values(values)
# Project wavefunction values to eigenbasis
if eigentrafo is True:
BT.transform_to_eigen(WFt)
# Fourier transform the values
WFt.set_values([fftn(value) for value in WFt.get_values()])
# Compute the prefactor
T = grid.get_extensions()
N = grid.get_number_nodes()
prefactor = product(array(T) / array(N).astype(floating)**2)
# Compute the autocorrelation
# TODO: Consider splitting into cases `fft` versus `fftn`
valueso = WFo.get_values()
valuest = WFt.get_values()
acs = [
prefactor * ifftn(sum(conjugate(valueso[n]) * valuest[n]))
for n in range(parameters["ncomponents"])
]
iom.save_autocorrelation(acs, timestep=step, blockid=blockid)
def plot_frames(PP,
iom,
blockid=0,
load=False,
eigentransform=False,
timerange=None,
view=None,
path='.'):
"""Plot the wave function for a series of timesteps.
:param iom: An :py:class:`IOManager` instance providing the simulation data.
"""
parameters = iom.load_parameters()
if not parameters["dimension"] == 2:
print("No wavefunction of two space dimensions, silent return!")
return
if PP is None:
PP = parameters
if load is True:
# TODO: Implement reshaping
raise NotImplementedError("Loading of 2D grids is not implemented")
else:
G = BlockFactory().create_grid(PP)
if eigentransform:
V = BlockFactory().create_potential(parameters)
BT = BasisTransformationWF(V)
BT.set_grid(G)
WF = WaveFunction(parameters)
WF.set_grid(G)
N = WF.get_number_components()
timegrid = iom.load_wavefunction_timegrid(blockid=blockid)
if timerange is not None:
if len(timerange) == 1:
I = (timegrid == timerange)
else:
I = ((timegrid >= timerange[0]) & (timegrid <= timerange[1]))
if any(I):
timegrid = timegrid[I]
else:
raise ValueError("No valid timestep remains!")
u, v = G.get_axes()
u = real(u.reshape(-1))
v = real(v.reshape(-1))
# View
if view is not None:
if view[0] is None:
view[0] = u.min()
if view[1] is None:
view[1] = u.max()
if view[2] is None:
view[2] = v.min()
if view[3] is None:
view[3] = v.max()
for step in timegrid:
print(" Plotting frame of timestep # {}".format(step))
# Load the data
wave = iom.load_wavefunction(blockid=blockid, timestep=step)
values = [wave[j, ...] for j in range(parameters["ncomponents"])]
WF.set_values(values)
# Transform the values to the eigenbasis
if eigentransform:
BT.transform_to_eigen(WF)
Psi = WF.get_values()
# Plot
fig = figure()
for level in range(N):
# Wavefunction data
z = Psi[level]
z = z.reshape(G.get_number_nodes())
fig.add_subplot(N, 1, level + 1)
plotcf2d(u, v, z, darken=0.3, limits=view)
fig.savefig(
os.path.join(
path,
"wavefunction_contour_block_%s_level_%d_timestep_%07d.png" %
(blockid, level, step)))
close(fig)
def compute_autocorrelation(iom, obsconfig=None, blockid=0, eigentrafo=True):
"""Compute the autocorrelation of a wavefunction timeseries.
:param iom: An :py:class:`IOManager` instance providing the simulation data.
:param obsconfig: Configuration parameters describing f.e. the inner product to use.
:type obsconfig: A :py:class:`ParameterProvider` instance.
Value has no effect in this class.
:param blockid: The data block from which the values are read.
:type blockid: Integer, Default is ``0``
:param eigentrafo: Whether to make a transformation into the eigenbasis.
:type eigentrafo: Boolean, default is ``True``.
"""
parameters = iom.load_parameters()
# Number of time steps we saved
timesteps = iom.load_wavefunction_timegrid(blockid=blockid)
nrtimesteps = timesteps.shape[0]
# Construct the grid from the parameters
grid = BlockFactory().create_grid(parameters)
# Basis transformator
if eigentrafo is True:
# The potential used
Potential = BlockFactory().create_potential(parameters)
BT = BasisTransformationWF(Potential)
BT.set_grid(grid)
# And two empty wavefunctions
WFo = WaveFunction(parameters)
WFo.set_grid(grid)
WFt = WaveFunction(parameters)
WFt.set_grid(grid)
# We want to save norms, thus add a data slot to the data file
iom.add_autocorrelation(parameters, timeslots=nrtimesteps, blockid=blockid)
# Preconfigure the
values = iom.load_wavefunction(timestep=0, blockid=blockid)
values = [values[j, ...] for j in range(parameters["ncomponents"])]
WFo.set_values(values)
# Project wavefunction values to eigenbasis
if eigentrafo is True:
BT.transform_to_eigen(WFo)
# Fourier transform the values
WFo.set_values([fftn(value) for value in WFo.get_values()])
# Iterate over all timesteps
for i, step in enumerate(timesteps):
print(" Computing autocorrelations of timestep %d" % step)
# Retrieve simulation data
values = iom.load_wavefunction(timestep=step, blockid=blockid)
values = [values[j, ...] for j in range(parameters["ncomponents"])]
WFt.set_values(values)
# Project wavefunction values to eigenbasis
if eigentrafo is True:
BT.transform_to_eigen(WFt)
# Fourier transform the values
WFt.set_values([fftn(value) for value in WFt.get_values()])
# Compute the prefactor
T = grid.get_extensions()
N = grid.get_number_nodes()
prefactor = product(array(T) / array(N).astype(floating)**2)
# Compute the autocorrelation
# TODO: Consider splitting into cases `fft` versus `fftn`
valueso = WFo.get_values()
valuest = WFt.get_values()
acs = [prefactor * ifftn(sum(conjugate(valueso[n]) * valuest[n])) for n in range(parameters["ncomponents"])]
iom.save_autocorrelation(acs, timestep=step, blockid=blockid)
def plot_frames(PP, iom, blockid=0, load=False, eigentransform=False, timerange=None, sparsify=10, view=None, interactive=False, path='.'):
"""Plot the wave function for a series of timesteps.
:param iom: An :py:class:`IOManager` instance providing the simulation data.
"""
parameters = iom.load_parameters()
if not parameters["dimension"] == 2:
print("No wavefunction of two space dimensions, silent return!")
return
if PP is None:
PP = parameters
if load is True:
# TODO: Implement reshaping
raise NotImplementedError("Loading of 2D grids is not implemented")
else:
G = BlockFactory().create_grid(PP)
if eigentransform:
V = BlockFactory().create_potential(parameters)
BT = BasisTransformationWF(V)
BT.set_grid(G)
WF = WaveFunction(parameters)
WF.set_grid(G)
N = WF.get_number_components()
timegrid = iom.load_wavefunction_timegrid(blockid=blockid)
if timerange is not None:
if len(timerange) == 1:
I = (timegrid == timerange)
else:
I = ((timegrid >= timerange[0]) & (timegrid <= timerange[1]))
if any(I):
timegrid = timegrid[I]
else:
raise ValueError("No valid timestep remains!")
u, v = G.get_nodes(split=True, flat=False)
u = real(u[::sparsify, ::sparsify])
v = real(v[::sparsify, ::sparsify])
# View
if view is not None:
if view[0] is None:
view[0] = u.min()
if view[1] is None:
view[1] = u.max()
if view[2] is None:
view[2] = v.min()
if view[3] is None:
view[3] = v.max()
for step in timegrid:
print(" Plotting frame of timestep # {}".format(step))
# Load the data
wave = iom.load_wavefunction(blockid=blockid, timestep=step)
values = [wave[j, ...] for j in range(parameters["ncomponents"])]
WF.set_values(values)
# Transform the values to the eigenbasis
if eigentransform:
BT.transform_to_eigen(WF)
Psi = WF.get_values()
for level in range(N):
# Wavefunction data
z = Psi[level]
z = z.reshape(G.get_number_nodes())[::sparsify, ::sparsify]
# View
if view is not None:
if view[4] is None:
view[4] = 0.0
if view[5] is None:
view[5] = 1.1 * abs(z).max()
# Plot
# if not interactive:
# mlab.options.offscreen = True
fig = mlab.figure(size=(800, 700))
surfcf(u, v, angle(z), abs(z), view=view)
mlab.draw()
if interactive:
mlab.show()
else:
mlab.savefig(os.path.join("wavefunction_surface_block_%s_level_%d_timestep_%07d.png" % (blockid, level, step)))
mlab.close(fig)
def compute_energy(iom, blockid=0, eigentrafo=True, iseigen=True):
"""
:param iom: An :py:class:`IOManager: instance providing the simulation data.
:param blockid: The data block from which the values are read. Default is `0`.
:param eigentrafo: Whether to make a transformation into the eigenbasis.
:type eigentrafo: Boolean, default is ``True``.
:param iseigen: Whether the data is assumed to be in the eigenbasis.
:type iseigen: Boolean, default is ``True``
"""
parameters = iom.load_parameters()
# Number of time steps we saved
timesteps = iom.load_wavefunction_timegrid(blockid=blockid)
nrtimesteps = timesteps.shape[0]
# Construct grid from the parameters
grid = BlockFactory().create_grid(parameters)
# The potential used
Potential = BlockFactory().create_potential(parameters)
# The operators
KO = KineticOperator(grid)
KO.calculate_operator(parameters["eps"])
opT = KO
if eigentrafo is True:
opV = Potential.evaluate_at(grid)
else:
if iseigen is True:
opV = Potential.evaluate_eigenvalues_at(grid, as_matrix=True)
else:
opV = Potential.evaluate_at(grid, as_matrix=True)
# Basis transformator
if eigentrafo is True:
BT = BasisTransformationWF(Potential)
BT.set_grid(grid)
# And two empty wavefunctions
WF = WaveFunction(parameters)
WF.set_grid(grid)
WF2 = WaveFunction(parameters)
WF2.set_grid(grid)
# We want to save norms, thus add a data slot to the data file
iom.add_energy(parameters, timeslots=nrtimesteps, blockid=blockid)
nst = Potential.get_number_components()
if eigentrafo is True:
# Iterate over all timesteps
for i, step in enumerate(timesteps):
print(" Computing energies of timestep # " + str(step))
# Retrieve simulation data
values = iom.load_wavefunction(timestep=step, blockid=blockid)
values = [ values[j,...] for j in xrange(parameters["ncomponents"]) ]
WF.set_values(values)
# Project wavefunction values to eigenbasis
BT.transform_to_eigen(WF)
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
WF2.set_values(tmp)
# Project this vector to the canonical basis
BT.transform_to_canonical(WF2)
# And calculate the energies of these components
ekinlist.append(WF2.kinetic_energy(opT, summed=True))
epotlist.append(WF2.potential_energy(opV, summed=True))
iom.save_energy((ekinlist, epotlist), timestep=step, blockid=blockid)
else:
# Iterate over all timesteps
for i, step in enumerate(timesteps):
print(" Computing energies of timestep # " + str(step))
# Retrieve simulation data
values = iom.load_wavefunction(timestep=step, blockid=blockid)
values = [ values[j,...] for j in xrange(parameters["ncomponents"]) ]
WF.set_values(values)
# And calculate the energies of these components
ekinlist = WF.kinetic_energy(opT, summed=False)
epotlist = WF.potential_energy(opV, summed=False)
iom.save_energy((ekinlist, epotlist), timestep=step, blockid=blockid)
def plot_frames(PP, iom, blockid=0, load=False, eigentransform=False, timerange=None, view=None, path='.'):
"""Plot the wave function for a series of timesteps.
:param iom: An :py:class:`IOManager` instance providing the simulation data.
"""
parameters = iom.load_parameters()
if not parameters["dimension"] == 2:
print("No wavefunction of two space dimensions, silent return!")
return
if PP is None:
PP = parameters
if load is True:
# TODO: Implement reshaping
raise NotImplementedError("Loading of 2D grids is not implemented")
else:
G = BlockFactory().create_grid(PP)
if eigentransform:
V = BlockFactory().create_potential(parameters)
BT = BasisTransformationWF(V)
BT.set_grid(G)
WF = WaveFunction(parameters)
WF.set_grid(G)
N = WF.get_number_components()
timegrid = iom.load_wavefunction_timegrid(blockid=blockid)
if timerange is not None:
if len(timerange) == 1:
I = (timegrid == timerange)
else:
I = ((timegrid >= timerange[0]) & (timegrid <= timerange[1]))
if any(I):
timegrid = timegrid[I]
else:
raise ValueError("No valid timestep remains!")
u, v = G.get_axes()
u = real(u.reshape(-1))
v = real(v.reshape(-1))
# View
if view is not None:
if view[0] is None:
view[0] = u.min()
if view[1] is None:
view[1] = u.max()
if view[2] is None:
view[2] = v.min()
if view[3] is None:
view[3] = v.max()
for step in timegrid:
print(" Plotting frame of timestep # {}".format(step))
# Load the data
wave = iom.load_wavefunction(blockid=blockid, timestep=step)
values = [wave[j, ...] for j in range(parameters["ncomponents"])]
WF.set_values(values)
# Transform the values to the eigenbasis
if eigentransform:
BT.transform_to_eigen(WF)
Psi = WF.get_values()
# Plot
fig = figure()
for level in range(N):
# Wavefunction data
z = Psi[level]
z = z.reshape(G.get_number_nodes())
fig.add_subplot(N, 1, level + 1)
plotcf2d(u, v, z, darken=0.3, limits=view)
fig.savefig(os.path.join(path, "wavefunction_contour_block_%s_level_%d_timestep_%07d.png" % (blockid, level, step)))
close(fig)
def compute_evaluate_wavepackets(pp, iom, blockid=0, eigentrafo=True):
"""Evaluate a homogeneous Hagedorn wavepacket on a given grid for each timestep.
:param pp: An :py:class:`ParameterProvider` instance providing the grid data.
:param iom: An :py:class:`IOManager` instance providing the simulation data.
:param blockid: The data block from which the values are read.
:param eigentrafo: Whether or not do an eigentransformation before evaluation is done.
"""
parameters = iom.load_parameters()
if pp is None:
pp = parameters
# Number of time steps we saved
timesteps = iom.load_wavepacket_timegrid(blockid=blockid)
nrtimesteps = timesteps.shape[0]
# Prepare the potential for basis transformations
Potential = BlockFactory().create_potential(parameters)
grid = BlockFactory().create_grid(pp)
# We want to save wavefunctions, thus add a data slot to the data file
d = {"ncomponents": parameters["ncomponents"],
"number_nodes": pp["number_nodes"],
"dimension": parameters["dimension"]}
iom.add_grid(d, blockid=blockid)
iom.add_wavefunction(d, timeslots=nrtimesteps, flat=True, blockid=blockid)
iom.save_grid(grid.get_nodes(), blockid=blockid)
# Initialize a Hagedorn wavepacket with the data
descr = iom.load_wavepacket_description(blockid=blockid)
HAWP = BlockFactory().create_wavepacket(descr)
# Basis transformator
if eigentrafo is True:
BT = BasisTransformationHAWP(Potential)
BT.set_matrix_builder(HAWP.get_innerproduct())
# Basis shapes
BS_descr = iom.load_wavepacket_basisshapes(blockid=blockid)
BS = {}
for ahash, descr in BS_descr.items():
BS[ahash] = BlockFactory().create_basis_shape(descr)
WF = WaveFunction(parameters)
WF.set_grid(grid)
# Iterate over all timesteps
for i, step in enumerate(timesteps):
print(" Evaluating homogeneous wavepacket at timestep %d" % step)
# Retrieve simulation data
params = iom.load_wavepacket_parameters(timestep=step, blockid=blockid, key=("q", "p", "Q", "P", "S", "adQ"))
hashes, coeffs = iom.load_wavepacket_coefficients(timestep=step, get_hashes=True, blockid=blockid)
# Configure the wavepacket
HAWP.set_parameters(params, key=("q", "p", "Q", "P", "S", "adQ"))
HAWP.set_basis_shapes([BS[int(ha)] for ha in hashes])
HAWP.set_coefficients(coeffs)
# Transform to the eigenbasis.
if eigentrafo is True:
BT.transform_to_eigen(HAWP)
# Evaluate the wavepacket
values = HAWP.evaluate_at(grid, prefactor=True)
WF.set_values(values)
# Save the wave function
iom.save_wavefunction(WF.get_values(), timestep=step, blockid=blockid)
def compute_energy(iom, blockid=0, eigentrafo=True, iseigen=True):
"""
:param iom: An :py:class:`IOManager: instance providing the simulation data.
:param blockid: The data block from which the values are read. Default is `0`.
:param eigentrafo: Whether to make a transformation into the eigenbasis.
:type eigentrafo: Boolean, default is ``True``.
:param iseigen: Whether the data is assumed to be in the eigenbasis.
:type iseigen: Boolean, default is ``True``
"""
parameters = iom.load_parameters()
# Number of time steps we saved
timesteps = iom.load_wavefunction_timegrid(blockid=blockid)
nrtimesteps = timesteps.shape[0]
# Construct grid from the parameters
grid = BlockFactory().create_grid(parameters)
# The potential used
Potential = BlockFactory().create_potential(parameters)
# The operators
KO = KineticOperator(grid)
KO.calculate_operator(parameters["eps"])
opT = KO
if eigentrafo is True:
opV = Potential.evaluate_at(grid)
else:
if iseigen is True:
opV = Potential.evaluate_eigenvalues_at(grid, as_matrix=True)
else:
opV = Potential.evaluate_at(grid, as_matrix=True)
# Basis transformator
if eigentrafo is True:
BT = BasisTransformationWF(Potential)
BT.set_grid(grid)
# And two empty wavefunctions
WF = WaveFunction(parameters)
WF.set_grid(grid)
WF2 = WaveFunction(parameters)
WF2.set_grid(grid)
# We want to save norms, thus add a data slot to the data file
iom.add_energy(parameters, timeslots=nrtimesteps, blockid=blockid)
nst = Potential.get_number_components()
if eigentrafo is True:
# Iterate over all timesteps
for i, step in enumerate(timesteps):
print(" Computing energies of timestep # " + str(step))
# Retrieve simulation data
values = iom.load_wavefunction(timestep=step, blockid=blockid)
values = [
values[j, ...] for j in xrange(parameters["ncomponents"])
]
WF.set_values(values)
# Project wavefunction values to eigenbasis
BT.transform_to_eigen(WF)
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
WF2.set_values(tmp)
# Project this vector to the canonical basis
BT.transform_to_canonical(WF2)
# And calculate the energies of these components
ekinlist.append(WF2.kinetic_energy(opT, summed=True))
epotlist.append(WF2.potential_energy(opV, summed=True))
iom.save_energy((ekinlist, epotlist),
timestep=step,
blockid=blockid)
else:
# Iterate over all timesteps
for i, step in enumerate(timesteps):
print(" Computing energies of timestep # " + str(step))
# Retrieve simulation data
values = iom.load_wavefunction(timestep=step, blockid=blockid)
values = [
values[j, ...] for j in xrange(parameters["ncomponents"])
]
WF.set_values(values)
# And calculate the energies of these components
ekinlist = WF.kinetic_energy(opT, summed=False)
epotlist = WF.potential_energy(opV, summed=False)
iom.save_energy((ekinlist, epotlist),
timestep=step,
blockid=blockid)
def plot_frames(PP, iom, blockid=0, load=False):
r"""
"""
parameters = iom.load_parameters()
if not parameters["dimension"] == 2:
print("No wavefunction of two space dimensions, silent return!")
return
if PP is None:
PP = parameters
if load is True:
# TODO: Implement reshaping
raise NotImplementedError("Loading of 2D grids is not implemented")
#G = iom.load_grid(blockid=blockid)
#G = grid.reshape((1, -1))
else:
G = BlockFactory().create_grid(PP)
V = BlockFactory().create_potential(parameters)
WF = WaveFunction(parameters)
WF.set_grid(G)
BT = BasisTransformationWF(V)
BT.set_grid(G)
timegrid = iom.load_wavefunction_timegrid(blockid=blockid)
u, v = G.get_nodes(split=True, flat=False)
u = real(u)
v = real(v)
N = WF.get_number_components()
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"]) ]
WF.set_values(values)
# Transform the values to the eigenbasis
# TODO: improve this:
if parameters["algorithm"] == "fourier":
BT.transform_to_eigen(WF)
else:
pass
Psi = WF.get_values()
for level in xrange(N):
z = Psi[level]
z = z.reshape(G.get_number_nodes())
# Plot the probability densities projected to the eigenbasis
fig = mlab.figure(size=(800,700))
surfcf(u, v, angle(z), abs(z))
#mlab.contour_surf(u, v, abs(z))
#mlab.outline()
#mlab.axes()
mlab.savefig("wavefunction_level_"+str(level)+"_timestep_"+(5-len(str(step)))*"0"+str(step)+".png")
mlab.close(fig)
print(" Plotting frames finished")
def plot_frames(PP,
iom,
blockid=0,
load=False,
eigentransform=False,
timerange=None,
sparsify=10,
view=None,
interactive=False,
path='.'):
"""Plot the wave function for a series of timesteps.
:param iom: An :py:class:`IOManager` instance providing the simulation data.
"""
parameters = iom.load_parameters()
if not parameters["dimension"] == 2:
print("No wavefunction of two space dimensions, silent return!")
return
if PP is None:
PP = parameters
if load is True:
# TODO: Implement reshaping
raise NotImplementedError("Loading of 2D grids is not implemented")
else:
G = BlockFactory().create_grid(PP)
if eigentransform:
V = BlockFactory().create_potential(parameters)
BT = BasisTransformationWF(V)
BT.set_grid(G)
WF = WaveFunction(parameters)
WF.set_grid(G)
N = WF.get_number_components()
timegrid = iom.load_wavefunction_timegrid(blockid=blockid)
if timerange is not None:
if len(timerange) == 1:
I = (timegrid == timerange)
else:
I = ((timegrid >= timerange[0]) & (timegrid <= timerange[1]))
if any(I):
timegrid = timegrid[I]
else:
raise ValueError("No valid timestep remains!")
u, v = G.get_nodes(split=True, flat=False)
u = real(u[::sparsify, ::sparsify])
v = real(v[::sparsify, ::sparsify])
# View
if view is not None:
if view[0] is None:
view[0] = u.min()
if view[1] is None:
view[1] = u.max()
if view[2] is None:
view[2] = v.min()
if view[3] is None:
view[3] = v.max()
for step in timegrid:
print(" Plotting frame of timestep # {}".format(step))
# Load the data
wave = iom.load_wavefunction(blockid=blockid, timestep=step)
values = [wave[j, ...] for j in range(parameters["ncomponents"])]
WF.set_values(values)
# Transform the values to the eigenbasis
if eigentransform:
BT.transform_to_eigen(WF)
Psi = WF.get_values()
for level in range(N):
# Wavefunction data
z = Psi[level]
z = z.reshape(G.get_number_nodes())[::sparsify, ::sparsify]
# View
if view is not None:
if view[4] is None:
view[4] = 0.0
if view[5] is None:
view[5] = 1.1 * abs(z).max()
# Plot
# if not interactive:
# mlab.options.offscreen = True
fig = mlab.figure(size=(800, 700))
surfcf(u, v, angle(z), abs(z), view=view)
mlab.draw()
if interactive:
mlab.show()
else:
mlab.savefig(
os.path.join(
"wavefunction_surface_block_%s_level_%d_timestep_%07d.png"
% (blockid, level, step)))
mlab.close(fig)