gkowal/amun-code
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

PYTHON: Rewrite vtkio.

Browse Source 
Signed-off-by: Grzegorz Kowal <grzegorz@amuncode.org>
This commit is contained in:
Grzegorz Kowal 2021-10-05 10:58:44 -03:00
parent e153cd417e
commit 1c4fa6035d
1 changed files with 38 additions and 110 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

148
python/amunpy/src/amunpy/vtkio.py
View File

@ -22,120 +22,67 @@
================================================================================ ================================================================================
submodule: vtkio.py module: VTKIO
This submodule provides a function to store a dataset as VTK image file. Module provides a function to store given dataset in the VTK image file.
The only requirements for this package are:
- numpy
- lz4 (for LZ4 compression)
-------------------------------------------------------------------------------- --------------------------------------------------------------------------------
""" """
import base64, numpy, struct, zlib, lz4.block, lzma
def WriteVTK(vtkfile, vname, data, origin = (0.0, 0.0, 0.0), \ def WriteVTK(vtkfile, vname, data, \
spacing = (1.0, 1.0, 1.0), points = False, encode = True, \ origin=(0, 0, 0), spacing=(1, 1, 1), \
compression=None, block_size = 32768, lz4mode = 'fast', \ compression=None, compression_level=19, encode=True, \
compression_level = 6, verbose = False): 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)): if isinstance(data, (list, tuple)):
# Check if all list or tuple components are arrays. if not all(isinstance(d, (numpy.ndarray)) for d in data):
# raise Exception('All input data components in WriteVTK must be arrays!')
allarrays = True
for d in data:
allarrays = allarrays & isinstance(d, (numpy.ndarray))
# Quit if not all elements are arrays. 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!')
if not allarrays:
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' dtype = 'vector'
ncomp = len(data)
# Get the number of vector components. ndims = data[0].ndim
# dims = data[0].shape
nc = len(data)
# Get the number of dimansions
#
nd = data[0].ndim
# Get the variable dimensions.
#
dm = data[0].shape
elif isinstance(data, (numpy.ndarray)): elif isinstance(data, (numpy.ndarray)):
# Define the type of input data.
#
dtype = 'scalar' dtype = 'scalar'
ncomp = 1
# Get the number of vector components. ndims = data.ndim
# dims = data.shape
nc = 1
# Get the number of dimansions
#
nd = data.ndim
# Get the variable dimensions.
#
dm = data.shape
else: else:
print('Unknown type of the input data!') raise Exception('Unknown type of the input data in WriteVTK!')
return
# Create the VTK output file and open it in the binary mode.
#
with open(vtkfile, 'wb') as vt: with open(vtkfile, 'wb') as vt:
# Initiate offset and array size.
#
offset = 0 offset = 0
# Prepare strong for dimensions.
#
sdims = '"' sdims = '"'
if points: if points:
for i in range(nd - 1, 0, -1): for i in range(ndims - 1, 0, -1):
sdims +='%d %d ' % (0, dm[i] - 1) sdims +='%d %d ' % (0, dims[i] - 1)
sdims += '%d %d" ' % (0, dm[0] - 1) sdims += '%d %d" ' % (0, dims[0] - 1)
else: else:
for i in range(nd - 1, 0, -1): for i in range(ndims - 1, 0, -1):
sdims +='%d %d ' % (0, dm[i]) sdims +='%d %d ' % (0, dims[i])
sdims += '%d %d" ' % (0, dm[0]) sdims += '%d %d" ' % (0, dims[0])
# Write the VTK header and data description.
#
string = '<?xml version="1.0"?>\n<VTKFile type="ImageData" version="1.0" byte_order="LittleEndian" header_type="UInt64"' string = '<?xml version="1.0"?>\n<VTKFile type="ImageData" version="1.0" byte_order="LittleEndian" header_type="UInt64"'
if compression == 'zlib': if compression == 'zlib':
string += ' compressor="vtkZLibDataCompressor"' string += ' compressor="vtkZLibDataCompressor"'
compression_level = min(max(compression_level, 0), 9)
elif compression == 'lzma': elif compression == 'lzma':
string += ' compressor="vtkLZMADataCompressor"' string += ' compressor="vtkLZMADataCompressor"'
compression_level = min(max(compression_level, 0), 9)
elif compression == 'lz4': elif compression == 'lz4':
string += ' compressor="vtkLZ4DataCompressor"' string += ' compressor="vtkLZ4DataCompressor"'
compression_level = min(max(compression_level, 0), 12)
lz4acceleration = max(lz4acceleration, 1)
string += '>\n <ImageData WholeExtent=%s' % sdims string += '>\n <ImageData WholeExtent=%s' % sdims
string += 'Origin="%e %e %e" ' % origin string += 'Origin="%e %e %e" ' % origin
string += 'Spacing="%e %e %e">\n' % spacing string += 'Spacing="%e %e %e">\n' % spacing
@ -144,7 +91,7 @@ def WriteVTK(vtkfile, vname, data, origin = (0.0, 0.0, 0.0), \
string += ' <PointData %ss="%s">\n' % (dtype, vname) string += ' <PointData %ss="%s">\n' % (dtype, vname)
else: else:
string += ' <CellData %ss="%s">\n' % (dtype, vname) string += ' <CellData %ss="%s">\n' % (dtype, vname)
if nc == 1: if ncomp == 1:
dmin = data.min() dmin = data.min()
dmax = data.max() dmax = data.max()
else: else:
@ -154,8 +101,8 @@ def WriteVTK(vtkfile, vname, data, origin = (0.0, 0.0, 0.0), \
dmin = numpy.sqrt(dd.min()) dmin = numpy.sqrt(dd.min())
dmax = numpy.sqrt(dd.max()) dmax = numpy.sqrt(dd.max())
string += ' <DataArray ' string += ' <DataArray '
if nc > 1: if ncomp > 1:
string += 'NumberOfComponents="{:d}" '.format(nc) string += 'NumberOfComponents="{:d}" '.format(ncomp)
string += 'type="Float32" Name="{}" RangeMin="{:e}" RangeMax="{:e}" format="appended" offset="{}">\n'.format(vname, dmin, dmax, offset) string += 'type="Float32" Name="{}" RangeMin="{:e}" RangeMax="{:e}" format="appended" offset="{}">\n'.format(vname, dmin, dmax, offset)
string += " </DataArray>\n" string += " </DataArray>\n"
if points: if points:
@ -166,8 +113,6 @@ def WriteVTK(vtkfile, vname, data, origin = (0.0, 0.0, 0.0), \
string += ' </ImageData>\n' string += ' </ImageData>\n'
vt.write(string.encode()) vt.write(string.encode())
# Append variable data.
#
if encode: if encode:
string = ' <AppendedData encoding="base64">\n' string = ' <AppendedData encoding="base64">\n'
else: else:
@ -175,40 +120,27 @@ def WriteVTK(vtkfile, vname, data, origin = (0.0, 0.0, 0.0), \
string += ' _' string += ' _'
vt.write(string.encode()) 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([ dims[0], dims[1], dims[2], ncomp ], dtype=numpy.float32)
if (dtype == 'vector'):
qt = numpy.zeros([ dm[0], dm[1], dm[2], nc ], dtype = numpy.float32)
for n, d in enumerate(data[:]): for n, d in enumerate(data[:]):
qt[:,:,:,n] = numpy.float32(d) qt[:,:,:,n] = numpy.float32(d)
else: else:
qt = numpy.float32(data) qt = numpy.float32(data)
barr = qt.tobytes() barr = qt.tobytes()
# Compress if desired.
#
if compression != None: if compression != None:
lz4compression = 10 - compression_level
if len(barr) < block_size: if len(barr) < block_size:
if compression == 'zlib': if compression == 'zlib':
carr = zlib.compress(barr, compression_level) carr = zlib.compress(barr, compression_level)
elif compression == 'lz4': 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': elif compression == 'lzma':
carr = lzma.compress(barr) 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)) head = struct.pack("QQQQ", 1, len(barr), 0, len(carr))
else: else:
nblocks = len(barr) // block_size nblocks = len(barr) // block_size
rsize = len(barr) % block_size rsize = len(barr) % block_size
if verbose: if verbose:
@ -231,7 +163,7 @@ def WriteVTK(vtkfile, vname, data, origin = (0.0, 0.0, 0.0), \
cctx = lzma cctx = lzma
for i in range(nblocks): for i in range(nblocks):
ie = min(len(barr), ib + block_size) 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)) head += struct.pack("Q", len(cblk))
carr += cblk carr += cblk
ib = ie ib = ie
@ -240,7 +172,7 @@ def WriteVTK(vtkfile, vname, data, origin = (0.0, 0.0, 0.0), \
cctx = lz4.block cctx = lz4.block
for i in range(nblocks): for i in range(nblocks):
ie = min(len(barr), ib + block_size) 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)) head += struct.pack("Q", len(cblk))
carr += cblk carr += cblk
ib = ie ib = ie
@ -251,16 +183,12 @@ def WriteVTK(vtkfile, vname, data, origin = (0.0, 0.0, 0.0), \
vt.write(head+carr) vt.write(head+carr)
else: else:
head = struct.pack("Q", len(barr)) head = struct.pack("Q", len(barr))
if encode: if encode:
vt.write(base64.b64encode(head+barr)) vt.write(base64.b64encode(head+barr))
else: else:
vt.write(head+barr) vt.write(head+barr)
# Close the appended data section.
#
string = '\n </AppendedData>\n' string = '\n </AppendedData>\n'
string += '</VTKFile>\n' string += '</VTKFile>\n'
vt.write(string.encode()) vt.write(string.encode())