def __init__(self, config):
self.TempPsi = None
self.Config = config
self.Logger = GetClassLogger(self)
try:
#Enable redirect
if hasattr(config.Propagation, "silent"):
self.Silent = config.Propagation.silent
else:
self.Silent = False
redirectStateOld = Redirect.redirect_stdout
Redirect.Enable(self.Silent)
#Create wavefunction
self.psi = CreateWavefunction(config)
self.Logger.debug("Creating Propagator...")
self.Propagator = CreatePropagator(config, self.psi)
#apply propagation config
config.Propagation.Apply(self)
#Disable redirect
if not redirectStateOld == Redirect.redirect_stdout:
Redirect.Disable()
except:
#Diasable redirect
Redirect.Disable()
raise
class Problem:
"""
This is the main class of pyprop. It reads a config object and sets up everything
to allow propagation. See the examples folder for examples how to use this class.
"""
def __init__(self, config):
self.TempPsi = None
self.Config = config
self.Logger = GetClassLogger(self)
try:
#Enable redirect
if hasattr(config.Propagation, "silent"):
self.Silent = config.Propagation.silent
else:
self.Silent = False
redirectStateOld = Redirect.redirect_stdout
Redirect.Enable(self.Silent)
#Create wavefunction
self.psi = CreateWavefunction(config)
self.Logger.debug("Creating Propagator...")
self.Propagator = CreatePropagator(config, self.psi)
#apply propagation config
config.Propagation.Apply(self)
#Disable redirect
if not redirectStateOld == Redirect.redirect_stdout:
Redirect.Disable()
except:
#Diasable redirect
Redirect.Disable()
raise
def GetGrid(self):
"""
Helper function to construct the grid in each rank of the wavefunction
This method is most useful for cartesian coordinates, and possibly other
non-compressed grids. For compressed grids (like spherical) the grid returned
here will typically just be the idices converted to double
"""
#set up grid
repr = self.psi.GetRepresentation()
grid = [repr.GetLocalGrid(i) for i in range(0, self.psi.GetRank())]
return grid
#Propagation-------------------------------------------------
def SetupStep(self, skipWavefunctionSetup=False):
"""
Runs the nescescary setup routines to allow propagation.
This function must be called before the first call to AdvanceStep()
or Advance().
"""
try:
#Enable redirect
redirectStateOld = Redirect.redirect_stdout
Redirect.Enable(self.Silent)
self.Logger.debug("Starting setup timestep...")
self.Logger.debug(" Setting up Propagator.")
if self.Propagator != None:
self.Propagator.SetupStep(self.TimeStep)
self.Logger.debug(" Setting up initial wavefunction")
if not skipWavefunctionSetup:
self.SetupWavefunction()
self.Logger.info("Setup timestep complete.")
#Disable redirect
if not redirectStateOld:
Redirect.Disable()
except:
#Diasable redirect
Redirect.Disable()
raise
def RestartPropagation(self, timestep, startTime, propagationTime):
"""
Resets propagation time, duration and timestep for all propagators,
allowing the problem object to be propagated again without having
to call SetupStep.
"""
self.TimeStep = timestep
self.PropagatedTime = startTime
self.StartTime = startTime
self.Duration = propagationTime
if self.Propagator != None:
self.Propagator.RestartPropagation(timestep, startTime, propagationTime)
def AdvanceStep(self):
"""
Advances the wavefunction one timestep.
most of the work is done in the propagator object. This function
is merley to keep a track of propagated time, and provide a simple
interface to the user.
"""
self.Propagator.AdvanceStep(self.PropagatedTime, self.TimeStep )
if self.Propagator.RenormalizeActive:
self.psi.Normalize()
if abs(self.TimeStep.real) < 1e-10:
self.PropagatedTime += abs(self.TimeStep)
else:
self.PropagatedTime += self.TimeStep.real
def MultiplyHamiltonian(self, srcPsi, dstPsi):
"""
Applies the Hamiltonian to the wavefunction H psi -> psi
"""
self.Propagator.MultiplyHamiltonian(srcPsi, dstPsi, self.PropagatedTime, self.TimeStep )
def Advance(self, yieldCount, duration=None, yieldEnd=False):
"""
Returns a generator for advancing the wavefunction a number of timesteps.
If duration is specified the wavefunction is propagated until propagated
time >= duration.
if yieldCount is a number, it is the number of times this routine will
yield control back to the caller.
If is a boolean, yieldCount=True will make this function yield every timestep
and yieldCount=False will make it yield one time. (at the last timestep (i think))
if yieldEnd is True, a yield will be given at the end of propagation, regardless if
it matches the wanted number of yields
"""
if duration == None:
duration = self.Duration
#Determine how often we should yield.
if yieldCount.__class__ is bool:
yieldStep = 1
else:
yieldStep = int((duration / abs(self.TimeStep)) / yieldCount)
#Modify the interrupt signal handler
if RedirectInterrupt:
try:
InterruptHandler.UnRegister()
except: pass
InterruptHandler.Register()
index = 0
if self.TimeStep.real == 0:
#negative imaginary time
stoppingCriterion = lambda: (self.StartTime + self.Duration - self.PropagatedTime) > 0.5 * abs(self.TimeStep)
else:
endTime = self.StartTime + sign(self.TimeStep) * self.Duration
#real time
stoppingCriterion = lambda: abs(self.PropagatedTime - endTime) > 0.5 * abs(self.TimeStep)
prevYield = numpy.NAN
while stoppingCriterion():
#next timestep
self.AdvanceStep()
#check keyboard interrupt
if RedirectInterrupt:
if InterruptHandler.IsInterrupted():
InterruptHandler.ProcessInterrupt()
index += 1
if index % yieldStep == 0:
yield self.PropagatedTime
prevYield = self.PropagatedTime
if yieldEnd and prevYield != self.PropagatedTime:
yield self.PropagatedTime
if RedirectInterrupt:
InterruptHandler.UnRegister()
def GetEnergy(self):
if isreal(self.TimeStep):
return self.GetEnergyExpectationValue()
else:
return self.GetEnergyImTime()
def GetTempPsi(self):
if self.TempPsi == None:
self.TempPsi = self.psi.CopyDeep()
return self.TempPsi
def GetEnergyExpectationValue(self):
"""
Calculates the total energy of the problem by finding the expectation value
of the Hamiltonian
"""
self.psi.Normalize()
#Make copy of wavefunction
tempPsi = self.GetTempPsi()
tempPsi.GetData()[:] = 0
#Apply the Hamiltonian to the wavefunction
self.MultiplyHamiltonian(self.psi, tempPsi)
#Calculate the inner product between the applied psi and the original psi
energy = self.psi.InnerProduct(tempPsi)
#Check that the energy is real
if not hasattr(self, "IgnoreWarningRealEnergy"):
if abs(imag(energy)) > 1e-10:
self.Logger.warning("Energy is not real (%s). Possible bug. Supressing further warnings of this type" % (energy))
self.IgnoreWarningRealEnergy = True
return energy.real
def GetEnergyImTime(self):
"""
Advance one timestep without normalization and imaginary time.
we can then find the energy by looking at the norm of the decaying
wavefunction.
This will only work when negative imaginary time is used to propagate
to the groundstate, and will only give a good estimate for the ground
state energy when the wavefunction is well converged to the ground state.
Because it uses the norm of the wavefunction to measure energy, it will
be very sensitive to differences in the shape of the wavefunction.
"""
if isreal(self.TimeStep):
raise "Can only find energy for imaginary time propagation"
renorm = self.Propagator.RenormalizeActive
self.psi.Normalize()
self.Propagator.RenormalizeActive = False
self.AdvanceStep()
self.Propagator.RenormalizeActive = renorm
norm = self.psi.GetNorm()
self.psi.GetData()[:] /= norm
energy = - log(norm**2) / (2 * abs(self.TimeStep))
return energy
#Initialization-----------------------------------------------
def ApplyConfigSection(self, configSection):
self.TimeStep = complex(configSection.timestep)
self.Duration = configSection.duration
self.PropagatedTime = 0
self.StartTime = 0
if hasattr(configSection, "start_time"):
self.StartTime = configSection.start_time
self.PropagatedTime = self.StartTime
def SetupWavefunction(self):
"""
Initializes the wavefunction according to the initially specified configuration
file.
The supported initial condition types are:
InitialConditionType.Function
see SetupWavefunctionFunction()
InitialConditionType.File
See SetupWavefunctionFile()
"""
type = self.Config.InitialCondition.type
if type == InitialConditionType.Function:
self.SetupWavefunctionFunction(self.Config)
elif type == InitialConditionType.File:
self.SetupWavefunctionFile(self.Config)
elif type == InitialConditionType.Class:
self.SetupWavefunctionClass(self.Config, self.psi)
elif type == InitialConditionType.Custom:
self.SetupWavefunctionCustom(self.Config)
elif type == None:
pass
else:
raise "Invalid InitialConditionType: " + self.Config.InitialCondition.type
def SetupWavefunctionClass(self, config, psi):
classname = config.InitialCondition.classname
#Create globals
glob = dict(ProjectNamespace)
glob.update(globals())
#try to
evaluator = None
try: evaluator = eval(classname + "()", glob)
except: pass
if evaluator == None:
try: evaluator = eval(classname + "_" + str(psi.GetRank()) + "()", glob)
except: pass
if evaluator == None:
raise "Invalid classname", classname
config.Apply(evaluator)
config.InitialCondition.Apply(evaluator)
evaluator.SetupWavefunction(psi)
def SetupWavefunctionFunction(self, config):
"""
Initializes the wavefunction from a function specified in the InitialCondition
section of config. The function refrence config.InitialCondition.function
is evaulated in all grid points of the wavefunction, and the wavefunction is
set up accordingly.
for example. if this is a rank 2 problem, then this will initialize the wavefunction
to a 2D gaussian wavepacket.
def func(x, conf):
return exp(- x[0]**2 + x[1]**2])
config.InitialCondition.function = func
prop.SetupWavefunctionFunction(config)
REMARK: This function should probably not be called on directly. Use SetupWavefunction()
instead. That function will automatically determine the type of initial condition to be used.
"""
func = config.InitialCondition.function
conf = config.InitialCondition
#TODO: We should get this class from Propagator in order to use a different
#evaluator for a compressed (i.e. spherical) grid
evalfunc = eval("core.SetWavefunctionFromGridFunction_" + str(self.psi.GetRank()))
evalfunc(self.psi, func, conf)
def SetupWavefunctionFile(self, config):
"""
Initializes the wavefunction from a file specified in the InitialCondition.filename
according to the format InitialCondition.format. If InitialCondition.format is not
specified, this routine should automatically try to figure out which format it it,
but this is not yet implemented.
See the functions LoadWavefunction*() SaveWavefunction*() for details on how to
load and save wavefunctions
REMARK: This function should probably not be called on directly. Use SetupWavefunction()
instead. That function will automatically determine the type of initial condition to be used.
REMARK2: Currently it is not possible to interpolate a wavefunction between different grids.
this means that exactly the same grid used for saving the wavefunction must be used when loading
it
"""
format = config.InitialCondition.format
filename = config.InitialCondition.filename
if format == WavefunctionFileFormat.Ascii:
self.LoadWavefunctionAscii(filename)
elif format == WavefunctionFileFormat.Binary:
self.LoadWavefunctionPickle(filename)
elif format == WavefunctionFileFormat.HDF:
datasetPath = str(config.InitialCondition.dataset)
self.LoadWavefunctionHDF(filename, datasetPath)
else:
raise "Invalid file format: " + format
def SetupWavefunctionCustom(self, config):
"""
Initializes the wavefunction from a function specified in the InitialCondition
section of config. The function refrence config.InitialCondition.function
is evaulated called once, with the wavefunction as the first parameter, and
the configSection as the second.
REMARK: This function should probably not be called on directly. Use SetupWavefunction()
instead. That function will automatically determine the type of initial condition to be used.
"""
func = config.InitialCondition.function
conf = config.InitialCondition
func(self.psi, conf)
#(de)serialization---------------------------------------------
def LoadWavefunctionData(self, newdata):
data = self.psi.GetData()
if newdata.shape != data.shape:
raise "Invalid shape on loaded wavefunction, got " + str(newdata.shape) + " expected " + str(data.shape)
data[:] = newdata
def SaveWavefunctionPickle(self, filename):
serialization.SavePickleArray(filename, self.psi.GetData())
def LoadWavefunctionPickle(self, filename):
arr = serialization.LoadPickleArray(filename)
self.LoadWavefunctionData(arr)
def LoadWavefunctionHDF(self, filename, datasetPath):
serialization.LoadWavefunctionHDF(filename, datasetPath, self.psi)
def SaveWavefunctionHDF(self, filename, datasetPath):
serialization.SaveWavefunctionHDF(filename, datasetPath, self.psi, conf=self.Config)
def SaveWavefunctionAscii(self, filename):
psiData = self.psi.GetData()
assert(len(psiData.shape) <= 1)
pylab.save(filename, transpose((psiData.real ,psiData.imag)), delimiter=' ')
def SaveWavefunctionFortran(self, filename):
psiData = self.psi.GetData()
assert(len(psiData.shape) <= 1)
fh = open(filename, "w")
for i in range(psiData.size):
fh.write("(%s, %s) " % (psiData[i].real, psiData[i].imag) )
fh.close()
def LoadWavefunctionAscii(self, filename):
r, c = pylab.load(filename, unpack=True)
arr = r + 1.0j*c
self.LoadWavefunctionData(arr)