diff --git a/uvw/data_array.py b/uvw/data_array.py index 1c400e2..ce1bc83 100644 --- a/uvw/data_array.py +++ b/uvw/data_array.py @@ -1,41 +1,44 @@ import numpy as np import functools import operator + 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)) for ax in spatial_axes: axes.remove(ax) - nb_components = functools.reduce(lambda x, y: x * data.shape[y], axes, 1) + nb_components = functools.reduce( + lambda x, y: x * data.shape[y], axes, 1) if components_order == 'C': axes.reverse() else: raise Exception('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): return self.attributes.__str__() diff --git a/uvw/vtk_files.py b/uvw/vtk_files.py index 29b7d3e..fb4478c 100644 --- a/uvw/vtk_files.py +++ b/uvw/vtk_files.py @@ -1,129 +1,139 @@ 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) # 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): 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, self.__class__.__name__, 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, "") + 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, rank=None): VTKFile.__init__(self, filename, self.__class__.__name__, 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') + 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, "") + extent = functools.reduce( + lambda x, y: x + "0 {} ".format(y), extent, "") self.writer.setDataNodeAttributes({ "WholeExtent": extent }) self.piece.setAttributes({ "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) class StructuredGrid(VTKFile): """VTK Structured grid (coordinates given by a single array of points)""" + def __init__(self, filename, points, shape, rank=None): VTKFile.__init__(self, filename, self.__class__.__name__, rank) if points.ndim != 2: raise '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, "") + 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')) + points_component.registerDataArray( + data_array.DataArray(points_3d, [0], 'points')) diff --git a/uvw/writer.py b/uvw/writer.py index 471c0cb..007a54c 100644 --- a/uvw/writer.py +++ b/uvw/writer.py @@ -1,107 +1,112 @@ import xml.dom import xml.dom.minidom as dom import functools import base64 import numpy as np + def setAttributes(node, attributes): """Set attributes of a node""" for item in attributes.items(): node.setAttribute(*item) + def encodeArray(array): # Mandatory number of bytes encoded as uint32 nbytes = array.nbytes bytes = base64.b64encode(np.array([nbytes], dtype=np.uint32)) bytes += base64.b64encode(array) return bytes + 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=dict()): """Register a sub-component""" sub_component = Component(name, self.node, self.writer) setAttributes(sub_component.node, attributes) return sub_component def registerDataArray(self, data_array, vtk_format='binary'): """Register a DataArray object""" array_component = Component('DataArray', self.node, self.writer) attributes = data_array.attributes attributes['format'] = vtk_format if vtk_format == 'append': raise 'Feature does not work' attributes['offset'] = str(self.writer.offset) array = data_array.flat_data self.writer.offset += array.nbytes self.writer.offset += self.writer.size_indicator_bytes self.writer.append_data_arrays.append(data_array) elif vtk_format == 'ascii': - data_as_str = functools.reduce(lambda x, y: x + str(y) + ' ', data_array.flat_data, "") - array_component.node.appendChild(self.document.createTextNode(data_as_str)) + data_as_str = functools.reduce( + lambda x, y: x + str(y) + ' ', data_array.flat_data, "") + array_component.node.appendChild( + self.document.createTextNode(data_as_str)) elif vtk_format == 'binary': array_component.node.appendChild( self.document.createTextNode( encodeArray(data_array.flat_data).decode('ascii'))) - setAttributes(array_component.node, attributes) class Writer: """Generic XML handler for VTK files""" + def __init__(self, vtk_format, 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.offset = 0 # Global offset + self.offset = 0 # Global offset self.size_indicator_bytes = np.dtype(np.uint32).itemsize self.append_data_arrays = [] def setDataNodeAttributes(self, attributes): """Set attributes for the entire dataset""" setAttributes(self.data_node, attributes) def registerPiece(self, attributes={}): """Register a piece element""" piece = Component('Piece', self.data_node, self) setAttributes(piece.node, attributes) return piece def registerAppend(self): append_node = Component('AppendedData', self.root, self) setAttributes(append_node.node, {'format': 'base64'}) self.root.appendChild(append_node.node) data_str = b"_" for data_array in self.append_data_arrays: data_str += encodeArray(data_array.flat_data) text = self.document.createTextNode(data_str.decode('ascii')) append_node.node.appendChild(text) def write(self, filename): with open(filename, 'w') as file: self.document.writexml(file, indent="\n ", addindent=" ") def __str__(self): """Print XML to string""" return self.document.toprettyxml() -