mirror of
https://github.com/wassname/simpeg.git
synced 2026-06-28 21:47:21 +08:00
Lars Ruthotto
54 lines
1.4 KiB
Python
54 lines
1.4 KiB
Python
import numpy as np
|
|
import unittest
|
|
import sys
|
|
sys.path.append('../')
|
|
from utils import mkvc, ndgrid
|
|
|
|
|
|
class TestSequenceFunctions(unittest.TestCase):
|
|
|
|
def setUp(self):
|
|
self.a = np.array([1, 2, 3])
|
|
self.b = np.array([1, 2])
|
|
self.c = np.array([1, 2, 3, 4])
|
|
|
|
def test_mkvc1(self):
|
|
x = mkvc(self.a)
|
|
self.assertTrue(x.shape, (3,))
|
|
|
|
def test_mkvc2(self):
|
|
x = mkvc(self.a, 2)
|
|
self.assertTrue(x.shape, (3, 1))
|
|
|
|
def test_mkvc3(self):
|
|
x = mkvc(self.a, 3)
|
|
self.assertTrue(x.shape, (3, 1, 1))
|
|
|
|
def test_ndgrid_2D(self):
|
|
XY = ndgrid([self.a, self.b])
|
|
|
|
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)
|
|
|
|
def test_ndgrid_3D(self):
|
|
XYZ = ndgrid([self.a, self.b, self.c])
|
|
|
|
X1_test = np.array([1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3, 1, 2, 3])
|
|
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(xtest and ytest and ztest)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|
|
|