Model data prep

pyproject.toml

@@ -34,17 +34,22 @@ python = ">=3.10, <3.13"
 click = "*"
 numpy = "^1.26.4"
 pandas = "^2.2.2"
-rasterra = "^0.5.9"
+rasterra = "^0.5.11"
 shapely = "^2.0.4"
 geopandas = "^0.14.4"
 xarray = "^2024.3.0"
 cdsapi = "^0.7.0"
 matplotlib = "^3.8.4"
 scikit-learn = "^1.4.2"
-rra-tools = "^1.0.6"
+rra-tools = "^1.0.10"
 netcdf4 = "^1.6.5"
 pyarrow = "^16.0.0"
 types-requests = "^2.31.0.20240406"
+types-tqdm = "^4.66.0.20240417"
+gcsfs = "^2024.6.0"
+zarr = "^2.18.2"
+types-pyyaml = "^6.0.12.20240311"
+dask = "^2024.5.2"
 
 [tool.poetry.group.dev.dependencies]
 mkdocstrings = {version = ">=0.23", extras = ["python"]}
@@ -90,6 +95,14 @@ ignore = [
     "RUF007",   # zip is idiomatic, this is a dumb check
     "RET505",   # Else after return, makes a lot of false positives
     "E501",     # Line too long, this is autoformatted
+    "PYI041",   # Use float instead of int | float; dumb rule
+    "T201",     # print is fine for now.
+    "RET504",   # Unnecessary assignment before return
+    "PLR0913",  # Too many arguments in function call, hard with CLIs.
+    "TRY201",   #
+    "PD010",    # I like stack and unstack
+    "FBT001",   # Boolean positional args are super common in clis
+    "FBT002",   # Boolean positional args are super common in clis
 ]
 
 [tool.ruff.lint.per-file-ignores]
