scripts/fake_data.py

"""Generate deterministic synthetic spherical data for pipeline smoke tests."""

import argparse
import json
from pathlib import Path

import numpy as np
import xarray as xr


VARIABLES = [
    "PRESsfc",
    "surface_temperature",
    *[f"air_temperature_{i}" for i in range(8)],
    *[f"specific_total_water_{i}" for i in range(8)],
    *[f"eastward_wind_{i}" for i in range(8)],
    *[f"northward_wind_{i}" for i in range(8)],
    "DSWRFtoa",
    "HGTsfc",
    "ocean_fraction",
]


def generate_data(output_dir: str, time_steps: int, latitude: int, longitude: int) -> Path:
    """Write a small NetCDF fixture and metadata to the configured data directory."""
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    data = np.zeros((time_steps, len(VARIABLES), latitude, longitude), dtype=np.float32)
    lat = np.linspace(-90, 90, latitude, dtype=np.float32)
    lon = np.linspace(0, 360, longitude, endpoint=False, dtype=np.float32)
    data[:, 0] = 1.0
    data[:, 1] = 280.0 + np.sin(np.deg2rad(lat))[None, :, None]
    data[:, -1] = 1.0
    dataset = xr.Dataset(
        {name: (("time", "lat", "lon"), data[:, index]) for index, name in enumerate(VARIABLES)},
        coords={"time": np.arange(time_steps), "lat": lat, "lon": lon},
        attrs={"dataset_type": "synthetic_smoke_data", "paper_reproduction": "false"},
    )
    netcdf_path = output_path / "synthetic_fv3gfs.nc"
    dataset.to_netcdf(netcdf_path)
    metadata = {
        "dataset_type": "synthetic_smoke_data",
        "paper_reproduction": False,
        "variables": VARIABLES,
        "shape": [time_steps, len(VARIABLES), latitude, longitude],
        "source": "scripts/generate_data.py",
    }
    (output_path / "synthetic_fv3gfs.json").write_text(
        json.dumps(metadata, indent=2) + "\n", encoding="utf-8"
    )
    return netcdf_path


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", default="conf/config.yaml")
    parser.add_argument("--output-dir")
    parser.add_argument("--time-steps", type=int)
    parser.add_argument("--latitude", type=int)
    parser.add_argument("--longitude", type=int)
    args = parser.parse_args()
    import yaml

    config = yaml.safe_load(open(args.config, encoding="utf-8"))["synthetic_data"]
    path = generate_data(
        args.output_dir or config["output_dir"],
        args.time_steps or config["time_steps"],
        args.latitude or config["latitude"],
        args.longitude or config["longitude"],
    )
    print(f"Generated synthetic smoke data: {path}")
    print("This fixture is for pipeline validation only, not paper reproduction.")


if __name__ == "__main__":
    main()