IndexCube now in Utils

SimPEG/tests/test_utils.py

@@ -2,7 +2,7 @@ import numpy as np
 import unittest
 import sys
 sys.path.append('../')
-from utils import mkvc, ndgrid
+from utils import mkvc, ndgrid, indexCube
 
 
 class TestSequenceFunctions(unittest.TestCase):
@@ -30,10 +30,8 @@ class TestSequenceFunctions(unittest.TestCase):
         X1_test = np.array([1, 2, 3, 1, 2, 3])
         X2_test = np.array([1, 1, 1, 2, 2, 2])
 
-        xtest = np.all(XY[:, 0] == X1_test)
-        ytest = np.all(XY[:, 1] == X2_test)
-
-        self.assertTrue(xtest and ytest)
+        self.assertTrue(np.all(XY[:, 0] == X1_test))
+        self.assertTrue(np.all(XY[:, 1] == X2_test))
 
     def test_ndgrid_3D(self):
         XYZ = ndgrid([self.a, self.b, self.c])
@@ -42,12 +40,27 @@ class TestSequenceFunctions(unittest.TestCase):
         X2_test = np.array([1, 1, 1, 2, 2, 2, 1, 1, 1, 2, 2, 2, 1, 1, 1, 2, 2, 2, 1, 1, 1, 2, 2, 2])
         X3_test = np.array([1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4])
 
-        xtest = np.all(XYZ[:, 0] == X1_test)
-        ytest = np.all(XYZ[:, 1] == X2_test)
-        ztest = np.all(XYZ[:, 2] == X3_test)
+        self.assertTrue(np.all(XYZ[:, 0] == X1_test))
+        self.assertTrue(np.all(XYZ[:, 1] == X2_test))
+        self.assertTrue(np.all(XYZ[:, 2] == X3_test))
 
-        self.assertTrue(xtest and ytest and ztest)
+    def test_indexCube_2D(self):
+        nN = np.array([3, 3])
+        self.assertTrue(np.all(indexCube('A', nN) == np.array([0, 1, 3, 4])))
+        self.assertTrue(np.all(indexCube('B', nN) == np.array([3, 4, 6, 7])))
+        self.assertTrue(np.all(indexCube('C', nN) == np.array([4, 5, 7, 8])))
+        self.assertTrue(np.all(indexCube('D', nN) == np.array([1, 2, 4, 5])))
+
+    def test_indexCube_3D(self):
+        nN = np.array([3, 3, 3])
+        self.assertTrue(np.all(indexCube('A', nN) == np.array([0, 1, 3, 4, 9, 10, 12, 13])))
+        self.assertTrue(np.all(indexCube('B', nN) == np.array([3, 4, 6, 7, 12, 13, 15, 16])))
+        self.assertTrue(np.all(indexCube('C', nN) == np.array([4, 5, 7, 8, 13, 14, 16, 17])))
+        self.assertTrue(np.all(indexCube('D', nN) == np.array([1, 2, 4, 5, 10, 11, 13, 14])))
+        self.assertTrue(np.all(indexCube('E', nN) == np.array([9, 10, 12, 13, 18, 19, 21, 22])))
+        self.assertTrue(np.all(indexCube('F', nN) == np.array([12, 13, 15, 16, 21, 22, 24, 25])))
+        self.assertTrue(np.all(indexCube('G', nN) == np.array([13, 14, 16, 17, 22, 23, 25, 26])))
+        self.assertTrue(np.all(indexCube('H', nN) == np.array([10, 11, 13, 14, 19, 20, 22, 23])))
 
 if __name__ == '__main__':
     unittest.main()
-

SimPEG/utils.py

@@ -124,6 +124,83 @@ def volTetra(xyz, A, B, C, D):
     return V/6
 
 
+def indexCube(nodes, nN):
+    """
+   Returns the index of nodes on the mesh.
+
+
+    Input:
+       nodes  - string of which nodes to return. e.g. 'ABCD'
+       nN     - size of the nodal grid
+
+    Output:
+       index  - index in the order asked e.g. 'ABCD' --> (A,B,C,D)
+
+      TWO DIMENSIONS:
+
+      node(i,j)          node(i,j+1)
+           A -------------- B
+           |                |
+           |    cell(i,j)   |
+           |        I       |
+           |                |
+          D -------------- C
+      node(i+1,j)        node(i+1,j+1)
+
+
+      THREE DIMENSIONS:
+
+            node(i,j,k+1)       node(i,j+1,k+1)
+                E --------------- F
+               /|               / |
+              / |              /  |
+             /  |             /   |
+      node(i,j,k)         node(i,j+1,k)
+           A -------------- B     |
+           |    H ----------|---- G
+           |   /cell(i,j)   |   /
+           |  /     I       |  /
+           | /              | /
+           D -------------- C
+      node(i+1,j,k)      node(i+1,j+1,k)
+
+
+         @author Rowan Cockett
+
+         Last modified on: 2013/03/07
+    """
+
+    assert type(nodes) == str, "Nodes must be a str variable: e.g. 'ABCD'"
+    assert type(nN) == np.ndarray, "Number of nodes must be an ndarray"
+    nodes = nodes.upper()
+    # Make sure that we choose from the possible nodes.
+    possibleNodes = 'ABCD' if nN.size == 2 else 'ABCDEFGH'
+    for node in nodes:
+        assert node in possibleNodes, "Nodes must be chosen from: '%s'" % possibleNodes
+    dim = nN.size
+    nC = nN - 1
+
+    if dim == 2:
+        ij = ndgrid(np.arange(nC[0]), np.arange(nC[1]))
+        i, j = ij[:, 0], ij[:, 1]
+    elif dim == 3:
+        ijk = ndgrid(np.arange(nC[0]), np.arange(nC[1]), np.arange(nC[2]))
+        i, j, k = ijk[:, 0], ijk[:, 1], ijk[:, 2]
+    else:
+        raise Exception('Only 2 and 3 dimensions supported.')
+
+    nodeMap = {'A': [0, 0, 0], 'B': [0, 1, 0], 'C': [1, 1, 0], 'D': [1, 0, 0],
+               'E': [0, 0, 1], 'F': [0, 1, 1], 'G': [1, 1, 1], 'H': [1, 0, 1]}
+    out = ()
+    for node in nodes:
+        shift = nodeMap[node]
+        if dim == 2:
+            out += (sub2ind(nN, np.c_[i+shift[0], j+shift[1]]).flatten(), )
+        elif dim == 3:
+            out += (sub2ind(nN, np.c_[i+shift[0], j+shift[1], k+shift[2]]).flatten(), )
+
+    return out
+
 def getSubArray(A, ind):
     """subArray"""
     return A[ind[0], :, :][:, ind[1], :][:, :, ind[2]]