test_vtkio.py 6.01 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
# -*- coding: utf-8 -*-
import unittest
import numpy
import os

import spam.helpers
import spam.datasets


class testAll(unittest.TestCase):

    def tearDown(self):
        try:
Emmanuel Roubin's avatar
Emmanuel Roubin committed
14
            # pass
15
            os.remove("./spam.vtk")
16
17
18
19
20
21
            os.remove("./snow-1.tif")
            os.remove("./snow-2.tif")
            os.remove("./snow-3.tif")
            os.remove("./snow-1.vtk")
            os.remove("./snow-1.h5")
            os.remove("./snow-1.xmf")
22
23
24
        except OSError:
            pass

25
26
27
    # This won't be run because it doesn't start with test_
    # However it is a good example of getting the same thing with both structured and unstructured VTKs
    def visualTest(self):
Emmanuel Roubin's avatar
Emmanuel Roubin committed
28
29
30
31
        gridFlat = numpy.zeros((len(grid[0].ravel()), 3))
        gridFlat[:, 0] = grid[0].ravel()
        gridFlat[:, 1] = grid[1].ravel()
        gridFlat[:, 2] = grid[2].ravel()
32

Emmanuel Roubin's avatar
Emmanuel Roubin committed
33
34
        spam.helpers.writeUnstructuredVTK(gridFlat, connectivity=scipy.spatial.Delaunay(gridFlat).simplices, pointData={'coords-xyz': gridFlat}, fileName='unstructured.vtk')
        spam.helpers.writeStructuredVTK(origin=grid[:, 0, 0, 0], pointData={'coords-xyz': gridFlat.reshape(grid.shape[1], grid.shape[2], grid.shape[3], 3)}, fileName='structured.vtk')
35

Emmanuel Roubin's avatar
Emmanuel Roubin committed
36
    def test_structured(self):
Emmanuel Roubin's avatar
Emmanuel Roubin committed
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
        # structured mesh

        # empty
        spam.helpers.writeStructuredVTK()

        # wrong cell dimensions
        cellScalars = numpy.random.rand(4, 5, 6)
        cellVectors = numpy.random.rand(3, 5, 6, 3)
        spam.helpers.writeStructuredVTK(cellData={'a': cellScalars, 'b': cellVectors})

        # wrong point dimensions
        pointScalars = numpy.random.rand(5, 6, 7)
        pointVectors = numpy.random.rand(4, 6, 7, 3)
        spam.helpers.writeStructuredVTK(pointData={'a': pointScalars, 'b': pointVectors})

        # wrong point VS cell dimensions
        pointScalars = numpy.random.rand(5, 6, 7)
        cellScalars = numpy.random.rand(3, 5, 6)
        spam.helpers.writeStructuredVTK(pointData={'a': pointScalars}, cellData={'b': cellScalars})

        # write 1 point and 1 cell
        pointScalars = numpy.random.rand(5, 6, 7)
        cellScalars = numpy.random.rand(4, 5, 6)
        spam.helpers.writeStructuredVTK(pointData={'a': pointScalars}, cellData={'b': cellScalars})

Emmanuel Roubin's avatar
Emmanuel Roubin committed
62
63
64
65
        # write a tensor
        tensor = numpy.random.rand(4, 5, 6, 3, 3)
        spam.helpers.writeStructuredVTK(pointData={'c': tensor})

66
67
68
69
        # write field with wrong format
        dummy = numpy.random.rand(4, 5, 6, 3, 2)
        spam.helpers.writeStructuredVTK(pointData={'d': dummy})

Emmanuel Roubin's avatar
Emmanuel Roubin committed
70
71
72
73
74
75
76
77
78
        # write 2 point and 2 cell
        cellScalars = numpy.random.rand(4, 5, 6)
        cellVectors = numpy.random.rand(4, 5, 6, 3)
        pointScalars = numpy.random.rand(5, 6, 7)
        pointVectors = numpy.random.rand(5, 6, 7, 3)
        spam.helpers.writeStructuredVTK(pointData={'a': pointScalars, 'aaa': pointVectors},
                                        cellData={'b': cellScalars, 'bbb': cellVectors})

        # read structured VTK
Emmanuel Roubin's avatar
Emmanuel Roubin committed
79

Emmanuel Roubin's avatar
Emmanuel Roubin committed
80
        # read all
Emmanuel Roubin's avatar
Emmanuel Roubin committed
81

82
        fields = spam.helpers.readStructuredVTK('./spam.vtk')
