### Install Core Dependencies

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/README.md

Install essential libraries including PyTorch Lightning, xarray, netcdf4, opencv-python, and the earthnet library.

```bash
python3 -m pip install pytorch_lightning==1.6.4
python3 -m pip install xarray netcdf4 opencv-python earthnet==0.3.9
```

--------------------------------

### Install Apex with CUDA 11.7

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/README.md

Installs the Apex library for PyTorch with specific CUDA version support. Ensure CUDA_HOME is set correctly.

```bash
CUDA_HOME=/usr/local/cuda python3 -m pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" pytorch-extension git+https://github.com/NVIDIA/apex.git
```

--------------------------------

### Execute Training Scripts via CLI

Source: https://context7.com/amazon-science/earth-forecasting-transformer/llms.txt

Provides command-line examples for training, testing, and resuming experiments for SEVIR, EarthNet2021, and ENSO datasets.

```bash
# SEVIR Precipitation Nowcasting
# Train from scratch
MASTER_ADDR=localhost MASTER_PORT=10001 \
python scripts/cuboid_transformer/sevir/train_cuboid_sevir.py \
    --gpus 2 \
    --cfg scripts/cuboid_transformer/sevir/earthformer_sevir_v1.yaml \
    --save sevir_experiment

# Test with pretrained checkpoint
python scripts/cuboid_transformer/sevir/train_cuboid_sevir.py \
    --gpus 1 \
    --pretrained \
    --test

# Resume training from checkpoint
python scripts/cuboid_transformer/sevir/train_cuboid_sevir.py \
    --gpus 2 \
    --cfg earthformer_sevir_v1.yaml \
    --ckpt_name last.ckpt \
    --save sevir_experiment

# EarthNet2021 Earth Surface Forecasting
MASTER_ADDR=localhost MASTER_PORT=10001 \
python scripts/cuboid_transformer/earthnet_w_meso/train_cuboid_earthnet.py \
    --gpus 2 \
    --cfg cfg.yaml \
    --save earthnet_experiment

# ENSO Forecasting
python scripts/cuboid_transformer/enso/train_cuboid_enso.py \
    --gpus 1 \
    --cfg earthformer_enso_v1.yaml \
    --save enso_experiment

# Download pretrained checkpoints
# Pretrained models are auto-downloaded when using --pretrained flag
# Manual download:
# wget https://earthformer.s3.amazonaws.com/pretrained_checkpoints/earthformer_sevir.pt
# wget https://earthformer.s3.amazonaws.com/pretrained_checkpoints/earthformer_earthnet2021.pt
# wget https://earthformer.s3.amazonaws.com/pretrained_checkpoints/earthformer_icarenso2021.pt
```

--------------------------------

### Define Experiment Configuration

Source: https://context7.com/amazon-science/earth-forecasting-transformer/llms.txt

Example YAML configuration file defining dataset parameters, optimization settings, model architecture, and logging preferences.

```yaml
# Example configuration file: earthformer_sevir_v1.yaml
dataset:
  dataset_name: "sevir"
  img_height: 384
  img_width: 384
  in_len: 13                           # Input sequence length
  out_len: 12                          # Output sequence length
  seq_len: 25                          # Total sequence length
  layout: "NTHWC"                      # Tensor layout
  sample_mode: "sequent"
  stride: 12
  metrics_list: ['csi', 'pod', 'sucr', 'bias']
  threshold_list: [16, 74, 133, 160, 181, 219]

optim:
  total_batch_size: 32
  micro_batch_size: 2                  # Per-GPU batch size
  seed: 0
  method: "adamw"
  lr: 0.001
  wd: 0.0
  gradient_clip_val: 1.0
  max_epochs: 100
  lr_scheduler_mode: "cosine"
  warmup_percentage: 0.2
  min_lr_ratio: 0.001
  early_stop: true
  early_stop_patience: 20

model:
  input_shape: [13, 384, 384, 1]
  target_shape: [12, 384, 384, 1]
  base_units: 128
  enc_depth: [1, 1]
  dec_depth: [1, 1]
  num_heads: 4
  attn_drop: 0.1
  proj_drop: 0.1
  ffn_drop: 0.1
  num_global_vectors: 8
  self_pattern: "axial"                # Attention pattern
  cross_self_pattern: "axial"
  cross_pattern: "cross_1x1"
  ffn_activation: "gelu"
  norm_layer: "layer_norm"
  pos_embed_type: "t+h+w"
  use_relative_pos: true
  initial_downsample_type: "stack_conv"
  initial_downsample_stack_conv_num_layers: 3
  initial_downsample_stack_conv_dim_list: [16, 64, 128]
  initial_downsample_stack_conv_downscale_list: [3, 2, 2]

trainer:
  precision: 32                        # 16 for mixed precision
  check_val_every_n_epoch: 1

logging:
  logging_prefix: "Cuboid_SEVIR"
  use_wandb: false
```

--------------------------------

### Install Earthformer Package

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/README.md

