From a01b8057133644be8ae8f0a8e29a7a0a91d0be63 Mon Sep 17 00:00:00 2001 From: rowanc1 Date: Thu, 19 Dec 2013 18:30:29 -0700 Subject: [PATCH] Updates to visualization, gridCC, and refine functionality --- SimPEG/mesh/QuadTreeMesh.py | 106 ++++++++++++++++++++++++------------ 1 file changed, 72 insertions(+), 34 deletions(-) diff --git a/SimPEG/mesh/QuadTreeMesh.py b/SimPEG/mesh/QuadTreeMesh.py index 3d4450ff..6fe6189d 100644 --- a/SimPEG/mesh/QuadTreeMesh.py +++ b/SimPEG/mesh/QuadTreeMesh.py @@ -75,11 +75,10 @@ class TreeFace(object): self.mesh.faceY.remove(self) - def viz(self, ax): + def viz(self, ax, text=True): if not self.isleaf: return - ax.plot(np.r_[self.x0[0],self.x0[0]+self.tangent[0]*self.sz], np.r_[self.x0[1], self.x0[1]+self.tangent[1]*self.sz],'rx-') - # if self.faceType is 'y': - ax.text(self.x0[0]+0.5*self.tangent[0]*self.sz, self.x0[1]+0.5*self.tangent[1]*self.sz,self.numFace) + ax.plot(np.r_[self.x0[0],self.x0[0]+self.tangent[0]*self.sz], np.r_[self.x0[1], self.x0[1]+self.tangent[1]*self.sz],'r-') + if text: ax.text(self.x0[0]+0.5*self.tangent[0]*self.sz, self.x0[1]+0.5*self.tangent[1]*self.sz,self.numFace) @@ -118,9 +117,16 @@ class TreeNode(object): else: return np.max([node.branchdepth for node in self.children.flatten('F')]) + @property + def gridCC(self): return self.x0 + 0.5*self.sz + + def refine(self, function=None): + if not self.isleaf and function is None: return + + if function is not None: + do = function(self.gridCC) > self.depth + if not do: return - def refine(self): - if not self.isleaf: return self.mesh.isNumbered = False self.children = np.empty((2,2),dtype=TreeNode) @@ -151,6 +157,11 @@ class TreeNode(object): self.mesh.cells.remove(self) + # pass the refine function to the children + if function is not None: + for child in self.children.flatten(): + child.refine(function) + @property def isleaf(self): return self.children is None @@ -170,7 +181,7 @@ class TreeNode(object): def vol(self): return self.sz.prod() - def viz(self, ax): + def viz(self, ax, text=True): if self.isleaf: x0, sz = self.x0, self.sz corners = np.c_[np.r_[x0[0] , x0[1] ], @@ -179,7 +190,7 @@ class TreeNode(object): np.r_[x0[0] , x0[1]+sz[1]], np.r_[x0[0] , x0[1] ]].T ax.plot(corners[:,0],corners[:,1], 'b') - if self.numCell is not None: + if self.numCell is not None and text: ax.text(x0[0]+sz[0]/2,x0[1]+sz[1]/2,'%d'%self.numCell) else: [node.viz(ax) for node in self.children.flatten('F')] @@ -189,17 +200,20 @@ class TreeNode(object): class QuadTreeMesh(object): """docstring for QuadTreeMesh""" - def __init__(self, cells, sz): + def __init__(self, h, x0=[0,0]): self.faces = set() self.faceX = set() self.faceY = set() self.cells = set() - self.children = np.empty(cells,dtype=TreeNode) - for i in range(cells[0]): - for j in range(cells[1]): + self.x0 = np.array(x0,dtype=float) + self.children = np.empty([hi.size for hi in h],dtype=TreeNode) + for i in range(h[0].size): + for j in range(h[1].size): fXm = None if i is 0 else self.children[i-1][j].faces['fXp'] fYm = None if j is 0 else self.children[i][j-1].faces['fYp'] - self.children[i][j] = TreeNode(self, x0=[i*sz[0],j*sz[1]], dim=2, depth=0, sz=sz, fXm=fXm, fYm=fYm) + x0i = (np.r_[x0[0], h[0][:i]]).sum() + x0j = (np.r_[x0[1], h[1][:j]]).sum() + self.children[i][j] = TreeNode(self, x0=[x0i, x0j], dim=len(h), depth=0, sz=[h[0][i], h[1][j]], fXm=fXm, fYm=fYm) isNumbered = utils.dependentProperty('_isNumbered', False, ['_faceDiv'], 'Setting this to False will delete all operators.') @@ -207,50 +221,66 @@ class QuadTreeMesh(object): def branchdepth(self): return np.max([node.branchdepth for node in self.children.flatten('F')]) + def refine(self, function): + for node in self.children.flatten(): + node.refine(function) def number(self): if self.isNumbered: return - sortCells = sorted(M.cells,key=SortByX0()) - sortFaceX = sorted(M.faceX,key=SortByX0()) - sortFaceY = sorted(M.faceY,key=SortByX0()) - nFx = len(sortFaceX) - for i, sc in enumerate(sortCells): sc.numCell = i - for i, sfx in enumerate(sortFaceX): sfx.numFace = i - for i, sfy in enumerate(sortFaceY): sfy.numFace = i + nFx + sortedCells = sorted(M.cells,key=SortByX0()) + sortedFaceX = sorted(M.faceX,key=SortByX0()) + sortedFaceY = sorted(M.faceY,key=SortByX0()) + nFx = len(sortedFaceX) + for i, sc in enumerate(sortedCells): sc.numCell = i + for i, sfx in enumerate(sortedFaceX): sfx.numFace = i + for i, sfy in enumerate(sortedFaceY): sfy.numFace = i + nFx - self.sortCells = sortCells - self.sortFaceX = sortFaceX - self.sortFaceY = sortFaceY + self.sortedCells = sortedCells + self.sortedFaceX = sortedFaceX + self.sortedFaceY = sortedFaceY self.isNumbered = True - def viz(self, ax=None): + def viz(self, ax=None, text=True): if ax is None: ax = plt.subplot(111) - [node.viz(ax) for node in self.cells] - [node.viz(ax) for node in self.faces] + # [node.viz(ax, text=text) for node in self.cells] + [node.viz(ax, text=text) for node in self.faces] @property def nC(self): return len(self.cells) @property def nF(self): return len(self.faces) + @property + def nFx(self): return len(self.facesX) + @property + def nFy(self): return len(self.facesY) + + @property + def gridCC(self): + if getattr(self, '_gridCC', None) is None: + self.number() + self._gridCC = np.empty((self.nC,self.dim)) + for ii, cell in enumerate(self.sortedCells): + self._gridCC[ii,:] = cell.gridCC + return self._gridCC @property def vol(self): self.number() - return np.array([cell.vol for cell in self.sortCells]) + return np.array([cell.vol for cell in self.sortedCells]) @property def area(self): self.number() - return np.concatenate(([face.area for face in self.sortFaceX],[face.area for face in self.sortFaceY])) + return np.concatenate(([face.area for face in self.sortedFaceX],[face.area for face in self.sortedFaceY])) @property def faceDiv(self): if getattr(self, '_faceDiv', None) is None: self.number() I, J, V = np.empty(0), np.empty(0), np.empty(0) - for cell in M.cells: + for cell in M.sortedCells: i, j, v = cell.faceIndex I, J, V = np.r_[I,i], np.r_[J,j], np.r_[V,v] @@ -263,16 +293,24 @@ class QuadTreeMesh(object): if __name__ == '__main__': - M = QuadTreeMesh([3,2],[1,2]) - for ii in range(1): - M.children[ii,ii].refine() - # M.children[ii,ii].children[0,0].refine() + M = QuadTreeMesh([np.ones(x) for x in [4,10]]) + def function(xc): + r = xc - np.r_[2.,6.] + dist = np.sqrt(r.dot(r)) + if dist < 1.0: + return 3 + if dist < 1.5: + return 2 + else: + return 1 + + M.refine(function) DIV = M.faceDiv plt.subplot(211) plt.spy(DIV) - M.viz(ax=plt.subplot(212)) + M.viz(ax=plt.subplot(212),text=False) # plt.gca().invert_yaxis() plt.show()