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:
        for k in BS:
            cj = vects[j]
            ck = vects[k]
            entry = 0.5 * squeeze(sum(conj(cj) * ck))
            MT[BS[j], BS[k]] = entry

    # Find eigenvalues and eigenvectors of the whole matrix
    M = MT + MV
    ew, ev = eigh(M)
    ind = argsort(ew)

    # Build the requested energy levels and states
    if parameters.has_key("eigenstates_indices"):
        states = parameters["eigenstates_indices"]
    else:
        # Groundstate only
        states = [0]

    BS = HAWP.get_basis_shapes(component=0)

    KEY = ("q","p","Q","P","S","adQ")

    print(70*"-")
    for state in states:
        if state > BS.get_basis_size():
            print("Warning: can not compute energy level "+state+" with basis size of "+str(BS))
            continue

        index = ind[state]

        coeffs = ev[:,index]
        energy = ew[index]

        print("Level: "+str(state))
        print("Energy: "+str(energy))
        print("Coefficients: \n")
        print(str(coeffs))
        print(70*"-")

        HAWP.set_coefficient_vector(coeffs.reshape((-1, 1)))

        # Save all the wavepacket data
        bid = IOM.create_block(groupid=gid)
        IOM.add_wavepacket(parameters, blockid=bid, key=KEY)
        IOM.save_wavepacket_description(HAWP.get_description(), blockid=bid)
        for shape in HAWP.get_basis_shapes():
            IOM.save_wavepacket_basisshapes(shape, blockid=bid)
        IOM.save_wavepacket_parameters(HAWP.get_parameters(key=KEY), timestep=0, blockid=bid, key=KEY)
        IOM.save_wavepacket_coefficients(HAWP.get_coefficients(), HAWP.get_basis_shapes(), timestep=0, blockid=bid)

    IOM.finalize()
def compute_eigenstate(parameters,
                       filename="eigenstates.hdf5",
                       computepq=True,
                       computePQ=True):
    r"""
    Special variables necessary in configuration:

    * eigenstate_of_level (default: 0)
    * eigenstates_indices (default: [0])
    * starting_point (default: (2, ..., 2))
    * hawp_template
    * innerproduct
    """
    D = parameters["dimension"]

    if "eigenstate_of_level" in parameters:
        N = parameters["eigenstate_of_level"]
    else:
        # Upper-most potential surface
        N = 0

    # Create output file now, in case this fails we did not waste computation time
    IOM = IOManager()
    IOM.create_file(filename)

    # 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()

    # Compute position and momentum
    if computepq:
        # 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.
        if "starting_point" in parameters:
            x0 = atleast_1d(parameters["starting_point"])
        else:
            x0 = 0.5 * 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)
    else:
        if "q0" in parameters:
            q0 = atleast_2d(parameters["q0"])
        else:
            q0 = zeros((D, 1))
        if "p0" in parameters:
            p0 = atleast_2d(parameters["p0"])
        else:
            p0 = zeros((D, 1))

    # Compute spreads
    if computePQ:
        # Q_0 = H^(-1/4)
        H = V.evaluate_hessian_at(q0)
        Q0 = inv(sqrtm(sqrtm(H)))
        # P_0 = i Q_0^(-1)
        P0 = 1.0j * inv(Q0)
    else:
        if "Q0" in parameters:
            Q0 = atleast_2d(parameters["Q0"])
        else:
            Q0 = identity(D)
        if "P0" in parameters:
            P0 = atleast_2d(parameters["P0"])
        else:
            P0 = 1.0j * inv(Q0)

    # The parameter set Pi
    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 check:")
    print("   P^T Q - Q^T P  =?=  0")
    print(dot(P0.T, Q0) - dot(Q0.T, P0))
    print("   Q^H P - P^H Q  =?=  2i")
    print(
        dot(transpose(conjugate(Q0)), P0) - dot(transpose(conjugate(P0)), Q0))

    # 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(component=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, component=N, as_packet=False)
        vects[i] = cnew

    for j in BS:
        for k in BS:
            cj = vects[j]
            ck = vects[k]
            entry = 0.5 * squeeze(sum(conjugate(cj) * ck))
            MT[BS[j], BS[k]] = entry

    # Find eigenvalues and eigenvectors of the whole matrix
    M = MT + MV
    ew, ev = eigh(M)
    ind = argsort(ew)

    # Build the requested energy levels and states
    if "eigenstates_indices" in parameters:
        states = parameters["eigenstates_indices"]
    else:
        # Groundstate only
        states = [0]

    BS = HAWP.get_basis_shapes(component=0)

    KEY = ("q", "p", "Q", "P", "S", "adQ")

    print(70 * "-")
    for state in states:
        if state > BS.get_basis_size():
            print(
                "Warning: can not compute energy level {} with basis size of {}"
                .format((state, BS)))
            continue

        index = ind[state]

        coeffs = ev[:, index]
        energy = ew[index]

        # Try to resolve ambiguities in sign
        imax = argmax(abs(coeffs))
        a = abs(angle(coeffs[imax]))
        if a > pi / 2.0:
            coeffs *= -1

        print("State: {}".format(state))
        print("Energy: {}".format(energy))
        print("Coefficients: \n")
        print(str(coeffs))
        print(70 * "-")

        HAWP.set_coefficient_vector(coeffs.reshape((-1, 1)))

        # Save all the wavepacket data
        bid = IOM.create_block(groupid=gid)
        IOM.add_wavepacket(parameters, blockid=bid, key=KEY)
        IOM.save_wavepacket(HAWP, 0, blockid=bid, key=KEY)

    IOM.finalize()

    # TODO: Find better criterion
    if norm(q0) > 1000:
        print("+----------------------------------+")
        print("| Run-away minimum?                |")
        print("| Maybe try different:             |")
        print("|   starting_point = [x0, y0, ...] |")
        print("+----------------------------------+")