Install the Earthformer package in editable mode after navigating to the project's root directory. Use --no-build-isolation for compatibility.

```bash
cd ROOT_DIR/earth-forecasting-transformer
python3 -m pip install -U -e . --no-build-isolation
```

--------------------------------

### Train Earthformer on N-body MNIST

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/nbody/README.md

Run this command to start training. Ensure configurations in cfg.yaml are adjusted as needed. This command utilizes two GPUs for training.

```bash
MASTER_ADDR=localhost MASTER_PORT=10001 python train_cuboid_nbody.py --gpus 2 --cfg cfg.yaml --ckpt_name last.ckpt --save tmp_nbody
```

--------------------------------

### Install PyTorch and Dependencies (CUDA 11.7)

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/README.md

Installs PyTorch version 1.13.1 with CUDA 11.7 support, along with PyTorch Lightning and other necessary libraries. This is for systems with CUDA 11.7 installed under /opt/cuda.

```bash
python3 -m pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 -f https://download.pytorch.org/whl/torch_stable.html
```

```bash
python3 -m pip install pytorch_lightning==1.6.4
```

```bash
python3 -m pip install xarray netcdf4 opencv-python earthnet==0.3.9
```

```bash
cd ROOT_DIR/earth-forecasting-transformer
python3 -m pip install -U -e . --no-build-isolation
```

```bash
CUDA_HOME=/opt/cuda python3 -m pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" pytorch-extension git+https://github.com/NVIDIA/apex.git
```

--------------------------------

### Train Earthformer on Moving MNIST

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/moving_mnist/README.md

Run this command to start training. Ensure configurations in cfg.yaml are adjusted as needed. This command utilizes two GPUs for training.

```bash
MASTER_ADDR=localhost MASTER_PORT=10001 python train_cuboid_mnist.py --gpus 2 --cfg cfg.yaml --ckpt_name last.ckpt --save tmp_mnist
```

--------------------------------

### Install PyTorch with CUDA Support

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/README.md

Install PyTorch version 1.12.1 with CUDA 11.6 support. Ensure your CUDA installation path and version match the command.

```bash
python3 -m pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 -f https://download.pytorch.org/whl/torch_stable.html
```

--------------------------------

### Install EarthNet2021 Dependencies

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/earthnet_w_meso/inference_tutorial_earthformer_earthnet2021.ipynb

Installs required Python packages for EarthNet2021, including PyTorch Lightning, EarthNet, OmegaConf, and specific GitHub repositories. It also downloads configuration and training script files.

```python
!python3 -m pip install -q "pytorch_lightning>=1.6.4,<1.8.0"
!python3 -m pip install -q earthnet
!python3 -m pip install -q omegaconf
!python3 -m pip install -q git+https://github.com/amazon-science/earth-forecasting-transformer.git
!python3 -m pip install -q --disable-pip-version-check --no-cache-dir git+https://github.com/NVIDIA/apex.git
!wget --quiet  https://raw.githubusercontent.com/amazon-science/earth-forecasting-transformer/main/scripts/cuboid_transformer/earthnet_w_meso/earthformer_earthnet_v1.yaml -O earthformer_earthnet_v1.yaml
!wget --quiet  https://raw.githubusercontent.com/amazon-science/earth-forecasting-transformer/main/scripts/cuboid_transformer/earthnet_w_meso/train_cuboid_earthnet.py -O train_cuboid_earthnet.py
```

--------------------------------

### Check CUDA Version

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/README.md

Verify the installed CUDA version on your machine. This is a prerequisite for installing compatible PyTorch versions.

```bash
nvcc --version
```

--------------------------------

### Initialize and Use SEVIRDataLoader

Source: https://context7.com/amazon-science/earth-forecasting-transformer/llms.txt

Configures the SEVIR dataset loader for precipitation nowcasting and demonstrates batch iteration and data preprocessing utilities.

```python
from earthformer.datasets.sevir.sevir_dataloader import SEVIRDataLoader
import datetime

# Initialize SEVIR DataLoader for VIL (Vertically Integrated Liquid)
dataloader = SEVIRDataLoader(
    data_types=['vil'],                      # Data modalities to load
    seq_len=25,                              # Total sequence length
    raw_seq_len=49,                          # Raw sequence length in dataset
    sample_mode='sequent',                   # 'random' or 'sequent'
    stride=12,                               # Stride between sequences
    batch_size=4,
    layout='NTHWC',                          # Output tensor layout
    # Date filtering
    start_date=datetime.datetime(2017, 1, 1),
    end_date=datetime.datetime(2019, 12, 31),
    # Data split (for train/val/test)
    num_shard=1,
    rank=0,
    # Preprocessing
    preprocess=True,
    rescale_method='01',                     # Scale values to [0, 1]
    shuffle=True,
    shuffle_seed=42,
)

# Iterate over batches
for batch in dataloader:
    vil_data = batch['vil']  # Shape: (batch_size, T, H, W, C)
    print(f"VIL batch shape: {vil_data.shape}")
    # Split into input/target for nowcasting
    x = vil_data[:, :13]    # Context: 65 minutes (13 frames * 5 min)
    y = vil_data[:, 13:]    # Target: 60 minutes (12 frames * 5 min)
    break

# Preprocessing utilities
preprocessed = SEVIRDataLoader.preprocess_data_dict(
    data_dict={'vil': vil_data},
    data_types=['vil'],
    layout='NTHWC',
    rescale='01'
)

# Reverse preprocessing for visualization
original = SEVIRDataLoader.process_data_dict_back(
    data_dict=preprocessed,
    data_types=['vil'],
    rescale='01'
)
```

