Last active
July 19, 2024 14:13
-
-
Save cphyc/75bb40e8be1e35ee7aa7d6310e90d93e to your computer and use it in GitHub Desktop.
Load data from Dyablo
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
"""Load Dyablo simulation from a hdf5-formatted dataset | |
Requires yt_experiments to be installed: | |
$ pip install git+https://github.com/cphyc/yt_experiments@octree-converter | |
Example: | |
>>> ds = load("./test_gravity_spheres_3D_iter0008987.h5") | |
... p = yt.SlicePlot(ds, "x", ("gas","density")) | |
... p.set_unit("density", "mp/cm**3") | |
... p.save("/tmp/") | |
""" | |
import yt | |
from yt_experiments.octree.converter import OctTree | |
from pathlib import Path | |
import h5py | |
import numpy as np | |
def readFile(filename: str | Path): | |
with h5py.File(filename) as f: | |
connectivity = f["connectivity"][:].astype(np.int64) | |
coords = f["coordinates"][:].astype(np.float64) | |
meta = dict(f["scalar_data"].attrs) | |
data = {} | |
for key in f: | |
if key in ("connectivity", "coordinates", "scalar_data"): | |
continue | |
tmp = f[key][:].astype(np.float64).reshape(-1, 1) | |
data["connect1", key] = tmp | |
return meta, connectivity, coords, data | |
def load(filename: str | Path): | |
meta, connectivity, coords, data = readFile(filename) | |
# Get domain extent | |
left_edge = coords.min(axis=0) | |
right_edge = coords.max(axis=0) | |
domain_width = right_edge - left_edge | |
# Get cell positions | |
cell_pos = (coords[connectivity[:, 6]] + coords[connectivity[:, 0]]) / 2 | |
cell_dx = coords[connectivity[:, 6]] - coords[connectivity[:, 0]] | |
# Make sure we have cubic cells | |
assert np.allclose(cell_dx, cell_dx[:, 0:1]) | |
cell_lvl = np.log2(domain_width[0] / cell_dx[:, 0]).astype(int) | |
yt.mylog.info("Creating octree") | |
oct = OctTree.from_list(cell_pos, cell_lvl, check=False) | |
yt.mylog.info("Creating refmask") | |
ref_mask, leaf_order = oct.get_refmask() | |
unit_time = meta.get("tstar", 1) | |
unit_length = meta.get("vstar", 1) * unit_time | |
unit_mass = meta.get("rhostar", 1) * unit_length**3 | |
# Somehow yt ignores unit_system when loading octrees | |
# and assumes the units are in cgs | |
unit_length *= 100 | |
unit_mass *= 1000 | |
ds = yt.load_octree( | |
ref_mask, | |
{("gas", key[1]): val[leaf_order] for key, val in data.items()}, | |
bbox=np.array([left_edge, right_edge]).T, | |
sim_time=meta["time"], | |
length_unit=unit_length, | |
time_unit=unit_time, | |
mass_unit=unit_mass, | |
unit_system="cgs", | |
num_zones=1, | |
) | |
ds.stream_handler.name = filename | |
return ds |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment