DOCSEARCHbyAlgolia
Submit Docs

Earth-2 Model Intercomparison Project (MIP)

repository·main·Indexed 42 days ago

https://github.com/nvidia/earth2mip

A Python framework for climate researchers and scientists to experiment with AI models for weather and climate forecasting. It provides a uniform interface for running inference and scoring model skill using metrics like ACC and RMSE. Earth-2 MIP supports various global weather and diagnostic models, including FourCastNet, Deep Learning Weather Prediction (DLWP), Pangu Weather, GraphCast, and ClimateNet.

Tokens
29.7K
Snippets
102
Records
117
Agent score
54%

What's inside Earth-2 MIP

  2. Overview of the earth2mip API modules

    main

    The earth2mip library provides a suite of modules for atmospheric modeling, inference, and data handling. The core API surface is organized into several functional areas:

    Inference and Forecasting

    • earth2mip.inference_ensemble: Tools for ensemble-based inference.
    • earth2mip.inference_medium_range: Tools for medium-range forecasting inference.
    • earth2mip.forecasts: Core forecasting logic.
    • earth2mip.make_job: Utilities for creating inference jobs.

    Data Loading and I/O

    • earth2mip.loaders: Modules for loading datasets.
    • earth2mip.datasets: Dataset management.
    • earth2mip.netcdf: NetCDF file handling.
    • earth2mip.filesystem: Filesystem utilities.
    • earth2mip.forecast_metrics_io: Input/Output for forecast metrics.

    Geospatial and Grid Operations

    • earth2mip.grid: Grid definitions and management.
    • earth2mip.regrid: Regridding operations.
    • earth2mip.geometry: Geometric primitives and operations.
    • earth2mip.geo_operator: Geospatial operators.

    Utilities and Diagnostics

    • earth2mip.diagnostics: Diagnostic tools.
    • earth2mip.ensemble_utils: Utilities for ensemble data.
    • earth2mip.score_ensemble_outputs: Scoring mechanisms for ensemble outputs.
    • earth2mip.time & earth2mip.time_collection: Time-related utilities.
    • earth2mip.time_loop: Time-stepping and loop logic.
    • earth2mip.weather_events: Detection or handling of weather events.
    • earth2mip.schema: Data schema definitions.
    • earth2mip.model_registry: Management of available models.

  3. Use the earth2mip.networks package for model inference

    main

    The earth2mip.networks package provides access to various deep learning model architectures used for Earth-2 MIP tasks. The package is organized into specialized modules for different model types, allowing users to instantiate and interact with specific neural network architectures for meteorological or climate forecasting.

    Supported model modules include:

    • earth2mip.networks.dlwp: Deep Learning Weather Prediction models.
    • earth2mip.networks.fcn: Fully Convolutional Network models.
    • earth2mip.networks.fcnv2_sm: FCNv2 Small models.
    • earth2mip.networks.graphcast: GraphCast architectures.
    • earth2mip.networks.pangu: Pangu-Weather architectures.

  4. What are diagnostic models in Earth-2 MIP

    main

    In Earth-2 MIP, a diagnostic model is a model that does not perform time-step predictions (prognostic). Instead, it predicts new quantities from existing ones at the same time step.

    Functionally, diagnostics act as earth2mip.geo_operator.GeoOperator instances and can be viewed as post-processors within the inference process. This allows you to extend foundational weather models to downstream tasks (like identifying tropical cyclones) without the expensive process of training a new time-integrating weather model.

  5. Implement a Machine Learning Module

    main

    At the root level, a Module is a machine learning model (typically a torch.nn.Module) that takes and returns plain torch tensors representing two-dimensional fields on a planet. To use a Module within the Earth2 MIP ecosystem, you must provide metadata about its inputs and outputs (channels, grid, etc.) using the earth2mip.schema.Model schema.

    Design Note: Earth2 MIP avoids using metadata-aware containers like xarray inside the modules to prevent mathematical complexity and maintain compatibility with standard ML workflows that expect multi-dimensional arrays.

    import torch.nn

