Python WaveBlocksND.BlockFactory類代碼示例

def read_data_inhomogeneous(iom, blockid=0):
    r"""
    :param iom: An :py:class:`IOManager` instance providing the simulation data.
    :param blockid: The data block from which the values are read.
    """
    BF = BlockFactory()

    parameters = iom.load_parameters()
    timegrid = iom.load_inhomogwavepacket_timegrid(blockid=blockid)

    # Basis shapes
    bsdescr = iom.load_inhomogwavepacket_basisshapes(blockid=blockid)
    BS = {}
    for ahash, descr in bsdescr.iteritems():
        BS[ahash] = BF.create_basis_shape(descr)

    # Plot the coefficients for all timesteps
    for j, step in enumerate(timegrid):
        hashes, coeffs = iom.load_inhomogwavepacket_coefficients(timestep=step, blockid=blockid, get_hashes=True)

        k = []

        for i in xrange(parameters["ncomponents"]):
            bs = BS[int(hashes[i])]
            ki = array([bs[node] for node in bs.get_node_iterator()])
            k.append(ki)

        plot_coefficients(k, coeffs, step, parameters["dt"], index=blockid)

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 compute_norm_inhawp(iom, blockid=0, eigentrafo=True):
    """Compute the norm of a wavepacket timeseries.
    This function is for inhomogeneous wavepackets.

    :param iom: An :py:class:`IOManager` instance providing the simulation data.
    :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_inhomogwavepacket_timegrid(blockid=blockid)
    nrtimesteps = timesteps.shape[0]

    # Basis transformator
    if eigentrafo is True:
        # The potential used
        Potential = BlockFactory().create_potential(parameters)
        BT = BasisTransformationHAWP(Potential)

    # We want to save norms, thus add a data slot to the data file
    iom.add_norm(parameters, timeslots=nrtimesteps, blockid=blockid)

    # Initialize a Hagedorn wavepacket with the data
    descr = iom.load_inhomogwavepacket_description(blockid=blockid)
    HAWP = BlockFactory().create_wavepacket(descr)

    if eigentrafo is True:
        BT.set_matrix_builder(HAWP.get_quadrature())

    # Basis shapes
    BS_descr = iom.load_inhomogwavepacket_basisshapes()
    BS = {}
    for ahash, descr in BS_descr.iteritems():
        BS[ahash] = BlockFactory().create_basis_shape(descr)

    # Iterate over all timesteps
    for i, step in enumerate(timesteps):
        print(" Computing norms of timestep "+str(step))

        # Retrieve simulation data
        params = iom.load_inhomogwavepacket_parameters(timestep=step, blockid=blockid)
        hashes, coeffs = iom.load_inhomogwavepacket_coefficients(timestep=step, get_hashes=True, blockid=blockid)

        # Configure the wavepacket
        HAWP.set_parameters(params)
        HAWP.set_basis_shape([ BS[int(ha)] for ha in hashes ])
        HAWP.set_coefficients(coeffs)

        # Transform to the eigenbasis.
        if eigentrafo is True:
            BT.transform_to_eigen(HAWP)

        # Measure norms in the eigenbasis
        norm = HAWP.norm()

        # Save the norms
        iom.save_norm(norm, timestep=step, blockid=blockid)

def read_data_inhomogeneous(iom, blockid=0):
    r"""
    :param iom: An :py:class:`IOManager` instance providing the simulation data.
    :param blockid: The data block from which the values are read.
    """
    parameters = iom.load_parameters()
    timegrid = iom.load_inhomogwavepacket_timegrid(blockid=blockid)
    time = timegrid * parameters["dt"]

    # The potential used
    Potential = BlockFactory().create_potential(parameters)

    # Basis transformator
    BT = BasisTransformationHAWP(Potential)

    # Basis shapes
    BS_descr = iom.load_wavepacket_basisshapes(blockid=blockid)
    BS = {}
    for ahash, descr in BS_descr.iteritems():
        BS[ahash] = BlockFactory().create_basis_shape(descr)

    # Initialize a Hagedorn wavepacket with the data
    descr = iom.load_inhomogwavepacket_description(blockid=blockid)
    HAWP = BlockFactory().create_wavepacket(descr)

    BT.set_matrix_builder(HAWP.get_quadrature())

    # Store the resulting coefficients here
    CI = [ [] for i in xrange(HAWP.get_number_components()) ]

    # Iterate over all timesteps, this is an *expensive* transformation
    for i, step in enumerate(timegrid):
        print(" Computing eigentransformation at timestep "+str(step))
        # Retrieve simulation data
        params = iom.load_inhomogwavepacket_parameters(timestep=step, blockid=blockid)
        hashes, coeffs = iom.load_inhomogwavepacket_coefficients(timestep=step, get_hashes=True, blockid=blockid)

        # Configure the wavepacket
        HAWP.set_parameters(params)
        HAWP.set_basis_shapes([ BS[int(ha)] for ha in hashes ])
        HAWP.set_coefficients(coeffs)

        # Transform to the eigenbasis.
        BT.transform_to_eigen(HAWP)

        for index, item in enumerate(HAWP.get_coefficients()):
            CI[index].append(item)

    CI = [ transpose(hstack(item)) for item in CI ]

    return time, CI

def read_data_inhomogeneous(iom, blockid=0):
    r"""
    :param iom: An :py:class:`IOManager` instance providing the simulation data.
    :param blockid: The data block from which the values are read.
    """
    parameters = iom.load_parameters()
    timegrid = iom.load_inhomogwavepacket_timegrid(blockid=blockid)
    dt = parameters["dt"] if "dt" in parameters else 1.0
    time = timegrid * dt

    # The potential used
    Potential = BlockFactory().create_potential(parameters)

    # Basis transformator
    BT = BasisTransformationHAWP(Potential)

    # Initialize a Hagedorn wavepacket with the data
    descr = iom.load_inhomogwavepacket_description(blockid=blockid)
    HAWP = BlockFactory().create_wavepacket(descr)

    BT.set_matrix_builder(HAWP.get_quadrature())

    # Store the resulting coefficients here
    CI = [[] for i in range(HAWP.get_number_components())]

    # Iterate over all timesteps, this is an *expensive* transformation
    for i, step in enumerate(timegrid):
        print(" Computing eigentransformation at timestep {}".format(step))
        # Retrieve simulation data
        HAWP = iom.load_inhomogwavepacket(timestep=step, blockid=blockid)

        # Transform to the eigenbasis.
        BT.transform_to_eigen(HAWP)

        for index, item in enumerate(HAWP.get_coefficients()):
            CI[index].append(item)

    CI = [transpose(hstack(item)) for item in CI]

    return time, CI

def read_data_inhomogeneous(iom, blockid=0, timerange=None, path='.'):
    r"""
    :param iom: An :py:class:`IOManager` instance providing the simulation data.
    :param blockid: The data block from which the values are read.
    """
    BF = BlockFactory()

    parameters = iom.load_parameters()
    timegrid = iom.load_inhomogwavepacket_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!")

    # Basis shapes
    bsdescr = iom.load_inhomogwavepacket_basisshapes(blockid=blockid)
    BS = {}
    for ahash, descr in bsdescr.items():
        BS[ahash] = BF.create_basis_shape(descr)

    # Plot the coefficients for all timesteps
    for j, step in enumerate(timegrid):
        allhashes, allcoeffs = iom.load_inhomogwavepacket_coefficients(timestep=step, blockid=blockid, get_hashes=True)

        k = []
        ck = []
        for ahash, coeffs in zip(allhashes, allcoeffs):
            bs = BS[int(ahash)]
            ki = array([bs[node] for node in bs.get_node_iterator(mode="mag")])
            ck.append(coeffs[ki])
            ki.sort()
            k.append(ki)

        dt = parameters["dt"] if "dt" in parameters else None
        plot_coefficients(k, ck, step, dt, blockid=blockid, path=path)

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 compute_autocorrelation_inhawp(iom, obsconfig, blockid=0, eigentrafo=True):
    """Compute the autocorrelation of a wavepacket timeseries.
    This function is for inhomogeneous wavepackets.

    :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.
    :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()
    BF = BlockFactory()

    # Number of time steps we saved
    timesteps = iom.load_inhomogwavepacket_timegrid(blockid=blockid)
    nrtimesteps = timesteps.shape[0]

    # Basis transformator
    if eigentrafo is True:
        # The potential used
        Potential = BF.create_potential(parameters)
        BT = BasisTransformationHAWP(Potential)

    # We want to save autocorrelations, thus add a data slot to the data file
    iom.add_autocorrelation(parameters, timeslots=nrtimesteps, blockid=blockid)

    # Initialize a Hagedorn wavepacket with the data
    descr = iom.load_inhomogwavepacket_description(blockid=blockid)
    HAWPo = BF.create_wavepacket(descr)
    HAWPt = BF.create_wavepacket(descr)

    if eigentrafo is True:
        BT.set_matrix_builder(HAWPo.get_innerproduct())

    # Basis shapes
    BS_descr = iom.load_inhomogwavepacket_basisshapes(blockid=blockid)
    BS = {}
    for ahash, descr in BS_descr.items():
        BS[ahash] = BF.create_basis_shape(descr)

    # Comfigure the original wavepacket
    # Retrieve simulation data
    params = iom.load_inhomogwavepacket_parameters(timestep=0, blockid=blockid)
    hashes, coeffs = iom.load_inhomogwavepacket_coefficients(timestep=0, get_hashes=True, blockid=blockid)
    # Configure the wavepacket
    HAWPo.set_parameters(params)
    HAWPo.set_basis_shapes([BS[int(ha)] for ha in hashes])
    HAWPo.set_coefficients(coeffs)

    # Set up the innerproduct for solving the integrals <phi_0 | phi_t>
    IP = BF.create_inner_product(obsconfig["innerproduct"])

    # Iterate over all timesteps
    for i, step in enumerate(timesteps):
        print(" Computing autocorrelations of timestep %d" % step)

        # Retrieve simulation data
        params = iom.load_inhomogwavepacket_parameters(timestep=step, blockid=blockid)
        hashes, coeffs = iom.load_inhomogwavepacket_coefficients(timestep=step, get_hashes=True, blockid=blockid)

        # Configure the wavepacket
        HAWPt.set_parameters(params)
        HAWPt.set_basis_shapes([BS[int(ha)] for ha in hashes])
        HAWPt.set_coefficients(coeffs)

        # Transform to the eigenbasis.
        if eigentrafo is True:
            BT.transform_to_eigen(HAWPt)

        # Measure autocorrelations in the eigenbasis
        acs = IP.quadrature(HAWPo, HAWPt, diagonal=True)

        # Save the autocorrelations
        iom.save_autocorrelation(acs, timestep=step, blockid=blockid)

    potvars = map(lambda x: ":math:`"+str(x)+"`", potvars)
    potvars = reduce(lambda a,b: a+", "+b, potvars)

    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"] = potdef
    P = BlockFactory().create_potential(params)

    if len(potdef["variables"]) == 1:
        # Plot the potential
        values = P.evaluate_eigenvalues_at(x)

        figure(figsize=(4,3))
        for value in values:
            plot(squeeze(x), squeeze(value))
        grid(True)
        xlabel(r"$x$")
        ylabel(r"$\lambda_i\left(x\right)$")
        xlim(min(x), max(x))
        savefig(potdef["name"] + ".png")

    elif len(potdef["variables"]) == 2:

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 load_from_file(filepath, blockid=0, timestep=0, sizeK=None):
    r"""Utility script to load wavepacket parameters and coefficients
    from another simulation result in a form suitable for the input
    configuration of a new simulation. This is (mainly) used
    to start simulations with previously computed eigenstates.

    :param filepath: The path to the `.hdf5` file from which data will be read.
    :param blockid: The `datablock` from which to read the data.
                    Default is the block with `blockid=0`.
    :param timestep: Load the data corresponding to the given `timestep`.
                     The default timestep is `0`.
    :param sizeK: Load at most 'sizeK' many coefficients. Note that the order
                  is defined by the linearization mapping :math:`\mu` of the
                  packet's current basis shape. We then pick the first `sizeK`
                  ones.
    """
    IOM = IOManager()
    IOM.open_file(filepath)

    # Check if we have data
    tg = IOM.load_wavepacket_timegrid(blockid=blockid)
    if timestep not in tg:
        raise ValueError("No data for timestep {}".format(timestep))

    # Load data and assemble packet
    BF = BlockFactory()

    # Basis shapes
    BS_descr = IOM.load_wavepacket_basisshapes(blockid=blockid)
    BS = {}
    for ahash, descr in BS_descr.items():
        BS[ahash] = BF.create_basis_shape(descr)

    # Create a packet
    wpd = IOM.load_wavepacket_description(blockid=blockid)
    HAWP = BF.create_wavepacket(wpd)

    # Data
    ha, ci = IOM.load_wavepacket_coefficients(blockid=blockid, timestep=timestep, get_hashes=True)
    Pi = IOM.load_wavepacket_parameters(blockid=blockid, timestep=timestep)

    HAWP.set_parameters(Pi)
    HAWP.set_basis_shapes([BS[int(h)] for h in ha])
    HAWP.set_coefficients(ci)

    # Reformat data
    C = []

    for n in range(HAWP.get_number_components()):
        B = HAWP.get_basis_shapes(component=n)
        cn = HAWP.get_coefficients(component=n)
        l = []
        for i in range(B.get_basis_size()):
            l.append((B[i], cn[i, 0]))
        C.append(l)

    if sizeK is not None:
        # We load at most 'sizeK' coefficients.
        # Note that this does NOT specify which
        # ones in terms of multi-indices.
        C = [c[:sizeK] for c in C]

    return Pi, C

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)

