🌧️ Downloading CHC-CMIP6 Daily Precipitation¶

#!/usr/bin/env python
"""
Download CHC-CMIP6 daily precipitation GeoTIFFs (CHIRPS-based),
clip to Ethiopia, convert to NetCDF, and fix units/variable names.

Example:
python download_chc_cmip6_precip_daily_to_netcdf.py \
  --period-tags 2030_SSP245 2030_SSP585 \
  --start-year 1983 --end-year 1984 \
  --outdir data/CHC_CMIP6 \
  --lat-min 3 --lat-max 15 --lon-min 33 --lon-max 48
"""

import argparse
import calendar
import os
from typing import Iterable, Tuple

import numpy as np
import requests
import rioxarray
import xarray as xr

BASE_URL = "https://data.chc.ucsb.edu/products/CHC_CMIP6"


# -----------------------------------------------------------------------------#
# Utilities
# -----------------------------------------------------------------------------#

def iter_dates(year: int) -> Iterable[Tuple[int, int, int]]:
    """Yield (year, month, day) for every day in a given year."""
    for month in range(1, 13):
        ndays = calendar.monthrange(year, month)[1]
        for day in range(1, ndays + 1):
            yield year, month, day


def download_file(url: str, dest_path: str, verbose: bool = False) -> bool:
    """
    Download file from URL to dest_path. 

    Returns True on success, False otherwise.
    """
    try:
        r = requests.get(url, stream=True, timeout=60)
        if r.status_code == 200:
            with open(dest_path, "wb") as f:
                for chunk in r.iter_content(1024 * 1024):
                    f.write(chunk)
            if verbose:
                print(f"[ok] Downloaded {os.path.basename(dest_path)}")
            return True
        else:
            if verbose:
                print(f"[warn] HTTP {r.status_code} for {url}")
            return False
    except Exception as exc:
        print(f"[err] Failed to download {url}: {exc}")
        return False


def subset_region(
    da: xr.DataArray,
    lat_min: float,
    lat_max: float,
    lon_min: float,
    lon_max: float,
) -> xr.DataArray:
    """
    Subset DataArray to a bounding box.

    Handles CHIRPS coordinate conventions (y descending).
    """
    # CHIRPS is in EPSG:4326
    da = da.rio.write_crs("EPSG:4326", inplace=True)

    # CHIRPS y is usually descending (north -> south), so use slice(lat_max, lat_min)
    return da.sel(y=slice(lat_max, lat_min), x=slice(lon_min, lon_max))


def fix_units_and_name(da: xr.DataArray) -> xr.DataArray:
    """Standardise variable to pr [mm/day]."""
    da.name = "pr"
    da.attrs["long_name"] = "Daily precipitation"
    da.attrs["units"] = "mm/day"
    da.attrs["standard_name"] = "precipitation_flux"
    return da


def combine_and_save(daily_arrays: list, out_path: str, compress: bool = True):
    """
    Combine daily DataArrays into a single NetCDF file.

    Parameters
    ----------
    daily_arrays : list
        List of xr.DataArray objects with time dimension
    out_path : str
        Output NetCDF file path
    compress : bool
        Enable zlib compression (default True)
    """
    if not daily_arrays:
        print(f"[warn] No data to save for {out_path}")
        return

    # Concatenate along time dimension
    ds = xr.concat(daily_arrays, dim="time")
    ds = ds.sortby("time")

    # Add metadata
    if isinstance(ds, xr.DataArray):
        ds = ds.to_dataset(name="pr")

    ds.attrs["title"] = "CHC-CMIP6 Daily Precipitation"
    ds.attrs["source"] = "Climate Hazards Center, UC Santa Barbara"
    ds.attrs["institution"] = "CHC-UCSB"
    ds.attrs["references"] = "https://data.chc.ucsb.edu/products/CHC_CMIP6"

    # Encoding for compression
    encoding = {}
    if compress:
        encoding = {"pr": {"zlib": True, "complevel": 4, "dtype": "float32"}}

    # Save
    ds.to_netcdf(out_path, encoding=encoding)
    print(f"[ok] Saved NetCDF → {out_path}")


