"""XDMF + H5 :class:`Storage`.
Reads and writes :class:`DataSource` objects to disk in the same byte
layout the v2 pressure pipeline produces today. Round-trip is the
contract: read a fixture, write it under a new key, the bytes match.
Layouts handled:
- *Timeseries* (the common case)::
/Triangles int32 (n_triangles, 3)
/Geometry float64 (n_vertices, 3)
/meta/time_steps float64 (n_timesteps,)
/meta/time_normalized float64 (n_timesteps,)
/{field}/t{T} float64 (n_elements,) per timestep
- *Time-aggregated* (stats)::
/Triangles, /Geometry as above
/{group}/{stat_name} float64 (n_elements,)
The file at ``<root>/<key>.h5`` together with the optional sidecar
``<root>/<key>.xdmf`` is the storage unit. ``read_data_source`` returns
a :class:`SurfaceDataSource` or :class:`PointsDataSource`; the kind is
inferred from the ``key`` filename prefix (``bodies.``/``cp_t.``/...
-> surface, ``points.`` -> points). ``write_data_source`` rewrites the
file from scratch using the existing ``cfdmod.io.xdmf`` helpers, so the
output format is exactly what the v2 pipeline produces.
"""
from __future__ import annotations
__all__ = ["XdmfH5Storage"]
import hashlib
import pathlib
from typing import Iterable
import h5py
import numpy as np
from cfdmod.adapters.xdmf_h5.field_store import H5FieldStore
from cfdmod.core.data_source import (
DataSource,
GroupsDataSource,
PointsDataSource,
SurfaceDataSource,
)
from cfdmod.core.errors import StorageKeyError
from cfdmod.core.field_meta import FieldMeta
from cfdmod.core.time_axis import TimeAxis
from cfdmod.core.topology import ElementMeta, Topology
from cfdmod.io import xdmf as _xdmf
# Root-level h5 attribute the freshness layer stamps an output's signature
# under. It lives in ``.attrs`` (not a dataset), so ``read_data_source``
# ignores it and the round-trip byte layout is unchanged.
_SIGNATURE_ATTR = "cfdmod_signature"
_RESERVED_ROOT_KEYS = frozenset({"Triangles", "Geometry", "Connectivity", "meta"})
# Geometry datasets embedded inside a stats group (so write_stats_xdmf can
# emit one Grid per group). They are topology, not stat fields, and must be
# skipped when reconstructing the field list on read.
_RESERVED_GROUP_DATASETS = frozenset({"Triangles", "Geometry", "Connectivity"})
# Synthetic h5 group holding stats that had no group prefix on the source
# (e.g. a Cp stats source with bare fields mean/rms/...). Stripped on read so
# the round-trip restores the original bare field names.
_BARE_STATS_GROUP = "stats"
def _kind_from_key(key: str) -> str:
"""Infer DataSource kind from the filename stem.
A ``points.*`` stem is read as a points (probe) source; every other
stem defaults to surface, which is the most common case (``bodies.*``,
``cp_t.*``, ``stats.*``, ...). The prefix is the only signal -- a probe
file must therefore be named ``points.*`` to be read back as points.
"""
stem = pathlib.Path(key).name
if stem.startswith("points."):
return "points"
return "surface"
def _derive_time_axis(time_steps: np.ndarray, time_normalized: np.ndarray) -> TimeAxis:
"""Reconstruct an affine :class:`TimeAxis` from the stored arrays.
The on-disk arrays are kept as-is; we only read them to pull the
three numbers (initial_time, timestep_size, n_timesteps) plus the
normalization offset. If the on-disk arrays are not strictly
uniform, we still take ``timestep_size`` from the first delta --
this matches what the v2 pipeline produces (regular sampling) and
avoids the cost of resampling on read.
"""
n = int(time_steps.shape[0])
if n == 0:
return TimeAxis(initial_time=0.0, timestep_size=0.0, n_timesteps=0)
if n == 1:
return TimeAxis(
initial_time=float(time_steps[0]),
timestep_size=1.0,
n_timesteps=1,
time_normalized_offset=float(time_steps[0] - time_normalized[0]),
)
dt = float(time_steps[1] - time_steps[0])
offset = float(time_steps[0] - time_normalized[0])
return TimeAxis(
initial_time=float(time_steps[0]),
timestep_size=dt,
n_timesteps=n,
time_normalized_offset=offset,
)
[docs]
class XdmfH5Storage:
""":class:`Storage` for the XDMF + H5 byte layout.
Args:
root: Directory under which keys resolve. ``read_data_source("bodies.foo")``
opens ``<root>/bodies.foo.h5``.
write_xdmf: When True, ``write_data_source`` also (re)generates
``<key>.xdmf`` next to the h5. Default True.
"""
__slots__ = ("_root", "_write_xdmf")
def __init__(self, root: pathlib.Path, *, write_xdmf: bool = True) -> None:
self._root = pathlib.Path(root)
self._write_xdmf = bool(write_xdmf)
# --- Path helpers ------------------------------------------------------
@property
def root(self) -> pathlib.Path:
return self._root
def h5_path(self, key: str) -> pathlib.Path:
return self._root / f"{key}.h5"
def xdmf_path(self, key: str) -> pathlib.Path:
return self._root / f"{key}.xdmf"
def keys(self) -> Iterable[str]:
if not self._root.exists():
return []
return sorted(p.stem for p in self._root.glob("*.h5"))
def __contains__(self, key: str) -> bool:
return self.h5_path(key).exists()
# --- Read --------------------------------------------------------------
def read_data_source(self, key: str) -> DataSource:
h5_path = self.h5_path(key)
if not h5_path.exists():
raise StorageKeyError(
f"XdmfH5Storage has no data source under key {key!r} ({h5_path})"
)
with h5py.File(h5_path, "r") as f:
if "Triangles" not in f or "Geometry" not in f:
raise ValueError(
f"{h5_path} is missing the standard /Triangles and /Geometry datasets; "
"this adapter only handles the cfdmod v2 layout."
)
triangles = np.asarray(f["Triangles"][:], dtype=np.int32)
vertices = np.asarray(f["Geometry"][:], dtype=np.float64)
has_meta = "meta" in f and "time_steps" in f["meta"]
time_steps = (
np.asarray(f["meta"]["time_steps"][:], dtype=np.float64) if has_meta else None
)
time_normalized = (
np.asarray(f["meta"]["time_normalized"][:], dtype=np.float64) if has_meta else None
)
field_groups: dict[str, str] = {}
time_keys: list[str] = []
time_aggregated = False
# Field groups are top-level groups other than 'meta'. Detect
# timeseries vs stats by inspecting one group's children.
for name in f.keys():
if name in _RESERVED_ROOT_KEYS:
continue
obj = f[name]
if not isinstance(obj, h5py.Group):
continue
children = [
k
for k in obj.keys()
if isinstance(obj[k], h5py.Dataset) and k not in _RESERVED_GROUP_DATASETS
]
if not children:
continue
# Stats layout: per-stat datasets directly under the group, no t-prefix.
# Timeseries layout: every dataset is t{float}.
if all(k.startswith("t") and _is_floatish(k[1:]) for k in children):
field_groups[name] = name
if not time_keys:
time_keys = sorted(children, key=lambda k: float(k[1:]))
else:
time_aggregated = True
for stat_name in children:
# Bare-stat sources are written under the synthetic
# "stats" group; strip it so the field name round-trips.
field_name = (
stat_name if name == _BARE_STATS_GROUP else f"{name}/{stat_name}"
)
field_groups[field_name] = f"{name}/{stat_name}"
# Topology + ElementMeta
kind = _kind_from_key(key)
if kind == "points":
topology = Topology.points(vertices)
else:
topology = Topology.triangles(triangles, vertices)
elements = ElementMeta()
# Time axis
if time_aggregated or not time_keys:
if has_meta and time_steps.shape[0] > 0 and not time_aggregated:
time = _derive_time_axis(time_steps, time_normalized)
else:
time = TimeAxis(initial_time=0.0, timestep_size=0.0, n_timesteps=0)
else:
if has_meta:
time = _derive_time_axis(time_steps, time_normalized)
else:
# Reconstruct from the keys themselves.
ts = np.array([float(k[1:]) for k in time_keys], dtype=np.float64)
time = _derive_time_axis(ts, ts)
store = H5FieldStore(
h5_path=h5_path,
field_groups=field_groups,
time_keys=[] if time_aggregated else time_keys,
n_elements=topology.n_elements,
time_aggregated=time_aggregated,
)
field_meta = {name: FieldMeta(name=name) for name in field_groups}
common = dict(
time=time,
topology=topology,
elements=elements,
fields=store,
field_meta=field_meta,
attrs={"source_path": str(h5_path)},
)
if kind == "points":
return PointsDataSource(**common)
return SurfaceDataSource(**common)
# --- Write -------------------------------------------------------------
def write_data_source(self, key: str, ds: DataSource) -> None:
# GroupsDataSource is special: it has no topology of its own,
# but it does carry parent_topology + parent_grouping. We
# broadcast per-group values back to the parent's triangles so
# the on-disk h5 is a regular surface timeseries that ParaView
# can render. The result matches the legacy run_cf output.
if isinstance(ds, GroupsDataSource):
ds = _groups_to_parent_surface(ds)
if ds.topology is None:
raise ValueError(
f"XdmfH5Storage cannot write a DataSource with no topology (kind={ds.kind!r})."
)
h5_path = self.h5_path(key)
h5_path.parent.mkdir(parents=True, exist_ok=True)
if h5_path.exists():
h5_path.unlink()
triangles = _connectivity_for_write(ds.topology)
vertices = np.asarray(ds.topology.vertices, dtype=np.float64)
_xdmf.write_timeseries_geometry(h5_path, triangles, vertices)
time_aggregated = ds.time.is_time_aggregated
if not time_aggregated:
time_steps = ds.time.times()
time_normalized = ds.time.times_normalized()
_xdmf.write_timeseries_meta(h5_path, time_steps, time_normalized)
# Resolve every field by reading via the source's own FieldStore.
# That makes the writeback work uniformly for MemoryFieldStore,
# H5FieldStore (with overlay), and any future store.
groups_for_xdmf: list[str] = []
for fname in sorted(ds.fields.keys()):
arr = ds.fields.read(fname)
if time_aggregated:
# Single dataset path: support 'group/stat' or bare 'stat'.
group, _, stat = fname.partition("/")
if not stat:
stat = group
group = _BARE_STATS_GROUP
_xdmf.write_stats_field(
h5_path,
group=group,
stat_name=stat,
values=np.asarray(arr, dtype=np.float64),
triangles=triangles,
vertices=vertices,
)
else:
# Timeseries: arr is (n_elements, n_timesteps).
if arr.ndim != 2:
raise ValueError(
f"field {fname!r} must be 2-D for a non-aggregated DataSource; "
f"got shape {arr.shape}"
)
ts = ds.time.times()
for i, t in enumerate(ts):
_xdmf.write_timeseries_step(
h5_path,
group=fname,
key=f"t{t}",
data=np.asarray(arr[:, i], dtype=np.float64),
)
groups_for_xdmf.append(fname)
if self._write_xdmf:
xdmf_path = self.xdmf_path(key)
if time_aggregated:
_xdmf.write_stats_xdmf(h5_path, xdmf_path)
elif groups_for_xdmf:
_xdmf.write_temporal_xdmf(h5_path, xdmf_path, groups_for_xdmf)
# --- Freshness ---------------------------------------------------------
[docs]
def digest(self, key: str, strategy: str = "size_mtime") -> str:
"""Change-detecting token for the ``<key>.h5`` (+ ``.xdmf``) pair.
- ``size_mtime`` (default): size + mtime of the file(s); no reads.
- ``content``: a blake2b hash of the file bytes, streamed.
- ``backend``: the local filesystem has no native token, so this
falls back to ``size_mtime`` (tagged so the fallback is visible).
"""
h5_path = self.h5_path(key)
if not h5_path.exists():
raise StorageKeyError(
f"XdmfH5Storage has no data source under key {key!r} ({h5_path})"
)
paths = [h5_path]
xdmf = self.xdmf_path(key)
if xdmf.exists():
paths.append(xdmf)
if strategy == "content":
h = hashlib.blake2b(digest_size=32)
for p in paths:
with open(p, "rb") as f:
for block in iter(lambda: f.read(1 << 20), b""):
h.update(block)
return f"content:{h.hexdigest()}"
# size_mtime and backend (backend has no FS-native token -> fall back)
prefix = "size_mtime" if strategy != "backend" else "backend_fs"
parts = [f"{p.stat().st_size}:{p.stat().st_mtime_ns}" for p in paths]
return f"{prefix}:" + "|".join(parts)
def read_signature(self, key: str) -> str | None:
h5_path = self.h5_path(key)
if not h5_path.exists():
return None
with h5py.File(h5_path, "r") as f:
raw = f.attrs.get(_SIGNATURE_ATTR)
if raw is None:
return None
return raw.decode("utf-8") if isinstance(raw, bytes) else str(raw)
def write_signature(self, key: str, signature: str) -> None:
h5_path = self.h5_path(key)
if not h5_path.exists():
raise StorageKeyError(f"cannot stamp signature: no h5 under key {key!r} ({h5_path})")
with h5py.File(h5_path, "a") as f:
f.attrs[_SIGNATURE_ATTR] = signature
def _groups_to_parent_surface(ds: GroupsDataSource) -> SurfaceDataSource:
"""Broadcast a GroupsDataSource back onto its parent surface.
Returns a :class:`SurfaceDataSource` over the parent's triangles
with each parent triangle taking the value of the group it belongs
to. Triangles in ungrouped territory (``-1``) get NaN.
"""
from cfdmod.adapters.memory import MemoryFieldStore
from cfdmod.core.field_meta import FieldMeta
parent_indices = ds.parent_grouping.indices
group_ids = ds.groupings[ds.parent_grouping.name].indices # row index -> group id
# Map group id -> row index in the groups source.
row_for_gid = {int(gid): row for row, gid in enumerate(group_ids)}
n_parent = ds.parent_topology.n_elements
out_arrays: dict[str, np.ndarray] = {}
out_meta: dict[str, FieldMeta] = {}
for fname in ds.fields.keys():
arr = np.asarray(ds.fields.read(fname), dtype=np.float64)
if arr.ndim == 2:
broadcast = np.full((n_parent, arr.shape[1]), np.nan, dtype=np.float64)
else:
broadcast = np.full(n_parent, np.nan, dtype=np.float64)
for tri in range(n_parent):
gid = int(parent_indices[tri])
if gid not in row_for_gid:
continue
broadcast[tri] = arr[row_for_gid[gid]]
out_arrays[fname] = broadcast
out_meta[fname] = ds.field_meta.get(fname, FieldMeta(name=fname))
return SurfaceDataSource(
time=ds.time,
topology=ds.parent_topology,
elements=ElementMeta(),
fields=MemoryFieldStore(out_arrays),
field_meta=out_meta,
)
def _connectivity_for_write(topology: Topology) -> np.ndarray:
"""Connectivity array as written under ``/Triangles``.
Triangle topologies write their connectivity directly. Point
topologies write a degenerate ``(n_points, 3)`` block of point
indices so the on-disk layout still has a ``/Triangles`` dataset --
this matches the v2 ``points.*.h5`` files in the fixtures.
"""
if topology.cell_type == "triangle":
return np.asarray(topology.connectivity, dtype=np.int32)
if topology.cell_type == "point":
n = topology.n_elements
col = np.arange(n, dtype=np.int32)
return np.stack([col, col, col], axis=1)
raise ValueError(f"XdmfH5Storage cannot write topology with cell_type={topology.cell_type!r}")
def _is_floatish(s: str) -> bool:
try:
float(s)
except ValueError:
return False
return True