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

Merge branch 'master' into reconnection

Browse Source
This commit is contained in:
Grzegorz Kowal 2021-10-05 11:01:17 -03:00
commit 346d263a83
8 changed files with 2437 additions and 2100 deletions

6
python/amunpy/setup.py

@ -5,7 +5,7 @@ with open("README.md", "r", encoding="utf-8") as fh:
setuptools.setup(
name="amunpy",
version="0.6.2",
version="0.9.1",
author="Grzegorz Kowal",
author_email="grzegorz@amuncode.org",
description="Python Interface for the AMUN code's snapshots",
@ -23,10 +23,8 @@ setuptools.setup(
package_dir={"": "src"},
packages=setuptools.find_packages(where="src"),
python_requires=">=3.6",
install_requires=['h5py', 'numpy'],
install_requires=['h5py', 'numpy', 'xxhash', 'lz4', 'zstandard'],
extras_require={
"digest": ['xxhash'],
"interpolation": ['scipy'],
"compression": ['lz4', 'zstandard'],
}
)

5
python/amunpy/src/amunpy/__init__.py

@ -12,10 +12,11 @@ See file LICENSE for more details.
from .amunxml import *
from .amunh5 import *
from .amunh5_deprecated import *
from .integrals import *
__all__ = [ 'AmunXML', 'amun_attribute', 'amun_coordinate', 'amun_dataset', \
'amun_dataset_vtk', 'amun_integrals' ]
__all__ = [ 'AmunXML', 'AmunH5', \
'amun_attribute', 'amun_coordinate', 'amun_dataset', 'amun_dataset_vtk', 'amun_integrals' ]
__author__ = "Grzegorz Kowal"
__copyright__ = "Copyright 2018-2021, Grzegorz Kowal <grzegorz@amuncode.org>"

1141
python/amunpy/src/amunpy/amun.py Normal file

File diff suppressed because it is too large Load Diff

859
python/amunpy/src/amunpy/amunh5.py

@ -24,798 +24,113 @@
module: AMUN
Python module with subroutines to read AMUN code HDF5 files.
The only requirements for this package are:
- h5py
- numpy
This module implements an interface class to read attributes, coordinates,
and datasets stored in AmunH5 format snapshots.
--------------------------------------------------------------------------------
"""
from .interpolation import interpolate
import h5py as h5
import numpy as np
import os.path as op
import sys
from .amun import Amun
class AmunH5(Amun):
"""AMUN H5 snapshot class"""
def amun_compatible(fname):
'''
Subroutine checks if the HDF5 file is AMUN compatible.
def __init_snapshot__(self):
"""
Sets the data format after verifying if the snapshot is stored
in AmunH5 format.
"""
import h5py
Arguments:
if not self.path_is_file:
raise Exception("AmunH5 requires a file not directory as the argument!")
fname - the HDF5 file name;
Return values:
True or False;
Examples:
comp = amun_compatible('p000010_00000.h5')
'''
with h5.File(fname, 'r') as f:
if 'codes' in f.attrs:
if f.attrs['code'].astype(str) == "AMUN":
return True
else:
print("'%s' contains attribute 'code'," % fname, \
" but it is not 'AMUN'!")
return False
elif 'attributes' in f and 'coordinates' in f and \
'variables' in f:
return True
else:
print("'%s' misses one of these groups:" % fname, \
"'attributes', 'coordinates' or 'variables'!")
return False
def amun_attribute(fname, aname):
'''
Subroutine to read global attributes from AMUN HDF5 snapshots.
Arguments:
fname - the HDF5 file name;
aname - the attribute name;
Return values:
ret - the value of the attribute or None;
Examples:
time = amun_attribute('p000010_00000.h5', 'time')
'''
if not amun_compatible(fname):
return None
with h5.File(fname, 'r') as f:
if aname in f['attributes'].attrs:
attr = f['attributes'].attrs[aname]
if attr.dtype.type is np.string_:
ret = np.squeeze(attr).astype(str)
else:
ret = np.squeeze(attr)
return ret
else:
print("Attribute '%s' cannot be found in '%s'!" % (aname, fname))
return None
def amun_coordinate(fname, iname):
'''
Subroutine to read coordinate items from AMUN HDF5 snapshots.
Arguments:
fname - the HDF5 file name;
iname - the item name;
Return values:
ret - the value of the item or None;
Examples:
bounds = amun_coordinate('p000010_00000.h5', 'bounds')
'''
if not amun_compatible(fname):
return None
with h5.File(fname, 'r') as f:
if iname in f['coordinates']:
return np.array(f['coordinates'][iname])
else:
print("Coordinate item '%s' not found in group 'coordinate' of '%s'!" % (iname, fname))
return None
def amun_dataset(fname, vname, shrink=1, interpolation='rebin', order=3, progress=False):
'''
Subroutine to read datasets from AMUN HDF5 snapshots.
Arguments:
fname - the HDF5 file name;
vname - the variable name;
shrink - the shrink factor (must be the power of 2 and not larger
than the block size);
progress - the progress bar switch;
Return values:
ret - the array of values for the variable;
Examples:
dn = amun_dataset('p000010_00000.h5', 'dens')
'''
if not amun_compatible(fname):
return None
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])
if ndims == 2:
bm[2] = 1
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
# 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
# check if the shrink parameter is correct (block dimensions should be
# divisible by the shrink factor)
#
shrink = max(1, int(shrink))
if shrink > 1:
if (nn % shrink) != 0:
print('The block dimension should be divisible by the shrink factor!')
return None
sh = shrink
while(sh > 2 and sh % 2 == 0):
sh = int(sh / 2)
if (sh % 2) != 0:
print('The shrink factor should be a power of 2!')
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())
# prepare dimensions of the output array and allocate it
#
dm = np.array(rm[0:ndims] * bm[0:ndims] * 2**(ml - 1) / shrink, \
dtype=np.int32)
ret = np.zeros(dm[::-1])
# 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')
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):
nn = 2**(ml - levels[l])
if nn <= shrink:
method = 'rebin'
else:
method = interpolation
cm = np.array(bm[0:ndims] * nn / shrink, dtype=np.int32)
ibeg = coords[0:ndims,l] * cm[0:ndims]
iend = ibeg + cm[0:ndims]
if ndims == 3:
ib, jb, kb = ibeg[0], ibeg[1], ibeg[2]
ie, je, ke = iend[0], iend[1], iend[2]
ret[kb:ke,jb:je,ib:ie] = interpolate(dataset[:,:,:,l], cm, ng, method=method, order=order)
else:
ib, jb = ibeg[0], ibeg[1]
ie, je = iend[0], iend[1]
ret[jb:je,ib:ie] = interpolate(dataset[0,:,:,l], cm, ng, method=method, order=order)
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()
return ret
def amun_dataset_vtk(fname, vname, label=None, compression='none', compression_level=19, 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)
with h5py.File(self.path, 'r') as h5:
if 'codes' in h5.attrs:
if h5.attrs['code'].astype(str) == "AMUN":
self.dataformat = 'AmunH5'
else:
dataset = g[vname][:,:,:,:]
raise Exception("'{}' contains attribute 'code', but its content is not 'AMUN'!".format(self.path))
elif 'attributes' in h5 and 'coordinates' in h5 and 'variables' in h5:
self.dataformat = 'AmunH5'
else:
raise Exception("{} misses one of these groups: 'attributes', 'coordinates' or 'variables'!".format(self.path))
# rescale all blocks to the effective resolution
#
for l in range(dblocks):
lv = levels[l] - 1
def __fill_attributes__(self):
"""
Reads attributes from the snapshot file and adds them to
the attributes' dictionary.
"""
import h5py
center = (bounds[0,:,l] + bounds[1,:,l]) / 2
base.createNodeBranch(center, lv)
base.setNodeData(center, lv, dataset[ng:-ng,ng:-ng,ng:-ng,l])
exclude_list = ['nseeds', 'seeds', 'dblocks', 'nproc', 'dims', 'dtn', 'last_id', 'mblocks']
nb += 1
with h5py.File(self.path, 'r') as h5:
for aname in h5['attributes'].attrs:
if not aname in exclude_list:
attr = h5['attributes'].attrs[aname]
if attr.dtype == 'float64' or attr.dtype == 'float32' or \
attr.dtype == 'int64' or attr.dtype == 'int32':
if len(attr) > 1:
self.attributes[aname] = attr.tolist()
else:
self.attributes[aname] = attr[0]
else:
self.attributes[aname] = attr[0].astype(str)
# 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 not 'nchunks' in self.attributes and 'nprocs' in self.attributes:
self.attributes['nchunks'] = self.attributes['nprocs']
del self.attributes['nprocs']
if 'rdims' in self.attributes:
self.attributes['xblocks'] = self.attributes['rdims'][0]
self.attributes['yblocks'] = self.attributes['rdims'][1]
if self.attributes['ndims'] == 3:
self.attributes['zblocks'] = self.attributes['rdims'][2]
del self.attributes['rdims']
if (progress):
sys.stdout.write('\n')
sys.stdout.flush()
if progress:
sys.stdout.write("Populating AMR structure\n")
base.populateNodeData()
def __fill_variables__(self):
"""
Reads the names of datasets stored in the snapshot and adds them
to the variables' dictionary.
"""
import h5py
if progress:
sys.stdout.write("Generating OverlappingAMR VTK files\n")
with h5py.File(self.path, 'r') as h5:
for variable in h5['variables']:
v = variable.strip()
self.variables[v] = v
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="{} {} {}" '.format(*base.lower) + \
'grid_description="XYZ">\n')
fmt = '{}_{:0' + str(len(str(ml))) + '}_{:0' + str(len(str(base.nodes))) + 'd}.vti'
def __fill_chunks__(self):
"""
Retrieves metadata about datablocks stored in the snapshot's chunks
and adds them to the chunks' dictionary.
"""
import h5py, numpy, os
m = 0
for lv in range(ml):
self.chunkname = 'p{:06d}'.format(self.attributes['isnap']) + '_{:05d}.h5'
for n in range(self.attributes['nchunks']):
self.chunks[n] = dict()
self.chunks[n]['filename'] = self.chunkname.format(n)
cname = os.path.join(self.dirname, self.chunks[n]['filename'])
if os.path.exists(cname):
with h5py.File(cname, 'r') as h5:
self.chunks[n]['dblocks'] = h5['attributes'].attrs['dblocks'][0]
cw = base.size / (rm * nn * 2**lv)
vtk.write(' <Block level="{}"'.format(lv) + \
' spacing="{} {} {}">\n'.format(*cw))
self.chunks[n]['levels'] = numpy.array(h5['coordinates']['levels'])
self.chunks[n]['bounds'] = numpy.array(h5['coordinates']['bounds'])
self.chunks[n]['coords'] = numpy.array(h5['coordinates']['coords'])
else:
raise Exception("Snapshot's chunk '{}' not present!".format(cname))
no = 0
for item in base.getNodesFromLevel(lv):
lo = np.array(item.index) * bm
up = lo + bm - 1
ll = np.stack((lo,up)).T.flatten()
if item.hasData:
vfile = op.join(opath, fmt.format(vname, lv, no))
WriteVTK(vfile, label, item.data, \
origin = (item.lower[0], item.lower[1], item.lower[2]), \
spacing = (cw[0], cw[1], cw[2]), \
compression=compression, compression_level=compression_level)
vtk.write(' <DataSet index="{}"'.format(no) + \
' amr_box = "{} {} {} {} {} {}"'.format(*ll) + \
' file = "{}"></DataSet>\n'.format(vfile))
no += 1
else:
vtk.write(' <DataSet index="{}"'.format(no) + \
' amr_box = "{} {} {} {} {} {}"'.format(*ll) + \
'></DataSet>\n')
m += 1
if progress:
sys.stdout.write('\r')
sys.stdout.write("Storing AMR block {} from {}".format(m, base.nodes))
sys.stdout.flush()
def __read_binary_data__(self, dataset_name, chunk_number):
"""
Gets the dataset array from a given snapshot's chunk.
"""
import h5py, numpy, os
vtk.write(' </Block>\n')
vtk.write(' </vtkOverlappingAMR>\n')
vtk.write('</VTKFile>')
if (progress):
sys.stdout.write('\n')
sys.stdout.flush()
cname = os.path.join(self.dirname, self.chunks[chunk_number]['filename'])
with h5py.File(cname, 'r') as h5:
return numpy.array(h5['variables'][dataset_name])

848
python/amunpy/src/amunpy/amunh5_deprecated.py Normal file

@ -0,0 +1,848 @@
"""
================================================================================
This file is part of the AMUN source code, a program to perform
Newtonian or relativistic magnetohydrodynamical simulations on uniform or
adaptive mesh.
Copyright (C) 2018-2021 Grzegorz Kowal <grzegorz@amuncode.org>
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
================================================================================
module: AMUN
Python module with subroutines to read AMUN code HDF5 files.
The only requirements for this package are:
- h5py
- numpy
--------------------------------------------------------------------------------
"""
#===============================================================================
'''
DEPRECATED FUNCTIONS
'''
def amun_compatible(fname):
'''
Subroutine checks if the HDF5 file is AMUN compatible.
Arguments:
fname - the HDF5 file name;
Return values:
True or False;
Examples:
comp = amun_compatible('p000010_00000.h5')
'''
from warnings import warn
import h5py as h5
warn('This function is deprecated', DeprecationWarning, stacklevel=2)
with h5.File(fname, 'r') as f:
if 'codes' in f.attrs:
if f.attrs['code'].astype(str) == "AMUN":
return True
else:
print("'%s' contains attribute 'code'," % fname, \
" but it is not 'AMUN'!")
return False
elif 'attributes' in f and 'coordinates' in f and \
'variables' in f:
return True
else:
print("'%s' misses one of these groups:" % fname, \
"'attributes', 'coordinates' or 'variables'!")
return False
def amun_attribute(fname, aname):
'''
Subroutine to read global attributes from AMUN HDF5 snapshots.
Arguments:
fname - the HDF5 file name;
aname - the attribute name;
Return values:
ret - the value of the attribute or None;
Examples:
time = amun_attribute('p000010_00000.h5', 'time')
'''
from warnings import warn
import h5py as h5
import numpy as np
warn('This function is deprecated', DeprecationWarning, stacklevel=2)
if not amun_compatible(fname):
return None
with h5.File(fname, 'r') as f:
if aname in f['attributes'].attrs:
attr = f['attributes'].attrs[aname]
if attr.dtype.type is np.string_:
ret = np.squeeze(attr).astype(str)
else:
ret = np.squeeze(attr)
return ret
else:
print("Attribute '%s' cannot be found in '%s'!" % (aname, fname))
return None
def amun_coordinate(fname, iname):
'''
Subroutine to read coordinate items from AMUN HDF5 snapshots.
Arguments:
fname - the HDF5 file name;
iname - the item name;
Return values:
ret - the value of the item or None;
Examples:
bounds = amun_coordinate('p000010_00000.h5', 'bounds')
'''
from warnings import warn
import h5py as h5
import numpy as np
warn('This function is deprecated', DeprecationWarning, stacklevel=2)
if not amun_compatible(fname):
return None
with h5.File(fname, 'r') as f:
if iname in f['coordinates']:
return np.array(f['coordinates'][iname])
else:
print("Coordinate item '%s' not found in group 'coordinate' of '%s'!" % (iname, fname))
return None
def amun_dataset(fname, vname, shrink=1, interpolation='rebin', order=3, progress=False):
'''
Subroutine to read datasets from AMUN HDF5 snapshots.
Arguments:
fname - the HDF5 file name;
vname - the variable name;
shrink - the shrink factor (must be the power of 2 and not larger
than the block size);
progress - the progress bar switch;
Return values:
ret - the array of values for the variable;
Examples:
dn = amun_dataset('p000010_00000.h5', 'dens')
'''
from .interpolation import interpolate
from warnings import warn
import h5py as h5
import numpy as np
import os, sys
warn('This function is deprecated', DeprecationWarning, stacklevel=2)
if not amun_compatible(fname):
return None
dname = os.path.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])
if ndims == 2:
bm[2] = 1
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
# 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
# check if the shrink parameter is correct (block dimensions should be
# divisible by the shrink factor)
#
shrink = max(1, int(shrink))
if shrink > 1:
if (nn % shrink) != 0:
print('The block dimension should be divisible by the shrink factor!')
return None
sh = shrink
while(sh > 2 and sh % 2 == 0):
sh = int(sh / 2)
if (sh % 2) != 0:
print('The shrink factor should be a power of 2!')
return None
# determine the actual maximum level from the blocks
#
levs = []
for n in range(nc):
fname = 'p%06d_%05d.h5' % (nr, n)
lname = os.path.join(dname, fname)
dblocks = amun_attribute(lname, 'dblocks')
if dblocks > 0:
levs = np.append(levs, [amun_coordinate(lname, 'levels')])
ml = int(levs.max())
# prepare dimensions of the output array and allocate it
#
dm = np.array(rm[0:ndims] * bm[0:ndims] * 2**(ml - 1) / shrink, \
dtype=np.int32)
ret = np.zeros(dm[::-1])
# iterate over all subdomain files
#
nb = 0
for n in range(nc):
fname = 'p%06d_%05d.h5' % (nr, n)
lname = os.path.join(dname, fname)
dblocks = amun_attribute(lname, 'dblocks')
if dblocks > 0:
levels = amun_coordinate(lname, 'levels')
coords = amun_coordinate(lname, 'coords')
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):
nn = 2**(ml - levels[l])
if nn <= shrink:
method = 'rebin'
else:
method = interpolation
cm = np.array(bm[0:ndims] * nn / shrink, dtype=np.int32)
ibeg = coords[0:ndims,l] * cm[0:ndims]
iend = ibeg + cm[0:ndims]
if ndims == 3:
ib, jb, kb = ibeg[0], ibeg[1], ibeg[2]
ie, je, ke = iend[0], iend[1], iend[2]
ret[kb:ke,jb:je,ib:ie] = interpolate(dataset[:,:,:,l], cm, ng, method=method, order=order)
else:
ib, jb = ibeg[0], ibeg[1]
ie, je = iend[0], iend[1]
ret[jb:je,ib:ie] = interpolate(dataset[0,:,:,l], cm, ng, method=method, order=order)
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()
return ret
def amun_dataset_vtk(fname, vname, label=None, compression=None, compression_level=19, 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
from warnings import warn
import numpy as np
import os, sys
warn('This function is deprecated', DeprecationWarning, stacklevel=2)
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 = os.path.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 = os.path.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 = os.path.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.createNodeBranch(center, lv)
base.setNodeData(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.populateNodeData()
if progress:
sys.stdout.write("Generating OverlappingAMR VTK files\n")
ofile = "{}_{:06d}.vthb".format(vname, nr)
opath = "{}_{:06d}".format(vname, nr)
if not os.path.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="{} {} {}" '.format(*base.lower) + \
'grid_description="XYZ">\n')
fmt = '{}_{:0' + str(len(str(ml))) + '}_{:0' + str(len(str(base.nodes))) + 'd}.vti'
m = 0
for lv in range(ml):
cw = base.size / (rm * nn * 2**lv)
vtk.write(' <Block level="{}"'.format(lv) + \
' spacing="{} {} {}">\n'.format(*cw))
no = 0
for item in base.getNodesFromLevel(lv):
lo = np.array(item.index) * bm
up = lo + bm - 1
ll = np.stack((lo,up)).T.flatten()
if item.hasData:
vfile = os.path.join(opath, fmt.format(vname, lv, no))
WriteVTK(vfile, label, item.data, \
origin = (item.lower[0], item.lower[1], item.lower[2]), \
spacing = (cw[0], cw[1], cw[2]), \
compression=compression, compression_level=compression_level)
vtk.write(' <DataSet index="{}"'.format(no) + \
' amr_box = "{} {} {} {} {} {}"'.format(*ll) + \
' file = "{}"></DataSet>\n'.format(vfile))
no += 1
else:
vtk.write(' <DataSet index="{}"'.format(no) + \
' amr_box = "{} {} {} {} {} {}"'.format(*ll) + \
'></DataSet>\n')
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()

1273
python/amunpy/src/amunpy/amunxml.py

File diff suppressed because it is too large Load Diff

255
python/amunpy/src/amunpy/interpolation.py

@ -22,14 +22,10 @@
================================================================================
module: AMUN
module: INTERPOLATION
Python module with subroutines to read AMUN code HDF5 files.
The only requirements for this package are:
- numpy
- scipy (optional, for data interpolation)
Support module for Amun snapshots for the block prolongation with different
types of interpolation.
--------------------------------------------------------------------------------
"""
@ -43,154 +39,129 @@ except ImportError:
def rebin(a, newshape):
'''
Subroutine changes the size of the input array to to new shape,
by copying cells or averaging them.
'''
assert len(a.shape) == len(newshape)
'''
Subroutine changes the size of the input array to to new shape,
by copying cells or averaging them.
'''
assert len(a.shape) == len(newshape)
m = a.ndim - 1
if (a.shape[m] > newshape[m]):
if a.ndim == 3:
nn = [newshape[0], int(a.shape[0] / newshape[0]),
newshape[1], int(a.shape[1] / newshape[1]),
newshape[2], int(a.shape[2] / newshape[2])]
return a.reshape(nn).mean(5).mean(3).mean(1)
m = a.ndim - 1
if (a.shape[m] > newshape[m]):
if a.ndim == 3:
nn = [newshape[0], a.shape[0] // newshape[0], \
newshape[1], a.shape[1] // newshape[1], \
newshape[2], a.shape[2] // newshape[2]]
return a.reshape(nn).mean(5).mean(3).mean(1)
else:
nn = [newshape[0], a.shape[0] // newshape[0], \
newshape[1], a.shape[1] // newshape[1]]
return a.reshape(nn).mean(3).mean(1)
else:
nn = [newshape[0], int(a.shape[0] / newshape[0]),
newshape[1], int(a.shape[1] / newshape[1])]
return a.reshape(nn).mean(3).mean(1)
else:
for n in range(a.ndim):
a = np.repeat(a, int(newshape[n] / a.shape[n]), axis=n)
return(a)
for n in range(a.ndim):
a = np.repeat(a, newshape[n] // a.shape[n], axis=n)
return(a)
def interpolate(a, newshape, nghosts, method='rebin', order=3):
'''
Subroutine rescales the block by interpolating its values.
'''
if (method == 'rebin' or not scipy_available):
def interpolate(a, newshape, nghosts=0, method=None, order=1):
'''
Subroutine rescales the block by interpolating its values.
'''
if method == None or method == 'rebin' or not scipy_available:
# calculate the indices in order to remove the ghost zones
#
if a.ndim == 3:
ib = nghosts
jb = nghosts
kb = nghosts
ie = a.shape[2] - nghosts
je = a.shape[1] - nghosts
ke = a.shape[0] - nghosts
if (a.shape[0] == 1):
kb = 0
ke = 1
ng = nghosts
if a.ndim == 3:
return rebin(a[ng:-ng,ng:-ng,ng:-ng], newshape)
else:
return rebin(a[ng:-ng,ng:-ng], newshape)
elif method == 'zoom':
zf = (newshape[1] // (a.shape[1] - 2 * nghosts))
ng = zf * nghosts
if a.ndim == 3:
return zoom(a, zf, order=order, grid_mode=True, mode='nearest')[ng:-ng,ng:-ng,ng:-ng]
else:
return zoom(a, zf, order=order, grid_mode=True, mode='nearest')[ng:-ng,ng:-ng]
elif method in [ 'monotonic', 'pchip' ]:
dims = np.arange(a.ndim)
q = a
for n in dims:
d2 = np.roll(q,-1, axis=0) + np.roll(q, 1, axis=0) - 2.0 * q
q = q - d2 / 24.0
d = np.array(q.shape)
xo = (np.arange(0.5, a.shape[n]) - nghosts) / (a.shape[n] - 2 * nghosts)
xn = np.arange(0.5, newshape[n]) / newshape[n]
u = q.reshape([d[0], q.size // d[0]])
f = np.zeros([newshape[n], q.size // d[0]])
for i in range(q.size // d[0]):
f[:,i] = pchip_interpolate(xo, u[:,i], xn)
d[0] = newshape[n]
f = f.reshape(d)
q = f.transpose(np.roll(dims, -1))
return q
elif method == 'spline':
dims = np.arange(a.ndim)
q = a
for n in dims:
d2 = np.roll(q,-1, axis=0) + np.roll(q, 1, axis=0) - 2.0 * q
q = q - d2 / 24.0
d = np.array(q.shape)
xo = (np.arange(0.5, a.shape[n]) - nghosts) / (a.shape[n] - 2 * nghosts)
xn = np.arange(0.5, newshape[n]) / newshape[n]
u = q.reshape([d[0], q.size // d[0]])
f = np.zeros([newshape[n], q.size // d[0]])
for i in range(q.size // d[0]):
t = splrep(xo, u[:,i], k=5, s=0.0)
f[:,i] = splev(xn, t)
d[0] = newshape[n]
f = f.reshape(d)
q = f.transpose(np.roll(dims, -1))
return q
return rebin(a[kb:ke,jb:je,ib:ie], newshape)
else:
ib = nghosts
jb = nghosts
ie = a.shape[1] - nghosts
je = a.shape[0] - nghosts
return rebin(a[jb:je,ib:ie], newshape)
dims = np.arange(a.ndim)
q = a
elif (method == 'zoom'):
for n in dims:
fc = int(newshape[1] / (a.shape[1] - 2 * nghosts))
ib, ie = fc * nghosts, - fc * nghosts
jb, je = fc * nghosts, - fc * nghosts
if a.ndim == 3:
kb, ke = fc * nghosts, - fc * nghosts
if (a.shape[0] == 1):
kb = 0
ke = 1
d2 = np.roll(q,-1, axis=0) + np.roll(q, 1, axis=0) - 2.0 * q
q = q - d2 / 24.0
return zoom(a, fc, order=order, grid_mode=True, mode='nearest')[kb:ke,jb:je,ib:ie]
else:
return zoom(a, fc, order=order, grid_mode=True, mode='nearest')[jb:je,ib:ie]
d = np.array(q.shape)
elif (method == 'monotonic' or method == 'pchip'):
xo = (np.arange(0.5, a.shape[n]) - nghosts) / (a.shape[n] - 2 * nghosts)
xn = np.arange(0.5, newshape[n]) / newshape[n]
dims = np.arange(a.ndim)
u = q.reshape([d[0], q.size // d[0]])
f = np.zeros([newshape[n], q.size // d[0]])
for i in range(q.size // d[0]):
t = interp1d(xo, u[:,i], kind=method)
f[:,i] = t(xn)
q = a
d[0] = newshape[n]
f = f.reshape(d)
for n in dims:
q = f.transpose(np.roll(dims, -1))
d2 = np.roll(q,-1, axis=0) + np.roll(q, 1, axis=0) - 2.0 * q
q = q - d2 / 24.0
d = np.array(q.shape)
xo = (np.arange(0.5, a.shape[n]) - nghosts) / (a.shape[n] - 2 * nghosts)
xn = np.arange(0.5, newshape[n]) / newshape[n]
u = q.reshape([d[0], int(q.size / d[0])])
f = np.zeros([newshape[n], int(q.size / d[0])])
for i in range(int(q.size / d[0])):
f[:,i] = pchip_interpolate(xo, u[:,i], xn)
d[0] = newshape[n]
f = f.reshape(d)
q = f.transpose(np.roll(dims, -1))
return q
elif (method == 'spline'):
dims = np.arange(a.ndim)
q = a
for n in dims:
d2 = np.roll(q,-1, axis=0) + np.roll(q, 1, axis=0) - 2.0 * q
q = q - d2 / 24.0
d = np.array(q.shape)
xo = (np.arange(0.5, a.shape[n]) - nghosts) / (a.shape[n] - 2 * nghosts)
xn = np.arange(0.5, newshape[n]) / newshape[n]
u = q.reshape([d[0], int(q.size / d[0])])
f = np.zeros([newshape[n], int(q.size / d[0])])
for i in range(int(q.size / d[0])):
t = splrep(xo, u[:,i], k=5, s=0.0)
f[:,i] = splev(xn, t)
d[0] = newshape[n]
f = f.reshape(d)
q = f.transpose(np.roll(dims, -1))
return q
else:
dims = np.arange(a.ndim)
q = a
for n in dims:
d2 = np.roll(q,-1, axis=0) + np.roll(q, 1, axis=0) - 2.0 * q
q = q - d2 / 24.0
d = np.array(q.shape)
xo = (np.arange(0.5, a.shape[n]) - nghosts) / (a.shape[n] - 2 * nghosts)
xn = np.arange(0.5, newshape[n]) / newshape[n]
u = q.reshape([d[0], int(q.size / d[0])])
f = np.zeros([newshape[n], int(q.size / d[0])])
for i in range(int(q.size / d[0])):
t = interp1d(xo, u[:,i], kind=method)
f[:,i] = t(xn)
d[0] = newshape[n]
f = f.reshape(d)
q = f.transpose(np.roll(dims, -1))
return q
return q

150
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 = '<?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 compression != None:
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())