# -----------------------------------------------------------------------------#
# Core processing
# -----------------------------------------------------------------------------#

def process_period(
    period_tag: str,
    start_year: int,
    end_year: int,
    outdir: str,
    lat_min: float,
    lat_max: float,
    lon_min: float,
    lon_max: float,
    verbose: bool = False,
):
    """
    Process a single CHC-CMIP6 period/scenario.

    Downloads daily GeoTIFFs, clips to region, converts to annual NetCDF.
    """
    dataset_dir = "chirps-v2"
    file_tag = "CHIRPS"

    # Temp directory for GeoTIFFs (inside outdir so it's easy to inspect/clean)
    tmp_root = os.path.join(outdir, "_tmp_chc_cmip6")
    os.makedirs(tmp_root, exist_ok=True)

    for year in range(start_year, end_year + 1):
        print(f"\n{'='*60}")
        print(f"[info] Processing {period_tag} {year}")
        print(f"{'='*60}")

        daily_list = []
        success_count = 0
        fail_count = 0

        year_url = f"{BASE_URL}/{period_tag}/{dataset_dir}/{year}"

        for y, m, d in iter_dates(year):
            fname = f"{period_tag}.{file_tag}.{y}.{m:02d}.{d:02d}.tif"
            url = f"{year_url}/{fname}"
            tmp_path = os.path.join(tmp_root, fname)

            # Download
            if not download_file(url, tmp_path, verbose=verbose):
                fail_count += 1
                continue

            try:
                # Open GeoTIFF
                da = rioxarray.open_rasterio(tmp_path).squeeze(drop=True)

                # Subset + standardise metadata
                da = subset_region(da, lat_min, lat_max, lon_min, lon_max)
                da = fix_units_and_name(da)

                # Load data into memory so we can safely close/delete the file
                da = da.load()

                # Add time dimension
                time_val = np.datetime64(f"{y:04d}-{m:02d}-{d:02d}")
                da = da.expand_dims(time=[time_val])

                daily_list.append(da)
                success_count += 1

                # Explicitly close raster handle
                try:
                    da.rio.close()
                except Exception:
                    pass

            except Exception as exc:
                print(f"[warn] Failed to process {fname}: {exc}")
                fail_count += 1

            finally:
                # Clean up temp file
                try:
                    if os.path.exists(tmp_path):
                        os.remove(tmp_path)
                except Exception as exc_rm:
                    print(f"[warn] Could not remove temp file {tmp_path}: {exc_rm}")

        # Summary for year
        total_days = 366 if calendar.isleap(year) else 365
        print(f"\n[summary] {period_tag} {year}: {success_count}/{total_days} days downloaded")

        if not daily_list:
            print(f"[warn] No valid data for {period_tag} {year}")
            continue

        # Output one NetCDF per year & scenario
        out_path = os.path.join(
            outdir,
            period_tag,
            f"{period_tag}_CHIRPS_{year}_daily.nc",
        )
        os.makedirs(os.path.dirname(out_path), exist_ok=True)
        combine_and_save(daily_list, out_path)


def process_all_periods(
    period_tags: list,
    start_year: int,
    end_year: int,
    outdir: str,
    lat_min: float,
    lat_max: float,
    lon_min: float,
    lon_max: float,
    verbose: bool = False,
):
    """Process multiple CHC-CMIP6 periods/scenarios."""
    print(f"\n{'#'*60}")
    print(f"# CHC-CMIP6 Daily Precipitation Download")
    print(f"# Scenarios: {', '.join(period_tags)}")
    print(f"# Years: {start_year} to {end_year}")
    print(f"# Region: lat [{lat_min}, {lat_max}], lon [{lon_min}, {lon_max}]")
    print(f"{'#'*60}\n")

    for tag in period_tags:
        process_period(
            period_tag=tag,
            start_year=start_year,
            end_year=end_year,
            outdir=outdir,
            lat_min=lat_min,
            lat_max=lat_max,
            lon_min=lon_min,
            lon_max=lon_max,
            verbose=verbose,
        )

    print(f"\n{'#'*60}")
    print(f"# Download complete!")
    print(f"# Output directory: {outdir}")
    print(f"{'#'*60}\n")


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