--------------------------------

### Initialize MovingMNIST Dataset

Source: https://context7.com/amazon-science/earth-forecasting-transformer/llms.txt

Demonstrates how to instantiate the MovingMNIST dataset and create PyTorch DataLoaders for training and testing.

```python
from earthformer.datasets.moving_mnist.moving_mnist import MovingMNIST
from torch.utils.data import DataLoader

# Initialize MovingMNIST dataset
train_dataset = MovingMNIST(
    root='./datasets/moving_mnist',
    train=True,
    split=1000,                              # Test set size
    download=True,                           # Auto-download if not exists
    transform=None,
    target_transform=None,
)

test_dataset = MovingMNIST(
    root='./datasets/moving_mnist',
    train=False,
    split=1000,
    download=True,
)

# Create data loaders
train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=16, shuffle=False)

# Each sample contains (input_seq, target_seq)
for seq, target in train_loader:
    # seq: first 10 frames, target: last 10 frames
    # Both have shape (batch, 10, 64, 64) normalized to [0, 1]
    print(f"Input: {seq.shape}, Target: {target.shape}")
    break
```

--------------------------------

### Configure CuboidSEVIRPLModule

Source: https://context7.com/amazon-science/earth-forecasting-transformer/llms.txt

Sets up the configuration dictionary for the PyTorch Lightning training module.

```python
from omegaconf import OmegaConf
import pytorch_lightning as pl
from pytorch_lightning import Trainer, seed_everything

# Assuming the training script is properly imported
# from scripts.cuboid_transformer.sevir.train_cuboid_sevir import CuboidSEVIRPLModule

# Configuration via YAML or programmatic setup
config = OmegaConf.create({
    "dataset": {
        "dataset_name": "sevir",
        "img_height": 384,
        "img_width": 384,
        "in_len": 13,
        "out_len": 12,
        "seq_len": 25,
        "layout": "NTHWC",
        "sample_mode": "sequent",
        "stride": 12,
    },
    "optim": {
        "total_batch_size": 32,
        "micro_batch_size": 2,
        "lr": 1e-3,
        "wd": 0.0,
        "max_epochs": 100,
        "lr_scheduler_mode": "cosine",
        "warmup_percentage": 0.2,
    },
    "model": {
        "input_shape": [13, 384, 384, 1],
        "target_shape": [12, 384, 384, 1],
        "base_units": 128,
        "enc_depth": [1, 1],
        "dec_depth": [1, 1],
        "num_heads": 4,
        "num_global_vectors": 8,
        "self_pattern": "axial",
        "cross_pattern": "cross_1x1",
    }
})

# Training command line usage:
# python train_cuboid_sevir.py --gpus 2 --cfg earthformer_sevir_v1.yaml --save my_experiment

# For inference with pretrained model:
# python train_cuboid_sevir.py --gpus 1 --pretrained --test
```

--------------------------------

### Upload SEVIR Experiment Logs

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/sevir/README.md

Uploads experiment records to TensorBoard for the SEVIR training run.

```bash
tensorboard dev upload --logdir ./experiments/tmp_sevir/lightning_logs --name 'tmp_sevir'
```

--------------------------------

### Upload Experiment Records to TensorBoard

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/earthnet_w_meso/README.md

Run this command to upload experiment logs to TensorBoard for visualization. Specify the log directory and a name for your experiment.

```bash
tensorboard dev upload --logdir ./experiments/tmp_earthnet_w_meso/lightning_logs --name 'tmp_earthnet_w_meso'
```

--------------------------------

### Upload Experiment Records to TensorBoard

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/moving_mnist/README.md

Execute this command to upload your experiment logs to TensorBoard for visualization. The --logdir argument should point to your experiment's output directory.

```bash
tensorboard dev upload --logdir ./experiments/tmp_mnist/lightning_logs --name 'tmp_mnist'
```

--------------------------------

### Upload Experiment Records to TensorBoard

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/nbody/README.md

Execute this command to upload experiment logs to TensorBoard for visualization. The log directory should match the save path used during training.

```bash
tensorboard dev upload --logdir ./experiments/tmp_nbody/lightning_logs --name 'tmp_nbody'
```

--------------------------------

### Upload SEVIR-LR Experiment Logs

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/sevir/README.md

Uploads experiment records to TensorBoard for the SEVIR-LR training run.

```bash
tensorboard dev upload --logdir ./experiments/tmp_sevirlr/lightning_logs --name 'tmp_sevirlr'
```

