test simple grid functions add Nodal Grid

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
 

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()