Daymet Daily V4 R1 from NASA Earthdata Cloud via OPeNDAP (Hyrax)

import datetime as dt
import json
import os
import time
from pathlib import Path

import earthaccess
import geopandas as gpd
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, PillowWriter
import numpy as np
import pydap
import requests
import rioxarray  # noqa: F401  (registers the .rio accessor)
import xarray as xr
from pydap.client import get_cmr_urls, open_url
from pydap.client import to_netcdf as dap_to_netcdf
from IPython.display import Image, display
import warnings
warnings.filterwarnings("ignore")

print("This notebook was produced with:")
for m in (earthaccess, pydap, xr, gpd, np):
    print(f"  {m.__name__:14s} {m.__version__}")

WORKDIR = Path("./daymet_tva")
(WORKDIR / "data").mkdir(parents=True, exist_ok=True)
print("\nOutputs ->", WORKDIR.resolve())

from earthaccess.exceptions import LoginStrategyUnavailable

try:
    auth = earthaccess.login(strategy="netrc", persist=True)
except LoginStrategyUnavailable:
    auth = earthaccess.login(strategy="interactive", persist=True)

# A requests.Session carrying EDL authorization; pass this to pydap / xarray.
my_session = auth.get_session()
print("Authenticated:", auth.authenticated)
print("Session:", type(my_session).__name__)

LOCAL_POLYGON = None   # e.g. "bnds/GRSM_BOUNDARY_POLYGON_fid17.shp" to use your own polygon
CACHE = WORKDIR / "tva_power_service_area.geojson"

TVA_SERVICE = (
    "https://services5.arcgis.com/bPacKTm9cauMXVfn/arcgis/rest/services/"
    "TVA_Power_Service_Area/FeatureServer/0/query"
)

if LOCAL_POLYGON:
    roi = gpd.read_file(LOCAL_POLYGON)
elif CACHE.exists():
    roi = gpd.read_file(CACHE)
else:
    params = {"where": "1=1", "outFields": "OID,Name", "returnGeometry": "true",
              "outSR": 4326, "f": "geojson"}
    r = requests.get(TVA_SERVICE, params=params, timeout=120)
    r.raise_for_status()
    roi = gpd.GeoDataFrame.from_features(r.json()["features"], crs="EPSG:4326")
    roi.to_file(CACHE, driver="GeoJSON")
    print("Cached ->", CACHE)

# dissolve multiple parts into a single region of interest
roi = roi.dissolve().reset_index(drop=True)[["geometry"]]

print("CRS              :", roi.crs)
print("Geographic bounds:", np.round(roi.total_bounds, 3))

DAYMET_PROJ = ("+proj=lcc +ellps=WGS84 +a=6378137 +b=6356752.314245 "
               "+lat_1=25 +lat_2=60 +lon_0=-100 +lat_0=42.5 +x_0=0 +y_0=0 "
               "+units=m +no_defs")

roi_wgs84 = roi.to_crs(epsg=4326)
roi_lcc = roi.to_crs(DAYMET_PROJ)

# bounding box for CMR search: [W, S, E, N]
bbox_wgs84 = [round(float(v), 4) for v in roi_wgs84.total_bounds]

# bounding box for OPeNDAP subsetting, in Daymet LCC metres, padded by 2 grid cells
PAD = 2000.0
xmin, ymin, xmax, ymax = roi_lcc.total_bounds
bbox_lcc = (xmin - PAD, ymin - PAD, xmax + PAD, ymax + PAD)

print("CMR search bbox (WGS84) :", bbox_wgs84)
print("Subset bbox (Daymet LCC):", tuple(round(v) for v in bbox_lcc))
print(f"Approx. subset size     : {(bbox_lcc[2]-bbox_lcc[0])/1000:.0f} x "
      f"{(bbox_lcc[3]-bbox_lcc[1])/1000:.0f} km "
      f"(~{(bbox_lcc[2]-bbox_lcc[0])/1000*(bbox_lcc[3]-bbox_lcc[1])/1000/1e3:.0f}k grid cells)")

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 5))

roi_wgs84.boundary.plot(ax=ax1, color="tab:blue", linewidth=1)
ax1.set_title("TVA power service area\nEPSG:4326 (degrees)")
ax1.set_xlabel("longitude"); ax1.set_ylabel("latitude")

roi_lcc.plot(ax=ax2, facecolor="0.85", edgecolor="tab:red", linewidth=1)
ax2.set_title("Same polygon in Daymet LCC\n(metres from 42.5N, 100W)")
ax2.set_xlabel("x (m)"); ax2.set_ylabel("y (m)")
for ax in (ax1, ax2):
    ax.set_aspect("equal"); ax.grid(alpha=0.3)
