tobias/entropice
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
entropice/steps/s1_1_era5/era5.py

577 lines
22 KiB
Python
Raw Normal View History

Restructure to steps
2025-10-21 18:42:01 +02:00
"""Download and preprocess ERA5 data.
Variables of Interest:
- 2 metre temperature (t2m) [instant]
- Total precipitation (tp) [accum]
- Snow Fall (sf) [accum]
- Snow cover (snowc) [instant]
- Snow depth (sde) [instant]
- Surface sensible heat flux (sshf) [accum]
- Lake ice bottom temperature (lblt) [instant]
Naming patterns:
- Instant Variables are downloaded already as statistically aggregated (lossy),
therefore their names get the aggregation as suffix
- Accumulation Variables are downloaded as totals, their names stay the same
Daily Variables (downloaded from hourly data):
- t2m_max
- t2m_min
- snowc_mean
- sde_mean
- lblt_max
- tp
- sf
- sshf
Derived Daily Variables:
- t2m_daily_avg
- t2m_daily_range
- t2m_daily_skew
- thawing_degree_days
- freezing_degree_days
- thawing_days
- freezing_days
- precipitation_occurrences
- snowfall_occurrences
- snow_isolation (snowc * sde)
Monthly Variables:
- t2m_monthly_max
- t2m_monthly_min
- tp_monthly_sum
- sf_monthly_sum
- snowc_monthly_mean
- sde_monthly_mean
- sshf_monthly_sum
- lblt_monthly_max
- t2m_monthly_avg
- t2m_monthly_range_avg
- t2m_monthly_skew_avg
- thawing_degree_days_monthly
- freezing_degree_days_monthly
- thawing_days_monthly
- freezing_days_monthly
- precipitation_occurrences_monthly TODO: Rename to precipitation_days_monthly?
- snowfall_occurrences_monthly TODO: Rename to snowfall_days_monthly?
- snow_isolation_monthly_mean
Yearly Variables:
- TODO
# TODO Variables:
- Day of first thaw (yearly)
- Day of last thaw (yearly)
- Thawing period length (yearly)
- Freezing period length (yearly)
Author: Tobias Hölzer
Date: 09. June 2025
"""
import os
import time
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Literal
import cyclopts
import dask.distributed as dd
import geopandas as gpd
import odc.geo
import odc.geo.xr
import pandas as pd
import shapely
import shapely.ops
import xarray as xr
from numcodecs.zarr3 import Blosc
from rich import pretty, print, traceback
from rich.progress import track
from shapely.geometry import LineString, Polygon
traceback.install(show_locals=True, suppress=[cyclopts, xr, pd])
pretty.install()
cli = cyclopts.App()
# TODO: Directly handle download on a grid level - this is more what the zarr access is indented to do
DATA_DIR = Path(os.environ.get("DATA_DIR", "data")) / "entropyc-rts"
ERA5_DIR = DATA_DIR / "era5"
DAILY_RAW_PATH = ERA5_DIR / "daily_raw.zarr"
def _get_grid_paths(
agg: Literal["daily", "monthly", "summer", "winter", "yearly"],
grid: Literal["hex", "healpix"],
level: int,
):
gridname = f"permafrost_{grid}{level}"
aligned_path = ERA5_DIR / f"{agg}_{gridname}.zarr"
return aligned_path
min_lat = 50
max_lat = 83.7 # Ensures the right Chunks Size (90 - 64 / 10 + 0.1)
min_time = "1990-01-01"
max_time = "2024-12-31"
today = time.strftime("%Y-%m-%d")
# ================
# === Download ===
# ================
def create_encoding(ds: xr.Dataset):
"""Create compression encoding for zarr dataset storage.
Creates Blosc compression configuration for all data variables and coordinates
in the dataset using zstd compression with level 9.
Args:
ds (xr.Dataset): The xarray Dataset to create encoding for.
Returns:
dict: Encoding dictionary with compression settings for each variable.
"""
# encoding = {var: {"compressors": BloscCodec(cname="zlib", clevel=9)} for var in ds.data_vars}
encoding = {var: {"compressors": Blosc(cname="zstd", clevel=9)} for var in [*ds.data_vars, *ds.coords]}
return encoding
def download_daily_aggregated():
"""Download and aggregate ERA5 data to daily resolution.
Downloads ERA5 reanalysis data from the DESTINE Earth Data Hub and aggregates
it to daily resolution. Includes temperature extremes, precipitation, snow,
and surface heat flux variables.
The function downloads hourly data and creates daily aggregates:
- Temperature: daily min/max
- Precipitation and snowfall: daily totals
- Snow cover and depth: daily means
- Surface heat flux: daily totals
- Lake ice temperature: daily max
The aggregated data is saved to a zarr file with compression.
"""
era5 = xr.open_dataset(
"https://data.earthdatahub.destine.eu/era5/reanalysis-era5-land-no-antartica-v0.zarr",
storage_options={"client_kwargs": {"trust_env": True}},
chunks={},
# chunks={},
engine="zarr",
).rename({"valid_time": "time"})
subset = {
"latitude": slice(max_lat, min_lat),
}
# Compute the clostest chunk-start to min_time, to avoid problems with cropped chunks at the start
tchunksize = era5.chunksizes["time"][0]
era5_chunk_starts = pd.date_range(era5.time.min().item(), era5.time.max().item(), freq=f"{tchunksize}h")
closest_chunk_start = era5_chunk_starts[
era5_chunk_starts.get_indexer([pd.to_datetime(min_time)], method="ffill")[0]
]
subset["time"] = slice(str(closest_chunk_start), max_time)
era5 = era5.sel(**subset)
daily_raw = xr.merge(
[
# Instant
era5.t2m.resample(time="1D").max().rename("t2m_max"),
era5.t2m.resample(time="1D").min().rename("t2m_min"),
era5.snowc.resample(time="1D").mean().rename("snowc_mean"),
era5.sde.resample(time="1D").mean().rename("sde_mean"),
era5.lblt.resample(time="1D").max().rename("lblt_max"),
# Accum
era5.tp.resample(time="1D").sum().rename("tp"),
era5.sf.resample(time="1D").sum().rename("sf"),
era5.sshf.resample(time="1D").sum().rename("sshf"),
]
)
# Assign attributes
daily_raw["t2m_max"].attrs = {"long_name": "Daily maximum 2 metre temperature", "units": "K"}
daily_raw["t2m_min"].attrs = {"long_name": "Daily minimum 2 metre temperature", "units": "K"}
daily_raw["tp"].attrs = {"long_name": "Daily total precipitation", "units": "m"}
daily_raw["sf"].attrs = {"long_name": "Daily total snow fall", "units": "m"}
daily_raw["snowc_mean"].attrs = {"long_name": "Daily mean snow cover", "units": "m"}
daily_raw["sde_mean"].attrs = {"long_name": "Daily mean snow depth", "units": "m"}
daily_raw["sshf"].attrs = {"long_name": "Daily total surface sensible heat flux", "units": "J/m²"}
daily_raw["lblt_max"].attrs = {"long_name": "Daily maximum lake ice bottom temperature", "units": "K"}
daily_raw = daily_raw.odc.assign_crs("epsg:4326")
daily_raw = daily_raw.drop_vars(["surface", "number", "depthBelowLandLayer"])
daily_raw.to_zarr(DAILY_RAW_PATH, mode="w", encoding=create_encoding(daily_raw), consolidated=False)
@cli.command
def download():
"""Download ERA5 data using Dask cluster for parallel processing.
Creates a local Dask cluster and downloads daily aggregated ERA5 data.
The cluster is configured with a single worker with 10 threads and 100GB
memory limit for optimal performance.
"""
with (
dd.LocalCluster(n_workers=1, threads_per_worker=10, memory_limit="100GB") as cluster,
dd.Client(cluster) as client,
):
print(client)
print(client.dashboard_link)
download_daily_aggregated()
print(f"Downloaded and aggregated ERA5 data to {DAILY_RAW_PATH.resolve()}.")
# ===========================
# === Spatial Aggregation ===
# ===========================
def _crosses_antimeridian(geom: Polygon) -> bool:
coords = shapely.get_coordinates(geom)
crosses_any_meridian = (coords[:, 0] > 0).any() and (coords[:, 0] < 0).any()
return crosses_any_meridian and abs(coords[:, 0]).max() > 90
def _split_antimeridian_cell(geom: Polygon) -> list[Polygon]:
# Assumes that it is a antimeridian hex
coords = shapely.get_coordinates(geom)
for i in range(coords.shape[0]):
if coords[i, 0] < 0:
coords[i, 0] += 360
geom = Polygon(coords)
antimeridian = LineString([[180, -90], [180, 90]])
polys = shapely.ops.split(geom, antimeridian)
return list(polys.geoms)
def _check_geobox(geobox):
x, y = geobox.shape
return x > 1 and y > 1
def extract_cell_data(idx: int, geom: Polygon) -> xr.Dataset:
"""Extract ERA5 data for a specific grid cell geometry.
Extracts and spatially averages ERA5 data within the bounds of a grid cell.
Handles antimeridian-crossing cells by splitting them appropriately.
Args:
idx (int): Index of the grid cell.
geom (Polygon): Polygon geometry of the grid cell.
Returns:
xr.Dataset: The computed cell dataset
"""
daily_raw = xr.open_zarr(DAILY_RAW_PATH, consolidated=False).set_coords("spatial_ref")
# cell.geometry is a shapely Polygon
if not _crosses_antimeridian(geom):
geoms = [geom]
# Split geometry in case it crossed antimeridian
else:
geoms = _split_antimeridian_cell(geom)
cell_data = []
for geom in geoms:
geom = odc.geo.Geometry(geom, crs="epsg:4326")
if not _check_geobox(daily_raw.odc.geobox.enclosing(geom)):
continue
# TODO: use mean for instant variables, sum for accum variables
cell_data.append(daily_raw.odc.crop(geom).drop_vars("spatial_ref").mean(["latitude", "longitude"]))
if len(cell_data) == 0:
return False
elif len(cell_data) == 1:
cell_data = cell_data[0]
else:
cell_data = xr.concat(cell_data, dim="part").mean("part")
cell_data = cell_data.expand_dims({"cell": [idx]}).compute()
return cell_data
@cli.command
def spatial_agg(
grid: Literal["hex", "healpix"],
level: int,
n_workers: int = 10,
executor: Literal["threads", "processes"] = "threads",
):
"""Perform spatial aggregation of ERA5 data to grid cells.
Loads a grid and spatially aggregates ERA5 data to each grid cell using
parallel processing. Creates an empty zarr file first, then fills it
with extracted data for each cell.
Args:
grid ("hex" | "healpix"): Grid type.
level (int): Grid resolution level.
n_workers (int, optional): Number of parallel workers to use. Defaults to 10.
executor ("threads" | "processes"): The type of parallel executor pool to use. Defaults to threads.
"""
gridname = f"permafrost_{grid}{level}"
daily_grid_path = _get_grid_paths("daily", grid, level)
grid = gpd.read_parquet(DATA_DIR / f"grids/{gridname}_grid.parquet")
# Create an empty zarr array with the right dimensions
daily_raw = xr.open_zarr(DAILY_RAW_PATH, consolidated=False).set_coords("spatial_ref")
assert {"latitude", "longitude", "time"} == set(daily_raw.dims), (
f"Expected dims ('latitude', 'longitude', 'time'), got {daily_raw.dims}"
)
assert daily_raw.odc.crs == "epsg:4326", f"Expected CRS 'epsg:4326', got {daily_raw.odc.crs}"
daily = (
xr.zeros_like(daily_raw.isel(latitude=0, longitude=0))
.expand_dims({"cell": [idx for idx, _ in grid.iterrows()]})
.chunk({"cell": min(len(grid), 1000), "time": len(daily_raw.time)}) # ~50MB chunks
)
daily.to_zarr(daily_grid_path, mode="w", consolidated=False, encoding=create_encoding(daily))
print(f"Created empty zarr at {daily_grid_path.resolve()} with shape {daily.sizes}.")
print(f"Starting spatial matching of {len(grid)} cells with {n_workers} workers...")
ExecutorCls = ThreadPoolExecutor if executor == "threads" else ProcessPoolExecutor
with ExecutorCls(max_workers=n_workers) as executor:
futures = {
executor.submit(extract_cell_data, idx, row.geometry): idx
for idx, row in grid.to_crs("epsg:4326").iterrows()
}
for future in track(as_completed(futures), total=len(futures), description="Processing cells"):
idx = futures[future]
try:
cell_data = future.result()
if not cell_data:
print(f"Cell {idx} did not overlap with ERA5 data.")
cell_data.to_zarr(daily_grid_path, region="auto", consolidated=False)
print(f"Successfully written cell {idx}")
except Exception as e:
print(f"{type(e)} processing cell {idx}: {e}")
print("Finished spatial matching.")
# ============================
# === Temporal Aggregation ===
# ============================
def daily_enrich(grid: Literal["hex", "healpix"], level: int) -> xr.Dataset:
"""Enrich daily ERA5 data with derived climate variables.
Loads spatially aligned ERA5 data and computes additional climate variables.
Creates derived variables including temperature statistics, degree days, and occurrence indicators.
Derived variables include:
- Daily average and range temperature
- Temperature skewness
- Thawing and freezing degree days
- Thawing and freezing day counts
- Precipitation and snowfall occurrences
- Snow isolation index
Args:
grid ("hex", "healpix"): Grid type.
level (int): Grid resolution level.
Returns:
xr.Dataset: Enriched dataset with original and derived variables.
"""
daily_grid_path = _get_grid_paths("daily", grid, level)
daily = xr.open_zarr(daily_grid_path, consolidated=False).set_coords("spatial_ref")
assert {"cell", "time"} == set(daily.dims), f"Expected dims ('cell', 'time'), got {daily.dims}"
# Formulas based on Groeke et. al. (2025) Stochastic Weather generation...
daily["t2m_avg"] = (daily.t2m_max + daily.t2m_min) / 2
daily.t2m_avg.attrs = {"long_name": "Daily average 2 metre temperature", "units": "K"}
daily["t2m_range"] = daily.t2m_max - daily.t2m_min
daily.t2m_range.attrs = {"long_name": "Daily range of 2 metre temperature", "units": "K"}
daily["t2m_skew"] = (daily.t2m_avg - daily.t2m_min) / daily.t2m_range
daily.t2m_skew.attrs = {"long_name": "Daily skewness of 2 metre temperature"}
daily["thawing_degree_days"] = (daily.t2m_avg - 273.15).clip(min=0)
daily.thawing_degree_days.attrs = {"long_name": "Thawing degree days", "units": "K"}
daily["freezing_degree_days"] = (273.15 - daily.t2m_avg).clip(min=0)
daily.freezing_degree_days.attrs = {"long_name": "Freezing degree days", "units": "K"}
daily["thawing_days"] = (daily.t2m_avg > 273.15).astype(int)
daily.thawing_days.attrs = {"long_name": "Thawing days"}
daily["freezing_days"] = (daily.t2m_avg < 273.15).astype(int)
daily.freezing_days.attrs = {"long_name": "Freezing days"}
daily["precipitation_occurrences"] = (daily.tp > 0).astype(int)
daily.precipitation_occurrences.attrs = {"long_name": "Precipitation occurrences"}
daily["snowfall_occurrences"] = (daily.sf > 0).astype(int)
daily.snowfall_occurrences.attrs = {"long_name": "Snowfall occurrences"}
daily["snow_isolation"] = daily.snowc_mean * daily.sde_mean
daily.snow_isolation.attrs = {"long_name": "Snow isolation"}
return daily
def monthly_aggregate(grid: Literal["hex", "healpix"], level: int):
"""Aggregate enriched daily ERA5 data to monthly resolution.
Takes the enriched daily ERA5 data and creates monthly aggregates using
appropriate statistical functions for each variable type. Temperature
variables use min/max/mean, accumulation variables use sums, and derived
variables use appropriate aggregations.
The aggregated monthly data is saved to a zarr file for further processing.
Args:
grid ("hex", "healpix"): Grid type.
level (int): Grid resolution level.
"""
daily = daily_enrich(grid, level)
assert {"cell", "time"} == set(daily.dims), f"Expected dims ('cell', 'time'), got {daily.dims}"
# Monthly aggregates
monthly = xr.merge(
[
# Original variables
daily.t2m_min.resample(time="1ME").min().rename("t2m_min"),
daily.t2m_max.resample(time="1ME").max().rename("t2m_max"),
daily.snowc_mean.resample(time="1ME").mean().rename("snowc_mean"),
daily.sde_mean.resample(time="1ME").mean().rename("sde_mean"),
daily.lblt_max.resample(time="1ME").max().rename("lblt_max"),
daily.tp.resample(time="1ME").sum().rename("tp"),
daily.sf.resample(time="1ME").sum().rename("sf"),
daily.sshf.resample(time="1ME").sum().rename("sshf"),
# Enriched variables
daily.t2m_avg.resample(time="1ME").mean().rename("t2m_avg"),
daily.t2m_range.resample(time="1ME").mean().rename("t2m_mean_range"),
daily.t2m_skew.resample(time="1ME").mean().rename("t2m_mean_skew"),
daily.thawing_degree_days.resample(time="1ME").sum().rename("thawing_degree_days"),
daily.freezing_degree_days.resample(time="1ME").sum().rename("freezing_degree_days"),
daily.thawing_days.resample(time="1ME").sum().rename("thawing_days"),
daily.freezing_days.resample(time="1ME").sum().rename("freezing_days"),
daily.precipitation_occurrences.resample(time="1ME").sum().rename("precipitation_occurrences"),
daily.snowfall_occurrences.resample(time="1ME").sum().rename("snowfall_occurrences"),
daily.snow_isolation.resample(time="1ME").mean().rename("snow_mean_isolation"),
]
)
monthly_grid_path = _get_grid_paths("monthly", grid, level)
monthly.to_zarr(monthly_grid_path, mode="w", encoding=create_encoding(monthly), consolidated=False)
def yearly_aggregate(monthly: xr.Dataset) -> xr.Dataset:
"""Aggregate monthly ERA5 data to yearly resolution.
Takes monthly aggregated data and creates yearly aggregates using a shifted
calendar (October to September) to better capture Arctic seasonal patterns.
Args:
monthly (xr.Dataset): The monthly aggregates
Returns:
xr.Dataset: The aggregated dataset
"""
return xr.merge(
[
# Original variables
monthly.t2m_min.resample(time="1YE").min().rename("t2m_min"),
monthly.t2m_max.resample(time="1YE").max().rename("t2m_max"),
monthly.snowc_mean.resample(time="1YE").mean().rename("snowc_mean"),
monthly.sde_mean.resample(time="1YE").mean().rename("sde_mean"),
monthly.lblt_max.resample(time="1YE").max().rename("lblt_max"),
monthly.tp.resample(time="1YE").sum().rename("tp"),
monthly.sf.resample(time="1YE").sum().rename("sf"),
monthly.sshf.resample(time="1YE").sum().rename("sshf"),
# Enriched variables
monthly.t2m_avg.resample(time="1YE").mean().rename("t2m_avg"),
# TODO: Check if this is correct -> use daily / hourly data instead for range and skew?
monthly.t2m_mean_range.resample(time="1YE").mean().rename("t2m_mean_range"),
monthly.t2m_mean_skew.resample(time="1YE").mean().rename("t2m_mean_skew"),
monthly.thawing_degree_days.resample(time="1YE").sum().rename("thawing_degree_days"),
monthly.freezing_degree_days.resample(time="1YE").sum().rename("freezing_degree_days"),
monthly.thawing_days.resample(time="1YE").sum().rename("thawing_days"),
monthly.freezing_days.resample(time="1YE").sum().rename("freezing_days"),
monthly.precipitation_occurrences.resample(time="1YE").sum().rename("precipitation_occurrences"),
monthly.snowfall_occurrences.resample(time="1YE").sum().rename("snowfall_occurrences"),
monthly.snow_mean_isolation.resample(time="1YE").mean().rename("snow_mean_isolation"),
]
)
def yearly_and_seasonal_aggregate(grid: Literal["hex", "healpix"], level: int):
"""Aggregate monthly ERA5 data to yearly resolution with seasonal splits.
Takes monthly aggregated data and creates yearly aggregates using a shifted
calendar (October to September) to better capture Arctic seasonal patterns.
Creates separate aggregates for full year, winter (Oct-Apr), and summer
(May-Sep) periods.
The first and last incomplete years are excluded from the analysis.
Winter months are defined as months 1-7 in the shifted calendar,
and summer months are 8-12.
The final dataset includes yearly, winter, and summer aggregates for all
climate variables, saved to a zarr file.
Args:
grid ("hex", "healpix"): Grid type.
level (int): Grid resolution level.
"""
monthly_grid_path = _get_grid_paths("monthly", grid, level)
monthly = xr.open_zarr(monthly_grid_path, consolidated=False).set_coords("spatial_ref")
assert {"cell", "time"} == set(monthly.dims), f"Expected dims ('cell', 'time'), got {monthly.dims}"
valid_years = slice(str(monthly.time.min().dt.year.item() + 1), str(monthly.time.max().dt.year.item()))
# Summer aggregates
summer = yearly_aggregate(monthly.sel(time=monthly.time.dt.month.isin([5, 6, 7, 8, 9])).sel(time=valid_years))
# Yearly aggregates (shifted by +8 months to start in Oktober, first and last years will be cropped)
monthly_shifted = monthly.copy()
monthly_shifted["time"] = monthly_shifted.get_index("time") + pd.DateOffset(months=8)
monthly_shifted = monthly_shifted.sel(time=valid_years)
yearly = yearly_aggregate(monthly_shifted)
# Winter aggregates (shifted by +8 months to start in Oktober, first and last years will be cropped)
monthly_shifted = monthly.copy().sel(time=monthly.time.dt.month.isin([1, 2, 3, 4, 10, 11, 12]))
monthly_shifted["time"] = monthly_shifted.get_index("time") + pd.DateOffset(months=8)
monthly_shifted = monthly_shifted.sel(time=valid_years)
winter = yearly_aggregate(monthly_shifted)
yearly_grid_path = _get_grid_paths("yearly", grid, level)
yearly.to_zarr(yearly_grid_path, mode="w", encoding=create_encoding(yearly), consolidated=False)
winter_grid_path = _get_grid_paths("winter", grid, level)
winter.to_zarr(winter_grid_path, mode="w", encoding=create_encoding(winter), consolidated=False)
summer_grid_path = _get_grid_paths("summer", grid, level)
summer.to_zarr(summer_grid_path, mode="w", encoding=create_encoding(summer), consolidated=False)
@cli.command
def temporal_agg(n_workers: int = 10):
"""Perform temporal aggregation of ERA5 data using Dask cluster.
Creates a Dask cluster and runs both monthly and yearly aggregation
functions to generate temporally aggregated climate datasets. The
processing uses parallel workers for efficient computation.
Args:
n_workers (int, optional): Number of Dask workers to use. Defaults to 10.
"""
with (
dd.LocalCluster(n_workers=n_workers, threads_per_worker=20, memory_limit="10GB") as cluster,
dd.Client(cluster) as client,
):
print(client)
print(client.dashboard_link)
monthly_aggregate()
yearly_and_seasonal_aggregate()
print("Enriched ERA5 data with additional features and aggregated it temporally.")
if __name__ == "__main__":
cli()
Reference in a new issue Copy permalink