class PersistenceModule(torch.nn.Module):
    def forward(self, x):
        return x

  6. Available submodules in earth2mip.initial_conditions

    main

    The earth2mip.initial_conditions package is organized into several submodules based on the data source or file format:

    • earth2mip.initial_conditions.base: Contains base classes and core logic for initial conditions.
    • earth2mip.initial_conditions.cds: Interface for Copernicus Climate Data Store (CDS) data.
    • earth2mip.initial_conditions.gfs: Interface for Global Forecast System (GFS) data.
    • earth2mip.initial_conditions.hdf5: Interface for loading data from HDF5 file formats.
    • earth2mip.initial_conditions.hrmip: Specific implementation for HRMIP (High-Resolution Model Intercomparison Project) data.
    • earth2mip.initial_conditions.ifs: Interface for Integrated Forecasting System (IFS) data.

  7. Understand Model Packages

    main

    A Model package is a directory used for reproducible loading of model checkpoints from disk or the cloud. It must contain:

    • A metadata.json file that follows the earth2mip.schema.Model schema.
    • Any other static data required for the model, such as model parameters or normalization constants.

  8. Understand Earth2 MIP core abstractions

    main

    Earth2 MIP uses layered abstractions to wrap machine learning models for weather and climate forecasting. The hierarchy moves from low-level mathematical operations to high-level forecast management:

    1. Machine Learning Module: The base layer (e.g., a torch.nn.Module) that performs the raw tensor transformations.
    2. Time Loop: A higher-level interface (earth2mip.time_loop.TimeLoop) that encapsulates the Module along with timestepping logic, data preprocessing, and output handling.
    3. Forecast: The highest level of abstraction (earth2mip.forecasts.Forecast) which represents a collection of states (a Forecast Array) where rows are initial times and columns are lead times. It encapsulates all initialization and time-handling logic, allowing the same scoring code to work for both streaming and static (disk-based) forecasts.

  9. Set up a Conda environment for Earth-2 MIP

    main

    It is highly encouraged to use a dedicated environment to avoid package conflicts. For a Conda-based setup:

    1. Create and activate a new environment with Python 3.10.
    2. Install Earth-2 MIP from source.
    3. (Optional) Install ipykernel to use the environment within Jupyter Lab.
    conda create --name earth2mip python=3.10
conda activate earth2mip

git clone git@github.com:NVIDIA/earth2mip.git
cd earth2mip && pip install .

# To use in Jupyter Lab:
pip install ipykernel
python -m ipykernel install --user --name=earth2mip-kernel

  10. Create a custom diagnostic model

    main

    You can implement custom diagnostics by subclassing earth2mip.diagnostic.base.DiagnosticBase and ensuring it meets the functional requirements of earth2mip.geo_operator.GeoOperator.

    A custom diagnostic must implement the following properties and methods:

    • in_channel_names: A list of input channel names (e.g., ['t2m']).
    • out_channel_names: A list of output channel names (e.g., ['t2m_c']).
    • in_grid: The input grid.LatLonGrid.
    • out_grid: The output grid.LatLonGrid.
    • __call__(x: torch.Tensor) -> torch.Tensor: The logic to transform the input tensor.
    • load_diagnostic: A class method to instantiate the diagnostic (often used for registry patterns).

    Note: Diagnostics must operate on grids and use channels consistent with the weather model they are attached to.

    class Kelvin2Celcius(DiagnosticBase):
    def __init__(self, grid: grid.LatLonGrid):
        super().__init__()
        self.grid = grid

    @property
    def in_channel_names(self) -> list[str]:
        return ['t2m']

    @property
    def out_channel_names(self) -> list[str]:
        return ['t2m_c']

    @property
    def in_grid(self) -> grid.LatLonGrid:
        return self.grid

    @property
    def out_grid(self) -> grid.LatLonGrid:
        return self.grid

    def __call__(self, x: torch.Tensor) -> torch.Tensor:
        return x - 273.15

    @classmethod
    def load_diagnostic(cls):
        pass

    @classmethod
    def load_diagnostic(
        cls, package: Optional[Package], grid: grid.LatLonGrid
    ):
        return cls(grid)

  11. Install Earth-2 MIP from source

    main

    Since Earth-2 MIP is not yet available on PyPI, you can install it by cloning the repository and using pip. Ensure you have a compatible Python version (3.10, 3.11, or 3.12).

    git clone git@github.com:NVIDIA/earth2mip.git

cd earth2mip && pip install .

  12. Access Earth-2 MIP examples as Jupyter notebooks and Python scripts

    main

    You can access a collection of Earth-2 MIP examples, including Jupyter notebooks and Python scripts, via the official documentation website. These examples demonstrate various functionalities and common workflows within the library.

    https://nvidia.github.io/earth2mip/examples/index.html