"""This manages the reading of the HDF5 output files for pyro.

"""
import importlib

import h5py

import pyro.mesh.boundary as bnd
from pyro.mesh.patch import Cartesian2d, CellCenterData2d, SphericalPolar
from pyro.particles import particles



[docs]
def read_bcs(f):
    """read in the boundary condition record from the HDF5 file"""
    try:
        gb = f["BC"]
    except KeyError:
        return None

    BCs = {}
    for name in gb:
        BCs[name] = gb[name]

    return BCs




[docs]
def read(filename):
    """read an HDF5 file and recreate the simulation object that holds the
    data and state of the simulation.

    """
    if not filename.endswith(".h5"):
        filename += ".h5"

    with h5py.File(filename, "r") as f:

        # read the simulation information -- this only exists if the
        # file was created as a simulation object
        try:
            solver_name = f.attrs["solver"]
            problem_name = f.attrs["problem"]
            t = f.attrs["time"]
            n = f.attrs["nsteps"]
        except KeyError:
            # this was just a patch written out
            solver_name = None

        # read in the grid info and create our grid
        grid = f["grid"].attrs

        try:
            coord_type = grid["coord_type"]
        except KeyError:
            coord_type = 0

        if coord_type == 1:
            grid_class = SphericalPolar
        else:
            grid_class = Cartesian2d

        myg = grid_class(grid["nx"], grid["ny"], ng=grid["ng"],
                         xmin=grid["xmin"], xmax=grid["xmax"],
                         ymin=grid["ymin"], ymax=grid["ymax"])

        # sometimes problems define custom BCs -- at the moment, we
        # are going to assume that these always map to BC.user.  We
        # need to read these in now, since the variable creation
        # requires it.
        custom_bcs = read_bcs(f)
        if custom_bcs is not None:
            if solver_name in ["compressible_fv4", "compressible_rk", "compressible_sdc"]:
                bc_solver = "compressible"
            else:
                bc_solver = solver_name
            bcmod = importlib.import_module(f"pyro.{bc_solver}.BC")
            for name, is_solid in custom_bcs.items():
                bnd.define_bc(name, bcmod.user, is_solid=is_solid)

        # read in the variable info -- start by getting the names
        gs = f["state"]
        names = []
        for n in gs:
            names.append(n)

        # create the CellCenterData2d object
        myd = CellCenterData2d(myg)

        for n in names:
            grp = gs[n]
            bc = bnd.BC(xlb=grp.attrs["xlb"], xrb=grp.attrs["xrb"],
                        ylb=grp.attrs["ylb"], yrb=grp.attrs["yrb"])
            myd.register_var(n, bc)

        myd.create()

        # auxiliary data
        for k in f["aux"].attrs:
            myd.set_aux(k, f["aux"].attrs[k])

        # restore the variable data
        for n in names:
            grp = gs[n]
            data = grp["data"]

            v = myd.get_var(n)
            v.v()[:, :] = data[:, :]

        # restore the particle data
        try:
            gparticles = f["particles"]
            particle_data = gparticles["particle_positions"]
            init_data = gparticles["init_particle_positions"]

            my_particles = particles.Particles(myd, None, len(particle_data),
                                            "array", particle_data, init_data)
        except KeyError:
            my_particles = None

        if solver_name is not None:
            solver = importlib.import_module(f"pyro.{solver_name}")

            sim = solver.Simulation(solver_name, problem_name, None, None)
            sim.n = n
            sim.cc_data = myd
            sim.cc_data.t = t
            sim.particles = my_particles

            sim.read_extras(f)

            # check if there are derived variables -- since we do a lot of
            # inheritance, we want to start with the deepest class and
            # work backwards through the inheritance until we find a derives
            # module
            for mod in [cls.__module__ for cls in type(sim).__mro__ if cls is not object]:
                try:
                    derives = importlib.import_module(f"{mod.replace('simulation', 'derives')}")
                    sim.cc_data.add_derived(derives.derive_primitives)

                except ModuleNotFoundError:
                    continue

                else:
                    break

    if solver_name is not None:
        return sim

    return myd