Added numpy_utils.py

pywafo/src/wafo/test/test_numpy_utils.py

@@ -0,0 +1,253 @@
+from numpy.testing import (
+    TestCase, assert_, assert_array_equal, assert_raises,)
+#     run_module_suite,
+#      assert_allclose, assert_array_max_ulp, assert_warns,
+#     assert_equal, assert_array_almost_equal, assert_almost_equal,
+#     )
+
+import unittest as local_unittest
+from wafo.numpy_utils import (rotation_matrix, rotate_2d, spaceline,
+                              args_flat, sub2index, index2sub, piecewise)
+import numpy as np
+
+
+class TestPiecewise(TestCase):
+    def test_condition_is_single_bool_list(self):
+        assert_raises(ValueError, piecewise, [0, 0], [True, False], [1])
+
+    def test_condition_is_list_of_single_bool_list(self):
+        x = piecewise([0, 0], [[True, False]], [1])
+        assert_array_equal(x, [1, 0])
+
+    def test_conditions_is_list_of_single_bool_array(self):
+        x = piecewise([0, 0], [np.array([True, False])], [1])
+        assert_array_equal(x, [1, 0])
+
+    def test_condition_is_single_int_array(self):
+        assert_raises(ValueError, piecewise, [0, 0], np.array([1, 0]), [1])
+
+    def test_condition_is_list_of_single_int_array(self):
+        x = piecewise([0, 0], [np.array([1, 0])], [1])
+        assert_array_equal(x, [1, 0])
+
+    def test_simple(self):
+        x = piecewise([0, 0], [[False, True]], [lambda x:-1])
+        assert_array_equal(x, [0, -1])
+
+        x = piecewise([1, 2], [[True, False], [False, True]], [3, 4])
+        assert_array_equal(x, [3, 4])
+
+    def test_default(self):
+        # No value specified for x[1], should be 0
+        x = piecewise([1, 2], [[True, False]], [2])
+        assert_array_equal(x, [2, 0])
+
+        # Should set x[1] to 3
+        x = piecewise([1, 2], [[True, False]], [2, 3])
+        assert_array_equal(x, [2, 3])
+
+    def test_0d(self):
+        x = np.array(3)
+        y = piecewise(x, [x > 3], [4, 0])
+        assert_(y.ndim == 0)
+        assert_(y == 0)
+
+        x = 5
+        y = piecewise(x, [[True], [False]], [1, 0])
+        assert_(y.ndim == 0)
+        assert_(y == 1)
+
+    def test_abs_function(self):
+        x = np.linspace(-2.5, 2.5, 6)
+        vals = piecewise((x,), [x < 0, x >= 0], [lambda x: -x, lambda x: x])
+        assert_array_equal(vals,
+                           [2.5,  1.5,  0.5,  0.5,  1.5,  2.5])
+
+    def test_abs_function_with_scalar(self):
+        x = np.array(-2.5)
+        vals = piecewise((x,), [x < 0, x >= 0], [lambda x: -x, lambda x: x])
+        assert_(vals == 2.5)
+
+    def test_otherwise_condition(self):
+        x = np.linspace(-2.5, 2.5, 6)
+        vals = piecewise((x,), [x < 0, ], [lambda x: -x, lambda x: x])
+        assert_array_equal(vals, [2.5,  1.5,  0.5,  0.5,  1.5,  2.5])
+
+    def test_passing_further_args_to_fun(self):
+        def fun0(x, y, scale=1.):
+            return -x*y/scale
+
+        def fun1(x, y, scale=1.):
+            return x*y/scale
+        x = np.linspace(-2.5, 2.5, 6)
+        vals = piecewise((x,), [x < 0, ], [fun0, fun1], args=(2.,), scale=2.)
+        assert_array_equal(vals, [2.5,  1.5,  0.5,  0.5,  1.5,  2.5])
+
+    def test_step_function(self):
+        x = np.linspace(-2.5, 2.5, 6)
+        vals = piecewise(x, [x < 0, x >= 0], [-1, 1])
+        assert_array_equal(vals, [-1., -1., -1.,  1.,  1.,  1.])
+
+    def test_step_function_with_scalar(self):
+        x = 1
+        vals = piecewise(x, [x < 0, x >= 0], [-1, 1])
+        assert_(vals == 1)
+
+    def test_function_with_two_args(self):
+        x = np.linspace(-2, 2, 5)
+        X, Y = np.meshgrid(x, x)
+        vals = piecewise(
+            (X, Y), [X * Y < 0, ], [lambda x, y: -x * y, lambda x, y: x * y])
+        assert_array_equal(vals, [[4.,  2., -0.,  2.,  4.],
+                                  [2.,  1., -0.,  1.,  2.],
+                                  [-0., -0.,  0.,  0.,  0.],
+                                  [2.,  1.,  0.,  1.,  2.],
+                                  [4.,  2.,  0.,  2.,  4.]])
+
+    def test_fill_value_and_function_with_two_args(self):
+        x = np.linspace(-2, 2, 5)
+        X, Y = np.meshgrid(x, x)
+        vals = piecewise((X, Y), [X * Y < -0.5, X * Y > 0.5],
+                         [lambda x, y: -x * y, lambda x, y: x * y],
+                         fill_value=np.nan)
+        nan = np.nan
+        assert_array_equal(vals, [[4.,   2.,  nan,   2.,   4.],
+                                  [2.,   1.,  nan,   1.,   2.],
+                                  [nan,  nan,  nan,  nan,  nan],
+                                  [2.,   1.,  nan,   1.,   2.],
+                                  [4.,   2.,  nan,   2.,   4.]])
+
+    def test_fill_value2_and_function_with_two_args(self):
+        x = np.linspace(-2, 2, 5)
+        X, Y = np.meshgrid(x, x)
+        vals = piecewise((X, Y), [X * Y < -0.5, X * Y > 0.5],
+                         [lambda x, y: -x * y, lambda x, y: x * y, np.nan])
+        nan = np.nan
+        assert_array_equal(vals, [[4.,   2.,  nan,   2.,   4.],
+                                  [2.,   1.,  nan,   1.,   2.],
+                                  [nan,  nan,  nan,  nan,  nan],
+                                  [2.,   1.,  nan,   1.,   2.],
+                                  [4.,   2.,  nan,   2.,   4.]])
+
+
+class TestRotationMatrix(TestCase):
+
+    def test_h0_p0_r0(self):
+        vals = rotation_matrix(heading=0, pitch=0, roll=0).tolist()
+        truevals = [[1., 0., 0.],
+                    [0., 1., 0.],
+                    [0., 0., 1.]]
+        self.assertListEqual(vals, truevals)
+
+    def test_h180_p0_r0(self):
+        vals = rotation_matrix(heading=180, pitch=0, roll=0).tolist()
+        truevals = [[-1.0, -1.2246467991473532e-16, 0.0],
+                    [1.2246467991473532e-16, -1.0, 0.0],
+                    [-0.0, 0.0, 1.0]]
+        self.assertListEqual(vals, truevals)
+
+    def test_h0_p180_r0(self):
+        vals = rotation_matrix(heading=0, pitch=180, roll=0).tolist()
+        truevals = [[-1.0, 0.0, 1.2246467991473532e-16],
+                    [-0.0, 1.0, 0.0],
+                    [-1.2246467991473532e-16, -0.0, -1.0]]
+        self.assertListEqual(vals, truevals)
+
+    def test_h0_p0_r180(self):
+        vals = rotation_matrix(heading=0, pitch=180, roll=0).tolist()
+        truevals = [[-1.0, 0.0, 1.2246467991473532e-16],
+                    [-0.0, 1.0, 0.0],
+                    [-1.2246467991473532e-16, -0.0, -1.0]]
+        self.assertListEqual(vals, truevals)
+
+
+class TestRotate2d(TestCase):
+
+    def test_rotate_0deg(self):
+        vals = list(rotate_2d(x=1, y=0, angle_deg=0))
+        truevals = [1.0, 0.0]
+        self.assertListEqual(vals, truevals)
+
+    def test_rotate_90deg(self):
+        vals = list(rotate_2d(x=1, y=0, angle_deg=90))
+        truevals = [6.123233995736766e-17, 1.0]
+        self.assertListEqual(vals, truevals)
+
+    def test_rotate_180deg(self):
+        vals = list(rotate_2d(x=1, y=0, angle_deg=180))
+        truevals = [-1.0, 1.2246467991473532e-16]
+        self.assertListEqual(vals, truevals)
+
+    def test_rotate_360deg(self):
+        vals = list(rotate_2d(x=1, y=0, angle_deg=360))
+        truevals = [1.0, -2.4492935982947064e-16]
+        self.assertListEqual(vals, truevals)
+
+
+class TestSpaceLine(TestCase):
+
+    def test_space_line(self):
+        vals = spaceline((2, 0, 0), (3, 0, 0), num=5).tolist()
+        truevals = [[2., 0., 0.],
+                    [2.25, 0., 0.],
+                    [2.5, 0., 0.],
+                    [2.75, 0., 0.],
+                    [3., 0., 0.]]
+        self.assertListEqual(vals, truevals)
+
+
+class TestArgsFlat(TestCase):
+
+    def test_1_vector_and_2_scalar_args(self):
+        x = [1, 2, 3]
+        pos, c_shape = args_flat(x, 2, 3)
+        truepos = [[1, 2, 3],
+                   [2, 2, 3],
+                   [3, 2, 3]]
+        truec_shape = [3, ]
+        self.assertListEqual(pos.tolist(), truepos)
+        self.assertListEqual(list(c_shape), truec_shape)
+
+    def test_1_vector_args(self):
+        pos1, c_shape1 = args_flat([1, 2, 3])
+        truepos1 = [[1, 2, 3]]
+        truec_shape1 = None
+        self.assertListEqual(pos1.tolist(), truepos1)
+        self.assertIs(c_shape1, truec_shape1)
+
+    def test_3_scalar_args(self):
+        pos1, c_shape1 = args_flat(1, 2, 3)
+        truepos1 = [[1, 2, 3]]
+        truec_shape1 = []
+        self.assertListEqual(pos1.tolist(), truepos1)
+        self.assertListEqual(list(c_shape1), truec_shape1)
+
+    def test_3_scalar_args_version2(self):
+        pos1, c_shape1 = args_flat([1], 2, 3)
+        truepos1 = [[1, 2, 3]]
+        truec_shape1 = [1, ]
+        self.assertListEqual(pos1.tolist(), truepos1)
+        self.assertListEqual(list(c_shape1), truec_shape1)
+
+
+class TestSub2index2Sub(TestCase):
+
+    def test_sub2index_and_index2sub(self):
+        shape = (3, 3, 4)
+        a = np.arange(np.prod(shape)).reshape(shape)
+        trueval = a[1, 2, 3]
+        order = 'C'
+        i = sub2index(shape, 1, 2, 3, order=order)
+        self.assertEquals(i, 23)
+
+        val = a.ravel(order)[i]
+        self.assertEquals(val, trueval)
+
+        sub = index2sub(shape, i, order=order)
+        for j, true_sub_j in enumerate([[1], [2], [3]]):
+            self.assertEquals(sub[j].tolist(), true_sub_j)
+
+
+if __name__ == '__main__':
+    runner = local_unittest.TextTestRunner()  # get_test_runner()
+    local_unittest.main(testRunner=runner)