diff --git a/tests/test_image_data.py b/tests/test_image_data.py new file mode 100644 index 0000000..073f20f --- /dev/null +++ b/tests/test_image_data.py @@ -0,0 +1,25 @@ +from uvw import ImageData, DataArray +import numpy as np + +def test_image_data(): + N = 4 + + x = np.linspace(0, 1, N*2) + y = np.linspace(1, 3, N+2) + z = np.linspace(0, 2, N) + + out_name = 'test_image_data.vti' + rect = ImageData(out_name, + [(0, 1), (1, 3), (0, 2)], + [N*2, N+2, N]) + + x, y, z = np.meshgrid(x, y, z, indexing='ij') + r = np.sqrt(x**2 + y**2 + z**2) + + rect.addPointData(DataArray(r, range(3), 'R')) + rect.write() + + output = open(out_name, 'r') + reference = open('test_image_data.ref', 'r') + + assert output.read() == reference.read() diff --git a/tests/test_image_data.ref b/tests/test_image_data.ref new file mode 100644 index 0000000..2e983f8 --- /dev/null +++ b/tests/test_image_data.ref @@ -0,0 +1,12 @@ + + + + + + AAYAAA==AAAAAAAA8D+3Ku22lSnwP6r9Q4Dno/A/0pQ1s1Bo8T83Q5SRkm3yPxAm0dSWqfM/+aJbJsAS9T/NO39mnqD2P2ZmZmZmZvY/2jpkSi2E9j8pnghHmdz2P/rs/W8Sbfc/Yafy9qwx+D/8aOCeoiX5PxI9+lHMQ/o/484kXQaH+z/NzMzMzMz8PwZ1+Mb74/w/z7dnHxop/T+3Aqch5pr9P9WEURVnN/4/Myr28RX8/j8j6U6gC+b/PwaaKxEXeQBAmpmZmZmZAUCLS/PNFqMBQA+dR/pvvwFAIXirrEvuAUCKMFjPGi8CQGgZZckggQJA2ktFXn3jAkC9PpQzN1UDQM3MzMzMzARAITcC5dTUBEADZpWu2uwEQASS1H6nFAVALUFX4eJLBUBIcuZMFpIFQLAnT9Kx5gVAl2TaXRFJBkAAAAAAAAAIQI2mYEj2BghAzuJpEc0bCEASwGOSYD4IQIQfrnh2bghAmNe1xr+rCEB//GVA2/UIQFNb2jpYTAlATZHrgsc68z+wFRkFb13zP8rQwQv5w/M/hZjuG2dq9D891KnU30r1P81r9fKQXvY/6Hme53me9z/fyMUh+wP5P/5fPWBez/g/n3xqEkTq+D9CnOEaSTr5P46zwh1+vfk/r7BiQulw+j9NUhfI1FD7P007/XAiWfw/bSmY6JWF/T/v6CTaO7b+P7J2Gmz6y/4/6OXxxNoM/z86//sj0nf/P1WQiImdBQBAHYYk7HliAEBD1GPzQNEAQG0F/sqaUAFABAMdCe1jAkDBnpAMAm0CQMp6VmImiAJAk9Jnbwu1AkDA+4kIM/MCQKi7lhH2QQNAd3mzoIygA0CT2ZrMFg4EQPC9tHYOeQVAWTmCONaABUAJdjWsHJgFQMsTdRCwvgVAayxcvj/0BUBmch9ZXzgGQMkB5tqKigZANVL4KSvqBkAjPsAs4JUIQBCv/g+snAhARzpEgASxCEDoThUtyNIIQHV03LvAAQlAWRRucqQ9CUBiGIZhGIYJQAjEBeuy2glAq6qqqqqq+j9qjHLFssP6Pz8J3CJADvs/tFaljsCI+z8MwzAMwzD8P7OirioxA/0/0T8oEI78/T+CL5pnMRn/P+/u7u7u7v4/xQxwwYUE/z/fMkTK8ET/P1xzttIqr/8/7qyba8sgAEBPxjO/Dn0AQLwiHr4n6wBADZ1+K8RpAUCW7MXsmesBQO33G8zr9AFA9k4ekcQQAkDJRzJozz4CQNdUum2EfgJAssI6JzDPAkDEz5ii/C8DQImE9z77nwNAKArgEn2UBEASHGUcm5wEQNs+OyHitARAbnkKuBndBEBaF0Oo5hQFQCA9GdDOWwVAVXKjDj6xBUBQDePDixQGQLa4OBAmYAdAIZow5ktnB0C9N+BasHwHQPBCqbwsoAdALTQMDoLRB0Cat0AlWxAIQAWGfmlPXAhA8bIP/OW0CEDZHS+nfEMKQKUTxXnZSQpAAwvfuuZcCkBIHbsIiXwKQCMI1ZiTqApArAwAbsngCkD66DHv3iQLQGCrSMd7dAtAqPSXm3fjAUDdEZm1zewBQPOgdAOzCAJAMYL0P9I2AkA+lkfWonYCQMNECmtwxwJA2RzPDWQoA0AuIQkUjpgDQNI1SiDOhwNAZJ4MpluQA0CYUlfl7akDQD6cjQdD1ANAKlRv7vAOBEBAazsyalkEQFc0/1oEswRARaGYvf4aBUD3xpe6moYFQJSflphdjgVArwBmgJWlBUCzLBcNEcwFQDcG1CqAAQZAHT0DPXdFBkBZn4Ufc5cGQBOkd6/d9gZAjLZPSiTJB0CJ2IucKtAHQKmIQi8x5QdAmhbDQhMICEDu1AD4lDgIQDJSoERldghASe3XeSDBCEDQ5W8xUxgJQBh4p4LwPQpA9Go9rU5ECkDg/oLwX1cKQJTgftkIdwpAiu7cgRyjCkDewEXHXdsKQJdNv+OAHwtAQvs/Ty1vC0DzWWFEK9gMQAt2a3n23QxA/g2ZIlHvDECIoKWHJgwNQN+JZKhUNA1AfSOJAa1nDUAvuaKR9aUNQEk/aBHq7g1A + + + + + _ + \ No newline at end of file diff --git a/uvw/__init__.py b/uvw/__init__.py index f7fa82f..ce55293 100644 --- a/uvw/__init__.py +++ b/uvw/__init__.py @@ -1,2 +1,2 @@ -from .vtk_files import VTKFile, RectilinearGrid +from .vtk_files import VTKFile, RectilinearGrid, ImageData from .data_array import DataArray diff --git a/uvw/vtk_files.py b/uvw/vtk_files.py index 6995d40..a6f0cc9 100644 --- a/uvw/vtk_files.py +++ b/uvw/vtk_files.py @@ -1,67 +1,99 @@ from . import writer from . import data_array import functools import numpy as np class VTKFile: """Generic VTK file""" def __init__(self, filename, filetype, rank=None): self.filename = filename self.rank = rank self.writer = writer.Writer(filetype) def addPointData(self, data_array): self.point_data.registerDataArray(data_array) def addCellData(self, data_array): self.cell_data.registerDataArray(data_array) def write(self): self.writer.registerAppend() self.writer.write(self.filename) +class ImageData(VTKFile): + """VTK Image data (coordinates are given by a range and constant spacing)""" + def __init__(self, filename, ranges, points, rank=None): + VTKFile.__init__(self, filename, 'ImageData', rank) + + # Computing spacings + spacings = [(x[1] - x[0]) / (n - 1) for x, n in zip(ranges, points)] + + # Filling in missing coordinates + for _ in range(len(points), 3): + points.append(1) + + # Setting extents, spacings and origin + extent = functools.reduce(lambda x, y: x + "0 {} ".format(y-1), points, "") + spacings = functools.reduce(lambda x, y: x + "{} ".format(y), spacings, "") + origins = functools.reduce(lambda x, y: x + "{} ".format(y[0]), ranges, "") + + self.writer.setDataNodeAttributes({ + 'WholeExtent': extent, + 'Spacing': spacings, + 'Origin': origins + }) + + self.piece = self.writer.registerPiece({ + "Extent": extent + }) + + # Registering data elements + self.point_data = self.piece.register('PointData') + self.cell_data = self.piece.register('CellData') + + class RectilinearGrid(VTKFile): """VTK Rectilinear grid (coordinates are given by 3 seperate ranges)""" def __init__(self, filename, coordinates, rank=None): VTKFile.__init__(self, filename, 'RectilinearGrid', rank) # Checking that we actually have a list or tuple if type(coordinates).__name__ == 'ndarray': coordinates = [coordinates] self.coordinates = list(coordinates) # Filling in missing coordinates for _ in range(len(self.coordinates), 3): self.coordinates.append(np.array([0.])) # Setting data extent extent = [] for coord in self.coordinates: if coord.ndim != 1: raise Exception('Coordinate array should have only one dimension') extent.append(coord.size-1) extent = functools.reduce(lambda x, y: x + "0 {} ".format(y), extent, "") self.writer.setDataNodeAttributes({ "WholeExtent": extent }) self.piece = self.writer.registerPiece({ "Extent": extent }) # Registering coordinates coordinate_component = self.piece.register('Coordinates') for coord, prefix in zip(self.coordinates, ('x', 'y', 'z')): array = data_array.DataArray(coord, [0], prefix + '_coordinates') coordinate_component.registerDataArray(array) # Registering data elements self.point_data = self.piece.register('PointData') self.cell_data = self.piece.register('CellData')