from pathlib import Pathimport eeimport geemapimport geopandas as gpdimport matplotlib as mplimport matplotlib.pyplot as pltimport numpy as npimport pandas as pdimport rioxarrayimport xarray as xrfrom matplotlib import font_managerfrom shapely.geometry import shapefrom xee import helpersimport xclim.indicators.atmos as xindexPROJECT = "giseryu"ROI_ASSET = "projects/giseryu/assets/HUANGTUGAOYUAN"SPINUP_START = "2009-01-01"ANALYSIS_START = "2010-01-01"ANALYSIS_END = "2020-01-01"GRID_DEGREES = 0.3HIGH_KBDI_THRESHOLD = 100.0OUTPUT_DIR = Path(r"E:\黄土高原_KBDI结果")available_fonts = {font.name for font in font_manager.fontManager.ttflist}chinese_candidates = ["Microsoft YaHei", "SimHei", "Noto Sans CJK SC", "Arial Unicode MS"]CHINESE_FONT = next((name for name in chinese_candidates if name in available_fonts), "DejaVu Sans")mpl.rcParams.update({    "font.family": "sans-serif",    "font.sans-serif": [CHINESE_FONT, "Arial", "DejaVu Sans"],    "axes.unicode_minus": False,    "svg.fonttype": "none",    "pdf.fonttype": 42,    "font.size": 9,    "axes.spines.top": False,    "axes.spines.right": False,    "figure.dpi": 120,})print(f"Chinese font: {CHINESE_FONT}")

FORCE_AUTH = Falseif FORCE_AUTH:    ee.Authenticate(force=True)try:    ee.Initialize(        project=PROJECT,        opt_url="https://earthengine-highvolume.googleapis.com",    )except Exception:    ee.Authenticate()    ee.Initialize(        project=PROJECT,        opt_url="https://earthengine-highvolume.googleapis.com",    )print(f"Earth Engine initialized with project: {PROJECT}")

roi_fc = ee.FeatureCollection(ROI_ASSET)roi = roi_fc.geometry().simplify(maxError=1000)feature_count = roi_fc.size().getInfo()area_km2 = roi.area(maxError=1000).divide(1e6).getInfo()bounds = roi.bounds(maxError=1000).getInfo()print(f"Feature count: {feature_count}")print(f"Area: {area_km2:,.2f} km2")print(f"Bounds: {bounds}")

study_map = geemap.Map(basemap="SATELLITE")study_map.centerObject(roi_fc, 5)study_map.addLayer(roi_fc, {"color": "yellow"}, "黄土高原研究区")study_map

era5 = (    ee.ImageCollection("ECMWF/ERA5/DAILY")    .filterDate(SPINUP_START, ANALYSIS_END)    .filterBounds(roi)    .select(        ["maximum_2m_air_temperature", "total_precipitation"],        ["max_temp", "pr"],    ))print(f"ERA5 image count including spin-up: {era5.size().getInfo()}")print(f"Bands: {ee.Image(era5.first()).bandNames().getInfo()}")

era5_map = geemap.Map()era5_map.centerObject(roi_fc, 5)mean_temperature = era5.select("max_temp").mean().subtract(273.15).clip(roi)annual_precipitation_gee = (    era5.filterDate(ANALYSIS_START, ANALYSIS_END)    .select("pr")    .sum()    .divide(10)    .multiply(1000)    .clip(roi))era5_map.addLayer(    mean_temperature,    {"min": 5, "max": 25, "palette": ["313695", "74add1", "ffffbf", "f46d43", "a50026"]},    "多年平均日最高温度（°C）",)era5_map.addLayer(    annual_precipitation_gee,    {"min": 200, "max": 800, "palette": ["fff7fb", "9ecae1", "3182bd", "08519c"]},    "多年平均年降水量（mm）",    False,)era5_map.addLayer(roi_fc, {"color": "black"}, "研究区边界")era5_map

roi_shapely = shape(roi.getInfo())grid_definition = helpers.fit_geometry(    geometry=roi_shapely,    grid_crs="EPSG:4326",    grid_scale=(GRID_DEGREES, -GRID_DEGREES),)dataset = xr.open_dataset(era5, engine="ee", **grid_definition)dataset = dataset.sortby("time") * 1dataset

