基于pyvista的批量绘图，解决内存泄漏的问题

 这个程序封装了一个类，用来读取三维数据，并通过pyvista进行显示和保存，可以读取批量的模型分别进行显示和图形保存。

#coding=utf-8

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import pyvista as pv
import sys
import gc

class MyDataSet:

    def __init__(self, ni, nj, nk, fileNameList):
        self._ni = ni
        self._nj = nj
        self._nk = nk
        self._fileNameList = fileNameList
        self.create_colorMap()
        self.create_x_y()
        self.p = pv.Plotter(off_screen=True)

		
    def create_x_y(self):
        xlength=19001
        ylength=8801
        self.x = np.zeros((self._nk, self._nj, self._ni))
        xxx=np.arange(0,xlength,148.4375)
        xxx=xxx.reshape(1,1,129)
        xxx=np.repeat(xxx,129, axis=1)
        xxx=np.repeat(xxx,201, axis=0)
        self.x=xxx   

        self.y = np.zeros_like(self.x)
        yyy=np.arange(0,ylength,68.75)
        yyy=yyy.reshape(1,129,1)
        yyy=np.repeat(yyy,129, axis=2)
        yyy=np.repeat(yyy,201, axis=0)
        self.y=yyy



    def read_data_txt(self, fileName):
        dat = np.loadtxt(fileName, skiprows=5)
        dat = dat.transpose()
        self.x = dat[3].reshape(self._nk, self._nj, self._ni)
        self.y = dat[4].reshape(self._nk, self._nj, self._ni)
        self.z = dat[5].reshape(self._nk, self._nj, self._ni)
        self.facies = dat[8].reshape(self._nk, self._nj, self._ni)
        self.facies[self.facies > 7] = 8
        print(self.z.min())
        print(self.z.max())
        with open(fileName, 'r') as f:
            line = f.readline()
            line = f.readline()
            line = f.readline()
            line = line.split()
            parm = []
            for i in range(len(line)):
                parm.append(float(line[i]))
            self.param = np.array(parm)

    def read_data_npz(self, fileName):
        self.fileName=fileName
        self.dat = np.load(fileName)
        print(self.dat.files)
        self.parms = self.dat['param']
        self.z = self.dat['topoz']
        self.facies = self.dat['code']
        self.facies[self.facies > 5] = 5
        print(self.z.min())
        print(self.z.max())

    def create_colorMap(self):
        #colors= ['blue','cyan','green','pink','magenta','purple','gold','red']
        #colors= ['blue','green','pink','magenta','orange','yellow','red', 'black']
        #colors = ['blue', 'green', 'aquamarine', 'mediumspringgreen', 'orange', 'lightpink','yellow', 'red', 'black']
        colors = ['blue', 'green', 'aquamarine', 'mediumspringgreen', 'orange']
        self.discmap = mpl.colors.ListedColormap(colors)

    def set_zscale(self, _zscale):
        self.p.set_scale(zscale=_zscale)

    def show_model_as_png(self):
        self.p.clear()
        self.p.set_scale(zscale=10)
        #self.p.add_slider_widget(self.set_zscale(),rng=np.range(1,100))
        self.p.add_axes()
        grid = pv.StructuredGrid(self.x, self.y, self.z)
        #grid.point_arrays["values"]=self.facies.flatten(order="F")
        grid.cell_data["values"] = self.facies[0:self._nk - 1, 0:self._nj - 1,
                                                 0:self._ni - 1].flatten(order="F")
        self.p.add_mesh(grid, cmap=self.discmap, show_scalar_bar=True)
        
        pngName=self.fileName[0:-4]+".png"
        self.p.show(interactive=True, auto_close=False, screenshot=pngName)
        self.p.screenshot(pngName)
        gc.collect()

    def show_model(self):
        self.p.set_scale(zscale=10)
        # self.p.add_slider_widget(self.set_zscale(),rng=np.range(1,100))
        self.p.add_axes()
        self.grid = pv.StructuredGrid(self.x, self.y, self.z)
        #grid.point_arrays["values"]=self.facies.flatten(order="F")
        self.grid.cell_data["values"] = self.facies[0:self._nk - 1, 0:self._nj - 1,
                                                 0:self._ni - 1].flatten(order="F")
        self.p.add_mesh(self.grid, cmap=self.discmap, show_scalar_bar=True)
        
        pngName=self.fileName[0:-4]+".png"
        self.p.show(interactive=True, auto_close=False, screenshot=pngName)
        #self.p.show(interactive=False, screenshot=pngName)
        self.p.screenshot(pngName, 'off_screen')
        self.p.close()
        gc.collect()


if __name__ == "__main__":
    fileNameList = []
    myDataSet = MyDataSet(129, 129, 201, fileNameList)
    #fileName = sys.argv[1]
    #myDataSet.read_data_npz(fileName)
    #myDataSet.show_model()

    for i in range(0,11000):
        fileName = '\case%d.npz' % i
        myDataSet.read_data_npz(fileName)
        #myDataSet.show_model()
        myDataSet.show_model_as_png()
        print(i)
        gc.collect()