@@ -142,6 +155,8 @@ exclude = [
  [[tool.mypy.overrides]]
  module = [
      "cdsapi.*",
+     "affine.*",
+     "gcsfs.*",
  ]
  ignore_missing_imports = true
 

src/climate_downscale/cli.py

@@ -1,6 +1,6 @@
 import click
 
-from climate_downscale import extract
+from climate_downscale import downscale, extract, generate
 
 
 @click.group()
@@ -13,7 +13,7 @@ def cdtask() -> None:
     """Entry point for running climate downscale tasks."""
 
 
-for module in [extract]:
+for module in [extract, downscale, generate]:
     runners = getattr(module, "RUNNERS", {})
     task_runners = getattr(module, "TASK_RUNNERS", {})
 

src/climate_downscale/cli_options.py

@@ -0,0 +1,259 @@
+from typing import ParamSpec, TypeVar
+
+import click
+from rra_tools.cli_tools import (
+    RUN_ALL,
+    ClickOption,
+    with_choice,
+    with_debugger,
+    with_input_directory,
+    with_num_cores,
+    with_output_directory,
+    with_progress_bar,
+    with_queue,
+    with_verbose,
+)
+
+_T = TypeVar("_T")
+_P = ParamSpec("_P")
+
+
+VALID_HISTORY_YEARS = [str(y) for y in range(1990, 2024)]
+VALID_REFERENCE_YEARS = VALID_HISTORY_YEARS[-5:]
+VALID_FORECAST_YEARS = [str(y) for y in range(2024, 2101)]
+
+
+def with_year(
+    *,
+    years: list[str],
+    allow_all: bool = False,
+) -> ClickOption[_P, _T]:
+    return with_choice(
+        "year",
+        "y",
+        allow_all=allow_all,
+        choices=years,
+        help="Year to extract data for.",
+    )
+
+
+VALID_MONTHS = [f"{i:02d}" for i in range(1, 13)]
+
+
+def with_month(
+    *,
+    allow_all: bool = False,
+) -> ClickOption[_P, _T]:
+    return with_choice(
+        "month",
+        "m",
+        allow_all=allow_all,
+        choices=VALID_MONTHS,
+        help="Month to extract data for.",
+    )
+
+
+VALID_ERA5_VARIABLES = [
+    "10m_u_component_of_wind",
+    "10m_v_component_of_wind",
+    "2m_dewpoint_temperature",
+    "2m_temperature",
+    "surface_net_solar_radiation",
+    "surface_net_thermal_radiation",
+    "surface_pressure",
+    "surface_solar_radiation_downwards",
+    "surface_thermal_radiation_downwards",
+    "total_precipitation",
+    "total_sky_direct_solar_radiation_at_surface",
+]
+
+
+def with_era5_variable(
+    *,
+    allow_all: bool = False,
+) -> ClickOption[_P, _T]:
+    return with_choice(
+        "era5-variable",
+        "x",
+        allow_all=allow_all,
+        choices=VALID_ERA5_VARIABLES,
+        help="Variable to extract.",
+    )
+
+
+VALID_ERA5_DATASETS = ["reanalysis-era5-land", "reanalysis-era5-single-levels"]
+
+
+def with_era5_dataset(
+    *,
+    allow_all: bool = False,
+) -> ClickOption[_P, _T]:
+    return with_choice(
+        "era5-dataset",
+        "d",
+        allow_all=allow_all,
+        choices=VALID_ERA5_DATASETS,
+        help="Dataset to extract.",
+    )
+
+
+VALID_CMIP6_SOURCES = [
+    "CAMS-CSM1-0",
+    "CanESM5",
+    "CNRM-ESM2-1",
+    "GFDL-ESM4",
+    "GISS-E2-1-G",
+    "MIROC-ES2L",
+    "MIROC6",
+    "MRI-ESM2-0",
+]
+
+
+def with_cmip6_source(
+    *,
+    allow_all: bool = False,
+) -> ClickOption[_P, _T]:
+    return with_choice(
+        "cmip6-source",
+        "s",
+        allow_all=allow_all,
+        choices=VALID_CMIP6_SOURCES,
+        help="CMIP6 source to extract.",
+    )
+
+
+VALID_CMIP6_EXPERIMENTS = [
+    "ssp119",
+    "ssp126",
+    "ssp245",
+    "ssp370",
+    "ssp585",
+]
+
+
+def with_cmip6_experiment(
+    *,
+    allow_all: bool = False,
+    allow_historical: bool = False,
+) -> ClickOption[_P, _T]:
+    choices = VALID_CMIP6_EXPERIMENTS[:]
+    if allow_historical:
+        choices.append("historical")
+    return with_choice(
+        "cmip6-experiment",
+        "e",
+        allow_all=allow_all,
+        choices=choices,
+        help="CMIP6 experiment to extract.",
+    )
+
+
+VALID_CMIP6_VARIABLES = [
+    "uas",
+    "vas",
+    "hurs",
+    "tas",
+    "tasmin",
+    "tasmax",
+    "pr",
+]
+
+
+def with_cmip6_variable(
+    *,
+    allow_all: bool = False,
+) -> ClickOption[_P, _T]:
+    return with_choice(
+        "cmip6-variable",
+        "x",
+        allow_all=allow_all,
+        choices=VALID_CMIP6_VARIABLES,
+        help="CMIP6 variable to extract.",
+    )
+
+
+def with_target_variable(
+    *,
+    variable_names: list[str],
+    allow_all: bool = False,
+) -> ClickOption[_P, _T]:
+    return with_choice(
+        "target-variable",
+        "t",
+        allow_all=allow_all,
+        choices=variable_names,
+        help="Variable to generate.",
+    )
+
+
+STRIDE = 30
+LATITUDES = [str(lat) for lat in range(-90, 90, STRIDE)]
+LONGITUDES = [str(lon) for lon in range(-180, 180, STRIDE)]
+
+
+def with_lat_start(
+    *,
+    allow_all: bool = False,
+) -> ClickOption[_P, _T]:
+    return with_choice(
+        "lat-start",
+        allow_all=allow_all,
+        choices=LATITUDES,
+        help="Latitude of the top-left corner of the tile.",
+    )
+
+
+def with_lon_start(
+    *,
+    allow_all: bool = False,
+) -> ClickOption[_P, _T]:
+    return with_choice(
+        "lon-start",
+        allow_all=allow_all,
+        choices=LONGITUDES,
+        help="Longitude of the top-left corner of the tile.",
+    )
+
+
+def with_overwrite() -> ClickOption[_P, _T]:
+    return click.option(
+        "--overwrite",
+        is_flag=True,
+        help="Overwrite existing files.",
+    )
+
+
+__all__ = [
+    "VALID_HISTORY_YEARS",
+    "VALID_REFERENCE_YEARS",
+    "VALID_FORECAST_YEARS",
+    "VALID_MONTHS",
+    "VALID_ERA5_VARIABLES",
+    "VALID_ERA5_DATASETS",
+    "VALID_CMIP6_SOURCES",
+    "VALID_CMIP6_EXPERIMENTS",
+    "VALID_CMIP6_VARIABLES",
+    "STRIDE",
+    "LATITUDES",
+    "LONGITUDES",
+    "with_year",
+    "with_month",
+    "with_era5_variable",
+    "with_era5_dataset",
+    "with_cmip6_source",
+    "with_cmip6_experiment",
+    "with_cmip6_variable",
+    "with_lat_start",
+    "with_lon_start",
+    "with_output_directory",
+    "with_queue",
+    "with_verbose",
+    "with_debugger",
+    "with_input_directory",
+    "with_num_cores",
+    "with_progress_bar",
+    "RUN_ALL",
+    "ClickOption",
+    "with_choice",
+    "with_overwrite",
+]