def compute_energy_inhawp(iom, blockid=0, eigentrafo=True, iseigen=True):
    """Compute the energies 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``.
    :param iseigen: Whether the data is assumed to be in the eigenbasis.
    :type iseigen: 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]

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

    # Basis transformator
    if eigentrafo is True:
        BT = BasisTransformationHAWP(Potential)

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

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

    # Inner product
    if HAWP.get_innerproduct() is None:
        IP = BF.create_inner_product(parameters["innerproduct"])
        HAWP.set_innerproduct(IP)

    if eigentrafo is True:
        BT.set_matrix_builder(HAWP.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)

    O = ObservablesHAWP()
    KEY = ("q", "p", "Q", "P", "S", "adQ")

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

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

        # Configure the wavepacket
        HAWP.set_parameters(params, key=KEY)
        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)

        # Compute the energies
        O.set_innerproduct(HAWP.get_innerproduct())
        O.set_gradient(HAWP.get_gradient_operator())
        ekin = O.kinetic_energy(HAWP)
        if iseigen is True:
            epot = O.potential_energy(HAWP, Potential.evaluate_eigenvalues_at)
        else:
            epot = O.potential_energy(HAWP, Potential.evaluate_at)

        iom.save_energy((ekin, epot), 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)
Exemple #5
0
def compute_autocorrelation_hawp(iom, obsconfig, blockid=0, eigentrafo=True):
    """Compute the autocorrelation of a wavepacket 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.
    :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_wavepacket_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 norms, 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_wavepacket_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_wavepacket_basisshapes(blockid=blockid)
    BS = {}
    for ahash, descr in BS_descr.items():
        BS[ahash] = BF.create_basis_shape(descr)

    # Comfigure the original wavepacket
    KEY = ("q", "p", "Q", "P", "S", "adQ")
    # Retrieve simulation data
    params = iom.load_wavepacket_parameters(timestep=0,
                                            blockid=blockid,
                                            key=KEY)
    hashes, coeffs = iom.load_wavepacket_coefficients(timestep=0,
                                                      get_hashes=True,
                                                      blockid=blockid)
    # Configure the wavepacket
    HAWPo.set_parameters(params, key=KEY)
    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"])

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

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

        # Retrieve simulation data
        paramst = iom.load_wavepacket_parameters(timestep=step,
                                                 blockid=blockid,
                                                 key=KEY)
        hashes, coeffs = iom.load_wavepacket_coefficients(timestep=step,
                                                          get_hashes=True,
                                                          blockid=blockid)

        # Configure the wavepacket
        HAWPt.set_parameters(paramst, key=KEY)
        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)