--------------------------------

### Run Persistence Baseline on EarthNet2021

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/baselines/persistence/earthnet/README.md

Executes the persistence baseline script with specified GPU and configuration settings. Ensure the working directory is set to the project root.

```bash
cd ROOT_DIR/earth-forecasting-transformer
MASTER_ADDR=localhost MASTER_PORT=10001 python ./scripts/baselines/persistence/earthnet/test_persistence_earthnet.py --gpus 2 --cfg ./scripts/baselines/persistence/earthnet/cfg.yaml --save tmp_earthnet_persistence
```

--------------------------------

### Test Earthformer with Pretrained Checkpoint

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/sevir/README.md

Loads a pretrained checkpoint for testing on the SEVIR dataset.

```bash
MASTER_ADDR=localhost MASTER_PORT=10001 python train_cuboid_sevir.py --gpus 2 --pretrained --save tmp_sevir
```

--------------------------------

### Upload Experiment Logs to TensorBoard

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/baselines/persistence/earthnet/README.md

Uploads local Lightning logs to TensorBoard for experiment tracking.

```bash
cd ROOT_DIR/earth-forecasting-transformer
tensorboard dev upload --logdir ./experiments/tmp_earthnet/lightning_logs --name 'tmp_earthnet'
```

--------------------------------

### Download SEVIR Dataset

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/README.md

Downloads the SEVIR dataset from AWS S3. The command should be executed from the project's root directory.

```bash
cd ROOT_DIR/earth-forecasting-transformer
python ./scripts/datasets/sevir/download_sevir.py --dataset sevir
```

--------------------------------

### Train Earthformer on SEVIR

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/sevir/README.md

Executes training on the SEVIR dataset using a specific configuration file.

```bash
MASTER_ADDR=localhost MASTER_PORT=10001 python train_cuboid_sevir.py --gpus 2 --cfg cfg_sevir.yaml --ckpt_name last.ckpt --save tmp_sevir
```

--------------------------------

### Train Earthformer on SEVIR-LR

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/sevir/README.md

Executes training on the SEVIR-LR dataset using the corresponding configuration file.

```bash
MASTER_ADDR=localhost MASTER_PORT=10001 python train_cuboid_sevir.py --gpus 2 --cfg cfg_sevirlr.yaml --ckpt_name last.ckpt --save tmp_sevirlr
```

--------------------------------

### Upload Experiment Logs to TensorBoard

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/enso/README.md

Uploads local experiment logs to TensorBoard for visualization.

```bash
tensorboard dev upload --logdir ./experiments/tmp_enso/lightning_logs --name 'tmp_enso'
```

--------------------------------

### Generate N-body MNIST Dataset

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/datasets/nbody/README.md

Run this command to generate the N-body MNIST dataset using a specified configuration file. Navigate to the project root directory first.

```bash
cd ROOT_DIR/earth-forecasting-transformer
python ./scripts/datasets/nbody/generate_nbody_dataset.py --cfg ./scripts/datasets/nbody/cfg.yaml
```

--------------------------------

### Load EarthNet2021 Configuration and Dataset

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/earthnet_w_meso/inference_tutorial_earthformer_earthnet2021.ipynb

Loads the EarthNet2021 configuration from a YAML file and initializes the IID test dataset. It sets up the random seed and defines dataset parameters based on the configuration.

```python
from omegaconf import OmegaConf
from earthformer.datasets.earthnet.earthnet_dataloader import EarthNet2021LightningDataModule, EarthNet2021TestDataset
from pytorch_lightning import Trainer, seed_everything, loggers as pl_loggers
from earthformer.datasets.earthnet.visualization import vis_earthnet_seq
import numpy as np
from pathlib import Path

config_file = "./earthformer_earthnet_v1.yaml"
config = OmegaConf.load(open(config_file, "r"))
in_len = config.layout.in_len
out_len = config.layout.out_len

seed = config.optim.seed
dataset_cfg = OmegaConf.to_object(config.dataset)

seed_everything(seed, workers=True)

micro_batch_size = 1

earthnet_iid_testset = EarthNet2021TestDataset(subset_name="iid",
                                               data_dir="./datasets/earthnet2021/iid_test_split",
                                               layout=config.dataset.layout,
                                               static_layout=config.dataset.static_layout,
                                               highresstatic_expand_t=config.dataset.highresstatic_expand_t,
                                               mesostatic_expand_t=config.dataset.mesostatic_expand_t,
                                               meso_crop=None,
                                               fp16=False)
total_num_test_samples = len(earthnet_iid_testset)
print("The total number samples in the IID test subset is", total_num_test_samples)
```

--------------------------------

### Test Earthformer with Pretrained Checkpoint

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/earthnet_w_meso/README.md

Load a pretrained checkpoint to test the Earthformer model. This command bypasses the training configuration and directly uses saved weights.

```bash
MASTER_ADDR=localhost MASTER_PORT=10001 python train_cuboid_earthnet.py --gpus 2 --pretrained --save tmp_earthnet_w_meso
```

