diff --git a/tests/conftest.py b/tests/conftest.py index 9317517..fabe522 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,40 +1,68 @@ import pytest import numpy as np from vtk.util.numpy_support import vtk_to_numpy as v2n -@pytest.fixture -def threeD_data(): +@pytest.fixture(params=[1, 2, 3]) +def field_data(request): N = 4 - x = np.linspace(0, 1, N*2) - y = np.linspace(0, 1, N+2) - z = np.linspace(0, 1, N) + dim = request.param + + if dim == 1: + x = np.linspace(0, 1, N*2) + coords = [x] + r = np.abs(x) + + if dim == 2: + x = np.linspace(0, 1, N*2) + y = np.linspace(0, 1, N+2) + coords = [x, y] + + xx, yy = np.meshgrid(x, y, indexing='ij', sparse=True) + r = np.sqrt(xx**2 + yy**2) + + if dim == 3: + x = np.linspace(0, 1, N*2) + y = np.linspace(0, 1, N+2) + z = np.linspace(0, 1, N) + coords = [x, y, z] + + xx, yy, zz = np.meshgrid(x, y, z, indexing='ij', sparse=True) + r = np.sqrt(xx**2 + yy**2 + zz**2) - xx, yy, zz = np.meshgrid(x, y, z, indexing='ij', sparse=True) - r = np.sqrt(xx**2 + yy**2 + zz**2) e_r = np.zeros([s - 1 for s in r.shape] + [3, 3]) - e_r[0, 0, 0, :, :] = np.array([[0, 1, 0], [1, 0, 0], [0, 1, 1]]) - e_r[-1, 0, 0, :, :] = np.eye(3) + e_r[..., :, :] = np.array([[0, 1, 0], [1, 0, 0], [0, 1, 1]]) + e_r[..., :, :] = np.eye(3) - return [x, y, z], r, e_r + return coords, r, e_r @pytest.fixture(params=[False, True]) def compression_fixture(request): return request @pytest.fixture(params=['ascii', 'binary', 'append']) def format_fixture(request): return request def get_vtk_data(reader, sstream): - reader.SetReadFromInputString(True) - reader.SetInputString(sstream.getvalue()) + if type(sstream) is str: + reader.SetFileName(sstream) + else: + reader.SetReadFromInputString(True) + reader.SetInputString(sstream.getvalue()) reader.Update() output = reader.GetOutput() return v2n(output.GetPointData().GetArray('point')), \ v2n(output.GetCellData().GetArray('cell')) + + +def transp(dim): + trans = list(range(dim+2)) + trans[-2], trans[-1] = trans[-1], trans[-2] + return trans + diff --git a/tests/test_edge_cases.py b/tests/test_edge_cases.py new file mode 100644 index 0000000..a46dce6 --- /dev/null +++ b/tests/test_edge_cases.py @@ -0,0 +1,63 @@ +import pytest +import numpy as np + +from uvw import ( + DataArray, + RectilinearGrid, + StructuredGrid, +) + + +def test_malformed_attributes(): + x = np.array([0, 1]) + + rect = RectilinearGrid('', x) + + with pytest.raises(ValueError): + rect.piece.register('', []) + + +def test_unsupported_format(): + x = np.array([0, 1]) + + rect = RectilinearGrid('', x) + with pytest.raises(ValueError): + rect.addPointData(DataArray(x, [0]), '#dummy') + + +def test_invalid_compression(): + x = np.array([0, 1]) + with pytest.raises(ValueError): + RectilinearGrid('', x, compression=100) + + +def test_invalid_file(): + x = np.array([0, 1]) + grid = RectilinearGrid('', x) + with pytest.raises(ValueError): + grid.writer.write(1) + + +def test_array_dimensions(): + x = np.array([0, 1]) + with pytest.raises(ValueError): + DataArray(x, range(2)) + + with pytest.raises(ValueError): + DataArray(x, range(1), components_order=2) + + +def test_rectilinear_fails(): + x = np.array([[0, 1], [2, 3]]) + with pytest.raises(ValueError): + RectilinearGrid('', x) + + x = np.array([0, 1]) + with pytest.raises(ValueError): + RectilinearGrid('', x, []) + + +def test_structured_fails(): + x = np.array([1, 2]) + with pytest.raises(ValueError): + StructuredGrid('', x, (1, 2)) diff --git a/tests/test_parallel.py b/tests/test_parallel.py index 53228ef..5437ba9 100644 --- a/tests/test_parallel.py +++ b/tests/test_parallel.py @@ -1,77 +1,106 @@ import numpy as np import pytest import vtk import os from mpi4py import MPI from vtk.util.numpy_support import vtk_to_numpy +from vtk import vtkXMLPRectilinearGridReader +from conftest import get_vtk_data, transp +from numpy import all from uvw.parallel import PRectilinearGrid from uvw import DataArray +def test_prectilinear_grid(field_data, compression_fixture, format_fixture): + coords, r, e_r = field_data + dim = r.ndim + out_name = 'test_prectilinear_grid.pvtr' + + compress = compression_fixture.param + format = format_fixture.param + with PRectilinearGrid(out_name, + coords, dim * [0], compression=compress) as rect: + rect.init_master(None) # useless here: for coverage only + rect.addPointData(DataArray(r, range(dim), 'point'), vtk_format=format) + rect.addCellData(DataArray(e_r, range(dim), 'cell'), vtk_format=format) + + reader = vtkXMLPRectilinearGridReader() + vtk_r, vtk_e_r = get_vtk_data(reader, out_name) + + vtk_r = vtk_r.reshape(r.shape, order='F') + vtk_e_r = vtk_e_r.reshape(e_r.shape, order='F').transpose(transp(dim)) + + assert all(vtk_r == r) + assert all(vtk_e_r == e_r) + + os.remove(out_name) + os.remove('test_prectilinear_grid_rank0.vtr') + + @pytest.mark.mpi(min_size=2) -def test_prectilinear_grid(compression_fixture, format_fixture): +def test_prectilinear_grid_mpi(compression_fixture, format_fixture): comm = MPI.COMM_WORLD rank = comm.Get_rank() N = 3 bounds = [ (0, 1), (1, 2) ] offsets = [ [0, 0, 0], [2*N, 0, 0], ] x = np.linspace(*bounds[rank], N*2) y = np.linspace(0, 1, N+2) z = np.linspace(0, 1, N) - out_name = 'test_prectilinear_grid.pvtr' + out_name = 'test_prectilinear_grid_mpi.pvtr' xx, yy, zz = np.meshgrid(x, y, z, indexing='ij', sparse=True) r = np.sqrt(xx**2 + yy**2 + zz**2) compress = compression_fixture.param format = format_fixture.param rect = PRectilinearGrid(out_name, (x, y, z), offsets[rank], compression=compress) rect.addPointData(DataArray(r, range(3), 'R'), vtk_format=format) rect.write() if rank == 0: reader = vtk.vtkXMLPRectilinearGridReader() reader.SetFileName(out_name) reader.Update() output = reader.GetOutput() vtk_r = vtk_to_numpy(output.GetPointData().GetArray('R')) else: vtk_r = None vtk_r = comm.bcast(vtk_r, root=0) vtk_r = vtk_r.reshape([4*N-1, N+2, N], order='F') i, j, k = [int(i) for i in offsets[rank]] # Adjusting for overlap if offsets[rank][0] != 0: i -= 1 sub_vtk = vtk_r[i:i+x.size, j:j+y.size, k:k+z.size] assert np.all(sub_vtk == r) if rank == 0: try: os.remove(out_name) - os.remove('test_prectilinear_grid_rank0.vtr') - os.remove('test_prectilinear_grid_rank1.vtr') - except: + os.remove('test_prectilinear_grid_mpi_rank0.vtr') + os.remove('test_prectilinear_grid_mpi_rank1.vtr') + except FileNotFoundError: pass if __name__ == '__main__': test_prectilinear_grid() diff --git a/tests/test_vtk_files.py b/tests/test_vtk_files.py index 6c61f4f..61e0912 100644 --- a/tests/test_vtk_files.py +++ b/tests/test_vtk_files.py @@ -1,99 +1,107 @@ import io import numpy as np from numpy import all, min, max from vtk import ( vtkXMLRectilinearGridReader, vtkXMLImageDataReader, vtkXMLStructuredGridReader, ) from vtk.util.numpy_support import vtk_to_numpy -from conftest import get_vtk_data +from conftest import get_vtk_data, transp from uvw import ( ImageData, RectilinearGrid, StructuredGrid, DataArray, ) -def test_rectilinear_grid(threeD_data, compression_fixture, format_fixture): - coords, r, e_r = threeD_data +def test_rectilinear_grid(field_data, compression_fixture, format_fixture): + coords, r, e_r = field_data + dim = r.ndim f = io.StringIO() compress = compression_fixture.param format = format_fixture.param - with RectilinearGrid(f, coords, compression=compress) as rect: - rect.addPointData(DataArray(r, range(3), 'point'), vtk_format=format) - rect.addCellData(DataArray(e_r, range(3), 'cell'), vtk_format=format) + rect = RectilinearGrid(f, coords, compression=compress) + rect.addPointData(DataArray(r, range(dim), 'point'), vtk_format=format) + rect.addCellData(DataArray(e_r, range(dim), 'cell'), vtk_format=format) + rect.write() reader = vtkXMLRectilinearGridReader() - vtk_r, vtk_e_r = get_vtk_data(reader, f) - vtk_r = vtk_r.reshape(r.shape, order='F') - vtk_e_r = vtk_e_r.reshape(e_r.shape, order='F').transpose((0, 1, 2, 4, 3)) + # Testing the xml pretty print output as well + pretty_sstream = io.StringIO(str(rect.writer)) - assert all(vtk_r == r) - assert all(vtk_e_r == e_r) + for ss in [f, pretty_sstream]: + vtk_r, vtk_e_r = get_vtk_data(reader, ss) + + vtk_r = vtk_r.reshape(r.shape, order='F') + vtk_e_r = vtk_e_r.reshape(e_r.shape, order='F').transpose(transp(dim)) + + assert all(vtk_r == r) + assert all(vtk_e_r == e_r) -def test_image_data(threeD_data, compression_fixture, format_fixture): - coords, r, e_r = threeD_data +def test_image_data(field_data, compression_fixture, format_fixture): + coords, r, e_r = field_data + dim = r.ndim f = io.StringIO() compress = compression_fixture.param format = format_fixture.param with ImageData( f, [(min(x), max(x)) for x in coords], [x.size for x in coords], compression=compress) as fh: - fh.addPointData(DataArray(r, range(3), 'point'), vtk_format=format) - fh.addCellData(DataArray(e_r, range(3), 'cell'), vtk_format=format) + fh.addPointData(DataArray(r, range(dim), 'point'), vtk_format=format) + fh.addCellData(DataArray(e_r, range(dim), 'cell'), vtk_format=format) reader = vtkXMLImageDataReader() vtk_r, vtk_e_r = get_vtk_data(reader, f) vtk_r = vtk_r.reshape(r.shape, order='F') - vtk_e_r = vtk_e_r.reshape(e_r.shape, order='F').transpose((0, 1, 2, 4, 3)) + vtk_e_r = vtk_e_r.reshape(e_r.shape, order='F').transpose(transp(dim)) assert all(vtk_r == r) assert all(vtk_e_r == e_r) def test_structured_grid(compression_fixture, format_fixture): f = io.StringIO() N = 5 r = np.linspace(0, 1, N) theta = np.linspace(0, 2*np.pi, 5*N) theta, r = np.meshgrid(theta, r, indexing='ij') x = r*np.cos(theta) y = r*np.sin(theta) points = np.vstack([x.ravel(), y.ravel()]).T compress = compression_fixture.param format = format_fixture.param grid = StructuredGrid(f, points, (N, 5*N), compression=compress) data = np.exp(-4*r**2) grid.addPointData(DataArray(data, reversed(range(2)), 'data'), vtk_format=format) grid.write() reader = vtkXMLStructuredGridReader() reader.SetReadFromInputString(True) reader.SetInputString(f.getvalue()) reader.Update() output = reader.GetOutput() vtk_data = vtk_to_numpy(output.GetPointData().GetArray('data')) vtk_data = vtk_data.reshape(data.shape, order='C') assert all(vtk_data == data) diff --git a/uvw/data_array.py b/uvw/data_array.py index f2da9b6..67d25ca 100644 --- a/uvw/data_array.py +++ b/uvw/data_array.py @@ -1,47 +1,50 @@ import functools class DataArray: """Class holding information on ndarray""" def __init__(self, data, spatial_axes, name='', components_order='C'): """ Data array constructor :param data: the numpy array containing the data (possibly a view) :param spatial_axes: a container of ints that indicate which axes of the array correspond to space dimensions (in order) :param name: the name of the data :param components_order: the order of the non-spatial axes of the array """ self.data = data axes = list(range(data.ndim)) spatial_axes = list(spatial_axes) if data.ndim < len(spatial_axes): - raise Exception('Dimensions of data smaller than space dimensions') + raise ValueError( + 'Dimensions of data smaller than space dimensions') for ax in spatial_axes: axes.remove(ax) nb_components = functools.reduce( lambda x, y: x * data.shape[y], axes, 1) if components_order == 'C': axes.reverse() + elif components_order == 'F': # pragma: no cover + pass else: - raise Exception('Unrecognized components order') + raise ValueError('Unrecognized components order') axes += spatial_axes # Hopefully this is a view self.flat_data = self.data.transpose(*axes).reshape(-1, order='F') self.attributes = { "Name": name, "type": str(self.flat_data.dtype).capitalize(), "NumberOfComponents": str(nb_components) } - def __str__(self): + def __str__(self): # pragma: no cover return self.attributes.__str__() diff --git a/uvw/vtk_files.py b/uvw/vtk_files.py index 9fda75f..d1cfd48 100644 --- a/uvw/vtk_files.py +++ b/uvw/vtk_files.py @@ -1,164 +1,164 @@ from . import writer from . import data_array import functools import numpy as np class VTKFile: """Generic VTK file""" def __init__(self, filename, filetype, compression): self.filetype = filetype self.filename = filename self.writer = writer.Writer(filetype, compression) # Center piece self.piece = self.writer.registerPiece() # Registering data elements self.point_data = self.piece.register('PointData') self.cell_data = self.piece.register('CellData') def addPointData(self, data_array, vtk_format='binary'): self.point_data.registerDataArray(data_array, vtk_format) def addCellData(self, data_array, vtk_format='binary'): self.cell_data.registerDataArray(data_array, vtk_format) def write(self): self.writer.registerAppend() self.writer.write(self.filename) def __enter__(self): return self def __exit__(self, *args): self.write() class ImageData(VTKFile): """VTK Image data (coordinates given by a range and constant spacing)""" def __init__(self, filename, ranges, points, compression=None): VTKFile.__init__(self, filename, 'ImageData', compression) # 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, spacing 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.setAttributes({ "Extent": extent }) class RectilinearGrid(VTKFile): """VTK Rectilinear grid (coordinates are given by 3 seperate ranges)""" def __init__(self, filename, coordinates, offsets=None, compression=None): VTKFile.__init__(self, filename, 'RectilinearGrid', compression) # 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( + raise ValueError( 'Coordinate array should have only one dimension' + ' (has {})'.format(coord.ndim)) extent.append(coord.size-1) if offsets is None: offsets = [0] * len(extent) elif len(offsets) == len(coordinates): offsets = list(offsets) offsets += [0] * (len(extent) - len(offsets)) else: - raise Exception( + raise ValueError( 'Size of offsets should ' 'match domain dimension {}'.format(len(coordinates))) def fold_extent(acc, couple): offset, extent = couple if offset != 0: - offset -= 1 + offset -= 1 # pragma: no cover return acc + "{} {} ".format(offset, offset+extent) # Create extent string with offsets self.extent = functools.reduce(fold_extent, zip(offsets, extent), "") self.writer.setDataNodeAttributes({ "WholeExtent": self.extent }) self.piece.setAttributes({ "Extent": self.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, vtk_format='append') class StructuredGrid(VTKFile): """VTK Structured grid (coordinates given by a single array of points)""" def __init__(self, filename, points, shape, compression=None): VTKFile.__init__(self, filename, 'StructuredGrid', compression) if points.ndim != 2: - raise 'Points should be a 2D array' + raise ValueError('Points should be a 2D array') # Completing the missing coordinates points_3d = np.zeros((points.shape[0], 3)) for i in range(points.shape[1]): points_3d[:, i] = points[:, i] extent = [n - 1 for n in shape] for i in range(len(extent), 3): extent.append(0) extent = functools.reduce( lambda x, y: x + "0 {} ".format(y), extent, "") self.writer.setDataNodeAttributes({ "WholeExtent": extent }) self.piece.setAttributes({ "Extent": extent }) points_component = self.piece.register('Points') points_component.registerDataArray( data_array.DataArray(points_3d, [0], 'points')) diff --git a/uvw/writer.py b/uvw/writer.py index a160781..4fd82f2 100644 --- a/uvw/writer.py +++ b/uvw/writer.py @@ -1,198 +1,200 @@ import xml.dom.minidom as dom import io import zlib import numpy as np from functools import reduce from base64 import b64encode def setAttributes(node, attributes): """Set attributes of a node""" for item in attributes.items(): node.setAttribute(*item) def encodeArray(array, level): """Encode array data and header in base64.""" def compress(array): """Compress array with zlib. Returns header and compressed data.""" raw_data = memoryview(array) # avoid a copy data_size = raw_data.nbytes max_block_size = 2**15 # Enough blocks to span whole data nblocks = data_size // max_block_size + 1 last_block_size = data_size % max_block_size # Compress regular blocks compressed_data = [ zlib.compress(raw_data[i*max_block_size:(i+1)*max_block_size], level) for i in range(nblocks-1) ] # Compress last (smaller) block compressed_data.append( zlib.compress(raw_data[-last_block_size:], level) ) # Header data (cf https://vtk.org/Wiki/VTK_XML_Formats#Compressed_Data) usize = max_block_size psize = last_block_size csize = [len(x) for x in compressed_data] header = np.array([nblocks, usize, psize] + csize, dtype=np.uint32) return header.tobytes(), b"".join(compressed_data) def raw(array): """Returns header and array data in bytes.""" header = np.array([array.nbytes], dtype=np.uint32) return header.tobytes(), memoryview(array) if level is not None: data = compress(array) else: data = raw(array) return "".join(map(lambda x: b64encode(x).decode(), data)) class Component: """Generic component class capable of registering sub-components""" def __init__(self, name, parent_node, writer): self.writer = writer self.document = writer.document self.node = self.document.createElement(name) parent_node.appendChild(self.node) def setAttributes(self, attributes): setAttributes(self.node, attributes) def register(self, name, attributes=None): """Register a sub-component""" if attributes is None: attributes = {} if type(attributes) != dict: - raise Exception( + raise ValueError( 'Cannot register attributes of type ' + str(type(attributes))) sub_component = Component(name, self.node, self.writer) setAttributes(sub_component.node, attributes) return sub_component def _addArrayNodeData(self, data_array, component, vtk_format): if vtk_format == 'ascii': sstream = io.StringIO() np.savetxt(sstream, data_array.flat_data, newline=' ') data_as_str = sstream.getvalue() # reduce(lambda x, y: x + str(y) + ' ', data_array.flat_data, "") elif vtk_format == 'binary': data_as_str = encodeArray(data_array.flat_data, self.writer.compression) elif vtk_format == 'append': self.writer.append_data_arrays[component] = data_array else: - raise Exception('Unsupported VTK Format "{}"'.format(vtk_format)) + raise ValueError('Unsupported VTK Format "{}"'.format(vtk_format)) if vtk_format != 'append': component.node.appendChild( self.document.createTextNode(data_as_str)) def _registerArrayComponent(self, array, name, vtk_format): attributes = array.attributes attributes['format'] = vtk_format return self.register(name, attributes) def registerDataArray(self, data_array, vtk_format='binary'): """Register a DataArray object""" component = self._registerArrayComponent(data_array, 'DataArray', vtk_format) self._addArrayNodeData(data_array, component, vtk_format) def registerPDataArray(self, data_array, vtk_format='binary'): """Register a DataArray object in p-file""" self._registerArrayComponent(data_array, 'PDataArray', vtk_format) class Writer: """Generic XML handler for VTK files""" def __init__(self, vtk_format, compression=None, vtk_version='0.1', byte_order='LittleEndian'): self.document = dom.getDOMImplementation() \ .createDocument(None, 'VTKFile', None) self.root = self.document.documentElement self.root.setAttribute('type', vtk_format) self.root.setAttribute('version', vtk_version) self.root.setAttribute('byte_order', byte_order) self.data_node = self.document.createElement(vtk_format) self.root.appendChild(self.data_node) self.size_indicator_bytes = np.dtype(np.uint32).itemsize self.append_data_arrays = {} if compression is not None and compression is not False: self.root.setAttribute('compressor', 'vtkZLibDataCompressor') if type(compression) is not int: compression = -1 else: if compression not in list(range(-1, 10)): - raise Exception(('compression level {} is not ' - 'recognized by zlib').format(compression)) + raise ValueError( + ('compression level {} is not ' + 'recognized by zlib').format(compression) + ) elif not compression: compression = None self.compression = compression def setDataNodeAttributes(self, attributes): """Set attributes for the entire dataset""" setAttributes(self.data_node, attributes) def registerPiece(self, attributes={}): """Register a piece element""" return self.registerComponent('Piece', self.data_node, attributes) def registerComponent(self, name, parent, attributes={}): """Register a Component to a parent Component with set attributes""" comp = Component(name, parent, self) setAttributes(comp.node, attributes) return comp def registerAppend(self): """Register AppendedData node (should only be called once)""" append_node = Component('AppendedData', self.root, self) append_node.setAttributes({'format': 'base64'}) self.root.appendChild(append_node.node) data_str = "_" offset = 0 for component, data in self.append_data_arrays.items(): component.setAttributes({'offset': str(offset)}) data_b64 = encodeArray(data.flat_data, self.compression) data_str += data_b64 offset += len(data_b64) text = self.document.createTextNode(data_str) append_node.node.appendChild(text) def write(self, fd): """Write to file descriptor""" - if type(fd) == str: + if type(fd) is str: with open(fd, 'w') as file: self.write(file) elif issubclass(type(fd), io.TextIOBase): self.document.writexml(fd, indent="\n ", addindent=" ") else: - raise RuntimeError("Expected a path or " - + "file handle, got {}".format(type(fd))) + raise ValueError("Expected a path or " + + "file descriptor, got {}".format(type(fd))) def __str__(self): """Print XML to string""" return self.document.toprettyxml()