""" This module provides a wrapper around a single HDF5 file written by Cactus (DataFile) and some helper classes for accessing the data. """ import h5py import numpy as np class _SliceLayout(object): ### public ### nb_dims = None coords = None ### private ### _origin = None _delta = None _shape = None def __init__(self, origin, delta, shape): if (len(origin) <= 0 or len(origin) != len(delta) or len(origin) != len(shape)): raise ValueError("Invalid slice description") self.nb_dims = len(origin) self._origin = origin self._delta = delta self._shape = shape self.coords = [] for i in range(self.nb_dims): self.coords.append(np.linspace(origin[i], origin[i] + delta[i] * shape[i], shape[i], endpoint = False)) class _DataSlice(object): layout = None data = None it = None time = None def __init__(self, data, layout, it, time): self.layout = layout self.data = data self.it = it self.time = time def __getitem__(self, key): return self.data[key] class _DataRefinementLevel(object): n = None ### private ### _ds = None """number of iterations per time step""" _stepsize_iter = None def __init__(self, n, ds): self._ds = ds self.n = n if ds._df._stepsize_iter is not None: self._stepsize_iter = ds._df._stepsize_iter / (1 << n) def __repr__(self): return 'rl%d@<%s>' % (self.n, self._ds.name) def __iter__(self): for it, t in self._ds._df.itertimes[self.n]: yield self._ds.slice(it, rl = self.n) def __reversed__(self): for it, t in reversed(self._ds._df.itertimes[self.n]): yield self._ds.slice(it, rl = self.n) #def _next(self): class _DataSet(object): name = None nb_components = None rl = None ### private ### _df = None _need_c = False def __init__(self, df, name): self._df = df self.name = name # detect the number of components first_iter = self._df.itertimes[0][0][0] try: self._slice_single_component(first_iter, 0) self.nb_components = 1 except KeyError: self._need_c = True c = 0 while True: try: self._slice_single_component(first_iter, 0, c) c += 1 except KeyError: break self.nb_components = c # detect the number of refinement levels self.rl = [] while True: try: self._slice_single_component(first_iter, len(self.rl)) self.rl.append(_DataRefinementLevel(len(self.rl), self)) except KeyError: break def __repr__(self): return '<%s>@<%s>' % (self.name, self._df.path) def _slice_single_component(self, it, rl, component = None): querystr = b'%s it=%d tl=0 rl=%d' % (self.name.encode('ascii'), it, rl) if component is None and self._need_c: component = 0 if component is not None: querystr += b' c=%d' % component data = self._df._f[querystr] layout = _SliceLayout(data.attrs['origin'], data.attrs['delta'], data.shape[::-1]) return _DataSlice(data, layout, it, data.attrs['time']) def _slice_merge_components(self, it, rl): data = [] for i in range(self.nb_components): data.append(self._slice_single_component(it, rl, i)) nb_dims = data[0].layout.nb_dims # get the spatial step sizes step = [None] * nb_dims for i in range(nb_dims): if len(data[0].layout.coords[i]) > 1: step[i] = abs(data[0].layout.coords[i][1] - data[0].layout.coords[i][0]) else: step[i] = 1.0 # calculate the extents of the result coord_max = [None] * nb_dims coord_min = [None] * nb_dims for d in data: for i in range(nb_dims): minval = d.layout.coords[i][0] maxval = d.layout.coords[i][-1] if coord_max[i] is None or coord_max[i] < maxval: coord_max[i] = maxval if coord_min[i] is None or coord_min[i] > minval: coord_min[i] = minval shape = [None] * nb_dims for i in range(nb_dims): shape[i] = int((coord_max[i] - coord_min[i]) / step[i]) + 1 res = np.empty(shape) for d in data: sl = [None] * nb_dims d_shape = d.data.shape[::-1] for i in range(nb_dims): start_idx = int((d.layout.coords[i][0] - coord_min[i]) / step[i]) sl[i] = slice(start_idx, start_idx + d_shape[i]) sl = tuple(sl[::-1]) res[sl] = d.data[:] layout = _SliceLayout(coord_min, step, res.shape[::-1]) return _DataSlice(res, layout, it, data[0].time) def slice(self, it = None, time = None, rl = None, component = None): if it is None and time is not None: it = self._df._iter_from_time(time) if component is not None or self.nb_components == 1: return self._slice_single_component(it, rl, component) return self._slice_merge_components(it, rl) class DataFile(object): ### public ### path = None datasets = None ### private ### # the h5py file object _f = None _closed = False # time duration of one iteration _dt_iter = None # number of iterations per coarsest-level time step _stepsize_iter = None _itertimes = None def __init__(self, path, mode = 'r'): self._f = h5py.File(path, mode) self.path = path try: for it in self._f.values(): if 'time' in it.attrs and it.attrs['time'] > 0.0: self._dt_iter = it.attrs['time'] / it.attrs['timestep'] break dt = self._f['Parameters and Global Attributes'].attrs['carpet_delta_time'] try: self._stepsize_iter = self._iter_from_time(dt) except ValueError: pass # retrieve the list of the datasets present in the file datasets_str = np.array(self._f['Parameters and Global Attributes']['Datasets']).item().decode('ascii') self.datasets = [] for ds in datasets_str.strip().split(): self.datasets.append(_DataSet(self, ds)) except: self.close() raise def _iter_from_time(self, time): if self._dt_iter is None: if time == 0.0: return 0 raise ValueError("No time step defined in the file: %s" % self.path) return int(time / self._dt_iter) def close(self): if self._f and not self._closed: self._f.close() self._closed = True def _get_itertimes(self): rls = [] datasets = list(map(lambda ds: (ds.attrs['level'], ds.attrs['timestep'], ds.attrs['time']), filter(lambda x: 'time' in x.attrs, self._f.values()))) i = 0 while True: ds_rl = list(map(lambda ds: (ds[1], ds[2]), filter(lambda x: x[0] == i, datasets))) if len(ds_rl) == 0: break rls.append(sorted(ds_rl, key = lambda x: x[0])) i += 1 return rls @property def itertimes(self): if not self._itertimes: self._itertimes = self._get_itertimes() return self._itertimes @property def rl(self): return self.datasets[0].rl def slice(self, *args, **kwargs): return self.datasets[0].slice(*args, **kwargs)