dataset["pr"] = dataset["pr"] * 1000.0dataset["pr"].attrs["units"] = "mm/day"dataset["max_temp"] = dataset["max_temp"] - 273.15dataset["max_temp"].attrs["units"] = "degC"analysis_dataset = dataset.sel(time=slice(ANALYSIS_START, "2019-12-31"))mean_annual_precipitation = (    analysis_dataset["pr"]    .resample(time="YE")    .sum(dim="time", skipna=True)    .mean(dim="time", skipna=True))mean_annual_precipitation.attrs["units"] = "mm/year"roi_gdf = geemap.ee_to_gdf(roi_fc)roi_gdf = roi_gdf.set_crs("EPSG:4326") if roi_gdf.crs is None else roi_gdf.to_crs("EPSG:4326")dataset = dataset.rio.write_crs("EPSG:4326")mean_annual_precipitation = mean_annual_precipitation.rio.write_crs("EPSG:4326")dataset_clip = dataset.rio.clip(roi_gdf.geometry, roi_gdf.crs, drop=True)precipitation_climatology = mean_annual_precipitation.rio.clip(roi_gdf.geometry, roi_gdf.crs, drop=True)analysis_clip = dataset_clip.sel(time=slice(ANALYSIS_START, "2019-12-31"))print(dataset_clip)

temperature_climatology = analysis_clip["max_temp"].mean("time", skipna=True)valid_fraction = analysis_clip["pr"].notnull().mean("time") * 100fig, axes = plt.subplots(1, 3, figsize=(15, 4.5), constrained_layout=True)precipitation_climatology.plot(    ax=axes[0], cmap="Blues", robust=True,    cbar_kwargs={"label": "Annual precipitation (mm/year)"},)temperature_climatology.plot(    ax=axes[1], cmap="RdYlBu_r", robust=True,    cbar_kwargs={"label": "Maximum temperature (°C)"},)valid_fraction.plot(    ax=axes[2], cmap="Greens", vmin=95, vmax=100,    cbar_kwargs={"label": "Valid observations (%)"},)titles = ["多年平均年降水量", "多年平均日最高温度", "有效观测覆盖率"]for ax, title in zip(axes, titles):    roi_gdf.boundary.plot(ax=ax, color="black", linewidth=0.7)    ax.set_title(title)    ax.set_xlabel("Longitude")    ax.set_ylabel("Latitude")plt.show()

monthly_precipitation = analysis_clip["pr"].groupby("time.month").sum("time", skipna=True) / 10monthly_temperature = analysis_clip["max_temp"].groupby("time.month").mean("time", skipna=True)regional_monthly_precipitation = monthly_precipitation.mean(("x", "y"), skipna=True)regional_monthly_temperature = monthly_temperature.mean(("x", "y"), skipna=True)months = np.arange(1, 13)fig, ax1 = plt.subplots(figsize=(9, 4.5))ax1.bar(months, regional_monthly_precipitation, color="#5DA5DA", width=0.72, label="Precipitation")ax1.set_xlabel("Month")ax1.set_ylabel("Precipitation (mm/month)", color="#2878B5")ax1.tick_params(axis="y", labelcolor="#2878B5")ax1.set_xticks(months)ax2 = ax1.twinx()ax2.plot(months, regional_monthly_temperature, color="#D9534F", marker="o", linewidth=1.8, label="Temperature")ax2.set_ylabel("Maximum temperature (°C)", color="#C43C39")ax2.tick_params(axis="y", labelcolor="#C43C39")ax2.spines["right"].set_visible(True)ax1.set_title("黄土高原月平均降水与最高温度季节循环")ax1.grid(axis="y", alpha=0.2)plt.tight_layout()plt.show()

kbdi_full = xindex.keetch_byram_drought_index(    pr=dataset_clip["pr"],    tasmax=dataset_clip["max_temp"],    pr_annual=precipitation_climatology,)kbdi = kbdi_full.sel(time=slice(ANALYSIS_START, "2019-12-31"))kbdi.name = "KBDI"kbdi.attrs["long_name"] = "Keetch-Byram drought index"kbdi.attrs["units"] = "mm"print(f"KBDI range: {float(kbdi.min()):.2f} to {float(kbdi.max()):.2f} mm")print(f"Missing fraction: {float(kbdi.isnull().mean() * 100):.3f}%")kbdi

annual_maximum = kbdi.resample(time="YE").max("time", skipna=True)median_annual_maximum = annual_maximum.median("time", skipna=True)temporal_p95 = kbdi.quantile(0.95, dim="time", skipna=True)high_kbdi_frequency = (kbdi >= HIGH_KBDI_THRESHOLD).mean("time") * 100absolute_maximum = kbdi.max("time", skipna=True)fig, axes = plt.subplots(2, 2, figsize=(12, 9), constrained_layout=True)maps = [median_annual_maximum, temporal_p95, high_kbdi_frequency, absolute_maximum]titles = [    "多年中位年最大KBDI",    "逐日KBDI的95百分位",    f"高KBDI日频率（≥{HIGH_KBDI_THRESHOLD:.0f} mm）",    "2010—2019年绝对最大KBDI",]cmaps = ["YlOrRd", "YlOrRd", "OrRd", "YlOrRd"]labels = ["KBDI (mm)", "KBDI (mm)", "Frequency (%)", "KBDI (mm)"]for ax, data, title, cmap, label in zip(axes.flat, maps, titles, cmaps, labels):    data.plot(ax=ax, cmap=cmap, robust=True, cbar_kwargs={"label": label})    roi_gdf.boundary.plot(ax=ax, color="black", linewidth=0.7)    ax.set_title(title)    ax.set_xlabel("Longitude")    ax.set_ylabel("Latitude")plt.show()

