File size: 3,317 Bytes
ae73c7f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | #!/usr/bin/env python3
"""
Convert The Well active_matter nested HDF5 into flat (T, C, H, W) files
for LocalWellHDF5 / run_full.
Default C=11 (full spatiotemporal state, float32, lossless gzip):
concentration, vx, vy, D00, D01, D10, D11, E00, E01, E10, E11
Does not reduce velocity to speed, does not drop tensors, does not
downcast precision. Truncated/corrupt source files are skipped.
Example:
python scripts/convert_well_active_matter.py \\
--src data/well/datasets/active_matter/data/train \\
--out data/well_simple \\
--max-trajs 16
"""
from __future__ import annotations
import argparse
from pathlib import Path
import h5py
import numpy as np
CHANNEL_NAMES = [
"concentration", "vx", "vy",
"D00", "D01", "D10", "D11",
"E00", "E01", "E10", "E11",
]
def convert_one(fp: Path, out_dir: Path, n_written: int, max_trajs: int) -> int:
try:
with h5py.File(fp, "r") as f:
conc = f["t0_fields/concentration"][:]
vel = f["t1_fields/velocity"][:]
D = f["t2_fields/D"][:]
E = f["t2_fields/E"][:]
meta = {}
for k in ("L", "alpha", "zeta"):
if f"scalars/{k}" in f:
meta[k] = float(f[f"scalars/{k}"][()])
except OSError as e:
print(f"SKIP truncated/corrupt: {fp.name} | {e}")
return n_written
N, T, H, W = conc.shape
for i in range(N):
if n_written >= max_trajs:
break
chans = [
conc[i],
vel[i, ..., 0], vel[i, ..., 1],
D[i, ..., 0, 0], D[i, ..., 0, 1], D[i, ..., 1, 0], D[i, ..., 1, 1],
E[i, ..., 0, 0], E[i, ..., 0, 1], E[i, ..., 1, 0], E[i, ..., 1, 1],
]
fields = np.stack(chans, axis=1).astype(np.float32)
assert fields.shape == (T, 11, H, W), fields.shape
out = out_dir / f"traj_{n_written:03d}.hdf5"
with h5py.File(out, "w") as g:
g.create_dataset("fields", data=fields, compression="gzip")
g.attrs["channel_names"] = np.array(CHANNEL_NAMES, dtype="S")
g.attrs["source_file"] = fp.name
for k, v in meta.items():
g.attrs[k] = v
print(f"wrote {out.name} {fields.shape}")
n_written += 1
return n_written
def main():
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument("--src", type=Path, required=True,
help="Directory of active_matter train *.hdf5 files")
ap.add_argument("--out", type=Path, required=True,
help="Output directory for flat traj_*.hdf5")
ap.add_argument("--max-trajs", type=int, default=16)
ap.add_argument("--clear-out", action="store_true",
help="Delete existing traj_*.hdf5 in --out first")
args = ap.parse_args()
args.out.mkdir(parents=True, exist_ok=True)
if args.clear_out:
for old in args.out.glob("traj_*.hdf5"):
old.unlink()
n = 0
for fp in sorted(args.src.glob("*.hdf5")):
n = convert_one(fp, args.out, n, args.max_trajs)
if n >= args.max_trajs:
break
print(f"total trajectories: {n}")
print(f"out_dir: {args.out}")
print("Train with: python -m src.run_full --data-root <out> --channels 11")
if __name__ == "__main__":
main()
|