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

PYTHON: Rewrite Amun module to support old and new snapshots.

Browse Source 
Signed-off-by: Grzegorz Kowal <grzegorz@amuncode.org>
This commit is contained in:
Grzegorz Kowal 2021-11-24 12:33:33 -03:00
parent a1cf8f9192
commit fd4029a196
5 changed files with 301 additions and 263 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
python/amunpy/setup.py
View File

@ -5,7 +5,7 @@ with open("README.md", "r", encoding="utf-8") as fh:
setuptools.setup( setuptools.setup(
name="amunpy", name="amunpy",
version="0.9.5", version="0.9.6",
author="Grzegorz Kowal", author="Grzegorz Kowal",
author_email="grzegorz@amuncode.org", author_email="grzegorz@amuncode.org",
description="Python Interface for the AMUN code's snapshots", description="Python Interface for the AMUN code's snapshots",

2
python/amunpy/src/amunpy/__init__.py
View File

@ -21,6 +21,6 @@ __all__ = [ 'AmunXML', 'AmunH5', 'WriteVTK', \
__author__ = "Grzegorz Kowal" __author__ = "Grzegorz Kowal"
__copyright__ = "Copyright 2018-2021, Grzegorz Kowal <grzegorz@amuncode.org>" __copyright__ = "Copyright 2018-2021, Grzegorz Kowal <grzegorz@amuncode.org>"
__version__ = "0.9.5" __version__ = "0.9.6"
__maintainer__ = "Grzegorz Kowal" __maintainer__ = "Grzegorz Kowal"
__email__ = "grzegorz@amuncode.org" __email__ = "grzegorz@amuncode.org"

471
python/amunpy/src/amunpy/amun.py
View File

@ -31,7 +31,7 @@
class Amun: class Amun:
"""AMUN snapshot base class""" """AMUN snapshot base class"""
def __init__(self, path): def __init__(self, path, version=-1):
import os import os
@ -49,6 +49,7 @@ class Amun:
self.path = path self.path = path
self.dataformat = None self.dataformat = None
self.version = version
self.chunkname = '{:d}' self.chunkname = '{:d}'
self.attributes = dict() self.attributes = dict()
self.chunks = dict() self.chunks = dict()
@ -169,7 +170,7 @@ class Amun:
for n in range(len(self.chunks)): for n in range(len(self.chunks)):
if not 'dims' in self.chunks[n]: if not 'dims' in self.chunks[n]:
self.chunks[n]['dims'] = [ self.attributes['bcells'] ]*self.attributes['ndims'] self.chunks[n]['dims'] = [ self.attributes['bcells'] ]*self.attributes['ndims']
self.chunks[n]['dims'].append( self.chunks[n]['dblocks'] ) self.chunks[n]['dims'][:0] = [ self.chunks[n]['dblocks'] ]
def __complete_variables__(self): def __complete_variables__(self):
@ -253,7 +254,7 @@ class Amun:
if dataset == 'mlev': if dataset == 'mlev':
dset = numpy.zeros(self.chunks[chunk_number]['dims']) dset = numpy.zeros(self.chunks[chunk_number]['dims'])
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] = self.chunks[chunk_number]['levels'][p] dset[p,...] = self.chunks[chunk_number]['levels'][p]
elif dataset == 'logd': elif dataset == 'logd':
dset = numpy.log10(self.__read_binary_data__('dens', chunk_number)) dset = numpy.log10(self.__read_binary_data__('dens', chunk_number))
elif dataset == 'logp': elif dataset == 'logp':
@ -308,13 +309,13 @@ class Amun:
dset = vy * bz - vz * by dset = vy * bz - vz * by
if self.attributes['resistivity'] > 0: if self.attributes['resistivity'] > 0:
if self.attributes['ndims'] == 3: if self.attributes['ndims'] == 3:
tmp = (numpy.roll(by, 1, axis=0) - numpy.roll(by, -1, axis=0)) tmp = (numpy.roll(by, 1, axis=1) - numpy.roll(by, -1, axis=1))
tmp += (numpy.roll(bz, -1, axis=1) - numpy.roll(bz, 1, axis=1)) tmp += (numpy.roll(bz, -1, axis=2) - numpy.roll(bz, 1, axis=2))
else: else:
tmp = (numpy.roll(bz, -1, axis=0) - numpy.roll(bz, 1, axis=0)) tmp = (numpy.roll(bz, -1, axis=1) - numpy.roll(bz, 1, axis=1))
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
tmp[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] tmp[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
dset -= 0.5 * self.attributes['resistivity'] * tmp dset -= 0.5 * self.attributes['resistivity'] * tmp
elif dataset == 'eley': elif dataset == 'eley':
bx = self.__read_binary_data__('magx', chunk_number) bx = self.__read_binary_data__('magx', chunk_number)
@ -324,13 +325,13 @@ class Amun:
dset = vz * bx - vx * bz dset = vz * bx - vx * bz
if self.attributes['resistivity'] > 0: if self.attributes['resistivity'] > 0:
if self.attributes['ndims'] == 3: if self.attributes['ndims'] == 3:
tmp = (numpy.roll(bx, -1, axis=0) - numpy.roll(bx, 1, axis=0)) tmp = (numpy.roll(bx, -1, axis=1) - numpy.roll(bx, 1, axis=1))
tmp += (numpy.roll(bz, 1, axis=2) - numpy.roll(bz, -1, axis=2)) tmp += (numpy.roll(bz, 1, axis=3) - numpy.roll(bz, -1, axis=3))
else: else:
tmp = (numpy.roll(bz, 1, axis=1) - numpy.roll(bz, -1, axis=1)) tmp = (numpy.roll(bz, 1, axis=2) - numpy.roll(bz, -1, axis=2))
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
tmp[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] tmp[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
dset -= 0.5 * self.attributes['resistivity'] * tmp dset -= 0.5 * self.attributes['resistivity'] * tmp
elif dataset == 'elez': elif dataset == 'elez':
bx = self.__read_binary_data__('magx', chunk_number) bx = self.__read_binary_data__('magx', chunk_number)
@ -340,14 +341,14 @@ class Amun:
dset = vx * by - vy * bx dset = vx * by - vy * bx
if self.attributes['resistivity'] > 0: if self.attributes['resistivity'] > 0:
if self.attributes['ndims'] == 3: if self.attributes['ndims'] == 3:
tmp = (numpy.roll(bx, 1, axis=2) - numpy.roll(bx, -1, axis=2))
tmp += (numpy.roll(by, -1, axis=3) - numpy.roll(by, 1, axis=3))
else:
tmp = (numpy.roll(bx, 1, axis=1) - numpy.roll(bx, -1, axis=1)) tmp = (numpy.roll(bx, 1, axis=1) - numpy.roll(bx, -1, axis=1))
tmp += (numpy.roll(by, -1, axis=2) - numpy.roll(by, 1, axis=2)) tmp += (numpy.roll(by, -1, axis=2) - numpy.roll(by, 1, axis=2))
else:
tmp = (numpy.roll(bx, 1, axis=0) - numpy.roll(bx, -1, axis=0))
tmp += (numpy.roll(by, -1, axis=1) - numpy.roll(by, 1, axis=1))
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
tmp[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] tmp[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
dset -= 0.5 * self.attributes['resistivity'] * tmp dset -= 0.5 * self.attributes['resistivity'] * tmp
elif dataset == 'elec': elif dataset == 'elec':
b1 = self.__read_binary_data__('magy', chunk_number) b1 = self.__read_binary_data__('magy', chunk_number)
@ -357,13 +358,13 @@ class Amun:
dtmp = v1 * b2 - v2 * b1 dtmp = v1 * b2 - v2 * b1
if self.attributes['resistivity'] > 0: if self.attributes['resistivity'] > 0:
if self.attributes['ndims'] == 3: if self.attributes['ndims'] == 3:
tmp = (numpy.roll(b1, 1, axis=0) - numpy.roll(b1, -1, axis=0)) tmp = (numpy.roll(b1, 1, axis=1) - numpy.roll(b1, -1, axis=1))
tmp += (numpy.roll(b2, -1, axis=1) - numpy.roll(b2, 1, axis=1)) tmp += (numpy.roll(b2, -1, axis=2) - numpy.roll(b2, 1, axis=2))
else: else:
tmp = (numpy.roll(b2, -1, axis=0) - numpy.roll(b2, 1, axis=0)) tmp = (numpy.roll(b2, -1, axis=1) - numpy.roll(b2, 1, axis=1))
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
tmp[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] tmp[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
dtmp -= 0.5 * self.attributes['resistivity'] * tmp dtmp -= 0.5 * self.attributes['resistivity'] * tmp
dset = dtmp**2 dset = dtmp**2
@ -372,13 +373,13 @@ class Amun:
dtmp = v2 * b1 - v1 * b2 dtmp = v2 * b1 - v1 * b2
if self.attributes['resistivity'] > 0: if self.attributes['resistivity'] > 0:
if self.attributes['ndims'] == 3: if self.attributes['ndims'] == 3:
tmp = (numpy.roll(b1, -1, axis=0) - numpy.roll(b1, 1, axis=0)) tmp = (numpy.roll(b1, -1, axis=1) - numpy.roll(b1, 1, axis=1))
tmp += (numpy.roll(b2, 1, axis=2) - numpy.roll(b2, -1, axis=2)) tmp += (numpy.roll(b2, 1, axis=3) - numpy.roll(b2, -1, axis=3))
else: else:
tmp = (numpy.roll(b2, 1, axis=1) - numpy.roll(b2, -1, axis=1)) tmp = (numpy.roll(b2, 1, axis=2) - numpy.roll(b2, -1, axis=2))
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
tmp[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] tmp[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
dtmp -= 0.5 * self.attributes['resistivity'] * tmp dtmp -= 0.5 * self.attributes['resistivity'] * tmp
dset += dtmp**2 dset += dtmp**2
@ -387,14 +388,14 @@ class Amun:
dtmp = v1 * b2 - v2 * b1 dtmp = v1 * b2 - v2 * b1
if self.attributes['resistivity'] > 0: if self.attributes['resistivity'] > 0:
if self.attributes['ndims'] == 3: if self.attributes['ndims'] == 3:
tmp = (numpy.roll(b1, 1, axis=2) - numpy.roll(b1, -1, axis=2))
tmp += (numpy.roll(b2, -1, axis=3) - numpy.roll(b2, 1, axis=3))
else:
tmp = (numpy.roll(b1, 1, axis=1) - numpy.roll(b1, -1, axis=1)) tmp = (numpy.roll(b1, 1, axis=1) - numpy.roll(b1, -1, axis=1))
tmp += (numpy.roll(b2, -1, axis=2) - numpy.roll(b2, 1, axis=2)) tmp += (numpy.roll(b2, -1, axis=2) - numpy.roll(b2, 1, axis=2))
else:
tmp = (numpy.roll(b1, 1, axis=0) - numpy.roll(b1, -1, axis=0))
tmp += (numpy.roll(b2, -1, axis=1) - numpy.roll(b2, 1, axis=1))
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
tmp[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] tmp[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
dtmp -= 0.5 * self.attributes['resistivity'] * tmp dtmp -= 0.5 * self.attributes['resistivity'] * tmp
dset += dtmp**2 dset += dtmp**2
dset = numpy.sqrt(dset) dset = numpy.sqrt(dset)
@ -406,13 +407,13 @@ class Amun:
wx = v1 * b2 - v2 * b1 wx = v1 * b2 - v2 * b1
if self.attributes['resistivity'] > 0: if self.attributes['resistivity'] > 0:
if self.attributes['ndims'] == 3: if self.attributes['ndims'] == 3:
tmp = (numpy.roll(b1, 1, axis=0) - numpy.roll(b1, -1, axis=0)) tmp = (numpy.roll(b1, 1, axis=1) - numpy.roll(b1, -1, axis=1))
tmp += (numpy.roll(b2, -1, axis=1) - numpy.roll(b2, 1, axis=1)) tmp += (numpy.roll(b2, -1, axis=2) - numpy.roll(b2, 1, axis=2))
else: else:
tmp = (numpy.roll(b2, -1, axis=0) - numpy.roll(b2, 1, axis=0)) tmp = (numpy.roll(b2, -1, axis=1) - numpy.roll(b2, 1, axis=1))
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
tmp[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] tmp[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
wx -= 0.5 * self.attributes['resistivity'] * tmp wx -= 0.5 * self.attributes['resistivity'] * tmp
b1 = self.__read_binary_data__('magx', chunk_number) b1 = self.__read_binary_data__('magx', chunk_number)
@ -420,13 +421,13 @@ class Amun:
wy = v2 * b1 - v1 * b2 wy = v2 * b1 - v1 * b2
if self.attributes['resistivity'] > 0: if self.attributes['resistivity'] > 0:
if self.attributes['ndims'] == 3: if self.attributes['ndims'] == 3:
tmp = (numpy.roll(b1, -1, axis=0) - numpy.roll(b1, 1, axis=0)) tmp = (numpy.roll(b1, -1, axis=1) - numpy.roll(b1, 1, axis=1))
tmp += (numpy.roll(b2, 1, axis=2) - numpy.roll(b2, -1, axis=2)) tmp += (numpy.roll(b2, 1, axis=3) - numpy.roll(b2, -1, axis=3))
else: else:
tmp = (numpy.roll(b2, 1, axis=1) - numpy.roll(b2, -1, axis=1)) tmp = (numpy.roll(b2, 1, axis=2) - numpy.roll(b2, -1, axis=2))
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
tmp[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] tmp[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
wy -= 0.5 * self.attributes['resistivity'] * tmp wy -= 0.5 * self.attributes['resistivity'] * tmp
b2 = self.__read_binary_data__('magy', chunk_number) b2 = self.__read_binary_data__('magy', chunk_number)
@ -434,14 +435,14 @@ class Amun:
wz = v1 * b2 - v2 * b1 wz = v1 * b2 - v2 * b1
if self.attributes['resistivity'] > 0: if self.attributes['resistivity'] > 0:
if self.attributes['ndims'] == 3: if self.attributes['ndims'] == 3:
tmp = (numpy.roll(b1, 1, axis=2) - numpy.roll(b1, -1, axis=2))
tmp += (numpy.roll(b2, -1, axis=3) - numpy.roll(b2, 1, axis=3))
else:
tmp = (numpy.roll(b1, 1, axis=1) - numpy.roll(b1, -1, axis=1)) tmp = (numpy.roll(b1, 1, axis=1) - numpy.roll(b1, -1, axis=1))
tmp += (numpy.roll(b2, -1, axis=2) - numpy.roll(b2, 1, axis=2)) tmp += (numpy.roll(b2, -1, axis=2) - numpy.roll(b2, 1, axis=2))
else:
tmp = (numpy.roll(b1, 1, axis=0) - numpy.roll(b1, -1, axis=0))
tmp += (numpy.roll(b2, -1, axis=1) - numpy.roll(b2, 1, axis=1))
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
tmp[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] tmp[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
wz -= 0.5 * self.attributes['resistivity'] * tmp wz -= 0.5 * self.attributes['resistivity'] * tmp
dset = [wx, wy, wz] dset = [wx, wy, wz]
@ -478,7 +479,7 @@ class Amun:
dset += tmp**2 dset += tmp**2
dset = 1.0 / numpy.sqrt(1.0 - dset) dset = 1.0 / numpy.sqrt(1.0 - dset)
elif dataset == 'divv': elif dataset == 'divv':
p = self.attributes['ndims'] - 1 p = self.attributes['ndims']
tmp = self.__read_binary_data__('velx', chunk_number) tmp = self.__read_binary_data__('velx', chunk_number)
dset = (numpy.roll(tmp, -1, axis=p) \ dset = (numpy.roll(tmp, -1, axis=p) \
- numpy.roll(tmp, 1, axis=p)) - numpy.roll(tmp, 1, axis=p))
@ -494,202 +495,202 @@ class Amun:
dset *= 0.5 dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'vorx': elif dataset == 'vorx':
if self.attributes['ndims'] == 3: if self.attributes['ndims'] == 3:
tmp = self.__read_binary_data__('vely', chunk_number) tmp = self.__read_binary_data__('vely', chunk_number)
dset = (numpy.roll(tmp, 1, axis=0) \ dset = (numpy.roll(tmp, 1, axis=1) \
- numpy.roll(tmp, -1, axis=0)) - numpy.roll(tmp, -1, axis=1))
tmp = self.__read_binary_data__('velz', chunk_number) tmp = self.__read_binary_data__('velz', chunk_number)
dset += (numpy.roll(tmp, -1, axis=1) \ dset += (numpy.roll(tmp, -1, axis=2) \
- numpy.roll(tmp, 1, axis=1)) - numpy.roll(tmp, 1, axis=2))
else: else:
tmp = self.__read_binary_data__('velz', chunk_number) tmp = self.__read_binary_data__('velz', chunk_number)
dset = (numpy.roll(tmp, -1, axis=0) \ dset = (numpy.roll(tmp, -1, axis=1) \
- numpy.roll(tmp, 1, axis=0)) - numpy.roll(tmp, 1, axis=1))
dset *= 0.5 dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'vory': elif dataset == 'vory':
if self.attributes['ndims'] == 3: if self.attributes['ndims'] == 3:
tmp = self.__read_binary_data__('velx', chunk_number) tmp = self.__read_binary_data__('velx', chunk_number)
dset = (numpy.roll(tmp, -1, axis=0) \ dset = (numpy.roll(tmp, -1, axis=1) \
- numpy.roll(tmp, 1, axis=0)) - numpy.roll(tmp, 1, axis=1))
tmp = self.__read_binary_data__('velz', chunk_number) tmp = self.__read_binary_data__('velz', chunk_number)
dset += (numpy.roll(tmp, 1, axis=2) \ dset += (numpy.roll(tmp, 1, axis=3) \
- numpy.roll(tmp, -1, axis=2)) - numpy.roll(tmp, -1, axis=3))
else: else:
tmp = self.__read_binary_data__('velz', chunk_number) tmp = self.__read_binary_data__('velz', chunk_number)
dset = (numpy.roll(tmp, 1, axis=1) \ dset = (numpy.roll(tmp, 1, axis=2) \
- numpy.roll(tmp, -1, axis=1)) - numpy.roll(tmp, -1, axis=2))
dset *= 0.5 dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'vorz': elif dataset == 'vorz':
if self.attributes['ndims'] == 3: if self.attributes['ndims'] == 3:
tmp = self.__read_binary_data__('velx', chunk_number)
dset = (numpy.roll(tmp, 1, axis=2) \
- numpy.roll(tmp, -1, axis=2))
tmp = self.__read_binary_data__('vely', chunk_number)
dset += (numpy.roll(tmp, -1, axis=3) \
- numpy.roll(tmp, 1, axis=3))
else:
tmp = self.__read_binary_data__('velx', chunk_number) tmp = self.__read_binary_data__('velx', chunk_number)
dset = (numpy.roll(tmp, 1, axis=1) \ dset = (numpy.roll(tmp, 1, axis=1) \
- numpy.roll(tmp, -1, axis=1)) - numpy.roll(tmp, -1, axis=1))
tmp = self.__read_binary_data__('vely', chunk_number) tmp = self.__read_binary_data__('vely', chunk_number)
dset += (numpy.roll(tmp, -1, axis=2) \ dset += (numpy.roll(tmp, -1, axis=2) \
- numpy.roll(tmp, 1, axis=2)) - numpy.roll(tmp, 1, axis=2))
else:
tmp = self.__read_binary_data__('velx', chunk_number)
dset = (numpy.roll(tmp, 1, axis=0) \
- numpy.roll(tmp, -1, axis=0))
tmp = self.__read_binary_data__('vely', chunk_number)
dset += (numpy.roll(tmp, -1, axis=1) \
- numpy.roll(tmp, 1, axis=1))
dset *= 0.5 dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'wvec': elif dataset == 'wvec':
if self.attributes['ndims'] == 3: if self.attributes['ndims'] == 3:
tmp = self.__read_binary_data__('velx', chunk_number) tmp = self.__read_binary_data__('velx', chunk_number)
wy = (numpy.roll(tmp, -1, axis=0) \ wy = (numpy.roll(tmp, -1, axis=1) \
- numpy.roll(tmp, 1, axis=0)) - numpy.roll(tmp, 1, axis=1))
wz = (numpy.roll(tmp, 1, axis=2) \
- numpy.roll(tmp, -1, axis=2))
tmp = self.__read_binary_data__('vely', chunk_number)
wz += (numpy.roll(tmp, -1, axis=3) \
- numpy.roll(tmp, 1, axis=3))
wx = (numpy.roll(tmp, 1, axis=1) \
- numpy.roll(tmp, -1, axis=1))
tmp = self.__read_binary_data__('velz', chunk_number)
wx += (numpy.roll(tmp, -1, axis=2) \
- numpy.roll(tmp, 1, axis=2))
wy += (numpy.roll(tmp, 1, axis=3) \
- numpy.roll(tmp, -1, axis=3))
else:
tmp = self.__read_binary_data__('velx', chunk_number)
wz = (numpy.roll(tmp, 1, axis=1) \ wz = (numpy.roll(tmp, 1, axis=1) \
- numpy.roll(tmp, -1, axis=1)) - numpy.roll(tmp, -1, axis=1))
tmp = self.__read_binary_data__('vely', chunk_number) tmp = self.__read_binary_data__('vely', chunk_number)
wz += (numpy.roll(tmp, -1, axis=2) \ wz += (numpy.roll(tmp, -1, axis=2) \
- numpy.roll(tmp, 1, axis=2)) - numpy.roll(tmp, 1, axis=2))
wx = (numpy.roll(tmp, 1, axis=0) \
- numpy.roll(tmp, -1, axis=0))
tmp = self.__read_binary_data__('velz', chunk_number) tmp = self.__read_binary_data__('velz', chunk_number)
wx += (numpy.roll(tmp, -1, axis=1) \ wx = (numpy.roll(tmp, -1, axis=1) \
- numpy.roll(tmp, 1, axis=1)) - numpy.roll(tmp, 1, axis=1))
wy += (numpy.roll(tmp, 1, axis=2) \ wy = (numpy.roll(tmp, 1, axis=2) \
- numpy.roll(tmp, -1, axis=2)) - numpy.roll(tmp, -1, axis=2))
else:
tmp = self.__read_binary_data__('velx', chunk_number)
wz = (numpy.roll(tmp, 1, axis=0) \
- numpy.roll(tmp, -1, axis=0))
tmp = self.__read_binary_data__('vely', chunk_number)
wz += (numpy.roll(tmp, -1, axis=1) \
- numpy.roll(tmp, 1, axis=1))
tmp = self.__read_binary_data__('velz', chunk_number)
wx = (numpy.roll(tmp, -1, axis=0) \
- numpy.roll(tmp, 1, axis=0))
wy = (numpy.roll(tmp, 1, axis=1) \
- numpy.roll(tmp, -1, axis=1))
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
h = 2 * self.cell_size[self.chunks[chunk_number]['levels'][p]] h = 2 * self.cell_size[self.chunks[chunk_number]['levels'][p]]
wx[...,p] /= h wx[p,...] /= h
wy[...,p] /= h wy[p,...] /= h
wz[...,p] /= h wz[p,...] /= h
dset = [wx, wy, wz] dset = [wx, wy, wz]
elif dataset == 'vort': elif dataset == 'vort':
if self.attributes['ndims'] == 3: if self.attributes['ndims'] == 3:
tmp = self.__read_binary_data__('velx', chunk_number) tmp = self.__read_binary_data__('velx', chunk_number)
wy = (numpy.roll(tmp, -1, axis=0) \ wy = (numpy.roll(tmp, -1, axis=1) \
- numpy.roll(tmp, 1, axis=0)) - numpy.roll(tmp, 1, axis=1))
wz = (numpy.roll(tmp, 1, axis=2) \
- numpy.roll(tmp, -1, axis=2))
tmp = self.__read_binary_data__('vely', chunk_number)
wz += (numpy.roll(tmp, -1, axis=3) \
- numpy.roll(tmp, 1, axis=3))
wx = (numpy.roll(tmp, 1, axis=1) \
- numpy.roll(tmp, -1, axis=1))
tmp = self.__read_binary_data__('velz', chunk_number)
wx += (numpy.roll(tmp, -1, axis=2) \
- numpy.roll(tmp, 1, axis=2))
wy += (numpy.roll(tmp, 1, axis=3) \
- numpy.roll(tmp, -1, axis=3))
else:
tmp = self.__read_binary_data__('velx', chunk_number)
wz = (numpy.roll(tmp, 1, axis=1) \ wz = (numpy.roll(tmp, 1, axis=1) \
- numpy.roll(tmp, -1, axis=1)) - numpy.roll(tmp, -1, axis=1))
tmp = self.__read_binary_data__('vely', chunk_number) tmp = self.__read_binary_data__('vely', chunk_number)
wz += (numpy.roll(tmp, -1, axis=2) \ wz += (numpy.roll(tmp, -1, axis=2) \
- numpy.roll(tmp, 1, axis=2)) - numpy.roll(tmp, 1, axis=2))
wx = (numpy.roll(tmp, 1, axis=0) \
- numpy.roll(tmp, -1, axis=0))
tmp = self.__read_binary_data__('velz', chunk_number) tmp = self.__read_binary_data__('velz', chunk_number)
wx += (numpy.roll(tmp, -1, axis=1) \ wx = (numpy.roll(tmp, -1, axis=1) \
- numpy.roll(tmp, 1, axis=1)) - numpy.roll(tmp, 1, axis=1))
wy += (numpy.roll(tmp, 1, axis=2) \ wy = (numpy.roll(tmp, 1, axis=2) \
- numpy.roll(tmp, -1, axis=2)) - numpy.roll(tmp, -1, axis=2))
else:
tmp = self.__read_binary_data__('velx', chunk_number)
wz = (numpy.roll(tmp, 1, axis=0) \
- numpy.roll(tmp, -1, axis=0))
tmp = self.__read_binary_data__('vely', chunk_number)
wz += (numpy.roll(tmp, -1, axis=1) \
- numpy.roll(tmp, 1, axis=1))
tmp = self.__read_binary_data__('velz', chunk_number)
wx = (numpy.roll(tmp, -1, axis=0) \
- numpy.roll(tmp, 1, axis=0))
wy = (numpy.roll(tmp, 1, axis=1) \
- numpy.roll(tmp, -1, axis=1))
dset = 0.5 * numpy.sqrt(wx**2 + wy**2 + wz**2) dset = 0.5 * numpy.sqrt(wx**2 + wy**2 + wz**2)
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'dvxdx': elif dataset == 'dvxdx':
p = self.attributes['ndims']
tmp = self.__read_binary_data__('velx', chunk_number)
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']):
dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'dvxdy':
p = self.attributes['ndims'] - 1 p = self.attributes['ndims'] - 1
tmp = self.__read_binary_data__('velx', chunk_number) tmp = self.__read_binary_data__('velx', chunk_number)
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p) dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5 dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'dvxdy':
p = self.attributes['ndims'] - 2
tmp = self.__read_binary_data__('velx', chunk_number)
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'dvxdz': elif dataset == 'dvxdz':
p = self.attributes['ndims'] - 3 p = self.attributes['ndims'] - 2
tmp = self.__read_binary_data__('velx', chunk_number) tmp = self.__read_binary_data__('velx', chunk_number)
if p >= 0: if p >= 0:
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p) dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5 dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
else: else:
dset = numpy.zeros_like(tmp) dset = numpy.zeros_like(tmp)
elif dataset == 'dvydx': elif dataset == 'dvydx':
p = self.attributes['ndims']
tmp = self.__read_binary_data__('vely', chunk_number)
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']):
dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'dvydy':
p = self.attributes['ndims'] - 1 p = self.attributes['ndims'] - 1
tmp = self.__read_binary_data__('vely', chunk_number) tmp = self.__read_binary_data__('vely', chunk_number)
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p) dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5 dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'dvydy':
p = self.attributes['ndims'] - 2
tmp = self.__read_binary_data__('vely', chunk_number)
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'dvydz': elif dataset == 'dvydz':
p = self.attributes['ndims'] - 3 p = self.attributes['ndims'] - 2
tmp = self.__read_binary_data__('vely', chunk_number) tmp = self.__read_binary_data__('vely', chunk_number)
if p >= 0: if p >= 0:
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p) dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5 dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
else: else:
dset = numpy.zeros_like(tmp) dset = numpy.zeros_like(tmp)
elif dataset == 'dvzdx': elif dataset == 'dvzdx':
p = self.attributes['ndims']
tmp = self.__read_binary_data__('velz', chunk_number)
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']):
dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'dvzdy':
p = self.attributes['ndims'] - 1 p = self.attributes['ndims'] - 1
tmp = self.__read_binary_data__('velz', chunk_number) tmp = self.__read_binary_data__('velz', chunk_number)
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p) dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5 dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'dvzdy':
p = self.attributes['ndims'] - 2
tmp = self.__read_binary_data__('velz', chunk_number)
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'dvzdz': elif dataset == 'dvzdz':
p = self.attributes['ndims'] - 3 p = self.attributes['ndims'] - 2
tmp = self.__read_binary_data__('velz', chunk_number) tmp = self.__read_binary_data__('velz', chunk_number)
if p >= 0: if p >= 0:
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p) dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5 dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
else: else:
dset = numpy.zeros_like(tmp) dset = numpy.zeros_like(tmp)
elif dataset == 'divb': elif dataset == 'divb':
p = self.attributes['ndims'] - 1 p = self.attributes['ndims']
tmp = self.__read_binary_data__('magx', chunk_number) tmp = self.__read_binary_data__('magx', chunk_number)
dset = (numpy.roll(tmp, -1, axis=p) \ dset = (numpy.roll(tmp, -1, axis=p) \
- numpy.roll(tmp, 1, axis=p)) - numpy.roll(tmp, 1, axis=p))
@ -705,198 +706,198 @@ class Amun:
dset *= 0.5 dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'curx': elif dataset == 'curx':
if self.attributes['ndims'] == 3: if self.attributes['ndims'] == 3:
tmp = self.__read_binary_data__('magy', chunk_number) tmp = self.__read_binary_data__('magy', chunk_number)
dset = (numpy.roll(tmp, 1, axis=0) \ dset = (numpy.roll(tmp, 1, axis=1) \
- numpy.roll(tmp, -1, axis=0)) - numpy.roll(tmp, -1, axis=1))
tmp = self.__read_binary_data__('magz', chunk_number) tmp = self.__read_binary_data__('magz', chunk_number)
dset += (numpy.roll(tmp, -1, axis=1) \ dset += (numpy.roll(tmp, -1, axis=2) \
- numpy.roll(tmp, 1, axis=1)) - numpy.roll(tmp, 1, axis=2))
else: else:
tmp = self.__read_binary_data__('magz', chunk_number) tmp = self.__read_binary_data__('magz', chunk_number)
dset = (numpy.roll(tmp, -1, axis=0) \ dset = (numpy.roll(tmp, -1, axis=1) \
- numpy.roll(tmp, 1, axis=0)) - numpy.roll(tmp, 1, axis=1))
dset *= 0.5 dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'cury': elif dataset == 'cury':
if self.attributes['ndims'] == 3: if self.attributes['ndims'] == 3:
tmp = self.__read_binary_data__('magx', chunk_number) tmp = self.__read_binary_data__('magx', chunk_number)
dset = (numpy.roll(tmp, -1, axis=0) \ dset = (numpy.roll(tmp, -1, axis=1) \
- numpy.roll(tmp, 1, axis=0)) - numpy.roll(tmp, 1, axis=1))
tmp = self.__read_binary_data__('magz', chunk_number) tmp = self.__read_binary_data__('magz', chunk_number)
dset += (numpy.roll(tmp, 1, axis=2) \ dset += (numpy.roll(tmp, 1, axis=3) \
- numpy.roll(tmp, -1, axis=2)) - numpy.roll(tmp, -1, axis=3))
else: else:
tmp = self.__read_binary_data__('magz', chunk_number) tmp = self.__read_binary_data__('magz', chunk_number)
dset = (numpy.roll(tmp, 1, axis=1) \ dset = (numpy.roll(tmp, 1, axis=2) \
- numpy.roll(tmp, -1, axis=1)) - numpy.roll(tmp, -1, axis=2))
dset *= 0.5 dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'curz': elif dataset == 'curz':
if self.attributes['ndims'] == 3: if self.attributes['ndims'] == 3:
tmp = self.__read_binary_data__('magx', chunk_number)
dset = (numpy.roll(tmp, 1, axis=2) \
- numpy.roll(tmp, -1, axis=2))
tmp = self.__read_binary_data__('magy', chunk_number)
dset += (numpy.roll(tmp, -1, axis=3) \
- numpy.roll(tmp, 1, axis=3))
else:
tmp = self.__read_binary_data__('magx', chunk_number) tmp = self.__read_binary_data__('magx', chunk_number)
dset = (numpy.roll(tmp, 1, axis=1) \ dset = (numpy.roll(tmp, 1, axis=1) \
- numpy.roll(tmp, -1, axis=1)) - numpy.roll(tmp, -1, axis=1))
tmp = self.__read_binary_data__('magy', chunk_number) tmp = self.__read_binary_data__('magy', chunk_number)
dset += (numpy.roll(tmp, -1, axis=2) \ dset += (numpy.roll(tmp, -1, axis=2) \
- numpy.roll(tmp, 1, axis=2)) - numpy.roll(tmp, 1, axis=2))
else:
tmp = self.__read_binary_data__('magx', chunk_number)
dset = (numpy.roll(tmp, 1, axis=0) \
- numpy.roll(tmp, -1, axis=0))
tmp = self.__read_binary_data__('magy', chunk_number)
dset += (numpy.roll(tmp, -1, axis=1) \
- numpy.roll(tmp, 1, axis=1))
dset *= 0.5 dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'jvec': elif dataset == 'jvec':
if self.attributes['ndims'] == 3: if self.attributes['ndims'] == 3:
tmp = self.__read_binary_data__('magx', chunk_number) tmp = self.__read_binary_data__('magx', chunk_number)
wy = (numpy.roll(tmp, -1, axis=0) \ wy = (numpy.roll(tmp, -1, axis=1) \
- numpy.roll(tmp, 1, axis=0)) - numpy.roll(tmp, 1, axis=1))
wz = (numpy.roll(tmp, 1, axis=2) \
- numpy.roll(tmp, -1, axis=2))
tmp = self.__read_binary_data__('magy', chunk_number)
wz += (numpy.roll(tmp, -1, axis=3) \
- numpy.roll(tmp, 1, axis=3))
wx = (numpy.roll(tmp, 1, axis=1) \
- numpy.roll(tmp, -1, axis=1))
tmp = self.__read_binary_data__('magz', chunk_number)
wx += (numpy.roll(tmp, -1, axis=2) \
- numpy.roll(tmp, 1, axis=2))
wy += (numpy.roll(tmp, 1, axis=3) \
- numpy.roll(tmp, -1, axis=3))
else:
tmp = self.__read_binary_data__('magx', chunk_number)
wz = (numpy.roll(tmp, 1, axis=1) \ wz = (numpy.roll(tmp, 1, axis=1) \
- numpy.roll(tmp, -1, axis=1)) - numpy.roll(tmp, -1, axis=1))
tmp = self.__read_binary_data__('magy', chunk_number) tmp = self.__read_binary_data__('magy', chunk_number)
wz += (numpy.roll(tmp, -1, axis=2) \ wz += (numpy.roll(tmp, -1, axis=2) \
- numpy.roll(tmp, 1, axis=2)) - numpy.roll(tmp, 1, axis=2))
wx = (numpy.roll(tmp, 1, axis=0) \
- numpy.roll(tmp, -1, axis=0))
tmp = self.__read_binary_data__('magz', chunk_number) tmp = self.__read_binary_data__('magz', chunk_number)
wx += (numpy.roll(tmp, -1, axis=1) \ wx = (numpy.roll(tmp, -1, axis=1) \
- numpy.roll(tmp, 1, axis=1)) - numpy.roll(tmp, 1, axis=1))
wy += (numpy.roll(tmp, 1, axis=2) \ wy = (numpy.roll(tmp, 1, axis=2) \
- numpy.roll(tmp, -1, axis=2)) - numpy.roll(tmp, -1, axis=2))
else:
tmp = self.__read_binary_data__('magx', chunk_number)
wz = (numpy.roll(tmp, 1, axis=0) \
- numpy.roll(tmp, -1, axis=0))
tmp = self.__read_binary_data__('magy', chunk_number)
wz += (numpy.roll(tmp, -1, axis=1) \
- numpy.roll(tmp, 1, axis=1))
tmp = self.__read_binary_data__('magz', chunk_number)
wx = (numpy.roll(tmp, -1, axis=0) \
- numpy.roll(tmp, 1, axis=0))
wy = (numpy.roll(tmp, 1, axis=1) \
- numpy.roll(tmp, -1, axis=1))
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
h = 2 * self.cell_size[self.chunks[chunk_number]['levels'][p]] h = 2 * self.cell_size[self.chunks[chunk_number]['levels'][p]]
wx[...,p] /= h wx[p,...] /= h
wy[...,p] /= h wy[p,...] /= h
wz[...,p] /= h wz[p,...] /= h
dset = [wx, wy, wz] dset = [wx, wy, wz]
elif dataset == 'curr': elif dataset == 'curr':
if self.attributes['ndims'] == 3: if self.attributes['ndims'] == 3:
tmp = self.__read_binary_data__('magx', chunk_number) tmp = self.__read_binary_data__('magx', chunk_number)
wy = (numpy.roll(tmp, -1, axis=0) \ wy = (numpy.roll(tmp, -1, axis=1) \
- numpy.roll(tmp, 1, axis=0)) - numpy.roll(tmp, 1, axis=1))
wz = (numpy.roll(tmp, 1, axis=2) \
- numpy.roll(tmp, -1, axis=2))
tmp = self.__read_binary_data__('magy', chunk_number)
wz += (numpy.roll(tmp, -1, axis=3) \
- numpy.roll(tmp, 1, axis=3))
wx = (numpy.roll(tmp, 1, axis=1) \
- numpy.roll(tmp, -1, axis=1))
tmp = self.__read_binary_data__('magz', chunk_number)
wx += (numpy.roll(tmp, -1, axis=2) \
- numpy.roll(tmp, 1, axis=2))
wy += (numpy.roll(tmp, 1, axis=3) \
- numpy.roll(tmp, -1, axis=3))
else:
tmp = self.__read_binary_data__('magx', chunk_number)
wz = (numpy.roll(tmp, 1, axis=1) \ wz = (numpy.roll(tmp, 1, axis=1) \
- numpy.roll(tmp, -1, axis=1)) - numpy.roll(tmp, -1, axis=1))
tmp = self.__read_binary_data__('magy', chunk_number) tmp = self.__read_binary_data__('magy', chunk_number)
wz += (numpy.roll(tmp, -1, axis=2) \ wz += (numpy.roll(tmp, -1, axis=2) \
- numpy.roll(tmp, 1, axis=2)) - numpy.roll(tmp, 1, axis=2))
wx = (numpy.roll(tmp, 1, axis=0) \
- numpy.roll(tmp, -1, axis=0))
tmp = self.__read_binary_data__('magz', chunk_number) tmp = self.__read_binary_data__('magz', chunk_number)
wx += (numpy.roll(tmp, -1, axis=1) \ wx = (numpy.roll(tmp, -1, axis=1) \
- numpy.roll(tmp, 1, axis=1)) - numpy.roll(tmp, 1, axis=1))
wy += (numpy.roll(tmp, 1, axis=2) \ wy = (numpy.roll(tmp, 1, axis=2) \
- numpy.roll(tmp, -1, axis=2)) - numpy.roll(tmp, -1, axis=2))
else:
tmp = self.__read_binary_data__('magx', chunk_number)
wz = (numpy.roll(tmp, 1, axis=0) \
- numpy.roll(tmp, -1, axis=0))
tmp = self.__read_binary_data__('magy', chunk_number)
wz += (numpy.roll(tmp, -1, axis=1) \
- numpy.roll(tmp, 1, axis=1))
tmp = self.__read_binary_data__('magz', chunk_number)
wx = (numpy.roll(tmp, -1, axis=0) \
- numpy.roll(tmp, 1, axis=0))
wy = (numpy.roll(tmp, 1, axis=1) \
- numpy.roll(tmp, -1, axis=1))
dset = 0.5 * numpy.sqrt(wx**2 + wy**2 + wz**2) dset = 0.5 * numpy.sqrt(wx**2 + wy**2 + wz**2)
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'dbxdx': elif dataset == 'dbxdx':
p = self.attributes['ndims']
tmp = self.__read_binary_data__('magx', chunk_number)
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']):
dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'dbxdy':
p = self.attributes['ndims'] - 1 p = self.attributes['ndims'] - 1
tmp = self.__read_binary_data__('magx', chunk_number) tmp = self.__read_binary_data__('magx', chunk_number)
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p) dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5 dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'dbxdy':
p = self.attributes['ndims'] - 2
tmp = self.__read_binary_data__('magx', chunk_number)
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'dbxdz': elif dataset == 'dbxdz':
p = self.attributes['ndims'] - 3 p = self.attributes['ndims'] - 2
tmp = self.__read_binary_data__('magx', chunk_number) tmp = self.__read_binary_data__('magx', chunk_number)
if p >= 0: if p >= 0:
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p) dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5 dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
else: else:
dset = numpy.zeros_like(tmp) dset = numpy.zeros_like(tmp)
elif dataset == 'dbydx': elif dataset == 'dbydx':
p = self.attributes['ndims']
tmp = self.__read_binary_data__('magy', chunk_number)
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']):
dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'dbydy':
p = self.attributes['ndims'] - 1 p = self.attributes['ndims'] - 1
tmp = self.__read_binary_data__('magy', chunk_number) tmp = self.__read_binary_data__('magy', chunk_number)
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p) dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5 dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'dbydy':
p = self.attributes['ndims'] - 2
tmp = self.__read_binary_data__('magy', chunk_number)
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'dbydz': elif dataset == 'dbydz':
p = self.attributes['ndims'] - 3 p = self.attributes['ndims'] - 2
tmp = self.__read_binary_data__('magy', chunk_number) tmp = self.__read_binary_data__('magy', chunk_number)
if p >= 0: if p >= 0:
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p) dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5 dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
else: else:
dset = numpy.zeros_like(tmp) dset = numpy.zeros_like(tmp)
elif dataset == 'dbzdx': elif dataset == 'dbzdx':
p = self.attributes['ndims']
tmp = self.__read_binary_data__('magz', chunk_number)
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']):
dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'dbzdy':
p = self.attributes['ndims'] - 1 p = self.attributes['ndims'] - 1
tmp = self.__read_binary_data__('magz', chunk_number) tmp = self.__read_binary_data__('magz', chunk_number)
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p) dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5 dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'dbzdy':
p = self.attributes['ndims'] - 2
tmp = self.__read_binary_data__('magz', chunk_number)
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
elif dataset == 'dbzdz': elif dataset == 'dbzdz':
p = self.attributes['ndims'] - 3 p = self.attributes['ndims'] - 2
tmp = self.__read_binary_data__('magz', chunk_number) tmp = self.__read_binary_data__('magz', chunk_number)
if p >= 0: if p >= 0:
dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p) dset = numpy.roll(tmp, -1, axis=p) - numpy.roll(tmp, 1, axis=p)
dset *= 0.5 dset *= 0.5
for p in range(self.chunks[chunk_number]['dblocks']): for p in range(self.chunks[chunk_number]['dblocks']):
dset[...,p] /= self.cell_size[self.chunks[chunk_number]['levels'][p]] dset[p,...] /= self.cell_size[self.chunks[chunk_number]['levels'][p]]
else: else:
dset = numpy.zeros_like(tmp) dset = numpy.zeros_like(tmp)
else: else:
@ -980,7 +981,7 @@ class Amun:
else: else:
method = interpolation method = interpolation
cm = bm[0:self.attributes['ndims']] * nl // shrink cm = bm[0:self.attributes['ndims']] * nl // shrink
il = self.chunks[n]['coords'][:,p][0:self.attributes['ndims']] * cm[0:self.attributes['ndims']] il = self.chunks[n]['coords'][p,:][0:self.attributes['ndims']] * cm[0:self.attributes['ndims']]
iu = il + cm[0:self.attributes['ndims']] iu = il + cm[0:self.attributes['ndims']]
if all(iu[:] > ll[:]) and all(il[:] < uu[:]): if all(iu[:] > ll[:]) and all(il[:] < uu[:]):
@ -994,14 +995,14 @@ class Amun:
if isinstance(dset, (list, tuple)): if isinstance(dset, (list, tuple)):
for di in range(3): for di in range(3):
if self.attributes['ndims'] == 3: if self.attributes['ndims'] == 3:
arr[di][ib[2]:ie[2],ib[1]:ie[1],ib[0]:ie[0]] = interpolate(dset[di][:,:,:,p], cm, self.attributes['nghosts'], method=method, order=order)[jb[2]:je[2],jb[1]:je[1],jb[0]:je[0]] arr[di][ib[2]:ie[2],ib[1]:ie[1],ib[0]:ie[0]] = interpolate(dset[di][p,:,:,:], cm, self.attributes['nghosts'], method=method, order=order)[jb[2]:je[2],jb[1]:je[1],jb[0]:je[0]]
else: else:
arr[di][ib[1]:ie[1],ib[0]:ie[0]] = interpolate(dset[di][:,:,p], cm, self.attributes['nghosts'], method=method, order=order)[jb[1]:je[1],jb[0]:je[0]] arr[di][ib[1]:ie[1],ib[0]:ie[0]] = interpolate(dset[di][p,:,:], cm, self.attributes['nghosts'], method=method, order=order)[jb[1]:je[1],jb[0]:je[0]]
else: else:
if self.attributes['ndims'] == 3: if self.attributes['ndims'] == 3:
arr[ib[2]:ie[2],ib[1]:ie[1],ib[0]:ie[0]] = interpolate(dset[:,:,:,p], cm, self.attributes['nghosts'], method=method, order=order)[jb[2]:je[2],jb[1]:je[1],jb[0]:je[0]] arr[ib[2]:ie[2],ib[1]:ie[1],ib[0]:ie[0]] = interpolate(dset[p,:,:,:], cm, self.attributes['nghosts'], method=method, order=order)[jb[2]:je[2],jb[1]:je[1],jb[0]:je[0]]
else: else:
arr[ib[1]:ie[1],ib[0]:ie[0]] = interpolate(dset[:,:,p], cm, self.attributes['nghosts'], method=method, order=order)[jb[1]:je[1],jb[0]:je[0]] arr[ib[1]:ie[1],ib[0]:ie[0]] = interpolate(dset[p,:,:], cm, self.attributes['nghosts'], method=method, order=order)[jb[1]:je[1],jb[0]:je[0]]
if progress: if progress:
cfile = self.chunkname.format(n) cfile = self.chunkname.format(n)
@ -1057,12 +1058,12 @@ class Amun:
lv = self.chunks[n]['levels'][p] - 1 lv = self.chunks[n]['levels'][p] - 1
center = (self.chunks[n]['bounds'][0,:,p] + self.chunks[n]['bounds'][1,:,p]) / 2 center = (self.chunks[n]['bounds'][p,0,:] + self.chunks[n]['bounds'][p,1,:]) / 2
base.createNodeBranch(center, lv) base.createNodeBranch(center, lv)
if isinstance(dset, (list, tuple)): if isinstance(dset, (list, tuple)):
base.setNodeData(center, lv, (dset[0][ng:-ng,ng:-ng,ng:-ng,p], dset[1][ng:-ng,ng:-ng,ng:-ng,p], dset[2][ng:-ng,ng:-ng,ng:-ng,p])) base.setNodeData(center, lv, (dset[0][p,ng:-ng,ng:-ng,ng:-ng], dset[1][p,ng:-ng,ng:-ng,ng:-ng], dset[2][p,ng:-ng,ng:-ng,ng:-ng]))
else: else:
base.setNodeData(center, lv, dset[ng:-ng,ng:-ng,ng:-ng,p]) base.setNodeData(center, lv, dset[p,ng:-ng,ng:-ng,ng:-ng])
if progress: if progress:
cfile = self.chunkname.format(n) cfile = self.chunkname.format(n)

