earth2mip
v0.1.0
Earth-2模型对比项目（MIP）是一个基于Python的AI框架，使气候研究人员和科学家能够探索和试验使用AI模型在天气和气候中。它提供了参考工作流程，以了解AI模型如何捕获地球大气的物理学以及它们如何与传统的数值天气预测模型一起工作。例如，该存储库提供了使用预训练的模型检查点运行推理的统一接口，并使用某些标准指标对此类模型的技能进行评分。该存储库的目的是促进天气和气候社区提出良好的事件参考基线，以测试模型，并与各种数据源一起使用。
一般版本后，Earth-2 MIP将在PYPI上安装。同时，可以从来源安装：
git clone [email protected]:NVIDIA/earth2mip.git
cd earth2mip && pip install .有关更多详细信息和其他选项，请参见安装文档。
Earth-2 MIP提供了一组示例，可以在示例文档页面上查看，可用于使用各种工作流程开始。这些示例可以作为Jupyer笔记本和Python脚本下载。源python脚本可以在示例文件夹中找到。
Earth-2 MIP提供了用于使用AI模型运行推理的高级API。例如，以下可用于使用气候数据存储（CDS）的初始状态运行Pangu天气：
python
>>> import datetime
>>> from earth2mip.networks import get_model
>>> from earth2mip.initial_conditions import cds
>>> from earth2mip.inference_ensemble import run_basic_inference
>>> time_loop = get_model( " e2mip://dlwp " , device= " cuda:0 " )
>>> data_source = cds.DataSource(time_loop.in_channel_names)
>>> ds = run_basic_inference(time_loop, n=10, data_source=data_source, time=datetime.datetime(2018, 1, 1))
>>> ds.chunk ()
< xarray.DataArray (time: 11, history: 1, channel: 69, lat: 721, lon: 1440) >
dask.array < xarray- < this-array > , shape=(11, 1, 69, 721, 1440), dtype=float32, chunksize=(11, 1, 69, 721, 1440), chunktype=numpy.ndarray >
Coordinates:
* lon (lon) float32 0.0 0.25 0.5 0.75 1.0 ... 359.0 359.2 359.5 359.8
* lat (lat) float32 90.0 89.75 89.5 89.25 ... -89.25 -89.5 -89.75 -90.0
* time (time) datetime64[ns] 2018-01-01 ... 2018-01-03T12:00:00
* channel (channel) < U5 ' z1000 ' ' z925 ' ' z850 ' ' z700 ' ... ' u10m ' ' v10m ' ' t2m '
Dimensions without coordinates: history您可以像这样获得ACC/RMSE：
>>> from earth2mip.inference_medium_range import score_deterministic
>>> import numpy as np
>>> scores = score_deterministic(time_loop,
data_source=data_source,
n=10,
initial_times=[datetime.datetime(2018, 1, 1)],
# fill in zeros for time-mean, will typically be grabbed from data.
time_mean=np.zeros((7, 721, 1440))
)
>>> scores
<xarray.Dataset>
Dimensions: (lead_time: 11, channel: 7, initial_time: 1)
Coordinates:
* lead_time (lead_time) timedelta64[ns] 0 days 00:00:00 ... 5 days 00:...
* channel (channel) <U5 't850' 'z1000' 'z700' ... 'z300' 'tcwv' 't2m'
Dimensions without coordinates: initial_time
Data variables:
acc (lead_time, channel) float64 1.0 1.0 1.0 ... 0.9686 0.9999
rmse (lead_time, channel) float64 0.0 2.469e-05 0.0 ... 7.07 2.998
initial_times (initial_time) datetime64[ns] 2018-01-01
>>> scores.rmse.sel(channel='z500')
<xarray.DataArray 'rmse' (lead_time: 11)>
array([ 0. , 150.83014446, 212.07880612, 304.98592282,
381.36510987, 453.31516952, 506.01464974, 537.11092269,
564.79603347, 557.22871627, 586.44691243])
Coordinates:
* lead_time (lead_time) timedelta64[ns] 0 days 00:00:00 ... 5 days 00:00:00
channel <U5 'z500'
这些笔记本用几个型号说明了如何使用的方法，只要兼容，它就可以作为引用您自己的检查站的参考。可能还有其他工作使其与Earth-2 MIP兼容。 Earth-2 MIP在模量中利用模型动物园来提供一组基础线模型。目的是使社区能够发展此模型集合，如下表所示。
| ID | 模型 | 建筑学 | 类型 | 参考 | 来源 | 尺寸 |
|---|---|---|---|---|---|---|
| FCN | 四castnet | 自适应傅里叶神经操作员 | 全球天气 | arxiv | 模量 | 300MB |
| DLWP | 深度学习天气预测 | 卷积编码器解码器 | 全球天气 | agu | 模量 | 50MB |
| pangu | Pangu天气（分层6 + 24小时） | 视觉变压器 | 全球天气 | 自然 | onnx | 2GB |
| pangu_6 | Pangu天气6小时型号 | 视觉变压器 | 全球天气 | 自然 | onnx | 1GB |
| Pangu_24 | Pangu天气24小时型号 | 视觉变压器 | 全球天气 | 自然 | onnx | 1GB |
| fcnv2_sm | Fourcastnet V2 | 球形谐波傅里叶神经操作员 | 全球天气 | arxiv | 模量 | 3.5GB |
| Graphcast | Graphcast，37级，0.25度 | 图神经网络 | 全球天气 | 科学 | github | 145MB |
| Graphcast_small | Graphcast，13级，1度 | 图神经网络 | 全球天气 | 科学 | github | 144MB |
| graphcast_operational | Graphcast，13级，0.25度 | 图神经网络 | 全球天气 | 科学 | github | 144MB |
| strad_afno | 四castnet降水 | 自适应傅里叶神经操作员 | 诊断 | arxiv | 模量 | 300MB |
| 气候 | 气候分割模型 | 卷积神经网络 | 诊断 | GMD | 模量 | 2MB |
* =即将到来
一些模型需要默认情况下未安装的其他依赖项。有关详细信息，请参阅安装说明。
注意：每个模型检查点可能都有其独特的许可。我们鼓励用户熟悉每个人，以了解对其特定用例的影响。
我们希望将您的模型集成到记分牌中，以展示社区！最好的方法是通过NVIDIA模量。您可以贡献您的模型（培训代码和模型检查点），我们可以确保将其作为参考集的一部分进行维护。
Earth-2 MIP是一种开源协作，其成功源于社区对进一步领域的贡献。感谢您为该项目做出贡献，以便其他人可以基于您的贡献。有关为Earth-2 MIP做出贡献的指南，请参考贡献指南。
这项工作的灵感是为了促进NVIDIA的团队之间的类似交往 - ML专家开发了新模型和气候科学领域的领域专家，评估了此类模型的技能。例如，通常必要的输入数据（例如归一化常数和超参数值）与模型权重并不包装。每个模型通常都会实现一个略有不同的接口。评分例程是特定于要评分的模型，并且在各组之间可能不一致。
Earth-2 MIP解决了这些挑战，并弥合了最经常评估ML模型的领域专家与ML专家生产它们之间的差距。与该领域的其他项目相比，Earth-2 MIP专注于即时评分模型。它具有适用于Jupyter书籍的Python API，适用于在许多GPU上分发的CLIS CLIS，以及一个灵活的插件框架，可让任何人使用自己的ML模型。更重要的是，Earth-2 MIP渴望促进气候研究界内的探索与协作，以评估AI模型在气候和天气模拟中的潜力。
请参阅文档页面，以获取有关Earth-2 MIP，功能，API等的深入信息。
Apache许可证2.0提供了Earth-2 MIP，请参见License.txt，以获取完整的许可文本。
