clawhdf5/examples/wasm-viewer/test/make_fixture.py

"""Write HDF5 and NetCDF-4 test files with h5py/netCDF4, and what libhdf5
reads back from them, for the clawhdf5-wasm tests.

    python make_fixture.py OUT_DIR

writes OUT_DIR/fixture.h5, OUT_DIR/fixture.nc and OUT_DIR/expected.json.
Both the Rust test (crates/clawhdf5-wasm/tests/h5py_interop.rs, native) and
the Node test (test.mjs, the built wasm package) compare against the same
expected.json, so the two check the same values.

Every expected value comes from h5py reading the file back (numpy slicing
for hyperslabs), never from the arrays that were written. Integers are
encoded as strings so JSON.parse keeps 64-bit values exact.
"""

import json
import os
import sys
import warnings
from pathlib import Path

import h5py
import netCDF4
import numpy as np

try:  # registers the LZ4/Zstd filters with libhdf5; optional
    import hdf5plugin
except ImportError:
    hdf5plugin = None

# netCDF4 1.7 trips numpy 2.5's shape-setting deprecation on assignment.
warnings.filterwarnings("ignore", category=DeprecationWarning)

out = Path(sys.argv[1])
out.mkdir(parents=True, exist_ok=True)
h5 = out / "fixture.h5"
nc = out / "fixture.nc"

rng = np.random.default_rng(7)
with h5py.File(h5, "w") as f:
    f.attrs["title"] = "wasm fixture"
    f.attrs["version"] = np.int64(3)
    f.attrs["scale"] = np.array([0.5, 2.0])
    f.attrs["big"] = np.uint64(2**63 + 5)
    f.attrs.create("vlen_note", "héllo", dtype=h5py.string_dtype())
    # No plain JavaScript form: listed with value null and its type.
    f.attrs["origin"] = np.array((1.5, 2), dtype=[("x", "<f8"), ("n", "<i4")])
    f.create_dataset(
        "grid", data=np.arange(60, dtype="<f8").reshape(6, 10) / 4,
        chunks=(4, 3), compression="gzip", shuffle=True,
    )
    f.create_dataset("f32_be", data=rng.standard_normal(7).astype(">f4"))
    f.create_dataset("f16", data=np.array([0.5, -1.25, 65504], dtype="<f2"))
    f.create_dataset("i8", data=np.array([-128, -1, 0, 127], dtype="i1"))
    f.create_dataset("i16_be", data=np.array([-32768, 5, 32767], dtype=">i2"))
    f.create_dataset("u16", data=np.array([0, 40000, 65535], dtype="<u2"))
    f.create_dataset("u32", data=np.array([0, 4_000_000_000], dtype="<u4"))
    f.create_dataset("i64", data=np.array([-(2**63), 2**53 + 1, 7], dtype="<i8"))
    f.create_dataset("u64", data=np.array([2**64 - 1, 1], dtype="<u8"))
    f.create_dataset("scalar", data=np.float64(2.5))
    f.create_dataset("fixed_str", data=np.array([b"alpha", b"be", b""], dtype="S5"))
    f.create_dataset(
        "vlen_str", data=["one", "двa", ""], dtype=h5py.string_dtype()
    )
    f.create_dataset("flags", data=np.array([True, False, True]))
    # An array datatype: four elements, each an i4[2].
    pairs = f.create_dataset("pairs", shape=(4,), dtype=np.dtype(("<i4", (2,))))
    pairs[...] = np.arange(8, dtype="<i4").reshape(4, 2)
    f.create_dataset(
        "cube", data=np.arange(2 * 5 * 6, dtype="<i4").reshape(2, 5, 6),
        chunks=(1, 2, 3), compression="gzip",
    )
    if hdf5plugin is not None:
        # LZ4 is built into clawhdf5-wasm; Zstd links C and is not.
        f.create_dataset("lz4", data=np.arange(40, dtype="<i4"), chunks=(10,),
                         **hdf5plugin.LZ4())
        f.create_dataset("zstd", data=np.arange(40, dtype="<i4"), chunks=(10,),
                         **hdf5plugin.Zstd())
    comp = np.zeros(2, dtype=[("x", "<f8"), ("n", "<i4")])
    f.create_dataset("table", data=comp)
    g = f.create_group("sensors")
    g.attrs["location"] = "lab"
    g.create_dataset("temp", data=np.array([21.5, 22.0, 22.25], dtype="<f4"))
    g.create_group("empty")
    f["alias"] = h5py.SoftLink("/sensors/temp")

with netCDF4.Dataset(nc, "w") as d:
    d.title = "nc fixture"
    d.createDimension("time", None)
    d.createDimension("x", 4)
    t = d.createVariable("time", "f8", ("time",))
    t.units = "days since 2000-01-01"
    v = d.createVariable("temp", "f4", ("time", "x"), zlib=True)
    t[:] = np.arange(3)
    v[:] = np.arange(12, dtype="f4").reshape(3, 4) + 0.5


def kind(dt):
    """The typed-array kind clawhdf5-wasm returns for a numpy dtype."""
    if dt.kind == "b" or h5py.check_enum_dtype(dt) is not None:
        return "strings"
    if h5py.check_string_dtype(dt) is not None or dt.kind == "S":
        return "strings"
    if dt.subdtype is not None:
        return kind(dt.subdtype[0])
    if dt.kind == "f":
        return "f64" if dt.itemsize == 8 else "f32"
    if dt.kind in "iu":
        return f"{dt.kind}{dt.itemsize * 8}"
    raise ValueError(dt)