year_labels = annual_maximum["time"].dt.year.valuesfig, axes = plt.subplots(2, 5, figsize=(16, 6.8), constrained_layout=True, sharex=True, sharey=True)shared_min = float(annual_maximum.quantile(0.02))shared_max = float(annual_maximum.quantile(0.98))for index, (ax, year) in enumerate(zip(axes.flat, year_labels)):    annual_maximum.isel(time=index).plot(        ax=ax,        cmap="YlOrRd",        vmin=shared_min,        vmax=shared_max,        add_colorbar=False,    )    roi_gdf.boundary.plot(ax=ax, color="black", linewidth=0.45)    ax.set_title(f"{int(year)}年")    ax.set_xlabel("")    ax.set_ylabel("")normalizer = mpl.colors.Normalize(vmin=shared_min, vmax=shared_max)colorbar = fig.colorbar(    mpl.cm.ScalarMappable(norm=normalizer, cmap="YlOrRd"),    ax=axes,    orientation="horizontal",    fraction=0.04,    pad=0.06,)colorbar.set_label("Annual maximum KBDI (mm)")fig.suptitle("黄土高原逐年最大KBDI空间分布", fontsize=13)plt.show()

regional_daily_mean = kbdi.mean(("x", "y"), skipna=True)fig, ax = plt.subplots(figsize=(13, 4.5))regional_daily_mean.plot(ax=ax, color="#C43C39", linewidth=0.8)ax.axhline(HIGH_KBDI_THRESHOLD, color="#4D4D4D", linestyle="--", linewidth=0.9, label=f"{HIGH_KBDI_THRESHOLD:.0f} mm")ax.set_title("黄土高原逐日空间平均KBDI")ax.set_xlabel("Date")ax.set_ylabel("KBDI (mm)")ax.grid(alpha=0.2)ax.legend()plt.tight_layout()plt.show()

day_of_year = regional_daily_mean.groupby("time.dayofyear")seasonal_median = day_of_year.median("time", skipna=True)seasonal_p10 = day_of_year.quantile(0.10, dim="time", skipna=True)seasonal_p90 = day_of_year.quantile(0.90, dim="time", skipna=True)fig, ax = plt.subplots(figsize=(9, 4.5))ax.fill_between(    seasonal_median["dayofyear"],    seasonal_p10,    seasonal_p90,    color="#F4A582",    alpha=0.35,    label="10th–90th percentile",)ax.plot(seasonal_median["dayofyear"], seasonal_median, color="#B2182B", linewidth=1.8, label="Median")ax.set_title("黄土高原KBDI多年平均季节循环")ax.set_xlabel("Day of year")ax.set_ylabel("KBDI (mm)")ax.set_xlim(1, 366)ax.grid(alpha=0.2)ax.legend()plt.tight_layout()plt.show()

regional_annual_maximum = annual_maximum.mean(("x", "y"), skipna=True)regional_annual_mean = kbdi.resample(time="YE").mean("time", skipna=True).mean(("x", "y"), skipna=True)regional_high_days = (kbdi >= HIGH_KBDI_THRESHOLD).resample(time="YE").sum("time").mean(("x", "y"), skipna=True)years = regional_annual_maximum["time"].dt.year.valuesfig, axes = plt.subplots(1, 3, figsize=(15, 4.2), constrained_layout=True)axes[0].plot(years, regional_annual_maximum, marker="o", color="#B2182B")axes[0].set_title("区域平均年最大KBDI")axes[0].set_ylabel("KBDI (mm)")axes[1].plot(years, regional_annual_mean, marker="o", color="#EF8A62")axes[1].set_title("区域平均年均KBDI")axes[1].set_ylabel("KBDI (mm)")axes[2].bar(years, regional_high_days, color="#D6604D")axes[2].set_title(f"年均高KBDI日数（≥{HIGH_KBDI_THRESHOLD:.0f} mm）")axes[2].set_ylabel("Days/year")for ax in axes:    ax.set_xlabel("Year")    ax.set_xticks(years)    ax.tick_params(axis="x", rotation=45)    ax.grid(axis="y", alpha=0.2)plt.show()

annual_summary = pd.DataFrame({    "year": years.astype(int),    "regional_mean_kbdi_mm": regional_annual_mean.values,    "regional_mean_annual_max_kbdi_mm": regional_annual_maximum.values,    f"regional_mean_days_kbdi_ge_{int(HIGH_KBDI_THRESHOLD)}": regional_high_days.values,})annual_summary.round(2)