| 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()) |
|
|