Emmanuel Roubin's avatar
Emmanuel Roubin committed
83
84
85
86
        self.assertAlmostEqual(sum(pointScalars.ravel() - fields['a']), 0, places=8)
        self.assertAlmostEqual(sum(cellScalars.ravel() - fields['b']), 0, places=8)
        self.assertAlmostEqual(sum(pointVectors.ravel() - fields['aaa'].ravel()), 0, places=8)
        self.assertAlmostEqual(sum(cellVectors.ravel() - fields['bbb'].ravel()), 0, places=8)
Emmanuel Roubin's avatar
Emmanuel Roubin committed
87

Emmanuel Roubin's avatar
Emmanuel Roubin committed
88
        # specific cases
89
90
91
92
        spam.helpers.readStructuredVTK('./spam.vtk', fieldAsked='a')
        spam.helpers.readStructuredVTK('./spam.vtk', fieldAsked=['a', 'bbb', 'b', 'aaa'])
        spam.helpers.readStructuredVTK('./spam.vtk', fieldAsked='spam')
        spam.helpers.readStructuredVTK('./spam.vtk', fieldAsked='aaa')
Emmanuel Roubin's avatar
Emmanuel Roubin committed
93
94
95
96
97

    def test_unstructured(self):
        mesh = spam.datasets.loadUnstructuredMesh()
        spam.helpers.writeUnstructuredVTK(mesh['points'], mesh['connectivity'],
                                          pointData=mesh['pointData'],
98
                                          cellData=mesh['cellData'])
99
        readPoints, readConnectivity, readPointData, readCellData = spam.helpers.readUnstructuredVTK('./spam.vtk')
Emmanuel Roubin's avatar
Emmanuel Roubin committed
100
        # check points
Emmanuel Roubin's avatar
Emmanuel Roubin committed
101
        self.assertAlmostEqual(numpy.sum(readPoints - mesh['points']), 0, places=5)
Emmanuel Roubin's avatar
Emmanuel Roubin committed
102
        # check connectivity
Emmanuel Roubin's avatar
Emmanuel Roubin committed
103
        self.assertAlmostEqual(numpy.sum(readConnectivity - mesh['connectivity']), 0, places=5)
Emmanuel Roubin's avatar
Emmanuel Roubin committed
104
105
        # check point data
        for a in mesh['pointData']:
Emmanuel Roubin's avatar
Emmanuel Roubin committed
106
            self.assertAlmostEqual(numpy.sum(readPointData[a] - mesh['pointData'][a]), 0, places=5)
Emmanuel Roubin's avatar
Emmanuel Roubin committed
107
108
        # check cell data
        for a in mesh['cellData']:
Emmanuel Roubin's avatar
Emmanuel Roubin committed
109
            self.assertAlmostEqual(numpy.sum(readCellData[a] - mesh['cellData'][a]), 0, places=5)
Emmanuel Roubin's avatar
Emmanuel Roubin committed
110
        # run swapAxes
Emmanuel Roubin's avatar
Emmanuel Roubin committed
111
        # spam.helpers.readUnstructuredVTK('spam.vtk',zFirst=True)
112

Emmanuel Roubin's avatar
Emmanuel Roubin committed
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
    def test_glyphs(self):

        # write glyphs
        coordinates = numpy.random.rand(7, 3)

        # good 1 field
        pointData = {'s': numpy.random.rand(7)}
        spam.helpers.writeGlyphsVTK(coordinates, pointData)

        # good 2 fields
        pointData = {'s': numpy.random.rand(7), 'v': numpy.random.rand(7, 3)}
        spam.helpers.writeGlyphsVTK(coordinates, pointData)

        # empty
        pointData = {}
        spam.helpers.writeGlyphsVTK(coordinates, pointData)

        # wrong dimensions
        pointData = {'s': numpy.random.rand(6)}
        spam.helpers.writeGlyphsVTK(coordinates, pointData)

134
    def test_TIFFtoVTK(self):
135
        # Make it smaller because this takes a long time
136
137
138
139
        im1 = spam.datasets.loadSnow()[0:20, 0:20, 0:20]
        im2 = spam.datasets.loadSnow()[0:20, 0:20, 2:22]
        im3 = spam.datasets.loadSnow()[0:20, 0:20, 4:24]

140
        import tifffile
141
142
143
144
145
        tifffile.imsave("./snow-1.tif", im1)
        tifffile.imsave("./snow-2.tif", im2)
        tifffile.imsave("./snow-3.tif", im3)

        spam.helpers.TIFFtoVTK("./snow-1.tif")
146
        # VTK = spam.helpers.readStructuredVTK("tmpsnow.vtk")  # can't read binary
147
        spam.helpers.TIFFtoVTK(["snow-1.tif", "snow-2.tif", "snow-3.tif"])
148

Emmanuel Roubin's avatar
Emmanuel Roubin committed
149

150
151
152

if __name__ == '__main__':
    unittest.main()