🌧️ Downloading ECMWF S2S Ensemble Precipitation Forecasts¶

#!/usr/bin/env python
"""
Download ECMWF S2S realtime **daily total precipitation (24h accumulations)**
for multiple ensemble members and compute the ensemble mean.

Example:
python download_ecmwf_s2s_tp_ensemble_dailymean.py \
  --outdir data/s2s_ecmwf \
  --outfile s2s_ecmwf_ensmean_tp_2025-11-03_ea.nc \
  --date 2025-11-03 \
  --lead-days 14 \
  --members 10 \
  --lat-min 3 --lat-max 15 --lon-min 33 --lon-max 48
"""

import argparse
import os
from datetime import datetime
from typing import List

import xarray as xr
from ecmwfapi import ECMWFDataServer


# --------------------------------------------------------------------------- #
# Helpers
# --------------------------------------------------------------------------- #

def build_daily_step_string(lead_days: int) -> str:
    """
    Build a MARS step string for 24-hour lead times.

    For lead_days = 5 → "24/48/72/96/120"
    (i.e. end of each 24-hour accumulation period).
    """
    if lead_days < 1:
        raise ValueError("lead_days must be >= 1")

    steps = [str(24 * (i + 1)) for i in range(lead_days)]
    return "/".join(steps)


def retrieve_s2s_tp_daily_member(
    date_str: str,
    lead_days: int,
    lat_min: float,
    lat_max: float,
    lon_min: float,
    lon_max: float,
    member: int,
    out_path: str,
    grid: str | None = None,
    fmt: str = "netcdf",
) -> bool:
    """
    Submit an ECMWF S2S request for **daily total precipitation** for ONE
    ensemble member (perturbed forecast, type=pf, number=member).

    Returns True if request appears to succeed, False if ECMWF returns an error
    (e.g. "data not found" or empty GRIB/NetCDF).
    """
    if lead_days < 1 or lead_days > 46:
        raise ValueError("lead_days must be between 1 and 46 for ECMWF S2S.")

    server = ECMWFDataServer()
    steps = build_daily_step_string(lead_days)

    # ECMWF area string is N/W/S/E
    area = f"{lat_max}/{lon_min}/{lat_min}/{lon_max}"

    request = {
        "class": "s2",
        "dataset": "s2s",
        "expver": "prod",
        "origin": "ecmf",
        "model": "glob",
        "levtype": "sfc",
        "stream": "enfo",
        "type": "pf",          # perturbed forecast (ensemble members)
        "number": str(member),
        "param": "tp",         # daily total precip (paramId 228228 internally)
        "date": date_str,      # YYYY-MM-DD, forecast start date
        "time": "00:00:00",
        "step": steps,         # e.g. "24/48/72/..."
        "area": area,
        "format": fmt,
        "target": out_path,
        # Avoid failing with "Expected N, got M" if some steps are missing
        "expect": "any",
    }

    if grid is not None:
        request["grid"] = grid  # e.g. "1.5/1.5" or "0.5/0.5"

    print(f"[info] Submitting S2S daily TP request to ECMWF for member={member}...")
    print("[info] Request:", request)

    try:
        server.retrieve(request)
    except Exception as exc:
        print(f"[warn] ECMWF request failed for member={member}: {exc}")
        return False

    # Quick file sanity check (sometimes an empty file is created)
    if (not os.path.exists(out_path)) or (os.path.getsize(out_path) < 500):
        print(f"[warn] Output for member={member} looks empty or too small: {out_path}")
        return False

    print(f"[info] Download finished for member={member}: {out_path}")
    return True