plt.tight_layout()

START = dt.datetime(2023, 4, 1)
END = dt.datetime(2023, 6, 30)
VARIABLES = ["tmax", "prcp"]      # any of tmax, tmin, prcp, srad, vp, swe, dayl
REGION = "na"                     # 'na' (North America), 'hi' (Hawaii), 'pr' (Puerto Rico)

YEARS = list(range(START.year, END.year + 1))
DT_FMT = "%Y-%m-%dT%H:%M:%SZ"
temporal_str = f"{START.strftime(DT_FMT)},{END.strftime(DT_FMT)}"

# string-based slice: works for datetime64 and cftime time axes alike
TIME_SLICE = slice(f"{START:%Y-%m-%d}", f"{END:%Y-%m-%d}")

print("time range :", temporal_str)
print("years      :", YEARS)
print("variables  :", VARIABLES)
print("region     :", REGION)

DOI = "10.3334/ORNLDAAC/2129"

# Resolve the collection concept ID from the DOI (handy if the collection is ever re-versioned)
doisearch = "https://cmr.earthdata.nasa.gov/search/collections.json?doi=" + DOI
CCID = requests.get(doisearch, timeout=60).json()["feed"]["entry"][0]["id"]
print("collection concept ID:", CCID)

cmr_urls = get_cmr_urls(ccid=CCID, time_range=[START, END], limit=1000)
print(f"{len(cmr_urls)} granule URLs returned by CMR for {START:%Y-%m-%d} to {END:%Y-%m-%d}")

# The collection bundles all variables and all three regions; keep only what we asked for.
def wanted(url, var, region, years):
    name = url.rsplit("/", 1)[-1]
    return any(f"daily_{region}_{var}_{y}.nc" in name for y in years)

urls = {
    var: sorted(u for u in cmr_urls if wanted(u, var, REGION, YEARS))
    for var in VARIABLES
}

for var, ulist in urls.items():
    print(f"\n{var}: {len(ulist)} granule(s)")
    for u in ulist:
        print("  ", u)

def dap4(url: str) -> str:
    # Convert an https Earthdata OPeNDAP URL into a DAP4 URL
    return url.replace("https://", "dap4://", 1)


sample_url = urls[VARIABLES[0]][0]
print(dap4(sample_url), "\n")

ds = xr.open_dataset(
    dap4(sample_url),
    engine="pydap",
    session=my_session,
    decode_coords="all",   # attaches lat/lon and the lambert_conformal_conic grid mapping
)
ds

# Daymet's native CRS straight from the granule metadata (compare with DAYMET_PROJ above)
print("CRS from file :", ds.rio.crs)
print("dims          :", dict(ds.sizes))
print("x range (m)   :", float(ds.x.min()), "->", float(ds.x.max()))
print("y range (m)   :", float(ds.y.min()), "->", float(ds.y.max()), "(descending: north to south)")
print("time          :", str(ds.time.values[0])[:10], "->", str(ds.time.values[-1])[:10])

def bbox_slices(dataset, bbox_lcc):
    # xarray .sel() slices for a Daymet granule, from an LCC bbox (xmin, ymin, xmax, ymax)
    xmin, ymin, xmax, ymax = bbox_lcc
    y_desc = bool(dataset.y[0] > dataset.y[-1])
    return {
        "x": slice(xmin, xmax),
        "y": slice(ymax, ymin) if y_desc else slice(ymin, ymax),
    }


sl = bbox_slices(ds, bbox_lcc)
probe = ds[VARIABLES[0]].sel(**sl).sel(time=TIME_SLICE)
n = int(np.prod(probe.shape))
print("subset shape      :", dict(zip(probe.dims, probe.shape)))
print("values            :", f"{n:,}")
print("approx. transfer  :", f"{n * 4 / 1e6:.0f} MB per variable")
print("full-granule size :", f"{int(np.prod(ds[VARIABLES[0]].shape)) * 4 / 1e9:.1f} GB")
print("reduction         :", f"{int(np.prod(ds[VARIABLES[0]].shape)) / n:,.0f}x smaller")

sl

subsets = {}

