Python mayavi.mlab.savefig() Examples

def run_plots(DataDirectory,Base_file):

    root = DataDirectory+Base_file
    filenames = get_filenames(root)
    counter = 0

    # create the plot for the initial raster
    initial_file = filenames[0]

    # read in the raster
    raster = IO.ReadRasterArrayBlocks(initial_file)

    f = mlab.figure(size=(1000,1000), bgcolor=(0.5,0.5,0.5))
    s = mlab.surf(raster, warp_scale=0.4, colormap='gist_earth', vmax=100)
    #mlab.outline(color=(0,0,0))

    #mlab.axes(s, color=(1,1,1), z_axis_visibility=True, y_axis_visibility=False, xlabel='', ylabel='', zlabel='', ranges=[0,500,0,1000,0,0])

    #@mlab.animate(delay=10)
    #def anim():
    # now loop through each file and update the z values
    for fname in filenames:
        this_rast = IO.ReadRasterArrayBlocks(fname)
        s.mlab_source.scalars = this_rast
        #f.scene.render()
        #
        mlab.savefig(fname[:-4]+'_3d.png')
        #mlab.clf()

    # for (x, y, z) in zip(xs, ys, zs):
    #     print('Updating scene...')
    #     plt.mlab_source.set(x=x, y=y, z=z)
    #     yield 

def plot_mcontour(self, ndim0, ndim1, z, show_mode):
        "use mayavi.mlab to plot contour."
        if not mayavi_installed:
            self.__logger.info("Mayavi is not installed on your device.")
            return
        #do 2d interpolation
        #get slice object
        s = np.s_[0:ndim0:1, 0:ndim1:1]
        x, y = np.ogrid[s]
        mx, my = np.mgrid[s]
        #use cubic 2d interpolation
        interpfunc = interp2d(x, y, z, kind='cubic')
        newx = np.linspace(0, ndim0, 600)
        newy = np.linspace(0, ndim1, 600)
        newz = interpfunc(newx, newy)
        #mlab
        face = mlab.surf(newx, newy, newz, warp_scale=2)
        mlab.axes(xlabel='x', ylabel='y', zlabel='z')
        mlab.outline(face)
        #save or show
        if show_mode == 'show':
            mlab.show()
        elif show_mode == 'save':
            mlab.savefig('mlab_contour3d.png')
        else:
            raise ValueError('Unrecognized show mode parameter : ' +
                             show_mode)

        return 

def plot_field(self, **kwargs):
        "plot scalar field for elf data"
        if not mayavi_installed:
            self.__logger.warning("Mayavi is not installed on your device.")
            return
        # set parameters
        vmin = kwargs['vmin'] if 'vmin' in kwargs else 0.0
        vmax = kwargs['vmax'] if 'vmax' in kwargs else 1.0
        axis_cut = kwargs['axis_cut'] if 'axis_cut' in kwargs else 'z'
        nct = kwargs['nct'] if 'nct' in kwargs else 5
        widths = kwargs['widths'] if 'widths' in kwargs else (1, 1, 1)
        elf_data, grid = self.expand_data(self.elf_data, self.grid, widths)
        #create pipeline
        field = mlab.pipeline.scalar_field(elf_data)  # data source
        mlab.pipeline.volume(field, vmin=vmin, vmax=vmax)  # put data into volumn to visualize
        #cut plane
        if axis_cut in ['Z', 'z']:
            plane_orientation = 'z_axes'
        elif axis_cut in ['Y', 'y']:
            plane_orientation = 'y_axes'
        elif axis_cut in ['X', 'x']:
            plane_orientation = 'x_axes'
        cut = mlab.pipeline.scalar_cut_plane(
            field.children[0], plane_orientation=plane_orientation)
        cut.enable_contours = True  # 开启等值线显示
        cut.contour.number_of_contours = nct
        mlab.show()
        #mlab.savefig('field.png', size=(2000, 2000))

        return 

