From fc5a846af36a30c1e3681712a22ba7c03881887c Mon Sep 17 00:00:00 2001 From: Grzegorz Kowal Date: Mon, 15 Mar 2021 15:54:26 -0300 Subject: [PATCH] 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 --- python/amunpy/src/amunpy/__init__.py | 3 +- python/amunpy/src/amunpy/amunh5.py | 370 +++++++++++++++++++++++++++ 2 files changed, 372 insertions(+), 1 deletion(-) diff --git a/python/amunpy/src/amunpy/__init__.py b/python/amunpy/src/amunpy/__init__.py index 4fdb14d..c3d36a1 100644 --- a/python/amunpy/src/amunpy/__init__.py +++ b/python/amunpy/src/amunpy/__init__.py @@ -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 " diff --git a/python/amunpy/src/amunpy/amunh5.py b/python/amunpy/src/amunpy/amunh5.py index 6c3abb2..2aef4be 100644 --- a/python/amunpy/src/amunpy/amunh5.py +++ b/python/amunpy/src/amunpy/amunh5.py @@ -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('\n') + vtk.write(' \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(' \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(' \n'.format(no, il, iu, jl, ju, kl, ku, vfile)) + no += 1 + else: + vtk.write(' \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(' \n') + vtk.write(' \n') + vtk.write('') + + if (progress): + sys.stdout.write('\n') + sys.stdout.flush()