#!/usr/bin/env python
""" Wrapper grid for some OpenGGCM convenience """
# look at what 'copy_on_transform' actually does... the fact
# that it's here is awkward
from __future__ import print_function, division
import os
import re
from operator import itemgetter
import numpy as np
try:
import h5py
HAS_H5PY = True
except ImportError:
HAS_H5PY = False
from viscid.readers.vfile_bucket import ContainerFile
from viscid.readers.hdf5 import FileHDF5, H5pyDataWrapper
from viscid import grid
from viscid import field
from viscid.coordinate import wrap_crds
base_hydro_names = ['rho', 'rhoux', 'rhouy', 'rhouz', 'e']
base_hydro_pretty_names = [r'$\rho$', r'$\rho u_x$', r'$\rho u_y$', r'$\rho u_z$',
r'$e$']
base_5m_names = ['rho_e', 'rhoux_e', 'rhouy_e', 'rhouz_e', 'e_e',
'rho_i', 'rhoux_i', 'rhouy_i', 'rhouz_i', 'e_i',
'Ex', 'Ey', 'Ez', 'Bx', 'By', 'Bz']
base_5m_pretty_names = [
r'$\rho_e$', r'$\rho u_{x,e}$', r'$\rho u_{y,e}$', r'$\rho u_{z,e}$', r'$e_e$',
r'$\rho_i$', r'$\rho u_{x,i}$', r'$\rho u_{y,i}$', r'$\rho u_{z,i}$', r'$e_i$',
r'$E_x$', r'$E_y$', r'$E_z$', r'$B_x$', r'$B_y$', r'$B_z$']
base_10m_names = ['rho_e', 'rhoux_e', 'rhouy_e', 'rhouz_e',
'pxx_e', 'pxy_e', 'pxz_e', 'pyy_e', 'pyz_e', 'pzz_e',
'rho_i', 'rhoux_i', 'rhouy_i', 'rhouz_i',
'pxx_i', 'pxy_i', 'pxz_i', 'pyy_i', 'pyz_i', 'pzz_i',
'Ex', 'Ey', 'Ez', 'Bx', 'By', 'Bz']
base_10m_pretty_names = [
r'$\rho_e$', r'$\rho u_{x,e}$', r'$\rho u_{y,e}$', r'$\rho u_{z,e}$',
r'$\mathcal{P}_{xx,e}$', r'$\mathcal{P}_{xy,e}$', r'$\mathcal{P}_{xz,e}$',
r'$\mathcal{P}_{yy,e}$', r'$\mathcal{P}_{yz,e}$', r'$\mathcal{P}_{zz,e}$',
r'$\rho_i$', r'$\rho u_{x,i}$', r'$\rho u_{y,i}$', r'$\rho u_{z,i}$',
r'$\mathcal{P}_{xx,i}$', r'$\mathcal{P}_{xy,i}$', r'$\mathcal{P}_{xz,i}$',
r'$\mathcal{P}_{yy,i}$', r'$\mathcal{P}_{yz,i}$', r'$\mathcal{P}_{zz,i}$',
r'$E_x$', r'$E_y$', r'$E_z$', r'$B_x$', r'$B_y$', r'$B_z$']
_type_info = {len(base_hydro_names): {'field_type': 'hydro',
'names': base_hydro_names,
'pretty_names': base_hydro_pretty_names},
len(base_5m_names): {'field_type': 'five-moment',
'names': base_5m_names,
'pretty_names': base_5m_pretty_names},
len(base_10m_names): {'field_type': 'ten-moment',
'names': base_10m_names,
'pretty_names': base_10m_pretty_names}}
[docs]def create_type_info(nr_fields):
type_info = dict()
type_info['field_type'] = 'unknown'
type_info['names'] = []
type_info['pretty_names'] = []
for ncomp in range(nr_fields):
type_info['names'].append('Q%d'%ncomp)
type_info['pretty_names'].append(r'$\rm{Q%d}$'%ncomp)
return type_info
[docs]class GkeyllGrid(grid.Grid):
""""""
def _assemble_vector(self, base_name, comp_names="xyz", suffix="",
forget_source=False, **kwargs):
opts = dict(forget_source=forget_source, **kwargs)
# caching behavior depends on self.longterm_field_caches
comps = [self[base_name + c + suffix] for c in comp_names]
return field.scalar_fields_to_vector(comps, name=base_name, **opts)
def _get_ux_e(self):
ux_e = self['rhoux_e'] / self['rho_e']
ux_e.name = 'ux_e'
ux_e.pretty_name = r'$u_{x,e}$'
return ux_e
def _get_uy_e(self):
uy_e = self['rhouy_e'] / self['rho_e']
uy_e.name = 'uy_e'
uy_e.pretty_name = r'$u_{y,e}$'
return uy_e
def _get_uz_e(self):
uz_e = self['rhouz_e'] / self['rho_e']
uz_e.name = 'uz_e'
uz_e.pretty_name = r'$u_{z,e}$'
return uz_e
def _get_Pxx_e(self):
Pxx_e = self['pxx_e'] - self['rhoux_e']*self['rhoux_e'] / self['rho_e']
Pxx_e.name = 'Pxx_e'
Pxx_e.pretty_name = r'$P_{xx,e}$'
return Pxx_e
def _get_Pyy_e(self):
Pyy_e = self['pyy_e'] - self['rhoux_e']*self['rhoux_e'] / self['rho_e']
Pyy_e.name = 'Pyy_e'
Pyy_e.pretty_name = r'$P_{yy,e}$'
return Pyy_e
def _get_Pzz_e(self):
Pzz_e = self['pzz_e'] - self['rhoux_e']*self['rhoux_e'] / self['rho_e']
Pzz_e.name = 'Pzz_e'
Pzz_e.pretty_name = r'$P_{zz,e}$'
return Pzz_e
def _get_Pxy_e(self):
Pxy_e = self['pxy_e'] - self['rhoux_e']*self['rhouy_e'] / self['rho_e']
Pxy_e.name = 'Pxy_e'
Pxy_e.pretty_name = r'$P_{xy,e}$'
return Pxy_e
def _get_Pxz_e(self):
Pxz_e = self['pxz_e'] - self['rhoux_e']*self['rhouz_e'] / self['rho_e']
Pxz_e.name = 'Pxz_e'
Pxz_e.pretty_name = r'$P_{xz,e}$'
return Pxz_e
def _get_Pyz_e(self):
Pyz_e = self['pyz_e'] - self['rhouy_e']*self['rhouz_e'] / self['rho_e']
Pyz_e.name = 'Pyz_e'
Pyz_e.pretty_name = r'$P_{yz,e}$'
return Pyz_e
def _get_ux_i(self):
ux_i = self['rhoux_i'] / self['rho_i']
ux_i.name = 'ux_i'
ux_i.pretty_name = r'$u_{x,i}$'
return ux_i
def _get_uy_i(self):
uy_i = self['rhouy_i'] / self['rho_i']
uy_i.name = 'uy_i'
uy_i.pretty_name = r'$u_{y,i}$'
return uy_i
def _get_uz_i(self):
uz_i = self['rhouz_i'] / self['rho_i']
uz_i.name = 'uz_i'
uz_i.pretty_name = r'$u_{z,i}$'
return uz_i
def _get_Pxx_i(self):
Pxx_i = self['pxx_i'] - self['rhoux_i']*self['rhoux_i'] / self['rho_i']
Pxx_i.name = 'Pxx_i'
Pxx_i.pretty_name = r'$P_{xx,i}$'
return Pxx_i
def _get_Pyy_i(self):
Pyy_i = self['pyy_i'] - self['rhoux_i']*self['rhoux_i'] / self['rho_i']
Pyy_i.name = 'Pyy_i'
Pyy_i.pretty_name = r'$P_{yy,i}$'
return Pyy_i
def _get_Pzz_i(self):
Pzz_i = self['pzz_i'] - self['rhoux_i']*self['rhoux_i'] / self['rho_i']
Pzz_i.name = 'Pzz_i'
Pzz_i.pretty_name = r'$P_{zz,i}$'
return Pzz_i
def _get_Pxy_i(self):
Pxy_i = self['pxy_i'] - self['rhoux_i']*self['rhouy_i'] / self['rho_i']
Pxy_i.name = 'Pxy_i'
Pxy_i.pretty_name = r'$P_{xy,i}$'
return Pxy_i
def _get_Pxz_i(self):
Pxz_i = self['pxz_i'] - self['rhoux_i']*self['rhouz_i'] / self['rho_i']
Pxz_i.name = 'Pxz_i'
Pxz_i.pretty_name = r'$P_{xz,i}$'
return Pxz_i
def _get_Pyz_i(self):
Pyz_i = self['pyz_i'] - self['rhouy_i']*self['rhouz_i'] / self['rho_i']
Pyz_i.name = 'Pyz_i'
Pyz_i.pretty_name = r'$P_{yz,i}$'
return Pyz_i
def _get_u_e(self):
# get from [ux_e, uy_e, uz_e]
return self._assemble_vector("u", suffix='_e',
_force_layout=self.force_vector_layout,
pretty_name='u_e')
def _get_u_i(self):
# get from [ux_i, uy_i, uz_i]
return self._assemble_vector("u", suffix='_i',
_force_layout=self.force_vector_layout,
pretty_name='u_i')
def _get_b(self):
# get from [Bx, By, Bz]
return self._assemble_vector("B", _force_layout=self.force_vector_layout,
pretty_name="b")
[docs]class GkeyllFile(FileHDF5, ContainerFile): # pylint: disable=abstract-method
""""""
_detector = r"^\s*(.*)_(q)_([0-9]+).(h5)\s*$"
_grid_type = GkeyllGrid
SAVE_ONLY = False
_fwrapper = None
_crds = None
_fld_templates = None
def __init__(self, fname, crds=None, fld_templates=None, **kwargs):
assert HAS_H5PY
self._crds = crds
self._fld_templates = fld_templates
super(GkeyllFile, self).__init__(fname, **kwargs)
[docs] @classmethod
def group_fnames(cls, fnames):
"""Group File names
The default implementation just returns fnames, but some file
types might do something fancy here
Parameters:
fnames (list): names that can be logically grouped, as in
a bunch of file names that are different time steps
of a given run
Returns:
A list of things that can be given to the constructor of
this class
"""
infolst = []
for name in fnames:
m = re.match(cls._detector, name)
grps = m.groups()
d = dict(runname=grps[0], ftype=grps[1],
fname=m.string)
try:
d["frame"] = int(grps[2])
except (TypeError, ValueError):
# grps[2] is none for "RUN.3d.xdmf" files
d["frame"] = -1
infolst.append(d)
# careful with sorting, only consecutive files will be grouped
infolst.sort(key=itemgetter("frame"))
infolst.sort(key=itemgetter("ftype"))
infolst.sort(key=itemgetter("runname"))
info_groups = []
info_group = [infolst[0]]
for info in infolst[1:]:
last = info_group[-1]
if info['runname'] == last['runname'] and \
info['ftype'] == last['ftype']:
info_group.append(info)
else:
info_groups.append(info_group)
info_group = [info]
info_groups.append(info_group)
# turn info_groups into groups of just file names
groups = []
for info_group in info_groups:
groups.append([info['fname'] for info in info_group])
return groups
[docs] @classmethod
def collective_name_from_group(cls, fnames):
fname0 = fnames[0]
basename = os.path.basename(fname0)
run = re.match(cls._detector, basename).group(1)
fldtype = re.match(cls._detector, basename).group(2)
new_basename = "{0}.{1}.STAR.h5".format(run, fldtype)
return os.path.join(os.path.dirname(fname0), new_basename)
[docs] def get_file_wrapper(self, filename):
if self._fwrapper is None:
# self._fwrapper = GGCMFortbinFileWrapper(filename)
return h5py.File(filename, 'r')
else:
raise NotImplementedError()
# assert (self._fwrapper.filename == filename or
# glob2(self._fwrapper.filename) == glob2(filename))
# return self._fwrapper
[docs] def set_file_wrapper(self, wrapper):
raise NotImplementedError("This must be done at file init")
[docs] def load(self, fname):
if isinstance(fname, list):
self._collection = fname
else:
self._collection = [fname]
fname0 = self._collection[0]
fname1 = self.collective_name(fname)
basename = os.path.basename(fname0)
self.set_info('run', re.match(self._detector, basename).group(1))
self.set_info('field_type', re.match(self._detector, basename).group(2))
super(GkeyllFile, self).load(fname1)
def _parse(self):
if len(self._collection) == 1:
# load a single file
if self._crds is None:
self._crds = self.make_crds(self.fname)
_grid = self._parse_file(self.fname, self)
self.add(_grid)
self.activate(0)
else:
# load each file, and add it to the bucket
if self._crds is None:
self._crds = self.make_crds(self._collection[0])
data_temporal = self._make_dataset(self, dset_type="temporal",
name="GkeyllTemporalCollection")
self._fld_templates = self._make_template(self._collection[0])
for fname in self._collection:
f = self._load_child_file(fname, index_handle=False,
file_type=type(self),
crds=self._crds,
fld_templates=self._fld_templates)
data_temporal.add(f)
data_temporal.activate(0)
self.add(data_temporal)
self.activate(0)
def _parse_file(self, filename, parent_node):
# we do minimal file parsing here for performance. we just
# make data wrappers from the templates we got from the first
# file in the group, and package them up into grids
# frame = int(re.match(self._detector, filename).group(3))
_grid = self._make_grid(parent_node, name="<GkeyllGrid>",
**self._grid_opts)
# FIXME: To get the time at load, we have to open all hdf5 files
# which defeats the purpose of making templates etc. in attempt to
# be lazy. Maybe there's a way to use frame above?
with h5py.File(filename, 'r') as f:
step = f['timeData'].attrs['vsStep']
time = f['timeData'].attrs['vsTime']
self.set_info("step", step)
self.time = time
_grid.time = time
_grid.set_crds(self._crds)
if self._fld_templates is None:
self._fld_templates = self._make_template(filename)
for i, meta in enumerate(self._fld_templates):
# FIXME: the transpose goes xyz -> zyx
h5_data = H5pyDataWrapper(self.fname, "StructGridField",
comp_dim=-1, comp_idx=i)
fld = field.wrap_field(h5_data, self._crds, meta['fld_name'],
center="cell", pretty_name=meta['pretty_name'])
_grid.add_field(fld)
_grid.set_info('field_type', self.get_info('field_type'))
return _grid
def _make_template(self, filename):
""""""
with self.get_file_wrapper(filename) as f:
shape = f["StructGridField"].shape
sshape = shape[:-1]
nr_fields = shape[-1]
try:
type_info = _type_info[nr_fields]
except KeyError:
try:
# if the two scalar corrections potentials are stored
type_info = dict(_type_info[nr_fields - 2])
type_info['names'].append('phi_E')
type_info['pretty_names'].append(r'$\Phi_E$')
type_info['names'].append('phi_B')
type_info['pretty_names'].append(r'$\Phi_B$')
except:
type_info = create_type_info(nr_fields)
template = []
# TODO: use nr_fields to figure out the names of the fields?
for i in range(nr_fields):
d = dict(fldnum=i, shape=sshape,
fld_name=type_info['names'][i],
pretty_name=type_info['pretty_names'][i])
template.append(d)
self.set_info("field_type", type_info['field_type'])
return template
@staticmethod
def _get_single_crd(h5file, idx, nr_crds):
gridType = h5file['StructGrid'].attrs['vsKind'].decode()
if gridType in ['uniform']:
if idx >= len(h5file['StructGrid'].attrs['vsNumCells']):
raise IndexError()
lower = h5file['StructGrid'].attrs['vsLowerBounds'][idx]
upper = h5file['StructGrid'].attrs['vsUpperBounds'][idx]
num = h5file['StructGrid'].attrs['vsNumCells'][idx]
return [lower, upper, num + 1]
elif gridType in ['rectilinear']:
crd_arr = h5file['StructGrid/axis%d'%idx][:]
return crd_arr
elif gridType in ['structured']:
if idx == 0:
crd_arr = h5file['StructGrid'][:, 0, 0, 0]
elif idx == 1:
crd_arr = h5file['StructGrid'][0, :, 0, 1]
elif idx == 2:
crd_arr = h5file['StructGrid'][0, 0, :, 2]
return crd_arr
else:
raise RuntimeError("Gkeyll StructGrid.vsKind not understood: {0}"
"".format(repr(gridType)))
[docs] def make_crds(self, fname):
with h5py.File(fname, 'r') as f:
clist = []
# FIXME: xyz
crds = "xyz"
nr_crds = len(f['StructGridField'].shape) - 1
for i in range(nr_crds):
try:
clist.append((crds[i], self._get_single_crd(f, i, nr_crds)))
except IndexError:
pass
if f['StructGrid'].attrs['vsKind'].decode() in ['uniform']:
# FIXME: this goes xyz -> zyx
crds = wrap_crds("uniform_cartesian", clist)
else:
crds = wrap_crds("nonuniform_cartesian", clist)
return crds
##
## EOF
##