Exemple #6
0
def compute_eigenstate(parameters, filename="eigenstates.hdf5", computepq=True, computePQ=True):
    r"""
    Special variables necessary in configuration:

    * eigenstate_of_level (default: 0)
    * eigenstates_indices (default: [0])
    * starting_point (default: (2, ..., 2))
    * hawp_template
    * innerproduct
    """
    D = parameters["dimension"]

    if "eigenstate_of_level" in parameters:
        N = parameters["eigenstate_of_level"]
    else:
        # Upper-most potential surface
        N = 0

    # Create output file now, in case this fails we did not waste computation time
    IOM = IOManager()
    IOM.create_file(filename)

    # 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()

    # Compute position and momentum
    if computepq:
        # 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.
        if "starting_point" in parameters:
            x0 = atleast_1d(parameters["starting_point"])
        else:
            x0 = 0.5 * 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)
    else:
        if "q0" in parameters:
            q0 = atleast_2d(parameters["q0"])
        else:
            q0 = zeros((D, 1))
        if "p0" in parameters:
            p0 = atleast_2d(parameters["p0"])
        else:
            p0 = zeros((D, 1))

    # Compute spreads
    if computePQ:
        # Q_0 = H^(-1/4)
        H = V.evaluate_hessian_at(q0)
        Q0 = inv(sqrtm(sqrtm(H)))
        # P_0 = i Q_0^(-1)
        P0 = 1.0j * inv(Q0)
    else:
        if "Q0" in parameters:
            Q0 = atleast_2d(parameters["Q0"])
        else:
            Q0 = identity(D)
        if "P0" in parameters:
            P0 = atleast_2d(parameters["P0"])
        else:
            P0 = 1.0j * inv(Q0)

    # The parameter set Pi
    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 check:")
    print("   P^T Q - Q^T P  =?=  0")
    print(dot(P0.T, Q0) - dot(Q0.T, P0))
    print("   Q^H P - P^H Q  =?=  2i")
    print(dot(transpose(conjugate(Q0)), P0) - dot(transpose(conjugate(P0)), Q0))

    # 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(component=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, component=N, as_packet=False)
        vects[i] = cnew

    for j in BS:
        for k in BS:
            cj = vects[j]
            ck = vects[k]
            entry = 0.5 * squeeze(sum(conjugate(cj) * ck))
            MT[BS[j], BS[k]] = entry

    # Find eigenvalues and eigenvectors of the whole matrix
    M = MT + MV
    ew, ev = eigh(M)
    ind = argsort(ew)

    # Build the requested energy levels and states
    if "eigenstates_indices" in parameters:
        states = parameters["eigenstates_indices"]
    else:
        # Groundstate only
        states = [0]

    BS = HAWP.get_basis_shapes(component=0)

    KEY = ("q", "p", "Q", "P", "S", "adQ")

    print(70 * "-")
    for state in states:
        if state > BS.get_basis_size():
            print("Warning: can not compute energy level {} with basis size of {}".format((state, BS)))
            continue

        index = ind[state]

        coeffs = ev[:, index]
        energy = ew[index]

        # Try to resolve ambiguities in sign
        imax = argmax(abs(coeffs))
        a = abs(angle(coeffs[imax]))
        if a > pi / 2.0:
            coeffs *= -1

        print("State: {}".format(state))
        print("Energy: {}".format(energy))
        print("Coefficients: \n")
        print(str(coeffs))
        print(70 * "-")

        HAWP.set_coefficient_vector(coeffs.reshape((-1, 1)))

        # Save all the wavepacket data
        bid = IOM.create_block(groupid=gid)
        IOM.add_wavepacket(parameters, blockid=bid, key=KEY)
        IOM.save_wavepacket(HAWP, 0, blockid=bid, key=KEY)

    IOM.finalize()

    # TODO: Find better criterion
    if norm(q0) > 1000:
        print("+----------------------------------+")
        print("| Run-away minimum?                |")
        print("| Maybe try different:             |")
        print("|   starting_point = [x0, y0, ...] |")
        print("+----------------------------------+")
def compute_energy_inhawp(iom, blockid=0, eigentrafo=True, iseigen=True):
    """Compute the energies 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``.
    :param iseigen: Whether the data is assumed to be in the eigenbasis.
    :type iseigen: 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]

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

    # Basis transformator
    if eigentrafo is True:
        BT = BasisTransformationHAWP(Potential)

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

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

    # Inner product
    if HAWP.get_innerproduct() is None:
        IP = BF.create_inner_product(parameters["innerproduct"])
        HAWP.set_innerproduct(IP)

    if eigentrafo is True:
        BT.set_matrix_builder(HAWP.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)

    O = ObservablesHAWP()
    KEY = ("q", "p", "Q", "P", "S", "adQ")

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

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

        # Configure the wavepacket
        HAWP.set_parameters(params, key=KEY)
        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)

        # Compute the energies
        O.set_innerproduct(HAWP.get_innerproduct())
        ekin = O.kinetic_energy(HAWP)
        if iseigen is True:
            epot = O.potential_energy(HAWP, Potential.evaluate_eigenvalues_at)
        else:
            epot = O.potential_energy(HAWP, Potential.evaluate_at)

        iom.save_energy((ekin, epot), timestep=step, blockid=blockid)