File size: 2,721 Bytes
6f3c6ef | 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 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | """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()
|