File size: 2,416 Bytes
a3441f1 | 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 | from __future__ import annotations
import argparse
from pathlib import Path
import h5py
import numpy as np
from common import DEFAULT_CONFIG, load_config, prepare_config
def generate_darcy_file(path: Path, samples: int, grid_size: int, seed: int, overwrite: bool) -> None:
if path.exists() and not overwrite:
print(f"Fake data already exists: {path}")
return
path.parent.mkdir(parents=True, exist_ok=True)
rng = np.random.default_rng(seed)
x = np.linspace(0.0, 1.0, grid_size, dtype=np.float32)
y = np.linspace(0.0, 1.0, grid_size, dtype=np.float32)
xx, yy = np.meshgrid(x, y, indexing="ij")
nu = rng.uniform(0.05, 0.25, size=(samples, grid_size, grid_size)).astype(np.float32)
tensor = np.empty((samples, 2, grid_size, grid_size), dtype=np.float32)
base = np.sin(np.pi * xx) * np.sin(np.pi * yy)
for idx in range(samples):
tensor[idx, 0] = base + 0.05 * rng.standard_normal(base.shape)
tensor[idx, 1] = tensor[idx, 0] * (1.0 + nu[idx]) + 0.01 * rng.standard_normal(base.shape)
with h5py.File(path, "w") as handle:
handle.create_dataset("tensor", data=tensor)
handle.create_dataset("nu", data=nu)
handle.create_dataset("x-coordinate", data=x)
handle.create_dataset("y-coordinate", data=y)
handle.create_dataset("t-coordinate", data=np.array([0.0, 1.0], dtype=np.float32))
print(f"Fake Darcy HDF5 written: {path}")
print(f"tensor={tensor.shape}, nu={nu.shape}")
def main() -> int:
parser = argparse.ArgumentParser(description="Generate a small PDENNEval FNO fake dataset.")
parser.add_argument("--config", default=str(DEFAULT_CONFIG), help="Path to conf/config.yaml")
parser.add_argument("--data-dir", default=None, help="Override datapipe.source.data_dir")
parser.add_argument("--samples", type=int, default=20)
parser.add_argument("--grid-size", type=int, default=16)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--overwrite", action="store_true")
args = parser.parse_args()
cfg = prepare_config(load_config(args.config), data_dir=args.data_dir)
data_dir = Path(cfg.datapipe.source.data_dir)
file_path = data_dir / cfg.datapipe.source.file_name
generate_darcy_file(file_path, args.samples, args.grid_size, args.seed, args.overwrite)
return 0
if __name__ == "__main__":
raise SystemExit(main())
|