Test pointerTree

SimPEG/Mesh/PointerTree.py

@@ -1,8 +1,52 @@
-import scurve
 from SimPEG import np, sp, Utils, Solver
 import matplotlib.pyplot as plt
 import matplotlib
 
+class ZCurve(object):
+    """
+        The Z-order curve is generated by interleaving the bits of an offset.
+
+        See:
+
+            https://github.com/cortesi/scurve
+            Aldo Cortesi <aldo@corte.si>
+
+    """
+    def __init__(self, dimension, bits):
+        """
+            dimension: Number of dimensions
+            bits: The number of bits per co-ordinate. Total number of points is
+            2**(bits*dimension).
+        """
+        self.dimension, self.bits = dimension, bits
+
+    def bitrange(self, x, width, start, end):
+        """
+            Extract a bit range as an integer.
+            (start, end) is inclusive lower bound, exclusive upper bound.
+        """
+        return x >> (width-end) & ((2**(end-start))-1)
+
+    def index(self, p):
+        p.reverse()
+        idx = 0
+        iwidth = self.bits * self.dimension
+        for i in range(iwidth):
+            bitoff = self.bits-(i/self.dimension)-1
+            poff = self.dimension-(i%self.dimension)-1
+            b = self.bitrange(p[poff], self.bits, bitoff, bitoff+1) << i
+            idx |= b
+        return idx
+
+    def point(self, idx):
+        p = [0]*self.dimension
+        iwidth = self.bits * self.dimension
+        for i in range(iwidth):
+            b = self.bitrange(idx, iwidth, i, i+1) << (iwidth-i-1)/self.dimension
+            p[i%self.dimension] |= b
+        p.reverse()
+        return p
+
 class Tree(object):
     def __init__(self, h_in, levels=3):
         assert type(h_in) is list, 'h_in must be a list'
@@ -24,7 +68,7 @@ class Tree(object):
         self._levelBits = int(np.ceil(np.sqrt(levels)))+1
 
 
-        self._z = scurve.zorder.ZOrder(self.dim,20)
+        self._z = ZCurve(self.dim, 20)
         self._treeInds = set()
         self._treeInds.add(0)
 
@@ -374,42 +418,3 @@ if __name__ == '__main__':
     # T.refine(function,recursive=True)
     # T.plotGrid(showIt=True)
     # BREAK
-    T = Tree([np.r_[1,2,1,5,2,3,1,1],8])
-    T._refineCell([0,0,0])
-    T._refineCell([4,4,1])
-    T._refineCell([0,0,1])
-    T._refineCell([2,2,2])
-    T.plotGrid(showIt=True)
-
-    T.number()
-    print sorted(T._treeInds) == [32, 40, 48, 60, 61, 62, 63, 68, 132, 224, 232, 240, 248]
-    print len(T._hangingFacesX) == 7
-    print T.nFx == 18
-    print T.vol == 1.0
-    print T.area
-
-    T.faceDiv
-
-    plt.subplot(211)
-    plt.spy(T.faceDiv)
-    T.plotGrid(ax=plt.subplot(212), showIt=True)
-
-    print T._getNextCell([4,0,1]) is None
-    print T._getNextCell([0,4,1]) == [T._index([4,4,2]), T._index([4,6,2])]
-    print T._getNextCell([0,2,2]) == [T._index([2,2,3]), T._index([2,3,3])]
-    print T._getNextCell([4,4,2]) == T._index([6,4,2])
-    print T._getNextCell([6,4,2]) is None
-    print T._getNextCell([2,0,2]) == T._index([4,0,1])
-    print T._getNextCell([4,0,1], positive=False) == [T._index([2,0,2]), [T._index([3,2,3]), T._index([3,3,3])]]
-    print T._getNextCell([3,3,3]) == T._index([4,0,1])
-    print T._getNextCell([3,2,3]) == T._index([4,0,1])
-    print T._getNextCell([2,2,3]) == T._index([3,2,3])
-    print T._getNextCell([3,2,3], positive=False) == T._index([2,2,3])
-
-
-    print T._getNextCell([0,0,2], direction=1) == T._index([0,2,2])
-    print T._getNextCell([0,2,2], direction=1, positive=False) == T._index([0,0,2])
-    print T._getNextCell([0,2,2], direction=1) == T._index([0,4,1])
-    print T._getNextCell([0,4,1], direction=1, positive=False) ==  [T._index([0,2,2]), [T._index([2,3,3]), T._index([3,3,3])]]
-
-

tests/mesh/test_pointerMesh.py

@@ -0,0 +1,52 @@
+from SimPEG.Mesh.PointerTree import Tree
+import numpy as np
+import unittest
+
+TOL = 1e-10
+
+class TestOcTreeObjects(unittest.TestCase):
+
+    def test_counts(self):
+
+        T = Tree([8,8])
+        T._refineCell([0,0,0])
+        T._refineCell([4,4,1])
+        T._refineCell([0,0,1])
+        T._refineCell([2,2,2])
+        T.number()
+        assert sorted(T._treeInds) == [2, 34, 66, 99, 107, 115, 123, 129, 257, 386, 418, 450, 482]
+        assert len(T._hangingFacesX) == 7
+        assert T.nFx == 18
+        assert T.vol.sum() == 1.0
+
+
+    def test_connectivity(self):
+        T = Tree([8,8])
+        T._refineCell([0,0,0])
+        T._refineCell([4,4,1])
+        T._refineCell([0,0,1])
+        T._refineCell([2,2,2])
+        T.number()
+        assert T._getNextCell([4,0,1]) is None
+        assert T._getNextCell([0,4,1]) == [T._index([4,4,2]), T._index([4,6,2])]
+        assert T._getNextCell([0,2,2]) == [T._index([2,2,3]), T._index([2,3,3])]
+        assert T._getNextCell([4,4,2]) == T._index([6,4,2])
+        assert T._getNextCell([6,4,2]) is None
+        assert T._getNextCell([2,0,2]) == T._index([4,0,1])
+        assert T._getNextCell([4,0,1], positive=False) == [T._index([2,0,2]), [T._index([3,2,3]), T._index([3,3,3])]]
+        assert T._getNextCell([3,3,3]) == T._index([4,0,1])
+        assert T._getNextCell([3,2,3]) == T._index([4,0,1])
+        assert T._getNextCell([2,2,3]) == T._index([3,2,3])
+        assert T._getNextCell([3,2,3], positive=False) == T._index([2,2,3])
+
+
+        assert T._getNextCell([0,0,2], direction=1) == T._index([0,2,2])
+        assert T._getNextCell([0,2,2], direction=1, positive=False) == T._index([0,0,2])
+        assert T._getNextCell([0,2,2], direction=1) == T._index([0,4,1])
+        assert T._getNextCell([0,4,1], direction=1, positive=False) ==  [T._index([0,2,2]), [T._index([2,3,3]), T._index([3,3,3])]]
+
+
+
+
+if __name__ == '__main__':
+    unittest.main()