def compute_ensemble_mean_tp(
    member_files: List[str],
    out_path: str,
) -> None:
    """
    Read all member NetCDFs, stack along 'member', optionally convert TP to
    mm/day (if still in m or kg m**-2), and write ensemble mean to out_path.
    """
    datasets = []
    member_ids = []

    for m, fpath in enumerate(member_files, start=1):
        if not os.path.exists(fpath):
            print(f"[warn] Member file missing, skipping: {fpath}")
            continue

        try:
            ds = xr.open_dataset(fpath)
        except Exception as exc:
            print(f"[warn] Failed to open {fpath}: {exc}")
            continue

        if "tp" not in ds:
            print(f"[warn] No 'tp' variable in {fpath}, skipping.")
            continue

        da = ds["tp"]

        units = str(da.attrs.get("units", "")).lower()

        # Convert only if still in m or kg m**-2
        if (
            "kg m-2" in units
            or "kg m**-2" in units
            or units.strip() == "m"
            or "m of water" in units
        ):
            da.attrs["units"] = "mm/day"
            print(f"[info] Converted TP from '{units}' to 'mm/day' for member {m}")
        elif "mm" in units:
            # Already in mm or mm/day → leave as is
            print(f"[info] TP already in mm units ('{units}') for member {m}, no scaling.")
        else:
            print(
                f"[warn] Unrecognized TP units '{units}' in {fpath}, "
                "leaving values unchanged."
            )

        # Add member dimension
        da = da.expand_dims({"member": [m]})
        datasets.append(da)
        member_ids.append(m)

    if not datasets:
        raise RuntimeError("No valid member files found to build ensemble mean.")

    tp_all = xr.concat(datasets, dim="member")
    tp_all["member"] = member_ids

    ens_mean = tp_all.mean(dim="member", skipna=True)
    ens_ds = ens_mean.to_dataset(name="tp")

    # Make sure units label is sensible
    if "units" not in ens_ds["tp"].attrs:
        ens_ds["tp"].attrs["units"] = "mm/day"

    ens_ds.to_netcdf(out_path)
    print(f"[info] Ensemble mean written to: {out_path}")
    print("[info] Variable: 'tp' (dims: time, lat, lon)")



# --------------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------------- #

def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description=(
            "Download ECMWF S2S realtime daily total precipitation (24h accum) "
            "for multiple ensemble members and compute ensemble mean."
        )
    )
    p.add_argument("--outdir", required=True, help="Output directory")
    p.add_argument("--outfile", required=True, help="Final ensemble-mean filename (NetCDF)")
    p.add_argument(
        "--date",
        required=True,
        help="Forecast initial date (YYYY-MM-DD) for S2S start (use valid S2S start, e.g. Monday/Thursday).",
    )
    p.add_argument(
        "--lead-days",
        type=int,
        required=True,
        help="Number of lead days (1–46) of daily totals to retrieve",
    )
    p.add_argument(
        "--members",
        type=int,
        required=True,
        help="Number of ensemble members (perturbed forecasts) to request (e.g. 10, 20, 50).",
    )
    p.add_argument("--lat-min", type=float, required=True)
    p.add_argument("--lat-max", type=float, required=True)
    p.add_argument("--lon-min", type=float, required=True)
    p.add_argument("--lon-max", type=float, required=True)
    p.add_argument(
        "--grid",
        default=None,
        help="Optional output grid resolution 'lat/lon', e.g. '1.5/1.5' or '0.5/0.5'",
    )
    p.add_argument(
        "--fmt",
        default="netcdf",
        choices=["netcdf", "grib"],
        help="Output format for member files (default: netcdf).",
    )
    return p.parse_args()


def main() -> None:
    args = parse_args()

    # Ensure output directory exists
    os.makedirs(args.outdir, exist_ok=True)
    ens_out_path = os.path.join(args.outdir, args.outfile)

    # Basic date sanity check
    try:
        datetime.strptime(args.date, "%Y-%m-%d")
    except ValueError as exc:
        raise SystemExit(f"Invalid --date '{args.date}', expected YYYY-MM-DD") from exc

    # Base path for member files (same prefix as final ensemble file)
    base_prefix = os.path.splitext(ens_out_path)[0]

    member_files: List[str] = []

    for m in range(1, args.members + 1):
        member_path = f"{base_prefix}_member{m:02d}.nc"
        ok = retrieve_s2s_tp_daily_member(
            date_str=args.date,
            lead_days=args.lead_days,
            lat_min=args.lat_min,
            lat_max=args.lat_max,
            lon_min=args.lon_min,
            lon_max=args.lon_max,
            member=m,
            out_path=member_path,
            grid=args.grid,
            fmt=args.fmt,
        )
        if ok:
            member_files.append(member_path)
        else:
            print(f"[warn] Skipping member={m} due to retrieval/conversion issues.")

    if not member_files:
        raise SystemExit("No member files were successfully downloaded. Aborting.")

    # Compute ensemble mean from the successfully downloaded members
    compute_ensemble_mean_tp(member_files, ens_out_path)


if __name__ == "__main__":
    main()

import xarray as xr
import matplotlib.pyplot as plt
import numpy as np

# Open the ensemble mean file
ds = xr.open_dataset('data/s2s_ensemble/s2s_ensmean_tp_ethiopia.nc')

# Display dataset information
print(ds)

