Aardvark Weather

Aardvark Weather is trained in multiple stages using specific scripts. Note that the training commands provided below are for illustrative purposes to demonstrate the training process and may require specific local data loading pipelines and compute infrastructure to execute.

Training Stages and Scripts

• Encoder Training: Handled via train_module.py (illustrated by train_encoder.sh).
• Processor Training: Includes pre-training via train_module.py (illustrated by train_processor.sh) and finetuning via finetune.py (illustrated by finetune.sh).
• Decoder Training: Handled via train_module.py (illustrated by train_decoder.sh).
• End-to-End Finetuning: Handled via e2e_train.py (illustrated by train_e2e.sh).

Trained model weights for the encoder, decoder, processor, and end-to-end finetuned models (one-day lead time) are available on Hugging Face. For the decoder and end-to-end models, weights are provided for both temperature and wind speed modalities.

https://huggingface.co/datasets/av555/aardvark-weather/tree/main/trained_model

There are two ways to access data for this project:

1. Sample Data: A local file data/sample_data_final.pkl is provided for immediate use with the demo notebooks.
2. Full Training Dataset: A machine learning-ready dataset containing observational data from 2007-2019 at 24-hour resolution is available on Hugging Face for those wishing to train their own end-to-end models.

https://huggingface.co/datasets/av555/aardvark-weather

The notebooks/ directory contains demonstrations for producing different types of forecasts:

• Full Model Forecasts: Use forecast_demo.ipynb to load the complete Aardvark Weather model, generate predictions from sample data, and visualize global gridded and station forecasts.
• Optimized Station Forecasts: Use e2e_finetune_demo.ipynb to generate optimized station forecasts using the end-to-end finetuned model.

SYNOPS data includes observations from land stations, marine platforms, and radiosonde profiles. Unlike the grid-based satellite data, these are often represented as point observations using scatter plots.

Available SYNOPS keys:

• x_context_hadisd_current / y_context_hadisd_current (HadISD)
• icoads_x_current / icoads_current (ICOADS)
• igra_x_current / igra_current (IGRA)

For HadISD, the longitude and latitude are stored in separate x_context and y_context keys.

# Example: Visualizing HadISD data
plot_channel = 0
plt.scatter(
    data["assimilation"]["x_context_hadisd_current"][plot_channel][0,0,:].cpu(),
    data["assimilation"]["x_context_hadisd_current"][plot_channel][0,1,:].cpu(),
    c = data["assimilation"]["y_context_hadisd_current"][plot_channel][0].T.cpu(), 
    cmap="magma"
)
plt.show()

The model outputs normalized predictions that must be unnormalised to represent physical values. This process requires loading the mean and std factors specific to the forecast variable from the norm_factors directory and applying a scaling factor (e.g., HADISD_SCALING_FACTOR = 10).

The formula used is: (prediction * std + mean) * scaling_factor.

# Load normalization factors
mean = np.load(f"../data/norm_factors/mean_hadisd_{local_forecast_var}.npy")
std = np.load(f"../data/norm_factors/std_hadisd_{local_forecast_var}.npy")

# Unnormalise
station_forecast_unnorm = (
    station_forecast.detach().cpu() * std + mean
) * HADISD_SCALING_FACTOR

Satellite data is stored under the data["assimilation"] key. The data is typically represented as tensors that need to be moved to the CPU and transposed for plotting.

Available satellite keys:

• hirs_current
• amsua_current
• amsub_current
• iasi_current
• ascat_current

Note: For amsua_current, the latitude array may require an offset (e.g., lat[:-1]) depending on the grid dimensions.

# Example: Visualizing HIRS data
import matplotlib.pyplot as plt

plot_channel = 11
plt.contourf(
    lon, 
    lat, 
    data["assimilation"]["hirs_current"][0, ..., plot_channel].cpu().T, 
    levels=100, 
    cmap="magma"
)
plt.show()

To inspect the data structure used for forecasting, you can load the sample pickle file provided in the repository. This file contains a single timeslice of all required observations.

import pickle

with open('../data/sample_data_final.pkl', 'rb') as fp:
    data = pickle.load(fp)

To generate weather forecasts, use the ConvCNPWeatherE2E class from the e2e_model module. You must specify the lead_time, paths to the trained encoder, processor, and decoder models, and the path to auxiliary data. The decoder path is variable-specific (e.g., ../trained_models/decoder/tas/).

Parameters:

• device: The device to run on (e.g., `

The following shell commands illustrate the training and finetuning steps for the various modules. These are intended for transparency regarding the training process and are not guaranteed to be executable without specific local infrastructure setup.

# Encoder
bash train_encoder.sh

# Processor
bash train_processor.sh
bash finetune.sh

# Decoder
bash train_decoder.sh

# End-to-end
bash train_e2e.sh