from __future__ import annotations
import shutil
from pathlib import Path
import h5py
import netCDF4 as nc
import numpy as np
import zarr
from generate.constants import (
CHUNK_ELEMS,
EXTRA_LARGE_H5,
EXTRA_LARGE_LOGICAL_BYTES,
EXTRA_LARGE_NC,
EXTRA_LARGE_ZARR,
FIXTURES_ROOT,
FLOAT32_BYTES,
LARGE_H5,
LARGE_NC,
LARGE_PER_FORMAT_BYTES,
LARGE_SUITE_LOGICAL_BYTES,
LARGE_ZARR,
SEED_LARGE,
)
from generate.util import ZARR_RAW_ARRAY_KWARGS, iter_slabs, progress, slab_count
def write_h5_f32_slabs(
path: Path,
*,
logical_bytes: int,
fixture: str,
suite_total_bytes: int | None,
quiet: bool,
seed: int,
) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
rng = np.random.default_rng(seed)
total = logical_bytes // FLOAT32_BYTES
shape = (total,)
n_slabs = slab_count(total)
gib = logical_bytes / (1024**3)
if not quiet:
print(f"HDF5: {path} ({gib:.2f} GiB logical f32, {n_slabs} slabs)")
with h5py.File(path, "w") as f:
f.attrs["tetration_fixture"] = fixture
f.attrs["tetration_logical_bytes"] = int(logical_bytes)
if suite_total_bytes is not None:
f.attrs["tetration_suite_total_bytes"] = int(suite_total_bytes)
dset = f.create_dataset(
"data",
shape=shape,
dtype=np.float32,
chunks=(CHUNK_ELEMS,),
)
for start, end in progress(
iter_slabs(total),
total=n_slabs,
desc=path.name,
quiet=quiet,
unit="slab",
):
dset[start:end] = rng.random(end - start, dtype=np.float32)
if not quiet:
print(f" finished {path}")
def write_nc_f32_slabs(
path: Path,
*,
logical_bytes: int,
fixture: str,
suite_total_bytes: int | None,
quiet: bool,
seed: int,
) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
rng = np.random.default_rng(seed)
total = logical_bytes // FLOAT32_BYTES
n_slabs = slab_count(total)
gib = logical_bytes / (1024**3)
if not quiet:
print(f"NetCDF: {path} ({gib:.2f} GiB logical f32, {n_slabs} slabs)")
with nc.Dataset(path, "w", format="NETCDF4") as ds:
ds.setncattr("tetration_fixture", fixture)
ds.setncattr("tetration_logical_bytes", int(logical_bytes))
if suite_total_bytes is not None:
ds.setncattr("tetration_suite_total_bytes", int(suite_total_bytes))
ds.createDimension("i", total)
var = ds.createVariable("data", "f4", ("i",), chunksizes=(CHUNK_ELEMS,))
for start, end in progress(
iter_slabs(total),
total=n_slabs,
desc=path.name,
quiet=quiet,
unit="slab",
):
var[start:end] = rng.random(end - start, dtype=np.float32)
if not quiet:
print(f" finished {path}")
def write_zarr_f32_slabs(
path: Path,
*,
logical_bytes: int,
fixture: str,
suite_total_bytes: int | None,
quiet: bool,
seed: int,
) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
if path.exists():
shutil.rmtree(path)
rng = np.random.default_rng(seed)
total = logical_bytes // FLOAT32_BYTES
shape = (total,)
n_slabs = slab_count(total)
gib = logical_bytes / (1024**3)
if not quiet:
print(f"Zarr: {path} ({gib:.2f} GiB logical f32, {n_slabs} slabs)")
root = zarr.open_group(str(path), mode="w")
attrs: dict[str, object] = {
"tetration_fixture": fixture,
"tetration_logical_bytes": int(logical_bytes),
}
if suite_total_bytes is not None:
attrs["tetration_suite_total_bytes"] = int(suite_total_bytes)
root.attrs.update(attrs)
arr = root.create_array(
"data",
shape=shape,
chunks=(CHUNK_ELEMS,),
dtype="f4",
**ZARR_RAW_ARRAY_KWARGS,
)
for start, end in progress(
iter_slabs(total),
total=n_slabs,
desc=path.name,
quiet=quiet,
unit="slab",
):
arr[start:end] = rng.random(end - start, dtype=np.float32)
if not quiet:
print(f" finished {path}")
def _write_h5_large(path: Path, *, quiet: bool) -> None:
write_h5_f32_slabs(
path,
logical_bytes=LARGE_PER_FORMAT_BYTES,
fixture="large_suite_h5",
suite_total_bytes=LARGE_SUITE_LOGICAL_BYTES,
quiet=quiet,
seed=SEED_LARGE,
)
def _write_nc_large(path: Path, *, quiet: bool) -> None:
write_nc_f32_slabs(
path,
logical_bytes=LARGE_PER_FORMAT_BYTES,
fixture="large_suite_nc",
suite_total_bytes=LARGE_SUITE_LOGICAL_BYTES,
quiet=quiet,
seed=SEED_LARGE + 1,
)
def _write_zarr_large(path: Path, *, quiet: bool) -> None:
write_zarr_f32_slabs(
path,
logical_bytes=LARGE_PER_FORMAT_BYTES,
fixture="large_suite_zarr",
suite_total_bytes=LARGE_SUITE_LOGICAL_BYTES,
quiet=quiet,
seed=SEED_LARGE + 2,
)
def generate_extra_large_h5(*, quiet: bool) -> None:
if not quiet:
gib = EXTRA_LARGE_LOGICAL_BYTES / (1024**3)
print(f"Generating extra-large HDF5 under {EXTRA_LARGE_H5} ({gib:.0f} GiB)")
write_h5_f32_slabs(
EXTRA_LARGE_H5 / "tensor_20gb.h5",
logical_bytes=EXTRA_LARGE_LOGICAL_BYTES,
fixture="extra_large_h5",
suite_total_bytes=None,
quiet=quiet,
seed=SEED_LARGE + 10,
)
if not quiet:
print(f"Done: {EXTRA_LARGE_H5}")
def generate_extra_large_netcdf(*, quiet: bool) -> None:
if not quiet:
gib = EXTRA_LARGE_LOGICAL_BYTES / (1024**3)
print(f"Generating extra-large NetCDF under {EXTRA_LARGE_NC} ({gib:.0f} GiB)")
write_nc_f32_slabs(
EXTRA_LARGE_NC / "tensor_20gb.nc",
logical_bytes=EXTRA_LARGE_LOGICAL_BYTES,
fixture="extra_large_nc",
suite_total_bytes=None,
quiet=quiet,
seed=SEED_LARGE + 11,
)
if not quiet:
print(f"Done: {EXTRA_LARGE_NC}")
def generate_extra_large_zarr(*, quiet: bool) -> None:
if not quiet:
gib = EXTRA_LARGE_LOGICAL_BYTES / (1024**3)
print(f"Generating extra-large Zarr under {EXTRA_LARGE_ZARR} ({gib:.0f} GiB)")
write_zarr_f32_slabs(
EXTRA_LARGE_ZARR / "tensor_20gb",
logical_bytes=EXTRA_LARGE_LOGICAL_BYTES,
fixture="extra_large_zarr",
suite_total_bytes=None,
quiet=quiet,
seed=SEED_LARGE + 12,
)
if not quiet:
print(f"Done: {EXTRA_LARGE_ZARR}")
def generate_large(*, quiet: bool) -> None:
suite_gib = LARGE_SUITE_LOGICAL_BYTES / (1024**3)
per_gib = LARGE_PER_FORMAT_BYTES / (1024**3)
if not quiet:
print(
f"Generating large suite under {FIXTURES_ROOT / 'large'} "
f"({suite_gib:.0f} GiB total ≈ {per_gib:.2f} GiB per format)"
)
_write_h5_large(LARGE_H5 / "tensor_large.h5", quiet=quiet)
_write_nc_large(LARGE_NC / "tensor_large.nc", quiet=quiet)
_write_zarr_large(LARGE_ZARR / "tensor_large", quiet=quiet)
if not quiet:
print(f"Done: {FIXTURES_ROOT / 'large'}")
def generate_large_h5(*, quiet: bool) -> None:
if not quiet:
gib = LARGE_PER_FORMAT_BYTES / (1024**3)
print(f"Generating large HDF5 under {LARGE_H5} ({gib:.2f} GiB logical f32)")
_write_h5_large(LARGE_H5 / "tensor_large.h5", quiet=quiet)
if not quiet:
print(f"Done: {LARGE_H5}")
def generate_large_netcdf(*, quiet: bool) -> None:
if not quiet:
gib = LARGE_PER_FORMAT_BYTES / (1024**3)
print(f"Generating large NetCDF under {LARGE_NC} ({gib:.2f} GiB logical f32)")
_write_nc_large(LARGE_NC / "tensor_large.nc", quiet=quiet)
if not quiet:
print(f"Done: {LARGE_NC}")
def generate_large_zarr(*, quiet: bool) -> None:
if not quiet:
gib = LARGE_PER_FORMAT_BYTES / (1024**3)
print(f"Generating large Zarr under {LARGE_ZARR} ({gib:.2f} GiB logical f32)")
_write_zarr_large(LARGE_ZARR / "tensor_large", quiet=quiet)
if not quiet:
print(f"Done: {LARGE_ZARR}")