diff --git a/Snakefile_cutouts b/Snakefile_cutouts index e15a757..2c8ed40 100644 --- a/Snakefile_cutouts +++ b/Snakefile_cutouts @@ -6,8 +6,10 @@ PyPSA-GB: Cutout Generation Workflow ║ ║ ║ This workflow acquires weather cutouts using a tiered strategy: ║ ║ 1. Data dir: check for cached copy in data/atlite/cutouts/ (instant) ║ -║ 2. Zenodo: download pre-built cutouts (~minutes per file) ║ -║ 3. Atlite/ERA5: full download from CDS API (~2-4 hours per file) ║ +║ 2. Earthmover: build from public Arraylake ERA5 (opt-in, ~minutes, ║ +║ no full-file download, any year 1940-present; needs a free token) ║ +║ 3. Zenodo: download pre-built cutouts (~minutes per file) ║ +║ 4. Atlite/ERA5: full download from CDS API (~2-4 hours per file) ║ ║ ║ ║ Snakemake skips re-downloading if output file already exists. ║ ║ Zenodo repository: https://zenodo.org/records/18325225 ║ diff --git a/config/cutouts_config.yaml b/config/cutouts_config.yaml index eaa0807..f04b4d4 100644 --- a/config/cutouts_config.yaml +++ b/config/cutouts_config.yaml @@ -79,13 +79,29 @@ download: zenodo: # Enable/disable Zenodo downloads (set to false to always use atlite) enabled: true - + # Verify MD5 checksum after download (recommended) verify_checksum: true - + # Zenodo record ID (update if a new version is published) record_id: "18325225" +# ──────────────────────────────────────────────────────────── +# EARTHMOVER ERA5 (opt-in, fastest tier — tried before Zenodo) +# ──────────────────────────────────────────────────────────── +# Builds the cutout directly from the public Arraylake ERA5 dataset +# (earthmover-public/era5): reads only the GB box for the year, so there is +# no ~700 MB per-year download and any year 1940-present is available. +# +# Requires the optional dependencies and a (free) Arraylake account: +# pip install arraylake "zarr>=3" icechunk pcodec numcodecs +# arraylake auth login # or set ARRAYLAKE_TOKEN +# +# Safe to leave enabled: if the deps/credentials are absent it is skipped +# automatically and acquisition falls through to Zenodo, then CDS. +earthmover: + enabled: false + # ──────────────────────────────────────────────────────────── # File Paths # ──────────────────────────────────────────────────────────── diff --git a/docs/source/getting_started/installation.md b/docs/source/getting_started/installation.md index b24d323..597a463 100644 --- a/docs/source/getting_started/installation.md +++ b/docs/source/getting_started/installation.md @@ -87,8 +87,9 @@ snakemake -n -p For renewable generation profiles, you need ERA5 weather data "cutouts". PyPSA-GB uses a **tiered acquisition strategy** to minimize download times: 1. **Data directory** - Check `data/atlite/cutouts/` for cached copy (instant) -2. **Zenodo** - Download pre-built cutouts from [Zenodo repository](https://zenodo.org/records/18325225) (~5-10 minutes per year, years 2010-2024) -3. **ERA5 API** - Full download via atlite as fallback (~2-4 hours per year) +2. **Earthmover** - Build from the public Arraylake ERA5 dataset, opt-in (~minutes, no full-file download, any year 1940-present) +3. **Zenodo** - Download pre-built cutouts from [Zenodo repository](https://zenodo.org/records/18325225) (~5-10 minutes per year, years 2010-2024) +4. **ERA5 API** - Full download via atlite as fallback (~2-4 hours per year) ### Quick Start (Recommended) @@ -109,6 +110,30 @@ snakemake -s Snakefile_cutouts --cores 1 **No CDS API credentials required** for years 2010-2024 when using Zenodo! +### Earthmover ERA5 (fastest, opt-in) + +The Earthmover tier builds cutouts directly from the public Arraylake ERA5 dataset +([`earthmover-public/era5`](https://docs.earthmover.io/)). It reads only the GB box for the +requested year, so there is **no ~700 MB per-year download** and **any year 1940-present** is +available (not just 2010-2024). It is tried before Zenodo when enabled. + +It needs a few optional packages and a free Arraylake account: + +```bash +pip install arraylake "zarr>=3" icechunk pcodec numcodecs +arraylake auth login # one-time; or set ARRAYLAKE_TOKEN + +# enable the tier +# config/cutouts_config.yaml: +# earthmover: +# enabled: true + +snakemake -s Snakefile_cutouts --cores 1 +``` + +If the optional packages or credentials are missing, the tier is skipped automatically and +acquisition falls through to Zenodo, then the CDS API — so it is safe to leave enabled. + ### Manual ERA5 Download (Advanced) For years outside 2010-2024 or if you prefer direct ERA5 download: diff --git a/scripts/utilities/download_cutouts.py b/scripts/utilities/download_cutouts.py index 984b403..e74f61a 100644 --- a/scripts/utilities/download_cutouts.py +++ b/scripts/utilities/download_cutouts.py @@ -220,14 +220,43 @@ def download_with_atlite(year, output_path): logger.info(f" atlite download complete for {year}") +def is_earthmover_available(): + """True if the Earthmover tier can be used (arraylake importable + credentials present).""" + try: + import arraylake # noqa: F401 + except ImportError: + return False + # Either ARRAYLAKE_TOKEN, or a cached `arraylake auth login` session. + if os.environ.get("ARRAYLAKE_TOKEN"): + return True + try: + from arraylake import Client + + Client() # raises if not logged in + return True + except Exception: + return False + + +def download_with_earthmover(year, output_path, bounds=None): + """Build the cutout from the public Earthmover ERA5 store (fast, no 700 MB download).""" + from scripts.utilities.earthmover_cutout import build_earthmover_cutout + + logger.info(f"[EARTHMOVER] Building cutout for {year} from earthmover-public/era5...") + build_earthmover_cutout(year, output_path, bounds=bounds) + logger.info(f" Earthmover build complete for {year}") + + def acquire_cutout(year, output_path, data_dir=None, enable_zenodo=True, - verify_checksum=True, zenodo_files=None): + verify_checksum=True, zenodo_files=None, + enable_earthmover=False, earthmover_bounds=None): """ Acquire a cutout file using a tiered strategy: 1. Check if it exists in data_dir (copy) - 2. Try downloading from Zenodo (fast) - 3. Fall back to atlite ERA5 download (slow) + 2. Build from Earthmover ERA5 if enabled (fast; reads only the GB box) + 3. Try downloading from Zenodo (fast) + 4. Fall back to atlite ERA5 download (slow) Note: Snakemake handles checking if output_path already exists, so we don't need to duplicate that logic here. @@ -246,11 +275,16 @@ def acquire_cutout(year, output_path, data_dir=None, enable_zenodo=True, Whether to verify MD5 checksums on Zenodo downloads. zenodo_files : dict or None Pre-fetched Zenodo file metadata. + enable_earthmover : bool + Whether to try the Earthmover ERA5 tier before Zenodo (needs the optional + arraylake deps and a free Arraylake account; skipped if unavailable). + earthmover_bounds : dict or None + {"north", "south", "west", "east"} GB box for the Earthmover build. Returns ------- str - The source of the cutout: "data_dir", "zenodo", or "atlite" + The source of the cutout: "data_dir", "earthmover", "zenodo", or "atlite" """ output_path = Path(output_path) filename = f"uk-{year}.nc" @@ -266,7 +300,15 @@ def acquire_cutout(year, output_path, data_dir=None, enable_zenodo=True, logger.info(f" Copied to {output_path}") return "data_dir" - # --- Step 2: Try Zenodo download --- + # --- Step 2: Try Earthmover (fast: reads only the GB box, no full-file download) --- + if enable_earthmover and is_earthmover_available(): + try: + download_with_earthmover(year, output_path, bounds=earthmover_bounds) + return "earthmover" + except Exception as e: + logger.warning(f" Earthmover build failed ({e}), falling back to Zenodo/atlite...") + + # --- Step 3: Try Zenodo download --- if enable_zenodo and is_available_on_zenodo(filename, zenodo_files): logger.info(f"[ZENODO] Cutout for {year} is available on Zenodo, downloading...") success = download_from_zenodo( diff --git a/scripts/utilities/earthmover_cutout.py b/scripts/utilities/earthmover_cutout.py new file mode 100644 index 0000000..2b7afd3 --- /dev/null +++ b/scripts/utilities/earthmover_cutout.py @@ -0,0 +1,157 @@ +"""Build an atlite ERA5 weather cutout from the public Earthmover/Arraylake ERA5 dataset. + +This is a fast alternative to the CDS ERA5 download (hours) and the Zenodo pre-built cutouts +(~700 MB/year): it reads only the GB box for the requested year directly from object storage and +lets atlite compute the cutout, producing a `uk-.nc` interchangeable with the other tiers. + +Source: the public Arraylake repo ``earthmover-public/era5`` (ERA5 hourly, 0.25°, 1940-present). +It is free but requires a (free) Arraylake account — either run ``arraylake auth login`` once, or +set the ``ARRAYLAKE_TOKEN`` environment variable. Optional dependencies (not in the base install): + + pip install arraylake "zarr>=3" icechunk pcodec numcodecs + +How it works: atlite's ``get_data_`` functions call +``atlite.datasets.era5.retrieve_data(variable=[CDS names], ...)`` and expect raw ERA5 short-name +variables back. The Earthmover ``single/temporal`` group holds exactly those variables, so we swap +``retrieve_data`` for one that slices the Earthmover store — atlite's feature maths and the on-disk +cutout schema are reused unchanged. +""" +from __future__ import annotations + +import itertools +import logging +import os + +import numpy as np +import pandas as pd +import xarray as xr + +logger = logging.getLogger(__name__) + +EARTHMOVER_REPO = "earthmover-public/era5" +EARTHMOVER_GROUP = "single/temporal" # surface (single-level) fields, time-contiguous chunking + +# CDS long name (what atlite requests) -> ERA5 short name (what the store holds). +_CDS_TO_SHORT = { + "10m_u_component_of_wind": "u10", + "10m_v_component_of_wind": "v10", + "100m_u_component_of_wind": "u100", + "100m_v_component_of_wind": "v100", + "forecast_surface_roughness": "fsr", + "surface_net_solar_radiation": "ssr", + "surface_solar_radiation_downwards": "ssrd", + "toa_incident_solar_radiation": "tisr", + "total_sky_direct_solar_radiation_at_surface": "fdir", + "2m_temperature": "t2m", + "soil_temperature_level_4": "stl4", + "2m_dewpoint_temperature": "d2m", + "runoff": "ro", + "geopotential": "z", +} +# atlite reads `.attrs["units"]`; backfill any the store omits. +_DEFAULT_UNITS = { + "u10": "m s**-1", "v10": "m s**-1", "u100": "m s**-1", "v100": "m s**-1", + "fsr": "m", "ssr": "J m**-2", "ssrd": "J m**-2", "tisr": "J m**-2", + "fdir": "J m**-2", "t2m": "K", "stl4": "K", "d2m": "K", "ro": "m", "z": "m**2 s**-2", +} + +_STORE_CACHE: dict[str, xr.Dataset] = {} + + +def _client(): + """Arraylake client. Uses ARRAYLAKE_TOKEN if set, else the cached `arraylake auth login`.""" + from arraylake import Client + + token = os.environ.get("ARRAYLAKE_TOKEN") + return Client(token=token) if token else Client() + + +def _open_store(repo: str = EARTHMOVER_REPO) -> xr.Dataset: + """Open the Earthmover ERA5 group, longitudes normalised to [-180, 180).""" + if repo not in _STORE_CACHE: + session = _client().get_repo(repo).readonly_session("main") + ds = xr.open_zarr(session.store, group=EARTHMOVER_GROUP, consolidated=False) + lon = ((ds["longitude"] + 180) % 360) - 180 # store is 0..360; atlite uses -180..180 + _STORE_CACHE[repo] = ds.assign_coords(longitude=lon).sortby("longitude") + return _STORE_CACHE[repo] + + +def _requested_times(updates: dict) -> pd.DatetimeIndex: + """The exact hours atlite asked for, from its year/month/day/time lists.""" + years = np.atleast_1d(updates["year"]) + months = np.atleast_1d(updates["month"]) + days = np.atleast_1d(updates.get("day", [f"{d:02d}" for d in range(1, 32)])) + hours = np.atleast_1d(updates.get("time", [f"{h:02d}:00" for h in range(24)])) + stamps = [f"{y}-{m}-{d} {hh}" for y, m, d, hh in itertools.product(years, months, days, hours)] + idx = pd.to_datetime(pd.Series(stamps), errors="coerce").dropna() + return pd.DatetimeIndex(idx.unique()).sort_values() + + +def _retrieve_data_earthmover(product=None, chunks=None, tmpdir=None, lock=None, **updates): + """Drop-in replacement for ``atlite.datasets.era5.retrieve_data`` sourcing from Earthmover.""" + ds = _open_store() + variables = np.atleast_1d(updates["variable"]).tolist() + short = [_CDS_TO_SHORT.get(v, v) for v in variables] + missing = [v for v, s in zip(variables, short) if s not in ds] + if missing: + raise KeyError(f"Earthmover ERA5 store lacks variables: {missing}") + + out = ds[short] + if "area" in updates: # area = [North, West, South, East]; store latitude is descending + n, w, s, e = updates["area"] + out = out.sel(latitude=slice(n, s), longitude=slice(w, e)) + + want = _requested_times(updates) + out = out.sel(valid_time=out.indexes["valid_time"].intersection(want)) + + for s in short: + out[s].attrs.setdefault("units", _DEFAULT_UNITS.get(s, "")) + # Eagerly load each (month, box) slice in the main thread: the store's year-long time chunks + # mean any read transiently decompresses full-year tiles, so loading per monthly request keeps + # memory bounded (one buffer at a time) instead of all reads coexisting under the synchronous + # scheduler. Also keeps arraylake reads off dask workers and yields uniform (single) chunks so + # atlite's write avoids "inconsistent chunks along y". + return out.load() + + +def _patch_atlite() -> None: + import atlite.datasets.era5 as era5 + + era5.retrieve_data = _retrieve_data_earthmover + + +def build_earthmover_cutout(year, output_path, bounds=None, features=("wind", "influx", "temperature"), + dx=0.25, dy=0.25): + """Build ``uk-.nc`` for the GB box from the Earthmover ERA5 store. + + Parameters + ---------- + year : int + output_path : str or Path + bounds : dict or None + {"north", "south", "west", "east"}; defaults to the GB box. + features : sequence of str + atlite features to prepare. wind + influx + temperature cover GB wind/solar + (cutout.wind()/cutout.pv()); runoff/height are not needed and not in the store. + """ + import atlite + import dask + + b = bounds or {"north": 61.0, "south": 49.5, "west": -11.0, "east": 2.5} + _patch_atlite() + cutout = atlite.Cutout( + path=str(output_path), + module="era5", + x=slice(b["west"], b["east"]), + y=slice(b["south"], b["north"]), + time=str(year), + dx=dx, + dy=dy, + ) + logger.info(f"Building Earthmover cutout {output_path} ({year}) grid={dict(cutout.coords.sizes)}") + # atlite runs get_data under dask; arraylake's sync() can't run from a dask worker, so use the + # synchronous scheduler. monthly_requests splits the read per month to cap peak memory. + with dask.config.set(scheduler="synchronous"): + cutout.prepare(features=list(features), monthly_requests=True) + logger.info(f" Earthmover cutout complete: {output_path}") + return cutout diff --git a/scripts/utilities/prepare_cutouts.py b/scripts/utilities/prepare_cutouts.py index adadbed..783dde8 100644 --- a/scripts/utilities/prepare_cutouts.py +++ b/scripts/utilities/prepare_cutouts.py @@ -40,7 +40,8 @@ DATA_DIR = "data/atlite/cutouts" -def prepare_cutouts(years, outputs, enable_zenodo=True, verify_checksum=True): +def prepare_cutouts(years, outputs, enable_zenodo=True, verify_checksum=True, + enable_earthmover=False, earthmover_bounds=None): """ Prepare cutouts for the given years using the tiered strategy. @@ -83,6 +84,8 @@ def prepare_cutouts(years, outputs, enable_zenodo=True, verify_checksum=True): enable_zenodo=enable_zenodo, verify_checksum=verify_checksum, zenodo_files=zenodo_files, + enable_earthmover=enable_earthmover, + earthmover_bounds=earthmover_bounds, ) sources[year] = source logger.info(f" Source: {source}") @@ -93,10 +96,11 @@ def prepare_cutouts(years, outputs, enable_zenodo=True, verify_checksum=True): logger.info(f"{'='*60}") for year, source in sources.items(): icon = { - "data_dir": "[CACHED] ", - "zenodo": "[ZENODO] ", - "atlite": "[ERA5] ", - }.get(source, "[?] ") + "data_dir": "[CACHED] ", + "earthmover": "[EARTHMVR] ", + "zenodo": "[ZENODO] ", + "atlite": "[ERA5] ", + }.get(source, "[?] ") logger.info(f" {icon} uk-{year}.nc") logger.info(f"{'='*60}\n") @@ -112,10 +116,17 @@ def prepare_cutouts(years, outputs, enable_zenodo=True, verify_checksum=True): enable_zenodo = zenodo_config.get("enabled", True) verify_checksum = zenodo_config.get("verify_checksum", True) + # Earthmover tier (opt-in; needs the optional arraylake deps + a free Arraylake account) + earthmover_config = snakemake_config.get("earthmover", {}) + enable_earthmover = earthmover_config.get("enabled", False) + earthmover_bounds = snakemake_config.get("era5", {}).get("bounds") + prepare_cutouts( years=years, outputs=outputs, enable_zenodo=enable_zenodo, verify_checksum=verify_checksum, + enable_earthmover=enable_earthmover, + earthmover_bounds=earthmover_bounds, )