diff --git a/SimPEG/Mesh/NewTreeMesh.py b/SimPEG/Mesh/NewTreeMesh.py index e18464cb..fa1aadd6 100644 --- a/SimPEG/Mesh/NewTreeMesh.py +++ b/SimPEG/Mesh/NewTreeMesh.py @@ -143,6 +143,28 @@ class TreeFace(TreeIndexer): V = 0.25 * (4.0*(p**2)*(q**2) - (a**2 + c**2 - b**2 - d**2)**2)**0.5 return V + @property + def children(self): + if self.isleaf: return None + ind = int(self.index) # can not get children of a fancy slice at the moment. + subInds = np.argwhere(self.F[:,PARENT] == ind).flatten() + return TreeFace(self.M, subInds) + + + def refine(self, function, level): + + int(self.index) # should only be able to refine one at a time. + + toLevel = function(self) + + if toLevel < level+1: + return + + inds, rows = self.M.refineFace(self.index) + for i in inds: + TreeFace(self.M, i).refine(function, level + 1) + + class TreeCell(TreeIndexer): _SubTree = TreeFace _pointer = '_cells' @@ -240,6 +262,9 @@ class TreeCell(TreeIndexer): @property def n7(self): return self.fZp.n3 @property + def nodes(self): + return [self.n0, self.n1, self.n2, self.n3, self.n4, self.n5, self.n6, self.n7] + @property def center(self): return (self.n0.vec + self.n1.vec + self.n2.vec + self.n3.vec + self.n4.vec + self.n5.vec + self.n6.vec + self.n7.vec)/8.0 @@ -279,6 +304,26 @@ class TreeCell(TreeIndexer): return (vol1 + vol2)/2.0 + @property + def children(self): + if self.isleaf: return None + ind = int(self.index) # can not get children of a fancy slice at the moment. + subInds = np.argwhere(self.C[:,PARENT] == ind).flatten() + return TreeCell(self.M, subInds) + + def refine(self, function, level): + + int(self.index) # should only be able to refine one at a time. + + toLevel = function(self) + + if toLevel < level+1: + return + + inds, rows = self.M.refineCell(self.index) + for i in inds: + TreeCell(self.M, i).refine(function, level + 1) + class TreeMesh(InnerProducts, BaseMesh): @@ -687,7 +732,7 @@ class TreeMesh(InnerProducts, BaseMesh): E2i, E2 = self.refineEdge(f[FEDGE2]) E3i, E3 = self.refineEdge(f[FEDGE3]) - nodeNums = self._edges[f[[FEDGE0, FEDGE1]],:][:,[ENODE0, ENODE1]] + nodeNums = [n.index for n in TreeFace(self, index).nodes] newNode, node = self.addNode(nodeNums) # Refine the inner edges @@ -1043,6 +1088,12 @@ class TreeMesh(InnerProducts, BaseMesh): return sp.vstack((xP, yP, zP)) return Pxxx + def refine(self, function): + if self.dim == 3: + TreeCell(self, 0).refine(function, 0) + elif self.dim == 2: + TreeFace(self, 0).refine(function, 0) + def plotGrid(self, ax=None, text=True, showIt=False): import matplotlib.pyplot as plt @@ -1050,22 +1101,63 @@ class TreeMesh(InnerProducts, BaseMesh): axOpts = {'projection':'3d'} if self.dim == 3 else {} if ax is None: ax = plt.subplot(111, **axOpts) + if self.dim == 3: + C = TreeCell(self, 'active') + + + # fZp + # | + # 6 ------eX3------ 7 + # /| | / | + # /eZ2 . / eZ3 + # eY2 | fYp eY3 | + # / | / fXp| + # 4 ------eX2----- 5 | + # |fXm 2 -----eX1--|---- 3 z + # eZ0 / | eY1 ^ y + # | eY0 . fYm eZ1 / | / + # | / | | / | / + # 0 ------eX0------1 o----> x + # | + # fZm + # + + n1, n2, n3, n4, n5 = 0, 1, 3, 2, 0 + eX = np.c_[C.nodes[n1].x, C.nodes[n2].x, C.nodes[n3].x, C.nodes[n4].x, C.nodes[n5].x, [np.nan]*self.nC] + eY = np.c_[C.nodes[n1].y, C.nodes[n2].y, C.nodes[n3].y, C.nodes[n4].y, C.nodes[n5].y, [np.nan]*self.nC] + eZ = np.c_[C.nodes[n1].z, C.nodes[n2].z, C.nodes[n3].z, C.nodes[n4].z, C.nodes[n5].z, [np.nan]*self.nC] + ax.plot(eX.flatten(), eY.flatten(), 'b-', zs=eZ.flatten()) + + n1, n2, n3, n4, n5 = 4, 5, 7, 6, 4 + eX = np.c_[C.nodes[n1].x, C.nodes[n2].x, C.nodes[n3].x, C.nodes[n4].x, C.nodes[n5].x, [np.nan]*self.nC] + eY = np.c_[C.nodes[n1].y, C.nodes[n2].y, C.nodes[n3].y, C.nodes[n4].y, C.nodes[n5].y, [np.nan]*self.nC] + eZ = np.c_[C.nodes[n1].z, C.nodes[n2].z, C.nodes[n3].z, C.nodes[n4].z, C.nodes[n5].z, [np.nan]*self.nC] + ax.plot(eX.flatten(), eY.flatten(), 'r-', zs=eZ.flatten()) + + ax.plot(self.gridN[:,0], self.gridN[:,1], 'bs', zs=self.gridN[:,2]) + + ax.set_xlabel('x') + ax.set_ylabel('y') + ax.set_zlabel('z') + if showIt: plt.show() + return ax + N = self._nodes E = self._edges C = self._faces - plt.plot(N[:,1], N[:,2], 'b.') + ax.plot(N[:,1], N[:,2], 'b.') activeCells = C[:,ACTIVE] == 1 for FEDGE in [FEDGE0, FEDGE1, FEDGE2, FEDGE3]: nInds = E[C[activeCells,FEDGE],:][:,[ENODE0,ENODE1]] eX = np.c_[N[nInds[:,0],NX], N[nInds[:,1],NX], [np.nan]*nInds.shape[0]] eY = np.c_[N[nInds[:,0],NY], N[nInds[:,1],NY], [np.nan]*nInds.shape[0]] - plt.plot(eX.flatten(), eY.flatten(), 'b-') + ax.plot(eX.flatten(), eY.flatten(), 'b-') gridCC = self.gridCC if text: [ax.text(cc[0], cc[1],i) for i, cc in enumerate(gridCC)] - plt.plot(gridCC[:,0], gridCC[:,1], 'r.') + ax.plot(gridCC[:,0], gridCC[:,1], 'r.') gridFx = self.gridFx gridFy = self.gridFy if text: @@ -1092,32 +1184,44 @@ if __name__ == '__main__': from SimPEG import Mesh, Utils import matplotlib.pyplot as plt - tM = TreeMesh([np.ones(3),np.ones(2)]) + # tM = TreeMesh([np.ones(3),np.ones(2)]) + tM = TreeMesh([1,1,1]) - tM.refineFace(0) - tM.refineFace(1) - tM.refineFace(3) - tM.refineFace(9) + tM.refine(lambda c:2) + # tM.refineCell(4) + + M = Mesh.TensorMesh([4,4,4]) + print tM.gridN - M.gridN + tM.plotGrid() + + plt.show() + + + + # tM.refineFace(0) + # tM.refineFace(1) + # tM.refineFace(3) + # tM.refineFace(9) # print tM._nodes[:,NUM] - tM.number() + # tM.number() # print tM._nodes[:,NUM] # print tM._edges[:,NUM] - print TreeFace(tM,'active').e3.n1.index + # print TreeFace(tM,'active').e3.n1.index - Mr = TreeMesh([2,1,1]) - Mr.refineCell(0) + # Mr = TreeMesh([2,1,1]) + # Mr.refineCell(0) - print Mr.vol + # print Mr.vol - M = TreeMesh([2,2,2]) - M.number() - M.refineCell(0) - M.refineCell(3) - assert M.isNumbered is False - C = TreeCell(M, 'active') - M.getEdgeInnerProduct() + # M = TreeMesh([2,2,2]) + # M.number() + # M.refineCell(0) + # M.refineCell(3) + # assert M.isNumbered is False + # C = TreeCell(M, 'active') + # M.getEdgeInnerProduct() # tM = TreeMesh([100,100,100]) # print tM.vol @@ -1141,3 +1245,18 @@ if __name__ == '__main__': # plt.figure(2) # plt.plot(SortByX0(tM.gridCC),'b.') # plt.show() + + + # M = TreeMesh([1,1,1]) + + # def refFunc(cell): + # n = 3 - np.sum((cell.center.flatten() - np.r_[0.5, 0.5, 0.5])**2)**0.5 * 2 + # print n, cell.center + # return n + # return 1 + # M.refine(refFunc) + + # M.plotGrid(text=False) + # plt.show() + + # print M.nC diff --git a/SimPEG/Tests/test_NewTreeMesh.py b/SimPEG/Tests/test_NewTreeMesh.py index 503834a4..f08091a5 100644 --- a/SimPEG/Tests/test_NewTreeMesh.py +++ b/SimPEG/Tests/test_NewTreeMesh.py @@ -1,5 +1,5 @@ from SimPEG.Mesh import TensorMesh -from SimPEG.Mesh.NewTreeMesh import TreeMesh +from SimPEG.Mesh.NewTreeMesh import TreeMesh, TreeCell import numpy as np import unittest import matplotlib.pyplot as plt @@ -75,6 +75,85 @@ class TestQuadTreeMesh(unittest.TestCase): ay = np.r_[0.5,0.5,1,1,0.5,0.5,0.5,0.5,1,1,1,1,1] self.assertTrue(np.linalg.norm((np.r_[ax,ay]-self.M.area)) < TOL) +class TestOcTreeConnectivity(unittest.TestCase): + + def setUp(self): + self.oM = TreeMesh([1,1,1]) + self.oM.refine(lambda c: 1) + + def test_setup(self): + C = TreeCell(self.oM, 0) + children = C.children + assert not C.isleaf + assert len(children.index) == 8 + # assert not TreeCell(self.oM, 0).isleaf + c0, c1, c2, c3, c4, c5, c6, c7 = [TreeCell(self.oM, i) for i in range(1,9)] + + + # .----------------.----------------. + # /| /| /| + # / | / | / | + # / | c6 / | c7 / | + # / | / | / | + # .----------------.----+-----------. | + # /| . ---------/|----.----------/|----. + # fZp / | /| / | /| / | /| + # | / | / | c4 / | / | c5 / | X | + # 6 ------eX3------ 7 / | / | / | / | / | / | + # /| | / | . -------------- .----------------. |/ | + # /eZ2 . / eZ3 | . ---+------|----.----+------|----. | + # eY2 | fYp eY3 | | /| .______|___/|____.______|___/|____. + # / | / fXp| | / | / c2 | / | / c3 | / | / + # 4 ------eX2----- 5 | | / | / | / | / | / | / + # |fXm 2 -----eX1--|---- 3 z . ---+---------- . ---+---------- . | / + # eZ0 / | eY1 ^ y | |/ | |/ | |/ + # | eY0 . fYm eZ1 / | / | . ----------|----.-----------|----. + # | / | | / | / | / c0 | / c1 | / + # 0 ------eX0------1 o----> x | / | / | / + # | | / | / | / + # fZm . -------------- . -------------- . + # + # + # fX fY fZ + # 2___________3 2___________3 2___________3 + # | e1 | | e1 | | e1 | + # | | | | | | + # e2 | x | e3 z e2 | x | e3 z e2 | x | e3 y + # | | ^ | | ^ | | ^ + # |___________| |___> y |___________| |___> x |___________| |___> x + # 0 e0 1 0 e0 1 0 e0 1 + # + + + # there are two faces for each edge + for ii, c in enumerate([c0, c1, c2, c3, c4, c5, c6, c7]): + assert c.fZm.e0.index == c.fYm.e0.index, "Cell %d: fZm.e0 and fYm.e0"%ii + assert c.fZm.e1.index == c.fYp.e0.index, "Cell %d: fZm.e1 and fYp.e0"%ii + assert c.fZp.e0.index == c.fYm.e1.index, "Cell %d: fZp.e0 and fYm.e1"%ii + assert c.fZp.e1.index == c.fYp.e1.index, "Cell %d: fZp.e1 and fYp.e1"%ii + assert c.fZm.e2.index == c.fXm.e0.index, "Cell %d: fZm.e2 and fXm.e0"%ii + assert c.fZm.e3.index == c.fXp.e0.index, "Cell %d: fZm.e3 and fXp.e0"%ii + assert c.fZp.e2.index == c.fXm.e1.index, "Cell %d: fZp.e2 and fXm.e1"%ii + assert c.fZp.e3.index == c.fXp.e1.index, "Cell %d: fZp.e3 and fXp.e1"%ii + assert c.fYm.e2.index == c.fXm.e2.index, "Cell %d: fYm.e2 and fXm.e2"%ii + assert c.fYm.e3.index == c.fXp.e2.index, "Cell %d: fYm.e3 and fXp.e2"%ii + assert c.fYp.e2.index == c.fXm.e3.index, "Cell %d: fYp.e2 and fXm.e3"%ii + assert c.fYp.e3.index == c.fXp.e3.index, "Cell %d: fYp.e3 and fXp.e3"%ii + + assert c0.eZ1.index == c1.eZ0.index + assert c0.eZ3.index == c1.eZ2.index + assert c2.eZ1.index == c3.eZ0.index + assert c2.eZ3.index == c3.eZ2.index + + assert c4.eZ1.index == c5.eZ0.index + assert c4.eZ3.index == c5.eZ2.index + assert c6.eZ1.index == c7.eZ0.index + assert c6.eZ3.index == c7.eZ2.index + + assert c0.n7.index == c7.n0.index + + + class SimpleOctreeOperatorTests(unittest.TestCase): @@ -107,6 +186,18 @@ class SimpleOctreeOperatorTests(unittest.TestCase): # self.assertAlmostEqual((self.tM2.getFaceInnerProduct() - self.oM2.getFaceInnerProduct()).toarray().sum(), 0) # self.assertAlmostEqual((self.tM2.getEdgeInnerProduct() - self.oM2.getEdgeInnerProduct()).toarray().sum(), 0) + def test_grids(self): + tM = TreeMesh([1,1,1]) + tM.refine(lambda c:2) + M = TensorMesh([4,4,4]) + self.assertAlmostEqual((tM.gridN - M.gridN).sum(), 0) + self.assertAlmostEqual((tM.gridCC - M.gridCC).sum(), 0) + self.assertAlmostEqual((tM.gridFx - M.gridFx).sum(), 0) + self.assertAlmostEqual((tM.gridFy - M.gridFy).sum(), 0) + self.assertAlmostEqual((tM.gridFz - M.gridFz).sum(), 0) + self.assertAlmostEqual((tM.gridEx - M.gridEx).sum(), 0) + self.assertAlmostEqual((tM.gridEy - M.gridEy).sum(), 0) + self.assertAlmostEqual((tM.gridEz - M.gridEz).sum(), 0) class TestOcTreeObjects(unittest.TestCase):