40
python/amunpy/src/amunpy/amunh5.py
View File

@ -54,6 +54,9 @@ class AmunH5(Amun):
self.dataformat = 'AmunH5' self.dataformat = 'AmunH5'
else: else:
raise Exception("{} misses one of these groups: 'attributes', 'coordinates' or 'variables'!".format(self.path)) raise Exception("{} misses one of these groups: 'attributes', 'coordinates' or 'variables'!".format(self.path))
if 'version' in h5.attrs:
if h5.attrs['version'].dtype == 'float32':
self.version = int(h5.attrs['version'])
def __fill_attributes__(self): def __fill_attributes__(self):
@ -62,6 +65,7 @@ class AmunH5(Amun):
the attributes' dictionary. the attributes' dictionary.
""" """
import h5py import h5py
import numpy
exclude_list = ['nseeds', 'seeds', 'dblocks', 'nproc', 'dims', 'dtn', 'last_id', 'mblocks'] exclude_list = ['nseeds', 'seeds', 'dblocks', 'nproc', 'dims', 'dtn', 'last_id', 'mblocks']
@ -69,14 +73,13 @@ class AmunH5(Amun):
for aname in h5['attributes'].attrs: for aname in h5['attributes'].attrs:
if not aname in exclude_list: if not aname in exclude_list:
attr = h5['attributes'].attrs[aname] attr = h5['attributes'].attrs[aname]
if attr.dtype == 'float64' or attr.dtype == 'float32' or \ if attr.dtype in [ 'float64', 'float32', 'int64', 'int32' ]:
attr.dtype == 'int64' or attr.dtype == 'int32': if isinstance(attr, numpy.ndarray):
if len(attr) > 1: self.attributes[aname] = numpy.squeeze(attr).tolist()
self.attributes[aname] = attr.tolist()
else: else:
self.attributes[aname] = attr[0] self.attributes[aname] = attr
else: else:
self.attributes[aname] = attr[0].astype(str) self.attributes[aname] = numpy.squeeze(attr).astype(str)
if not 'nchunks' in self.attributes and 'nprocs' in self.attributes: if not 'nchunks' in self.attributes and 'nprocs' in self.attributes:
self.attributes['nchunks'] = self.attributes['nprocs'] self.attributes['nchunks'] = self.attributes['nprocs']
@ -116,11 +119,11 @@ class AmunH5(Amun):
cname = os.path.join(self.dirname, self.chunks[n]['filename']) cname = os.path.join(self.dirname, self.chunks[n]['filename'])
if os.path.exists(cname): if os.path.exists(cname):
with h5py.File(cname, 'r') as h5: with h5py.File(cname, 'r') as h5:
self.chunks[n]['dblocks'] = h5['attributes'].attrs['dblocks'][0] self.chunks[n]['dblocks'] = numpy.squeeze(h5['attributes'].attrs['dblocks'])
if self.chunks[n]['dblocks'] > 0: if self.chunks[n]['dblocks'] > 0:
self.chunks[n]['levels'] = numpy.array(h5['coordinates']['levels']) self.chunks[n]['levels'] = numpy.array(h5['coordinates']['levels'])
self.chunks[n]['bounds'] = numpy.array(h5['coordinates']['bounds']) self.chunks[n]['bounds'] = self.__swap__(numpy.array(h5['coordinates']['bounds']))
self.chunks[n]['coords'] = numpy.array(h5['coordinates']['coords']) self.chunks[n]['coords'] = self.__swap__(numpy.array(h5['coordinates']['coords']))
else: else:
self.chunks[n]['levels'] = None self.chunks[n]['levels'] = None
self.chunks[n]['coords'] = None self.chunks[n]['coords'] = None
@ -129,6 +132,19 @@ class AmunH5(Amun):
raise Exception("Snapshot's chunk '{}' not present!".format(cname)) raise Exception("Snapshot's chunk '{}' not present!".format(cname))
def __swap__(self, arr):
"""
Function swaps the array for version before 1.0.
"""
pass
import numpy
if self.version < 1:
return numpy.transpose(arr, numpy.roll(numpy.arange(arr.ndim), 1))
else:
return arr
def __read_binary_data__(self, dataset_name, chunk_number): def __read_binary_data__(self, dataset_name, chunk_number):
""" """
Gets the dataset array from a given snapshot's chunk. Gets the dataset array from a given snapshot's chunk.
@ -137,8 +153,4 @@ class AmunH5(Amun):
cname = os.path.join(self.dirname, self.chunks[chunk_number]['filename']) cname = os.path.join(self.dirname, self.chunks[chunk_number]['filename'])
with h5py.File(cname, 'r') as h5: with h5py.File(cname, 'r') as h5:
dset = numpy.array(h5['variables'][dataset_name]) return self.__swap__(numpy.squeeze(numpy.array(h5['variables'][dataset_name])))
if self.attributes['ndims'] == 3:
return dset
else:
return dset[0,:,:,:]