def parse_args():
    p = argparse.ArgumentParser(
        description=(
            "Download CHC-CMIP6 CHIRPS-based daily precipitation, "
            "subset to region, convert to NetCDF."
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Single scenario, single year (quick test)
  python download_chc_cmip6_precip_daily.py \\
      --period-tags 2030_SSP245 \\
      --start-year 1983 --end-year 1983 \\
      --outdir data/CHC_CMIP6 \\
      --lat-min 3 --lat-max 15 --lon-min 33 --lon-max 48

  # Multiple scenarios, multiple years
  python download_chc_cmip6_precip_daily.py \\
      --period-tags 2030_SSP245 2030_SSP585 2050_SSP245 \\
      --start-year 1983 --end-year 2020 \\
      --outdir data/CHC_CMIP6 \\
      --lat-min 3 --lat-max 15 --lon-min 33 --lon-max 48
        """
    )
    p.add_argument(
        "--period-tags",
        nargs="+",
        required=True,
        help="Period tags: 2030_SSP245 2030_SSP585 2050_SSP245 2050_SSP585 etc.",
    )
    p.add_argument(
        "--start-year", 
        type=int, 
        default=1983,
        help="Start year (default: 1983)"
    )
    p.add_argument(
        "--end-year", 
        type=int, 
        default=1983,
        help="End year (default: 1983)"
    )
    p.add_argument(
        "--outdir", 
        required=True,
        help="Output directory for NetCDF files"
    )
    p.add_argument("--lat-min", type=float, default=3.0, help="Min latitude")
    p.add_argument("--lat-max", type=float, default=15.0, help="Max latitude")
    p.add_argument("--lon-min", type=float, default=33.0, help="Min longitude")
    p.add_argument("--lon-max", type=float, default=48.0, help="Max longitude")
    p.add_argument(
        "--verbose", "-v",
        action="store_true",
        help="Print verbose download messages"
    )
    return p.parse_args()


def main():
    args = parse_args()

    process_all_periods(
        period_tags=args.period_tags,
        start_year=args.start_year,
        end_year=args.end_year,
        outdir=args.outdir,
        lat_min=args.lat_min,
        lat_max=args.lat_max,
        lon_min=args.lon_min,
        lon_max=args.lon_max,
        verbose=args.verbose,
    )


if __name__ == "__main__":
    main()

import xarray as xr
import matplotlib.pyplot as plt

# Open a downloaded file
ds = xr.open_dataset('data/CHC_CMIP6/2030_SSP245/2030_SSP245_CHIRPS_1983_daily.nc')

# Display dataset information
print(ds)
print(f"\nDimensions: {dict(ds.dims)}")
print(f"Variables: {list(ds.data_vars)}")
print(f"Time range: {ds.time.values[0]} to {ds.time.values[-1]}")
print(f"Precipitation range: {float(ds.pr.min()):.2f} to {float(ds.pr.max()):.2f} mm/day")

# Plot annual mean precipitation
annual_mean = ds.pr.mean(dim='time')

fig, ax = plt.subplots(figsize=(10, 8))
annual_mean.plot(ax=ax, cmap='Blues', cbar_kwargs={'label': 'mm/day'})
ax.set_title('CHC-CMIP6 (2030_SSP245) Annual Mean Precipitation 1983')
plt.savefig('chc_cmip6_annual_mean.png', dpi=150, bbox_inches='tight')
plt.show()

# Monthly climatology
monthly_mean = ds.pr.groupby('time.month').mean(dim='time')

fig, axes = plt.subplots(3, 4, figsize=(16, 12))
for i, month in enumerate(range(1, 13)):
    ax = axes.flatten()[i]
    monthly_mean.sel(month=month).plot(ax=ax, cmap='Blues', add_colorbar=False, vmin=0, vmax=10)
    ax.set_title(f'Month {month}')
    ax.set_xlabel('')
    ax.set_ylabel('')
plt.suptitle('CHC-CMIP6 Monthly Mean Precipitation (mm/day)')
plt.tight_layout()
plt.savefig('chc_cmip6_monthly.png', dpi=150, bbox_inches='tight')
plt.show()

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

# Load two scenarios
ds_ssp245 = xr.open_dataset('data/CHC_CMIP6/2050_SSP245/2050_SSP245_CHIRPS_2000_daily.nc')
ds_ssp585 = xr.open_dataset('data/CHC_CMIP6/2050_SSP585/2050_SSP585_CHIRPS_2000_daily.nc')

# Compute annual means
mean_ssp245 = ds_ssp245.pr.mean(dim='time')
mean_ssp585 = ds_ssp585.pr.mean(dim='time')

# Compute difference (SSP585 - SSP245)
diff = mean_ssp585 - mean_ssp245

# Plot comparison
fig, axes = plt.subplots(1, 3, figsize=(18, 5))

# SSP245
mean_ssp245.plot(ax=axes[0], cmap='Blues', vmin=0, vmax=8, 
                  cbar_kwargs={'label': 'mm/day'})
axes[0].set_title('2050_SSP245 Annual Mean')

# SSP585
mean_ssp585.plot(ax=axes[1], cmap='Blues', vmin=0, vmax=8,
                  cbar_kwargs={'label': 'mm/day'})
axes[1].set_title('2050_SSP585 Annual Mean')

# Difference
diff.plot(ax=axes[2], cmap='RdBu', center=0, vmin=-2, vmax=2,
          cbar_kwargs={'label': 'mm/day difference'})
axes[2].set_title('Difference (SSP585 - SSP245)')

plt.tight_layout()
plt.savefig('scenario_comparison.png', dpi=150, bbox_inches='tight')
plt.show()

# Print statistics
print(f"SSP245 mean: {float(mean_ssp245.mean()):.2f} mm/day")
print(f"SSP585 mean: {float(mean_ssp585.mean()):.2f} mm/day")
print(f"Difference: {float(diff.mean()):.2f} mm/day")

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

# Load all years for a scenario
scenario = "2050_SSP245"
files = sorted(glob.glob(f'data/CHC_CMIP6/{scenario}/*.nc'))

# Open as multi-file dataset
ds = xr.open_mfdataset(files, combine='by_coords')
print(ds)

# Compute annual totals
annual_total = ds.pr.resample(time='YE').sum()

# Spatial mean time series
spatial_mean = annual_total.mean(dim=['y', 'x'])

# Plot time series
plt.figure(figsize=(12, 5))
spatial_mean.plot(marker='o', linewidth=2, color='steelblue')
plt.axhline(y=float(spatial_mean.mean()), color='red', linestyle='--', 
            label=f'Mean: {float(spatial_mean.mean()):.0f} mm/year')
plt.xlabel('Year')
plt.ylabel('Annual Precipitation (mm)')
plt.title(f'{scenario} Annual Precipitation Trend - Ethiopia')
plt.legend()
plt.grid(True, alpha=0.3)
plt.savefig('precipitation_trend.png', dpi=150, bbox_inches='tight')
plt.show()

# Compute linear trend
years = np.arange(len(spatial_mean))
coeffs = np.polyfit(years, spatial_mean.values, 1)
trend_mm_per_decade = coeffs[0] * 10

print(f"Linear trend: {trend_mm_per_decade:.1f} mm/decade")