wiend 0.1.0__py3-none-any.whl

wiend/__init__.py

@@ -0,0 +1,11 @@
+__version__ = "0.1.0"
+
+from wiend.download import download, download_and_downscale
+from wiend.downscale import downscale
+
+__all__ = [
+    "__version__",
+    "download",
+    "downscale",
+    "download_and_downscale",
+]

wiend/data_manager.py

@@ -0,0 +1,267 @@
+"""Data access and cache helpers for wiend."""
+
+from __future__ import annotations
+
+import json
+import os
+import shutil
+import urllib.error
+import urllib.parse
+import urllib.request
+from pathlib import Path
+from typing import Any
+
+import xarray as xr
+
+from wiend.validation import (
+    check_dims_and_coords,
+    check_lat_lon_1d_ascending,
+)
+
+EDH_DEFAULT_PATH = "era5/reanalysis-era5-single-levels-v0.zarr"
+DEFAULT_GITHUB_REPOSITORY = "florianscheiber/wiend"
+DEFAULT_WEIBULL_RELEASE_TAG = "latest"
+
+
+def get_cache_dir(cache_dir: str | Path | None = None) -> Path:
+    """Return the local cache directory, defaulting to ~/.wiend."""
+    if cache_dir is not None:
+        return Path(cache_dir).expanduser().resolve()
+    return (Path.home() / ".wiend").resolve()
+
+
+def load_edh_key(edh_key: str | None = None) -> str:
+    """Load EDH token from argument or environment."""
+    if edh_key is not None:
+        return edh_key
+
+    try:
+        from dotenv import load_dotenv
+    except ImportError:
+        load_dotenv = None
+
+    if load_dotenv is not None:
+        load_dotenv()
+
+    key = os.getenv("EDH_TOKEN") or os.getenv("EDH_KEY")
+    if not key:
+        raise ValueError("No EDH token found. Pass edh_key or set EDH_TOKEN/EDH_KEY.")
+    return key
+
+
+def build_edh_url(edh_path: str, edh_key: str) -> str:
+    """Build full EDH URL from logical path and token."""
+    if edh_path.startswith("http://") or edh_path.startswith("https://"):
+        return edh_path
+    normalized_path = edh_path.lstrip("/")
+    return f"https://edh:{edh_key}@api.earthdatahub.destine.eu/{normalized_path}"
+
+
+def load_wind_from_zarr(path_input: str | Path) -> xr.DataArray:
+    """Load wind data from Zarr and validate expected structure."""
+    ds = xr.open_zarr(str(path_input))
+
+    if "wind" not in ds:
+        raise KeyError(f"Zarr missing variable 'wind'. Found: {list(ds.data_vars)}")
+
+    wind = ds["wind"].astype("float32")
+    required_dims = ("time", "lat", "lon")
+    required_coords = ("time", "lat", "lon")
+    check_dims_and_coords(wind, required_dims, required_coords)
+    check_lat_lon_1d_ascending(wind["lat"], wind["lon"])
+    return wind
+
+
+def load_weibull_params_from_zarr(
+    path_input: str | Path,
+    *,
+    bbox: tuple[float, float, float, float] | None = None,
+    crop: bool = True,
+) -> tuple[xr.DataArray, xr.DataArray]:
+    """Load Weibull parameters k and a with optional bbox crop."""
+    ds = xr.open_zarr(str(path_input))
+
+    if crop:
+        if bbox is None:
+            raise ValueError("bbox must be provided when crop=True.")
+        lat_min, lat_max, lon_min, lon_max = bbox
+        ds = ds.sel(lat=slice(lat_min, lat_max), lon=slice(lon_min, lon_max))
+
+    if "k" not in ds or "a" not in ds:
+        raise KeyError(f"Params Zarr {path_input} missing 'k' or 'a'.")
+
+    k = ds["k"]
+    a = ds["a"]
+
+    if k.sizes["lat"] == 0 or k.sizes["lon"] == 0:
+        raise ValueError(f"Cropping produced an empty grid for {path_input} and bbox={bbox}.")
+
+    required_dims = ("lat", "lon")
+    required_coords = ("lat", "lon")
+    check_dims_and_coords(k, required_dims, required_coords)
+    check_dims_and_coords(a, required_dims, required_coords)
+    check_lat_lon_1d_ascending(k["lat"], k["lon"])
+    check_lat_lon_1d_ascending(a["lat"], a["lon"])
+    return k, a
+
+
+def _spacing_token(spacing: float | None) -> str | None:
+    if spacing is None:
+        return None
+    return f"spacing-{str(spacing).replace('.', '-')}"
+
+
+def _github_api_release_url(repository: str, release_tag: str) -> str:
+    owner, repo = repository.split("/", 1)
+    owner_q = urllib.parse.quote(owner, safe="")
+    repo_q = urllib.parse.quote(repo, safe="")
+    if release_tag == "latest":
+        return f"https://api.github.com/repos/{owner_q}/{repo_q}/releases/latest"
+    tag_q = urllib.parse.quote(release_tag, safe="")
+    return f"https://api.github.com/repos/{owner_q}/{repo_q}/releases/tags/{tag_q}"
+
+
+def _github_release_payload(
+    repository: str,
+    release_tag: str,
+    github_token: str | None = None,
+) -> dict[str, Any]:
+    url = _github_api_release_url(repository, release_tag)
+    headers = {"Accept": "application/vnd.github+json"}
+    if github_token:
+        headers["Authorization"] = f"Bearer {github_token}"
+    req = urllib.request.Request(url, headers=headers)
+    try:
+        with urllib.request.urlopen(req) as response:
+            body = response.read().decode("utf-8")
+    except urllib.error.HTTPError as exc:
+        raise RuntimeError(f"Failed to fetch GitHub release metadata from {url}: {exc}") from exc
+    return json.loads(body)
+
+
+def _select_asset(
+    assets: list[dict[str, Any]],
+    *,
+    role_key: str,
+    spacing: float | None,
+) -> dict[str, Any]:
+    matches: list[dict[str, Any]] = []
+    for asset in assets:
+        name = str(asset.get("name", "")).lower()
+        if role_key not in name:
+            continue
+        matches.append(asset)
+
+    if not matches:
+        raise RuntimeError(f"No release asset found for role={role_key!r}.")
+    if len(matches) > 1:
+        names = ", ".join(str(a.get("name", "<unknown>")) for a in matches)
+        raise RuntimeError(f"Multiple release assets matched role={role_key!r}: {names}")
+    return matches[0]
+
+
+def _download_file(url: str, target: Path, github_token: str | None = None) -> None:
+    target.parent.mkdir(parents=True, exist_ok=True)
+    headers: dict[str, str] = {}
+    if github_token:
+        headers["Authorization"] = f"Bearer {github_token}"
+
+    req = urllib.request.Request(url, headers=headers)
+    try:
+        with urllib.request.urlopen(req) as response, target.open("wb") as file_obj:
+            shutil.copyfileobj(response, file_obj)
+    except urllib.error.HTTPError as exc:
+        raise RuntimeError(f"Failed to download asset from {url}: {exc}") from exc
+
+
+def _materialize_release_asset(
+    asset: dict[str, Any],
+    *,
+    cache_dir: Path,
+    github_token: str | None = None,
+) -> Path:
+    name = str(asset["name"])
+    asset_dir = cache_dir / "assets"
+    archive_path = asset_dir / name
+    extract_dir = cache_dir / "datasets" / Path(name).stem
+
+    if extract_dir.exists():
+        return _resolve_dataset_path(extract_dir)
+
+    if not archive_path.exists():
+        download_url = str(asset["browser_download_url"])
+        _download_file(download_url, archive_path, github_token=github_token)
+
+    suffix = archive_path.suffix.lower()
+    if suffix == ".zip":
+        shutil.unpack_archive(str(archive_path), str(extract_dir))
+        return _resolve_dataset_path(extract_dir)
+
+    # Fallback: non-archive asset, expose directly.
+    return archive_path
+
+
+def _resolve_dataset_path(extract_dir: Path) -> Path:
+    """
+    Resolve the actual dataset root from an extracted release asset directory.
+    """
+    zarr_dirs = list(extract_dir.rglob("*.zarr"))
+    if len(zarr_dirs) == 1:
+        return zarr_dirs[0]
+    if len(zarr_dirs) > 1:
+        raise RuntimeError(
+            f"Ambiguous extracted asset {extract_dir}: multiple .zarr directories found."
+        )
+
+    children = [p for p in extract_dir.iterdir() if p.name != "__MACOSX"]
+    if len(children) == 1:
+        return children[0]
+    return extract_dir
+
+
+def resolve_weibull_reference_paths(
+    *,
+    spacing: float | None,
+    cache_dir: str | Path | None = None,
+    github_repository: str = DEFAULT_GITHUB_REPOSITORY,
+    release_tag: str = DEFAULT_WEIBULL_RELEASE_TAG,
+    github_token: str | None = None,
+) -> tuple[Path, Path]:
+    """
+    Resolve baseline/manipulator Weibull parameter paths from GitHub release assets.
+    """
+    cache_path = get_cache_dir(cache_dir)
+    cache_path.mkdir(parents=True, exist_ok=True)
+
+    payload = _github_release_payload(
+        github_repository,
+        release_tag,
+        github_token=github_token,
+    )
+    assets = payload.get("assets")
+    if not isinstance(assets, list):
+        raise RuntimeError("GitHub release payload does not include an assets list.")
+
+    baseline_asset = _select_asset(
+        assets,
+        role_key="baseline",
+        spacing=spacing,
+    )
+    manipulator_asset = _select_asset(
+        assets,
+        role_key="corrected",
+        spacing=spacing,
+    )
+
+    baseline_path = _materialize_release_asset(
+        baseline_asset,
+        cache_dir=cache_path,
+        github_token=github_token,
+    )
+    manipulator_path = _materialize_release_asset(
+        manipulator_asset,
+        cache_dir=cache_path,
+        github_token=github_token,
+    )
+
+    return baseline_path, manipulator_path