for var, ulist in urls.items():
    t0 = time.time()
    pieces = []
    for u in ulist:
        print(f"opening {u.rsplit('/', 1)[-1]} ...", flush=True)
        with xr.open_dataset(dap4(u), engine="pydap", session=my_session,
                             decode_coords="all") as gds:
            sl = bbox_slices(gds, bbox_lcc)
            piece = gds[[var]].sel(**sl).sel(time=TIME_SLICE)
            pieces.append(piece.load())          # <- data transferred here
    da = xr.concat(pieces, dim="time").sortby("time") if len(pieces) > 1 else pieces[0]
    da = da.rio.write_crs(DAYMET_PROJ)
    subsets[var] = da
    print(f"  {var}: {dict(da.sizes)} in {time.time() - t0:.1f} s\n")

subset = xr.merge(subsets.values())
subset.attrs["title"] = "Daymet V4 R1 daily subset, TVA power service area"
subset.attrs["source"] = "NASA Earthdata Cloud OPeNDAP (Hyrax), collection " + CCID
subset.attrs["history"] = f"Subset via DAP4 on {dt.date.today():%Y-%m-%d}"
subset

sl

out_nc = WORKDIR / "data" / (
    f"daymet_v4_tva_{'_'.join(VARIABLES)}_{START:%Y%m%d}_{END:%Y%m%d}.nc"
)
encoding = {v: {"zlib": True, "complevel": 4} for v in VARIABLES}
subset.to_netcdf(out_nc, encoding=encoding)
print(f"wrote {out_nc}  ({out_nc.stat().st_size / 1e6:.1f} MB on disk)")

from shapely.geometry import mapping

clipped = subset.rio.clip(
    roi_lcc.geometry.apply(mapping), roi_lcc.crs, drop=False, all_touched=True
)

out_clip = out_nc.with_name(out_nc.stem + "_tvaclip.nc")
clipped.to_netcdf(out_clip, encoding=encoding)
print("wrote", out_clip)

frac = float(clipped[VARIABLES[0]].isel(time=0).notnull().mean())
print(f"{frac:.0%} of the bounding-box cells fall inside the TVA polygon")

day = clipped.time.values[len(clipped.time) // 2]

fig, axes = plt.subplots(1, 2, figsize=(14, 5.5))
for ax, var, cmap, label in zip(
    axes, VARIABLES, ["turbo", "Blues"],
    ["daily max temperature (deg C)", "precipitation (mm/day)"],
):
    clipped[var].sel(time=day).plot(
        ax=ax, cmap=cmap, robust=True, cbar_kwargs={"label": label}
    )
    roi_lcc.boundary.plot(ax=ax, color="black", linewidth=0.8)
    ax.set_title(f"Daymet {var} - {str(day)[:10]}")
    ax.set_xlabel("x (m)"); ax.set_ylabel("y (m)")
    ax.set_aspect("equal")
plt.tight_layout()
fig.savefig(WORKDIR / "tva_daily_maps.png", dpi=150, bbox_inches="tight")

# Season summaries over the clipped region, exported as Cloud-friendly GeoTIFFs
summaries = {
    f"{VARIABLES[0]}_mean": clipped[VARIABLES[0]].mean("time", keep_attrs=True),
    f"{VARIABLES[0]}_days_above_32C": (clipped[VARIABLES[0]] > 32).sum("time").where(
        clipped[VARIABLES[0]].isel(time=0).notnull()
    ),
}
if "prcp" in VARIABLES:
    summaries["prcp_total"] = clipped["prcp"].sum("time", keep_attrs=True).where(
        clipped["prcp"].isel(time=0).notnull()
    )

for name, da in summaries.items():
    da = da.rio.write_crs(DAYMET_PROJ)
    tif = WORKDIR / "data" / f"daymet_v4_tva_{name}_{START:%Y%m%d}_{END:%Y%m%d}.tif"
    da.rio.to_raster(tif)
    print("wrote", tif.name)

fig, axes = plt.subplots(1, len(summaries), figsize=(6 * len(summaries), 5))
for ax, (name, da) in zip(np.atleast_1d(axes), summaries.items()):
    da.plot(ax=ax, robust=True, cmap="turbo", cbar_kwargs={"label": name})
    roi_lcc.boundary.plot(ax=ax, color="black", linewidth=0.8)
    ax.set_title(f"{name}\n{START:%Y-%m-%d} to {END:%Y-%m-%d}")
    ax.set_aspect("equal"); ax.set_xlabel("x (m)"); ax.set_ylabel("y (m)")
plt.tight_layout()
fig.savefig(WORKDIR / "tva_season_summaries.png", dpi=150, bbox_inches="tight")

### for fun: daily precipitation GIF from all 91 days, with one fixed color scale so the frames are comparable

prcp = clipped["prcp"]  

# Keep one scale for all 91 frames; avoids color changes from day to day.
vmin = 0
vmax = float(prcp.quantile(0.99))
norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)

