读取npz格式的模型并调用pyvista进行三维显示

读取npz格式的模型并调用pyvista进行三维显示

In [1]: import numpy as np

In [6]: dat = np.load('case2.npz')

　In [4]: dat.files
　Out[4]: ['param', 'topoz', 'code']

In [7]: dat['topoz'].shape
Out[7]: (201, 129, 129)

In [8]: dat['code'].shape
Out[8]: (202, 129, 129)

In [9]: dat['param'].shape
Out[9]: (14,)

 

#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=23000
        ylength=23000
        self.x = np.zeros((self._nk, self._nj, self._ni))
        xxx=np.arange(0,xlength,xlength/self._ni)
        xxx=xxx.reshape(1,1,self._ni)
        xxx=np.repeat(xxx,self._nj, axis=1)
        xxx=np.repeat(xxx,self._nk, axis=0)
        self.x=xxx   

        self.y = np.zeros_like(self.x)
        yyy=np.arange(0,ylength,ylength/self._nj)
        yyy=yyy.reshape(1,self._nj,1)
        yyy=np.repeat(yyy,self._ni, axis=2)
        yyy=np.repeat(yyy,self._nk, 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.view_vector([0.1,-0.5,0.1])  #set the view point
        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,489):
        fileName = 'Z:\jupiter\jupiter_deeplearn\jupiter_carbsims_case0\case%d.npz' % i
        myDataSet.read_data_npz(fileName)
        #myDataSet.show_model()
        myDataSet.show_model_as_png()
        print(i)
        gc.collect()
    #param=myDataSet.param
    #z = myDataSet.z
    #facies=myDataSet.facies
    #np.savez("case1.npz", param=param, z=z, facies=facies)