earth2mip
v0.1.0
Le projet d'intercomparaison du modèle Earth-2 (MIP) est un cadre d'IA basé sur Python qui permet aux chercheurs et aux scientifiques du climat d'explorer et d'expérimenter l'utilisation de modèles d'IA pour la météo et le climat. Il fournit des workflows de référence pour comprendre comment les modèles d'IA capturent la physique de l'atmosphère terrestre et comment ils peuvent fonctionner avec des modèles de prévision météorologiques numériques traditionnels. Par exemple, le repo fournit une interface uniforme pour exécuter l'inférence en utilisant les points de contrôle du modèle pré-formé et la notation de la compétence de ces modèles en utilisant certaines mesures standard. Ce référentiel est destiné à faciliter la communauté météorologique et climatique pour proposer une bonne référence de référence des événements pour tester les modèles contre et utiliser avec une variété de sources de données.
Le MIP Earth-2 sera installable sur PYPI lors de la libération générale. En attendant, on peut installer à partir de la source:
git clone [email protected]:NVIDIA/earth2mip.git
cd earth2mip && pip install .Voir la documentation d'installation pour plus de détails et autres options.
Earth-2 MIP fournit un ensemble d'exemples qui peuvent être consultés sur la page de documentation des exemples qui peut être utilisé pour commencer avec divers workflows. Ces exemples peuvent être téléchargés à la fois sous forme de carnets Jupyer et de scripts Python. Les scripts Python source peuvent être trouvés dans les dossiers des exemples.
Le MIP Earth-2 fournit des API de haut niveau pour faire fonctionner l'inférence avec les modèles d'IA. Par exemple, les éléments suivants peuvent être utilisés pour exécuter la météo Pangu en utilisant un état initial du magasin de données climatiques (CDS):
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: historyEt vous pouvez obtenir ACC / RMSE comme ceci:
>>> 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'
Ces cahiers illustrent la manière d'utiliser avec quelques modèles et cela peut servir de référence pour apporter votre propre point de contrôle tant qu'il est compatible. Il peut y avoir des travaux supplémentaires pour le rendre compatible avec le MIP Earth-2. Le MIP Earth-2 exploite le zoo du modèle dans le module pour fournir un ensemble de référence de modèles de base. L'objectif est de permettre à la communauté de développer cette collection de modèles comme indiqué dans le tableau ci-dessous.
| IDENTIFIANT | Modèle | Architecture | Taper | Référence | Source | Taille |
|---|---|---|---|---|---|---|
| FCN | Falcastnet | Opérateur neuronal de Fourier adaptatif | météo mondial | Arxiv | module | 300 Mo |
| dlwp | Prédiction météorologique en profondeur | Coder convolutionnel | météo mondial | Agu | module | 50 Mo |
| pangu | Météo pangu (hiérarchique 6 + 24 heures) | Transformateur de vision | météo mondial | Nature | onnx | 2 Go |
| pangu_6 | Modèle 6hr de temps Pangu | Transformateur de vision | météo mondial | Nature | onnx | 1 Go |
| Pangu_24 | Modèle Pangu Weather 24hr | Transformateur de vision | météo mondial | Nature | onnx | 1 Go |
| fcnv2_sm | Fourcastnet V2 | Opérateur neuronal de Fourier des harmones sphériques | météo mondial | Arxiv | module | 3,5 Go |
| graphape | Graphcast, 37 niveaux, 0,25 degrés | Réseau de neurones graphiques | météo mondial | Science | github | 145 Mo |
| graphcast_small | Graphcast, 13 niveaux, 1 deg | Réseau de neurones graphiques | météo mondial | Science | github | 144 Mo |
| graphcast_opération | Graphcast, 13 niveaux, 0,25 degrés | Réseau de neurones graphiques | météo mondial | Science | github | 144 Mo |
| précipitation_afno | Précipitation Fourcastnet | Opérateur neuronal de Fourier adaptatif | diagnostique | Arxiv | module | 300 Mo |
| climatenet | Modèle de segmentation de climatenet | Réseau neuronal convolutionnel | diagnostique | DPI | module | 2 Mo |
* = à venir bientôt
Certains modèles nécessitent des dépendances supplémentaires non installées par défaut. Reportez-vous aux instructions d'installation pour plus de détails.
Remarque : chaque point de contrôle du modèle peut avoir sa propre licence unique. Nous encourageons les utilisateurs à se familiariser avec chacun pour comprendre les implications pour leur cas d'utilisation particulier.
Nous voulons intégrer votre modèle dans le tableau de bord pour montrer la communauté! La meilleure façon de le faire est via le module Nvidia. Vous pouvez contribuer votre modèle (à la fois le code de formation et le point de contrôle du modèle) et nous pouvons nous assurer qu'il est maintenu dans le cadre de l'ensemble de référence.
Earth-2 MIP est une collaboration open source et son succès est enracinée dans la contribution communautaire à plus de vastes sur le terrain. Merci d'avoir contribué au projet afin que d'autres puissent s'appuyer sur votre contribution. Pour obtenir des conseils sur la contribution au MIP Earth-2, veuillez vous référer aux directives contributives.
Ce travail est inspiré pour faciliter des engagements similaires entre les équipes ici à NVIDIA - les experts de la ML développant de nouveaux modèles et les experts du domaine en science du climat évaluant les compétences de ces modèles. Par exemple, les données d'entrée souvent nécessaires telles que les constantes de normalisation et les valeurs d'hyperparamètre ne sont pas emballées aux côtés des poids du modèle. Chaque modèle implémente généralement une interface légèrement différente. Les routines de notation sont spécifiques au modèle noté et peuvent ne pas être cohérents entre les groupes.
Le MIP Earth-2 relève ces défis et combler l'écart entre les experts du domaine qui évaluent le plus souvent les modèles ML, et les experts ML qui les produisent. Par rapport à d'autres projets de cet espace, Earth-2 MIP se concentre sur les modèles de notation à la volée. Il a des API Python adaptées à une itération rapide dans un livre de jupyter, des CLI pour les modèles de notation distribués sur de nombreux GPU et un cadre de plugin flexible qui permet à quiconque d'utiliser ses propres modèles ML. Plus important encore, le MIP Earth-2 aspire à faciliter l'exploration et la collaboration au sein de la communauté de la recherche climatique pour évaluer le potentiel des modèles d'IA dans les simulations climatiques et météorologiques.
Veuillez consulter la page de documentation pour des informations approfondies sur la MIP de la Terre-2, les fonctionnalités, les API, etc.
Earth-2 MIP est fourni dans la licence Apache 2.0, veuillez consulter Licence.txt pour le texte complet de la licence.