fig, ax = plt.subplots(figsize=(7, 6))
colorbar = fig.colorbar(
    mpl.cm.ScalarMappable(norm=norm, cmap="Blues"),
    ax=ax,
    label="precipitation (mm/day)",
)

def draw_frame(i):
    ax.clear()

    day = prcp.time.values[i]
    prcp.isel(time=i).plot(
        ax=ax,
        cmap="Blues",
        norm=norm,
        add_colorbar=False,
    )

    roi_lcc.boundary.plot(ax=ax, color="black", linewidth=0.8)
    ax.set_title(f"Daymet precipitation — {str(day)[:10]}")
    ax.set_xlabel("x (m)")
    ax.set_ylabel("y (m)")
    ax.set_aspect("equal")

animation = FuncAnimation(
    fig,
    draw_frame,
    frames=prcp.sizes["time"],
    interval=125,  # milliseconds per frame
)

gif_path = WORKDIR / "tva_daily_precipitation.gif"
animation.save(
    gif_path,
    writer=PillowWriter(fps=8),
    dpi=120,
)

plt.close(fig)
print(f"Saved: {gif_path}")

Image(filename=str(gif_path))

area_mean = clipped[VARIABLES[0]].mean(dim=("x", "y"))
area_max = clipped[VARIABLES[0]].max(dim=("x", "y"))
area_min = clipped[VARIABLES[0]].min(dim=("x", "y"))

fig, axes = plt.subplots(2, 1, figsize=(12, 7), sharex=True)

axes[0].fill_between(clipped.time, area_min, area_max, alpha=0.25,
                     color="tab:red", label="spatial min-max")
axes[0].plot(clipped.time, area_mean, color="tab:red", lw=1.6, label="area mean")
axes[0].set_ylabel("tmax (deg C)")
axes[0].set_title("Daymet daily maximum temperature, TVA power service area")
axes[0].legend(); axes[0].grid(alpha=0.3)

if "prcp" in VARIABLES:
    axes[1].bar(clipped.time, clipped["prcp"].mean(dim=("x", "y")),
                width=1.0, color="tab:blue")
    axes[1].set_ylabel("area-mean prcp (mm/day)")
    axes[1].set_title("Daymet daily precipitation, area average")
    axes[1].grid(alpha=0.3)
axes[1].set_xlabel("date")
plt.tight_layout()
fig.savefig(WORKDIR / "tva_timeseries.png", dpi=150, bbox_inches="tight")

print(f"Area-mean tmax  : {float(area_mean.mean()):.2f} deg C")
print(f"Hottest day     : {str(area_mean.idxmax().values)[:10]} "
      f"({float(area_mean.max()):.2f} deg C area mean)")
if "prcp" in VARIABLES:
    tot = float(clipped['prcp'].mean(dim=('x', 'y')).sum())
    print(f"Area-mean total precipitation: {tot:.0f} mm")

# Single-pixel ("Daymet single pixel extraction" style) time series for selected sites
sites = {
    "Knoxville, TN": (-83.9207, 35.9606),
    "Chattanooga, TN": (-85.3097, 35.0456),
    "Huntsville, AL": (-86.5861, 34.7304),
}
site_pts = gpd.GeoDataFrame(
    {"site": list(sites)},
    geometry=gpd.points_from_xy(*zip(*sites.values())),
    crs="EPSG:4326",
).to_crs(DAYMET_PROJ)

fig, ax = plt.subplots(figsize=(12, 4.5))
for name, pt in zip(site_pts.site, site_pts.geometry):
    ts = clipped[VARIABLES[0]].sel(x=pt.x, y=pt.y, method="nearest")
    ax.plot(ts.time, ts, lw=1.2, label=f"{name} ({float(ts.mean()):.1f} deg C mean)")
ax.set_ylabel("tmax (deg C)"); ax.set_xlabel("date")
ax.set_title("Daymet nearest-pixel daily tmax")
ax.legend(); ax.grid(alpha=0.3)
plt.tight_layout()
fig.savefig(WORKDIR / "tva_site_timeseries.png", dpi=150, bbox_inches="tight")

