PYTHON: Rewrite vtkio.
Signed-off-by: Grzegorz Kowal <grzegorz@amuncode.org>
This commit is contained in:
parent
e153cd417e
commit
1c4fa6035d
@ -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 = '<?xml version="1.0"?>\n<VTKFile type="ImageData" version="1.0" byte_order="LittleEndian" header_type="UInt64"'
|
||||
if compression == 'zlib':
|
||||
string += ' compressor="vtkZLibDataCompressor"'
|
||||
compression_level = min(max(compression_level, 0), 9)
|
||||
elif compression == 'lzma':
|
||||
string += ' compressor="vtkLZMADataCompressor"'
|
||||
compression_level = min(max(compression_level, 0), 9)
|
||||
elif compression == 'lz4':
|
||||
string += ' compressor="vtkLZ4DataCompressor"'
|
||||
compression_level = min(max(compression_level, 0), 12)
|
||||
lz4acceleration = max(lz4acceleration, 1)
|
||||
string += '>\n <ImageData WholeExtent=%s' % sdims
|
||||
string += 'Origin="%e %e %e" ' % origin
|
||||
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)
|
||||
else:
|
||||
string += ' <CellData %ss="%s">\n' % (dtype, vname)
|
||||
if nc == 1:
|
||||
if ncomp == 1:
|
||||
dmin = data.min()
|
||||
dmax = data.max()
|
||||
else:
|
||||
@ -154,8 +101,8 @@ def WriteVTK(vtkfile, vname, data, origin = (0.0, 0.0, 0.0), \
|
||||
dmin = numpy.sqrt(dd.min())
|
||||
dmax = numpy.sqrt(dd.max())
|
||||
string += ' <DataArray '
|
||||
if nc > 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 += " </DataArray>\n"
|
||||
if points:
|
||||
@ -166,8 +113,6 @@ def WriteVTK(vtkfile, vname, data, origin = (0.0, 0.0, 0.0), \
|
||||
string += ' </ImageData>\n'
|
||||
vt.write(string.encode())
|
||||
|
||||
# Append variable data.
|
||||
#
|
||||
if encode:
|
||||
string = ' <AppendedData encoding="base64">\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 </AppendedData>\n'
|
||||
string += '</VTKFile>\n'
|
||||
vt.write(string.encode())
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user