From 13325e63fc3b986c228b03c62da9d52317e794ef Mon Sep 17 00:00:00 2001 From: Luz Angelica Caudillo Mata Date: Tue, 2 Jul 2013 12:59:02 -0700 Subject: [PATCH] Implements the first draft of the TensorMesh class with basic visualization functions. It works for 1, 2 and 3 dimensions. Provides nodal, cell centers and staggered grids. --- code/TensorMesh.py | 257 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 257 insertions(+) create mode 100644 code/TensorMesh.py diff --git a/code/TensorMesh.py b/code/TensorMesh.py new file mode 100644 index 00000000..3a061561 --- /dev/null +++ b/code/TensorMesh.py @@ -0,0 +1,257 @@ +#------------------------------------------------------------------------------- +# Packages +#------------------------------------------------------------------------------- +import numpy as np +import matplotlib.pyplot as plt +from mpl_toolkits.mplot3d import Axes3D + +#------------------------------------------------------------------------------- +# Class definition +#------------------------------------------------------------------------------- +class TensorMesh: + """ + Define nodal, cell-centered and staggered tensor meshes for 1, 2 and 3 + dimensions. + """ + + # ---------------------- Properties -------------------------------------- + h = None # Array Spacing or cell-sizes in each direction + x0 = None # Array Origin (x1,x2,x3) + dim = None # Int Dimension + n = None # Array Number of cells in each direction + nC = None # Int Total number of cells + nE = None # Int Total number of edges + nF = None # Int Total number of faces + + # ----------------------- Methods ---------------------------------------- + def __init__(self,h,x0): + """Compute number of edges,faces and cell-centers """ + + # Assign values to properties + self.h = h + self.x0 = x0 + + + #Compute derived properties + self.dim = np.size(x0) + + #Compute the num of cells in each direction + self.n = np.zeros((self.dim,1)) + + for d in range(self.dim): + self.n[d] = np.size(h[d]) + + # Compute the number of cell-centers + self.nC = np.prod(self.n) + + # Compute the number of edges (makes sense only for 3D) + # Equivalent to: + # nEdges = n[0] * (n[1]+1) * (n[2]+1) + # + (n[0]+1) * ny[1] * (nz[2]+1) + # + (n[0]+1) * (ny[1]+1)* (nz[2]) + + if self.dim == 3: + self.nE = np.prod(np.kron(np.ones((3,1)),self.n.T)+np.ones((3,3))-np.eye(3),1) + print self.nE + + # Compute the number of faces (makes sense only for 2 and 3D) + # Equivalent to + # nFaces = (n[0]+1) * n[1] * n[2] + # + n[0] * (ny[1]+1) * nz[2] + # + n[0] * ny[1] * (nz[2]+1) + if self.dim >=2: + self.nF = np.prod(np.kron(np.ones((self.dim,1)),self.n.T)+np.eye(self.dim),1) + print self.nF + + def xin(self,i): + """Construct the 1D nodal mesh from the ith-component of h. Return an array.""" + return np.insert(np.cumsum(self.h[i-1]),0,0.0) + self.x0[i-1] + + def xic(self,i): + """Construct the 1D cell-centerd mesh from the ith-component of h. Return an array.""" + return .5*( np.insert(np.cumsum(self.h[i-1][:,0:-1]),0,0.0) + np.cumsum(self.h[i-1])) + + def getNodalGrid(self): + """Construct nodal grid for 1, 2 and 3 dimensions""" + if self.dim==1: + return [self.xin(1)] + elif self.dim==2: + return self.ndgrid([self.xin(1),self.xin(2)]) + elif self.dim==3: + return self.ndgrid([self.xin(1),self.xin(2),self.xin(3)]) + + def ndgrid(self, xin): + """Form tensorial grid for 1, 2 and 3 dimensions. Return X1,X2,X3 arrays depending on the dimension""" + ei = lambda i : np.ones((np.size(xin[i-1]),1)) + + if self.dim==1: + return [xin] + elif self.dim==2: + X1 = np.kron(ei(2),xin[0]).reshape(-1,1) + X2 = np.kron(xin[1],ei(1).T).reshape(-1,1) + return X1,X2 + elif self.dim==3: + X1 = np.kron(ei(3),np.kron(ei(2),xin[0])).reshape(-1,1) + X2 = np.kron(ei(3).T,np.kron(xin[1],ei(1).T)).reshape(-1,1) + X3 = np.kron(xin[2],np.kron(ei(2),ei(1))).T.reshape(-1,1) + + return X1,X2,X3 + + def getCellCenteredGrid(self): + """Construct cell-centered grid for 1, 2 and 3 dimensions.""" + + if self.dim==1: + return [self.xic(1)] + elif self.dim==2: + return self.ndgrid([self.xic(1),self.xic(2)]) + elif self.dim==3: + return self.ndgrid([self.xic(1),self.xic(2),self.xic(3)]) + + def getFaceStgGrid(self,direction): + """Construct the face staggered grids for 2 and 3 dimensions.""" + + if self.dim==1: + print 'Error: dimension must be larger than 1' + elif self.dim==2: + if direction == 1: + return self.ndgrid([self.xin(1),self.xic(2)]) + elif direction == 2: + return self.ndgrid([self.xic(1),self.xin(2)]) + else: + print 'Error: direction must be equal to 1 or 2' + elif self.dim==3: + if direction == 1: + return self.ndgrid([self.xin(1),self.xic(2),self.xic(3)]) + elif direction == 2: + return self.ndgrid([self.xic(1),self.xin(2),self.xic(3)]) + elif direction == 3: + return self.ndgrid([self.xic(1),self.xic(2),self.xin(3)]) + else: + print 'Error: direction must be equal to 1, 2 or 3' + + def getEdgeStgGrid(self,direction): + """Construct the edge staggered grids for 3 dimension case.""" + if self.dim != 3: + print 'Error: dimension must be equal to 3' + else: + if direction == 1: + return self.ndgrid([self.xic(1),self.xin(2),self.xin(3)]) + elif direction == 2: + return self.ndgrid([self.xin(1),self.xic(2),self.xin(3)]) + elif direction == 3: + return self.ndgrid([self.xin(1),self.xin(2),self.xic(3)]) + else: + print 'Error: direction must be equal to 1, 2 or 3' + + def plotImage(self,I): + + if self.dim==1: + fig = plt.figure(1) + fig.clf() + ax=plt.subplot(111) + if np.size(I)==self.n[0]: + print 'cell-centered image' + xx = self.getCellCenteredGrid() + ax.plot(xx[0],I,'ro') + elif np.size(I)==self.n[0]+1: + print 'nodal image' + xx = self.getNodalGrid() + ax.plot(xx[0],I,'bs') + + fig.show() + + + def plotGrid(self): + """Plot the nodal, cell-centered and staggered grids for 1,2 and 3 dimensions.""" + if self.dim == 1: + fig = plt.figure(1) + fig.clf() + ax = plt.subplot(111) + xn = self.getNodalGrid() + xc = self.getCellCenteredGrid() + print xn + ax.hold(True) + ax.plot(xn,np.ones(np.shape(xn)),'bs') + ax.plot(xc,np.ones(np.shape(xc)),'ro') + ax.plot(xn,np.ones(np.shape(xn)),'k--') + ax.grid(True) + ax.hold(False) + ax.set_xlabel('x1') + fig.show() + elif self.dim == 2: + fig = plt.figure(2) + fig.clf() + ax = plt.subplot(111) + xn = self.getNodalGrid() + xc = self.getCellCenteredGrid() + xs1 = self.getFaceStgGrid(1) + xs2 = self.getFaceStgGrid(2) + + ax.hold(True) + ax.plot(xn[0],xn[1],'bs') + ax.plot(xc[0],xc[1],'ro') + ax.plot(xs1[0],xs1[1],'g>') + ax.plot(xs2[0],xs2[1],'g^') + ax.grid(True) + ax.hold(False) + ax.set_xlabel('x1') + ax.set_ylabel('x2') + fig.show() + elif self.dim == 3: + fig = plt.figure(3) + fig.clf() + ax = fig.add_subplot(111, projection='3d') + xn = self.getNodalGrid() + xc = self.getCellCenteredGrid() + xfs1 = self.getFaceStgGrid(1) + xfs2 = self.getFaceStgGrid(2) + xfs3 = self.getFaceStgGrid(3) + xes1 = self.getEdgeStgGrid(1) + xes2 = self.getEdgeStgGrid(2) + xes3 = self.getEdgeStgGrid(3) + + ax.hold(True) + ax.plot(xn[0],xn[1],'bs',zs=xn[2]) + ax.plot(xc[0],xc[1],'ro',zs=xc[2]) + ax.plot(xfs1[0],xfs1[1],'g>',zs=xfs1[2]) + ax.plot(xfs2[0],xfs2[1],'g<',zs=xfs2[2]) + ax.plot(xfs3[0],xfs3[1],'g^',zs=xfs3[2]) + ax.plot(xes1[0],xes1[1],'k>',zs=xes1[2]) + ax.plot(xes2[0],xes2[1],'k<',zs=xes2[2]) + ax.plot(xes3[0],xes3[1],'k^',zs=xes3[2]) + ax.grid(True) + ax.hold(False) + ax.set_xlabel('x1') + ax.set_ylabel('x2') + ax.set_zlabel('x3') + fig.show() + + + +if __name__ == '__main__': + print('Welcome to tensor mesh!') + + testDim = 1 + h1 = 0.3*np.ones((1,7)) + h1[:,0] = 0.5 + h1[:,-1] = 0.6 + h2 = .5* np.ones((1,4)) + h3 = .4* np.ones((1,6)) + x0 = np.zeros((3,1)) + + if testDim == 1: + h = [h1] + x0 = x0[0] + elif testDim==2: + h = [h1,h2] + x0 = x0[0:2] + else: + h = [h1,h2,h3] + + I = np.linspace(0,1,8) + M = TensorMesh(h,x0) + + xn = M.plotGrid() + + + \ No newline at end of file