Python BlockFactory.create_grid方法代碼示例

# 需要導入模塊: from WaveBlocksND import BlockFactory [as 別名]
# 或者: from WaveBlocksND.BlockFactory import create_grid [as 別名]
def plot_frames(PP, iom, blockid=0, load=False, limits=None):
    r"""
    """
    parameters = iom.load_parameters()
    BF = BlockFactory()

    if not parameters["dimension"] == 2:
        print("No wavepacket 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 = BF.create_grid(PP)

    u, v = map(squeeze, G.get_axes())

    V = BF.create_potential(parameters)
    BT = BasisTransformationHAWP(V)

    wpd = iom.load_wavepacket_description(blockid=blockid)
    HAWP = BF.create_wavepacket(wpd)

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

    timegrid = iom.load_wavepacket_timegrid(blockid=blockid)

    N = HAWP.get_number_components()

    for step in timegrid:
        print(" Plotting frame of timestep # " + str(step))

        hi, ci = iom.load_wavepacket_coefficients(timestep=step, get_hashes=True, blockid=blockid)
        Pi = iom.load_wavepacket_parameters(timestep=step, blockid=blockid)

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

        psi = HAWP.evaluate_at(G, prefactor=True, component=0)

        fig = figure()

        for level in xrange(N):
            z = psi[level]
            z = z.reshape(G.get_number_nodes())

            subplot(N,1,level+1)
            #plotcm(z.reshape(G.get_number_nodes()), darken=0.3)
            plotcf2d(u, v, z, darken=0.3, limits=limits)

        savefig("wavepacket_block_"+str(blockid)+"_level_"+str(level)+"_timestep_"+(5-len(str(step)))*"0"+str(step)+".png")
        close(fig)

    print(" Plotting frames finished")

# 需要導入模塊: from WaveBlocksND import BlockFactory [as 別名]
# 或者: from WaveBlocksND.BlockFactory import create_grid [as 別名]
def plot_frames(PP, iom, blockid=0, load=False, eigentransform=False, timerange=None, view=None, path='.'):
    """Plot the wavepacket for a series of timesteps.

    :param iom: An :py:class:`IOManager` instance providing the simulation data.
    """
    parameters = iom.load_parameters()
    BF = BlockFactory()

    if not parameters["dimension"] == 2:
        print("No two-dimensional wavepacket, 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 = BF.create_grid(PP)

    if eigentransform:
        V = BF.create_potential(parameters)
        BT = BasisTransformationHAWP(V)

    timegrid = iom.load_wavepacket_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))

        HAWP = iom.load_wavepacket(step, blockid=blockid)
        N = HAWP.get_number_components()

        # Transform the values to the eigenbasis
        if eigentransform:
            BT.transform_to_eigen(HAWP)

        psi = HAWP.evaluate_at(G.get_nodes(), prefactor=True, component=0)

        # Plot
        fig = figure()

        for level in range(N):
            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, "wavepacket_block_%s_level_%d_timestep_%07d.png" % (blockid, level, step)))
        close(fig)