#!/usr/bin/env python3 """Generate OneScience-compatible synthetic ERA5 HDF5 data.""" from __future__ import annotations import argparse from pathlib import Path import h5py import numpy as np try: from common import PROJECT_ROOT, SYNTHETIC_GENERATOR_VERSION, channel_order, load_config, resolve_path, write_json except ModuleNotFoundError: # supports ``python -m scripts.fake_data`` from scripts.common import PROJECT_ROOT, SYNTHETIC_GENERATOR_VERSION, channel_order, load_config, resolve_path, write_json def make_field(name: str, step: int, year: int, lat: np.ndarray, lon: np.ndarray, rng: np.random.Generator) -> np.ndarray: latr, lonr = np.deg2rad(lat)[:, None], np.deg2rad(lon)[None, :] wave = np.cos(latr) * np.sin(lonr + step * 0.05) noise = rng.normal(0, 1, wave.shape).astype(np.float32) if name == "sea_ice_cover": return np.clip(0.55 - 0.45 * np.cos(latr) + 0.03 * noise, 0, 1).astype(np.float32) if name == "sea_surface_temperature": return (273.15 + 26 * np.cos(latr) + 1.5 * wave + 0.1 * noise).astype(np.float32) base, _, suffix = name.rpartition("_") level = int(suffix) ratio = level / 1000.0 if base == "geopotential": # Hydrostatic log-pressure profile calibrated against the bundled ERA5 # statistics (about 464 km2/s2 at 1 hPa and 0.7 km2/s2 at 1000 hPa). vertical = 733.0 + 67000.0 * np.log(1000.0 / level) return (vertical + 1800.0 * (1.0 - ratio) * np.sin(latr) ** 2 + 120.0 * wave + noise).astype(np.float32) if base == "temperature": # Piecewise standard-atmosphere profile captures stratospheric warming; # a monotone 220->288 K profile is not physically valid above 10 hPa. anchors_hpa = np.asarray([1.0, 10.0, 100.0, 500.0, 1000.0]) anchors_k = np.asarray([261.0, 229.0, 207.0, 253.0, 281.0]) vertical = np.interp(np.log(level), np.log(anchors_hpa), anchors_k) return (vertical + 4.0 * wave + 0.1 * noise).astype(np.float32) if base == "specific_humidity": vertical = 4.0e-6 + 7.0e-3 * ratio**3 return np.maximum((vertical * (0.8 + 0.2 * np.cos(latr)) + 1e-7 * noise), 1e-8).astype(np.float32) if base == "u_component_of_wind": return (12 * (1 - ratio) * wave + 0.2 * noise).astype(np.float32) if base == "v_component_of_wind": return (8 * (1 - ratio) * np.sin(2 * latr) * np.cos(lonr) + 0.2 * noise).astype(np.float32) if name.startswith("specific_cloud_ice_water_content_"): profile = 8e-6 * np.exp(-0.5 * ((level - 450.0) / 180.0) ** 2) return np.maximum(profile * (0.7 + 0.3 * np.cos(latr)) + 1e-7 * noise, 0).astype(np.float32) if name.startswith("specific_cloud_liquid_water_content_"): profile = 1.2e-5 * np.exp(-0.5 * ((level - 750.0) / 160.0) ** 2) return np.maximum(profile * (0.7 + 0.3 * np.cos(latr)) + 1e-7 * noise, 0).astype(np.float32) raise ValueError(f"Unsupported channel: {name}") def generate_year( path: Path, channels: list[str], timesteps: int, height: int, width: int, year: int, seed: int, time_step_hours: int, ) -> dict: lat = np.linspace(90, -90, height, dtype=np.float32) lon = np.linspace(0, 360, width, endpoint=False, dtype=np.float32) path.parent.mkdir(parents=True, exist_ok=True) sums = np.zeros(len(channels), dtype=np.float64) sq = np.zeros(len(channels), dtype=np.float64) with h5py.File(path, "w") as handle: fields = handle.create_dataset("fields", shape=(timesteps, len(channels), height, width), dtype="f4", chunks=(1, 1, min(height, 32), min(width, 64)), compression="gzip", compression_opts=1) fields.attrs["variables"] = np.asarray(channels, dtype=h5py.string_dtype("utf-8")) fields.attrs["time_step"] = time_step_hours fields.attrs["synthetic"] = True fields.attrs["generator_version"] = SYNTHETIC_GENERATOR_VERSION for t in range(timesteps): rng = np.random.default_rng(seed + year * 1009 + t) for i, name in enumerate(channels): value = make_field(name, t, year, lat, lon, rng) fields[t, i] = value sums[i] += value.sum(dtype=np.float64) sq[i] += np.square(value, dtype=np.float64).sum(dtype=np.float64) count = timesteps * height * width means = sums / count stds = np.sqrt(np.maximum(sq / count - means**2, 1e-12)) handle.create_dataset("global_means", data=means[None, :, None, None].astype("f4")) handle.create_dataset("global_stds", data=stds[None, :, None, None].astype("f4")) return {"path": str(path), "shape": [timesteps, len(channels), height, width], "mean_std_min": float(stds.min())} def generate_static(path: Path, height: int, width: int) -> dict: """Generate deterministic ERA5-like static topography and land mask. These variables are auxiliary model inputs rather than flattened dynamic channels. Values use the physical units expected by the official Gin profiles: geopotential at surface in m² s⁻² and land-sea mask in [0, 1]. """ import xarray as xr lat = np.linspace(90.0, -90.0, height, dtype=np.float32) lon = np.linspace(0.0, 360.0, width, endpoint=False, dtype=np.float32) latr, lonr = np.deg2rad(lat)[:, None], np.deg2rad(lon)[None, :] # Smooth continent-like mask and non-negative terrain height. This is # intentionally synthetic; real ERA5 static fields should be preferred. mask = (0.5 + 0.5 * np.sin(2.0 * latr) * np.cos(3.0 * lonr) > 0.52).astype(np.float32) elevation_m = np.maximum( 0.0, 2200.0 * mask * (0.35 + 0.65 * np.cos(latr) ** 2) + 250.0 * np.sin(latr) ** 2, ).astype(np.float32) geopotential = (9.80665 * elevation_m).astype(np.float32) ds = xr.Dataset( { "geopotential_at_surface": (("longitude", "latitude"), geopotential.T), "land_sea_mask": (("longitude", "latitude"), mask.T), }, coords={"latitude": lat, "longitude": lon}, attrs={"synthetic": "true", "source": "neuralgcm_develop.fake_data"}, ) ds["geopotential_at_surface"].attrs["units"] = "m**2 s**-2" ds["land_sea_mask"].attrs["units"] = "dimensionless" path.parent.mkdir(parents=True, exist_ok=True) ds.to_netcdf(path) return {"path": str(path), "shape": [height, width], "variables": list(ds.data_vars)} def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--config", default=str(PROJECT_ROOT / "conf/config.yaml")) parser.add_argument("--output-dir") parser.add_argument("--years", nargs="*", type=int) parser.add_argument("--timesteps", type=int) parser.add_argument( "--forecast-steps", type=int, help="number of future 6-hour frames; defaults to data.virtual.forecast_steps", ) parser.add_argument("--height", type=int) parser.add_argument("--width", type=int) args = parser.parse_args() config = load_config(args.config) out = resolve_path(args.output_dir or config["data"]["data_dir"], args.config) virtual = config["data"]["virtual"] time_step_hours = int(config["data"].get("time_step_hours", 6)) if time_step_hours <= 0: raise ValueError("data.time_step_hours must be positive") years = args.years or sorted(set(config["data"]["train_years"] + config["data"]["val_years"] + config["data"]["test_years"])) channels = channel_order(config) # The explicit value always wins, followed by the virtual-data setting, # then the inference default. The project defaults to a compact two-day # window; 60 six-hour steps exercise the full official 15-day capability. forecast_steps = int( args.forecast_steps if args.forecast_steps is not None else virtual.get("forecast_steps", config.get("inference", {}).get("prediction_steps", 8)) ) if forecast_steps <= 0: raise ValueError("--forecast-steps must be positive") input_steps = int(config["data"].get("input_steps", 1)) configured_timesteps = int(virtual["timesteps_per_year"]) # The synthetic file contains one initial frame plus the requested future # frames. Explicit --timesteps remains available for tiny smoke tests. explicit_timesteps = args.timesteps is not None timesteps = int(args.timesteps) if explicit_timesteps else max( configured_timesteps, input_steps + forecast_steps ) if timesteps <= 0: raise ValueError("--timesteps must be positive") if not explicit_timesteps and timesteps < input_steps + forecast_steps: raise ValueError( f"timesteps={timesteps} is too short for input_steps={input_steps} " f"and forecast_steps={forecast_steps}; need at least {input_steps + forecast_steps}" ) complete_window = timesteps >= input_steps + forecast_steps forecast_horizon_days = forecast_steps * time_step_hours / 24 if explicit_timesteps and not complete_window: print( f"Warning: timesteps={timesteps} provides only a smoke window; " f"{input_steps + forecast_steps} frames are required for the full " f"{forecast_horizon_days:g}-day configured forecast." ) records = [ generate_year( out / "data" / f"{year}.h5", channels, timesteps, args.height or virtual["height"], args.width or virtual["width"], year, int(virtual["seed"]), time_step_hours, ) for year in years ] static_record = generate_static(out / "static.nc", args.height or virtual["height"], args.width or virtual["width"]) write_json(out / "metadata" / "dataset_card.json", {"name": "neuralgcm-synthetic-era5", "format": "OneScience ERA5Dataset HDF5", "channels": channels, "files": records, "static_file": static_record, "native_model_grid": config["model"]["grid_shape"], "input_steps": input_steps, "forecast_steps": forecast_steps, "time_step_hours": time_step_hours, "forecast_horizon_hours": forecast_steps * time_step_hours, "forecast_horizon_days": forecast_horizon_days, "official_forecast_capability_days": [2, 15], "official_15_day_window_complete": forecast_horizon_days >= 15, "forecast_window_complete": complete_window}) print(f"Generated {len(records)} years under {out}") if __name__ == "__main__": main()