当前位置：首页 > news >正文

neural_gcm模型进行气象预测教程

news 来源：原创 2025/5/10 4:49:53

简介

NeuralGCM （General circulation models）是Google开发的一种新型的天气和气候模型，它结合了传统的物理建模与机器学习技术，相关论文于2024年发表在nature上。它既可以做短期的天气预报，也可以做长期的气候预测。

安装以及初始化：

我的环境是在windows11上的WSL 2.0 Ubuntu虚拟机，，显卡为4070，cuda版本为12.6，conda里的python版本为3.10，安装步骤如下：

conda create --name neural_gcm_env python=3.10
pip install neuralgcm
pip install gcsfs
pip install "jax[cuda126]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

安装完毕后进行环境测试：

（1）测试gcsfs

import gcsfs
import jax
import numpy as np
import pickle
import xarrayfrom dinosaur import horizontal_interpolation
from dinosaur import spherical_harmonic
from dinosaur import xarray_utils
import neuralgcmgcs = gcsfs.GCSFileSystem(token='anon')
print("finish test")

（2）测试jax能否正常使用GPU

import jax
import jax.numpy as jnp# 检查 JAX 是否能找到 GPU 设备
devices = jax.devices()
print(f"JAX devices: {devices}")# 检查 GPU 是否为默认设备类型
default_device = jax.default_backend()
print(f"JAX default backend: {default_device}")# 尝试在 GPU 上运行一些简单的计算
key = jax.random.PRNGKey(0)
x = jax.random.normal(key, (1000, 1000))
y = x @ x
print(f"Result on device: {y.device}")# 打印 JAX 的配置信息
print(jax.config.values)

预测

下载模型和数据

进入网站（https://neuralgcm.readthedocs.io/en/latest/checkpoints.html），下载预训练模型

根据官方推荐，用于天气预报我们使用1.4° stochastic model(44M大小)，下载后将模型文件放在当前目录下。然后写代码测试：

import gcsfs
import jax
import numpy as np
import pickle
import xarrayfrom dinosaur import horizontal_interpolation
from dinosaur import spherical_harmonic
from dinosaur import xarray_utils
import neuralgcmgcs = gcsfs.GCSFileSystem(token='anon')model_name = 'models_v1_stochastic_1_4_deg.pkl'  #@param ['v1/deterministic_0_7_deg.pkl', 'v1/deterministic_1_4_deg.pkl', 'v1/deterministic_2_8_deg.pkl', 'v1/stochastic_1_4_deg.pkl', 'v1_precip/stochastic_precip_2_8_deg.pkl', 'v1_precip/stochastic_evap_2_8_deg'] {type: "string"}with open(f'/home/ruiduobao/research/Neural_GCM_work/models_v1_stochastic_1_4_deg.pkl', 'rb') as f:ckpt = pickle.load(f)model = neuralgcm.PressureLevelModel.from_checkpoint(ckpt)

然后从google cloud下载ERA5数据来进行预测（这里也可以换成自己下载的era5数据，由于只是测试，就使用了GOOGLE的数据库匿名下载，就不进行官网下载了）：

era5_path = 'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3'
full_era5 = xarray.open_zarr(gcs.get_mapper(era5_path), chunks=None)demo_start_time = '2020-02-14'
demo_end_time = '2020-02-18'
data_inner_steps = 24  # process every 24th hoursliced_era5 = (full_era5[model.input_variables + model.forcing_variables].pipe(xarray_utils.selective_temporal_shift,variables=model.forcing_variables,time_shift='24 hours',).sel(time=slice(demo_start_time, demo_end_time, data_inner_steps)).compute()
)

对齐网格

然后再进行网格对齐，代码为：

era5_grid = spherical_harmonic.Grid(latitude_nodes=full_era5.sizes['latitude'],longitude_nodes=full_era5.sizes['longitude'],latitude_spacing=xarray_utils.infer_latitude_spacing(full_era5.latitude),longitude_offset=xarray_utils.infer_longitude_offset(full_era5.longitude),
)
regridder = horizontal_interpolation.ConservativeRegridder(era5_grid, model.data_coords.horizontal, skipna=True
)
eval_era5 = xarray_utils.regrid(sliced_era5, regridder)
eval_era5 = xarray_utils.fill_nan_with_nearest(eval_era5)

预测

inner_steps = 1  # save model outputs once every 24 hours
outer_steps = 4 * 24 // inner_steps  # total of 4 days
timedelta = np.timedelta64(1, 'h') * inner_steps
times = (np.arange(outer_steps) * inner_steps)  # time axis in hours# initialize model state
inputs = model.inputs_from_xarray(eval_era5.isel(time=0))
input_forcings = model.forcings_from_xarray(eval_era5.isel(time=0))
rng_key = jax.random.key(42)  # optional for deterministic models
initial_state = model.encode(inputs, input_forcings, rng_key)# use persistence for forcing variables (SST and sea ice cover)
all_forcings = model.forcings_from_xarray(eval_era5.head(time=1))# make forecast
final_state, predictions = model.unroll(initial_state,all_forcings,steps=outer_steps,timedelta=timedelta,start_with_input=True,
)
predictions_ds = model.data_to_xarray(predictions, times=times)