wiend/download.py

@@ -0,0 +1,242 @@
+"""Download API for wiend."""
+
+from __future__ import annotations
+
+import json
+import shutil
+from pathlib import Path
+
+import xarray as xr
+
+from wiend.data_manager import (
+    DEFAULT_GITHUB_REPOSITORY,
+    DEFAULT_WEIBULL_RELEASE_TAG,
+    EDH_DEFAULT_PATH,
+    build_edh_url,
+    load_edh_key,
+)
+from wiend.downscale import downscale
+from wiend.validation import normalize_wind_da, parse_bbox
+
+
+def _get_wind_da(
+    *,
+    bbox: tuple[float, float, float, float],
+    start: str | None = None,
+    end: str | None = None,
+    edh_key: str | None = None,
+    edh_path: str | None = None,
+    height: str = "10m",
+    u_var: str | None = None,
+    v_var: str | None = None,
+    bbox_order: str = "lonlat",
+    time_dim: str = "time",
+    chunks: dict[str, int] | None = None,
+) -> xr.DataArray:
+    key = load_edh_key(edh_key)
+    path = edh_path or EDH_DEFAULT_PATH
+    url = build_edh_url(path, key)
+
+    ds_src = xr.open_dataset(
+        url,
+        engine="zarr",
+        chunks=chunks or {},
+        decode_coords="all",
+        mask_and_scale=False,
+    )
+
+    if u_var is None:
+        preferred_u = ["u10", "u10m", "u"] if height == "10m" else ["u100", "u100m", "u"]
+        for name in preferred_u:
+            if name in ds_src.data_vars:
+                u_var = name
+                break
+
+    if v_var is None:
+        preferred_v = ["v10", "v10m", "v"] if height == "10m" else ["v100", "v100m", "v"]
+        for name in preferred_v:
+            if name in ds_src.data_vars:
+                v_var = name
+                break
+
+    if u_var is None or v_var is None:
+        raise ValueError(f"Could not infer u/v variables. Available: {list(ds_src.data_vars)}")
+
+    u_da = normalize_wind_da(ds_src[u_var], time_dim=time_dim)
+    v_da = normalize_wind_da(ds_src[v_var], time_dim=time_dim)
+
+    lon_min, lat_min, lon_max, lat_max = parse_bbox(bbox, bbox_order=bbox_order)
+    u_da = u_da.sel(lat=slice(lat_min, lat_max), lon=slice(lon_min, lon_max))
+    v_da = v_da.sel(lat=slice(lat_min, lat_max), lon=slice(lon_min, lon_max))
+
+    if start is not None or end is not None:
+        u_da = u_da.sel(time=slice(start, end))
+        v_da = v_da.sel(time=slice(start, end))
+
+    u_da, v_da = xr.align(u_da, v_da, join="inner")
+    wind_da = (u_da**2 + v_da**2) ** 0.5
+    return wind_da.rename("wind").astype("float32")
+
+
+def download(
+    *,
+    bbox: tuple[float, float, float, float],
+    start: str,
+    end: str,
+    out: str | Path | None = None,
+    edh_key: str | None = None,
+    edh_path: str | None = None,
+    height: str = "10m",
+    u_var: str | None = None,
+    v_var: str | None = None,
+    bbox_order: str = "lonlat",
+    time_dim: str = "time",
+    overwrite: bool = False,
+    edh_time_chunk: int = 168,
+    tile_lat: int = 50,
+    tile_lon: int = 50,
+    min_valid: int = 100,
+    compute: bool = True,
+    write_summary: bool = False,
+) -> xr.DataArray:
+    """Download ERA5 wind data and optionally write it to Zarr."""
+    wind_da = _get_wind_da(
+        bbox=bbox,
+        start=start,
+        end=end,
+        edh_key=edh_key,
+        edh_path=edh_path,
+        height=height,
+        u_var=u_var,
+        v_var=v_var,
+        bbox_order=bbox_order,
+        time_dim=time_dim,
+        chunks={"time": int(edh_time_chunk)},
+    )
+
+    if compute:
+        wind_da = wind_da.compute()
+
+    if out is None:
+        return wind_da
+
+    out_path = Path(out)
+    if out_path.exists():
+        if not overwrite:
+            raise ValueError(f"Output {out_path} exists; set overwrite=True to replace.")
+        if out_path.is_dir():
+            shutil.rmtree(out_path)
+        else:
+            out_path.unlink()
+
+    ds_out = wind_da.to_dataset()
+    ds_out = ds_out.assign_coords(
+        lat=ds_out["lat"].astype("float32"), lon=ds_out["lon"].astype("float32")
+    )
+    ds_out.attrs.update(
+        {
+            "start": start,
+            "end": end,
+            "bbox": list(map(float, bbox)),
+            "height": height,
+            "edh_path": edh_path or EDH_DEFAULT_PATH,
+            "time_chunk": int(edh_time_chunk),
+            "min_valid": int(min_valid),
+            "source": "era5_edh",
+            "tile_lat": int(tile_lat),
+            "tile_lon": int(tile_lon),
+        }
+    )
+
+    chunks_time = min(int(edh_time_chunk), int(ds_out.sizes.get("time", edh_time_chunk)))
+    chunks_lat = min(int(tile_lat), int(ds_out.sizes.get("lat", tile_lat)))
+    chunks_lon = min(int(tile_lon), int(ds_out.sizes.get("lon", tile_lon)))
+    encoding = {
+        "wind": {"dtype": "float32", "chunks": (chunks_time, chunks_lat, chunks_lon)},
+        "lat": {"dtype": "float32"},
+        "lon": {"dtype": "float32"},
+    }
+
+    for var_name in ds_out.data_vars:
+        ds_out[var_name].encoding = {}
+    for coord_name in ds_out.coords:
+        ds_out.coords[coord_name].encoding = {}
+    ds_out.encoding = {}
+
+    ds_out.to_zarr(str(out_path), mode="w", compute=True, encoding=encoding)
+
+    if write_summary:
+        summary = {
+            "bbox": list(map(float, bbox)),
+            "start": start,
+            "end": end,
+            "out": str(out_path),
+            "tile_lat": int(tile_lat),
+            "tile_lon": int(tile_lon),
+            "min_valid": int(min_valid),
+            "height": height,
+            "written_vars": list(ds_out.data_vars.keys()),
+        }
+        summary_path = Path(str(out_path) + ".summary.json")
+        summary_path.write_text(json.dumps(summary, indent=2), encoding="utf-8")
+
+    return wind_da
+
+
+def download_and_downscale(
+    *,
+    bbox: tuple[float, float, float, float],
+    start: str,
+    end: str,
+    out: str | Path,
+    mapping_method: str,
+    params_baseline_path: str | Path | None = None,
+    params_manipulator_path: str | Path | None = None,
+    spacing: float | None = None,
+    cache_dir: str | Path | None = None,
+    github_repository: str = DEFAULT_GITHUB_REPOSITORY,
+    release_tag: str = DEFAULT_WEIBULL_RELEASE_TAG,
+    github_token: str | None = None,
+    edh_key: str | None = None,
+    edh_path: str | None = None,
+    height: str = "10m",
+    u_var: str | None = None,
+    v_var: str | None = None,
+    bbox_order: str = "lonlat",
+    time_dim: str = "time",
+    overwrite: bool = False,
+    no_crop: bool = False,
+    edh_time_chunk: int = 168,
+) -> Path:
+    """Download ERA5 wind and run downscaling in one call."""
+    wind_da = download(
+        bbox=bbox,
+        start=start,
+        end=end,
+        out=None,
+        edh_key=edh_key,
+        edh_path=edh_path,
+        height=height,
+        u_var=u_var,
+        v_var=v_var,
+        bbox_order=bbox_order,
+        time_dim=time_dim,
+        edh_time_chunk=edh_time_chunk,
+        compute=True,
+    )
+    return downscale(
+        input_data=wind_da,
+        out=out,
+        bbox=bbox,
+        mapping_method=mapping_method,
+        params_baseline_path=params_baseline_path,
+        params_manipulator_path=params_manipulator_path,
+        spacing=spacing,
+        cache_dir=cache_dir,
+        github_repository=github_repository,
+        release_tag=release_tag,
+        github_token=github_token,
+        bbox_order=bbox_order,
+        no_crop=no_crop,
+        overwrite=overwrite,
+    )