基于ERA5与SPI的机器学习干旱预测Notebook管线

简介面向气候科学与机器学习交叉方向的开源项目资源聚焦干旱预测这一具体场景提供一套端到端的机器学习流程实现用于统一数据的创建、方法之间的横向比较与效果评估。适合具备 Python 基础、希望把机器学习应用到气候数据分析中的研究者与高年级学生也可作为课程设计或论文实验的参考框架。压缩包约 49.31MB其中 Jupyter Notebook 与 Python 源码为主要组成文件总数与类型明细上游暂未提供。代码按功能拆分在 src 目录下的不同类中分别承担数据预处理、统一格式转换、模型训练与评估等任务配套 environment.yml 可通过 conda 一次性创建 Python 3.7 运行环境降低复现门槛。资源描述中还提到三个入口点以及说明流程设计目标的博客文章便于按模块顺序理解整体架构。目前已有 143 人学习下载适合想系统了解气候数据机器学习流水线组织方式并动手改造实验的读者。1. 连续三个月少雨就算干旱吗ml_drought 想用模型回答这个问题去年入夏后某地四十多天没下过一场透雨群里每天都有人问是不是要旱了。可降水量距平、土壤墒情、潜在蒸散这几条曲线摆在一起谁也说不清这到底算不算气象干旱更别说提前一个月判断它会不会继续恶化。ml_drought这个项目就是冲着这种模糊判断来的它用 Copernicus 气候数据服务里的 ERA5 再分析月尺度数据先按不同时间窗口算出标准化降水指数 SPI 当标签再用机器学习模型预测下一期干旱落在哪个等级整套流程用 Jupyter Notebook 一步步写出来下载、清洗、特征构造、训练、评估各自独立成篇改一个参数就能换区域重跑。它适合两类人一类是想入门气候数据处理的开发者另一类是手里有遥感或气象格点数据、想接一套机器学习基线做二次开发的人。整条链路没有黑匣子每一格都能单独执行、单独检查。2. 项目定位与技术栈一条用 Notebook 串起来的 ERA5 干旱预测管线2.1 它真正解决的是干旱预测里最混乱的那一段干旱这件事麻烦不在于模型而在于「什么叫干旱」本身没有唯一答案。气象干旱看降水亏缺农业干旱看土壤含水量水文干旱看径流和地下水位同一个地方在三种定义下可能一个是正常、一个是重旱。学术界折中的办法是 SPI把某段时间的降水量按历史同期分布做标准化落到标准正态上然后按阈值切等级——大于 -0.5 算正常-0.5 到 -1.0 是轻旱-1.0 到 -1.5 是中旱-1.5 到 -2.0 是重旱低于 -2.0 是特旱。这套阈值不需要人工标注天然适合当机器学习的标签。ml_drought的定位就在这儿它不发明新的干旱定义而是把「算 SPI」和「预测 SPI」这两件事拆开前者用统计方法老老实实做标准化后者交给模型。项目关键词里的 2019 大致对应数据快照和数据截止年份也就是说训练样本覆盖到 2019 年前后的月尺度序列。这样一来模型要学的其实是给定本月及之前若干个月的降水、气温、土壤墒情下一刻的 SPI 等级会往哪边走。这个思路值得说清楚因为很多人第一反应是直接拿降水回归去预测干旱结果模型学到的只是季节性——夏天降水多、冬天降水少RMSE 看着不大但对真正的旱情转折毫无分辨力。SPI 把季节性先抹掉了模型才被迫去学滞后累积和土壤记忆这类真信号。2.2 技术栈xarray 管数据、scipy 管标准化、scikit-learn 管模型仓库主干是 Notebook依赖也集中在几个成熟库上属于典型的科研向 Python 数据栈层次主要库在项目里的职责数据获取cdsapi按变量、年份、区域批量拉取 ERA5 月平均再分析数据数据容器xarraynetCDF4读写 nc 文件按 time/lat/lon 维度切片、重采样数值计算numpypandasscipy.stats滚动累积、Gamma 分布拟合、正态分位数变换建模scikit-learn随机森林、梯度提升、逻辑回归基线交叉验证与指标可视化matplotlibseabornSPI 时间序列、混淆矩阵、特征重要性运行环境condajupyterlab隔离依赖Notebook 交互式调试为什么数据层用 xarray 而不是直接 pandas因为 ERA5 的原始文件是四维的time, lat, lon 加上变量本身一个文件里塞着几十年的格点场用 xarray 可以写成ds[tp].sel(time2015-07, methodnearest)这种人类可读的切片还能惰性加载——只算你要的那块不把整个文件读进内存。等降维到一个格点或一个区域平均之后再.to_dataframe()转成 pandas后面的特征工程就顺手了。为什么模型层停在 scikit-learn不上 LSTM样本量决定的。单格点、月尺度、三十年也就三百六十条记录就算扩展到一片小区域、再堆上多个滞后特征量级也就几千行。这个规模下树模型的表现稳定、训练几秒钟、特征重要性可直接解释硬上深度网络大概率过拟合还得额外解释一堆超参。后续真想上时序模型也是在这条基线之上做增量而不是一上来就替换。2.3 目录结构与执行顺序这类以 Notebook 为主的仓库常见做法是按处理阶段编号让每个文件都能独立跑完并落一份中间结果ml_drought/ ├── data/ │ ├── raw/ # cdsapi 下载的原始 .nc只读不改 │ └── interim/ # 降维后的单点/区域 csv、parquet ├── notebooks/ │ ├── 01_download.ipynb # 按区域和年份批量拉 ERA5 │ ├── 02_preprocess.ipynb # 单位换算、时间轴对齐、缺测检查 │ ├── 03_spi_label.ipynb # 滚动累积 分月 Gamma 拟合生成标签 │ ├── 04_features.ipynb # 滞后项、季节编码、土壤墒情拼特征表 │ └── 05_model.ipynb # 训练、时序交叉验证、混淆矩阵 ├── src/ │ └── drought.py # 被 Notebook import 的稳定函数 ├── figures/ # 导出的图和表 └── environment.yml执行顺序就是编号顺序中间任何一步都能单独重跑。src/drought.py是关键设计一旦某个函数在 Notebook 里调通了就把它挪进srcNotebook 只负责调用和看图。这么做的好处是后来写脚本批处理时不用复制粘贴单元格直接from src.drought import compute_spi就行。很多 Notebook 项目烂尾就是因为所有逻辑都活在.ipynb的 JSON 里改一处要翻十屏。3. 环境与数据准备从 CDS 下载到第一个能跑的 Notebook3.1 建环境与依赖口径先用 conda 建一个干净环境Python 版本别追新3.9 或 3.10 兼容性最省心。xarray、netCDF4、scipy 这几个包在老版本上经常互相打架装完先 import 一遍确认没报错再往下走。# 建环境指定 Python 小版本避免 netCDF4 轮子不匹配 conda create -n mldrought python3.9 -y conda activate mldrought # 数据与计算 pip install numpy pandas xarray netCDF4 scipy cdsapi # 建模与可视化 pip install scikit-learn matplotlib seaborn # 交互环境 pip install jupyterlab这段没什么玄学但有两点值得强调。一是netCDF4一定要走 pip 或 conda 的预编译轮子源码编译会去找系统里的 HDF5 库报错信息极其难读。二是把版本写进environment.yml再提交否则半年后别人 clone 下来scikit-learn 换了默认参数同一份代码跑出不同结果排查起来能耗掉一整天。环境好了之后先在 Jupyter 里建个空 Notebook执行import xarray as xr; print(xr.__version__)能出版本号就说明数据栈通了。这一步看着多余但它把「环境问题」和「代码问题」提前隔离开了后面出岔子时你可以直接排除前者。3.2 CDS 账号与下载脚本ERA5 的数据要走 CDS 的 API得先注册账号拿到 key写进家目录下的配置文件。这个文件长这样两行就够# ~/.cdsapirc url: CDS 控制台给出的 API 地址 key: 你的账号密钥配置文件权限建议chmod 600别顺手提交到仓库里。配好之后cdsapi.Client()能直接读到不用在代码里明文写 key。下载这段是整个项目里最需要讲清楚的地方因为参数错一个后面全白搭import cdsapi c cdsapi.Client() # 区域给我一块矩形框顺序是 北, 西, 南, 东别搞反 area [45, 110, 30, 125] c.retrieve( reanalysis-era5-single-levels-monthly-means, { product_type: monthly_averaged_reanalysis, # 降水、2 米气温、潜在蒸发、第一层土壤体积含水量 variable: [ total_precipitation, 2m_temperature, potential_evaporation, volumetric_soil_water_layer_1, ], year: [str(y) for y in range(1990, 2020)], # 覆盖到 2019 month: [f{m:02d} for m in range(1, 13)], time: 00:00, # 月平均产品固定取 00 时 area: area, format: netcdf, }, data/raw/era5_monthly.nc, )逐项说下参数。product_type选monthly_averaged_reanalysis这是月平均产品和逐小时的reanalysis不是一回事后者体量大几十倍且没有必要。area的顺序是北、西、南、东写成南北颠倒是最常见的翻车点之一跑完发现数据落在海里。time对月平均产品固定填00:00就行它只是格式要求。一次拉三十年、四个变量请求体不算大但服务端排队可能要几分钟到十几分钟脚本里最好加个重试和日志别干等着。如果你的研究区域很大或者想把时间跨度拉到 1979 年建议按十年拆成多个请求分别下载再在预处理里用xr.concat(..., dimtime)拼起来。单次请求过大容易在服务端超时重下又得从头排。3.3 跑通第一篇 Notebook 的检查点数据到手别急着往下写特征先花十分钟做四项体检后面能省下好几小时import xarray as xr ds xr.open_dataset(data/raw/era5_monthly.nc) print(ds) # 看维度、变量名、坐标范围 print(ds[tp].isel(time0).values.shape) # 首月格点形状对不对 print(ds[time].values[[0, -1]]) # 时间起止是否符合预期 print(ds[tp].isnull().mean().item()) # 缺测比例超过 5% 就要查原因第一项看ds的打印结果重点确认变量名有没有被加上奇怪的前后缀以及时间维度是不是叫time、格点维度是不是latitude/longitude——xarray 后面所有的.sel()都依赖这个命名。第二项确认单月格点形状和你的区域匹配如果拉成了全球场多半是area没生效。第三项确认时间轴首尾闰年和月份缺失的问题都会在这里暴露。第四项看缺测率ERA5 本身质量不错出现大面积 NaN 通常是区域框到了海上或者请求参数写错了。最后画一张月降水的气候态曲线十二个月的平均看看峰值是不是落在当地的雨季。这一步是「肉眼校验」比任何断言都直接如果雨季出现在冬天那单位换算或者时间轴几乎肯定出了问题。4. 核心代码拆解SPI 标签、特征表与时间序列交叉验证4.1 单位与时间轴先把口径对齐再谈建模ERA5 的单位是有坑的。降水tp在月平均产品里是「平均每天的累积量」单位是米要转成毫米得先乘 1000再乘当月天数才是月总降水。气温t2m是开尔文。潜在蒸发pev是通量向上为负所以要取负号再乘 1000。import numpy as np import pandas as pd import xarray as xr def load_point(nc_path: str, lat: float, lon: float) - pd.DataFrame: ds xr.open_dataset(nc_path) # 用最近邻取一个格点区域平均的话改成 .mean(dim(latitude, longitude)) pt ds.sel(latitudelat, longitudelon, methodnearest) df pd.DataFrame({ # 米/天 - 毫米/月先乘 1000再乘当月天数 tp: pt[tp].values * 1000.0 * pt[time].dt.days_in_month.values, # 开尔文 - 摄氏度 t2m: pt[t2m].values - 273.15, # 潜在蒸发向上为负取反后变成毫米/月的正数 pet: -pt[pev].values * 1000.0 * pt[time].dt.days_in_month.values, swvl1: pt[swvl1].values, }, indexpd.to_datetime(pt[time].values)) df.index df.index.to_period(M).to_timestamp(M) # 统一到月末 return df.sort_index()这段的重点全在注释里。days_in_month必须乘否则二月和七月会被当成同样长度的月份滚动累积出来的 SPI 在季节交替处会有系统性偏差。pet的符号如果搞反算 SPEI 时水分子减会变成加标签直接失真。to_period(M)是为了把时间戳统一到月末避免不同来源的数据一个是月初一个是月末resample时对不齐。4.2 SPI 标签滚动累积加逐月 Gamma 拟合SPI 的计算分三步先做窗口累积再对每个日历月单独拟合 Gamma 分布最后把累积降水映射到标准正态。第三步是关键一定不能在全样本上拟合。from scipy.stats import gamma, norm def compute_spi(series: pd.Series, scale: int 3, baseline_end: int 2010) - pd.Series: scale: 累积窗口(月)baseline_end: 拟合基准期的最后一年 # 1) 滚动累积min_periods 保证前 scale-1 个位置是 NaN 而不是偏小的值 acc series.rolling(scale, min_periodsscale).sum() out pd.Series(indexseries.index, dtypefloat) for month in range(1, 13): m_mask acc.index.month month # 2) 只用基准期数据拟合避免把未来信息带进标签 base acc[m_mask (acc.index.year baseline_end)].dropna() if len(base) 30: # 样本太少分布拟合不可信 continue p_zero float((base 0).mean()) # 零降水的经验概率 pos base[base 0] a, loc, b gamma.fit(pos, floc0) # 固定下界为 0 x acc[m_mask].dropna() cdf p_zero (1 - p_zero) * gamma.cdf(x, a, locloc, scaleb) cdf cdf.clip(1e-6, 1 - 1e-6) # 防止 ppf 出 ±inf out.loc[x.index] norm.ppf(cdf) return out逐段解释。滚动累积用min_periodsscale意思是窗口不满就留空否则序列开头会拿一两个月的降水当成三个月的累积算出来的 SPI 偏高这部分假数据会污染训练集。按month分组是必须的降水的分布在一年里随季节变化把所有月份混在一起拟合等于假设七月和一月服从同一个分布。基准期用baseline_end截断是整个项目里最重要的一行防御。如果拿全样本含未来年份去拟合 Gamma 参数那么 2010 年的 SPI 值里就掺进了 2019 年才有的信息模型在验证集上会显得特别准上线后立刻掉下来。这类数据泄漏在时序任务里极其隐蔽因为它不在特征里而在标签的构造过程里。p_zero处理的是干旱区的零降水问题Gamma 分布要求正值碰上连续无降水的月份直接拟合会失败或者给出离谱的参数。用「零值经验概率 正值部分条件分布」的混合写法是文献里的常规做法工程上够用。最后clip一下再ppf是因为norm.ppf(0)和norm.ppf(1)会返回负无穷和正无穷一个inf进了特征表后面所有标准化和树模型都会报错。算完 SPI 之后切标签就是查表def spi_to_class(spi: pd.Series) - pd.Series: # 按 SPI 阈值切干旱等级-0.5 以下才算进入干旱区间 bins [-np.inf, -2.0, -1.5, -1.0, -0.5, np.inf] labels [extreme, severe, moderate, mild, normal] return pd.cut(spi, binsbins, labelslabels, rightFalse)rightFalse表示左闭右开保证 -1.0 只会落进一个桶里。pd.cut返回的是有序分类变量后面喂给 sklearn 前记得.cat.codes或者直接当字符串处理两种都行但要保持一致。4.3 特征表滞后项、季节编码与土壤墒情到了建模这一步特征的设计比模型选型重要得多。干旱是有记忆的前几个月的降水和土壤含水量直接决定了下个月会不会继续恶化所以滞后项必须做。def build_features(df: pd.DataFrame, spi: pd.Series, windows(1, 3, 6, 12)) - pd.DataFrame: feat pd.DataFrame(indexdf.index) # 1) 多尺度降水累积不同窗口对应气象/农业/水文干旱 for w in windows: feat[ftp_sum_{w}] df[tp].rolling(w, min_periodsw).sum() feat[fpet_sum_{w}] df[pet].rolling(w, min_periodsw).sum() feat[ft2m_mean_{w}] df[t2m].rolling(w, min_periodsw).mean() # 2) 水分收支比单看降水更接近农业干旱 feat[water_bal_3] feat[tp_sum_3] - feat[pet_sum_3] feat[swvl1] df[swvl1] # 3) 季节编码用正余弦避免 12 月和 1 月被当成距离最远 month df.index.month.values feat[month_sin] np.sin(2 * np.pi * month / 12) feat[month_cos] np.cos(2 * np.pi * month / 12) # 4) SPI 滞后项这是模型的记忆全部来自过去不算泄漏 for lag in (1, 2, 3, 6): feat[fspi_lag{lag}] spi.shift(lag) return feat四个板块各有理由。多尺度累积对应不同的干旱类型1 个月窗口抓到的是短时急旱6 到 12 个月窗口对应的是缓慢累积的水文干旱模型能自己学出哪个月份该看哪个尺度。water_bal_3是降水减蒸散的简化水分收支比单看降水更贴近实际墒情变化。季节编码用正余弦而不是 one-hot是因为月份本身有循环结构。one-hot 会让模型认为十二月和一月毫无关系而正余弦编码在二维平面上把它们放回相邻位置树模型虽然不直接理解距离但在分裂时能更平滑地处理季节过渡。滞后项这里要划清界限spi_lag1是上个月的 SPI属于历史信息用它合理但绝不能把spi当期值放进特征那等于把答案抄进题干。判断标准很简单——这条特征在预测时刻能不能拿到。拿不到就不能进没有例外。4.4 时间序列交叉验证与评估口径建模这段最容易犯错的地方不是模型是切分。时序数据绝不能随机train_test_split因为随机切分会让模型在训练时见过 2018 年 3 月、测试时预测 2018 年 4 月相邻月份高度相关分数虚高得离谱。from sklearn.ensemble import RandomForestClassifier from sklearn.pipeline import Pipeline from sklearn.preprocessing import StandardScaler from sklearn.impute import SimpleImputer from sklearn.model_selection import TimeSeriesSplit from sklearn.metrics import f1_score, cohen_kappa_score, confusion_matrix X feat.dropna() # 前 12 行因窗口不足天然为空直接丢 y label.loc[X.index].cat.codes pipe Pipeline([ (imputer, SimpleImputer(strategymedian)), # 少量缺测用中位数兜底 (scaler, StandardScaler()), # 线性基线需要树模型可留 (clf, RandomForestClassifier( n_estimators400, min_samples_leaf3, # 样本少叶子别切太细 class_weightbalanced, # 干旱样本天然稀少 random_state42, )), ]) tscv TimeSeriesSplit(n_splits5, gap0) scores [] for tr, te in tscv.split(X): pipe.fit(X.iloc[tr], y.iloc[tr]) pred pipe.predict(X.iloc[te]) scores.append(f1_score(y.iloc[te], pred, averagemacro)) print(confusion_matrix(y.iloc[te], pred)) print(macro-F1 各折:, np.round(scores, 3), 均值:, np.mean(scores).round(3))参数逐个说。TimeSeriesSplit(n_splits5)是滚动前进式切分第一折用前 20% 训练、紧接着的一段测试第二折把测试段并入训练再往后推模拟的正是「用过去预测未来」的真实场景。gap0在这里够用因为标签是下一期而不是当期如果你改成预测未来第 3 个月就要把gap设成 2把重叠部分掐掉。min_samples_leaf3是针对小样本的保守设置叶子节点至少三个样本避免模型记住个别年份的噪声。class_weightbalanced必须加正常月份的样本远多于重旱月份不加权的话模型全预测「正常」也能拿到很高的准确率而准确率在这个场景下毫无意义。评估指标别只看 accuracy。averagemacro的 F1 对每个类别一视同仁重旱这一类预测得差会立刻把分数拉下来这正是我们关心的。再补一个 Cohens kappa它扣掉了随机猜测的成分两个指标一起看能判断模型是不是真的学到了东西。最后打印混淆矩阵重点看「重旱被判成正常」这一格——这类错误在业务上代价最高哪怕总体 F1 好看这一格数字大也得回去改特征。还应该跑一个持续性基线直接拿本月的 SPI 等级当作下个月的预测。如果随机森林打不过这个基线说明特征里根本没提供超出「昨天怎样今天还怎样」的信息得回头补土壤墒情或者更长的滞后窗口。5. 避坑与排查五条把 Notebook 跑崩的实录5.1 降水单位差了一千倍SPI 曲线却看着挺正常现象SPI 时序图整体形态像模像样但所有值挤在 ±0.3 之间几乎没有干旱等级被触发。原因tp没从米换算成毫米。SPI 是标准化指标量纲整体缩放后分布形状不变标准化之后仍然落在合理范围内所以图看着「正常」但阈值切分全部失效。解决在预处理里固定一行断言assert df[tp].mean() 500月降水的量级不该超过几百毫米。凡是涉及单位换算的字段都在落盘前加一条量级检查比事后看图的效率高得多。5.2 潜在蒸发符号搞反水分子减变成了相加现象换成 SPEI 之后干旱样本比例骤降湿润年份反而被判成重旱。原因ERA5 的通量约定是向上为负pev原始值大多是负数。直接tp - pev实际上是在做加法。解决在读取阶段就取反并转成正的毫米数同时在 Notebook 里画一张pet的月度箱线图正常应该是全正且夏季偏高。符号错这类问题不会报错只会让结论悄悄反掉属于最该防的一类。5.3 用全样本拟合 Gamma验证分数虚高现象交叉验证 macro-F1 高得超出预期但把模型拿去预测最后两年表现断崖式下跌。原因SPI 在计算时用了整个时间跨度拟合分布参数等于标签里混进了未来信息属于典型的数据泄漏。解决把基准期固定成一个明确的时间段比如截到 2010 年之后所有年份都用这套参数做变换。真要严谨就在滚动验证的每一折里只用该折训练段的数据重新拟合代价是计算时间翻几倍但结论可靠。5.4 滚动窗口不满就出值序列开头一片假干旱现象SPI 序列最前面一两年异常偏低模型在这些年份上误差明显更大。原因rolling(12)没设min_periods前 11 个月拿不足一年的降水当整年累积累积值偏小标准化后自然偏负。解决所有滚动操作统一写上min_periodswindow让不足窗口的位置留 NaN然后在建模前dropna()。丢掉的这十几行对总样本量影响很小但能避免一整段被污染的标签。5.5 一次 open_dataset 把内存吃满现象Notebook 跑到特征工程那格突然变慢内存监控一路飙红最后内核重启。原因xr.open_dataset默认惰性加载但一旦执行.values或者转 DataFrame整个四维场就被读进内存几十年的全球格点轻松几个 G。解决切片要在.values之前做完顺序是先.sel(time..., latitude..., longitude...)再取值。区域大的话用.chunk({time: 120})开 dask 分块或者干脆在下载阶段就把区域裁小别指望在内存里做裁剪。6. 二次开发把一次性实验改成能换区域跑的工具Notebook 跑通只是起点。真正拿去用的时候你大概率要换区域、换时间窗、换模型对比这时候还在单元格里改坐标就太难受了。我的做法是把稳定逻辑全部抽进src/drought.pyNotebook 只留调用和可视化然后加一个config.yaml管参数region: {north: 45, west: 110, south: 30, east: 125} period: {start: 1990, end: 2019} baseline_end: 2010 spi_windows: [1, 3, 6, 12] model: {type: random_forest, n_estimators: 400, min_samples_leaf: 3}配套写一个run.py读配置、调load_point/compute_spi/build_features、训练、把结果和图分别写进outputs/{region}_{timestamp}/。这样一次完整实验就是一条命令参数组合也能批量跑。改造时我最推荐先做的一件事是缓存中间结果。compute_spi里的 Gamma 拟合在多年份、多格点循环下很慢把算好的 SPI 和特征表存成 parquet键用「区域 窗口 基准期」拼出来重跑时命中缓存就直接读。我一开始嫌麻烦没做后来调特征的时候每次都要等几分钟血泪经验。验证方法上别只用交叉验证的均值。留出一整段连续年份比如最后五年完全不参与任何调参最后只跑一次看混淆矩阵里「重旱被判成正常」那一格的数量。这个数才是业务上真正关心的也比任何单一指标更能说明模型能不能用。想再进一步可以把标签从分类换成回归直接预测 SPI 数值再按阈值切等级好处是保留了「旱到什么程度」的连续信息代价是要重新设计损失函数和评估口径。还有一个具体技巧换区域时不要只看指标先画两张图——该区域的 SPI-3 时间序列以及模型预测等级和真实等级的堆叠对比。图比数字诚实区域性气候差异比如季风区和高原区在数字上可能表现差不多在图上却一眼能看出模型是不是只学会了复读上个月。从那以后我每次动特征表之前都会先把当前版本的指标、混淆矩阵和两张图存进outputs/打个快照再开始改。没有基线做对照的调参本质上是在赌运气。希望帮到你。本文还有配套的精品资源点击获取