PYTHON: Add amun_dataset_vtk() function.
This function works only with HDF5 files and its purpose is to convert the dataset and store it as the OverlappedAMR VTK files, which can be read by ParaView. Signed-off-by: Grzegorz Kowal <grzegorz@amuncode.org>
This commit is contained in:
parent
a2c3666ef2
commit
fc5a846af3
@ -14,7 +14,8 @@ from .amunxml import *
|
||||
from .amunh5 import *
|
||||
from .integrals import *
|
||||
|
||||
__all__ = [ 'AmunXML', 'amun_attribute', 'amun_coordinate', 'amun_dataset', 'amun_integrals' ]
|
||||
__all__ = [ 'AmunXML', 'amun_attribute', 'amun_coordinate', 'amun_dataset', \
|
||||
'amun_dataset_vtk', 'amun_integrals' ]
|
||||
|
||||
__author__ = "Grzegorz Kowal"
|
||||
__copyright__ = "Copyright 2018-2021, Grzegorz Kowal <grzegorz@amuncode.org>"
|
||||
|
||||
@ -453,3 +453,373 @@ def amun_dataset(fname, vname, shrink=1, interpolation='rebin', progress=False):
|
||||
sys.stdout.flush()
|
||||
|
||||
return ret
|
||||
|
||||
|
||||
def amun_dataset_vtk(fname, vname, label=None, compression='none', progress=False):
|
||||
'''
|
||||
Subroutine to convert a dataset specified by argument 'vname' from
|
||||
the AMUN HDF5 snapshot to OverlappedAMR VTK file.
|
||||
|
||||
Arguments:
|
||||
|
||||
fname - the HDF5 file name;
|
||||
vname - the variable name;
|
||||
label - the variable label (long name);
|
||||
compression - the compression type: 'lz4', 'zlib', 'lzma'
|
||||
progress - the progress bar switch;
|
||||
|
||||
Examples:
|
||||
|
||||
dn = amun_dataset_vtk('p000010_00000.h5', 'dens')
|
||||
|
||||
'''
|
||||
from .octree import OcBase, OcNode
|
||||
from .vtkio import WriteVTK
|
||||
import os
|
||||
|
||||
if not amun_compatible(fname):
|
||||
return None
|
||||
|
||||
if amun_attribute(fname, 'ndims') < 3:
|
||||
print('Subroutine amun_dataset_vtk() supports only 3D domains.')
|
||||
return None
|
||||
|
||||
if label == None:
|
||||
label = vname
|
||||
|
||||
dname = op.dirname(fname)
|
||||
|
||||
if progress:
|
||||
sys.stdout.write("Data file path:\n '%s'\n" % (dname))
|
||||
|
||||
# get attributes necessary to reconstruct the domain
|
||||
#
|
||||
eqsys = amun_attribute(fname, 'eqsys')
|
||||
eos = amun_attribute(fname, 'eos')
|
||||
nr = amun_attribute(fname, 'isnap')
|
||||
nc = amun_attribute(fname, 'nprocs')
|
||||
nl = amun_attribute(fname, 'nleafs')
|
||||
if eos == 'adi':
|
||||
gm = amun_attribute(fname, 'adiabatic_index')
|
||||
|
||||
# get block dimensions and the maximum level
|
||||
#
|
||||
ndims = amun_attribute(fname, 'ndims')
|
||||
nn = amun_attribute(fname, 'ncells')
|
||||
bm = np.array([nn, nn, nn])
|
||||
ng = amun_attribute(fname, 'nghosts')
|
||||
ml = amun_attribute(fname, 'maxlev')
|
||||
|
||||
# get the base block dimensions
|
||||
#
|
||||
rm = amun_attribute(fname, 'bdims')
|
||||
if rm is None:
|
||||
rm = amun_attribute(fname, 'domain_base_dims')
|
||||
if rm is None:
|
||||
rm = amun_attribute(fname, 'rdims')
|
||||
if rm is None:
|
||||
return None
|
||||
|
||||
# get domain bounds
|
||||
#
|
||||
xmin = amun_attribute(fname, 'xmin')
|
||||
ymin = amun_attribute(fname, 'ymin')
|
||||
zmin = amun_attribute(fname, 'zmin')
|
||||
xlen = amun_attribute(fname, 'xmax') - xmin
|
||||
ylen = amun_attribute(fname, 'ymax') - ymin
|
||||
zlen = amun_attribute(fname, 'zmax') - zmin
|
||||
|
||||
# build the list of supported variables
|
||||
#
|
||||
variables = []
|
||||
with h5.File(fname, 'r') as f:
|
||||
for var in f['variables'].keys():
|
||||
variables.append(var)
|
||||
|
||||
# add derived variables if possible
|
||||
#
|
||||
variables.append('level')
|
||||
if 'velx' in variables and 'vely' in variables and 'velz' in variables:
|
||||
variables.append('velo')
|
||||
variables.append('divv')
|
||||
variables.append('vort')
|
||||
if 'magx' in variables and 'magy' in variables and 'magz' in variables:
|
||||
variables.append('magn')
|
||||
variables.append('divb')
|
||||
variables.append('curr')
|
||||
if (eqsys == 'hd' or eqsys == 'mhd') and eos == 'adi' \
|
||||
and 'pres' in variables:
|
||||
variables.append('eint')
|
||||
if 'dens' in variables and 'pres' in variables:
|
||||
variables.append('temp')
|
||||
if (eqsys == 'hd' or eqsys == 'mhd') \
|
||||
and 'dens' in variables \
|
||||
and 'velx' in variables \
|
||||
and 'vely' in variables \
|
||||
and 'velz' in variables:
|
||||
variables.append('ekin')
|
||||
if (eqsys == 'mhd' or eqsys == 'srmhd') \
|
||||
and 'magx' in variables \
|
||||
and 'magy' in variables \
|
||||
and 'magz' in variables:
|
||||
variables.append('emag')
|
||||
if eqsys == 'hd' and 'ekin' in variables and 'eint' in variables:
|
||||
variables.append('etot')
|
||||
if eqsys == 'mhd' and 'eint' in variables \
|
||||
and 'ekin' in variables \
|
||||
and 'emag' in variables:
|
||||
variables.append('etot')
|
||||
if (eqsys == 'srhd' or eqsys == 'srmhd') and 'velo' in variables:
|
||||
variables.append('lore')
|
||||
|
||||
# check if the requested variable is in the variable list
|
||||
#
|
||||
if not vname in variables:
|
||||
print('The requested variable cannot be extracted from the file datasets!')
|
||||
return None
|
||||
|
||||
# determine the actual maximum level from the blocks
|
||||
#
|
||||
levs = []
|
||||
for n in range(nc):
|
||||
fname = 'p%06d_%05d.h5' % (nr, n)
|
||||
lname = op.join(dname, fname)
|
||||
dblocks = amun_attribute(lname, 'dblocks')
|
||||
if dblocks > 0:
|
||||
levs = np.append(levs, [amun_coordinate(lname, 'levels')])
|
||||
ml = int(levs.max())
|
||||
|
||||
# create octree base
|
||||
base = OcBase([xmin, ymin, zmin], [xlen, ylen, zlen], rm)
|
||||
|
||||
# iterate over all subdomain files
|
||||
#
|
||||
nb = 0
|
||||
for n in range(nc):
|
||||
fname = 'p%06d_%05d.h5' % (nr, n)
|
||||
lname = op.join(dname, fname)
|
||||
dblocks = amun_attribute(lname, 'dblocks')
|
||||
if dblocks > 0:
|
||||
levels = amun_coordinate(lname, 'levels')
|
||||
coords = amun_coordinate(lname, 'coords')
|
||||
bounds = amun_coordinate(lname, 'bounds')
|
||||
dx = amun_coordinate(lname, 'dx')
|
||||
dy = amun_coordinate(lname, 'dy')
|
||||
dz = amun_coordinate(lname, 'dz')
|
||||
with h5.File(lname, 'r') as f:
|
||||
g = f['variables']
|
||||
if vname == 'level':
|
||||
dataset = np.zeros(g[variables[0]].shape)
|
||||
for l in range(dblocks):
|
||||
dataset[:,:,:,l] = levels[l]
|
||||
elif vname == 'velo':
|
||||
dataset = np.sqrt(g['velx'][:,:,:,:]**2 \
|
||||
+ g['vely'][:,:,:,:]**2 \
|
||||
+ g['velz'][:,:,:,:]**2)
|
||||
elif vname == 'magn':
|
||||
dataset = np.sqrt(g['magx'][:,:,:,:]**2 \
|
||||
+ g['magy'][:,:,:,:]**2 \
|
||||
+ g['magz'][:,:,:,:]**2)
|
||||
elif vname == 'eint':
|
||||
dataset = 1.0 / (gm - 1.0) * g['pres'][:,:,:,:]
|
||||
elif vname == 'ekin':
|
||||
dataset = 0.5 * g['dens'][:,:,:,:] * (g['velx'][:,:,:,:]**2 \
|
||||
+ g['vely'][:,:,:,:]**2 \
|
||||
+ g['velz'][:,:,:,:]**2)
|
||||
elif vname == 'emag':
|
||||
dataset = 0.5 * (g['magx'][:,:,:,:]**2 \
|
||||
+ g['magy'][:,:,:,:]**2 \
|
||||
+ g['magz'][:,:,:,:]**2)
|
||||
elif vname == 'etot':
|
||||
dataset = 1.0 / (gm - 1.0) * g['pres'][:,:,:,:] \
|
||||
+ 0.5 * g['dens'][:,:,:,:] * (g['velx'][:,:,:,:]**2 \
|
||||
+ g['vely'][:,:,:,:]**2 \
|
||||
+ g['velz'][:,:,:,:]**2)
|
||||
if eqsys == 'mhd':
|
||||
dataset += 0.5 * (g['magx'][:,:,:,:]**2 \
|
||||
+ g['magy'][:,:,:,:]**2 \
|
||||
+ g['magz'][:,:,:,:]**2)
|
||||
elif vname == 'temp':
|
||||
dataset = g['pres'][:,:,:,:] / g['dens'][:,:,:,:]
|
||||
elif vname == 'lore':
|
||||
dataset = 1.0 / np.sqrt(1.0 - (g['velx'][:,:,:,:]**2 \
|
||||
+ g['vely'][:,:,:,:]**2 \
|
||||
+ g['velz'][:,:,:,:]**2))
|
||||
elif vname == 'divv':
|
||||
dataset = np.zeros(g['velx'].shape)
|
||||
fields = [ 'velx', 'vely', 'velz' ]
|
||||
h = (dx, dy, dz)
|
||||
for i in range(ndims):
|
||||
v = fields[i]
|
||||
dataset += 0.5 * (np.roll(g[v][:,:,:,:], -1, axis=2) \
|
||||
- np.roll(g[v][:,:,:,:], 1, axis=2)) \
|
||||
/ h[i][levels[:] - 1]
|
||||
elif vname == 'divb':
|
||||
dataset = np.zeros(g['magx'].shape)
|
||||
fields = [ 'magx', 'magy', 'magz' ]
|
||||
h = (dx, dy, dz)
|
||||
for i in range(ndims):
|
||||
v = fields[i]
|
||||
dataset += 0.5 * (np.roll(g[v][:,:,:,:], -1, axis=2) \
|
||||
- np.roll(g[v][:,:,:,:], 1, axis=2)) \
|
||||
/ h[i][levels[:] - 1]
|
||||
elif vname == 'vort':
|
||||
if ndims == 3:
|
||||
wx = 0.5 * (np.roll(g['velz'][:,:,:,:], -1, axis=1) \
|
||||
- np.roll(g['velz'][:,:,:,:], 1, axis=1)) \
|
||||
/ dy[levels[:]-1] \
|
||||
- 0.5 * (np.roll(g['vely'][:,:,:,:], -1, axis=0) \
|
||||
- np.roll(g['vely'][:,:,:,:], 1, axis=0)) \
|
||||
/ dz[levels[:]-1]
|
||||
wy = 0.5 * (np.roll(g['velx'][:,:,:,:], -1, axis=0) \
|
||||
- np.roll(g['velx'][:,:,:,:], 1, axis=0)) \
|
||||
/ dz[levels[:]-1] \
|
||||
- 0.5 * (np.roll(g['velz'][:,:,:,:], -1, axis=2) \
|
||||
- np.roll(g['velz'][:,:,:,:], 1, axis=2)) \
|
||||
/ dx[levels[:]-1]
|
||||
wz = 0.5 * (np.roll(g['vely'][:,:,:,:], -1, axis=2) \
|
||||
- np.roll(g['vely'][:,:,:,:], 1, axis=2)) \
|
||||
/ dx[levels[:]-1] \
|
||||
- 0.5 * (np.roll(g['velx'][:,:,:,:], -1, axis=1) \
|
||||
- np.roll(g['velx'][:,:,:,:], 1, axis=1)) \
|
||||
/ dy[levels[:]-1]
|
||||
else:
|
||||
wx = 0.5 * (np.roll(g['velz'][:,:,:,:], -1, axis=1) \
|
||||
- np.roll(g['velz'][:,:,:,:], 1, axis=1)) \
|
||||
/ dy[levels[:]-1]
|
||||
wy = - 0.5 * (np.roll(g['velz'][:,:,:,:], -1, axis=2) \
|
||||
- np.roll(g['velz'][:,:,:,:], 1, axis=2)) \
|
||||
/ dx[levels[:]-1]
|
||||
wz = 0.5 * (np.roll(g['vely'][:,:,:,:], -1, axis=2) \
|
||||
- np.roll(g['vely'][:,:,:,:], 1, axis=2)) \
|
||||
/ dx[levels[:]-1] \
|
||||
- 0.5 * (np.roll(g['velx'][:,:,:,:], -1, axis=1) \
|
||||
- np.roll(g['velx'][:,:,:,:], 1, axis=1)) \
|
||||
/ dy[levels[:]-1]
|
||||
dataset = np.sqrt(wx * wx + wy * wy + wz * wz)
|
||||
elif vname == 'curr':
|
||||
if ndims == 3:
|
||||
wx = 0.5 * (np.roll(g['magz'][:,:,:,:], -1, axis=1) \
|
||||
- np.roll(g['magz'][:,:,:,:], 1, axis=1)) \
|
||||
/ dy[levels[:]-1] \
|
||||
- 0.5 * (np.roll(g['magy'][:,:,:,:], -1, axis=0) \
|
||||
- np.roll(g['magy'][:,:,:,:], 1, axis=0)) \
|
||||
/ dz[levels[:]-1]
|
||||
wy = 0.5 * (np.roll(g['magx'][:,:,:,:], -1, axis=0) \
|
||||
- np.roll(g['magx'][:,:,:,:], 1, axis=0)) \
|
||||
/ dz[levels[:]-1] \
|
||||
- 0.5 * (np.roll(g['magz'][:,:,:,:], -1, axis=2) \
|
||||
- np.roll(g['magz'][:,:,:,:], 1, axis=2)) \
|
||||
/ dx[levels[:]-1]
|
||||
wz = 0.5 * (np.roll(g['magy'][:,:,:,:], -1, axis=2) \
|
||||
- np.roll(g['magy'][:,:,:,:], 1, axis=2)) \
|
||||
/ dx[levels[:]-1] \
|
||||
- 0.5 * (np.roll(g['magx'][:,:,:,:], -1, axis=1) \
|
||||
- np.roll(g['magx'][:,:,:,:], 1, axis=1)) \
|
||||
/ dy[levels[:]-1]
|
||||
else:
|
||||
wx = 0.5 * (np.roll(g['magz'][:,:,:,:], -1, axis=1) \
|
||||
- np.roll(g['magz'][:,:,:,:], 1, axis=1)) \
|
||||
/ dy[levels[:]-1]
|
||||
wy = - 0.5 * (np.roll(g['magz'][:,:,:,:], -1, axis=2) \
|
||||
- np.roll(g['magz'][:,:,:,:], 1, axis=2)) \
|
||||
/ dx[levels[:]-1]
|
||||
wz = 0.5 * (np.roll(g['magy'][:,:,:,:], -1, axis=2) \
|
||||
- np.roll(g['magy'][:,:,:,:], 1, axis=2)) \
|
||||
/ dx[levels[:]-1] \
|
||||
- 0.5 * (np.roll(g['magx'][:,:,:,:], -1, axis=1) \
|
||||
- np.roll(g['magx'][:,:,:,:], 1, axis=1)) \
|
||||
/ dy[levels[:]-1]
|
||||
dataset = np.sqrt(wx * wx + wy * wy + wz * wz)
|
||||
else:
|
||||
dataset = g[vname][:,:,:,:]
|
||||
|
||||
# rescale all blocks to the effective resolution
|
||||
#
|
||||
for l in range(dblocks):
|
||||
|
||||
lv = levels[l] - 1
|
||||
|
||||
center = (bounds[0,:,l] + bounds[1,:,l]) / 2
|
||||
base.createBranch(center, lv)
|
||||
base.setLeafData(center, lv, dataset[ng:-ng,ng:-ng,ng:-ng,l])
|
||||
|
||||
nb += 1
|
||||
|
||||
# print progress bar if desired
|
||||
#
|
||||
if progress:
|
||||
sys.stdout.write('\r')
|
||||
sys.stdout.write("Reading '%s' from '%s': block %d from %d" \
|
||||
% (vname, fname, nb, nl))
|
||||
sys.stdout.flush()
|
||||
|
||||
if (progress):
|
||||
sys.stdout.write('\n')
|
||||
sys.stdout.flush()
|
||||
|
||||
if progress:
|
||||
sys.stdout.write("Populating AMR structure\n")
|
||||
base.populateBranches()
|
||||
|
||||
if progress:
|
||||
sys.stdout.write("Generating OverlappedAMR VTK files\n")
|
||||
|
||||
ofile = "{}_{:06d}.vthb".format(vname, nr)
|
||||
opath = "{}_{:06d}".format(vname, nr)
|
||||
if not op.exists(opath):
|
||||
os.makedirs(opath)
|
||||
with open(ofile, 'w') as vtk:
|
||||
vtk.write('<VTKFile type="vtkOverlappingAMR" version="1.1" byte_order="LittleEndian" header_type="UInt64">\n')
|
||||
vtk.write(' <vtkOverlappingAMR origin="{} {} {}" grid_description="XYZ">\n'.format(*base.lower))
|
||||
|
||||
m = 0
|
||||
for lv in range(ml):
|
||||
|
||||
sm = np.array(rm) * 2**lv
|
||||
dh = np.array([xlen, ylen, zlen]) / np.array(sm * bm)
|
||||
vtk.write(' <Block level="{}" spacing="{} {} {}">\n'.format(lv, dh[0], dh[1], dh[2]))
|
||||
dxb = base.size[0] / sm[0]
|
||||
dyb = base.size[1] / sm[1]
|
||||
dzb = base.size[2] / sm[2]
|
||||
no = 0
|
||||
for i in range(sm[0]):
|
||||
x = base.lower[0] + (i + 0.5) * dxb
|
||||
for j in range(sm[1]):
|
||||
y = base.lower[1] + (j + 0.5) * dyb
|
||||
for k in range(sm[2]):
|
||||
z = base.lower[2] + (k + 0.5) * dzb
|
||||
|
||||
item = base.getItem([x, y, z], lv)
|
||||
|
||||
if isinstance(item, OcNode):
|
||||
il = i * bm[0]
|
||||
jl = j * bm[1]
|
||||
kl = k * bm[2]
|
||||
iu = il + bm[0] - 1
|
||||
ju = jl + bm[1] - 1
|
||||
ku = kl + bm[2] - 1
|
||||
if item.hasData:
|
||||
vfile = op.join(opath, '{}_{}_{}.vti'.format(var, lv, no))
|
||||
WriteVTK(vfile, label, item.data, \
|
||||
origin = (item.lower[0], item.lower[1], item.lower[2]), \
|
||||
spacing = (dh[0], dh[1], dh[2]), compression=compression)
|
||||
vtk.write(' <DataSet index="{}" amr_box = "{} {} {} {} {} {}" file = "{}"></DataSet>\n'.format(no, il, iu, jl, ju, kl, ku, vfile))
|
||||
no += 1
|
||||
else:
|
||||
vtk.write(' <DataSet index="{}" amr_box = "{} {} {} {} {} {}"></DataSet>\n'.format(no, il, iu, jl, ju, kl, ku))
|
||||
|
||||
m += 1
|
||||
|
||||
if progress:
|
||||
sys.stdout.write('\r')
|
||||
sys.stdout.write("Storing AMR block {} from {}".format(m, base.nodes))
|
||||
sys.stdout.flush()
|
||||
|
||||
vtk.write(' </Block>\n')
|
||||
vtk.write(' </vtkOverlappingAMR>\n')
|
||||
vtk.write('</VTKFile>')
|
||||
|
||||
if (progress):
|
||||
sys.stdout.write('\n')
|
||||
sys.stdout.flush()
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user