earth2mip
v0.1.0
El Proyecto de Intercomparación del Modelo Earth-2 (MIP) es un marco de IA basado en Python que permite a los investigadores y científicos del clima explorar y experimentar con el uso de modelos de IA para el clima y el clima. Proporciona flujos de trabajo de referencia para comprender cómo los modelos de IA capturan la física de la atmósfera de la Tierra y cómo pueden trabajar con los modelos tradicionales de pronóstico del clima numérico. Por ejemplo, el repositorio proporciona una interfaz uniforme para ejecutar una inferencia utilizando puntos de control de modelo previamente capacitados y calificar la habilidad de dichos modelos utilizando ciertas métricas estándar. Este repositorio está destinado a facilitar el clima y la comunidad climática para que se les ocurra una buena línea de referencia de los eventos para probar los modelos y usar con una variedad de fuentes de datos.
Earth-2 MIP será instalable en PYPI tras el lanzamiento general. Mientras tanto, uno puede instalar desde la fuente:
git clone [email protected]:NVIDIA/earth2mip.git
cd earth2mip && pip install .Consulte la documentación de instalación para obtener más detalles y otras opciones.
Earth-2 MIP proporciona un conjunto de ejemplos que se pueden ver en la página de documentación de ejemplos que se puede utilizar para comenzar con varios flujos de trabajo. Estos ejemplos se pueden descargar como cuadernos Jupyer y scripts de Python. Los scripts de Python de origen se pueden encontrar en las carpetas de ejemplos.
Earth-2 MIP proporciona API de alto nivel para ejecutar inferencia con modelos AI. Por ejemplo, lo siguiente se puede usar para ejecutar el clima Pangu utilizando un estado inicial del almacén de datos climático (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: historyY puedes obtener ACC/RMSE así:
>>> 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'
Estos cuadernos ilustran cómo usar con algunos modelos y esto puede servir como referencia para traer su propio punto de control siempre que sea compatible. Puede haber trabajo adicional para hacerlo compatible con la Tierra-2 MIP. Earth-2 MIP aprovecha el zoológico del modelo en el módulo para proporcionar un conjunto de referencia de modelos de línea base. El objetivo es permitir que la comunidad aumente esta colección de modelos como se muestra en la tabla a continuación.
| IDENTIFICACIÓN | Modelo | Arquitectura | Tipo | Referencia | Fuente | Tamaño |
|---|---|---|---|---|---|---|
| FCN | Fourcastnet | Operador neural adaptativo de Fourier | clima global | Arxiv | módulo | 300MB |
| DLWP | Predicción del clima de aprendizaje profundo | Codificador convolucional | clima global | Agua | módulo | 50 MB |
| atar | Clima Pangu (jerárquico 6 + 24 h) | Transformador de visión | clima global | Naturaleza | ONNX | 2GB |
| pangu_6 | Modelo de clima panguero 6 horas | Transformador de visión | clima global | Naturaleza | ONNX | 1GB |
| pangu_24 | Modelo de clima panguero 24 horas | Transformador de visión | clima global | Naturaleza | ONNX | 1GB |
| FCNV2_SM | Fourcastnet v2 | Armónicos esféricos operador neural de Fourier | clima global | Arxiv | módulo | 3.5GB |
| gráfico | Graphcast, 37 niveles, 0.25 grados | Red neuronal gráfica | clima global | Ciencia | github | 145 MB |
| Graphcast_small | Graphcast, 13 niveles, 1 gr. | Red neuronal gráfica | clima global | Ciencia | github | 144 MB |
| Graphcast_operational | Graphcast, 13 niveles, 0.25 grados | Red neuronal gráfica | clima global | Ciencia | github | 144 MB |
| precipitación_afno | Precipitación de FourCastnet | Operador neural adaptativo de Fourier | diagnóstico | Arxiv | módulo | 300MB |
| climateneta | Modelo de segmentación climateneta | Red neuronal convolucional | diagnóstico | GMD | módulo | 2MB |
* = próximamente
Algunos modelos requieren dependencias adicionales no instaladas de forma predeterminada. Consulte las instrucciones de instalación para obtener más detalles.
Nota : Cada punto de control de modelo puede tener su propia licencia única. Alentamos a los usuarios a familiarizarse con cada uno para comprender las implicaciones para su caso de uso particular.
¡Queremos integrar su modelo en el marcador para mostrar a la comunidad! La mejor manera de hacerlo es a través del módulo Nvidia. Puede contribuir con su modelo (tanto el código de entrenamiento como el punto de control del modelo) y podemos asegurarnos de que se mantenga como parte del conjunto de referencias.
Earth-2 MIP es una colaboración de código abierto y su éxito se basa en la contribución de la comunidad para promover el campo. Gracias por contribuir al proyecto para que otros puedan aprovechar su contribución. Para obtener orientación sobre la contribución a la Tierra-2 MIP, consulte las pautas contribuyentes.
Este trabajo se inspira para facilitar compromisos similares entre los equipos aquí en NVIDIA: los expertos en ML que desarrollan nuevos modelos y los expertos en dominios en ciencia climática que evalúan la habilidad de dichos modelos. Por ejemplo, a menudo los datos de entrada necesarios, como las constantes de normalización y los valores de hiperparameter, no se empaquetan junto con los pesos del modelo. Cada modelo generalmente implementa una interfaz ligeramente diferente. Las rutinas de puntuación son específicas para el modelo que se califica y puede no ser consistente entre los grupos.
Earth-2 MIP aborda estos desafíos y une la brecha entre los expertos en dominios que con mayor frecuencia están evaluando modelos de ML, y los expertos en ML que los producen. En comparación con otros proyectos en este espacio, Earth-2 MIP se centra en la puntuación de los modelos sobre la marcha. Tiene API de Python adecuadas para la rápida iteración en un libro de Jupyter, clis para calificar modelos distribuidos a través de muchas GPU y un marco de complementos flexible que permite a cualquier persona usar sus propios modelos ML. Más importante aún, Earth-2 MIP aspira para facilitar la exploración y la colaboración dentro de la comunidad de investigación climática para evaluar el potencial de los modelos de IA en las simulaciones climáticas y climáticas.
Consulte la página de documentación para obtener información en profundidad sobre Earth-2 MIP, funcionalidad, API, etc.
Earth-2 MIP se proporciona bajo la licencia APACHE 2.0, consulte License.txt para ver el texto completo de la licencia.