def flat(a, k):
    a = np.asarray(a)
    if k == "strings":
        if a.dtype.kind == "b":
            return ["TRUE" if x else "FALSE" for x in a.ravel()]
        return [x.decode() if isinstance(x, bytes) else str(x) for x in a.ravel()]
    if k.startswith(("i", "u")):
        return [str(int(x)) for x in a.ravel()]
    return [float(x) for x in a.ravel()]


def entry(ds, slab=None):
    k = kind(ds.dtype)
    data = ds[()]
    e = {"kind": k, "shape": list(np.shape(data)), "values": flat(data, k)}
    if slab:
        start, count, stride = slab
        idx = tuple(slice(s, s + (c - 1) * st + 1, st)
                    for s, c, st in zip(start, count, stride))
        part = ds[idx]
        e["slab"] = {"start": start, "count": count, "stride": stride,
                     "shape": list(part.shape), "values": flat(part, k)}
    return e


def attr(v):
    v = np.asarray(v) if not isinstance(v, (str, bytes)) else v
    if isinstance(v, np.ndarray) and v.dtype.names:
        return {"raw": "compound"}
    if isinstance(v, bytes):
        return {"string": v.decode()}
    if isinstance(v, str):
        return {"string": v}
    if v.dtype.kind in "iu":
        return {"int": [str(int(x)) for x in v.ravel()], "scalar": v.ndim == 0}
    if v.dtype.kind == "f":
        return {"float": [float(x) for x in v.ravel()], "scalar": v.ndim == 0}
    if v.dtype.kind in "OSU":
        items = [x.decode() if isinstance(x, bytes) else str(x) for x in v.ravel()]
        return {"string": items[0]} if v.ndim == 0 else {"strings": items}
    raise ValueError(v.dtype)


# Attributes the reader returns but the comparison leaves out: netCDF-4's
# internal ones (a leading underscore), and dimension-scale bookkeeping.
SKIP_ATTRS = ["DIMENSION_LIST", "REFERENCE_LIST", "CLASS", "NAME"]


def describe(path):
    expected = {"datasets": {}, "errors": {}, "attrs": {}, "lists": {},
                "skip_attrs": SKIP_ATTRS}
    with h5py.File(path, "r") as f:
        def walk(key, obj):
            expected["attrs"][key] = {
                k: attr(obj.attrs[k]) for k in obj.attrs
                if not k.startswith("_") and k not in SKIP_ATTRS}
            if isinstance(obj, h5py.Group):
                members = {n: obj.get(n) for n in obj}
                groups = [n for n, o in members.items() if isinstance(o, h5py.Group)]
                sets = [n for n, o in members.items() if isinstance(o, h5py.Dataset)]
                expected["lists"][key] = {"groups": sorted(groups),
                                          "datasets": sorted(sets)}
                for n, o in members.items():
                    walk(key.rstrip("/") + "/" + n, o)
            elif obj.dtype.names:
                expected["errors"][key] = "compound"
            else:
                expected["datasets"][key] = entry(obj, slab_for(obj))
                if key == "/zstd":
                    # Refused by the wasm build (no zstd); a native build
                    # that unifies in clawhdf5-format/zstd reads it.
                    expected["datasets"][key]["unavailable"] = "unsupported filter: 32015"

        walk("/", f)
    return expected


def slab_for(obj):
    """A strided hyperslab inside the dataset's extent, or None."""
    if obj.ndim == 0 or obj.shape[0] < 2 or 0 in obj.shape:
        return None
    start = [1] + [0] * (obj.ndim - 1)
    count = [max(1, (obj.shape[0] - 1) // 2)] + [
        max(1, (n + 1) // 2) for n in obj.shape[1:]]
    stride = [2 if obj.shape[0] > 2 else 1] + [2] * (obj.ndim - 1)
    count = [min(c, (n - s - 1) // st + 1)
             for c, s, st, n in zip(count, start, stride, obj.shape)]
    return (start, count, stride)


json.dump({"fixture.h5": describe(h5), "fixture.nc": describe(nc)},
          open(out / "expected.json", "w"), indent=1, ensure_ascii=False)


def write_big(path, megabytes):
    """A large file for the range-request tests (`openUrl`): `/big`, about
    `megabytes` MB of float64 in 1 MiB chunks, written after a small
    dataset and a group, so listing and reading `/small` touch a few blocks
    of the file and a window of `/big` one chunk. Returns what h5py reads
    back."""
    n = megabytes * 1_000_000 // 8
    chunk = 1 << 17
    with h5py.File(path, "w") as f:
        f.attrs["note"] = "large file for range reads"
        f.create_dataset("small", data=np.array([1.5, -2.0, 3.25]))
        g = f.create_group("meta")
        g.attrs["units"] = "m"
        g.create_dataset("ids", data=np.arange(10, dtype="<i4"))
        big = f.create_dataset("big", shape=(n,), dtype="<f8", chunks=(chunk,))
        for s in range(0, n, 1 << 22):
            e = min(n, s + (1 << 22))
            big[s:e] = np.arange(s, e, dtype="<f8") * 0.5
    start = n // 2 + 12_345
    with h5py.File(path, "r") as f:
        return {
            "size": path.stat().st_size,
            "list": {"groups": ["meta"], "datasets": ["big", "small"]},
            "small": [float(x) for x in f["small"][()]],
            "ids": [str(int(x)) for x in f["meta/ids"][()]],
            "big_shape": list(f["big"].shape),
            "window": {"start": start, "count": 10,
                       "values": [float(x) for x in f["big"][start:start + 10]]},
        }


big_mb = int(os.environ.get("WASM_BIG_MB", "0"))
if big_mb > 0:
    json.dump(write_big(out / "big.h5", big_mb), open(out / "big.json", "w"), indent=1)