DOCSEARCHbyAlgolia
Submit Docs

Aurora

repository·main·Indexed 21 days ago

https://github.com/microsoft/aurora

A machine learning foundation model for Earth system forecasting designed to predict atmospheric variables such as temperature, air pollution, and ocean waves. Aurora provides specialized versions for medium- and high-resolution weather prediction, air pollution prediction, and ocean wave prediction. It can be run locally using AuroraSmallPretrained or AuroraHighRes, integrated with ECMWF's ai-models framework via the ai-models-aurora plugin, or executed on Azure AI Foundry.

Tokens
16.6K
Snippets
55
Records
67
Agent score
75%

What's inside microsoft-aurora

  1. Overview of Aurora Foundation Model

    main

    Aurora is a machine learning foundation model designed for predicting atmospheric and environmental variables (e.g., temperature, wind speed). It is pretrained on diverse weather and climate data and can be fine-tuned for specialized tasks such as:

    • Weather prediction
    • Air pollution modelling
    • Ocean wave forecasting

    Aurora 1.5 Improvements: Aurora 1.5 extends the original architecture with:

    • 22 new single-level output variables: Including radiation fluxes, precipitation, and 100-m winds.
    • Variable lead-time embeddings: Enables predictions at any lead time as fine as one hour.
    • Ensemble version: Includes stochastic noise injection to generate physically plausible ensemble members for probabilistic forecasting.

  2. What is Aurora and its available model versions?

    main

    Aurora is a machine learning foundation model for Earth system forecasting. It is designed to be trained generally on large datasets and then adapted to specialized atmospheric forecasting tasks with minimal task-specific data.

    Aurora provides four specialized versions:

    1. Medium-resolution weather prediction
    2. High-resolution weather prediction
    3. Air pollution prediction
    4. Ocean wave prediction

  3. Configure noise accumulation in Aurora 1.5 Ensemble

    main

    When using rollout with fine_lead_times in an ensemble model, you can control how noise is handled via the use_noise_accumulation parameter:

    • use_noise_accumulation=True (Default): Keeps noise correlated across sub-steps for smoother intra-step transitions while using independent effective noise between main steps. This matches the training regimen.
    • use_noise_accumulation=False: Draws independent noise at each sub-step. This is not recommended.

  4. Understand the intended and out-of-scope use cases for Aurora

    main

    Primary Use Cases

    Aurora 1.5 is designed for medium-range deterministic and probabilistic weather prediction. Its foundation-model architecture allows it to be adapted to specialized atmospheric forecasting tasks with limited task-specific data.

    Key features of Aurora 1.5 include:

    • Expanded surface variables (26 inputs, 7 outputs).
    • Variable lead-time embeddings for predictions as fine as one hour.
    • An ensemble version (AuroraV1p5Ensemble) for probabilistic forecasting and uncertainty quantification.

    Out-of-Scope Use Cases

    Aurora is not designed for:

    • Direct operational decision-making without expert review.
    • Applications requiring guaranteed forecast accuracy.
    • Non-environmental prediction tasks.
    • Safety-critical planning or automated decision pipelines without domain validation.

    Note: It is not recommended to use Aurora in commercial or real-world applications without extensive further testing and development.

  5. Disable LoRA in Aurora for more realistic predictions

    main

    In the fine-tuned models (like Aurora or AuroraHighRes), you can choose between optimal long-term Mean Squared Error (MSE) and prediction realism by toggling LoRA:

    • LoRA Enabled (Default): Provides optimal long-term MSE, but predictions may be slightly blurrier.
    • LoRA Disabled: Provides more realistic predictions at the expense of slightly higher long-term MSE.

    To disable LoRA, set use_lora=False in the constructor and use strict=False in load_checkpoint().

    from aurora import Aurora