--------------------------------

### Create Conda Environment

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/README.md

Set up a new Conda environment for the project with Python 3.9. This helps manage dependencies and avoid conflicts.

```bash
conda create -n earthformer python=3.9
conda activate earthformer
```

--------------------------------

### Initialize CuboidSelfAttentionLayer with Local Strategy

Source: https://context7.com/amazon-science/earth-forecasting-transformer/llms.txt

Creates a core attention layer implementing the cuboid decomposition strategy with local attention. Requires input features and optional global vectors.

```python
import torch
from earthformer.cuboid_transformer.cuboid_transformer import CuboidSelfAttentionLayer

# Create cuboid self-attention layer with local strategy
attention_layer = CuboidSelfAttentionLayer(
    dim=128,                             # Input/output dimension
    num_heads=4,                         # Number of attention heads
    cuboid_size=(2, 7, 7),              # (T, H, W) size of each cuboid
    shift_size=(0, 0, 0),               # Shift for shifted window attention
    strategy=('l', 'l', 'l'),           # 'l'=local, 'd'=dilated for each axis
    padding_type='ignore',               # 'ignore', 'zeros', or 'nearest'
    qkv_bias=False,
    attn_drop=0.1,
    proj_drop=0.1,
    use_relative_pos=True,              # Use relative positional encoding
    use_global_vector=True,             # Enable global memory vectors
    use_global_self_attn=False,
)

# Input: (B, T, H, W, C)
x = torch.randn(2, 8, 28, 28, 128)
global_vectors = torch.randn(2, 8, 128)  # (B, num_global, C)

# Forward pass returns updated features and global vectors
output, updated_global = attention_layer(x, global_vectors)
print(f"Output shape: {output.shape}")  # (2, 8, 28, 28, 128)
```

--------------------------------

### Test Earthformer with Pretrained Checkpoint

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/enso/README.md

Loads a pretrained checkpoint to perform testing on the ICAR-ENSO dataset.

```bash
MASTER_ADDR=localhost MASTER_PORT=10001 python train_cuboid_enso.py --gpus 2 --pretrained --save tmp_enso
```

--------------------------------

### Train Earthformer on ICAR-ENSO

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/enso/README.md

Executes the training process using the specified configuration file and GPU count.

```bash
MASTER_ADDR=localhost MASTER_PORT=10001 python train_cuboid_enso.py --gpus 2 --cfg cfg.yaml --ckpt_name last.ckpt --save tmp_enso
```

--------------------------------

### Load EarthNet2021x IID Test Dataset

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/earthnet_w_meso/inference_tutorial_earthformer_earthnet2021x.ipynb

Loads a subset of the EarthNet2021x IID test split using a configuration file. Ensure the configuration file and dataset path are correctly specified.

```python
from omegaconf import OmegaConf
from earthformer.datasets.earthnet.earthnet21x_dataloader import EarthNet2021xLightningDataModule, EarthNet2021xTestDataset
from pytorch_lightning import Trainer, seed_everything, loggers as pl_loggers
from earthformer.datasets.earthnet.visualization import vis_earthnet_seq
import numpy as np
from pathlib import Path

config_file = "./earthformer_earthnet_v1.yaml"
config = OmegaConf.load(open(config_file, "r"))
in_len = config.layout.in_len
out_len = config.layout.out_len

seed = config.optim.seed
dataset_cfg = OmegaConf.to_object(config.dataset)

seed_everything(seed, workers=True)

micro_batch_size = 1

earthnet_iid_testset = EarthNet2021xTestDataset(subset_name="iid",
                                               data_dir="/Net/Groups/BGI/work_1/scratch/s3/earthnet/earthnet2021x/iid/",
                                               layout=config.dataset.layout,
                                               static_layout=config.dataset.static_layout,
                                               highresstatic_expand_t=config.dataset.highresstatic_expand_t,
                                               mesostatic_expand_t=config.dataset.mesostatic_expand_t,
                                               meso_crop=None,
                                               fp16=False)
total_num_test_samples = len(earthnet_iid_testset)
print("The total number samples in the IID test subset is", total_num_test_samples)
```

--------------------------------

### Download N-body MNIST Dataset

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/datasets/nbody/README.md

Use this command to download the N-body MNIST dataset from AWS S3. Ensure you are in the correct directory before running.

```bash
cd ROOT_DIR/earth-forecasting-transformer
python ./scripts/datasets/nbody/download_nbody_paper.py
```

--------------------------------

### Process and Visualize EarthNet2021x Sample

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/earthnet_w_meso/inference_tutorial_earthformer_earthnet2021x.ipynb

Processes a single EarthNet2021x sample by cleaning NaN values and clipping data, then visualizes the high-resolution dynamic input and output sequences.

