1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
|
"""Unit tests for the array related extension code.
"""
# Author: Prabhu Ramachandran <prabhu_r@users.sf.net>
# Copyright (c) 2005, Enthought, Inc.
# License: BSD Style.
import pytest
import unittest
import numpy
from tvtk.array_handler import ID_TYPE_CODE, set_id_type_array_py
try:
from tvtk.array_ext import set_id_type_array
except ModuleNotFoundError: # not compiled
pytest.skip("array_ext not found", allow_module_level=True)
class TestArrayExt(unittest.TestCase):
def check(self, set_id_type_array):
N = 5
a = numpy.zeros((N, 4), ID_TYPE_CODE)
a[:, 1] = 1
a[:, 2] = 2
a[:, 3] = 3
def diff_arr(x, y):
return numpy.sum(numpy.ravel(x) - numpy.ravel(y[:, 1:]))
# Test contiguous arrays.
b = numpy.zeros((N, 5), ID_TYPE_CODE)
set_id_type_array(a, b)
self.assertEqual(diff_arr(a, b), 0)
# Test non-contiguous arrays.
b = numpy.zeros((N, 3), ID_TYPE_CODE)
set_id_type_array(a[:, ::2], b)
self.assertEqual(diff_arr(a[:, ::2], b), 0)
# Test 1D array.
b = numpy.zeros(N*5, ID_TYPE_CODE)
set_id_type_array(a, b)
self.assertEqual(diff_arr(a, numpy.reshape(b, (N, 5))), 0)
# Test assertions.
d = a.astype('d')
b = numpy.zeros((N, 5), ID_TYPE_CODE)
self.assertRaises(AssertionError, set_id_type_array,
d, b)
# B should b contiguous.
b = numpy.zeros((N, 10), ID_TYPE_CODE)
self.assertRaises(AssertionError, set_id_type_array,
a, b[:, ::2])
self.assertRaises(AssertionError, set_id_type_array,
a[0], b)
# Test size check assertion.
b = numpy.zeros((N, 4), ID_TYPE_CODE)
self.assertRaises(AssertionError, set_id_type_array,
a, b)
b = numpy.zeros(N*6, ID_TYPE_CODE)
self.assertRaises(AssertionError, set_id_type_array,
a, b)
# This should work!
set_id_type_array(a, b[:N*5])
self.assertEqual(diff_arr(a, numpy.reshape(b[:N*5], (N, 5))), 0)
def test_set_id_type_array(self):
self.check(set_id_type_array)
def test_set_id_type_array_py(self):
self.check(set_id_type_array_py)
if __name__ == "__main__":
unittest.main()
|
