| """Generate a tiny, structurally realistic CorrDiff NPZ dataset.""" |
|
|
| import argparse |
| from pathlib import Path |
|
|
| import numpy as np |
| import yaml |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--config", default=str(ROOT / "conf/config.yaml")) |
| parser.add_argument("--output") |
| args = parser.parse_args() |
| config = yaml.safe_load(Path(args.config).read_text(encoding="utf-8")) |
| data = config["data"] |
| count = data["fake_samples"] |
| rng = np.random.default_rng(config["seed"]) |
| coarse = rng.normal(size=(count, *data["input_shape"])).astype("float32") |
| |
| y = np.linspace(-1, 1, data["target_shape"][1], dtype="float32") |
| x = np.linspace(-1, 1, data["target_shape"][2], dtype="float32") |
| yy, xx = np.meshgrid(y, x, indexing="ij") |
| target = np.empty((count, *data["target_shape"]), dtype="float32") |
| coarse_signal = coarse.mean(axis=(2, 3)) |
| for sample in range(count): |
| for channel in range(data["target_shape"][0]): |
| target[sample, channel] = coarse_signal[sample, channel] + 0.3 * np.sin( |
| (channel + 1) * np.pi * xx |
| ) + 0.2 * np.cos((channel + 1) * np.pi * yy) |
| target += rng.normal(0, 0.05, target.shape).astype("float32") |
| output = Path(args.output) if args.output else ROOT / data["path"] |
| output.parent.mkdir(parents=True, exist_ok=True) |
| np.savez_compressed(output, input=coarse, target=target, |
| protocol=np.asarray(data["protocol"]), data_source=np.asarray("synthetic")) |
| print(f"saved={output} input={coarse.shape} target={target.shape}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
