test_vtkio.py

# -*- coding: utf-8 -*-
import unittest
import numpy
import os

import spam.helpers
import spam.datasets


class testAll(unittest.TestCase):

    def tearDown(self):
        try:
            # pass
            os.remove("./spam.vtk")
            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")
        except OSError:
            pass

    # 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):
        gridFlat = numpy.zeros((len(grid[0].ravel()), 3))
        gridFlat[:, 0] = grid[0].ravel()
        gridFlat[:, 1] = grid[1].ravel()
        gridFlat[:, 2] = grid[2].ravel()

        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')

    def test_structured(self):
        # 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})

        # write a tensor
        tensor = numpy.random.rand(4, 5, 6, 3, 3)
        spam.helpers.writeStructuredVTK(pointData={'c': tensor})

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

        # 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

        # read all

        fields = spam.helpers.readStructuredVTK('./spam.vtk')
        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)

        # specific cases
        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')

    def test_unstructured(self):
        mesh = spam.datasets.loadUnstructuredMesh()
        spam.helpers.writeUnstructuredVTK(mesh['points'], mesh['connectivity'],
                                          pointData=mesh['pointData'],
                                          cellData=mesh['cellData'])
        readPoints, readConnectivity, readPointData, readCellData = spam.helpers.readUnstructuredVTK('./spam.vtk')
        # check points
        self.assertAlmostEqual(numpy.sum(readPoints - mesh['points']), 0, places=5)
        # check connectivity
        self.assertAlmostEqual(numpy.sum(readConnectivity - mesh['connectivity']), 0, places=5)
        # check point data
        for a in mesh['pointData']:
            self.assertAlmostEqual(numpy.sum(readPointData[a] - mesh['pointData'][a]), 0, places=5)
        # check cell data
        for a in mesh['cellData']:
            self.assertAlmostEqual(numpy.sum(readCellData[a] - mesh['cellData'][a]), 0, places=5)
        # run swapAxes
        # spam.helpers.readUnstructuredVTK('spam.vtk',zFirst=True)

    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)

    def test_TIFFtoVTK(self):
        # Make it smaller because this takes a long time
        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]

        import tifffile
        tifffile.imsave("./snow-1.tif", im1)
        tifffile.imsave("./snow-2.tif", im2)
        tifffile.imsave("./snow-3.tif", im3)

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



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