def index_slices(dataset, bbox_lcc):
    # Convert an LCC bbox into DAP4 integer dim_slices for /x and /y
    xmin, ymin, xmax, ymax = bbox_lcc
    xv = np.asarray(dataset.x); yv = np.asarray(dataset.y)
    xi = np.where((xv >= xmin) & (xv <= xmax))[0]
    yi = np.where((yv >= ymin) & (yv <= ymax))[0]
    return {"/x": (int(xi.min()), int(xi.max())), "/y": (int(yi.min()), int(yi.max()))}


dim_slices = index_slices(ds, bbox_lcc)
print("dim_slices:", dim_slices)

BULK = True # Turn off to False to avoid downloading.
if BULK:
    bulk_years = list(range(2015, 2025))

    bulk_cmr_urls = get_cmr_urls(
        ccid=CCID,
        time_range=[
            dt.datetime(bulk_years[0], 1, 1),
            dt.datetime(bulk_years[-1], 12, 31),
        ],
        limit=1000,
    )

    # Leave these as https:// URLs for dap_to_netcdf.
    bulk_urls = [
        u for u in bulk_cmr_urls
        if wanted(u, "tmax", REGION, bulk_years)
    ]

    print(f"streaming {len(bulk_urls)} granules")
    bulk_dir = Path(WORKDIR) / "data" / "bulk"
    bulk_dir.mkdir(parents=True, exist_ok=True)
    print("Writing to:", bulk_dir.resolve())
    probe = bulk_dir / ".permission_check"
    probe.write_text("ok")
    probe.unlink()

    tile_root = Path(WORKDIR) / "data" / "tiles"
    final_dir = Path(WORKDIR) / "data" / "bulk"
    tile_root.mkdir(parents=True, exist_ok=True)
    final_dir.mkdir(parents=True, exist_ok=True)
    
    # Recover the existing spatial bounds.
    y0, y1 = dim_slices["/y"]
    x0, x1 = dim_slices["/x"]
    
    strip_height = 48  # lower this to 48 if a strip still times out
    
    for url in bulk_urls:
        year_name = url.rsplit("/", 1)[-1]
        print(f"\nDownloading {year_name}")
    
        tile_files = []
    
        for ys in range(y0, y1, strip_height):
            ye = min(ys + strip_height, y1)
            tile_dir = tile_root / f"{year_name}_y{ys}_{ye}"
            tile_dir.mkdir(parents=True, exist_ok=True)
    
            print(f"  y indices {ys}:{ye}")
            dap_to_netcdf(
                [url],
                session=my_session,
                output_path=str(tile_dir),
                keep_variables=["/time", "/y", "/x", "/lat", "/lon", "/tmax"],
                dim_slices={
                    "/y": (ys, ye),
                    "/x": (x0, x1),
                },
            )
            tile_files.extend(tile_dir.glob("*.nc4"))
        # Reassemble that year's y-strips into one local NetCDF file.
        with xr.open_mfdataset(
            [str(p) for p in tile_files],
            combine="by_coords",
            decode_coords="all",
        ) as ds:
            ds.load().to_netcdf(final_dir / f"{year_name}.nc4")
else:
    print("BULK is False - set it to True to stream the multi-year archive.")

# Open the finished yearly files
print (f"The final datasets are stored here: {final_dir}.")
ds_bulk = xr.open_mfdataset(
    str(final_dir / "*.nc4"),
    combine="by_coords",
    decode_coords="all",
)
ds_bulk

BULK = True

if BULK:
    bulk_years = list(range(2015, 2025))

    bulk_cmr_urls = get_cmr_urls(
        ccid=CCID,
        time_range=[
            dt.datetime(bulk_years[0], 1, 1),
            dt.datetime(bulk_years[-1], 12, 31),
        ],
        limit=1000,
    )

    # Leave these as https:// URLs for dap_to_netcdf.
    bulk_urls = [
        u for u in bulk_cmr_urls
        if wanted(u, "tmax", REGION, bulk_years)
    ]

    smoke_dir = Path(WORKDIR) / "data" / "smoke_test"
    smoke_dir.mkdir(parents=True, exist_ok=True)

    tiny_slices = {
    "/time": (0, 1),
    "/y": (0, 9),
    "/x": (0, 9),
    }

    dap_to_netcdf(
    bulk_urls[:1],            # one HTTPS URL only
    session=my_session,
    output_path=str(smoke_dir),
    keep_variables=["/time", "/y", "/x", "/tmax"],
    dim_slices=tiny_slices,
    )

    print("done. Open the result with:")
    print("  xr.open_mfdataset(str(WORKDIR / 'data' / 'smoke_test' / '*.nc'), decode_coords='all')")
else:
    print("BULK is False - set it to True to stream the multi-year archive.")