30分钟上手AIFS ENS v2.0Jupyter Notebook实操教程【免费下载链接】aifs-ens-2.0项目地址: https://ai.gitcode.com/hf_mirrors/ecmwf/aifs-ens-2.0AIFS ENS v2.0是一款基于AI的天气预报模型能够利用ECMWF开放数据和先进的深度学习技术生成高精度的气象预测。本教程将带你通过Jupyter Notebook快速掌握AIFS ENS v2.0的安装配置和基础使用即使是没有深厚气象学背景的新手也能在半小时内完成你的第一次天气预报实验。 准备工作环境要求与依赖安装系统要求AIFS ENS v2.0对运行环境有一定要求确保你的系统满足以下条件Python版本3.11或3.12硬件要求Ampere架构及以上的GPU已在Colab L4/A100和ECMWF ATOS HPCF AC集群测试通过磁盘空间至少需要数GB的可用空间快速安装依赖根据你的运行环境选择合适的安装命令。如果使用ColabL4或A100运行时可以直接运行以下命令!pip install anemoi-inference[huggingface]0.8.3 anemoi-models0.11.2 anemoi-utils0.4.35.post3 !pip install -q torch2.7.0 torch-geometric2.6.1 !pip install -q earthkit-regrid0.5.1 ecmwf-opendata0.3.29 earthkit-data1.0.0 !pip install -q flash-attn2.7.4.post1如果你在ECMWF的ATOS HPCF通过JupyterHub访问GPU资源选择ECMWF ATOS HPCF - GPU partition后运行%pip install anemoi-inference[huggingface]0.8.3 anemoi-models0.11.2 anemoi-utils0.4.35.post3 %pip install -q torch2.7.0 torch-geometric2.6.1 %pip install -q earthkit-regrid0.5.1 ecmwf-opendata0.3.29 earthkit-data1.0.0 %ip install -q --force-reinstall --no-deps flash-attn https://github.com/cathalobrien/get-flash-attn/releases/download/v0.1-alpha/flash_attn-2.8.3cu12torch2.7cxx11abiFALSE-cp312-cp312-linux_x86_64.whl环境验证安装完成后运行以下代码验证环境是否配置正确import platform import torch print(Platform:, platform.platform()) print(Torch:, torch.__version__) print(CUDA available:, torch.cuda.is_available()) if not torch.cuda.is_available(): raise RuntimeError(CUDA GPU not available. This notebook requires a GPU.) gpu_name torch.cuda.get_device_name(0) major, minor torch.cuda.get_device_capability(0) print(GPU:, gpu_name) print(fCompute capability: {major}.{minor}) if major 8: raise RuntimeError(FlashAttention requires Ampere GPUs or newer) import flash_attn print(FlashAttention:, flash_attn.__version__) 数据准备获取初始条件导入必要库首先导入运行所需的Python库import datetime from collections import defaultdict import numpy as np import earthkit.data as ekd import earthkit.regrid as ekr from anemoi.inference.runners.simple import SimpleRunner from anemoi.inference.outputs.printer import print_state from ecmwf.opendata import Client as OpendataClient # 设置数据缓存策略 ekd.config.set({cache-policy: user})定义参数列表AIFS ENS v2.0需要多种气象参数作为输入定义要从ECMWF开放数据获取的参数列表PARAM_SFC [10u, 10v, 2d, 2t, msl, skt, sp, tcw, sd] PARAM_SFC_FC [lsm, z, slor, sdor] PARAM_SOIL [vsw,sot] PARAM_WAVE [wmb, h1012, h1214, h1417, h1721, h2125, h2530, mwd, cdww, mwp, swh] PARAM_PL [gh, t, u, v, w, q] LEVELS [1000, 925, 850, 700, 600, 500, 400, 300, 250, 200, 150, 100, 50, 10] SOIL_LEVELS [1,2]获取最新可用数据设置数据源并获取最新可用的日期SOURCE ecmwf # 其他选项: azure, aws, ecmwf 或 google DATE OpendataClient(SOURCE).latest() print(Initial date is, DATE)下载并处理数据创建数据下载和插值函数从ECMWF开放数据API获取并预处理初始条件def get_open_data(param, levelist[], number None, **kwargs): fields defaultdict(list) # 获取当前日期和前一天的数据 for date in [DATE - datetime.timedelta(hours6), DATE]: if number is None: data ekd.from_source(ecmwf-open-data, datedate, paramparam, levelistlevelist, source SOURCE, **kwargs) else: kwargs.setdefault(stream, enfo) data ekd.from_source(ecmwf-open-data, datedate, paramparam, levelistlevelist, number[number], source SOURCE, **kwargs) for f in data: # 将数据从-180~180转换为0~360 assert f.to_numpy().shape (721,1440) values np.roll(f.to_numpy(), -f.shape[1] // 2, axis1) # 插值到N320网格 values ekr.interpolate(values, {grid: (0.25, 0.25)}, {grid: N320}) # 添加到字段列表 name f{f.metadata(param)}_{f.metadata(levelist)} if levelist else f.metadata(param) fields[name].append(values) # 为每个参数创建单一矩阵 for param, values in fields.items(): fields[param] np.stack(values) return fields # 初始化字段字典和集合数 fields {} number 1 # 设置为1-50可获取不同的集合初始条件现在依次下载各类数据# 下载表面字段 fields.update(get_open_data(paramPARAM_SFC, numbernumber, levtypesfc)) fields.update(get_open_data(paramPARAM_SFC_FC, levtypesfc)) assert all(p in fields for p in [*PARAM_SFC, *PARAM_SFC_FC]), Missing parameters: %s % (set(PARAM_SFC) set(PARAM_SFC_FC) - set(fields.keys())) # 下载波浪字段 fields.update(get_open_data(paramPARAM_WAVE, streamwave if not number else waef, numbernumber)) assert all(p in fields for p in PARAM_WAVE), Missing parameters: %s % (set(PARAM_WAVE) - set(fields.keys())) # 下载土壤字段 soilget_open_data(paramPARAM_SOIL,levelistSOIL_LEVELS, numbernumber) soil_names [f{p}_{lev} for p in PARAM_SOIL for lev in SOIL_LEVELS] assert all(p in soil for p in soil_names), Missing parameters: %s % (set(soil_names) - set(soil.keys())) # 下载压力水平字段 fields.update(get_open_data(paramPARAM_PL, levelistLEVELS, numbernumber)) PRESSURE_NAMES [f{p}_{lev} for p in PARAM_PL for lev in LEVELS] assert all(p in fields for p in PRESSURE_NAMES), Missing parameters: %s % (set(PRESSURE_NAMES) - set(fields.keys()))数据转换对下载的数据进行必要的转换使其符合模型输入要求# 将平均波向转换为正弦和余弦分量 mwd fields.pop(mwd) mwd_rad np.deg2rad(mwd) fields[cos_mwd] np.cos(mwd_rad) fields[sin_mwd] np.sin(mwd_rad) # 重命名土壤字段 mapping {sot_1: stl1, sot_2: stl2, vsw_1: swvl1,vsw_2: swvl2} for k,v in soil.items(): fields[mapping[k]]v # 移除未使用的特定湿度字段 fields.pop(q_10, None) # 应用陆海掩码 mask np.equal(ekd.from_source(file, lsm.grib)[0].to_numpy(flattenTrue), 0) fields[sd][:, mask] np.nan fields[swvl1][:, mask] np.nan fields[swvl2][:,mask] np.nan # 将位势高度转换为位势 for level in LEVELS: gh fields.pop(fgh_{level}) fields[fz_{level}] gh * 9.80665 模型加载与预测创建初始状态准备模型输入状态input_state dict(dateDATE, fieldsfields)加载模型从Hugging Face加载预训练模型 checkpointcheckpoint {huggingface:ecmwf/aifs-ens-2.0} # 创建运行器 runner SimpleRunner(checkpoint)内存优化提示如果遇到内存问题可以设置以下环境变量import os os.environ[PYTORCH_CUDA_ALLOC_CONF]expandable_segments:True os.environ[ANEMOI_INFERENCE_NUM_CHUNKS]16运行预测生成12小时的天气预报可通过修改LEAD_TIME调整预测长度LEAD_TIME 12 # 预测时长小时 states [] for state in runner.run(input_stateinput_state, lead_timeLEAD_TIME): states.append(state) print_state(state)运行后会输出类似以下的结果 date2026-04-23T12:00:00 latitudes(542080,) longitudes(542080,) fields120 100u shape(542080,) min-26.7593 max28.8329 q_700 shape(542080,) min0 max0.01222 t_925 shape(542080,) min231.897 max310.795 v_850 shape(542080,) min-40.2831 max29.4611 z_925 shape(542080,) min1964.48 max9248.86 date2026-04-23T18:00:00 latitudes(542080,) longitudes(542080,) fields120 100u shape(542080,) min-28.3416 max30.0988 q_700 shape(542080,) min0 max0.0121639 t_925 shape(542080,) min234.108 max310.616 v_850 shape(542080,) min-37.9228 max31.7181 z_925 shape(542080,) min774.796 max9244.69 结果可视化可选安装可视化依赖!pip install -q matplotlib cartopy绘制100u100米风速预报结果import matplotlib.pyplot as plt import cartopy.crs as ccrs import cartopy.feature as cfeature import matplotlib.tri as tri def fix(lons): # 将经度从0-360转换为-180-180 return np.where(lons 180, lons - 360, lons) latitudes state[latitudes] longitudes state[longitudes] values state[fields][100u] fig, ax plt.subplots(figsize(11, 6), subplot_kw{projection: ccrs.PlateCarree()}) ax.coastlines() ax.add_feature(cfeature.BORDERS, linestyle:) triangulation tri.Triangulation(fix(longitudes), latitudes) contourax.tricontourf(triangulation, values, levels20, transformccrs.PlateCarree(), cmapRdBu) cbar fig.colorbar(contour, axax, orientationvertical, shrink0.7, label100u) plt.title(100m winds (100u) at {}.format(state[date])) plt.show()⚠️ 注意事项GPU非确定性GPU操作可能不具有位确定性重复运行可能产生略有不同的数值结果。如需强制确定性可设置import torch torch.backends.cudnn.benchmark False torch.backends.cudnn.deterministic True torch.use_deterministic_algorithms(True)⚠️ 启用确定性会显著增加运行时间输入数据差异本教程使用的ECMWF开放数据与ECMWF业务系统使用的数据存在重投影路径差异可能导致预测结果略有不同。 项目文件说明run_AIFS_ENS_v2.0.ipynb主Jupyter Notebook文件包含完整的预测流程inference.yaml推理配置文件aifs-ens-crps-2.0.ckpt模型权重checkpoint仅包含模型权重不含优化器状态lsm.grib陆海掩码数据 模型训练历程AIFS ENS v2.0的训练历程展示了从ENS 1.0到ENS 2.0的优化过程该图表显示了ENS 1.0和ENS 2.0的微调时间线包括关键的套件周期如Suite cycle 50r1帮助理解模型的迭代优化过程。通过本教程你已经掌握了AIFS ENS v2.0的基本使用方法。要进一步深入可以尝试调整预测时长、修改输入参数或探索不同的可视化方式。如有问题可查阅项目文档或提交issue获取帮助。【免费下载链接】aifs-ens-2.0项目地址: https://ai.gitcode.com/hf_mirrors/ecmwf/aifs-ens-2.0创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考