```python
idx = 1

def process_sample(sample):
    # High-resolution Earth surface data. The channels are [blue, green, red, nir, cloud]
    highresdynamic = sample['highresdynamic']
    highresstatic = sample['highresstatic']

    # The meso-scale data. The channels are ["precipitation", "pressure", "temp mean", "temp min", "temp max"]
    mesodynamic = sample['mesodynamic']
    mesostatic = sample['mesostatic']

    highresdynamic = np.nan_to_num(highresdynamic, nan=0.0, posinf=1.0, neginf=0.0)
    highresdynamic = np.clip(highresdynamic, a_min=0.0, a_max=1.0)
    mesodynamic = np.nan_to_num(mesodynamic, nan=0.0)
    highresstatic = np.nan_to_num(highresstatic, nan=0.0)
    mesostatic = np.nan_to_num(mesostatic, nan=0.0)
    return highresdynamic, highresstatic, mesodynamic, mesostatic

highresdynamic, highresstatic, mesodynamic, mesostatic = process_sample(earthnet_iid_testset[idx])

highresdynamic_in, highresdynamic_out = highresdynamic[:in_len, ...], highresdynamic[in_len:, ...]
highresstatic_in, highresstatic_out = highresstatic[:in_len, ...], highresstatic[in_len:, ...]
mesodynamic_in, mesodynamic_out = mesodynamic[:in_len, ...], mesodynamic[in_len:, ...]
mesostatic_in, mesostatic_out = mesostatic[:in_len, ...], mesostatic[in_len:, ...]

fig_rgb = vis_earthnet_seq(context_np=np.expand_dims(highresdynamic_in, axis=0),
                       target_np=np.expand_dims(highresdynamic_out, axis=0),
                       pred_np=None,
                       layout='N' + config.dataset.layout)
```

--------------------------------

### Initialize CuboidTransformerModel for Precipitation Nowcasting

Source: https://context7.com/amazon-science/earth-forecasting-transformer/llms.txt

Initializes the main Earthformer model with specified input/output shapes, encoder-decoder configurations, and attention parameters. Suitable for precipitation nowcasting tasks.

```python
import torch
from earthformer.cuboid_transformer.cuboid_transformer import CuboidTransformerModel

# Initialize the Cuboid Transformer for precipitation nowcasting
model = CuboidTransformerModel(
    input_shape=(13, 384, 384, 1),      # (T_in, H, W, C) - 13 input frames
    target_shape=(12, 384, 384, 1),     # (T_out, H, W, C) - 12 output frames
    base_units=128,                      # Base channel dimension
    block_units=None,                    # Auto-compute (doubles at each level)
    scale_alpha=1.0,                     # Channel scaling factor
    # Encoder configuration
    enc_depth=[1, 1],                    # Depth at each encoder level
    enc_attn_patterns="axial",           # Attention pattern: 'axial', 'full', etc.
    # Decoder configuration
    dec_depth=[1, 1],                    # Depth at each decoder level
    dec_self_attn_patterns="axial",
    dec_cross_attn_patterns="cross_1x1",
    # Attention parameters
    num_heads=4,
    attn_drop=0.1,
    proj_drop=0.1,
    ffn_drop=0.1,
    # Global vectors for long-range dependencies
    num_global_vectors=8,
    use_dec_self_global=True,
    use_dec_cross_global=True,
    # Downsampling/upsampling
    downsample=2,
    downsample_type="patch_merge",
    upsample_type="upsample",
    # Initial spatial downsampling via stacked convolutions
    initial_downsample_type="stack_conv",
    initial_downsample_stack_conv_num_layers=3,
    initial_downsample_stack_conv_dim_list=[16, 64, 128],
    initial_downsample_stack_conv_downscale_list=[3, 2, 2],
    initial_downsample_stack_conv_num_conv_list=[2, 2, 2],
    # Positional embedding
    pos_embed_type="t+h+w",
    use_relative_pos=True,
    z_init_method="zeros",
)

# Forward pass: input shape (B, T_in, H, W, C)
x = torch.randn(2, 13, 384, 384, 1)  # Batch of 2 sequences
output = model(x)  # Shape: (2, 12, 384, 384, 1)
print(f"Input shape: {x.shape}, Output shape: {output.shape}")
```

--------------------------------

### Train Earthformer on EarthNet2021

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/earthnet_w_meso/README.md

Use this command to train the Earthformer model on the EarthNet2021 dataset. Ensure to adjust configurations in the `cfg.yaml` file and specify GPU usage and a checkpoint name.

```bash
MASTER_ADDR=localhost MASTER_PORT=10001 python train_cuboid_earthnet.py --gpus 2 --cfg cfg.yaml --ckpt_name last.ckpt --save tmp_earthnet_w_meso
```

--------------------------------

### Download EarthNet2021x Dataset

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/README.md

Downloads the EarthNet2021x dataset, which uses the .nc file format. This command uses the earthnet_toolket.

```python
import earthnet as en
en.download(dataset="earthnet2021x", splits="all", save_directory="./datasets/earthnet2021x")
```

--------------------------------

### Compute SEVIR Skill Scores

Source: https://context7.com/amazon-science/earth-forecasting-transformer/llms.txt

