Spaces:
Running on Zero
Running on Zero
| import os | |
| os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") | |
| os.environ.setdefault("WANDB_MODE", "disabled") | |
| os.environ.setdefault("MPLBACKEND", "Agg") | |
| import spaces # noqa: E402 (must precede torch / fme imports) | |
| import asyncio # noqa: E402 | |
| import logging # noqa: E402 | |
| import shutil # noqa: E402 | |
| import tempfile # noqa: E402 | |
| import time # noqa: E402 | |
| import cartopy.crs as ccrs # noqa: E402 | |
| import cartopy.feature as cfeature # noqa: E402 | |
| import gradio as gr # noqa: E402 | |
| import matplotlib # noqa: E402 | |
| matplotlib.use("Agg") | |
| import matplotlib.pyplot as plt # noqa: E402 | |
| import numpy as np # noqa: E402 | |
| import xarray as xr # noqa: E402 | |
| from huggingface_hub import hf_hub_download # noqa: E402 | |
| import pipeline as P # noqa: E402 | |
| logging.basicConfig(level=logging.INFO) | |
| REPO = "allenai/HiRO-ACEv2" | |
| DATA_DIR = os.path.abspath("hiro_data") | |
| FILES = [ | |
| "ACE2S-X-SHiELD.ckpt", | |
| "HiROv2.ckpt", | |
| "initial_conditions/ic_2023.nc", | |
| "forcing_data/forcing_2023.nc", | |
| ] | |
| for f in FILES: | |
| hf_hub_download(REPO, f, local_dir=DATA_DIR) | |
| IC_PATH = os.path.join(DATA_DIR, "initial_conditions/ic_2023.nc") | |
| FORCING_DIR = os.path.join(DATA_DIR, "forcing_data") | |
| MODELS = P.Models( | |
| os.path.join(DATA_DIR, "ACE2S-X-SHiELD.ckpt"), os.path.join(DATA_DIR, "HiROv2.ckpt") | |
| ) | |
| # Pre-fetch Natural Earth coastlines so plotting needs no network at request time. | |
| try: | |
| for feat in (cfeature.COASTLINE, cfeature.BORDERS, cfeature.STATES): | |
| list(feat.with_scale("50m").geometries())[:1] | |
| except Exception as e: # pragma: no cover | |
| print("Could not prefetch Natural Earth data:", e) | |
| MONTHS = ["January", "February", "March", "April", "May", "June", "July", | |
| "August", "September", "October", "November", "December"] | |
| # name, display label, unit conversion, colormap, units | |
| PLOT_VARS = [ | |
| ("PRATEsfc", "Precipitation rate", lambda x: x * 86400.0, "Blues", "mm/day"), | |
| ("WSPD10m", "10 m wind speed", lambda x: x, "viridis", "m/s"), | |
| ("PRMSL", "Sea-level pressure", lambda x: x, "RdBu_r", "hPa"), | |
| ] | |
| def _region(lat_center, lon_center, size): | |
| half = size / 2.0 | |
| lat0 = float(np.clip(lat_center - half, -66.0, 70.0 - size)) | |
| lat1 = lat0 + size | |
| lon_center = ((float(lon_center) + 180.0) % 360.0) - 180.0 # -> [-180, 180) | |
| lon0, lon1 = lon_center - half, lon_center + half | |
| if lon1 <= 0: # fully west of Greenwich -> use 0..360 convention | |
| lon0, lon1 = lon0 + 360.0, lon1 + 360.0 | |
| elif lon0 < 0 < lon1: | |
| pass # crosses the prime meridian; fme handles negative starts | |
| return lat0, lat1, lon0, lon1 | |
| def _subset_lon(da, lon0, lon1): | |
| lon = da["lon"].values | |
| lon_wrapped = np.where(lon > 180, lon - 360, lon) if lon0 < 0 else lon | |
| da = da.assign_coords(lon=lon_wrapped).sortby("lon") | |
| return da.sel(lon=slice(lon0, lon1)) | |
| def _fields(ds): | |
| out = {} | |
| out["PRATEsfc"] = ds["PRATEsfc"] | |
| out["WSPD10m"] = np.hypot(ds["UGRD10m"], ds["VGRD10m"]) | |
| out["PRMSL"] = ds["PRMSL"] | |
| out["UGRD10m"], out["VGRD10m"] = ds["UGRD10m"], ds["VGRD10m"] | |
| return out | |
| def make_figure(coarse, fine, title): | |
| proj = ccrs.PlateCarree() | |
| fig, axes = plt.subplots(len(PLOT_VARS), 2, figsize=(12, 4.6 * len(PLOT_VARS)), | |
| subplot_kw={"projection": proj}, constrained_layout=True) | |
| cf, ff = _fields(coarse), _fields(fine) | |
| for row, (name, label, conv, cmap, units) in enumerate(PLOT_VARS): | |
| c = conv(cf[name].values) | |
| f = conv(ff[name].values) | |
| if name == "PRATEsfc": | |
| vmin, vmax = 0.0, max(float(np.nanpercentile(f, 99.5)), 1.0) | |
| else: | |
| both = np.concatenate([c.ravel(), f.ravel()]) | |
| vmin, vmax = float(np.nanpercentile(both, 0.5)), float(np.nanpercentile(both, 99.5)) | |
| for col, (da, data, which) in enumerate( | |
| [(cf[name], c, "ACE2S coarse (~100 km)"), (ff[name], f, "HiROv2 downscaled (3 km)")] | |
| ): | |
| ax = axes[row, col] | |
| m = ax.pcolormesh(da["lon"].values, da["lat"].values, data, cmap=cmap, | |
| vmin=vmin, vmax=vmax, transform=proj, shading="auto") | |
| ax.add_feature(cfeature.COASTLINE.with_scale("50m"), linewidth=0.7) | |
| ax.add_feature(cfeature.BORDERS.with_scale("50m"), linewidth=0.4, alpha=0.6) | |
| ax.set_extent([float(da["lon"].min()), float(da["lon"].max()), | |
| float(da["lat"].min()), float(da["lat"].max())], crs=proj) | |
| ax.set_title(f"{label} — {which}", fontsize=11) | |
| fig.colorbar(m, ax=axes[row, :].tolist(), shrink=0.85, label=f"{label} [{units}]") | |
| fig.suptitle(title, fontsize=13) | |
| tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False) | |
| fig.savefig(tmp.name, dpi=110) | |
| plt.close(fig) | |
| return tmp.name | |
| def _duration(month="February", lead_days=5.25, lat_center=44.0, lon_center=-124.0, | |
| region_size=16, n_ens=1, seed=0, *args, **kwargs): | |
| steps = max(1, int(round(float(lead_days) * 4))) | |
| patches = 1 if int(region_size) <= 16 else 4 | |
| # measured on ZeroGPU: ~16 s fixed + 0.25 s / ACE step + 6.5 s per HiRO patch x member | |
| est = 16 + 0.25 * steps + 6.5 * patches * int(n_ens) | |
| return int(np.ceil(1.15 * est)) | |
| def _run(month, lead_days, lat_center, lon_center, region_size, n_ens, seed): | |
| t0 = time.perf_counter() | |
| # ZeroGPU runs this in a worker thread without an asyncio loop; fme's zarr | |
| # reader calls asyncio.get_event_loop(), so give the thread its own loop. | |
| try: | |
| asyncio.get_event_loop() | |
| except RuntimeError: | |
| asyncio.set_event_loop(asyncio.new_event_loop()) | |
| month_idx = MONTHS.index(month) + 1 | |
| ic_time = f"2023-{month_idx:02d}-01T00:00:00" | |
| n_steps = max(1, int(round(float(lead_days) * 4))) | |
| work = tempfile.mkdtemp(prefix="hiro_") | |
| try: | |
| ace_zarr = P.run_ace(MODELS, IC_PATH, FORCING_DIR, ic_time, n_steps, | |
| os.path.join(work, "ace"), int(seed)) | |
| coarse_ds = xr.open_zarr(ace_zarr, decode_times=xr.coders.CFDatetimeCoder(use_cftime=True)) | |
| if "sample" in coarse_ds.dims: | |
| coarse_ds = coarse_ds.isel(sample=0) | |
| event = coarse_ds["time"].values[-1] | |
| event_str = event.strftime("%Y-%m-%dT%H:%M:%S") | |
| t_ace = time.perf_counter() - t0 | |
| lat0, lat1, lon0, lon1 = _region(float(lat_center), float(lon_center), float(region_size)) | |
| hiro_zarr = P.run_hiro(MODELS, ace_zarr, event_str, lat0, lat1, lon0, lon1, | |
| int(n_ens), os.path.join(work, "hiro")) | |
| fine_ds = xr.open_zarr(hiro_zarr, decode_times=False).isel(time=0).load() | |
| fine_ds = fine_ds.rename({k: v for k, v in {"latitude": "lat", "longitude": "lon"}.items() | |
| if k in fine_ds.variables}) | |
| coarse = coarse_ds[P.VARS].isel(time=-1).load() | |
| coarse = _subset_lon(coarse.sel(lat=slice(lat0, lat1)), lon0, lon1) | |
| t_total = time.perf_counter() - t0 | |
| return coarse, fine_ds, event_str, t_ace, t_total | |
| finally: | |
| shutil.rmtree(work, ignore_errors=True) | |
| def downscale(month: str = "February", lead_days: float = 5.25, lat_center: float = 44.0, | |
| lon_center: float = -124.0, region_size: int = 16, n_ens: int = 1, | |
| seed: int = 0, progress=gr.Progress(track_tqdm=True)): | |
| """Simulate global climate with ACE2S-X-SHiELD and downscale a region to 3 km with HiROv2. | |
| ACE2S is run forward (6-hour steps) from the 2023 X-SHiELD initial condition on the | |
| first day of `month` for `lead_days`; the final global ~100 km state is then | |
| downscaled with the HiROv2 diffusion model over a square box centred on | |
| (`lat_center`, `lon_center`). | |
| Args: | |
| month: Month of 2023 whose 1st-day initial condition starts the simulation. | |
| lead_days: Simulated days after the initial condition (multiples of 0.25). | |
| lat_center: Box centre latitude in degrees (model trained on 66S-70N). | |
| lon_center: Box centre longitude in degrees (-180..180). | |
| region_size: Box side length in degrees (16 = one HiRO patch, 32 = 2x2 patches). | |
| n_ens: Number of HiRO ensemble members to generate (member 0 is plotted). | |
| seed: Random seed for the stochastic ACE2S rollout. | |
| Returns: | |
| Comparison figure (PNG path), NetCDF file with the 3 km fields, and a status string. | |
| """ | |
| coarse, fine, event_str, t_ace, t_total = _run( | |
| month, lead_days, lat_center, lon_center, region_size, n_ens, seed) | |
| fine0 = fine.isel(ensemble=0) if "ensemble" in fine.dims else fine | |
| title = (f"Simulated {event_str.replace('T', ' ')} UTC | ACE2S initialised " | |
| f"{month} 1, 2023 + {lead_days:g} days") | |
| fig_path = make_figure(coarse, fine0, title) | |
| nc = tempfile.NamedTemporaryFile(suffix="_hiro_3km.nc", delete=False) | |
| fine.to_netcdf(nc.name) | |
| status = (f"Valid time {event_str} · fine grid {fine0.sizes.get('lat')}×{fine0.sizes.get('lon')} " | |
| f"· ACE {t_ace:.1f}s, total GPU {t_total:.1f}s") | |
| return fig_path, nc.name, status | |
| CSS = """ | |
| #col-container { max-width: 1150px; margin: 0 auto; } | |
| .dark .gradio-container { color: var(--body-text-color); } | |
| """ | |
| with gr.Blocks(title="HiRO-ACEv2") as demo: | |
| with gr.Column(elem_id="col-container"): | |
| gr.Markdown( | |
| "# HiRO-ACEv2 — 3 km climate downscaling\n" | |
| "Runs the **ACE2S-X-SHiELD** stochastic climate emulator globally at ~100 km, then " | |
| "downscales a region to **3 km** precipitation, 10 m winds and sea-level pressure " | |
| "with the **HiROv2** diffusion model. This is a *climate simulation*, not a weather " | |
| "forecast. " | |
| "[Model](https://huggingface.co/allenai/HiRO-ACEv2) · " | |
| "[Paper](http://arxiv.org/abs/2610.17693) · [Code](https://github.com/ai2cm/ace)" | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| month = gr.Dropdown(MONTHS, value="February", label="Initial condition (2023)") | |
| lead_days = gr.Slider(0.25, 14, value=5.25, step=0.25, | |
| label="Simulated days after the initial condition") | |
| lat_center = gr.Slider(-58, 62, value=44.0, step=0.5, label="Region centre latitude") | |
| lon_center = gr.Slider(-180, 180, value=-124.0, step=0.5, | |
| label="Region centre longitude") | |
| run = gr.Button("Simulate & downscale", variant="primary") | |
| with gr.Accordion("Advanced settings", open=False): | |
| region_size = gr.Radio([16, 32], value=16, label="Region size (degrees)") | |
| n_ens = gr.Slider(1, 4, value=1, step=1, label="HiRO ensemble members") | |
| seed = gr.Number(value=0, precision=0, label="ACE2S seed") | |
| with gr.Column(scale=2): | |
| fig = gr.Image(label="ACE2S (left) vs HiROv2 3 km (right)", type="filepath") | |
| status = gr.Textbox(label="Status", interactive=False) | |
| nc_file = gr.File(label="3 km fields (NetCDF)") | |
| inputs = [month, lead_days, lat_center, lon_center, region_size, n_ens, seed] | |
| run.click(downscale, inputs=inputs, outputs=[fig, nc_file, status], api_name="downscale") | |
| gr.Examples( | |
| examples=[ | |
| ["February", 5.25, 44.0, -124.0], # Pacific Northwest (authors' WA AR region) | |
| ["September", 3.0, 22.0, -80.0], # Caribbean / Florida | |
| ["July", 7.0, 28.0, 86.0], # Himalaya monsoon | |
| ["January", 4.0, 46.0, 10.0], # Alps | |
| ["August", 6.0, 33.0, 135.0], # Japan | |
| ], | |
| inputs=[month, lead_days, lat_center, lon_center], | |
| outputs=[fig, nc_file, status], | |
| fn=downscale, | |
| cache_examples=True, | |
| cache_mode="lazy", | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch(show_error=True, theme=gr.themes.Citrus(), css=CSS, mcp_server=True) | |