diff --git a/python/amunpy/src/amunpy/vtkio.py b/python/amunpy/src/amunpy/vtkio.py index b6d1493..f574c0c 100644 --- a/python/amunpy/src/amunpy/vtkio.py +++ b/python/amunpy/src/amunpy/vtkio.py @@ -22,120 +22,67 @@ ================================================================================ - submodule: vtkio.py + module: VTKIO - This submodule provides a function to store a dataset as VTK image file. - - The only requirements for this package are: - - - numpy - - lz4 (for LZ4 compression) + Module provides a function to store given dataset in the VTK image file. -------------------------------------------------------------------------------- """ -import base64, numpy, struct, zlib, lz4.block, lzma -def WriteVTK(vtkfile, vname, data, origin = (0.0, 0.0, 0.0), \ - spacing = (1.0, 1.0, 1.0), points = False, encode = True, \ - compression=None, block_size = 32768, lz4mode = 'fast', \ - compression_level = 6, verbose = False): +def WriteVTK(vtkfile, vname, data, \ + origin=(0, 0, 0), spacing=(1, 1, 1), \ + compression=None, compression_level=19, encode=True, \ + lz4mode='default', lz4acceleration=1, block_size=32768, \ + points=False, verbose=False): + + import base64, numpy, struct, zlib, lz4.block, lzma - # Separate cases when the input is a vector field and scalar field. - # if isinstance(data, (list, tuple)): - # Check if all list or tuple components are arrays. - # - allarrays = True - for d in data: - allarrays = allarrays & isinstance(d, (numpy.ndarray)) + if not all(isinstance(d, (numpy.ndarray)) for d in data): + raise Exception('All input data components in WriteVTK must be arrays!') - # Quit if not all elements are arrays. - # - if not allarrays: + if not all(data[0].shape == d.shape for d in data): + raise Exception('All input data components in WriteVTK must have the same dimensions!') - print('All input data components must be arrays!') - return - - # Check if the components have the same dimensions - # - sameshapes = True - for d in data: - sameshapes = sameshapes & (data[0].shape == d.shape) - - # Quit if elements have different dimensions. - # - if not sameshapes: - print('All input data component must have the same dimensions!') - return - - # Define the type of input data. - # dtype = 'vector' - - # Get the number of vector components. - # - nc = len(data) - - # Get the number of dimansions - # - nd = data[0].ndim - - # Get the variable dimensions. - # - dm = data[0].shape + ncomp = len(data) + ndims = data[0].ndim + dims = data[0].shape elif isinstance(data, (numpy.ndarray)): - # Define the type of input data. - # dtype = 'scalar' - - # Get the number of vector components. - # - nc = 1 - - # Get the number of dimansions - # - nd = data.ndim - - # Get the variable dimensions. - # - dm = data.shape + ncomp = 1 + ndims = data.ndim + dims = data.shape else: - print('Unknown type of the input data!') - return + raise Exception('Unknown type of the input data in WriteVTK!') - # Create the VTK output file and open it in the binary mode. - # with open(vtkfile, 'wb') as vt: - - # Initiate offset and array size. - # offset = 0 - - # Prepare strong for dimensions. - # sdims = '"' if points: - for i in range(nd - 1, 0, -1): - sdims +='%d %d ' % (0, dm[i] - 1) - sdims += '%d %d" ' % (0, dm[0] - 1) + for i in range(ndims - 1, 0, -1): + sdims +='%d %d ' % (0, dims[i] - 1) + sdims += '%d %d" ' % (0, dims[0] - 1) else: - for i in range(nd - 1, 0, -1): - sdims +='%d %d ' % (0, dm[i]) - sdims += '%d %d" ' % (0, dm[0]) + for i in range(ndims - 1, 0, -1): + sdims +='%d %d ' % (0, dims[i]) + sdims += '%d %d" ' % (0, dims[0]) - # Write the VTK header and data description. - # string = '\n 1: - string += 'NumberOfComponents="{:d}" '.format(nc) + if ncomp > 1: + string += 'NumberOfComponents="{:d}" '.format(ncomp) string += 'type="Float32" Name="{}" RangeMin="{:e}" RangeMax="{:e}" format="appended" offset="{}">\n'.format(vname, dmin, dmax, offset) string += " \n" if points: @@ -166,8 +113,6 @@ def WriteVTK(vtkfile, vname, data, origin = (0.0, 0.0, 0.0), \ string += ' \n' vt.write(string.encode()) - # Append variable data. - # if encode: string = ' \n' else: @@ -175,40 +120,27 @@ def WriteVTK(vtkfile, vname, data, origin = (0.0, 0.0, 0.0), \ string += ' _' vt.write(string.encode()) - # Convert the input data to Float32. In the case of vector data concatenate the components first. - # - if (dtype == 'vector'): - qt = numpy.zeros([ dm[0], dm[1], dm[2], nc ], dtype = numpy.float32) + if dtype == 'vector': + qt = numpy.zeros([ dims[0], dims[1], dims[2], ncomp ], dtype=numpy.float32) for n, d in enumerate(data[:]): qt[:,:,:,n] = numpy.float32(d) - else: qt = numpy.float32(data) barr = qt.tobytes() - # Compress if desired. - # if compression != None: - - lz4compression = 10 - compression_level - if len(barr) < block_size: - if compression == 'zlib': carr = zlib.compress(barr, compression_level) elif compression == 'lz4': - carr = lz4.block.compress(barr, mode = lz4mode, acceleration = lz4compression, compression = compression_level, store_size = False) + carr = lz4.block.compress(barr, mode=lz4mode, acceleration=lz4acceleration, compression=compression_level, store_size=False) elif compression == 'lzma': carr = lzma.compress(barr) - # Prepare the number of blocks, the size of the last partial block, - # and the size of the compressed data. - # head = struct.pack("QQQQ", 1, len(barr), 0, len(carr)) else: - nblocks = len(barr) // block_size rsize = len(barr) % block_size if verbose: @@ -231,7 +163,7 @@ def WriteVTK(vtkfile, vname, data, origin = (0.0, 0.0, 0.0), \ cctx = lzma for i in range(nblocks): ie = min(len(barr), ib + block_size) - cblk = cctx.compress(barr[ib:ie], preset = compression_level) + cblk = cctx.compress(barr[ib:ie], preset=compression_level) head += struct.pack("Q", len(cblk)) carr += cblk ib = ie @@ -240,7 +172,7 @@ def WriteVTK(vtkfile, vname, data, origin = (0.0, 0.0, 0.0), \ cctx = lz4.block for i in range(nblocks): ie = min(len(barr), ib + block_size) - cblk = cctx.compress(barr[ib:ie], mode = lz4mode, acceleration = lz4compression, compression = compression_level, store_size = False) + cblk = cctx.compress(barr[ib:ie], mode=lz4mode, acceleration=lz4acceleration, compression=compression_level, store_size=False) head += struct.pack("Q", len(cblk)) carr += cblk ib = ie @@ -251,16 +183,12 @@ def WriteVTK(vtkfile, vname, data, origin = (0.0, 0.0, 0.0), \ vt.write(head+carr) else: - head = struct.pack("Q", len(barr)) - if encode: vt.write(base64.b64encode(head+barr)) else: vt.write(head+barr) - # Close the appended data section. - # string = '\n \n' string += '\n' vt.write(string.encode())