def mayavi_scraper(block, block_vars, gallery_conf):
    """Scrape Mayavi images.

    Parameters
    ----------
    block : tuple
        A tuple containing the (label, content, line_number) of the block.
    block_vars : dict
        Dict of block variables.
    gallery_conf : dict
        Contains the configuration of Sphinx-Gallery

    Returns
    -------
    rst : str
        The ReSTructuredText that will be rendered to HTML containing
        the images. This is often produced by :func:`figure_rst`.
    """
    from mayavi import mlab
    image_path_iterator = block_vars['image_path_iterator']
    image_paths = list()
    e = mlab.get_engine()
    for scene, image_path in zip(e.scenes, image_path_iterator):
        mlab.savefig(image_path, figure=scene)
        # make sure the image is not too large
        scale_image(image_path, image_path, 850, 999)
        if 'images' in gallery_conf['compress_images']:
            optipng(image_path, gallery_conf['compress_images_args'])
        image_paths.append(image_path)
    mlab.close(all=True)
    return figure_rst(image_paths, gallery_conf['src_dir']) 

def PlotNewtonRaphsonConvergence(self, increment=None, figure=None, show_plot=True, save=False, filename=None):
        """Plots convergence of Newton-Raphson for a given increment"""

        if self.fem_solver is None:
            raise ValueError("FEM solver not set for post-processing")

        if increment == None:
            increment = len(self.fem_solver.NRConvergence)-1

        import matplotlib.pyplot as plt
        if figure is None:
            figure = plt.figure()

        plt.plot(np.log10(self.fem_solver.NRConvergence['Increment_'+str(increment)]),'-ko')
        axis_font = {'size':'18'}
        plt.xlabel(r'$No\;\; of\;\; Iterations$', **axis_font)
        plt.ylabel(r'$log_{10}|Residual|$', **axis_font)
        plt.grid('on')

        if save:
            if filename is None:
                warn("No filename provided. I am going to write one in the current directory")
                filename = PWD(__file__) + '/output.eps'
            plt.savefig(filename, format='eps', dpi=500)

        if show_plot:
            plt.show() 

def _save_all_fired(self):
        file_dialog = DirectoryDialog(action = 'open', title = 'Select Directory')
        if file_dialog.open() == OK:
            out_path = file_dialog.path
            if self.display_all:
                items = [(None, k) for k in xrange(self._max_weight + 1)]
            else:
                items = product(self._names, xrange(self._max_weight + 1))
            for name, k in items:
                if name is not None:
                    self.selected_mesh = name
                self.weight_index = k
                mlab.savefig(path.join(out_path, "%s_%03d.png" % (name, k))) 

def plot_contour(self, ndim0, ndim1, z, show_mode):
        '''
        ndim0: int, point number on x-axis
        ndim1: int, point number on y-axis
        z    : 2darray, values on plane perpendicular to z axis
        '''
        #do 2d interpolation
        #get slice object
        s = np.s_[0:ndim0:1, 0:ndim1:1]
        x, y = np.ogrid[s]
        self.__logger.info('z shape = %s, x shape = %s, y shape = %s',
                           str(z.shape), str(x.shape), str(y.shape))
        mx, my = np.mgrid[s]
        #use cubic 2d interpolation
        interpfunc = interp2d(x, y, z, kind='cubic')
        newx = np.linspace(0, ndim0, 600)
        newy = np.linspace(0, ndim1, 600)
        #-----------for plot3d---------------------
        ms = np.s_[0:ndim0:600j, 0:ndim1:600j]  # |
        newmx, newmy = np.mgrid[ms]             # |
        #-----------for plot3d---------------------
        newz = interpfunc(newx, newy)

        #plot 2d contour map
        fig2d_1, fig2d_2, fig2d_3 = plt.figure(), plt.figure(), plt.figure()
        ax1 = fig2d_1.add_subplot(1, 1, 1)
        extent = [np.min(newx), np.max(newx), np.min(newy), np.max(newy)]
        img = ax1.imshow(newz, extent=extent, origin='lower')
        #coutour plot
        ax2 = fig2d_2.add_subplot(1, 1, 1)
        cs = ax2.contour(newx.reshape(-1), newy.reshape(-1), newz, 20, extent=extent)
        ax2.clabel(cs)
        plt.colorbar(mappable=img)
        # contourf plot
        ax3 = fig2d_3.add_subplot(1, 1, 1)
        ax3.contourf(newx.reshape(-1), newy.reshape(-1), newz, 20, extent=extent)

        #3d plot
        fig3d = plt.figure(figsize=(12, 8))
        ax3d = fig3d.add_subplot(111, projection='3d')
        ax3d.plot_surface(newmx, newmy, newz, cmap=plt.cm.RdBu_r)

        #save or show
        if show_mode == 'show':
            plt.show()
        elif show_mode == 'save':
            fig2d_1.savefig('surface2d.png', dpi=500)
            fig2d_2.savefig('contour2d.png', dpi=500)
            fig2d_3.savefig('contourf2d.png', dpi=500)
            fig3d.savefig('surface3d.png', dpi=500)
        else:
            raise ValueError('Unrecognized show mode parameter : ' +
                             show_mode)

        return 

