diff --git a/SimPEG/Mesh/TreeMesh.py b/SimPEG/Mesh/TreeMesh.py index 9d616480..82b2e5a7 100644 --- a/SimPEG/Mesh/TreeMesh.py +++ b/SimPEG/Mesh/TreeMesh.py @@ -43,10 +43,12 @@ class TreeObject(object): children = None #: children of the tree object num = None + depth = 0 - def __init__(self, mesh, parent): + def __init__(self, mesh, parent, depth): self.mesh = mesh self._parent = parent + self.depth = depth @property def parent(self): return self._parent @@ -60,11 +62,18 @@ class TreeObject(object): @property def center(self): return self.x0 + @property + def branchdepth(self): + if self.isleaf: + return self.depth + else: + return np.max([node.branchdepth for node in self.children.flatten('F')]) + class TreeNode(TreeObject): """docstring for TreeNode""" - def __init__(self, mesh, x0=[0,0], depth=0, parent=None): - TreeObject.__init__(self, mesh, parent) + def __init__(self, mesh, x0=[0,0], parent=None): + TreeObject.__init__(self, mesh, parent, 0) self.x0 = np.array(x0, dtype=float) self.mesh.nodes.add(self) @@ -74,10 +83,9 @@ class TreeEdge(TreeObject): def __init__(self, mesh, x0=[0,0], edgeType=None, sz=[1,], depth=0, node0=None, node1=None, parent=None): - TreeObject.__init__(self, mesh, parent) + TreeObject.__init__(self, mesh, parent, depth) self.x0 = np.array(x0, dtype=float) - self.depth = depth self.edgeType = edgeType self.sz = np.array(sz, dtype=float) @@ -114,9 +122,9 @@ class TreeEdge(TreeObject): elif self.edgeType is 'y': return np.r_[0,1.,0] elif self.edgeType is 'z': return np.r_[0,0,1.] - def plotGrid(self, ax, text=False): + def plotGrid(self, ax, text=False, lineOpts={'color':'r', 'ls': '-'}): line = np.c_[self.node0.x0, self.node1.x0].T - ax.plot(line[:,0], line[:,1],'r-', zs=line[:,2]) + ax.plot(line[:,0], line[:,1], zs=line[:,2], **lineOpts) class TreeFace(TreeObject): @@ -125,10 +133,9 @@ class TreeFace(TreeObject): node0=None, node1=None, edge0=None, edge1=None, edge2=None, edge3=None, parent=None): - TreeObject.__init__(self, mesh, parent) + TreeObject.__init__(self, mesh, parent, depth) self.x0 = np.array(x0, dtype=float) - self.depth = depth self.faceType = faceType self.sz = np.array(sz, dtype=float) @@ -272,7 +279,7 @@ class TreeFace(TreeObject): for O in order: i, j = O['c'] - x0r = 0.5*i*self.tangent0*self.sz[0] + 0.5*j*self.tangent1*self.sz[1] + x0r = self.x0 + 0.5*i*self.tangent0*self.sz[0] + 0.5*j*self.tangent1*self.sz[1] e0, e1, e2, e3 = getEdge(O['e0']), getEdge(O['e1']), getEdge(O['e2']), getEdge(O['e3']) self.children[i,j] = TreeFace(self.mesh, x0=x0r, faceType=self.faceType, depth=self.depth+1, sz=0.5*self.sz, parent=self, edge0=e0, edge1=e1, edge2=e2, edge3=e3) @@ -310,11 +317,10 @@ class TreeCell(TreeObject): fYm=None, fYp=None, fZm=None, fZp=None, parent=None): - TreeObject.__init__(self, mesh, parent) + TreeObject.__init__(self, mesh, parent, depth) self.x0 = np.array(x0, dtype=float) self.sz = np.array(sz, dtype=float) - self.depth = depth if self.dim == 2: # # 2___________3 @@ -422,13 +428,6 @@ class TreeCell(TreeObject): mesh.cells.add(self) - @property - def branchdepth(self): - if self.isleaf: - return self.depth - else: - return np.max([node.branchdepth for node in self.children.flatten('F')]) - @property def center(self): return self.x0 + 0.5*self.sz diff --git a/SimPEG/Tests/test_TreeMesh.py b/SimPEG/Tests/test_TreeMesh.py index 3c10b64a..e906ec3e 100644 --- a/SimPEG/Tests/test_TreeMesh.py +++ b/SimPEG/Tests/test_TreeMesh.py @@ -7,19 +7,13 @@ class TestOcTreeObjects(unittest.TestCase): def setUp(self): self.M = TreeMesh([2,1,1]) + self.M.number() + self.Mr = TreeMesh([2,1,1]) self.Mr.children[0,0,0].refine() self.Mr.number() def test_counts(self): - ax = plt.subplot(111,projection='3d') - # self.Mr.plotGrid(showIt=False,plotC=True,plotEy=True) - - cell = self.Mr.sortedCells[1] - [cell.edges[e].plotGrid(ax) for e in cell.edges] - cell.plotGrid(ax) - plt.show() - self.assertTrue(self.M.nC == 2) self.assertTrue(self.M.nFx == 3) self.assertTrue(self.M.nFy == 4) @@ -37,22 +31,14 @@ class TestOcTreeObjects(unittest.TestCase): self.assertTrue(self.Mr.nFz == 14) self.assertTrue(self.Mr.nF == 41) - cell = self.Mr.sortedCells[1] - self.assertTrue(cell.edges['eX0'].edgeType=='x') - self.assertTrue(cell.edges['eX1'].edgeType=='x') - self.assertTrue(cell.edges['eX2'].edgeType=='x') - self.assertTrue(cell.edges['eX3'].edgeType=='x') - self.assertTrue(cell.edges['eY0'].edgeType=='y') - self.assertTrue(cell.edges['eY1'].edgeType=='y') - self.assertTrue(cell.edges['eY2'].edgeType=='y') - self.assertTrue(cell.edges['eY3'].edgeType=='y') - self.assertTrue(cell.edges['eZ0'].edgeType=='z') - self.assertTrue(cell.edges['eZ1'].edgeType=='z') - self.assertTrue(cell.edges['eZ2'].edgeType=='z') - self.assertTrue(cell.edges['eZ3'].edgeType=='z') - print self.Mr.nN + for cell in self.Mr.sortedCells: + for e in cell.edges: + self.assertTrue(cell.edges[e].edgeType==e[1].lower()) + + # print self.Mr.nEx # self.assertTrue(self.Mr.nN == 22) + # self.assertTrue(self.Mr.nEx == 22) def test_pointersM(self): c0 = self.M.children[0,0,0] @@ -60,12 +46,16 @@ class TestOcTreeObjects(unittest.TestCase): c0fXp = c0.faces['fXp'] c0fYm = c0.faces['fYm'] c0fYp = c0.faces['fYp'] + c0fZm = c0.faces['fZm'] + c0fZp = c0.faces['fZp'] c1 = self.M.children[1,0,0] c1fXm = c1.faces['fXm'] c1fXp = c1.faces['fXp'] c1fYm = c1.faces['fYm'] c1fYp = c1.faces['fYp'] + c1fZm = c1.faces['fZm'] + c1fZp = c1.faces['fZp'] self.assertTrue(c0fXp is c1fXm) self.assertTrue(c0fXp.edges['e0'] is c1fXm.edges['e0']) @@ -75,68 +65,112 @@ class TestOcTreeObjects(unittest.TestCase): self.assertTrue(c0fYp is not c1fYm) self.assertTrue(c0fXm is not c1fXm) + # Test connectivity of shared edges + self.assertTrue(c0fZp.edges['e3'] is not c1fZp.edges['e0']) + self.assertTrue(c0fZp.edges['e3'] is not c1fZp.edges['e1']) + self.assertTrue(c0fZp.edges['e3'] is c1fZp.edges['e2']) + self.assertTrue(c0fZp.edges['e3'] is not c1fZp.edges['e3']) + + self.assertTrue(c0fZm.edges['e3'] is not c1fZm.edges['e0']) + self.assertTrue(c0fZm.edges['e3'] is not c1fZm.edges['e1']) + self.assertTrue(c0fZm.edges['e3'] is c1fZm.edges['e2']) + self.assertTrue(c0fZm.edges['e3'] is not c1fZm.edges['e3']) + + self.assertTrue(c0fYp.edges['e3'] is not c1fYp.edges['e0']) + self.assertTrue(c0fYp.edges['e3'] is not c1fYp.edges['e1']) + self.assertTrue(c0fYp.edges['e3'] is c1fYp.edges['e2']) + self.assertTrue(c0fYp.edges['e3'] is not c1fYp.edges['e3']) + + self.assertTrue(c0fYm.edges['e3'] is not c1fYm.edges['e0']) + self.assertTrue(c0fYm.edges['e3'] is not c1fYm.edges['e1']) + self.assertTrue(c0fYm.edges['e3'] is c1fYm.edges['e2']) + self.assertTrue(c0fYm.edges['e3'] is not c1fYm.edges['e3']) + + self.assertTrue(c0fZm.edges['e3'] is c1fXm.edges['e0']) + self.assertTrue(c0fZp.edges['e3'] is c1fXm.edges['e1']) + self.assertTrue(c0fYm.edges['e3'] is c1fXm.edges['e2']) + self.assertTrue(c0fYp.edges['e3'] is c1fXm.edges['e3']) + + self.assertTrue(c0fZm.edges['e3'] is c0fXp.edges['e0']) + self.assertTrue(c0fZp.edges['e3'] is c0fXp.edges['e1']) + self.assertTrue(c0fYm.edges['e3'] is c0fXp.edges['e2']) + self.assertTrue(c0fYp.edges['e3'] is c0fXp.edges['e3']) + + self.assertTrue(c1fZm.edges['e2'] is c1fXm.edges['e0']) + self.assertTrue(c1fZp.edges['e2'] is c1fXm.edges['e1']) + self.assertTrue(c1fYm.edges['e2'] is c1fXm.edges['e2']) + self.assertTrue(c1fYp.edges['e2'] is c1fXm.edges['e3']) + + self.assertTrue(c1fZm.edges['e2'] is c0fXp.edges['e0']) + self.assertTrue(c1fZp.edges['e2'] is c0fXp.edges['e1']) + self.assertTrue(c1fYm.edges['e2'] is c0fXp.edges['e2']) + self.assertTrue(c1fYp.edges['e2'] is c0fXp.edges['e3']) + + def test_pointersMr(self): + ax = plt.subplot(111, projection='3d') + self.Mr.plotGrid(ax=ax,showIt=False,plotC=True,plotEy=True, text=False) + + cell = self.Mr.sortedCells[1] + [cell.edges[e].plotGrid(ax,lineOpts={'color':'b','ls':'-'}) for e in cell.edges] + cell.plotGrid(ax) + plt.show() + + def q(s): + c = self.Mr.sortedCells[int(s[1])] + if len(s) == 2: return c + if s[2] == 'f': return c.faces[s[2:]] + if s[2] == 'e': return c.edges[s[2:]] + c0 = self.Mr.sortedCells[0] c0fXm = c0.faces['fXm'] - c0fXp = c0.faces['fXp'] - c0fYm = c0.faces['fYm'] - c0fYp = c0.faces['fYp'] - c0fZm = c0.faces['fZm'] - c0fZp = c0.faces['fZp'] - self.assertTrue(np.all(c0.center==np.r_[0.125,0.25,0.25])) + c0eX0 = c0.edges['eX0'] + self.assertTrue(c0 is q('c0')) + self.assertTrue(c0fXm is q('c0fXm')) + self.assertTrue(c0eX0 is q('c0eX0')) - c1 = self.Mr.sortedCells[1] - c1fXm = c1.faces['fXm'] - c1fXp = c1.faces['fXp'] - c1fYm = c1.faces['fYm'] - c1fYp = c1.faces['fYp'] - c1fZm = c1.faces['fZm'] - c1fZp = c1.faces['fZp'] - self.assertTrue(np.all(c1.center==np.r_[0.375,0.25,0.25])) + self.assertTrue(q('c0').depth == 1) + self.assertTrue(q('c1').depth == 1) + self.assertTrue(q('c2').depth == 0) - c2 = self.Mr.sortedCells[2] - c2fXm = c2.faces['fXm'] - c2fXp = c2.faces['fXp'] - c2fYm = c2.faces['fYm'] - c2fYp = c2.faces['fYp'] - c2fZm = c2.faces['fZm'] - c2fZp = c2.faces['fZp'] - self.assertTrue(np.all(c2.center==np.r_[0.75,0.5,0.5])) + # Make sure we know where the center of the cells are. + self.assertTrue(np.all(q('c0').center == np.r_[0.125,0.25,0.25])) + self.assertTrue(np.all(q('c1').center == np.r_[0.375,0.25,0.25])) + self.assertTrue(np.all(q('c2').center == np.r_[0.75,0.5,0.5])) + self.assertTrue(np.all(q('c3').center == np.r_[0.125,0.75,0.25])) + self.assertTrue(np.all(q('c4').center == np.r_[0.375,0.75,0.25])) + self.assertTrue(np.all(q('c5').center == np.r_[0.125,0.25,0.75])) + self.assertTrue(np.all(q('c6').center == np.r_[0.375,0.25,0.75])) + self.assertTrue(np.all(q('c7').center == np.r_[0.125,0.75,0.75])) + self.assertTrue(np.all(q('c8').center == np.r_[0.375,0.75,0.75])) - c4 = self.Mr.sortedCells[4] - c4fXm = c4.faces['fXm'] - c4fXp = c4.faces['fXp'] - c4fYm = c4.faces['fYm'] - c4fYp = c4.faces['fYp'] - c4fZm = c4.faces['fZm'] - c4fZp = c4.faces['fZp'] - self.assertTrue(np.all(c4.center==np.r_[0.375,0.75,0.25])) + self.assertTrue(np.all(q('c0fXm').center == np.r_[0,0.25,0.25])) + self.assertTrue(np.all(q('c0fXp').center == np.r_[0.25,0.25,0.25])) + self.assertTrue(q('c0fXp') is q('c1fXm')) + self.assertTrue(np.all(q('c1fXp').center == np.r_[0.5,0.25,0.25])) + self.assertTrue(np.all(q('c2fXm').center == np.r_[0.5,0.5,0.5])) + self.assertTrue(q('c2fXm').branchdepth == 1) + self.assertTrue(q('c1fXp').parent is q('c2fXm')) + self.assertTrue(q('c2fXm').children[0,0] is q('c1fXp')) - c6 = self.Mr.sortedCells[6] - c6fXm = c6.faces['fXm'] - c6fXp = c6.faces['fXp'] - c6fYm = c6.faces['fYm'] - c6fYp = c6.faces['fYp'] - c6fZm = c6.faces['fZm'] - c6fZp = c6.faces['fZp'] - self.assertTrue(np.all(c6.center==np.r_[0.375,0.25,0.75])) + self.assertTrue(q('c0fXp') is q('c1fXm')) + self.assertTrue(q('c0fYp') is not q('c1fYm')) + self.assertTrue(q('c0fXm') is not q('c1fXm')) - self.assertTrue(c0fXp is c1fXm) - self.assertTrue(c0fYp is not c1fYm) - self.assertTrue(c0fXm is not c1fXm) + self.assertTrue(q('c1fXp') is q('c2fXm').children[0,0]) + self.assertTrue(q('c1fXp').parent is q('c2fXm')) - self.assertTrue(c1fXp is c2fXm.children[0,0]) - self.assertTrue(c1fXp.parent is c2fXm) + self.assertTrue(q('c1fYp') is q('c4fYm')) + self.assertTrue(q('c1fZp') is q('c6fZm')) - self.assertTrue(c1fYp is c4fYm) - self.assertTrue(c1fZp is c6fZm) + self.assertTrue(q('c6fXp') is q('c2fXm').children[0,1]) + self.assertTrue(q('c6fXp').parent is q('c2fXm')) - self.assertTrue(c6fXp is c2fXm.children[0,1]) - self.assertTrue(c6fXp.parent is c2fXm) + self.assertTrue(q('c4fXp') is q('c2fXm').children[1,0]) + self.assertTrue(q('c4fXp').parent is q('c2fXm')) - self.assertTrue(c4fXp is c2fXm.children[1,0]) - self.assertTrue(c4fXp.parent is c2fXm) + #TEST EDGES! def test_gridCC(self): x = np.r_[0.25,0.75]