Experimental script to extract Orion AMR output

orion_to_hdf5.py

import numpy as np
import h5py


def parse_multi_tuple(string):
    string = string.replace(' ', '')
    string = string.replace(')(', '), (')
    return eval(string)


class Grid(object):

    def __init__(self):
        self.xmin, self.xmax = None, None
        self.ymin, self.ymax = None, None
        self.zmin, self.zmax = None, None


class Fab(object):

    def __init__(self):
        self.imin, self.imax, self.itype = None, None, None
        self.jmin, self.jmax, self.jtype = None, None, None
        self.kmin, self.kmax, self.ktype = None, None, None
        self.xmin, self.xmax = None, None
        self.ymin, self.ymax = None, None
        self.zmin, self.zmax = None, None
        self.data = {}

    def read_data(self, filename, offset, components):

        print "reading ", filename

        fabsize = self.nx * self.ny * self.nz

        f = open(filename, 'rb')
        f.seek(offset)

        header = f.readline().strip()

        p1 = header.find('((') + 2
        p2 = header.find(', ', p1)
        n_bytes = int(header[p1:p2])

        p3 = header.find('(', p2) + 1
        p4 = header.find('))', p3)
        bytes = [int(x) for x in header[p3:p4].split()]

        p5 = header.find('(', p4) + 1
        p6 = header.find(', ', p5)
        n_bytes = int(header[p5:p6])

        p7 = header.find('(', p6) + 1
        p8 = header.find('))', p7)
        bytes = [int(x) for x in header[p7:p8].split()]

        if bytes == range(1, n_bytes + 1):
            endian = '>'
        elif bytes == range(n_bytes, 0, -1):
            endian = '<'
        else:
            raise Exception("Unexpected bytes: %s" % str(bytes))

        n_components = int(header.split()[-1])

        pos = f.tell()

        for component in components:
            f.seek(pos + component[0] * n_bytes * fabsize)
            array = np.fromstring(f.read()[:n_bytes * fabsize],
                                  dtype='%sf%i' % (endian, n_bytes))
            self.data[component[1]] = array.reshape(self.nx, self.ny, self.nz)


class Level(object):

    def __init__(self):

        # The level number
        self.level = 0

        # Minimum and maximum position in index space
        self.imin, self.imax = None, None
        self.jmin, self.jmax = None, None
        self.kmin, self.kmax = None, None

        # The grids
        self.grids = []

        # The fabs
        self.fabs = []


