def compute_evaluate_wavepackets(iom, basis="eigen", blockid=0): """Evaluate an in homogeneous Hagdorn wavepacket on a given grid for each timestep. :param iom: An ``IOManager`` instance providing the simulation data. :param basis: The basis where the evaluation is done. Can be 'eigen' or 'canonical'. :param blockid: The data block from which the values are read. """ parameters = iom.load_parameters() # Number of time steps we saved timesteps = iom.load_inhomogwavepacket_timegrid(blockid=blockid) nrtimesteps = timesteps.shape[0] # Prepare the potential for basis transformations Potential = PotentialFactory().create_potential(parameters) # Retrieve simulation data if iom.has_grid(blockid=blockid): grid = iom.load_grid(blockid=blockid) else: grid = iom.load_grid(blockid="global") params = iom.load_inhomogwavepacket_parameters(blockid=blockid) coeffs = iom.load_inhomogwavepacket_coefficients(blockid=blockid) # A data transformation needed by API specification params = [ [ params[j][i,:] for j in xrange(parameters["ncomponents"]) ] for i in xrange(nrtimesteps) ] coeffs = [ [ coeffs[i,j,:] for j in xrange(parameters["ncomponents"]) ] for i in xrange(nrtimesteps) ] # We want to save wavefunctions, thus add a data slot to the data file iom.add_wavefunction(parameters, timeslots=nrtimesteps, blockid=blockid) # Hack for allowing data blocks with different basis size than the global one # todo: remove when we got local parameter sets parameters.update_parameters({"basis_size": coeffs[0][0].shape[0]}) HAWP = HagedornWavepacketInhomogeneous(parameters) HAWP.set_quadrature(None) WF = WaveFunction(parameters) WF.set_grid(grid) # Iterate over all timesteps for i, step in enumerate(timesteps): print(" Evaluating inhomogeneous wavepacket at timestep "+str(step)) # Configure the wavepacket HAWP.set_parameters(params[i]) HAWP.set_coefficients(coeffs[i]) # Project to the eigenbasis if desired if basis == "eigen": HAWP.project_to_eigen(Potential) # 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): """ :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)