script/fake_data.py

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

import numpy as np

if __package__ in (None, ""):
    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from script.utils import DEFAULT_CONFIG, load_config


SPLIT_OFFSETS = {"train": 0, "val": 100_000, "test": 200_000}

# Matches the official Earthformer SEVIR config
# (scripts/cuboid_transformer/sevir/earthformer_sevir_v1.yaml):
#   dataset.img_height/img_width = 384, in_len = 13, out_len = 12,
#   seq_len = 25, interval_real_time = 5, sample_mode = "sequent",
#   stride = 12, metrics_list = ['csi', 'pod', 'sucr', 'bias'],
#   threshold_list = [16, 74, 133, 160, 181, 219].
SEVIR_VIL_REFERENCE = {
    "dataset": "SEVIR VIL",
    "spatial_shape": [384, 384, 1],
    "sequence": {"input_frames": 13, "output_frames": 12},
    "seq_len": 25,
    "sample_mode": "sequent",
    "stride": 12,
    "frame_interval_minutes": 5,
    "thresholds": [16, 74, 133, 160, 181, 219],
    "layout": "train/val/test NPZ splits; each holds inputs [N,13,H,W,1] and targets [N,12,H,W,1]",
}


def generate_sequence(height: int, width: int, input_length: int, output_length: int, seed: int) -> np.ndarray:
    """Generate a continuous, temporally coherent 25-frame SEVIR-like VIL window.

    The window is sampled exactly like official SEVIR "sequent" mode: a single
    continuous sequence of seq_len = input_length + output_length frames is
    produced and later split into the first 13 input and the last 12 target
    frames, matching the official 13 -> 12 (65 -> 60 minute) task at 5-minute
    intervals.
    """
    rng = np.random.default_rng(seed)
    total = input_length + output_length
    yy, xx = np.mgrid[:height, :width]
    background = rng.normal(0.0, 0.008, (height, width)).astype(np.float32)
    background = (background + np.roll(background, 1, 0) + np.roll(background, 1, 1)) / 3.0
    cell_count = int(rng.integers(2, 5))
    cells = []
    for _ in range(cell_count):
        cells.append(
            (
                rng.uniform(0.15 * width, 0.85 * width),
                rng.uniform(0.15 * height, 0.85 * height),
                rng.uniform(-0.45, 0.45),
                rng.uniform(-0.45, 0.45),
                rng.uniform(max(1.2, width / 18), max(2.0, width / 8)),
                rng.uniform(max(1.2, height / 18), max(2.0, height / 8)),
                rng.uniform(0.45, 0.95),
                rng.uniform(-0.035, 0.035),
                rng.uniform(0, np.pi),
            )
        )
    frames = np.empty((total, height, width, 1), dtype=np.float32)
    for time in range(total):
        frame = np.maximum(background * (0.8 + 0.2 * np.sin(time / 5)), 0.0)
        for cx, cy, vx, vy, sx, sy, amplitude, growth, angle in cells:
            dx, dy = xx - (cx + vx * time), yy - (cy + vy * time)
            ca, sa = np.cos(angle), np.sin(angle)
            xr, yr = ca * dx + sa * dy, -sa * dx + ca * dy
            scale = np.clip(1.0 + growth * time, 0.55, 1.6)
            intensity = amplitude * np.exp(-0.5 * ((xr / (sx * scale)) ** 2 + (yr / (sy * scale)) ** 2))
            lifecycle = np.clip(1.0 + growth * time, 0.35, 1.25)
            frame += intensity.astype(np.float32) * lifecycle
        noise = rng.normal(0.0, 0.004, (height, width)).astype(np.float32)
        frames[time, ..., 0] = np.clip(frame + noise, 0.0, 1.0)
    return frames


def generate_split(samples: int, height: int, width: int, input_length: int, output_length: int, seed: int) -> tuple[np.ndarray, np.ndarray]:
    sequences = np.stack(
        [generate_sequence(height, width, input_length, output_length, seed + index) for index in range(samples)]
    )
    return sequences[:, :input_length], sequences[:, input_length:]


def parse_args() -> argparse.Namespace:
    config_parser = argparse.ArgumentParser(add_help=False)
    config_parser.add_argument("--config", default=str(DEFAULT_CONFIG))
    config_args, _ = config_parser.parse_known_args()
    defaults = load_config(config_args.config)
    data = defaults["data"]
    parser = argparse.ArgumentParser(description="Generate deterministic synthetic SEVIR-like VIL sequences")
    parser.add_argument("--config", default=config_args.config)
    parser.add_argument("--output-dir", default=data["data_dir"])
    parser.add_argument("--height", type=int, default=int(data["height"]))
    parser.add_argument("--width", type=int, default=int(data["width"]))
    parser.add_argument("--train-samples", type=int, default=int(data["train_samples"]))
    parser.add_argument("--val-samples", type=int, default=int(data["val_samples"]))
    parser.add_argument("--test-samples", type=int, default=int(data["test_samples"]))
    parser.add_argument("--seed", type=int, default=int(defaults["train"]["seed"]))
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    config = load_config(args.config)
    data = config["data"]
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    split_sizes = {"train": args.train_samples, "val": args.val_samples, "test": args.test_samples}
    for split, samples in split_sizes.items():
        if samples <= 0:
            raise ValueError(f"{split} samples must be positive")
        inputs, targets = generate_split(
            samples,
            args.height,
            args.width,
            int(data["input_length"]),
            int(data["output_length"]),
            args.seed + SPLIT_OFFSETS[split],
        )
        np.savez_compressed(output_dir / f"{split}.npz", inputs=inputs, targets=targets)
    metadata = {
        "synthetic": True,
        "official_sevir": False,
        "description": "Deterministic synthetic SEVIR-like VIL; not official SEVIR data",
        "protocol": "synthetic_sevir",
        "reference": SEVIR_VIL_REFERENCE,
        "seq_len": int(data["input_length"]) + int(data["output_length"]),
        "sample_mode": "sequent",
        "stride": SEVIR_VIL_REFERENCE["stride"],
        "frame_interval_minutes": int(data["frame_interval_minutes"]),
        "input_frames": int(data["input_length"]),
        "output_frames": int(data["output_length"]),
        "shape": [args.height, args.width, 1],
        "normalization": "unit [0,1] float32",
        "seed": args.seed,
        "splits": split_sizes,
    }
    (output_dir / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8")
    print(json.dumps({"output_dir": str(output_dir), "metadata": metadata}, indent=2))


if __name__ == "__main__":
    main()