def show_selected_grasps_with_color(m, ply_name_, title, obj_):
    m_good = m[m[:, 1] <= 0.4]
    m_good = m_good[np.random.choice(len(m_good), size=25, replace=True)]
    m_bad = m[m[:, 1] >= 1.8]
    m_bad = m_bad[np.random.choice(len(m_bad), size=25, replace=True)]
    collision_grasp_num = 0

    if save_fig or show_fig:
        # fig 1: good grasps
        mlab.figure(bgcolor=(1, 1, 1), fgcolor=(0.7, 0.7, 0.7), size=(1000, 1000))
        mlab.pipeline.surface(mlab.pipeline.open(ply_name_))
        for a in m_good:
            # display_gripper_on_object(obj, a[0])  # real gripper
            collision_free = display_grasps(a[0], obj_, color='d')  # simulated gripper
            if not collision_free:
                collision_grasp_num += 1

        if save_fig:
            mlab.savefig("good_"+title+".png")
            mlab.close()
        elif show_fig:
            mlab.title(title, size=0.5)

        # fig 2: bad grasps
        mlab.figure(bgcolor=(1, 1, 1), fgcolor=(0.7, 0.7, 0.7), size=(1000, 1000))
        mlab.pipeline.surface(mlab.pipeline.open(ply_name_))

        for a in m_bad:
            # display_gripper_on_object(obj, a[0])  # real gripper
            collision_free = display_grasps(a[0], obj_, color=(1, 0, 0))
            if not collision_free:
                collision_grasp_num += 1

        if save_fig:
            mlab.savefig("bad_"+title+".png")
            mlab.close()
        elif show_fig:
            mlab.title(title, size=0.5)
            mlab.show()
    elif generate_new_file:
        # only to calculate collision:
        collision_grasp_num = 0
        ind_good_grasp_ = []
        for i_ in range(len(m)):
            collision_free = display_grasps(m[i_][0], obj_, color=(1, 0, 0))
            if not collision_free:
                collision_grasp_num += 1
            else:
                ind_good_grasp_.append(i_)
        collision_grasp_num = str(collision_grasp_num)
        collision_grasp_num = (4-len(collision_grasp_num))*" " + collision_grasp_num
        print("collision_grasp_num =", collision_grasp_num, "| object name:", title)
        return ind_good_grasp_ 

