PYTHON: Add function to read datablocks.

Signed-off-by: Grzegorz Kowal <grzegorz@amuncode.org>
This commit is contained in:
Grzegorz Kowal 2020-05-01 15:02:44 -03:00
parent 084750da46
commit e86b9f7339

View File

@ -134,6 +134,18 @@ class AmunXML:
self.variables = variables
def readdblocks(self, no, var):
import numpy as np
dfile = op.join(self.path, 'datablock_%s_%06d.bin' % (var, no))
if op.exists(dfile):
return np.fromfile(dfile, dtype='float64')
else:
print("Cannot find file %s!" % dfile)
return None
def dataset(self, var):
import numpy as np
@ -199,336 +211,181 @@ class AmunXML:
for p in range(nb):
dset[...,p] = level[ids[p]]
elif var == 'velo':
dfile = op.join(self.path, 'datablock_velx_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
tmp = self.readdblocks(n, 'velx')
dset = tmp**2
dfile = op.join(self.path, 'datablock_vely_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
tmp = self.readdblocks(n, 'vely')
dset += tmp**2
dfile = op.join(self.path, 'datablock_velz_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
tmp = self.readdblocks(n, 'velz')
dset += tmp**2
dset = np.reshape(np.sqrt(dset), cm)
elif var == 'magn':
dfile = op.join(self.path, 'datablock_magx_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
tmp = self.readdblocks(n, 'magx')
dset = tmp**2
dfile = op.join(self.path, 'datablock_magy_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
tmp = self.readdblocks(n, 'magy')
dset += tmp**2
dfile = op.join(self.path, 'datablock_magz_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
tmp = self.readdblocks(n, 'magz')
dset += tmp**2
dset = np.reshape(np.sqrt(dset), cm)
elif var == 'ekin':
dfile = op.join(self.path, 'datablock_velx_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
tmp = self.readdblocks(n, 'velx')
dset = tmp**2
dfile = op.join(self.path, 'datablock_vely_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
tmp = self.readdblocks(n, 'vely')
dset += tmp**2
dfile = op.join(self.path, 'datablock_velz_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
tmp = self.readdblocks(n, 'velz')
dset += tmp**2
dfile = op.join(self.path, 'datablock_dens_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
dset = 0.5 * np.reshape(tmp * dset, cm)
tmp = self.readdblocks(n, 'dens')
dset *= tmp
dset *= 0.5
dset = np.reshape(dset, cm)
elif var == 'emag':
dfile = op.join(self.path, 'datablock_magx_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
tmp = self.readdblocks(n, 'magx')
dset = tmp**2
dfile = op.join(self.path, 'datablock_magy_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
tmp = self.readdblocks(n, 'magy')
dset += tmp**2
dfile = op.join(self.path, 'datablock_magz_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
tmp = self.readdblocks(n, 'magz')
dset += tmp**2
dset = 0.5 * np.reshape(dset, cm)
dset *= 0.5
dset = np.reshape(dset, cm)
elif var == 'eint':
dfile = op.join(self.path, 'datablock_pres_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
dset = 1.0 / (self.attributes('gamma') - 1.0) * tmp
dset = self.readdblocks(n, 'pres')
dset *= 1.0 / (self.attributes('gamma') - 1.0)
dset = np.reshape(dset, cm)
elif var == 'temp':
dfile = op.join(self.path, 'datablock_pres_%06d.bin' % (n))
if op.exists(dfile):
dset = np.fromfile(dfile, dtype='float64')
else:
return None
dfile = op.join(self.path, 'datablock_dens_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
dset = dset / tmp
dset = self.readdblocks(n, 'pres')
tmp = self.readdblocks(n, 'pres')
dset /= tmp
dset = np.reshape(dset, cm)
elif var == 'etot':
dfile = op.join(self.path, 'datablock_velx_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
tmp = self.readdblocks(n, 'velx')
dset = tmp**2
dfile = op.join(self.path, 'datablock_vely_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
tmp = self.readdblocks(n, 'vely')
dset += tmp**2
dfile = op.join(self.path, 'datablock_velz_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
tmp = self.readdblocks(n, 'velz')
dset += tmp**2
dfile = op.join(self.path, 'datablock_dens_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
dset = tmp * dset
dfile = op.join(self.path, 'datablock_pres_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
tmp = self.readdblocks(n, 'dens')
dset *= tmp
tmp = self.readdblocks(n, 'pres')
dset += 2.0 / (self.attributes('gamma') - 1.0) * tmp
if 'magn' in self.variables:
dfile = op.join(self.path, 'datablock_magx_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
tmp = self.readdblocks(n, 'magx')
dset = tmp**2
tmp = self.readdblocks(n, 'magy')
dset += tmp**2
dfile = op.join(self.path, 'datablock_magy_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
tmp = self.readdblocks(n, 'magz')
dset += tmp**2
dfile = op.join(self.path, 'datablock_magz_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
dset += tmp**2
dset = 0.5 * np.reshape(dset, cm)
dset *= 0.5
dset = np.reshape(dset, cm)
elif var == 'lorentz':
dfile = op.join(self.path, 'datablock_velx_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
tmp = self.readdblocks(n, 'velx')
dset = tmp**2
dfile = op.join(self.path, 'datablock_vely_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
tmp = self.readdblocks(n, 'vely')
dset += tmp**2
dfile = op.join(self.path, 'datablock_velz_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.fromfile(dfile, dtype='float64')
else:
return None
tmp = self.readdblocks(n, 'velz')
dset += tmp**2
dset = 1.0 / np.sqrt(1.0 - np.reshape(dset, cm))
dset = 1.0 / np.sqrt(1.0 - dset)
dset = np.reshape(dset, cm)
elif var == 'divv':
m = nd - 1
dfile = op.join(self.path, 'datablock_velx_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.reshape(np.fromfile(dfile, dtype='float64'), cm)
else:
return None
dset = (np.roll(tmp, -1, axis = m) - np.roll(tmp, 1, axis = m))
tmp = np.reshape(self.readdblocks(n, 'velx'), cm)
dset = (np.roll(tmp, -1, axis = m) \
- np.roll(tmp, 1, axis = m))
m -= 1
dfile = op.join(self.path, 'datablock_vely_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.reshape(np.fromfile(dfile, dtype='float64'), cm)
else:
return None
dset += (np.roll(tmp, -1, axis = m) - np.roll(tmp, 1, axis = m))
tmp = np.reshape(self.readdblocks(n, 'vely'), cm)
dset += (np.roll(tmp, -1, axis = m) \
- np.roll(tmp, 1, axis = m))
m -= 1
if m >= 0:
dfile = op.join(self.path, 'datablock_velz_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.reshape(np.fromfile(dfile, dtype='float64'), cm)
else:
return None
dset += (np.roll(tmp, -1, axis = m) - np.roll(tmp, 0, axis = m))
tmp = np.reshape(self.readdblocks(n, 'velz'), cm)
dset += (np.roll(tmp, -1, axis = m) \
- np.roll(tmp, 0, axis = m))
dset *= 0.5
for p in range(nb):
ii = ids[p]
dset[...,p] = 0.5 * dset[...,p] / self.cell_size[level[ii]]
dset[...,p] /= self.cell_size[level[ids[p]]]
elif var == 'divb':
m = nd - 1
dfile = op.join(self.path, 'datablock_magx_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.reshape(np.fromfile(dfile, dtype='float64'), cm)
else:
return None
dset = (np.roll(tmp, -1, axis = m) - np.roll(tmp, 1, axis = m))
tmp = np.reshape(self.readdblocks(n, 'magx'), cm)
dset = (np.roll(tmp, -1, axis = m) \
- np.roll(tmp, 1, axis = m))
m -= 1
dfile = op.join(self.path, 'datablock_magy_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.reshape(np.fromfile(dfile, dtype='float64'), cm)
else:
return None
dset += (np.roll(tmp, -1, axis = m) - np.roll(tmp, 1, axis = m))
tmp = np.reshape(self.readdblocks(n, 'magy'), cm)
dset += (np.roll(tmp, -1, axis = m) \
- np.roll(tmp, 1, axis = m))
m -= 1
if m >= 0:
dfile = op.join(self.path, 'datablock_magz_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.reshape(np.fromfile(dfile, dtype='float64'), cm)
else:
return None
dset += (np.roll(tmp, -1, axis = m) - np.roll(tmp, 0, axis = m))
tmp = np.reshape(self.readdblocks(n, 'magz'), cm)
dset += (np.roll(tmp, -1, axis = m) \
- np.roll(tmp, 0, axis = m))
dset *= 0.5
for p in range(nb):
ii = ids[p]
dset[...,p] = 0.5 * dset[...,p] / self.cell_size[level[ii]]
dset[...,p] /= self.cell_size[level[ids[p]]]
elif var == 'vort':
if nd == 3:
dfile = op.join(self.path, 'datablock_velx_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.reshape(np.fromfile(dfile, dtype='float64'), cm)
else:
return None
wy = (np.roll(tmp, -1, axis = 0) - np.roll(tmp, 1, axis = 0))
wz = (np.roll(tmp, 1, axis = 1) - np.roll(tmp, -1, axis = 1))
dfile = op.join(self.path, 'datablock_vely_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.reshape(np.fromfile(dfile, dtype='float64'), cm)
else:
return None
wz += (np.roll(tmp, -1, axis = 2) - np.roll(tmp, 1, axis = 2))
wx = (np.roll(tmp, 1, axis = 0) - np.roll(tmp, -1, axis = 0))
dfile = op.join(self.path, 'datablock_velz_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.reshape(np.fromfile(dfile, dtype='float64'), cm)
else:
return None
wx += (np.roll(tmp, -1, axis = 1) - np.roll(tmp, 1, axis = 1))
wy += (np.roll(tmp, 1, axis = 2) - np.roll(tmp, -1, axis = 2))
tmp = np.reshape(self.readdblocks(n, 'velx'), cm)
wy = (np.roll(tmp, -1, axis = 0) \
- np.roll(tmp, 1, axis = 0))
wz = (np.roll(tmp, 1, axis = 1) \
- np.roll(tmp, -1, axis = 1))
tmp = np.reshape(self.readdblocks(n, 'vely'), cm)
wz += (np.roll(tmp, -1, axis = 2) \
- np.roll(tmp, 1, axis = 2))
wx = (np.roll(tmp, 1, axis = 0) \
- np.roll(tmp, -1, axis = 0))
tmp = np.reshape(self.readdblocks(n, 'velz'), cm)
wx += (np.roll(tmp, -1, axis = 1) \
- np.roll(tmp, 1, axis = 1))
wy += (np.roll(tmp, 1, axis = 2) \
- np.roll(tmp, -1, axis = 2))
else:
dfile = op.join(self.path, 'datablock_velx_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.reshape(np.fromfile(dfile, dtype='float64'), cm)
else:
return None
wz = (np.roll(tmp, 1, axis = 0) - np.roll(tmp, -1, axis = 0))
dfile = op.join(self.path, 'datablock_vely_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.reshape(np.fromfile(dfile, dtype='float64'), cm)
else:
return None
wz += (np.roll(tmp, -1, axis = 1) - np.roll(tmp, 1, axis = 1))
dfile = op.join(self.path, 'datablock_velz_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.reshape(np.fromfile(dfile, dtype='float64'), cm)
else:
return None
wx = (np.roll(tmp, -1, axis = 0) - np.roll(tmp, 1, axis = 0))
wy = (np.roll(tmp, 1, axis = 1) - np.roll(tmp, -1, axis = 1))
tmp = np.reshape(self.readdblocks(n, 'velx'), cm)
wz = (np.roll(tmp, 1, axis = 0) \
- np.roll(tmp, -1, axis = 0))
tmp = np.reshape(self.readdblocks(n, 'vely'), cm)
wz += (np.roll(tmp, -1, axis = 1) \
- np.roll(tmp, 1, axis = 1))
tmp = np.reshape(self.readdblocks(n, 'velz'), cm)
wx = (np.roll(tmp, -1, axis = 0) \
- np.roll(tmp, 1, axis = 0))
wy = (np.roll(tmp, 1, axis = 1) \
- np.roll(tmp, -1, axis = 1))
dset = 0.5 * np.sqrt(wx**2 + wy**2 + wz**2)
for p in range(nb):
ii = ids[p]
dset[...,p] = dset[...,p] / self.cell_size[level[ii]]
dset[...,p] /= self.cell_size[level[ids[p]]]
elif var == 'curr':
if nd == 3:
dfile = op.join(self.path, 'datablock_magx_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.reshape(np.fromfile(dfile, dtype='float64'), cm)
else:
return None
wy = (np.roll(tmp, -1, axis = 0) - np.roll(tmp, 1, axis = 0))
wz = (np.roll(tmp, 1, axis = 1) - np.roll(tmp, -1, axis = 1))
dfile = op.join(self.path, 'datablock_magy_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.reshape(np.fromfile(dfile, dtype='float64'), cm)
else:
return None
wz += (np.roll(tmp, -1, axis = 2) - np.roll(tmp, 1, axis = 2))
wx = (np.roll(tmp, 1, axis = 0) - np.roll(tmp, -1, axis = 0))
dfile = op.join(self.path, 'datablock_magz_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.reshape(np.fromfile(dfile, dtype='float64'), cm)
else:
return None
wx += (np.roll(tmp, -1, axis = 1) - np.roll(tmp, 1, axis = 1))
wy += (np.roll(tmp, 1, axis = 2) - np.roll(tmp, -1, axis = 2))
tmp = np.reshape(self.readdblocks(n, 'magx'), cm)
wy = (np.roll(tmp, -1, axis = 0) \
- np.roll(tmp, 1, axis = 0))
wz = (np.roll(tmp, 1, axis = 1) \
- np.roll(tmp, -1, axis = 1))
tmp = np.reshape(self.readdblocks(n, 'magy'), cm)
wz += (np.roll(tmp, -1, axis = 2) \
- np.roll(tmp, 1, axis = 2))
wx = (np.roll(tmp, 1, axis = 0) \
- np.roll(tmp, -1, axis = 0))
tmp = np.reshape(self.readdblocks(n, 'magz'), cm)
wx += (np.roll(tmp, -1, axis = 1) \
- np.roll(tmp, 1, axis = 1))
wy += (np.roll(tmp, 1, axis = 2) \
- np.roll(tmp, -1, axis = 2))
else:
dfile = op.join(self.path, 'datablock_magx_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.reshape(np.fromfile(dfile, dtype='float64'), cm)
else:
return None
wz = (np.roll(tmp, 1, axis = 0) - np.roll(tmp, -1, axis = 0))
dfile = op.join(self.path, 'datablock_magy_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.reshape(np.fromfile(dfile, dtype='float64'), cm)
else:
return None
wz += (np.roll(tmp, -1, axis = 1) - np.roll(tmp, 1, axis = 1))
dfile = op.join(self.path, 'datablock_magz_%06d.bin' % (n))
if op.exists(dfile):
tmp = np.reshape(np.fromfile(dfile, dtype='float64'), cm)
else:
return None
wx = (np.roll(tmp, -1, axis = 0) - np.roll(tmp, 1, axis = 0))
wy = (np.roll(tmp, 1, axis = 1) - np.roll(tmp, -1, axis = 1))
tmp = np.reshape(self.readdblocks(n, 'magx'), cm)
wz = (np.roll(tmp, 1, axis = 0) \
- np.roll(tmp, -1, axis = 0))
tmp = np.reshape(self.readdblocks(n, 'magy'), cm)
wz += (np.roll(tmp, -1, axis = 1) \
- np.roll(tmp, 1, axis = 1))
tmp = np.reshape(self.readdblocks(n, 'magz'), cm)
wx = (np.roll(tmp, -1, axis = 0) \
- np.roll(tmp, 1, axis = 0))
wy = (np.roll(tmp, 1, axis = 1) \
- np.roll(tmp, -1, axis = 1))
dset = 0.5 * np.sqrt(wx**2 + wy**2 + wz**2)
for p in range(nb):
ii = ids[p]
dset[...,p] = dset[...,p] / self.cell_size[level[ii]]
dset[...,p] /= self.cell_size[level[ids[p]]]
else:
dfile = op.join(self.path, 'datablock_%s_%06d.bin' % (var, n))
if op.exists(dfile):
dset = np.fromfile(dfile, dtype='float64')
else:
return None
dset = self.readdblocks(n, var)
dset = np.reshape(dset, cm)
for p in range(nb):