49
python/amunpy/src/amunpy/amunxml.py
View File

@ -49,6 +49,8 @@ class AmunXML(Amun):
if os.path.exists(mfile): if os.path.exists(mfile):
tree = ET.parse(mfile) tree = ET.parse(mfile)
root = tree.getroot() root = tree.getroot()
if self.version < 0:
self.version = int(float(root.attrib['version']))
if root.tag == 'AMUNFile': if root.tag == 'AMUNFile':
self.dataformat = 'AmunXML' self.dataformat = 'AmunXML'
else: else:
@ -144,43 +146,66 @@ class AmunXML(Amun):
mset = self.__read_binary_meta('fields') mset = self.__read_binary_meta('fields')
n = mset.size // self.attributes['nleafs'] n = mset.size // self.attributes['nleafs']
mset = numpy.reshape(mset, [n,self.attributes['nleafs']]) if self.version > 0:
mset = numpy.reshape(mset, [self.attributes['nleafs'],n])
else:
mset = self.__swap__(numpy.reshape(mset, [n, self.attributes['nleafs']]))
index = dict() index = dict()
level = dict() level = dict()
coord = dict() coord = dict()
for n in range(self.attributes['nleafs']): for n in range(self.attributes['nleafs']):
index[mset[0,n]] = n index[mset[n,0]] = n
level[mset[0,n]] = mset[ 1,n] level[mset[n,0]] = mset[n, 1]
coord[mset[0,n]] = mset[2:5,n] coord[mset[n,0]] = mset[n,2:5]
bounds = self.__read_binary_meta('bounds', dtype='float64') bounds = self.__read_binary_meta('bounds', dtype='float64')
bounds = numpy.reshape(bounds, [2, 3, self.attributes['nleafs']]) if self.version > 0:
bounds = numpy.reshape(bounds, [self.attributes['nleafs'], 2, 3])
else:
bounds = self.__swap__(numpy.reshape(bounds, [2, 3, self.attributes['nleafs']]))
for n in range(self.attributes['nchunks']): for n in range(self.attributes['nchunks']):
if self.chunks[n]['dblocks'] > 0: nblk = self.chunks[n]['dblocks']
if nblk > 0:
ids = self.__read_binary_data('ids', n, dtype='int32') ids = self.__read_binary_data('ids', n, dtype='int32')
self.chunks[n]['levels'] = numpy.array([level[ids[p]] for p in range(self.chunks[n]['dblocks'])]) self.chunks[n]['levels'] = numpy.array([level[ids[p]] for p in range(nblk)])
self.chunks[n]['coords'] = numpy.array([coord[ids[p]] for p in range(self.chunks[n]['dblocks'])]).T self.chunks[n]['coords'] = numpy.array([coord[ids[p]] for p in range(nblk)])
ii = [ index[ids[p]] for p in range(self.chunks[n]['dblocks']) ] ii = [ index[ids[p]] for p in range(nblk) ]
self.chunks[n]['bounds'] = numpy.array([bounds[:,:,ii[p]].T for p in range(self.chunks[n]['dblocks'])]).T self.chunks[n]['bounds'] = numpy.array([bounds[ii[p],:,:] for p in range(nblk)])
else: else:
self.chunks[n]['levels'] = None self.chunks[n]['levels'] = None
self.chunks[n]['coords'] = None self.chunks[n]['coords'] = None
self.chunks[n]['bounds'] = None self.chunks[n]['bounds'] = None
def __swap__(self, arr):
"""
Function swaps the array for version before 1.0.
"""
pass
import numpy
if self.version < 1:
return numpy.transpose(arr, numpy.roll(numpy.arange(arr.ndim), 1))
else:
return arr
def __read_binary_data__(self, dataset_name, chunk_number): def __read_binary_data__(self, dataset_name, chunk_number):
""" """
Gets the dataset array from a given snapshot's chunk. Gets the dataset array from a given snapshot's chunk.
""" """
import numpy import numpy
dset = self.__read_binary_data(dataset_name, chunk_number) dims = self.chunks[chunk_number]['dims']
return numpy.reshape(dset, self.chunks[chunk_number]['dims']) if self.version < 1:
dims = numpy.roll(self.chunks[chunk_number]['dims'], -1)
dset = numpy.reshape(self.__read_binary_data(dataset_name, chunk_number), dims)
return self.__swap__(dset)
def __check_digest(self, filename, digest, data): def __check_digest(self, filename, digest, data):