Add a new rewritten API.

1 files changed, 136 insertions, 0 deletions

diff --git a/simdata.py b/simdata.py
new file mode 100644
index 0000000..dd5398e
--- /dev/null
+++ b/simdata.py
@@ -0,0 +1,136 @@
+"""
+This module provides a wrapper around a Cactus simulation output, i.e. a directory of HDF5 files.
+"""
+
+import os
+
+from . import datafile
+
+
+class _RefinedSlice(object):
+    # simulation time on this slice
+    time = -1
+    # simulation iteration on this slice
+    it   = -1
+
+    # a list of arrays of coordinate points along each direction [x, y, z]
+    c = None
+
+    # arrays of coordinate points along each direction
+    x = None
+    y = None
+    z = None
+
+    rl = None
+
+    # private
+    _sd = None
+    _rl = None
+
+    def __init__(self, sd, rl, time = -1, iteration = -1):
+        self._sd = sd
+        self._rl = rl
+        self.rl  = rl
+
+        if time < 0 and iteration < 0:
+            raise TypeError('Neither time nor iteration provided')
+
+        if time >= 0:
+            self.it = sd._df._iter_from_time(time)
+        else:
+            self.it = iteration
+
+        try:
+            s = sd._df.slice(it = self.it, rl = rl.n)
+            self.time = s.time
+        except KeyError:
+            if time < 0:
+                raise IndexError('No such iteration: ' + str(iteration))
+            else:
+                raise IndexError('No such time: ' + str(time))
+
+        # pick a representative datafile and get the grid properties from it
+        if 'H' in self._sd.df:
+            df = self._sd.df['H']
+        else:
+            df = None
+            for key in self._sd.df:
+                if '.' in key:
+                    continue
+                df = self._sd.df[key]
+                break
+            if df is None:
+                raise IndexError
+
+        ds = df.slice(self.it, rl = rl.n)
+        self.c = ds.layout.coords
+        self.x, self.y, self.z = self.c
+
+    def __getitem__(self, key):
+        if key in self._sd.df:
+            return self._sd.df[key].slice(self.it, rl = self._rl.n)
+        raise IndexError
+
+
+class _RefinementLevel(object):
+    # refinement level number
+    n  = -1
+
+    # private
+    _sd = None
+
+    def __init__(self, sd, n):
+        self._sd = sd
+        self.n   = n
+
+    def slice(self, iteration = -1, time = -1):
+        return _RefinedSlice(self._sd, self, time, iteration)
+
+
+class SimulationData(object):
+    "The full path to the simulation directory"
+    dirname = None
+
+    "A dictionary of { basename : DataFile } items representing the data files. The basename is the file name without the .h5 extension"
+    df = None
+
+    "A dictionary of all the { gridfunction : [datafiles it is located in] } pairs for this set of data"
+    gf = None
+
+    # per-refinement level parameters
+    rl = None
+
+    _df = None
+
+    def __init__(self, dirname):
+        self.dirname = os.path.abspath(dirname)
+        self.df = {}
+        self.gf = {}
+
+        # open all the hdf5 files in the dir
+        for f in os.listdir(dirname):
+            if (not f.endswith('.h5') or f.startswith('checkpoint') or f.endswith('.d.h5')):
+                continue
+
+            self.df[f[:-3]] = datafile.DataFile('%s/%s' % (dirname, f))
+        if len(self.df) == 0:
+            raise ValueError('No HDF5 data files in the directory.')
+
+        for df in self.df.values():
+            for ds in df.datasets:
+                if ds.name in self.gf:
+                    self.gf[ds.name].append(df)
+                self.gf[ds.name] = [df]
+
+        # pick a representative datafile and get the grid properties from it
+        if 'H' in self.df:
+            df = self.df['H']
+        else:
+            df = self.df.values()[0]
+
+        self._df = df
+
+        # get the refinement levels, iterations and times
+        self.rl = []
+        for rl in df.rl:
+            self.rl.append(_RefinementLevel(self, len(self.rl)))