diff --git a/SimPEG/Mesh/NewTreeMesh.py b/SimPEG/Mesh/NewTreeMesh.py index a8741c0f..6e068b87 100644 --- a/SimPEG/Mesh/NewTreeMesh.py +++ b/SimPEG/Mesh/NewTreeMesh.py @@ -18,7 +18,8 @@ def SortByX0(grid): class TreeMesh(object): - def __init__(self, hx, hy): + def __init__(self, h): + hx, hy = h nx = np.r_[0,hx.cumsum()] ny = np.r_[0,hy.cumsum()] vnC = [nx.size-1, ny.size-1] @@ -54,11 +55,12 @@ class TreeMesh(object): @property def isNumbered(self): - return self._numberedCC and self._numberedEx and self._numberedEy + return self._numberedCC and self._numberedN and self._numberedEx and self._numberedEy @isNumbered.setter def isNumbered(self, value): assert value is False, 'Can only set to False.' self._numberedCC = False + self._numberedN = False self._numberedEx = False self._numberedEy = False for name in ['vol', 'area', 'edge', 'gridCC', 'gridN', 'gridEx', 'gridEy', 'gridEz', 'gridFx', 'gridFy', 'gridFz']: @@ -269,6 +271,22 @@ class TreeMesh(object): return self._vol + @property + def gridN(self): + N = self._nodes + activeNodes = N[:,ACTIVE] == 1 + Nx = N[activeNodes,NX] + Ny = N[activeNodes,NY] + + P = SortByX0(np.c_[Nx, Ny]) + if not self._numberedN: + cnt = np.zeros(P.size, dtype=int) + cnt[P] = np.arange(P.size) + self._nodes[activeNodes, NUM] = cnt + self._numberedN = True + + return np.c_[Nx, Ny][P, :] + @property def gridCC(self): N = self._nodes @@ -287,7 +305,7 @@ class TreeMesh(object): self._faces[activeCells, NUM] = cnt self._numberedCC = True - return np.c_[Cx,Cy][P, :] + return np.c_[Cx, Cy][P, :] @property def gridEx(self): @@ -306,7 +324,7 @@ class TreeMesh(object): self._edges[activeEdges, NUM] = cnt self._numberedEx = True - return np.c_[Ex,Ey][P, :] + return np.c_[Ex, Ey][P, :] @property def gridEy(self): @@ -325,7 +343,7 @@ class TreeMesh(object): self._edges[activeEdges, NUM] = cnt + self.nEx self._numberedEy = True - return np.c_[Ex,Ey][P, :] + return np.c_[Ex, Ey][P, :] @property @@ -384,6 +402,7 @@ class TreeMesh(object): self._faces[:,NUM] = -1 # self._cells[:,NUM] = -1 self.gridCC + self.gridN self.gridEx self.gridEy diff --git a/SimPEG/Tests/test_NewTreeMesh.py b/SimPEG/Tests/test_NewTreeMesh.py new file mode 100644 index 00000000..430766ca --- /dev/null +++ b/SimPEG/Tests/test_NewTreeMesh.py @@ -0,0 +1,59 @@ +from SimPEG.Mesh import TensorMesh +from SimPEG.Mesh.NewTreeMesh import TreeMesh +import numpy as np +import unittest +import matplotlib.pyplot as plt + +TOL = 1e-10 + +class TestQuadTreeMesh(unittest.TestCase): + + def setUp(self): + M = TreeMesh([np.ones(x) for x in [3,2]]) + M.refineFace(0) + self.M = M + M.number() + # M.plotGrid(showIt=True) + + def test_MeshSizes(self): + self.assertTrue(self.M.nC==9) + self.assertTrue(self.M.nF==25) + self.assertTrue(self.M.nFx==12) + self.assertTrue(self.M.nFy==13) + self.assertTrue(self.M.nE==25) + self.assertTrue(self.M.nEx==13) + self.assertTrue(self.M.nEy==12) + + def test_gridCC(self): + x = np.r_[0.25,0.75,1.5,2.5,0.25,0.75,0.5,1.5,2.5] + y = np.r_[0.25,0.25,0.5,0.5,0.75,0.75,1.5,1.5,1.5] + self.assertTrue(np.linalg.norm((np.c_[x,y]-self.M.gridCC).flatten()) == 0) + + def test_gridN(self): + x = np.r_[0,0.5,1,2,3,0,0.5,1,0,0.5,1,2,3,0,1,2,3] + y = np.r_[0,0,0,0,0,.5,.5,.5,1,1,1,1,1,2,2,2,2] + self.assertTrue(np.linalg.norm((np.c_[x,y]-self.M.gridN).flatten()) == 0) + + def test_gridFx(self): + x = np.r_[0.0,0.5,1.0,2.0,3.0,0.0,0.5,1.0,0.0,1.0,2.0,3.0] + y = np.r_[0.25,0.25,0.25,0.5,0.5,0.75,0.75,0.75,1.5,1.5,1.5,1.5] + self.assertTrue(np.linalg.norm((np.c_[x,y]-self.M.gridFx).flatten()) == 0) + + def test_gridFy(self): + x = np.r_[0.25,0.75,1.5,2.5,0.25,0.75,0.25,0.75,1.5,2.5,0.5,1.5,2.5] + y = np.r_[0,0,0,0,0.5,0.5,1,1,1,1,2,2,2] + self.assertTrue(np.linalg.norm((np.c_[x,y]-self.M.gridFy).flatten()) == 0) + + def test_gridEx(self): + x = np.r_[0.25,0.75,1.5,2.5,0.25,0.75,0.25,0.75,1.5,2.5,0.5,1.5,2.5] + y = np.r_[0,0,0,0,0.5,0.5,1,1,1,1,2,2,2] + self.assertTrue(np.linalg.norm((np.c_[x,y]-self.M.gridEx).flatten()) == 0) + + def test_gridEy(self): + x = np.r_[0.0,0.5,1.0,2.0,3.0,0.0,0.5,1.0,0.0,1.0,2.0,3.0] + y = np.r_[0.25,0.25,0.25,0.5,0.5,0.75,0.75,0.75,1.5,1.5,1.5,1.5] + self.assertTrue(np.linalg.norm((np.c_[x,y]-self.M.gridEy).flatten()) == 0) + + +if __name__ == '__main__': + unittest.main()