py: Read in a Genesis grid directly with xarray

read_grids.py

import numpy as np
import xarray as xa
import h5py as h5
from pathlib import Path
from tqdm import tqdm

GRIDS_DIR = Path("PATH_GOES_HERE")

def grid_path(snapshot: int, sub: int):
    return GRIDS_DIR / f"snapshot_{snapshot:03d}.den.{sub}"

def read_grid(snapshot: int):
    with h5.File(grid_path(snapshot, 0), 'r') as fp:
        boxsize = fp.attrs['BoxSize']
        nfiles = fp.attrs['Num_files']
        dim = fp.attrs['Ngrid_X']
           
    grid = xa.DataArray(coords=[np.linspace(0, boxsize, dim, endpoint=False)]*3, dims=('x', 'y', 'z'))
    for ii in tqdm(range(nfiles)):
        with h5.File(grid_path(snapshot, ii), 'r') as fp:
            xstart = fp.attrs['Local_x_start']
            nx = fp.attrs['Local_nx']
            grid[xstart:xstart+nx, ...] = fp['Density'][:].reshape([nx, dim, dim])
    
    return grid