def save_figures(image_path, fig_count, gallery_conf):
    """Save all open matplotlib figures of the example code-block

    Parameters
    ----------
    image_path : str
        Path where plots are saved (format string which accepts figure number)
    fig_count : int
        Previous figure number count. Figure number add from this number

    Returns
    -------
    list of strings containing the full path to each figure
    """
    figure_list = []

    fig_managers = matplotlib._pylab_helpers.Gcf.get_all_fig_managers()
    for fig_mngr in fig_managers:
        # Set the fig_num figure as the current figure as we can't
        # save a figure that's not the current figure.
        fig = plt.figure(fig_mngr.num)
        kwargs = {}
        to_rgba = matplotlib.colors.colorConverter.to_rgba
        for attr in ['facecolor', 'edgecolor']:
            fig_attr = getattr(fig, 'get_' + attr)()
            default_attr = matplotlib.rcParams['figure.' + attr]
            if to_rgba(fig_attr) != to_rgba(default_attr):
                kwargs[attr] = fig_attr

        current_fig = image_path.format(fig_count + fig_mngr.num)
        fig.savefig(current_fig, **kwargs)
        figure_list.append(current_fig)

    if gallery_conf.get('find_mayavi_figures', False):
        from mayavi import mlab
        e = mlab.get_engine()
        last_matplotlib_fig_num = len(figure_list)
        total_fig_num = last_matplotlib_fig_num + len(e.scenes)
        mayavi_fig_nums = range(last_matplotlib_fig_num, total_fig_num)

        for scene, mayavi_fig_num in zip(e.scenes, mayavi_fig_nums):
            current_fig = image_path.format(mayavi_fig_num)
            mlab.savefig(current_fig, figure=scene)
            # make sure the image is not too large
            scale_image(current_fig, current_fig, 850, 999)
            figure_list.append(current_fig)
        mlab.close(all=True)

    return figure_list 

def plot(self, figname=None, size=(300, 325), view=(30, 30), color=(1.0, 1.0, 1.0)):

        fig = mlab.figure(size=size)
        figure = mlab.gcf()
        fig.scene.disable_render = True
        figure.scene.background = (0.0, 0.0, 0.0)
        mlab.view(0, 90, distance=0.2)
        assert (self.structure.natom > 0)

        x = self.structure.positions[:, 0]
        y = self.structure.positions[:, 1]
        z = self.structure.positions[:, 2]
        cr = covalent_radius(self.structure.symbols)
        s = np.apply_along_axis(np.linalg.norm, 1, self.structure.positions)

        mlab.points3d(x, y, z, s, scale_factor=1.0, resolution=8, opacity=1.0,
                      color=color,
                      scale_mode='none')

        if self.structure.is_crystal:
            frame, line1, line2, line3 = self.structure.get_cell().get_path()

            mlab.plot3d(frame[:, 0], frame[:, 1], frame[:, 2], tube_radius=.02, color=(1, 1, 1))
            mlab.plot3d(line1[:, 0], line1[:, 1], line1[:, 2], tube_radius=.02, color=(1, 1, 1))
            mlab.plot3d(line2[:, 0], line2[:, 1], line2[:, 2], tube_radius=.02, color=(1, 1, 1))
            mlab.plot3d(line3[:, 0], line3[:, 1], line3[:, 2], tube_radius=.02, color=(1, 1, 1))
        else:
            for i in range(self.structure.natom - 1):
                for j in range(i + 1, self.structure.natom):
                    vector = self.structure.positions[i] - self.structure.positions[j]
                    mvector = np.linalg.norm(vector)
                    uvector = 1.0 / mvector * vector
                    if 2 * mvector < covalent_radius(self.structure.symbols[i]) + \
                            covalent_radius(self.structure.symbols[j]):
                        pair = np.concatenate(
                            (self.structure.positions[i] - 0.1 * uvector,
                             self.structure.positions[j] + 0.1 * uvector)).reshape((-1, 3))
                        mlab.plot3d(pair[:, 0], pair[:, 1], pair[:, 2], tube_radius=0.15, opacity=1.0, color=(1, 1, 1))

        mlab.view(distance=12.0)
        fig.scene.disable_render = False
        if figname is not None:
            mlab.savefig(figname)
        return figure