Commands¶

datasets create¶

You will need a CDS account to download ERA5 data with anemoi. By default, a sample dataset will be downloaded which should be small enough to fit on your personal computer. The full datasets are available on Isambard.

uv run imp datasets create

Anemoi tracks which date groups have been downloaded, so an interrupted download can be resumed by simply rerunning this command.

To create the synthetic dataset, use:

uv run imp datasets create --config-name synthetic

datasets inspect¶

uv run imp datasets inspect

Prints basic properties of each dataset. With the --verbose option it will also print statistical summaries of the variables.

datasets plot¶

uv run imp datasets plot --dataset samp-sicsouth-osisaf-25p0km-2020-2024-24h-v1 --timestep 0

Creates one static PNG per variable for the selected timestep of a configured downloaded dataset. Plots are written to ${base_path}/data/input_plots in a subdirectory named after the dataset. Running without --dataset will plot every configured dataset. This is useful for inspecting raw inputs without running model training or evaluation.

Use --timestep to select another dataset index and the normal --config-name to choose the dataset configuration.

Pass --video to create an animation over consecutive timesteps instead of a single static plot:

uv run imp datasets plot --dataset samp-sicsouth-osisaf-25p0km-2020-2024-24h-v1 --video --timestep 0 --n-steps 30

This animates --n-steps consecutive timesteps starting at --timestep, writing one video file per variable to the same input_plots directory.

train¶

Standard (non-synthetic) runs use Weights & Biases. Generate an API key, then authenticate:

export WANDB_API_KEY=<your_api_key>
wandb login

Trains the model end-to-end:

uv run imp train

PYTORCH_ENABLE_MPS_FALLBACK=1 uv run imp train

Multistage training¶

For EncodeProcessDecode models, pass --multistage to train each component separately before finetuning.

uv run imp train --multistage

Checkpoints are saved to ${BASE_DIR}/training/wandb/run-<date>-<id>/checkpoints/<name>.ckpt, where BASE_DIR is the base path defined in your config.

See Train in stages for a full walkthrough.

Weights & Biases logging¶

To disable logging to W&B, set either loggers.wandb.offline=true or the WANDB_MODE=offline environment variable

uv run imp train loggers.wandb.offline=true
WANDB_MODE=offline uv run imp train

Run data such as metrics and figures will still be written locally, but will not be uploaded.

Synthetic experiments do not use W&B. Use the synthetic configuration, which saves metrics and plotting artefacts locally under ${BASE_DIR}/report:

uv run imp train --config-name synthetic

sweep initialise¶

uv run imp sweep initialise --sweep-yaml example.sweep.yaml --config-name baseline/02_cnn_unet_cnn

Creates a W&B sweep and initialises a local Optuna study directory; hyperparameters are sampled per trial at runtime. See Run a hyperparameter sweep for the full workflow.

sweep trial¶

uv run imp sweep trial --sweep-path <path to sweep directory created above>

Runs a single hyperparameter trial as part of a W&B sweep. See Run a hyperparameter sweep for the full workflow.

sweep summarise¶

uv run imp sweep summarise --sweep-path <path to sweep directory created above>

Reads the local Optuna study and reports the number of completed trials, plus the value and hyperparameters for the best trial. This works without a W&B connection. See Run a hyperparameter sweep for the full workflow.

evaluate¶

uv run imp evaluate --checkpoint PATH_TO_A_CHECKPOINT

Visualisations¶

To plot static images or animations of the raw input data, add the following to your local config:

evaluate:
  callbacks:
    plotting:
      make_input_plots: true

Output directories, styling, and animation parameters can be altered by changing config.evaluate.callbacks.plotting.plot_spec. Any of these can be overridden at the command line.