#.........這裏部分代碼省略.........

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 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 compute_eigenstate(parameters):
    r"""
    Special variables necessary in configuration:

    * eigenstate_of_level (default: 0)
    * states_indices (default: [0])
    """
    D = parameters["dimension"]

    if parameters.has_key("eigenstate_of_level"):
        N = parameters["eigenstate_of_level"]
    else:
        # Upper-most potential surface
        N = 0

    # Create output file now, in case this fails we did not waste computations
    IOM = IOManager()
    IOM.create_file("eigenstates.hdf5")

    # Save the simulation parameters
    IOM.add_parameters()
    IOM.save_parameters(parameters)

    gid = IOM.create_group()

    BF = BlockFactory()
    # Create the potential
    V = BF.create_potential(parameters)
    V.calculate_local_quadratic()

    # Minimize the potential to find q0
    f = lambda x: real((squeeze(V.evaluate_at(x)[N])))
    # Start with an offset because exact 0.0 values can give
    # issues, especially with the Hessian evaluation. This way
    # the minimizer will always stay away from zero a tiny bit.
    # The current starting point can give issues if the potential
    # is stationary at the point (2, ..., 2) but that is less likely.
    x0 = 2.0*ones(D)
    q0 = fmin(f, x0, xtol=1e-12)
    q0 = q0.reshape((D,1))

    # We are at the minimum with no momentum
    p0 = zeros_like(q0)

    # Compute spreads now
    # Q_0 = H^(-1/4)
    H = V.evaluate_hessian_at(q0)
    Q0 = inv(sqrtm(sqrtm(H)))
    # Take P_00 = i Q_0^(-1)
    P0 = 1.0j * inv(Q0)

    #
    print(70*"-")
    print("Parameter values are:")
    print("---------------------")
    print(" q0:")
    print(str(q0))
    print(" p0:")
    print(str(p0))
    print(" Q0:")
    print(str(Q0))
    print(" P0:")
    print(str(P0))
    # Consistency check
    print(" consistency:")
    print(str(conj(Q0)*P0 - conj(P0)*Q0))
    print(70*"-")

    # Next find the new coefficients c'
    HAWP = BF.create_wavepacket(parameters["hawp_template"])

    # Set the parameter values
    Pi = HAWP.get_parameters()
    Pi[0] = q0
    Pi[1] = p0
    Pi[2] = Q0
    Pi[3] = P0
    HAWP.set_parameters(Pi)

    # Next compute the matrix M_ij = <phi_i | T + V | phi_j>
    # The potential part
    HQ = BF.create_inner_product(parameters["innerproduct"])

    opV = lambda x, q, entry: V.evaluate_at(x, entry=entry)
    MV = HQ.build_matrix(HAWP, operator=opV)

    # The kinetic part
    MT = zeros_like(MV, dtype=complexfloating)
    GR = GradientHAWP()
    BS = HAWP.get_basis_shapes(N)

    vects = {}
    for i in BS:
        z = zeros_like(HAWP.get_coefficient_vector(), dtype=complexfloating)
        HAWP.set_coefficient_vector(z)
        HAWP.set_coefficient(N, i, 1.0)
        Kn, cnew = GR.apply_gradient(HAWP, N)
        vects[i] = cnew

    for j in BS:
#.........這裏部分代碼省略.........