Add wrapper classes for data files and gridfunctions

1 files changed, 179 insertions, 0 deletions

diff --git a/cactus_utils.py b/cactus_utils.py
index af5959e..14fc40c 100644
--- a/cactus_utils.py
+++ b/cactus_utils.py
@@ -976,6 +976,35 @@ def l2norm(sd, it, gf, maxrl = None, minrl = None):
     return sum(results)
 
 
+class Dataset:
+    """
+    A wrapper around the h5py Dataset that holds a reference to the underlying
+    file and closes it on destruct.
+    """
+    ### private ###
+    # the wrapped dataset
+    _dataset = None
+    # the underlying File object
+    _file = None
+
+    def __init__(self, dataset, file):
+        self.__dict__['_dataset'] = dataset
+        self.__dict__['_file']    = file
+
+    def __getattr__(self, name):
+        attr = self._dataset.__getattribute__(name)
+        if callable(attr):
+            def call_wrapper(*args, **kwargs):
+                return attr(*args, **kwargs)
+            return call_wrapper
+        else:
+            return attr
+
+    def __del__(self):
+        if self._file is not None:
+            self._file.close()
+
+
 class RefinedSlice:
     # simulation time on this slice
     time = -1
@@ -1084,6 +1113,156 @@ class RefinementLevel:
                 self.dx[0], self.dx[1], self.dx[2], self.dt)
 
 
+class GridFunction:
+    """
+    A wrapper around a grid function
+    """
+    # name of the function (without the thorn prefix)
+    name  = None
+    # the thorn this function belongs to
+    thorn = None
+
+    # A list of all the refinement levels on which this
+    # gridfunction has values, starting with reflevel 0.
+    # Each entry is a list of (iteration, time) tuples sorted in chronological
+    # order.
+    reflevels = None
+
+    ### private ###
+    # the data file this function is located in
+    _df = None
+
+    # Keys into the data file used for extracting the data.
+    # The value is a list running over reflevels, each list entry is a dict with
+    # keys being iterations (as integers) and values are the strings for
+    # indexing the datafile with
+    _keys = None
+
+    def __init__(self, datafile, fileobj, name):
+        self._df = datafile
+
+        self.thorn, self.name = name.split('::')
+
+        reflevels = {}
+
+        # extract identifiers of all the slices on which this function has
+        # values
+        for key in fileobj:
+            val = fileobj[key]
+            if not 'name' in val.attrs or val.attrs['name'] != name:
+                continue
+
+            level = val.attrs['level']
+            if not level in reflevels:
+                reflevels[level] = {}
+
+            reflevels[level][val.attrs['timestep']] = (val.attrs['time'], key)
+
+        # parse the extracted data
+        nb_reflevels   = max(reflevels.keys()) + 1
+        self.reflevels = [None] * nb_reflevels
+        self._keys     = [None] * nb_reflevels
+
+        for i in xrange(nb_reflevels):
+            self.reflevels[i] = []
+            self._keys[i]     = {}
+
+            rl = reflevels[i]
+
+            for iteration in sorted(rl.keys()):
+                time, key = rl[iteration]
+                self.reflevels[i].append((iteration, time))
+                self._keys[i][iteration] = key
+
+    def __repr__(self):
+        return '<Cactus grid function \'%s\'>' % '::'.join((self.thorn, self.name))
+
+    def _iteration_from_time(self, reflevel, time):
+        rl = self.reflevels[reflevel]
+        for it, t in rl:
+            if t == time:
+                return it
+        raise IndexError("No iteration for time %g" % time)
+
+    def data(self, reflevel, iteration = None, time = None):
+        """
+        Get the data for this gridfunction on the specified refinement level and
+        iteration or time.
+        """
+        if iteration is None and time is None:
+            raise ValueError("Neither time nor iteration specified")
+
+        rl = self.reflevels[reflevel]
+        if time is not None:
+            iteration = self._iteration_from_time(reflevel, time)
+        elif iteration == "last":
+            iteration = rl[-1][0]
+        elif iteration == "first":
+            iteration = rl[0][0]
+
+        key = self._keys[reflevel][iteration]
+        return self._df[key]
+
+
+class GridFunctionDict(dict):
+    """
+    A thin wrapper around a dict allowing to match gridfunctions by name only
+    (without the thorn prefix)
+    """
+    def __getitem__(self, key):
+        try:
+            return super(GridFunctionDict, self).__getitem__(key)
+        except KeyError:
+            for gf in self.values():
+                if gf.name == key:
+                    return gf
+            raise KeyError(key)
+
+    def __contains__(self, key):
+        try:
+            self[key]
+            return True
+        except KeyError:
+            return False
+
+
+class DataFile:
+    """
+    A wrapper around a single HDF data file.
+    """
+    # a dict of all gridfunctions, indexed by the full gridfunction name
+    # (including the thorn prefix)
+    gridfunctions = None
+
+    ### private ###
+    # full path to the file
+    _path = None
+
+    _PARPREFIX = 'Parameters and Global Attributes'
+
+    def __init__(self, path):
+        self._path        = path
+
+        self.gridfunctions = GridFunctionDict()
+
+        with h5py.File(path, 'r') as f:
+            for gf in f[self._PARPREFIX]['Datasets'].value.split():
+                self.gridfunctions[gf] = GridFunction(self, f, gf)
+
+    def __getitem__(self, key):
+        f = None
+        try:
+            f = h5py.File(self._path, 'r')
+            return Dataset(f[key], f)
+        except:
+            if f is not None:
+                f.close()
+            raise
+
+    def __repr__(self):
+        return '%s.DataFile(\'%s\')' % (self.__module__, self._path)
+
+
 class SimulationData:
     # directory name
     dirname = None