可以看到，几乎把我的显存吃满了，速度也非常快。outer_steps = 2 * 24 代表了从2月14日预测未来2天 (2* 24 小时) 的气象数据，inner_steps=1代表1小时输出一次结果。几分钟的时间就出结果了。

然后再看看温度的预测效果：

完整代码为：

import gcsfs
import jax
import numpy as np
import pickle
import xarray
from tqdm import tqdm
from dinosaur import horizontal_interpolation
from dinosaur import spherical_harmonic
from dinosaur import xarray_utils
import neuralgcm#Google cloud匿名访问
gcs = gcsfs.GCSFileSystem(token='anon')#模型加载
model_name = 'models_v1_stochastic_1_4_deg.pkl'  #@param ['v1/deterministic_0_7_deg.pkl', 'v1/deterministic_1_4_deg.pkl', 'v1/deterministic_2_8_deg.pkl', 'v1/stochastic_1_4_deg.pkl', 'v1_precip/stochastic_precip_2_8_deg.pkl', 'v1_precip/stochastic_evap_2_8_deg'] {type: "string"}with open(f'/home/ruiduobao/research/Neural_GCM_work/models_v1_stochastic_1_4_deg.pkl', 'rb') as f:ckpt = pickle.load(f)model = neuralgcm.PressureLevelModel.from_checkpoint(ckpt)#下载气象数据
era5_path = 'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3'
full_era5 = xarray.open_zarr(gcs.get_mapper(era5_path), chunks=None)demo_start_time = '2022-02-05'
demo_end_time = '2022-02-07'
data_inner_steps = 24  # process every 24th hoursliced_era5 = (full_era5[model.input_variables + model.forcing_variables].pipe(xarray_utils.selective_temporal_shift,variables=model.forcing_variables,time_shift='24 hours',).sel(time=slice(demo_start_time, demo_end_time, data_inner_steps)).compute()
)
print("over")#网格对齐
era5_grid = spherical_harmonic.Grid(latitude_nodes=full_era5.sizes['latitude'],longitude_nodes=full_era5.sizes['longitude'],latitude_spacing=xarray_utils.infer_latitude_spacing(full_era5.latitude),longitude_offset=xarray_utils.infer_longitude_offset(full_era5.longitude),
)
regridder = horizontal_interpolation.ConservativeRegridder(era5_grid, model.data_coords.horizontal, skipna=True
)
eval_era5 = xarray_utils.regrid(sliced_era5, regridder)
eval_era5 = xarray_utils.fill_nan_with_nearest(eval_era5)#预测
inner_steps = 1  # save model outputs once every 24 hours
outer_steps = 2 * 24 // inner_steps  # total of 4 days
timedelta = np.timedelta64(1, 'h') * inner_steps
times = (np.arange(outer_steps) * inner_steps)  # time axis in hours# initialize model state
inputs = model.inputs_from_xarray(eval_era5.isel(time=0))
input_forcings = model.forcings_from_xarray(eval_era5.isel(time=0))
rng_key = jax.random.key(42)  # optional for deterministic models
initial_state = model.encode(inputs, input_forcings, rng_key)# use persistence for forcing variables (SST and sea ice cover)
all_forcings = model.forcings_from_xarray(eval_era5.head(time=1))# make forecast
final_state, predictions = model.unroll(initial_state,all_forcings,steps=outer_steps,timedelta=timedelta,start_with_input=True,
)
predictions_ds = model.data_to_xarray(predictions, times=times)#保存数据
# 保存预测结果到本地文件
output_path = '/home/ruiduobao/research/Neural_GCM_work/predictions.nc'  # 指定保存路径
predictions_ds.to_netcdf(output_path)print(f"Predictions saved to {output_path}")

总结

NeuralGCM是一个精度非常高的AI气象模型（参考论文），

安装简单（三个pip就可以安装成功了），

上手容易（对比传统数值模型），

效率对比传统模型（比如WRF）快得多。

‍

参考

Kochkov D, Yuval J, Langmore I, et al. Neural general circulation models for weather and climate[J]. Nature, 2024, 632(8027): 1060-1066.

neuralgcm.https://github.com/neuralgcm/neuralgcm

neuralgcm文档.https://neuralgcm.readthedocs.io/en/latest/inference_demo.html

和鲸社区.https://www.heywhale.com/mw/project/67b0457f2248a2c0ca311a31?run=

简介

安装以及初始化：

预测

下载模型和数据

对齐网格

预测

总结

参考

相关文章：