气候指数计算教程：SPI、SPEI 与 PET 的实现

!pip install climate-indices -i https://pypi.mirrors.ustc.edu.cn/simple/

import numpy as np
import pandas as pd
import xarray as xr

# 1. 设定时间范围（30年，月频率）
times = pd.date_range("1991-01-01", "2020-12-31", freq="MS")
lats = np.array([30.0, 31.0])
lons = np.array([110.0, 111.0])

# --- 修复部分 ---
# 将 times.month 转换为 numpy 数组，再进行维度扩展 [:, None, None]
# 这样是为了将 (360,) 形状的数组变成 (360, 1, 1)，以便与 (360, 2, 2) 的格点数据对齐
month_values = times.month.values  # 转换为 numpy 数组
seasonal_cycle = np.sin(2 * np.pi * month_values / 12)
seasonal_cycle_3d = seasonal_cycle[:, None, None] 
# ----------------

# 2. 模拟降水数据 (单位: mm) - 包含季节性
np.random.seed(42)
# 降水通常也有季节性，这里简单模拟
precip_base = 50 + 40 * seasonal_cycle_3d 
precip_data = np.maximum(0, precip_base + np.random.normal(0, 15, (len(times), len(lats), len(lons))))

# 3. 模拟气温数据 (单位: Celsius) - 修复后的写法
temp_data = 15 + 10 * seasonal_cycle_3d + np.random.normal(0, 2, (len(times), len(lats), len(lons)))

# 4. 封装进 xarray Dataset
ds = xr.Dataset(
    {
        "precip": (["time", "lat", "lon"], precip_data),
        "temp": (["time", "lat", "lon"], temp_data),
    },
    coords={
        "time": times,
        "lat": lats,
        "lon": lons,
    },
)

print("数据集成功创建！")
print(ds)

from climate_indices import indices
from climate_indices import compute
# 提取其中一个格点的数据
selected_lat = ds.lat.values[0]
temp_series = ds.temp.sel(lat=selected_lat, lon=110.0).values

# 在 climate-indices 2.0.0 中，计算 PET 的方式如下：
# 我们需要指定 pet_type (如 thornthwaite)
pet = indices.pet(
    temperature_celsius=temp_series,
    latitude_degrees=selected_lat,
    data_start_year=1991# 对应我们模拟数据的开始年份
)

print(f"PET 计算成功！前5个月数值: {pet[:5]}")

precip_series = ds.precip.sel(lat=30.0, lon=110.0).values

# 计算 3 个月尺度的 SPI
spi_3 = indices.spi(
    values=precip_series,
    scale=3,
    distribution=indices.Distribution.gamma,
    data_start_year=1991,
    calibration_year_initial=1991,
    calibration_year_final=2020,
    periodicity=compute.Periodicity.monthly  # <--- 新增这个参数
)

print(f"SPI-3 计算成功！前5个月数值:\n{spi_3[:5]}")

spei_6 = indices.spei(precips_mm=precip_series,
                    pet_mm=pet,
                    scale=6, # 6个月尺度
                    distribution=indices.Distribution.gamma,
                    periodicity=compute.Periodicity.monthly,
                    data_start_year=1991,
                    calibration_year_initial=1991,
                    calibration_year_final=2020,
                    )
print(f"前10个月的 SPEI-6 结果: {spei_3[:10]}")

import matplotlib.pyplot as plt

plt.figure(figsize=(12, 6))

# 设置绘图时间轴
plot_times = ds.time.values

plt.plot(plot_times, spi_3, label='SPI-3 (Precipitation)', alpha=0.7, color='blue')
plt.plot(plot_times, spei_6, label='SPEI-6 (Water Balance)', alpha=0.7, color='red')

# 添加 0 基准线和旱涝阈值
plt.axhline(0, color='black', linestyle='-', linewidth=1)
plt.axhline(-1, color='orange', linestyle='--', label='Moderate Drought')
plt.axhline(-1.5, color='red', linestyle='--', label='Severe Drought')

plt.title("Climate Indices Comparison (Synthetic Data)")
plt.xlabel("Year")
plt.ylabel("Index Value")
plt.legend()
plt.grid(True, alpha=0.3)
plt.show()