From fe547476c9ea0b7855df80b811b7d83fa304b03c Mon Sep 17 00:00:00 2001 From: Rowan Cockett Date: Fri, 26 Jul 2013 15:56:21 -0700 Subject: [PATCH] IndexCube now in Utils --- SimPEG/tests/test_utils.py | 33 +++++++++++----- SimPEG/utils.py | 77 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 100 insertions(+), 10 deletions(-) diff --git a/SimPEG/tests/test_utils.py b/SimPEG/tests/test_utils.py index 5d05b710..59b70147 100644 --- a/SimPEG/tests/test_utils.py +++ b/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() - diff --git a/SimPEG/utils.py b/SimPEG/utils.py index 55acf9b0..8aa7f27b 100644 --- a/SimPEG/utils.py +++ b/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]]