Calculates standard nowcasting metrics like CSI, POD, and Bias using the SEVIRSkillScore class.

```python
import torch
from earthformer.metrics.sevir import SEVIRSkillScore

# Initialize skill score metric
skill_score = SEVIRSkillScore(
    layout="NTHWC",
    mode="0",                                # "0": aggregate all, "1": per-timestep
    seq_len=12,
    preprocess_type="sevir",
    threshold_list=(16, 74, 133, 160, 181, 219),  # VIL thresholds (mm)
    metrics_list=("csi", "pod", "sucr", "bias"),
    eps=1e-4,
)

# Update with predictions and ground truth
pred = torch.randn(4, 12, 384, 384, 1)     # Predicted sequence
target = torch.randn(4, 12, 384, 384, 1)   # Ground truth sequence

skill_score.update(pred, target)

# Compute final scores
scores = skill_score.compute()

# Access scores by threshold and metric
for threshold in [16, 74, 133]:
    csi = scores[threshold]['csi']
    pod = scores[threshold]['pod']
    print(f"Threshold {threshold}: CSI={csi:.4f}, POD={pod:.4f}")

# Average across all thresholds
avg_csi = scores['avg']['csi']
print(f"Average CSI: {avg_csi:.4f}")
```

--------------------------------

### Load Earthformer Pretrained Model

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/earthnet_w_meso/inference_tutorial_earthformer_earthnet2021.ipynb

Loads a pretrained Earthformer model for EarthNet2021. Ensure the save directory exists and the download utility is available. The model is loaded onto the CPU initially.

```python
import torch
import os
from train_cuboid_earthnet import CuboidEarthNet2021PLModule
from earthformer.utils.utils import download

save_dir = "./experiments"

pl_module = CuboidEarthNet2021PLModule(
    total_num_steps=None,
    save_dir="./experiments",
    oc_file=config_file
)

pretrained_checkpoint_url = "https://earthformer.s3.amazonaws.com/pretrained_checkpoints/earthformer_earthnet2021.pt"
local_checkpoint_path = os.path.join(save_dir, "earthformer_earthnet2021.pt")
download(url=pretrained_checkpoint_url, path=local_checkpoint_path)

state_dict = torch.load(local_checkpoint_path, map_location=torch.device("cpu"))
pl_module.torch_nn_module.load_state_dict(state_dict=state_dict)
```

--------------------------------

### Download EarthNet2021 Dataset

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/README.md

Downloads the EarthNet2021 dataset using the earthnet_toolket. Specify the dataset name and save directory.

```python
import earthnet as en
en.download(dataset="earthnet2021", splits="all", save_directory="./datasets/earthnet2021")
```

--------------------------------

### Download EarthNet2021 IID Subset

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/earthnet_w_meso/inference_tutorial_earthformer_earthnet2021.ipynb

Downloads the IID test subset of EarthNet2021. It customizes the download progress bar and limits the download to a single file for demonstration purposes. Ensure the target directory exists.

```python
import earthnet.download_links
import earthnet.download
from earthnet.download import DownloadProgressBar
import earthnet as en
import os

class NewDownloadProgressBar(DownloadProgressBar):
    def __init__(self, **kwargs):
        if "disable" in kwargs:
            kwargs["disable"] = True
            super().__init__(**kwargs)
        else:
            super().__init__(disable=True, **kwargs)

earthnet.download.DownloadProgressBar = NewDownloadProgressBar
earthnet.download_links.DOWNLOAD_LINKS["iid"] = earthnet.download_links.DOWNLOAD_LINKS["iid"][:1]
print(earthnet.download_links.DOWNLOAD_LINKS["iid"])
en.Downloader.get("./datasets/earthnet2021", "iid")
```

--------------------------------

### Visualize Earthformer Predictions (RGB)

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/earthnet_w_meso/inference_tutorial_earthformer_earthnet2021.ipynb

Visualizes the predicted sequence against the target and context sequences using RGB channels. The layout is determined by the configuration. This function is useful for qualitative assessment of the model's performance.

```python
fig_rgb = vis_earthnet_seq(context_np=np.expand_dims(highresdynamic_in, axis=0),
                           target_np=np.expand_dims(highresdynamic_out, axis=0),
                           pred_np=pred_seq_np,
                           layout='N' + config.dataset.layout)
```

--------------------------------

### Visualize Earthformer Predictions (NDVI)

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/earthnet_w_meso/inference_tutorial_earthformer_earthnet2021.ipynb

Visualizes the predicted sequence against the target and context sequences specifically for the NDVI variable. This allows for a focused evaluation of the model's ability to forecast Normalized Difference Vegetation Index.

```python
fig_ndvi = vis_earthnet_seq(context_np=np.expand_dims(highresdynamic_in, axis=0),
                           target_np=np.expand_dims(highresdynamic_out, axis=0),
                           pred_np=pred_seq_np,
                           layout='N' + config.dataset.layout,
                           variable="ndvi")
```

--------------------------------

### Visualize NDVI Sequence

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/earthnet_w_meso/inference_tutorial_earthformer_earthnet2021.ipynb

