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
82
83
84
85
86
87
|
import unittest
import numpy as np
from tvtk.api import tvtk
from mayavi.core.utils import DataSetHelper
from vtk.numpy_interface import dataset_adapter as dsa
class TestDataSetHelper(unittest.TestCase):
def _make_data(self, factor=1.0):
id = tvtk.ImageData(dimensions=(2, 2, 1), origin=(0, 0, 0),
spacing=(1, 1, 1))
id.point_data.scalars = np.arange(4, dtype=float)*factor
id.point_data.scalars.name = 'ps'
id.point_data.vectors = np.ones((4, 3), dtype=float)*factor
id.point_data.vectors.name = 'pv'
id.cell_data.scalars = np.arange(4, dtype=float)*0.5*factor
id.cell_data.scalars.name = 'cs'
id.cell_data.vectors = np.ones((4, 3), dtype=float)*0.5*factor
id.cell_data.vectors.name = 'cv'
return id
def test_get_range_works_for_simple_datasets(self):
# Given
id = self._make_data()
# When
dsh = DataSetHelper(id)
# Then
name, rng = dsh.get_range('scalars', 'point')
self.assertEqual(name, 'ps')
self.assertEqual(rng, [0.0, 3.0])
name, rng = dsh.get_range('scalars', 'cell')
self.assertEqual(name, 'cs')
self.assertEqual(rng, [0.0, 1.5])
name, rng = dsh.get_range('vectors', 'point')
self.assertEqual(name, 'pv')
self.assertEqual(rng, [0.0, np.sqrt(3.0)])
name, rng = dsh.get_range('vectors', 'cell')
self.assertEqual(name, 'cv')
self.assertEqual(rng, [0.0, np.sqrt(3.0)*0.5])
def test_get_range_works_for_multiblock_datasets(self):
# Given
id1 = self._make_data()
id2 = self._make_data(factor=2.0)
mb = tvtk.MultiBlockDataSet()
mb.set_block(0, id1)
mb.set_block(1, id2)
# When
dsh = DataSetHelper(mb)
# Then
name, rng = dsh.get_range('scalars', 'point')
self.assertEqual(name, 'ps')
self.assertEqual(rng, [0.0, 6.0])
name, rng = dsh.get_range('scalars', 'cell')
self.assertEqual(name, 'cs')
self.assertEqual(rng, [0.0, 3.0])
name, rng = dsh.get_range('vectors', 'point')
self.assertEqual(name, 'pv')
self.assertEqual(rng, [0.0, np.sqrt(3.0)*2.0])
name, rng = dsh.get_range('vectors', 'cell')
self.assertEqual(name, 'cv')
self.assertEqual(rng, [0.0, np.sqrt(3.0)])
def test_get_range_works_for_VTKNoneArray(self):
id_ = tvtk.ImageData(dimensions=(2, 2, 1), origin=(0, 0, 0),
spacing=(1, 1, 1))
id_.point_data.scalars = np.arange(4, dtype=float)
dsh = DataSetHelper(id_)
self.assertFalse(dsh._composite)
self.assertEqual(dsh.get_range(), (None, [0., 1.]))
# XXX there is some wackiness here, no idea why this changes!
self.assertEqual(dsh.get_range(), ('point_scalars', [0., 3.]))
if __name__ == '__main__':
unittest.main()
|