class AMR(object):

    def __init__(self, dirname, components=['density']):

        # Open file
        f = file('%s/Header' % dirname, 'rb')

        # Read version number
        version = f.readline().strip()

        # Read number of components
        n_quantities = int(f.readline().strip())

        # Read in component names
        quantities = [f.readline().strip() for i in range(n_quantities)]

        # Make list of wanted quantities, and their indices
        quantities = [(quantities.index(c), c) for c in components]

        # Read in number of dimensions
        ndim = int(f.readline().strip())
        if ndim != 3:
            raise Exception("Number of dimensions is not 3")

        # Read in time
        creation_time = float(f.readline().strip())

        # Read in maximum level of refinement
        max_level = int(f.readline().strip())

        # Create list of levels
        self.levels = [Level() for i in range(max_level + 1)]

        # Read in position of box corners
        self.xmin, self.ymin, self.zmin = [float(x) for x in f.readline().strip().split()]
        self.xmax, self.ymax, self.zmax = [float(x) for x in f.readline().strip().split()]

        # Read in refinement ratios
        self.refinement_ratios = [int(x) for x in f.readline().strip().split()]

        # Read in next line
        line = f.readline().strip()

        # Split into groups of ndim values
        elements = line.replace(' ', '').replace('((', '(').replace('))', ')')[1:-1].split(')(')
        for level in self.levels:
            level.idxlo = [int(x) for x in elements[3 * i].split(', ')]
            level.idxhi = [int(x) for x in elements[3 * i + 1].split(', ')]
            level.periodicity = [int(x) for x in elements[3 * i + 2].split(', ')]

        # Read in number of steps on each level
        levelsteps = [int(x) for x in f.readline().strip().split()]

        # Read in grid spacing on each level
        # gridspacing = np.zeros((self.ndim, self.maxlevel+1))
        # for level in self.levels:
        #     level.gridspacing = [float(x) for x in f.readline().strip().split()]
        for level in self.levels:
            f.readline()

        # Read in coordinate type
        coordtype = int(f.readline().strip())
        if coordtype != 0:
            raise Exception("coordtype should be zero")

        # Skip dummy line
        f.readline()

        # First part done. Now need to read in individual levels and grids

        # Initialize list of levels

        # Loop through levels
        for level in self.levels:

            level_num, nfabs, creation_time = f.readline().strip().split()
            level.number = int(level_num)
            nfabs = int(nfabs)

            # Initialize grids
            level.fabs = [Fab() for ifab in range(nfabs)]

            levelsteps = int(f.readline().strip())

            for fab in level.fabs:
                fab.xmin, fab.xmax = [float(x) for x in f.readline().split()]
                fab.ymin, fab.ymax = [float(y) for y in f.readline().split()]
                fab.zmin, fab.zmax = [float(z) for z in f.readline().split()]

            n_quantities_check = 0
            nfiles = 0
            nfilecomp = []

            # Read filename header file
            fname = f.readline().strip()

            fh = open("%s/%s_H" % (dirname, fname))

            fh.readline()
            fh.readline()

            # Read the number of components in multifab files
            nfabcomp = int(fh.readline())

            if nfabcomp != n_quantities:
                raise Exception("Only some of the components included in multifab file")

            fh.readline()

            # Read the number of boxes
            nfabs_check = int(fh.readline().strip()[1:].split()[0])

            print nfabs_check, nfabs, type(nfabs_check), type(nfabs)

            if nfabs_check != nfabs:
                raise Exception("Number of fabs in multifab file does not match known number")

            # Loop through the fabs
            for fab in level.fabs:
                values = parse_multi_tuple(fh.readline())
                fab.imin, fab.jmin, fab.kmin = values[0]
                fab.imax, fab.jmax, fab.kmax = values[1]
                fab.itype, fab.jtype, fab.ktype = values[2]
                fab.nx = fab.imax - fab.imin + 1
                fab.ny = fab.jmax - fab.jmin + 1
                fab.nz = fab.kmax - fab.kmin + 1

            fh.readline()
            fh.readline()

            for fab in level.fabs:
                string = fh.readline().split(':')[1]
                filename = "%s/Level_%i/%s" % (dirname, level.number, string.split()[0].strip())
                offset = int(string.split()[1])
                fab.read_data(filename, offset, quantities)

    def tohdf5(self, filename):

        f = h5py.File(filename, 'w')

        g_geo = f.create_group("Geometry")
        g_geo.attrs['nlevels'] = len(self.levels)
        g_phy = f.create_group("Physics")
        g_phy.attrs['nlevels'] = len(self.levels)

        for ilevel, level in enumerate(self.levels):

            gl_geo = g_geo.create_group("Level %i" % ilevel)
            gl_geo.attrs['n_fabs'] = len(level.fabs)
            gl_phy = g_phy.create_group("Level %i" % ilevel)
            gl_phy.attrs['n_fabs'] = len(level.fabs)

            for ifab, fab in enumerate(level.fabs):

                # Geometry
                gf_geo = gl_geo.create_group("Fab %i" % ifab)
                gf_geo.attrs['xmin'] = fab.xmin
                gf_geo.attrs['xmax'] = fab.xmax
                gf_geo.attrs['ymin'] = fab.ymin
                gf_geo.attrs['ymax'] = fab.ymax
                gf_geo.attrs['zmin'] = fab.zmin
                gf_geo.attrs['zmax'] = fab.zmax
                quantity = fab.data.keys()[0]
                gf_geo.attrs['n1'] = fab.data[quantity].shape[0]
                gf_geo.attrs['n2'] = fab.data[quantity].shape[1]
                gf_geo.attrs['n3'] = fab.data[quantity].shape[2]

                # Physics
                gf_phy = gl_phy.create_group("Fab %i" % ifab)
                for quantity in fab.data:
                    gf_phy.create_dataset(quantity, data=fab.data[quantity][np.newaxis, :, :, :], compression=True)

        f.close()