Visualizes the Normalized Difference Vegetation Index (NDVI) sequence for a given EarthNet2021 sample. NDVI quantifies vegetation using near-infrared and red light reflectance.

```python
fig_ndvi = vis_earthnet_seq(context_np=np.expand_dims(highresdynamic_in, axis=0),
                            target_np=np.expand_dims(highresdynamic_out, axis=0),
                            pred_np=None,
                            layout='N' + config.dataset.layout,
                            variable="ndvi")
```

--------------------------------

### Plot NDVI Predictions from NetCDF using xarray

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/earthnet_w_meso/inference_tutorial_earthformer_earthnet2021x.ipynb

Opens a NetCDF file containing Earthformer predictions and plots the NDVI predictions over time using xarray's imshow functionality. The plot is wrapped into 10 columns for better visualization.

```python
import xarray as xr
mc = xr.open_dataset("/Net/Groups/BGI/people/vbenson/EarthNet/earthnet-models-pytorch/experiments/en21x/persistence/preds/iid/30TYR/30TYR_2018-06-04_2018-10-31_313_441_4409_4537_4_84_68_148.nc")
mc.ndvi_pred.plot.imshow(col = "time", col_wrap = 10, vmin = 0.0, vmax = 1.0)
```

--------------------------------

### Perform Inference with Earthformer

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/earthnet_w_meso/inference_tutorial_earthformer_earthnet2021.ipynb

Moves the loaded Earthformer model to CUDA, sets it to evaluation mode, and performs inference on input data. The predictions and loss are then detached and moved to the CPU for further processing.

```python
pl_module.torch_nn_module.cuda()
pl_module.torch_nn_module.eval()

with torch.no_grad():
    pred_seq, loss, in_seq, target_seq, mask = pl_module({"highresdynamic": torch.tensor(np.expand_dims(highresdynamic, axis=0)).cuda(), 
                                                          "highresstatic": torch.tensor(np.expand_dims(highresstatic, axis=0)).cuda(),
                                                          "mesodynamic": torch.tensor(np.expand_dims(mesodynamic, axis=0)).cuda(),
                                                          "mesostatic": torch.tensor(np.expand_dims(mesostatic, axis=0)).cuda()})
    pred_seq_np = pred_seq.detach().cpu().numpy()

    print("loss=", loss.detach().cpu().numpy())
```

--------------------------------

### Process and Visualize Earth Surface Data

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/scripts/cuboid_transformer/earthnet_w_meso/inference_tutorial_earthformer_earthnet2021.ipynb

Processes a sample from the EarthNet2021 test set, handling NaNs and clipping values. Visualizes the high-resolution Earth surface data sequence.

```python
idx = 1

def process_sample(sample):
    # High-resolution Earth surface data. The channels are [blue, green, red, nir, cloud]
    highresdynamic = sample['highresdynamic']
    highresstatic = sample['highresstatic']

    # The meso-scale data. The channels are ["precipitation", "pressure", "temp mean", "temp min", "temp max"]
    mesodynamic = sample['mesodynamic']
    mesostatic = sample['mesostatic']

    highresdynamic = np.nan_to_num(highresdynamic, nan=0.0, posinf=1.0, neginf=0.0)
    highresdynamic = np.clip(highresdynamic, a_min=0.0, a_max=1.0)
    mesodynamic = np.nan_to_num(mesodynamic, nan=0.0)
    highresstatic = np.nan_to_num(highresstatic, nan=0.0)
    mesostatic = np.nan_to_num(mesostatic, nan=0.0)
    return highresdynamic, highresstatic, mesodynamic, mesostatic

earthnet_iid_testset = None # Placeholder for actual dataset
config = None # Placeholder for actual config
in_len = 10 # Placeholder for actual input length

highresdynamic, highresstatic, mesodynamic, mesostatic = process_sample(earthnet_iid_testset[idx])

highresdynamic_in, highresdynamic_out = highresdynamic[:in_len, ...], highresdynamic[in_len:, ...]
highresstatic_in, highresstatic_out = highresstatic[:in_len, ...], highresstatic[in_len:, ...]
mesodynamic_in, mesodynamic_out = mesodynamic[:in_len, ...], mesodynamic[in_len:, ...]
mesostatic_in, mesostatic_out = mesostatic[:in_len, ...], mesostatic[in_len:, ...]

fig_rgb = vis_earthnet_seq(context_np=np.expand_dims(highresdynamic_in, axis=0),
                       target_np=np.expand_dims(highresdynamic_out, axis=0),
                       pred_np=None,
                       layout='N' + config.dataset.layout)
```

--------------------------------

### Extract ICAR-ENSO Dataset

Source: https://github.com/amazon-science/earth-forecasting-transformer/blob/main/README.md

Extracts the ICAR-ENSO dataset after downloading the zip file. The zip file should be placed in the 'datasets/' directory.

```bash
unzip datasets/enso_round1_train_20210201.zip -d datasets/icar_enso_2021
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.