File size: 1,754 Bytes
5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c | 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 | """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")
# Correlated targets make the regression/backward smoke test meaningful.
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()
|