# Disable LoRA for more realistic samples.
model = Aurora(use_lora=False) 
model.load_checkpoint(strict=False)

  6. Control rollout input clipping in Aurora 1.5

    main

    To prevent continuous rollouts from drifting into unrealistic physical states, rollout includes an option apply_rollout_input_clipping (defaults to True). This clips certain surface variables to known minimum/maximum bounds (e.g., cloud cover in [0, 1]) only when a Batch is being fed back into the model.

    Note that this clipping does not prevent the model's forward pass from generating values outside these bounds; it only ensures the feedback loop matches the training regimen. You can override the default clipping behavior using the rollout_input_clipping keyword argument in rollout.

  7. Responsible AI usage and limitations of Aurora

    main

    Aurora is a research forecasting model and should not be treated as an operational weather service.

    Intended Use

    Aurora is intended for:

    • Medium- and high-resolution weather prediction
    • Air pollution forecasting
    • Ocean wave prediction
    • Global environmental modelling

    Out-of-scope Use Cases

    Aurora is not designed for:

    • Direct operational decision-making without expert review
    • Applications requiring guaranteed forecast accuracy
    • Non-environmental prediction tasks
    • Safety-critical planning or automated decision pipelines without domain validation

    Key Limitations

    • No Accuracy Guarantees: As a neural network, predictions may be inaccurate, especially if inputs are out-of-distribution or differ from the training set.
    • Inherited Biases: The model may inherit biases present in its training datasets (e.g., ERA5, CMIP6, CAMS).
    • Research-only Version: The published models are streamlined versions of internal models and may deviate subtly from original performance.
    • Reliability Degradation: Reliability can degrade in out-of-distribution conditions like rare extremes or regime shifts.

  8. Responsible AI considerations for Aurora

    main

    Aurora is a research forecasting model, not an operational weather service. Its reliability may degrade in out-of-distribution conditions (e.g., rare extremes or regime shifts).

    Risk Mitigation for Developers:

    • Communicate Intent: Clearly state that the model is for research evaluation and reproducibility.
    • Expert Review: Require domain-expert review before any real-world decisions are informed by outputs.
    • Input Validation: Ensure initial conditions come from credible sources (e.g., established meteorological agencies).
    • Benchmarking: Compare performance against accepted physical modeling systems for the specific geography and horizon.
    • Human-in-the-loop: Integrate Aurora as a decision-support component for expert analysis rather than an autonomous trigger for actions.

  9. Use Aurora 0.25° 12-Hour Pretrained for 12-hour lead times

    main

    Use Aurora12hPretrained specifically when you require predictions with a 12-hour lead time. It requires the same standard variables and pressure levels as the 0.25° Pretrained model.

    from aurora import Aurora12hPretrained

model = Aurora12hPretrained()
model.load_checkpoint()

  10. Deploying Aurora via Azure AI Foundry

    main

    Aurora is available in the Azure AI model catalog. For the latest version, Aurora-1.5, Azure AI Foundry uses a custom model and container image approach rather than traditional MLflow-based deployments.

    Note: This repository provides the client-side code to call endpoints deployed with Foundry, but it does not support building your own endpoint locally. It is recommended to use the pre-deployed models in the Azure AI model catalog.

  11. Use Aurora 0.25° Wave for ocean wave forecasts

    main

    The AuroraWave model is fine-tuned on HRES-WAM ocean wave data. It is designed for predictions combining HRES T0 meteorological variables and HRES-WAM analysis ocean wave variables.

    Warning: For optimal performance, ensure all meteorological variables are from HRES T0 and all ocean wave variables are from HRES-WAM analysis. Specific postprocessing may be required (see Supplementary Information Section C.5).

    Required Variables:

    • Surface-level: 2t, 10u, 10v, swh, mwd, mwp, pp1d, shww, mdww, mpww, shts, mdts, mpts, swh1, mwd1, mwp1, swh2, mwd2, mwp2, 10u_wave, 10v_wave, wind
    • Static: lsm, slt, z, wmb, lat_mask
    • Atmospheric: t, u, v, q, z
    • Pressure levels (hPa): 50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000
    from aurora import AuroraWave

model = AuroraWave()
model.load_checkpoint()

  12. Extend Aurora with new variables

    main

    You can add new variables to Aurora by updating the surf_vars, static_vars, and atmos_vars keyword arguments in the AuroraPretrained constructor.

    1. Define new variables

    Pass the new variable names as part of the tuples for the relevant category.

    2. Set normalization statistics

    You must manually set the normalization means (locations) and standard deviations (scales) for every new variable using aurora.normalisation.locations and aurora.normalisation.scales.

    3. Handle initialization

    By default, patch embeddings for new variables in the encoder are initialized randomly, which perturbs existing variable predictions. To prevent this, you can initialize them to zero using model.encoder.{surf,atmos}_token_embeds.weights.

    4. Learning rates

    For better efficiency, use a separate, higher learning rate (e.g., 1e-3) for the patch embeddings of the new variables and a lower rate (e.g., 3e-4) for the rest of the parameters.

    from aurora import AuroraPretrained
from aurora.normalisation import locations, scales

model = AuroraPretrained(
    surf_vars=("2t", "10u", "10v", "msl", "new_surf_var"),
    static_vars=("lsm", "z", "slt", "new_static_var"),
    atmos_vars=("z", "u", "v", "t", "q", "new_atmos_var"),
)
model.load_checkpoint(strict=False)

# Normalisation means:
locations["new_surf_var"] = 0.0
locations["new_static_var"] = 0.0
locations["new_atmos_var"] = 0.0

# Normalisation standard deviations:
scales["new_surf_var"] = 1.0
scales["new_static_var"] = 1.0
scales["new_atmos_var"] = 1.0