# Check dimensions
print(f"Lead times: {len(ds.time) if 'time' in ds.dims else 'N/A'}")
print(f"Latitude range: {float(ds.latitude.min()):.2f} to {float(ds.latitude.max()):.2f}")
print(f"Longitude range: {float(ds.longitude.min()):.2f} to {float(ds.longitude.max()):.2f}")

# Check precipitation range
print(f"Precipitation range: {float(ds.tp.min()):.1f} to {float(ds.tp.max()):.1f} mm/day")

# Plot Week 1 ensemble mean precipitation
fig, ax = plt.subplots(figsize=(10, 8))
week1_mean = ds.tp.isel(time=slice(0, 7)).mean(dim='time')
week1_mean.plot(ax=ax, cmap='Blues', vmin=0, vmax=20)
ax.set_title('S2S Ensemble Mean: Week 1 Daily Precipitation')
plt.savefig('s2s_ensemble_week1.png', dpi=150, bbox_inches='tight')
plt.show()

# Time series for a point
lat_point, lon_point = 9.0, 38.7  # Addis Ababa
point_data = ds.tp.sel(latitude=lat_point, longitude=lon_point, method='nearest')
point_data.plot(marker='o', figsize=(12, 4), color='steelblue')
plt.title(f'S2S Ensemble Mean Precipitation for Addis Ababa')
plt.ylabel('Precipitation (mm/day)')
plt.xlabel('Lead Time')
plt.grid(True, alpha=0.3)
plt.savefig('s2s_ensemble_timeseries.png', dpi=150, bbox_inches='tight')
plt.show()

import xarray as xr
import numpy as np
import glob

# Load all member files
member_files = sorted(glob.glob('data/s2s_ensemble/*_member*.nc'))
print(f"Found {len(member_files)} member files")

# Stack all members
datasets = []
for i, f in enumerate(member_files, 1):
    ds = xr.open_dataset(f)
    ds = ds.expand_dims({'member': [i]})
    datasets.append(ds)

# Combine all members
ds_all = xr.concat(datasets, dim='member')
print(ds_all)

# Compute ensemble statistics
ens_mean = ds_all.tp.mean(dim='member')
ens_std = ds_all.tp.std(dim='member')
ens_median = ds_all.tp.median(dim='member')
ens_p10 = ds_all.tp.quantile(0.1, dim='member')
ens_p25 = ds_all.tp.quantile(0.25, dim='member')
ens_p75 = ds_all.tp.quantile(0.75, dim='member')
ens_p90 = ds_all.tp.quantile(0.9, dim='member')

# Compute probability of exceeding threshold
threshold = 10  # mm/day
prob_exceed = (ds_all.tp > threshold).mean(dim='member') * 100  # percentage

# Create output dataset
ds_out = xr.Dataset({
    'tp_mean': ens_mean,
    'tp_std': ens_std,
    'tp_median': ens_median,
    'tp_p10': ens_p10,
    'tp_p25': ens_p25,
    'tp_p75': ens_p75,
    'tp_p90': ens_p90,
    'prob_exceed_10mm': prob_exceed,
})

# Add attributes
ds_out.tp_mean.attrs = {'units': 'mm/day', 'long_name': 'Ensemble mean precipitation'}
ds_out.tp_std.attrs = {'units': 'mm/day', 'long_name': 'Ensemble standard deviation'}
ds_out.prob_exceed_10mm.attrs = {'units': '%', 'long_name': 'Probability of exceeding 10 mm/day'}

# Save
ds_out.to_netcdf('data/s2s_ensemble/s2s_probabilistic_tp.nc')
print("Saved probabilistic products")

# Visualize probability of exceeding threshold
import matplotlib.pyplot as plt

fig, axes = plt.subplots(1, 4, figsize=(16, 4))

# Week 1-4 probability maps
for i, week in enumerate([0, 7, 14, 21]):
    if week + 7 <= len(prob_exceed.time):
        weekly_prob = prob_exceed.isel(time=slice(week, week+7)).mean(dim='time')
        im = weekly_prob.plot(ax=axes[i], cmap='YlOrRd', vmin=0, vmax=100, add_colorbar=False)
        axes[i].set_title(f'Week {i+1}')
        axes[i].set_xlabel('')
        axes[i].set_ylabel('')

plt.suptitle('Probability of Precipitation > 10 mm/day (%)')
plt.tight_layout()
plt.savefig('s2s_probability_map.png', dpi=150, bbox_inches='tight')
plt.show()