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 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211
|
"""Tests for the ArraySource data source for MayaVi.
"""
# Author: Prabhu Ramachandran <prabhu_r@users.sf.net>
# Copyright (c) 2005, Enthought, Inc.
# License: BSD Style.
# Standard library imports.
from os.path import abspath
from io import BytesIO
import copy
import numpy
# Enthought library imports.
from traits.api import TraitError
# Local imports.
from common import TestCase
class TestArraySource(TestCase):
def check_input_validation(self, obj):
"""Tests if only the correct forms of input arrays are supported."""
# These should work.
obj.scalar_data = numpy.zeros((2,2), 'd')
obj.scalar_data = numpy.zeros((2,2,2), 'd')
obj.scalar_data = None
obj.vector_data = numpy.zeros((2,2,3), 'd')
obj.vector_data = numpy.zeros((2,2,2,3), 'd')
obj.vector_data = None
# These should not.
self.assertRaises(TraitError, setattr, obj, 'scalar_data', [1,2,3])
self.assertRaises(TraitError, setattr, obj, 'scalar_data',
numpy.zeros((2,2,2,3), 'd'))
obj.scalar_data = None
self.assertRaises(TraitError, setattr, obj, 'vector_data', [[1,2,3]])
self.assertRaises(TraitError, setattr, obj, 'vector_data',
numpy.zeros((2,2,2,1), 'd'))
obj.vector_data = None
obj.scalar_data = numpy.zeros((2,2), 'd')
self.assertRaises(TraitError, setattr, obj, 'vector_data',
numpy.zeros((4,4,3), 'd'))
obj.vector_data = numpy.zeros((2,2,3), 'd')
self.assertRaises(TraitError, setattr, obj, 'scalar_data',
numpy.zeros((4,3), 'i'))
self.assertRaises(TraitError, setattr, obj, 'scalar_data',
numpy.zeros((2,2,2), 'i'))
obj.scalar_data = numpy.zeros((2,2), 'f')
# Clean up the object so it can be used for further testing.
obj.scalar_data = obj.vector_data = None
def make_2d_data(self):
s = numpy.array([[0, 1],[2, 3]], 'd')
v = numpy.array([[[1,1,1], [1,0,0]],[[0,1,0], [0,0,1]]], 'd')
tps = numpy.transpose
s, v = tps(s), tps(v, (1, 0, 2))
return s, v
def make_3d_data(self):
s = numpy.array([[[0, 1],[2, 3]],
[[4, 5],[6, 7]]], 'd')
v = numpy.array([[[[0,0,0],
[1,0,0]],
[[0,1,0],
[1,1,0]]],
[[[0,0,1],
[1,0,1]],
[[0,1,1],
[1,1,1]]]], 'd')
tps = numpy.transpose
s, v = tps(s), tps(v, (2, 1, 0, 3))
return s, v
def check(self):
script = self.script
s = script.engine.current_scene
d1, d2 = s.children
s1, v1 = d1.children[0].children[1:]
expect = list(self.make_2d_data())
tps = numpy.transpose
expect[0] = tps(expect[0])
expect[1] = tps(expect[1], (1, 0, 2))
sc1 = s1.actor.mapper.input.point_data.scalars.to_array()
assert numpy.allclose(sc1.flatten(), expect[0].flatten())
vec1 = s1.actor.mapper.input.point_data.vectors.to_array()
assert numpy.allclose(vec1.flatten(), expect[1].flatten())
s2, v2 = d2.children[0].children[1:]
expect = list(self.make_3d_data())
tps = numpy.transpose
expect[0] = tps(expect[0])
expect[1] = tps(expect[1], (2, 1, 0, 3))
sc2 = s2.actor.mapper.input.point_data.scalars.to_array()
assert numpy.allclose(sc2.flatten(), expect[0].flatten())
vec2 = s2.actor.mapper.input.point_data.vectors.to_array()
assert numpy.allclose(vec2.flatten(), expect[1].flatten())
def test(self):
self.main()
def do(self):
############################################################
# Imports.
script = self.script
from mayavi.sources.array_source import ArraySource
from mayavi.modules.outline import Outline
from mayavi.modules.surface import Surface
from mayavi.modules.vectors import Vectors
############################################################
# Create a new scene and set up the visualization.
s = self.new_scene()
d = ArraySource()
self.check_input_validation(d)
sc, vec = self.make_2d_data()
d.origin = (-1, -1, 0)
d.scalar_data = sc
d.vector_data = vec
script.add_source(d)
# Create an outline for the data.
o = Outline()
script.add_module(o)
# View the data.
s = Surface()
script.add_module(s)
v = Vectors()
script.add_module(v)
# Add a 3D data source
d = ArraySource()
sc, vec = self.make_3d_data()
d.scalar_data = sc
d.vector_data = vec
script.add_source(d)
# Create an outline for the data.
o = Outline()
script.add_module(o)
# View a slice.
s = Surface()
script.add_module(s)
v = Vectors()
script.add_module(v)
# Set the scene to a suitable view.
s.scene.z_plus_view()
c = s.scene.camera
c.azimuth(-30)
c.elevation(30)
self.check()
############################################################
# Test if saving a visualization and restoring it works.
bg = s.scene.background
# Save visualization.
f = BytesIO()
f.name = abspath('test.mv2') # We simulate a file.
script.save_visualization(f)
f.seek(0) # So we can read this saved data.
# Remove existing scene.
engine = script.engine
engine.close_scene(s)
# Load visualization
script.load_visualization(f)
s = engine.current_scene
# Set the scene to a suitable view.
s.scene.z_plus_view()
c = s.scene.camera
c.azimuth(-30)
c.elevation(30)
s.scene.background = bg
self.check()
############################################################
# Test if the MayaVi2 visualization can be deepcopied.
# Pop the source object.
sources = s.children
s.children = []
# Add it back to see if that works without error.
s.children.extend(sources)
s.scene.reset_zoom()
self.check()
# Now deepcopy the source and replace the existing one with
# the copy. This basically simulates cutting/copying the
# object from the UI via the right-click menu on the tree
# view, and pasting the copy back.
sources1 = copy.deepcopy(sources)
s.children[:] = sources
s.scene.reset_zoom()
self.check()
# If we have come this far, we are golden!
if __name__ == "__main__":
t = TestArraySource()
t.test()
|