Merge branch 'master' into reconnection
This commit is contained in:
commit
5055e05001
252
python/amun.py
252
python/amun.py
@ -35,6 +35,7 @@
|
||||
import h5py as h5
|
||||
import numpy as np
|
||||
import os.path as op
|
||||
import sys
|
||||
|
||||
def amun_attribute(fname, aname):
|
||||
'''
|
||||
@ -86,7 +87,7 @@ def amun_attribute(fname, aname):
|
||||
#
|
||||
return ret
|
||||
|
||||
def amun_dataset(fname, vname):
|
||||
def amun_dataset(fname, vname, progress = False):
|
||||
'''
|
||||
Subroutine to reads dataset from AMUN HDF5 snapshots.
|
||||
|
||||
@ -118,27 +119,77 @@ def amun_dataset(fname, vname):
|
||||
print('It seems this HDF5 file is corrupted or not compatible with the AMUN format!')
|
||||
return False
|
||||
|
||||
# build the list of supported variables
|
||||
# get the file path
|
||||
#
|
||||
variables = []
|
||||
for var in f['variables'].keys():
|
||||
variables.append(var)
|
||||
dname = op.dirname(fname)
|
||||
|
||||
# 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 False
|
||||
if progress:
|
||||
sys.stdout.write("Data file path:\n '%s'\n" % (dname))
|
||||
|
||||
# get attributes necessary to reconstruct the domain
|
||||
#
|
||||
g = f['attributes'].attrs
|
||||
|
||||
# get the set of equations used to perform the simulation
|
||||
# and the equation of state
|
||||
#
|
||||
eqsys = g.get('eqsys')[0].astype(str)
|
||||
eos = g.get('eos')[0].astype(str)
|
||||
|
||||
# get the snapshot number, the number of domain files, and the number of blocks
|
||||
#
|
||||
nr = g.get('isnap')[0]
|
||||
nc = g.get('nprocs')[0]
|
||||
nl = g.get('nleafs')[0]
|
||||
if eos == 'adi':
|
||||
gm = g.get('gamma')[0]
|
||||
|
||||
# build the list of supported variables
|
||||
#
|
||||
variables = []
|
||||
for var in f['variables'].keys():
|
||||
variables.append(var)
|
||||
|
||||
# add derived variables if possible
|
||||
#
|
||||
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:
|
||||
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 False
|
||||
|
||||
# prepare array to hold data
|
||||
#
|
||||
@ -157,9 +208,10 @@ def amun_dataset(fname, vname):
|
||||
|
||||
# iterate over all subdomain files
|
||||
#
|
||||
dname = op.dirname(fname)
|
||||
nb = 0
|
||||
for n in range(nc):
|
||||
lname = op.join(dname, 'p%06d_%05d.h5' % (nr, n))
|
||||
fname = 'p%06d_%05d.h5' % (nr, n)
|
||||
lname = op.join(dname, fname)
|
||||
f = h5.File(lname, 'r')
|
||||
g = f['attributes'].attrs
|
||||
dblocks = g.get('dblocks')[0]
|
||||
@ -167,29 +219,171 @@ def amun_dataset(fname, vname):
|
||||
g = f['coordinates']
|
||||
levels = g['levels'][()]
|
||||
coords = g['coords'][()]
|
||||
g = f['variables']
|
||||
if ndims == 3:
|
||||
dataset = g[vname][ng:-ng,ng:-ng,ng:-ng,:]
|
||||
for l in range(dblocks):
|
||||
nn = 2**(ml - levels[l])
|
||||
cm = bm * nn
|
||||
ibeg = coords[:,l] * cm[:]
|
||||
iend = ibeg + cm[:]
|
||||
dx = g['dx'][()]
|
||||
dy = g['dy'][()]
|
||||
dz = g['dz'][()]
|
||||
g = f['variables']
|
||||
if 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])
|
||||
cm = bm[0:ndims] * nn
|
||||
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] = np.repeat(np.repeat(np.repeat(dataset[:,:,:,l], nn, axis = 0), nn, axis = 1), nn, axis = 2)
|
||||
else:
|
||||
dataset = g[vname][0,ng:-ng,ng:-ng,:]
|
||||
for l in range(dblocks):
|
||||
nn = 2**(ml - levels[l])
|
||||
cm = bm[0:ndims] * nn
|
||||
ibeg = coords[:,l] * cm[:]
|
||||
iend = ibeg + cm[:]
|
||||
ds = dataset[ng:-ng,ng:-ng,ng:-ng,l]
|
||||
for p in range(ndims):
|
||||
ds = np.repeat(ds, nn, axis = p)
|
||||
ret[kb:ke,jb:je,ib:ie] = ds
|
||||
else:
|
||||
ib, jb = ibeg[0], ibeg[1]
|
||||
ie, je = iend[0], iend[1]
|
||||
ret[jb:je,ib:ie] = np.repeat(np.repeat(dataset[:,:,l], nn, axis = 0), nn, axis = 1)
|
||||
ds = dataset[0,ng:-ng,ng:-ng,l]
|
||||
for p in range(ndims):
|
||||
ds = np.repeat(ds, nn, axis = p)
|
||||
ret[jb:je,ib:ie] = ds
|
||||
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()
|
||||
|
||||
f.close()
|
||||
|
||||
if (progress):
|
||||
sys.stdout.write('\n')
|
||||
sys.stdout.flush()
|
||||
|
||||
# return the dataset
|
||||
#
|
||||
return ret
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user