Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+resources/id4nvwbetz/AtmoRep_id4nvwbetz.mod filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,8 @@
+__pycache__/
+*.py[cod]
+*.pdf
+result/
+weight/*
+!weight/.gitkeep
+.pytest_cache/
+.ipynb_checkpoints/

README.md

@@ -0,0 +1,171 @@
+---
+license: other
+language:
+- en
+- zh
+tags:
+- OneScience
+- Earth Science
+- Weather Forecasting
+- ERA5
+- Representation Learning
+frameworks: PyTorch
+datasets:
+- OneScience/ERA5
+---
+
+<p align="center">
+  <strong><span style="font-size: 30px;">AtmoRep</span></strong>
+</p>
+
+# Model Overview
+
+AtmoRep is a stochastic atmospheric dynamics model based on large-scale representation learning that captures the distribution of atmospheric states through masked-token training and ensemble outputs.
+
+Paper: *AtmoRep: A stochastic model of atmosphere dynamics using large scale representation learning*
+
+https://arxiv.org/abs/2308.13280
+
+# Model Description
+
+This directory retains the official vorticity single-field model weights, configuration, and normalization, and provides a tiny AtmoRep-style model for local training and inference verification.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Official Resource Validation | Load official `.mod` weights and verify configuration. |
+| Local Rapid Verification | Run masked-token training and inference with the tiny model. |
+| ERA5 Atmospheric Representation Learning | Subsequently interface with official GRIB or Zarr data. |
+
+# Usage
+
+## 1. OneCode
+
+[Click to experience intelligent one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation & Usage
+
+**Hardware Requirements**
+
+- The tiny model runs on CPU.
+- GPU is recommended for the official model and real-data inference.
+
+### Download the Model Package
+
+```bash
+hf download --model OneScience-Group/AtmoRep --local-dir ./AtmoRep
+cd AtmoRep
+```
+
+### Set Up the Runtime Environment
+
+**DCU Environment**
+
+```bash
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+**GPU Environment**
+
+```bash
+conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+If dependencies are missing from the official paths, install them with:
+
+```bash
+pip install zarr wandb cfgrib xarray dask netCDF4 torchinfo
+```
+
+### Weights & Data
+
+The current directory preserves:
+
+```text
+resources/id4nvwbetz/AtmoRep_id4nvwbetz.mod
+resources/id4nvwbetz/model_id4nvwbetz.json
+resources/data/normalization/vorticity/
+```
+
+Before re-downloading official resources please note: the current vendor snapshot is not a Git checkout, and the existing download script cannot re-clone into a non-empty `vendor/atmorep-official`. The resources bundled with this package do not require re-downloading; the script is only suitable for an empty target directory.
+
+```bash
+bash scripts/download_official_resources.sh .
+```
+
+### Tiny Training
+
+```bash
+python scripts/train.py
+```
+
+This command performs multi-epoch masked-token training using independent train/validation fake Datasets, including DataLoader, AdamW, validation, learning rate scheduling, early stopping, best/latest checkpointing, and training history. Default parameters reside in `conf/config.yaml`.
+
+Resuming training:
+
+```bash
+python scripts/train.py --resume weight/training/latest.pth --epochs 20
+```
+
+Training artifacts are `weight/training/latest.pth`, `best.pth`, and `history.json`; the inference-compatible weights `weight/tiny_atmorep.pth` are also updated. Each Dataset sample consists of `fields [T,V,H,W]` and a non-empty `mask [N]`; train/validation use different seeds. This is a complete training pipeline for the tiny model, not a reproduction of the paper's 3.5-billion-parameter official model training.
+
+### Tiny Inference
+
+```bash
+python scripts/inference.py
+```
+
+Inference results are saved as:
+
+```text
+result/prediction.pt
+result/target.pt
+```
+
+These include the ensemble, ensemble mean, ensemble std, and mask.
+
+### Result Inspection
+
+```bash
+python scripts/result.py
+```
+
+This command generates `result/metrics.json` and `result/comparison.png`. The reported metrics are ensemble/mean/spread RMSE in normalized token space, not the paper's physical-unit RMSE, ACC, CRPS, or spread-skill.
+
+### Paper vs. Current Implementation I/O
+
+| Item | Paper / Official Model | Tiny Smoke Model |
+| --- | --- | --- |
+| Input | ERA5 local 4D neighborhoods, 5 mode levels, multiple physical fields | `[B,4,1,8,8]` single-field random tensor |
+| Token | Variable-correlated 4D tokens with absolute space-time and level conditioning | `1×4×4` patches, 16 tokens, relative coordinates with single-level conditioning |
+| Output | Multi-head ensembles supporting reconstruction, nowcasting, and interpolation | 4-member masked-token ensemble |
+| Training | Large-scale masked-token distribution learning | Multi-epoch Dataset training with independent validation and checkpoint resumption |
+| Weights | `resources/id4nvwbetz` official vorticity weights | `weight/tiny_atmorep.pth`; the two are mutually incompatible |
+
+The complete tiny execution flow is `train.py -> inference.py -> result.py`. The result analysis reads `weight/training/history.json` when present; inference output includes `ensemble`, `ensemble_mean`, `ensemble_std`, `mask`, and `target`. Random data is generated by the Dataset by index; there is currently no independent fake dataset that can be used with the official Zarr sampler. The model package is distributed without local training weights or `result/` artifacts — these are created at the paths described above after running the commands.
+
+### Official Real-Data Inference
+
+A directly runnable official real-data inference command is not yet available. It further requires ERA5 vorticity GRIB/Zarr data, the ecCodes environment, and parameterization of the file paths in the official `evaluate.py`.
+
+### Real Data
+
+The official vorticity model requires ERA5 vorticity, model levels 96/105/114/123/137, hourly temporal axes, and 0.25° global GRIB/Zarr data.
+
+# OneScience Official Information
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citation & License
+
+- Official code is MIT License.
+- Official model weights are declared CC BY 4.0.
+- ERA5 is subject to Copernicus/ECMWF data terms.

conf/config.yaml

@@ -0,0 +1,13 @@
+training:
+  epochs: 10
+  batch_size: 8
+  train_samples: 64
+  validation_samples: 16
+  learning_rate: 0.001
+  weight_decay: 0.000001
+  mask_fraction: 0.25
+  patience: 3
+  seed: 2026
+  num_workers: 0
+  device: cpu
+  checkpoint_dir: weight/training

configuration.json

@@ -0,0 +1 @@
+{"framework":"Pytorch","task":"atmospheric-representation-learning","implementation":"tiny-smoke-and-official-resources"}

model/__init__.py

@@ -0,0 +1,3 @@
+from .tiny_atmorep import TinyAtmoRep, ensemble_statistical_loss
+
+__all__ = ["TinyAtmoRep", "ensemble_statistical_loss"]

model/fake_data.py

@@ -0,0 +1,23 @@
+"""Deterministic fake dataset for the tiny AtmoRep training pipeline."""
+
+import torch
+from torch.utils.data import Dataset
+
+from .tiny_atmorep import TinyAtmoRepConfig
+
+
+class FakeAtmoRepDataset(Dataset):
+    def __init__(self, config: TinyAtmoRepConfig, samples: int, seed: int, mask_fraction: float) -> None:
+        self.config, self.samples, self.seed = config, samples, seed
+        self.num_tokens = (config.input_shape[0] // config.patch_shape[0]) * (config.input_shape[2] // config.patch_shape[1]) * (config.input_shape[3] // config.patch_shape[2])
+        self.mask_count = max(1, round(self.num_tokens * mask_fraction))
+
+    def __len__(self):
+        return self.samples
+
+    def __getitem__(self, index):
+        generator = torch.Generator().manual_seed(self.seed + index)
+        fields = torch.randn(self.config.input_shape, generator=generator)
+        mask = torch.zeros(self.num_tokens, dtype=torch.bool)
+        mask[torch.randperm(self.num_tokens, generator=generator)[:self.mask_count]] = True
+        return fields, mask

model/tiny_atmorep.py

@@ -0,0 +1,132 @@
+"""Independent tiny single-field AtmoRep-style fallback for pipeline validation."""
+
+from __future__ import annotations
+
+import math
+from dataclasses import asdict, dataclass
+
+import torch
+from torch import Tensor, nn
+
+
+@dataclass(frozen=True)
+class TinyAtmoRepConfig:
+    input_shape: tuple[int, int, int, int] = (4, 1, 8, 8)
+    patch_shape: tuple[int, int, int] = (1, 4, 4)
+    embed_dim: int = 32
+    num_heads: int = 4
+    num_layers: int = 2
+    ensemble_size: int = 4
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+class TinyAtmoRep(nn.Module):
+    """Masked-token transformer with four-dimensional token conditioning.
+
+    Inputs use ``[batch, time, variable, latitude, longitude]``. This fallback
+    intentionally supports one field only; level is supplied as token metadata.
+    """
+
+    def __init__(self, config: TinyAtmoRepConfig | None = None) -> None:
+        super().__init__()
+        self.config = config or TinyAtmoRepConfig()
+        time, variables, height, width = self.config.input_shape
+        pt, ph, pw = self.config.patch_shape
+        if variables != 1:
+            raise ValueError("TinyAtmoRep is a single-field fallback (V must equal 1)")
+        if time % pt or height % ph or width % pw:
+            raise ValueError("input_shape must be divisible by patch_shape")
+
+        self.grid_shape = (time // pt, height // ph, width // pw)
+        self.patch_dim = pt * ph * pw
+        self.patch_embed = nn.Conv3d(
+            1, self.config.embed_dim, kernel_size=self.config.patch_shape,
+            stride=self.config.patch_shape,
+        )
+        self.condition_embed = nn.Sequential(
+            nn.Linear(4, self.config.embed_dim), nn.GELU(),
+            nn.Linear(self.config.embed_dim, self.config.embed_dim),
+        )
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, self.config.embed_dim))
+        layer = nn.TransformerEncoderLayer(
+            d_model=self.config.embed_dim,
+            nhead=self.config.num_heads,
+            dim_feedforward=4 * self.config.embed_dim,
+            dropout=0.0,
+            activation="gelu",
+            batch_first=True,
+            norm_first=True,
+        )
+        self.encoder = nn.TransformerEncoder(layer, self.config.num_layers)
+        self.ensemble_heads = nn.ModuleList(
+            nn.Linear(self.config.embed_dim, self.patch_dim)
+            for _ in range(self.config.ensemble_size)
+        )
+        nn.init.normal_(self.mask_token, std=0.02)
+
+    @property
+    def num_tokens(self) -> int:
+        return math.prod(self.grid_shape)
+
+    def token_conditions(self, batch_size: int, level: float, device: torch.device) -> Tensor:
+        """Return normalized [time, level, latitude, longitude] per token."""
+        nt, nh, nw = self.grid_shape
+        axes = [torch.linspace(-1.0, 1.0, n, device=device) for n in (nt, nh, nw)]
+        time, latitude, longitude = torch.meshgrid(*axes, indexing="ij")
+        model_level = torch.full_like(time, float(level) / 137.0)
+        conditions = torch.stack((time, model_level, latitude, longitude), dim=-1)
+        return conditions.reshape(1, self.num_tokens, 4).expand(batch_size, -1, -1)
+
+    def tokenize(self, fields: Tensor) -> Tensor:
+        self._validate_fields(fields)
+        volume = fields.permute(0, 2, 1, 3, 4)
+        pt, ph, pw = self.config.patch_shape
+        patches = volume.unfold(2, pt, pt).unfold(3, ph, ph).unfold(4, pw, pw)
+        return patches.permute(0, 2, 3, 4, 1, 5, 6, 7).reshape(
+            fields.shape[0], self.num_tokens, self.patch_dim
+        )
+
+    def forward(self, fields: Tensor, mask: Tensor, level: float = 137.0) -> Tensor:
+        self._validate_fields(fields)
+        if mask.shape != (fields.shape[0], self.num_tokens) or mask.dtype != torch.bool:
+            raise ValueError(f"mask must be bool [B, {self.num_tokens}]")
+        tokens = self.patch_embed(fields.permute(0, 2, 1, 3, 4)).flatten(2).transpose(1, 2)
+        conditions = self.token_conditions(fields.shape[0], level, fields.device)
+        tokens = tokens + self.condition_embed(conditions)
+        tokens = torch.where(mask.unsqueeze(-1), self.mask_token.expand_as(tokens), tokens)
+        encoded = self.encoder(tokens)
+        return torch.stack([head(encoded) for head in self.ensemble_heads], dim=1)
+
+    def _validate_fields(self, fields: Tensor) -> None:
+        expected = self.config.input_shape
+        if fields.ndim != 5 or tuple(fields.shape[1:]) != expected:
+            raise ValueError(f"fields must have shape [B, {expected}], got {tuple(fields.shape)}")
+
+
+def ensemble_statistical_loss(
+    predictions: Tensor,
+    targets: Tensor,
+    mask: Tensor,
+    statistical_weight: float = 0.1,
+) -> tuple[Tensor, dict[str, Tensor]]:
+    """Combine masked ensemble MSE with ensemble mean/spread statistics."""
+    if predictions.ndim != 4 or targets.ndim != 3:
+        raise ValueError("predictions must be [B,E,N,P] and targets [B,N,P]")
+    selected = mask[:, None, :, None].expand_as(predictions)
+    expanded_targets = targets[:, None].expand_as(predictions)
+    ensemble_mse = (predictions[selected] - expanded_targets[selected]).square().mean()
+
+    ensemble_mean = predictions.mean(dim=1)
+    ensemble_std = predictions.std(dim=1, unbiased=False)
+    target_std = targets.std(dim=-1, unbiased=False, keepdim=True).expand_as(targets)
+    masked = mask.unsqueeze(-1).expand_as(targets)
+    mean_loss = (ensemble_mean[masked] - targets[masked]).square().mean()
+    spread_loss = (ensemble_std[masked] - target_std[masked]).square().mean()
+    stats_loss = mean_loss + spread_loss
+    total = ensemble_mse + statistical_weight * stats_loss
+    return total, {
+        "ensemble_mse": ensemble_mse.detach(),
+        "statistical": stats_loss.detach(),
+    }

resources/data/normalization/vorticity/global_normalization_mean_var_vorticity_ml105.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6d509500b9b870f44f12e998e7e50ab3c5cb0646867d5c2789b63af51ad0836b
+size 8640

resources/data/normalization/vorticity/global_normalization_mean_var_vorticity_ml114.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e2fac57f8d31533a7a4c825e2240c6751a816be93ab79ebff7007368749b537e
+size 8640

resources/data/normalization/vorticity/global_normalization_mean_var_vorticity_ml123.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c5f5d843b9bbf5d876d1faa287ba09fa2556ecf9b3437bc6289f159d34a288de
+size 8640

resources/data/normalization/vorticity/global_normalization_mean_var_vorticity_ml137.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a980bfa8920ae1ddeca783852811b1d49aad80de695e57baa5dcab1b35249e0f
+size 8640

resources/data/normalization/vorticity/global_normalization_mean_var_vorticity_ml96.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e4344b3f123303183e2a2af5d58ab001982406225256e75ca158153a0d41f5a7
+size 8640

resources/id4nvwbetz/AtmoRep_id4nvwbetz.mod

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fec447ac33bf6774dc552689199822f2803d772783a473a9452ce558c5d4ff55
+size 2719496689

resources/id4nvwbetz/model_id4nvwbetz.json

@@ -0,0 +1 @@
+{"num_accs_per_task": 1, "with_hvd": true, "hvd_rank": 0, "hvd_size": 16, "back_passes_per_step": 4, "comment": "", "file_format": "grib", "data_dir": "../data/", "level_type": "ml", "fields": [["vorticity", [1, 2048, [], 0], [96, 105, 114, 123, 137], [12, 6, 12], [3, 9, 9], [0.5, 0.9, 0.1, 0.05]]], "fields_prediction": [["vorticity", 1.0]], "fields_targets": [], "years_train": [1980, 1981, 1982, 1983, 1984, 1985, 1986, 1987, 1988, 1989, 1990, 1991, 1992, 1993, 1994, 1995, 1996, 1997, 1998, 1999, 2000, 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008, 2009, 2010, 2011, 2012, 2013, 2014, 2015, 2016, 2017], "years_test": [2018], "month": null, "geo_range_sampling": [[-90.0, 90.0], [0.0, 360.0]], "time_sampling": 1, "data_smoothing": 0, "file_shape": [-1, 721, 1440], "num_t_samples": 744, "num_patches_per_t_train": 4, "num_patches_per_t_test": 1, "torch_seed": 5371822430435701752, "batch_size_test": 24, "batch_size_start": 6, "batch_size_max": 6, "batch_size_delta": 8, "num_epochs": 2048, "num_files_train": 2, "num_files_test": 2, "loader_num_workers": 8, "size_token_info": 8, "size_token_info_net": 16, "grad_checkpointing": true, "with_cls": false, "with_layernorm": true, "coupling_num_heads_per_field": 1, "dropout_rate": 0.05, "learnable_mask": false, "with_qk_lnorm": true, "num_neighborhoods": 1, "encoder_num_layers": 10, "encoder_num_heads": 16, "encoder_num_mlp_layers": 2, "encoder_att_type": "dense", "decoder_num_layers": 10, "decoder_num_heads": 16, "decoder_num_mlp_layers": 2, "decoder_self_att": false, "decoder_cross_att_ratio": 0.5, "decoder_cross_att_rate": 1.0, "decoder_att_type": "dense", "net_tail_num_nets": 16, "net_tail_num_layers": 0, "ensemble_weighted": false, "ensemble_weighted_T": 50.0, "losses": ["mse_ensemble", "stats"], "lr_start": 1e-05, "lr_max": 8e-05, "lr_min": 6.000000000000001e-05, "weight_decay": 0.025, "lr_decay_rate": 1.025, "lr_start_epochs": 3, "BERT_strategy": "BERT", "BERT_window": false, "BERT_fields_synced": false, "BERT_mr_max": 2, "log_test_num_ranks": 0, "save_grads": false, "profile": false, "test_initial": true, "rng_seed": null, "with_wandb": true, "slurm_job_id": "7749847", "wandb_id": "4nvwbetz", "lat_sampling_weighted": true}

scripts/download_official_resources.sh

@@ -0,0 +1,20 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+ROOT="${1:?usage: download_official_resources.sh ROOT_DIR}"
+COMMIT="055f858e68f5e0151eb6fece9f6b574d3da4af8d"
+REPOSITORY="https://github.com/clessig/atmorep.git"
+MODEL_URL="https://datapub.fz-juelich.de/atmorep/models/model_id4nvwbetz.tar.gz"
+
+mkdir -p "$ROOT/vendor" "$ROOT/resources"
+if [[ ! -d "$ROOT/vendor/atmorep-official/.git" ]]; then
+  git clone --no-checkout "$REPOSITORY" "$ROOT/vendor/atmorep-official"
+fi
+git -C "$ROOT/vendor/atmorep-official" fetch origin "$COMMIT"
+git -C "$ROOT/vendor/atmorep-official" checkout --detach "$COMMIT"
+curl -fL --retry 3 --continue-at - -o "$ROOT/resources/model_id4nvwbetz.tar.gz" "$MODEL_URL"
+gzip -t "$ROOT/resources/model_id4nvwbetz.tar.gz"
+tar -xzf "$ROOT/resources/model_id4nvwbetz.tar.gz" -C "$ROOT/resources"
+sha256sum "$ROOT/resources/model_id4nvwbetz.tar.gz" \
+  "$ROOT/resources/id4nvwbetz/AtmoRep_id4nvwbetz.mod" \
+  "$ROOT/resources/id4nvwbetz/model_id4nvwbetz.json"

scripts/inference.py

@@ -0,0 +1,57 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+import torch
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+from model.tiny_atmorep import TinyAtmoRep, TinyAtmoRepConfig
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint", type=Path, default=ROOT / "weight" / "tiny_atmorep.pth")
+    parser.add_argument("--output", type=Path, default=ROOT / "result" / "prediction.pt")
+    parser.add_argument("--seed", type=int, default=17)
+    args = parser.parse_args()
+
+    payload = torch.load(args.checkpoint, map_location="cpu", weights_only=True)
+    config = TinyAtmoRepConfig(**payload["config"])
+    model = TinyAtmoRep(config)
+    model.load_state_dict(payload["model"])
+    model.eval()
+    torch.manual_seed(args.seed)
+    fields = torch.randn(1, *config.input_shape)
+    mask = torch.zeros(1, model.num_tokens, dtype=torch.bool)
+    mask[:, 1::4] = True
+    with torch.inference_mode():
+        ensemble = model(fields, mask, level=137.0)
+    target = model.tokenize(fields)
+    result = {
+        "ensemble": ensemble,
+        "ensemble_mean": ensemble.mean(dim=1),
+        "ensemble_std": ensemble.std(dim=1, unbiased=False),
+        "mask": mask,
+        "target": target,
+        "input_shape": tuple(fields.shape),
+    }
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    torch.save(target, args.output.parent / "target.pt")
+    torch.save(result, args.output)
+    print(json.dumps({
+        "output": str(args.output),
+        "ensemble_shape": list(ensemble.shape),
+        "mean_shape": list(result["ensemble_mean"].shape),
+        "finite": bool(torch.isfinite(ensemble).all()),
+        "bytes": args.output.stat().st_size,
+    }, indent=2))
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,46 @@
+"""Show the saved tiny AtmoRep training and inference summaries."""
+
+import json
+from pathlib import Path
+
+import torch
+import matplotlib
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+if __name__ == "__main__":
+    history_path = ROOT / "weight" / "training" / "history.json"
+    train_history = json.loads(history_path.read_text()) if history_path.exists() else []
+    prediction = torch.load(ROOT / "result" / "prediction.pt", map_location="cpu", weights_only=True)
+    target = torch.load(ROOT / "result" / "target.pt", map_location="cpu", weights_only=True)
+    ensemble = prediction["ensemble"]
+    target = target["target"] if isinstance(target, dict) else target
+    mean = ensemble.mean(dim=1)
+    std = ensemble.std(dim=1, unbiased=False)
+    ensemble_rmse = torch.sqrt((ensemble - target[:, None]).square().mean())
+    mean_rmse = torch.sqrt((mean - target).square().mean())
+    target_std = target.std(dim=-1, unbiased=False, keepdim=True).expand_as(target)
+    spread_rmse = torch.sqrt((std - target_std).square().mean())
+    summary = {
+        "train_history": train_history,
+        "ensemble_shape": list(ensemble.shape),
+        "target_shape": list(target.shape),
+        "ensemble_rmse": float(ensemble_rmse),
+        "mean_rmse": float(mean_rmse),
+        "spread_rmse": float(spread_rmse),
+        "finite": bool(torch.isfinite(ensemble).all()),
+    }
+    (ROOT / "result" / "metrics.json").write_text(json.dumps(summary, indent=2) + "\n")
+    fig, axes = plt.subplots(1, 4, figsize=(12, 3))
+    for ax, data, title in zip(axes, (target[0, 0], mean[0, 0], (mean - target)[0, 0].abs(), std[0, 0]), ("Target", "Mean", "Abs error", "Spread")):
+        image = ax.imshow(data.reshape(4, 4).numpy(), cmap="viridis")
+        ax.set_title(title)
+        plt.colorbar(image, ax=ax, shrink=0.75)
+    plt.tight_layout()
+    plt.savefig(ROOT / "result" / "comparison.png", dpi=150)
+    plt.close()
+    print(json.dumps(summary, indent=2))

scripts/train.py

@@ -0,0 +1,50 @@
+"""Train tiny AtmoRep with train/validation epochs and resumable checkpoints."""
+
+import argparse, json, random, sys
+from pathlib import Path
+import numpy as np
+import torch, yaml
+from torch.utils.data import DataLoader
+
+ROOT = Path(__file__).resolve().parents[1]; sys.path.insert(0, str(ROOT))
+from model.fake_data import FakeAtmoRepDataset
+from model.tiny_atmorep import TinyAtmoRep, TinyAtmoRepConfig, ensemble_statistical_loss
+
+
+def args():
+    p0=argparse.ArgumentParser(add_help=False); p0.add_argument("--config",type=Path,default=ROOT/"conf/config.yaml"); k,_=p0.parse_known_args(); c=yaml.safe_load(k.config.read_text())["training"]
+    p=argparse.ArgumentParser(parents=[p0]);
+    for name, typ in (("epochs", int), ("batch-size", int), ("train-samples", int), ("validation-samples", int), ("seed", int), ("patience", int), ("num-workers", int), ("learning-rate", float), ("weight-decay", float), ("mask-fraction", float)):
+        p.add_argument("--" + name, type=typ, default=c[name.replace("-", "_")])
+    p.add_argument("--checkpoint-dir",type=Path,default=ROOT/c["checkpoint_dir"]); p.add_argument("--resume",type=Path); p.add_argument("--device",choices=("cpu","cuda"),default=c["device"]); return p.parse_args()
+
+
+def run_epoch(model, loader, device, optimizer=None):
+    model.train(optimizer is not None); totals=[]
+    context=torch.enable_grad() if optimizer else torch.inference_mode()
+    with context:
+        for fields,mask in loader:
+            fields,mask=fields.to(device),mask.to(device); pred=model(fields,mask,level=137.0); target=model.tokenize(fields); loss,_=ensemble_statistical_loss(pred,target,mask)
+            if optimizer: optimizer.zero_grad(set_to_none=True); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),1.0); optimizer.step()
+            totals.append(float(loss.detach()))
+    return float(np.mean(totals))
+
+
+def main():
+    a=args(); device=torch.device(a.device); random.seed(a.seed); np.random.seed(a.seed); torch.manual_seed(a.seed); a.checkpoint_dir.mkdir(parents=True,exist_ok=True)
+    config=TinyAtmoRepConfig(); model=TinyAtmoRep(config).to(device); opt=torch.optim.AdamW(model.parameters(),lr=a.learning_rate,weight_decay=a.weight_decay); sched=torch.optim.lr_scheduler.ReduceLROnPlateau(opt,patience=max(1,a.patience//2),factor=.5)
+    train=DataLoader(FakeAtmoRepDataset(config,a.train_samples,a.seed,a.mask_fraction),batch_size=a.batch_size,shuffle=True,num_workers=a.num_workers); val=DataLoader(FakeAtmoRepDataset(config,a.validation_samples,a.seed+100000,a.mask_fraction),batch_size=a.batch_size,num_workers=a.num_workers)
+    start=0; best=float("inf"); history=[]
+    if a.resume:
+        q=torch.load(a.resume,map_location=device,weights_only=False); model.load_state_dict(q["model"]); opt.load_state_dict(q["optimizer"]); sched.load_state_dict(q["scheduler"]); start=q["epoch"]+1; best=q["best_val_loss"]; history=q["history"]
+    stale=0
+    for epoch in range(start,a.epochs):
+        tr=run_epoch(model,train,device,opt); va=run_epoch(model,val,device); sched.step(va); row={"epoch":epoch,"train_loss":tr,"validation_loss":va,"learning_rate":opt.param_groups[0]["lr"]}; history.append(row); improved=va<best; best=min(best,va); stale=0 if improved else stale+1
+        train_config = {key: str(value) if isinstance(value, Path) else value for key, value in vars(a).items()}
+        payload={"model":model.state_dict(),"optimizer":opt.state_dict(),"scheduler":sched.state_dict(),"config":config.to_dict(),"epoch":epoch,"step":len(train)*(epoch+1),"best_val_loss":best,"history":history,"train_config":train_config}
+        torch.save(payload,a.checkpoint_dir/"latest.pth"); torch.save(payload,ROOT/"weight/tiny_atmorep.pth")
+        if improved: torch.save(payload,a.checkpoint_dir/"best.pth")
+        (a.checkpoint_dir/"history.json").write_text(json.dumps(history,indent=2)+"\n"); print(json.dumps(row))
+        if stale>=a.patience: break
+    print(json.dumps({"status":"completed","epochs_completed":len(history),"best_validation_loss":best},indent=2))
+if __name__=="__main__": main()

vendor/atmorep-official/.github/ISSUE_TEMPLATE/bug_report.md

@@ -0,0 +1,33 @@
+---
+name: Bug report
+about: Create a report to help us improve
+title: "[BUG]"
+labels: ''
+assignees: ''
+
+---
+
+**Describe the bug**
+A clear and concise description of what the bug is.
+
+**To Reproduce**
+Steps to reproduce the behavior:
+1. Activate '...'
+2. Change '....'
+3. Run '....'
+4. See error
+
+**Expected behavior**
+A clear and concise description of what you expected to happen.
+
+**Screenshots**
+If applicable, add screenshots to help explain your problem.
+
+**Hardware and environment:**
+ - System: [e.g. ATOS, LUMI, JSC]
+ - Runtime: [e.g. interactive, sbatch]
+ - Special branch: [e.g. develop, main]
+ - Python Version: [e.g. 3.11]
+
+**Additional context**
+Add any other context about the problem here.

vendor/atmorep-official/.github/ISSUE_TEMPLATE/feature_request.md

@@ -0,0 +1,20 @@
+---
+name: Feature request
+about: Suggest an idea for this project
+title: "[FEATURE]"
+labels: ''
+assignees: ''
+
+---
+
+**Is your feature request related to a problem? Please describe.**
+A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
+
+**Describe the solution you'd like**
+A clear and concise description of what you want to happen. Please use bullet points and action items when possible 
+
+**Describe alternatives you've considered**
+A clear and concise description of any alternative solutions or features you've considered.
+
+**Additional context**
+Add any other context or screenshots about the feature request here.

vendor/atmorep-official/.github/ISSUE_TEMPLATE/scientific-discussion.md

@@ -0,0 +1,20 @@
+---
+name: Scientific discussion
+about: discuss scientific developments
+title: "[SCIENTIFIC]"
+labels: ''
+assignees: ''
+
+---
+
+**Describe the topic**
+A clear and concise description of the scientific topic you want to discuss.
+
+**Describe the technical developments**
+A clear and concise description of the technical advancements needed to make this happen. Ideally as bullet points
+
+**Additional context**
+Add any other context or screenshots about the discussed topic here. 
+
+** Related issues**
+link here the issues related to this scientific discussion.

vendor/atmorep-official/.gitignore

@@ -0,0 +1,167 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# data files
+*.zarr
+
+# image files
+*.png
+

vendor/atmorep-official/CITATION.cff

@@ -0,0 +1,56 @@
+# This CITATION.cff file was generated with cffinit.
+# Visit https://bit.ly/cffinit to generate yours today!
+
+cff-version: 1.2.0
+title: Atmorep
+message: >-
+  If you use this software, please cite it using the
+  metadata from this file.
+type: software
+authors:
+  - given-names: Christian
+    family-names: Lessig
+    email: christian.lessig@ecmwf.int
+    affiliation: European Centre for Medium-Range Weather Forecasts (ECMWF)
+  - given-names: Ilaria
+    family-names: Luise
+    email: ilaria.luise@cern.ch
+    affiliation: European Organization for Nuclear Research (CERN)
+   - given-names: Martin
+    family-names: Schultz
+    email: m.schultz@fz-juelich.de
+    orcid: 'https://orcid.org/0000-0003-3455-774X'
+    affiliation: Forschungszentrum Jülich (FZJ)
+  - given-names: Michael
+    family-names: Langguth
+    email: m.langguth@fz-juelich.de
+    orcid: 'https://orcid.org/0000-0003-3354-5333'
+    affiliation: Forschungszentrum Jülich (FZJ)
+identifiers:
+  - type: url
+    value: 'https://arxiv.org/abs/2308.13280'
+    description: corresponding Preprint
+repository-code: 'https://isggit.cs.uni-magdeburg.de/atmorep/atmorep'
+url: 'https://www.atmorep.org'
+abstract: >-
+  AtmoRep is a novel, task-independent stochastic computer 
+  model of atmospheric dynamics that can provide skillful 
+  results for a wide range of applications. AtmoRep uses 
+  large-scale representation learning from artificial 
+  intelligence to determine a general description of the 
+  highly complex, stochastic dynamics of the atmosphere 
+  from the best available estimate of the system's historical 
+  trajectory as constrained by observations. This is enabled 
+  by a novel self-supervised learning objective and a unique 
+  ensemble that samples from the stochastic model with a 
+  variability informed by the one in the historical record. 
+  Our work establishes that large-scale neural networks can
+  provide skillful, task-independent models of atmospheric
+  dynamics. With this, they provide a novel means to make
+  the large record of atmospheric observations accessible
+  for applications and for scientific inquiry, complementing
+  existing simulations based on first principles.
+license: MIT
+commit: b0da5b32ec70295914bbb486dbcb77885671dc45
+version: 2.0 (preprint)
+date-released: '2023-11-28'

vendor/atmorep-official/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Otto-von-Guericke-Universitaet Magdeburg, Forschungszentrum Juelich GmbH, European Organization for Nuclear Research
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

vendor/atmorep-official/README.md

@@ -0,0 +1,200 @@
+
+# IMPORTANT NOTE: 
+<p align="center">
+<img src="https://github.com/user-attachments/assets/50518092-5386-4ef5-98f3-e2abf147556f" />
+</p>
+
+Please note that that the folder is **not maintained anymore since March 1st**. \
+Please use the **WeatherGenerator** code instead: https://github.com/ecmwf/WeatherGenerator 
+
+
+
+# AtmoRep
+
+This repository contains the source code for the [AtmoRep](https://www.atmorep.org) models for large scale representation learning of atmospheric dynamics as well as links to the pre-trained models and the required model input data.
+
+The pre-print for the work is available on ArXiv: https://arxiv.org/abs/2308.13280.
+
+```
+@misc{Lessig2023atmorep,
+  title = {AtmoRep: A stochastic model of atmosphere dynamics using large scale representation learning},
+	author = {Christian Lessig and Ilaria Luise and Bing Gong and Michael Langguth and Scarlet Stadler and Martin Schultz},
+  eprint = {2308.13280},
+  primaryclass = {physics.ao-ph},
+  url = {https://arxiv.org/abs/2308.13280},
+  year = {2023},
+```
+
+# Starter README
+
+## 1. Pull code
+
+`````
+%> wget git@github.com:clessig/atmorep.git
+`````
+This creates a directory ``atmorep`` with the code that contains the source code including the python scripts for model training and evaluation.
+
+After following the steps described below, the final directory structure will look as follows:
+````
+└── atmorep/
+    ├── atmorep/
+    │     └── ... 
+    ├── data/                         <- top level data directory
+    │    ├── normalisation/           <- directory for data normalisations 
+    │    ├── vorticity/
+    │    │       ├── ml105/           <- model levels with monthly GRIB files
+    │    │       │     ├── era5_vorticity_y2021_m03_ml137.grib   <- grib data file
+    │    │       │     ├── ...
+    │    │       ├── ml114/
+    │    │       ├── ml123/
+    │    │       ├── ml137/
+    │    │       ├── ml96/
+    .    .       .
+    │    ├── temperature/
+    .    .
+    ├── models
+    │    ├── id4nvwbetz   <- Directory containing model weights and config
+    │    │       ├──  model_id4nvwbetz.json     
+    │    │       └──  AtmoRep_id4nvwbetz.mod
+    │    ├── id<model_id>
+    .    .
+    └── results
+         ├── id4nvwbetz
+         ...
+````
+The directories ``data``, ``models``, and ``results`` need to be created if they do not exist. All directories might be large and should thus be on a directory with sufficient storage space; in this case they can be soft-linked to the default ones above or they can be set in ``atmorep/config/config``. 
+
+## 2. Download the data 
+
+### 2.1 Download pre-trained models
+
+Models can be downloaded from: https://datapub.fz-juelich.de/atmorep/trained-models.html
+
+An example for downloading the pre-trained models is given here, in this case for the vorticity model.
+
+`````
+% atmorep/> mkdir models
+% atmorep/> cd models
+% atmorep/data/> wget https://datapub.fz-juelich.de/atmorep/models/model_id4nvwbetz.tar.gz
+% atmorep/data/> tar xvzf model_id4nvwbetz.tar.gz
+% atmorep/data/> ls id4nvwbetz
+AtmoRep_id4nvwbetz.mod  model_id4nvwbetz.json
+`````
+
+
+### 2.2 Download model input data (ERA5)
+
+The input data in the required structure can be downloaded from the [Jülich datapub](https://datapub.fz-juelich.de/atmorep/era5-data.html) server. Direct link to WebDAV [https://datapub.fz-juelich.de/atmorep/data/](https://datapub.fz-juelich.de/atmorep/data/). Alternatively, it can be directly downloaded from MARS using the following [script](https://www.atmorep.org/code/mars_era5_download.py).
+
+#### Download a subset of files
+
+All data files (fields and normalizations) should be downloaded into the ``data`` directory. Un-taring the files will generate the correct folder structure. For example (we will use the vorticity example also below to run the first model so it is recommended to download it as a first step):
+`````
+% atmorep/> mkdir data
+% atmorep/> cd data
+% atmorep/data/> wget https://datapub.fz-juelich.de/atmorep/data/vorticity/ml137/era5_vorticity_y2021_ml137.tar
+% atmorep/data/> tar xvf era5_vorticity_y2021_ml137.tar
+% atmorep/data/> ls -lah vorticity/ml137/
+total 18G
+era5_vorticity_y2021_m01_ml137.grib
+era5_vorticity_y2021_m02_ml137.grib
+...
+era5_vorticity_y2021_m12_ml137.grib
+`````
+For efficiency reasons, AtmoRep takes monthly ERA5 data as input. Therefore, each tar file contains 12 GRIB files of about 1.5 GBytes each.
+
+Coefficients for data normalization per field and level can be downloaded here: https://datapub.fz-juelich.de/atmorep/data/normalization/. They should also be located in the ```data``` directory:
+`````
+% atmorep/data/> wget https://datapub.fz-juelich.de/atmorep/data/normalization/normalization_vorticity_ml137.tar.gz
+% atmorep/data/> tar xvzf normalization_vorticity_ml137.tar.gz
+`````
+
+## 3. Install python packages
+
+Create a python environment, e.g.
+
+`````
+% atmorep/> python3 -m venv pyenv
+`````
+
+and activate the environment:
+
+`````
+% atmorep/> source pyenv/bin/activate
+`````
+conda is also possible, no environment is strictly required although we would recommend it. Please make sure to use a recent python version (we tested with python3.10).
+Then install the AtmoRep package:
+`````
+% atmorep/>
+% atmorep/> pip install -e .
+`````
+
+torch is currently not included (since it is often available or has particular dependencies, e.g. a specific Cuda version). In the simplest case, it can just be installed by:
+
+`````
+% atmorep/> pip install torch
+`````
+We require torch 2.x. (A container solution allows to run even on systems where torch 2.x is not available.)
+
+## 4. Run model:
+Pre-trained models can normally be run by:
+`````
+% atmorep/> python atmorep/core/evaluate.py
+`````
+You can easily adapt the configuration by selecting the corresponding _model_id_ in ``evaluate.py`` (see below). It defaults to the single-field configuration of vorticity, of which we have downloaded the data above.
+
+Depending on your compute hardware, you might also have to run the computations by submitting the job using a batch system or allocate a compute node in interactive mode (if an interactive seesion is possible, then this is recommended). If you run an interactive session you will likely need to use the following:
+`````
+%  atmorep/> export CUDA_VISIBLE_DEVICES=0,1,2,3
+%  atmorep/> MASTER_ADDR="$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)"
+`````
+
+The default evaluation mode is currently global forecast. The output will be (similar to) this:
+````
+devices : ['cuda:0', 'cuda:1', 'cuda:2', 'cuda:3']
+Wandb run: atmorep-ztvyw7k6-8932958
+Running Evaluate.evaluate with mode = global_forecast
+Loaded AtmoRep id=4nvwbetz, ignoring/missing 2 elements.
+Loaded model id = 4nvwbetz at epoch = -2.
+Number of batches per global forecast: 14
+INFO:: data stats vorticity : 5.374998363549821e-05 / 0.9978392720222473
+num_accs_per_task : 1
+with_hvd : True
+hvd_rank : 0
+
+...
+
+wandb_id : ztvyw7k6
+dates : [[2021, 2, 10, 12]]
+token_overlap : [0, 0]
+forecast_num_tokens : 1
+validation loss for strategy=forecast at epoch 0 : 0.12402566522359848
+validation loss for vorticity : 0.12402566522359848
+wandb: Waiting for W&B process to finish... (success).
+wandb: 
+wandb: Run history:
+wandb:        val. loss forecast ▁
+wandb: val., forecast, vorticity ▁
+wandb: 
+wandb: Run summary:
+wandb:        val. loss forecast 0.12403
+wandb: val., forecast, vorticity 0.12403
+wandb: 
+wandb: You can sync this run to the cloud by running:
+wandb: wandb sync /p/project/atmo-rep/lessig/atmorep/atmorep/lessig-cleanup/atmorep/wandb/offline-run-20231124_095428-ztvyw7k6
+````
+For the vorticity example above, we evaluate with ``global_forecast`` for a specific date and using only a single model level:
+````
+mode, options = 'global_forecast', { 'fields[0][2]' : [137],
+                                     'dates' : [ [2021, 2, 10, 12] ],
+                                     'token_overlap' : [0, 0],
+                                     'forecast_num_tokens' : 1, 
+                                     'attention' : False}
+````
+We perform a 3 hour forecast, since 1 token is 3 hours wide. Another mode is the BERT masked token model mode used for pre-training:
+`````
+mode, options = 'BERT', {'years_test' : [2021], 'fields[0][2]' : [123, 137]}
+`````
+Again, we chose some custom options by using two levels instead of the five ones that are default and were used during pre-training and by using 2021 as the test year (since we downloaded the data).
+
+The generated model output (stored in ``./results/id{wandbid}``) for the ```global_forecast``` example can be post-processed into a spatial map with the [following code](https://www.atmorep.org/code/plot_forecast.py). The run_id at the top needs to be replaced by the wandb_id of your run, it can be read off from the console output. Results will be stored as ``example_0000{0,1,2}.png``. The code is also an as-simple-as-possible example with many parameters hard-coded, see our analysis code for a proper handling. 

vendor/atmorep-official/atmorep/config/config.py

@@ -0,0 +1,15 @@
+import os 
+from pathlib import Path
+
+fpath = os.path.dirname(os.path.realpath(__file__))
+
+path_models = Path( fpath, '../../models/')
+path_results = Path( fpath, '../../results')
+path_plots = Path( fpath, '../results/plots/')
+
+grib_index = { 'vorticity' : 'vo', 'divergence' : 'd', 'geopotential' : 'z',
+                'orography' : 'z', 'temperature': 't', 'specific_humidity' : 'q',
+                'mean_top_net_long_wave_radiation_flux' : 'mtnlwrf',
+                'velocity_u' : 'u', 'velocity_v': 'v', 'velocity_z' : 'w',
+                'total_precip' : 'tp', 'radar_precip' : 'yw_hourly',
+                't2m' : 't_2m', 'u_10m' : 'u_10m', 'v_10m' : 'v_10m',  }

vendor/atmorep-official/atmorep/core/atmorep_model.py

@@ -0,0 +1,590 @@
+####################################################################################################
+#
+#  Copyright (C) 2022
+#
+####################################################################################################
+#
+#  project     : atmorep
+#
+#  author      : atmorep collaboration
+# 
+#  description :
+#
+#  license     :
+#
+####################################################################################################
+
+import torch
+import numpy as np
+import code
+import os
+# code.interact(local=locals())
+
+# import horovod.torch as hvd
+
+import atmorep.utils.utils as utils
+from atmorep.utils.utils import identity
+from atmorep.utils.utils import NetMode
+from atmorep.utils.utils import get_model_filename
+
+from atmorep.transformer.transformer_base import prepare_token
+from atmorep.transformer.transformer_base import checkpoint_wrapper
+
+from atmorep.datasets.multifield_data_sampler import MultifieldDataSampler
+
+from atmorep.transformer.transformer_encoder import TransformerEncoder
+from atmorep.transformer.transformer_decoder import TransformerDecoder
+from atmorep.transformer.tail_ensemble import TailEnsemble
+
+
+####################################################################################################
+class AtmoRepData( torch.nn.Module) :
+
+  def __init__( self, net) :
+    '''Wrapper class for AtmoRep that handles data loading'''
+
+    super( AtmoRepData, self).__init__()
+    
+    self.data_loader_test = None
+    self.data_loader_train = None
+    self.data_loader_iter = None
+
+    self.net = net
+
+    # ensure that all data loaders have the same seed and hence load the same data
+    self.rng_seed = net.cf.rng_seed 
+    if not self.rng_seed :
+      self.rng_seed = int(torch.randint( 100000000, (1,))) 
+
+  ###################################################
+  def set_data( self, mode : NetMode, times_pos, batch_size = -1, num_loader_workers = -1) :
+
+    cf = self.net.cf
+    if batch_size < 0 :
+      batch_size = cf.batch_size_train if mode == NetMode.train else cf.batch_size_test
+    
+    dataset = self.dataset_train if mode == NetMode.train else self.dataset_test
+    dataset.set_data( times_pos, batch_size)
+    
+    self._set_data( dataset, mode, batch_size, num_loader_workers)
+
+  ###################################################
+  def set_global( self, mode : NetMode, times, batch_size = -1, num_loader_workers = -1) :
+
+    cf = self.net.cf
+    if batch_size < 0 :
+      batch_size = cf.batch_size_train if mode == NetMode.train else cf.batch_size_test
+    dataset = self.dataset_train if mode == NetMode.train else self.dataset_test
+    dataset.set_global( times, batch_size, cf.token_overlap)
+
+    self._set_data( dataset, mode, batch_size, num_loader_workers)
+
+  ###################################################
+  def set_location( self, mode : NetMode, pos, years, months, num_t_samples_per_month, 
+                          batch_size = -1, num_loader_workers = -1) :
+
+    cf = self.net.cf
+    if batch_size < 0 :
+      batch_size = cf.batch_size_train if mode == NetMode.train else cf.batch_size_test
+    
+    dataset = self.dataset_train if mode == NetMode.train else self.dataset_test
+    dataset.set_location( pos, years, months, num_t_samples_per_month, batch_size)
+
+    self._set_data( dataset, mode, batch_size, num_loader_workers)
+
+  ###################################################
+  def _set_data( self, dataset, mode : NetMode, batch_size = -1, loader_workers = -1) :
+    '''Private implementation for set_data, set_global'''
+
+    cf = self.net.cf
+    if loader_workers < 0 :
+      loader_workers = cf.num_loader_workers
+
+    loader_params = { 'batch_size': None, 'batch_sampler': None, 'shuffle': False, 
+                      'num_workers': loader_workers, 'pin_memory': True}
+    
+    if mode == NetMode.train :
+      self.data_loader_train = torch.utils.data.DataLoader( dataset, **loader_params, 
+                                                            sampler = None)
+    elif mode == NetMode.test :
+      self.data_loader_test = torch.utils.data.DataLoader( dataset, **loader_params, 
+                                                           sampler = None)
+    else :
+      assert False
+
+  ###################################################
+  def normalizer( self, field, vl_idx, lats_idx, lons_idx ) :
+
+    if isinstance( field, str) :
+      for fidx, field_info in enumerate(self.cf.fields) :
+        if field == field_info[0] :
+          break
+      assert fidx < len(self.cf.fields), 'invalid field'
+      normalizer = self.dataset_train.datasets[fidx].normalizer
+
+    elif isinstance( field, int) :
+      normalizer = self.dataset_train.normalizers[field][vl_idx]
+      if len(normalizer.shape) > 2:
+        normalizer = np.take( np.take( normalizer, lats_idx, -2), lons_idx, -1)
+    else :
+      assert False, 'invalid argument type (has to be index to cf.fields or field name)'
+    
+    year_base = self.dataset_train.year_base
+
+    return normalizer, year_base
+
+  ###################################################
+  def mode( self, mode : NetMode) :
+    
+    if mode == NetMode.train :
+      self.data_loader_iter = iter(self.data_loader_train)
+      self.net.train()
+    elif mode == NetMode.test :
+      self.data_loader_iter = iter(self.data_loader_test)
+      self.net.eval()
+    else :
+      assert False
+
+    self.cur_mode = mode
+
+  ###################################################
+  def len( self, mode : NetMode) :
+    if mode == NetMode.train :
+      return len(self.data_loader_train)
+    elif mode == NetMode.test :
+      return len(self.data_loader_test)
+    else :
+      assert False
+
+  ###################################################
+  def next( self) :
+    return next(self.data_loader_iter)
+
+  ###################################################
+  def forward( self, xin) :
+    pred = self.net.forward( xin)
+    return pred
+
+  ###################################################
+  def get_attention( self, xin) :
+    attn = self.net.get_attention( xin)
+    return attn
+
+  ###################################################
+  def create( self, pre_batch, devices, create_net = True, pre_batch_targets = None,
+              load_pretrained=True) :
+
+    if create_net :
+      self.net.create( devices, load_pretrained)
+
+    self.pre_batch = pre_batch
+    self.pre_batch_targets = pre_batch_targets
+
+    cf = self.net.cf
+    loader_params = { 'batch_size': None, 'batch_sampler': None, 'shuffle': False, 
+                      'num_workers': cf.num_loader_workers, 'pin_memory': True}
+
+    self.dataset_train = MultifieldDataSampler( cf.file_path, cf.fields, cf.years_train,
+                                                cf.batch_size,
+                                                pre_batch, cf.n_size, cf.num_samples_per_epoch,
+                                                with_shuffle = (cf.BERT_strategy != 'global_forecast'), 
+                                                with_source_idxs = True, 
+                                                compute_weights = (cf.losses.count('weighted_mse') > 0) )
+    self.data_loader_train = torch.utils.data.DataLoader( self.dataset_train, **loader_params,
+                                                          sampler = None)
+
+    self.dataset_test = MultifieldDataSampler( cf.file_path, cf.fields, cf.years_val,
+                                               cf.batch_size_validation,
+                                               pre_batch, cf.n_size, cf.num_samples_validate,
+                                               with_shuffle = (cf.BERT_strategy != 'global_forecast'),
+                                               with_source_idxs = True, 
+                                               compute_weights = (cf.losses.count('weighted_mse') > 0) )                                               
+    self.data_loader_test = torch.utils.data.DataLoader( self.dataset_test, **loader_params,
+                                                          sampler = None)
+
+    return self
+
+####################################################################################################
+class AtmoRep( torch.nn.Module) :
+
+  def __init__(self, cf) :
+    '''Constructor'''
+ 
+    super( AtmoRep, self).__init__()
+
+    self.cf = cf
+
+  ###################################################
+  def create( self, devices, load_pretrained=True) :
+    '''Create network'''
+
+    cf = self.cf
+    self.devices = devices
+    self.fields_coupling_idx = []
+
+    self.fields_index = {}
+    for ifield, field_info in enumerate(cf.fields) :
+      self.fields_index[ field_info[0] ] = ifield 
+    
+    # # embedding network for global/auxiliary token infos
+    # TODO: only for backward compatibility, remove
+    self.embed_token_info = torch.nn.Linear( cf.size_token_info, cf.size_token_info_net)
+    torch.nn.init.constant_( self.embed_token_info.weight, 0.0)
+
+    self.embeds_token_info = torch.nn.ModuleList()
+    for ifield, field_info in enumerate( cf.fields) :
+      
+      self.embeds_token_info.append( torch.nn.Linear( cf.size_token_info, cf.size_token_info_net))
+      
+      if len(field_info[1]) > 4 and load_pretrained :
+        # TODO: inconsistent with embeds_token_info -> version that can handle both
+        #       we could imply use the file name: embed_token_info vs embeds_token_info
+        name = 'AtmoRep' + '_embeds_token_info'
+        if not os.path.exists(get_model_filename( name, field_info[1][4][0], field_info[1][4][1])):
+          name = 'AtmoRep' + '_embed_token_info'
+
+        mloaded = torch.load( get_model_filename( name, field_info[1][4][0], field_info[1][4][1]))
+      
+        if "weight" not in mloaded.keys(): #TODO: get rid of this
+          mloaded["weight"] = mloaded["0.weight"]
+          mloaded["bias"] = mloaded["0.bias"]
+          del mloaded["0.weight"]
+          del mloaded["0.bias"]
+
+        self.embeds_token_info[-1].load_state_dict( mloaded)
+        print( 'Loaded embed_token_info from id = {}.'.format( field_info[1][4][0] ) )
+      else :
+        # initalization
+        torch.nn.init.constant_( self.embeds_token_info[-1].weight, 0.0)
+        self.embeds_token_info[-1].bias.data.fill_(0.0)
+
+    # embedding and encoder
+
+    self.embeds = torch.nn.ModuleList()
+    self.encoders = torch.nn.ModuleList()
+
+    for field_idx, field_info in enumerate(cf.fields) : 
+
+      # encoder
+      self.encoders.append( TransformerEncoder( cf, field_idx, True).create())
+      # load pre-trained model if specified
+      if len(field_info[1]) > 4 and load_pretrained :
+        self.load_block( field_info, 'encoder', self.encoders[-1])
+      self.embeds.append( self.encoders[-1].embed)
+     
+      # indices of coupled fields for efficient access in forward
+      self.fields_coupling_idx.append( [field_idx])
+      for field_coupled in field_info[1][2] : 
+        if 'axial' in cf.encoder_att_type :
+          self.fields_coupling_idx[field_idx].append( self.fields_index[field_coupled] )
+        else :
+          for _ in range(cf.coupling_num_heads_per_field) :
+            self.fields_coupling_idx[field_idx].append( self.fields_index[field_coupled] )
+
+    # decoder 
+
+    self.decoders = torch.nn.ModuleList()
+    self.field_pred_idxs = []
+    for field in cf.fields_prediction :
+
+      for ifield, field_info in enumerate(cf.fields) : 
+        if field_info[0] == field[0] :
+          self.field_pred_idxs.append( ifield)
+          break
+
+      self.decoders.append( TransformerDecoder( cf, field_info ) )
+      # load pre-trained model if specified
+      if len(field_info[1]) > 4 and load_pretrained :
+        self.load_block( field_info, 'decoder', self.decoders[-1])
+
+    # tail networks
+    
+    self.tails = torch.nn.ModuleList()
+    for ifield, field in enumerate(cf.fields_prediction) :
+
+      field_idx = self.field_pred_idxs[ifield]
+      field_info = cf.fields[field_idx]
+      self.tails.append( TailEnsemble( cf, field_info[1][1], np.prod(field_info[4]) ).create())
+      # load pre-trained model if specified
+      if len(field_info[1]) > 4 and load_pretrained:
+        self.load_block( field_info, 'tail', self.tails[-1])
+
+    # set devices
+
+    for field_idx, field_info in enumerate(cf.fields) :
+      # find determined device, use default if nothing specified
+      device = self.devices[0]
+      if len(field_info[1]) > 3 :
+        assert field_info[1][3] < 4, 'Only single node model parallelism supported'
+        print(devices, field_info[1][3])
+        assert field_info[1][3] < len(devices), 'Per field device id larger than max devices'
+        device = self.devices[ field_info[1][3] ]
+      # set device
+      self.embeds[field_idx].to(device)
+      self.encoders[field_idx].to(device)
+
+    for field_idx, field in enumerate(cf.fields_prediction) :
+      field_info = cf.fields[ self.field_pred_idxs[field_idx] ]
+      device = self.devices[0]
+      if len(field_info[1]) > 3 :
+        device = self.devices[ field_info[1][3] ]
+      self.decoders[field_idx].to(device)
+      self.tails[field_idx].to(device)
+
+    # embed_token_info on device[0] since it is shared by all fields, potentially sub-optimal
+    self.embed_token_info.to(devices[0])  # TODO: only for backward compatibility, remove
+    self.embeds_token_info.to(devices[0])
+
+    self.checkpoint = identity
+    if cf.grad_checkpointing :
+      self.checkpoint = checkpoint_wrapper
+
+    return self
+
+  ###################################################
+  def load_block( self, field_info, block_name, block ) :
+
+    # name = self.__class__.__name__ + '_' + block_name + '_' + field_info[0]
+    name = 'AtmoRep_' + block_name + '_' + field_info[0]
+
+    b_loaded = torch.load( get_model_filename(name, field_info[1][4][0], field_info[1][4][1]))
+
+    # in coupling mode, proj_out of attention heads needs separate treatment: only the pre-trained
+    # part can be loaded
+    keys_del = []
+    for name, param in block.named_parameters():
+      if 'proj_out' in name :
+        for k in b_loaded.keys() :
+          if name == k :
+            if param.shape[0] != param.shape[1] :   # non-square proj_out indicate deviation from pre-training
+              with torch.no_grad() :
+                # load pre-trained part
+                param[ : , : b_loaded[k].shape[1] ] = b_loaded[k]
+                # initalize remaining part to small random value
+                param[ : , b_loaded[k].shape[1] : ] = 0.01 * torch.rand( param.shape[0],
+                                                                         param.shape[1] - b_loaded[k].shape[1])
+                keys_del += [ k ]
+
+      #for backward compatibility. solved in new runs
+      if 'proj_heads_other' in name: 
+        for k in b_loaded.keys() :
+          if name == k :
+            if b_loaded[name].shape[0] == 0:
+              keys_del += [ name ]
+      
+    for k in keys_del :
+      del b_loaded[k]
+    
+    # use strict=False so that differing blocks, e.g. through coupling, are ignored
+    mkeys, _ = block.load_state_dict( b_loaded, False)
+
+    # missing keys = keys that are not pre-trained are initalized to small value
+    [mkeys.remove(k) for k in keys_del]  # remove proj_out keys so that they are not over-written
+    [utils.init_weights_uniform( block.state_dict()[k], 0.01) for k in mkeys]
+
+    print( 'Loaded {} for {} from id = {} (ignoring/missing {} elements).'.format( block_name,
+                                              field_info[0], field_info[1][4][0], len(mkeys) ) )
+
+  ###################################################
+  def translate_weights(self, mloaded, mkeys, ukeys):
+    '''
+    Function used for backward compatibility 
+    '''
+    cf = self.cf
+
+    #encoder:
+    for layer in range(cf.encoder_num_layers) :
+  
+      #shape([16, 3, 128, 2048])
+      mw  = torch.cat([mloaded[f'encoders.0.heads.{layer}.heads_self.{head}.proj_{k}.weight'] for head in range(cf.encoder_num_heads) for k in ["qs", "ks", "vs"]])
+      mloaded[f'encoders.0.heads.{layer}.proj_heads.weight'] = mw
+
+      for head in range(cf.encoder_num_heads):
+        del mloaded[f'encoders.0.heads.{layer}.heads_self.{head}.proj_qs.weight']
+        del mloaded[f'encoders.0.heads.{layer}.heads_self.{head}.proj_ks.weight']
+        del mloaded[f'encoders.0.heads.{layer}.heads_self.{head}.proj_vs.weight']
+
+      #cross attention
+      if f'encoders.0.heads.{layer}.heads_other.0.proj_qs.weight' in ukeys:
+        mw  = torch.cat([mloaded[f'encoders.0.heads.{layer}.heads_other.{head}.proj_{k}.weight'] for head in range(cf.encoder_num_heads) for k in ["qs", "ks", "vs"]])
+       
+        for i in range(cf.encoder_num_heads):
+          del mloaded[f'encoders.0.heads.{layer}.heads_other.{head}.proj_qs.weight']
+          del mloaded[f'encoders.0.heads.{layer}.heads_other.{head}.proj_ks.weight']
+          del mloaded[f'encoders.0.heads.{layer}.heads_other.{head}.proj_vs.weight']
+
+        mloaded[f'encoders.0.heads.{layer}.proj_heads_other.0.weight'] = mw  
+     
+    #decoder
+    for iblock in range(0, 19, 2) : 
+      mw  = torch.cat([mloaded[f'decoders.0.blocks.{iblock}.heads.{head}.proj_{k}.weight'] for head in range(8) for k in ["qs", "ks", "vs"]])
+      mloaded[f'decoders.0.blocks.{iblock}.proj_heads.weight'] = mw 
+
+      qs = [mloaded[f'decoders.0.blocks.{iblock}.heads_other.{head}.proj_qs.weight'] for head in range(8)]
+      mw  = torch.cat([mloaded[f'decoders.0.blocks.{iblock}.heads_other.{head}.proj_{k}.weight'] for head in range(8) for k in ["ks", "vs"]])
+      
+      mloaded[f'decoders.0.blocks.{iblock}.proj_heads_o_q.weight']  = torch.cat([*qs])
+      mloaded[f'decoders.0.blocks.{iblock}.proj_heads_o_kv.weight'] = mw
+
+      #self.num_samples_validate
+      decoder_dim = self.decoders[0].blocks[iblock].ln_q.weight.shape #128
+      mloaded[f'decoders.0.blocks.{iblock}.ln_q.weight'] = torch.tensor(np.ones(decoder_dim))
+      mloaded[f'decoders.0.blocks.{iblock}.ln_k.weight'] = torch.tensor(np.ones(decoder_dim))
+      mloaded[f'decoders.0.blocks.{iblock}.ln_q.bias'] = torch.tensor(np.ones(decoder_dim))
+      mloaded[f'decoders.0.blocks.{iblock}.ln_k.bias'] = torch.tensor(np.ones(decoder_dim))
+
+      for i in range(8):
+        del mloaded[f'decoders.0.blocks.{iblock}.heads.{i}.proj_qs.weight']
+        del mloaded[f'decoders.0.blocks.{iblock}.heads.{i}.proj_ks.weight']
+        del mloaded[f'decoders.0.blocks.{iblock}.heads.{i}.proj_vs.weight']
+        del mloaded[f'decoders.0.blocks.{iblock}.heads_other.{i}.proj_qs.weight']
+        del mloaded[f'decoders.0.blocks.{iblock}.heads_other.{i}.proj_ks.weight']
+        del mloaded[f'decoders.0.blocks.{iblock}.heads_other.{i}.proj_vs.weight']
+    
+    return mloaded
+
+  ###################################################
+  @staticmethod
+  def load( model_id, devices, cf = None, epoch = -2, load_pretrained=False) :
+    '''Load network from checkpoint'''
+
+    if not cf : 
+      cf = utils.Config()
+      cf.load_json( model_id)
+
+    model = AtmoRep( cf).create( devices, load_pretrained=False)
+    mloaded = torch.load( utils.get_model_filename( model, model_id, epoch) )
+    mkeys, ukeys = model.load_state_dict( mloaded, False )
+    if (f'encoders.0.heads.0.proj_heads.weight') in mkeys:
+      mloaded = model.translate_weights(mloaded, mkeys, ukeys)
+      mkeys, ukeys = model.load_state_dict( mloaded, False )
+
+    if len(mkeys) > 0 :
+      print( f'Loaded AtmoRep: ignoring {len(mkeys)} elements: {mkeys}')
+
+    # TODO: remove, only for backward 
+    if model.embeds_token_info[0].weight.abs().max() == 0. :
+      model.embeds_token_info = torch.nn.ModuleList()
+      
+    return model
+    
+  ###################################################
+  def save( self, epoch = -2) :
+    '''Save network '''
+
+    # save entire network
+    torch.save( self.state_dict(), utils.get_model_filename( self, self.cf.wandb_id, epoch) )
+
+    # save parts also separately
+
+    # name = self.__class__.__name__ + '_embed_token_info'
+    # torch.save( self.embed_token_info.state_dict(),
+    #             utils.get_model_filename( name, self.cf.wandb_id, epoch) )
+    name = self.__class__.__name__ + '_embeds_token_info'
+    torch.save( self.embeds_token_info.state_dict(),
+                utils.get_model_filename( name, self.cf.wandb_id, epoch) )
+
+    for ifield, enc in enumerate(self.encoders) :
+      name = self.__class__.__name__ + '_encoder_' + self.cf.fields[ifield][0]
+      torch.save( enc.state_dict(), utils.get_model_filename( name, self.cf.wandb_id, epoch) )
+
+    for ifield, dec in enumerate(self.decoders) :
+      name = self.__class__.__name__ + '_decoder_' + self.cf.fields_prediction[ifield][0]
+      torch.save( dec.state_dict(), utils.get_model_filename( name, self.cf.wandb_id, epoch) )
+
+    for ifield, tail in enumerate(self.tails) :
+      name = self.__class__.__name__ + '_tail_' + self.cf.fields_prediction[ifield][0]
+      torch.save( tail.state_dict(), utils.get_model_filename( name, self.cf.wandb_id, epoch) )
+
+  ###################################################
+  def forward( self, xin) :
+    '''Evaluate network'''
+
+    # embedding
+    cf = self.cf
+    
+    fields_embed = self.get_fields_embed(xin)
+    
+    # attention maps (if requested)
+    atts = [ [] for _ in cf.fields ]
+
+    # encoder
+    embeds_layers = [[] for i in self.field_pred_idxs]
+    for ib in range(self.cf.encoder_num_layers) :
+      fields_embed, att = self.forward_encoder_block( ib, fields_embed) 
+      [embeds_layers[idx].append( fields_embed[i]) for idx,i in enumerate(self.field_pred_idxs)]
+      [atts[i].append( att[i]) for i,_ in enumerate(cf.fields) ]
+        
+    # encoder-decoder coupling / token transformations
+    (decoders_in, embeds_layers) = self.encoder_to_decoder( embeds_layers)
+
+    preds = []
+    for idx,i in enumerate(self.field_pred_idxs) :
+    
+      # decoder
+      token_seq_embed, att = self.decoders[idx]( (decoders_in[idx], embeds_layers[idx]) )
+      
+      # tail net
+      tail_in = self.decoder_to_tail( idx, token_seq_embed)
+      pred = self.checkpoint( self.tails[idx], tail_in)
+      
+      preds.append( pred)
+      [atts[i].append( a) for a in att]
+
+    return preds, atts
+
+  ###################################################
+  def forward_encoder_block( self, iblock, fields_embed) :
+    ''' evaluate one block (attention and mlp) '''
+
+    # double buffer for commutation-invariant result (w.r.t evaluation order of transformers)
+    fields_embed_cur, atts = [], []
+
+    # attention heads
+    for ifield in range( len(fields_embed)) :
+      d = fields_embed[ifield].device
+      fields_in =[fields_embed[i].to(d,non_blocking=True) for i in self.fields_coupling_idx[ifield]]
+      # unpack list in argument for checkpointing
+      y, att = self.checkpoint( self.encoders[ifield].heads[iblock], *fields_in)
+      fields_embed_cur.append( y)
+      atts.append( att)
+      
+    # MLPs 
+    for ifield in range( len(fields_embed)) :
+      fields_embed_cur[ifield] = self.checkpoint( self.encoders[ifield].mlps[iblock], 
+                                                  fields_embed_cur[ifield] )
+  
+    return fields_embed_cur, atts
+
+  ###################################################
+  def get_fields_embed( self, xin ) :
+    if 0 == len(self.embeds_token_info) :   # TODO: only for backward compatibility, remove
+      emb_net_ti = self.embed_token_info
+      return [prepare_token( field_data, emb_net, emb_net_ti )
+                              for fidx,(field_data,emb_net) in enumerate(zip( xin, self.embeds))]
+    else :
+      embs_net_ti = self.embeds_token_info
+      return [prepare_token( field_data, emb_net, embs_net_ti[fidx] )
+                                for fidx,(field_data,emb_net) in enumerate(zip( xin, self.embeds))]
+    
+  ###################################################
+
+  def get_attention( self, xin) : 
+
+    cf = self.cf
+    attn = []
+    fields_embed = self.get_fields_embed(xin)
+    #either accumulated attention or last layer attention:
+    blocks = list(range(self.cf.encoder_num_layers)) if cf.attention_mode == 'accum' else [self.cf.encoder_num_layers-1]
+    for idx, ifield in enumerate(self.field_pred_idxs) :      
+      d = fields_embed[ifield].device
+      fields_in =[fields_embed[i].to(d,non_blocking=True) for i in self.fields_coupling_idx[ifield]]
+      attn_field = self.encoders[ifield].heads[blocks[0]].get_attention(fields_in)
+      if cf.attention_mode == 'accum':
+        for iblock in blocks[1:]:
+          attn_layer = self.encoders[ifield].heads[iblock].get_attention(fields_in)
+          attn_field = attn_field + attn_layer
+        attn_field = torch.sum(attn_field, dim = 0, keepdim=True)
+      attn.append(attn_field)
+#      print("att FINAL", ifield, len(attn), attn[0].shape)
+    return attn

vendor/atmorep-official/atmorep/core/evaluate.py

@@ -0,0 +1,74 @@
+ ####################################################################################################
+#
+#  Copyright (C) 2022
+#
+####################################################################################################
+#
+#  project     : atmorep
+#
+#  author      : atmorep collaboration
+# 
+#  description :
+#
+#  license     :
+#
+####################################################################################################
+
+from atmorep.core.evaluator import Evaluator
+import time
+
+if __name__ == '__main__':
+
+  # arXiv 2023: models for individual fields
+  #model_id = '4nvwbetz'     # vorticity
+  #model_id = 'oxpycr7w'     # divergence
+  #model_id = '1565pb1f'     # specific_humidity
+  #model_id = '3kdutwqb'     # total precip
+  #model_id = 'dys79lgw'     # velocity_u
+  #model_id = '22j6gysw'     # velocity_v
+  #model_id = '15oisw8d'     # velocity_z
+  #model_id = '3qou60es'     # temperature 
+  #model_id = '2147fkco'     # temperature (also 2147fkco)
+  
+  # new runs 2024
+  #model_id='j8dwr5qj' #velocity_u
+  #model_id='0tlnm5up' #velocity_v
+  #model_id='v63l01zu' #specific humidity 
+  #model_id='9l1errbo' #velocity_z
+  model_id='7ojls62c' #temperature 1024 
+  
+  # supported modes: test, forecast, fixed_location, temporal_interpolation, global_forecast,
+  #                  global_forecast_range
+  # options can be used to over-write parameters in config; some modes also have specific options, 
+  # e.g. global_forecast where a start date can be specified
+
+  #Add 'attention' : True to options to store the attention maps. NB. supported only for single field runs. 
+  
+  # BERT masked token model
+  #mode, options = 'BERT', {'years_val' : [2021], 'num_samples_validate' : 128, 'with_pytest' : True}
+
+  # BERT forecast mode
+  #mode, options = 'forecast', {'forecast_num_tokens' : 2, 'num_samples_validate' : 128, 'with_pytest' : True }
+
+  #temporal interpolation 
+  #idx_time_mask: list of relative time positions of the masked tokens within the cube wrt num_tokens[0]  
+  #mode, options = 'temporal_interpolation', {'idx_time_mask': [5,6,7], 'num_samples_validate' : 128, 'with_pytest' : True}
+
+  # BERT forecast with patching to obtain global forecast
+  mode, options = 'global_forecast', { 
+                                      #'dates' : [[2021, 2, 10, 12]]
+                                      'dates' : [
+                                         [2021, 1, 10, 12] , [2021, 1, 11, 0], [2021, 1, 11, 12], [2021, 1, 12, 0], #[2021, 1, 12, 12], [2021, 1, 13, 0], 
+                                         [2021, 4, 10, 12], [2021, 4, 11, 0], [2021, 4, 11, 12], [2021, 4, 12, 0], #[2021, 4, 12, 12], [2021, 4, 13, 0], 
+                                         [2021, 7, 10, 12], [2021, 7, 11, 0], [2021, 7, 11, 12], [2021, 7, 12, 0], #[2021, 7, 12, 12], [2021, 7, 13, 0], 
+                                         [2021, 10, 10, 12], [2021, 10, 11, 0], [2021, 10, 11, 12], #[2021, 10, 12, 0], [2021, 10, 12, 12], [2021, 10, 13, 0]
+                                       ], 
+                                      'token_overlap' : [0, 0],
+                                      'forecast_num_tokens' : 2,
+                                      'with_pytest' : True }
+
+  file_path = '/gpfs/scratch/ehpc03/era5_y1979_2021_res025_chunk8.zarr'
+  
+  now = time.time()
+  Evaluator.evaluate( mode, model_id, file_path, options)
+  print("time", time.time() - now)

vendor/atmorep-official/atmorep/core/evaluator.py

@@ -0,0 +1,240 @@
+####################################################################################################
+#
+#  Copyright (C) 2022
+#
+####################################################################################################
+#
+#  project     : atmorep
+#
+#  author      : atmorep collaboration
+# 
+#  description :
+#
+#  license     :
+#
+####################################################################################################
+
+import numpy as np
+import os
+import code
+import pytest
+import datetime
+
+import wandb
+
+from atmorep.core.trainer import Trainer_BERT
+from atmorep.utils.utils import Config
+from atmorep.utils.utils import setup_ddp
+from atmorep.utils.utils import setup_wandb
+from atmorep.utils.utils import init_torch
+from atmorep.utils.utils import NetMode
+import atmorep.utils.utils as utils
+
+import atmorep.config.config as config
+
+class Evaluator( Trainer_BERT) :
+
+  ##############################################
+  def __init__( self, cf, devices) :
+    Trainer_BERT.__init__( self, cf, devices)
+
+  ##############################################
+  def parse_args( cf, args) :
+
+    # set/over-write options as desired
+    for (key,val) in args.items() :
+      if '[' in key :  # handle lists, e.g. fields[0][2]
+        key_split = key.split( '[')
+        k, v = key_split[0], key_split[1:]
+        v = [int(a[0]) for a in v]
+        utils.list_replace_rec( getattr( cf, k), v, val)
+      else :
+        setattr( cf, key, val)
+
+  ##############################################
+  @staticmethod
+  def run( cf, model_id, model_epoch, devices) :
+    
+    cf.with_mixed_precision = True
+
+    # set/over-write options as desired
+    evaluator = Evaluator.load( cf, model_id, model_epoch, devices)
+
+    if 0 == cf.par_rank :
+      cf.print()
+      cf.write_json( wandb)
+    evaluator.validate( 0, cf.BERT_strategy)
+
+  ##############################################
+  @staticmethod
+  def evaluate( mode, model_id, file_path, args = {}, model_epoch=-2) :
+
+    devices = init_torch()
+    with_ddp = True 
+    par_rank, par_size = setup_ddp( with_ddp)
+    
+    cf = Config().load_json( model_id)
+  
+    cf.num_accs_per_task = len(devices)
+    cf.file_path = file_path
+    cf.with_wandb = True
+    cf.with_ddp = with_ddp
+    cf.par_rank = par_rank
+    cf.par_size = par_size
+    cf.losses = cf.losses
+    # overwrite old config
+    cf.attention = False
+    setup_wandb( cf.with_wandb, cf, par_rank, '', mode='offline')
+    if 0 == cf.par_rank :
+      print( 'Running Evaluate.evaluate with mode =', mode)
+
+    # if not hasattr( cf, 'num_loader_workers'):
+    cf.num_loader_workers = 12 #cf.loader_num_workers
+    cf.rng_seed = None 
+    
+    #backward compatibility
+    if not hasattr( cf, 'n_size'):
+      cf.n_size = [36, 0.25*9*6, 0.25*9*12]
+      #cf.n_size = [36, 0.25*27*2, 0.25*27*4] 
+    if not hasattr(cf, 'num_samples_per_epoch'):
+      cf.num_samples_per_epoch = 1024
+    if not hasattr(cf, 'with_mixed_precision'):
+      cf.with_mixed_precision = False
+    if not hasattr(cf, 'with_pytest'):
+      cf.with_pytest = False
+    if not hasattr(cf, 'batch_size'):
+      cf.batch_size = cf.batch_size_max
+    if not hasattr(cf, 'batch_size_validation'):
+      cf.batch_size_validation = cf.batch_size_max
+    if not hasattr(cf, 'years_val'):
+      cf.years_val = cf.years_test
+
+    func = getattr( Evaluator, mode)
+    func( cf, model_id, model_epoch, devices, args)
+    
+    if cf.with_pytest:
+      fields = [field[0] for field in cf.fields_prediction]
+      for field in fields:
+        pytest.main(["-x", "-s", "./atmorep/tests/validation_test.py", "--field", field, "--model_id", cf.wandb_id, "--strategy", cf.BERT_strategy])
+
+  ##############################################
+  @staticmethod
+  def BERT( cf, model_id, model_epoch, devices, args = {}) :
+
+    cf.lat_sampling_weighted = False
+    cf.BERT_strategy = 'BERT'
+    cf.log_test_num_ranks = 4
+    cf.num_samples_validate = 10 #28 #1472 
+    Evaluator.parse_args( cf, args)
+    utils.check_num_samples(cf.num_samples_validate, cf.batch_size)
+    Evaluator.run( cf, model_id, model_epoch, devices)
+
+  ##############################################
+  @staticmethod
+  def forecast( cf, model_id, model_epoch, devices, args = {}) :
+
+    cf.lat_sampling_weighted = False
+    cf.BERT_strategy = 'forecast'
+    cf.log_test_num_ranks = 4
+    cf.forecast_num_tokens = 1  # will be overwritten when user specified
+    cf.num_samples_validate = 128 #128 
+    Evaluator.parse_args( cf, args)
+    utils.check_num_samples(cf.num_samples_validate, cf.batch_size)
+    Evaluator.run( cf, model_id, model_epoch, devices)
+  
+  ##############################################
+  @staticmethod
+  def global_forecast( cf, model_id, model_epoch, devices, args = {}) :
+    
+    cf.BERT_strategy = 'global_forecast'
+    cf.batch_size_test = 24
+    cf.num_loader_workers = 12 #1
+    cf.log_test_num_ranks = 1
+
+    #TODO: temporary solution. Add support for batch_size > 1 
+    cf.batch_size_validation = 1 #64
+    cf.batch_size = 1
+    
+    if not hasattr(cf, 'num_samples_validate'):
+      cf.num_samples_validate = 196 
+    #if not hasattr(cf,'with_mixed_precision'):
+    cf.with_mixed_precision = True
+
+    Evaluator.parse_args( cf, args)
+
+    dates = args['dates']
+    evaluator = Evaluator.load( cf, model_id, model_epoch, devices)
+    evaluator.model.set_global( NetMode.test, np.array( dates))
+    if 0 == cf.par_rank :
+      cf.print()
+      cf.write_json( wandb)
+    evaluator.validate( 0, cf.BERT_strategy)
+
+  ##############################################
+  @staticmethod
+  def global_forecast_range( cf, model_id, model_epoch, devices, args = {}) :
+
+    cf.forecast_num_tokens = 2
+    cf.BERT_strategy = 'global_forecast'
+    cf.token_overlap = [0, 0]
+
+    cf.batch_size_test = 24
+    cf.num_loader_workers = 1
+    cf.log_test_num_ranks = 1
+    cf.batch_size_start = 14
+    if not hasattr(cf, 'num_samples_validate'):
+      cf.num_samples_validate = 196 
+
+    Evaluator.parse_args( cf, args)
+
+    if 0 == cf.par_rank :
+      cf.print()
+      cf.write_json( wandb)
+
+    # generate temporal sequence
+    dates = [ ]
+    num_steps = 31*2 
+    cur_date = [2018, 1, 1, 0+6] #6h models
+    for _ in range(num_steps) :
+      tdate = datetime.datetime( cur_date[0], cur_date[1], cur_date[2], cur_date[3])
+      tdate += datetime.timedelta( hours = 12 )
+      cur_date = [tdate.year, tdate.month, tdate.day, tdate.hour]
+      dates += [cur_date]
+
+    evaluator = Evaluator.load( cf, model_id, model_epoch, devices)
+    evaluator.model.set_global( NetMode.test, np.array( dates))
+    evaluator.evaluate( 0, cf.BERT_strategy)
+
+  ##############################################
+  @staticmethod
+  def temporal_interpolation( cf, model_id, model_epoch, devices, args = {}) :
+
+    # set/over-write options
+    cf.BERT_strategy = 'temporal_interpolation'
+    cf.log_test_num_ranks = 4
+    cf.num_samples_validate = 128
+    Evaluator.parse_args( cf, args)
+    utils.check_num_samples(cf.num_samples_validate, cf.batch_size)
+    Evaluator.run( cf, model_id, model_epoch, devices)
+
+  ##############################################
+  @staticmethod
+  def fixed_location( cf, model_id, model_epoch, devices, args = {}) :
+
+    # set/over-write options
+    cf.BERT_strategy = 'BERT'
+    cf.num_files_test = 2
+    cf.num_patches_per_t_test = 2
+    cf.log_test_num_ranks = 4
+
+    pos = [ 33.55 , 18.25 ]
+    years = [2018]
+    months = list(range(1,12+1))
+    num_t_samples_per_month = 2
+
+    evaluator = Evaluator.load( cf, model_id, model_epoch, devices)
+    evaluator.model.set_location( NetMode.test, pos, years, months, num_t_samples_per_month)
+    if 0 == cf.par_rank :
+      cf.print()
+      cf.write_json( wandb)
+    evaluator.evaluate( 0)

vendor/atmorep-official/atmorep/core/train.py

@@ -0,0 +1,251 @@
+####################################################################################################
+#
+#  Copyright (C) 2022
+#
+####################################################################################################
+#
+#  project     : atmorep
+#
+#  author      : atmorep collaboration
+# 
+#  description :
+#
+#  license     :
+#
+####################################################################################################
+
+import torch
+import os
+import sys
+import traceback
+import pdb
+import wandb
+
+from atmorep.core.trainer import Trainer_BERT
+from atmorep.utils.utils import Config
+from atmorep.utils.utils import setup_ddp
+from atmorep.utils.utils import setup_wandb
+from atmorep.utils.utils import init_torch
+
+
+####################################################################################################
+def train_continue( wandb_id, epoch, Trainer, epoch_continue = -1) :
+
+  devices = init_torch()
+  with_ddp = True
+  par_rank, par_size = setup_ddp( with_ddp)
+
+  cf = Config().load_json( wandb_id)
+
+  cf.num_accs_per_task = len(devices)   # number of GPUs / accelerators per task
+  cf.with_ddp = with_ddp
+  cf.par_rank = par_rank
+  cf.par_size = par_size
+  cf.optimizer_zero = False
+  cf.attention = False
+  # name has changed but ensure backward compatibility
+  if hasattr( cf, 'loader_num_workers') :
+    cf.num_loader_workers = cf.loader_num_workers
+  if not hasattr( cf, 'n_size'):
+    cf.n_size = [36, 0.25*9*6, 0.25*9*12] 
+  if not hasattr(cf, 'num_samples_per_epoch'):
+    cf.num_samples_per_epoch = 1024
+  if not hasattr(cf, 'num_samples_validate'):
+    cf.num_samples_validate = 128
+  if not hasattr(cf, 'with_mixed_precision'):
+    cf.with_mixed_precision = True
+  if not hasattr(cf, 'years_val'):
+    cf.years_val = cf.years_test
+ 
+  # any parameter in cf can be overwritten when training is continued, e.g. we can increase the 
+  # masking rate 
+  # cf.fields = [ [ 'specific_humidity', [ 1, 2048, [ ], 0 ], 
+  #                               [ 96, 105, 114, 123, 137 ], 
+  #                               [12, 6, 12], [3, 9, 9], [0.5, 0.9, 0.1, 0.05] ] ]
+  
+  setup_wandb( cf.with_wandb, cf, par_rank, project_name='train', mode='offline')  
+  # resuming a run requires online mode, which is not available everywhere
+  #setup_wandb( cf.with_wandb, cf, par_rank, wandb_id = wandb_id)  
+  
+  if cf.with_wandb and 0 == cf.par_rank :
+    cf.write_json( wandb)
+    cf.print()
+
+  if -1 == epoch_continue :
+    epoch_continue = epoch
+
+  # run
+  trainer = Trainer.load( cf, wandb_id, epoch, devices)
+  print( 'Loaded run \'{}\' at epoch {}.'.format( wandb_id, epoch))
+  trainer.run( epoch_continue)
+
+####################################################################################################
+def train() :
+
+  devices = init_torch()
+  with_ddp = True
+  par_rank, par_size = setup_ddp( with_ddp)
+
+  # torch.cuda.set_sync_debug_mode(1)
+  torch.backends.cuda.matmul.allow_tf32 = True
+
+  cf = Config()
+  # parallelization
+  cf.with_ddp = with_ddp
+  cf.num_accs_per_task = len(devices)   # number of GPUs / accelerators per task
+  cf.par_rank = par_rank
+  cf.par_size = par_size
+  
+  # format: list of fields where for each field the list is 
+  # [ name , 
+  #   [ dynamic or static field { 1, 0 }, embedding dimension, , device id ],
+  #   [ vertical levels ],
+  #   [ num_tokens],
+  #   [ token size],
+  #   [ total masking rate, rate masking, rate noising, rate for multi-res distortion]
+  # ]
+
+  # cf.fields = [ [ 'temperature', [ 1, 1024, [ ], 0 ], 
+  #                              [ 96, 105, 114, 123, 137 ], 
+  #                              [12, 2, 4], [3, 27, 27], [0.5, 0.9, 0.2, 0.05], 'local' ] ]
+  # cf.fields_prediction = [ [cf.fields[0][0], 1.] ]
+ 
+  cf.fields = [ [ 'velocity_u', [ 1, 1024, [ ], 0 ], 
+                                [ 96, 105, 114, 123, 137 ], 
+                                 [12, 3, 6], [3, 18, 18], [0.5, 0.9, 0.2, 0.05] ] ]
+
+  cf.fields_prediction = [ [cf.fields[0][0], 1.] ]
+
+  
+  # cf.fields = [ [ 'velocity_v', [ 1, 1024, [ ], 0 ], 
+  #                               [ 96, 105, 114, 123, 137 ], 
+  #                               [12, 3, 6], [3, 18, 18], [0.25, 0.9, 0.1, 0.05] ] ]
+
+  # cf.fields = [ [ 'velocity_z', [ 1, 1024, [ ], 0 ], 
+  #                               [ 96, 105, 114, 123, 137 ], 
+  #                               [12, 3, 6], [3, 18, 18], [0.25, 0.9, 0.1, 0.05] ] ]
+
+  # cf.fields = [ [ 'specific_humidity', [ 1, 1024, [ ], 0 ], 
+  #                               [ 96, 105, 114, 123, 137 ], 
+  #                               [12, 3, 6], [3, 18, 18], [0.25, 0.9, 0.1, 0.05] ] ]
+  
+  cf.fields_targets = []
+
+  cf.years_train = list( range( 1979, 2021))
+  cf.years_val = [2021]  #[2018] 
+  cf.month = None
+  cf.geo_range_sampling = [[ -90., 90.], [ 0., 360.]]
+  cf.time_sampling = 1   # sampling rate for time steps
+  # random seeds
+  cf.torch_seed = torch.initial_seed()
+  # training params
+  cf.batch_size_validation = 1 #64
+  cf.batch_size = 96
+  cf.num_epochs = 400 #128
+  cf.num_samples_per_epoch = 4096*12
+  cf.num_samples_validate = 128*12
+  cf.num_loader_workers = 8
+  
+  # additional infos
+  cf.size_token_info = 8
+  cf.size_token_info_net = 16
+  cf.grad_checkpointing = True
+  cf.with_cls = False
+  # network config
+  cf.with_mixed_precision = True
+  cf.with_layernorm = True
+  cf.coupling_num_heads_per_field = 1
+  cf.dropout_rate = 0.05
+  cf.with_qk_lnorm = False
+  # encoder
+  cf.encoder_num_layers = 6
+  cf.encoder_num_heads = 16
+  cf.encoder_num_mlp_layers = 2
+  cf.encoder_att_type = 'dense'
+  # decoder
+  cf.decoder_num_layers = 6
+  cf.decoder_num_heads = 16
+  cf.decoder_num_mlp_layers = 2
+  cf.decoder_self_att = False
+  cf.decoder_cross_att_ratio = 0.5
+  cf.decoder_cross_att_rate = 1.0
+  cf.decoder_att_type = 'dense'
+  # tail net
+  cf.net_tail_num_nets = 16
+  cf.net_tail_num_layers = 0
+  # loss
+  cf.losses = ['mse_ensemble', 'stats'] # mse, mse_ensemble, stats, crps, weighted_mse
+  # training
+  cf.optimizer_zero = False
+  cf.lr_start = 5. * 10e-7
+  cf.lr_max = 0.00005*3
+  cf.lr_min = 0.00004 #0.00002
+  cf.weight_decay = 0.05 #0.1
+  cf.lr_decay_rate = 1.025
+  cf.lr_start_epochs = 3
+  cf.model_log_frequency = 256 #save checkpoint every X batches
+  # BERT
+  # strategies: 'BERT', 'forecast', 'temporal_interpolation'
+  cf.BERT_strategy = 'BERT'
+  cf.forecast_num_tokens = 2      # only needed / used for BERT_strategy 'forecast
+  cf.BERT_fields_synced = False   # apply synchronized / identical masking to all fields 
+                                  # (fields need to have same BERT params for this to have effect)
+  cf.BERT_mr_max = 2              # maximum reduction rate for resolution
+  
+  # debug / output
+  cf.log_test_num_ranks = 0
+  cf.save_grads = False
+  cf.profile = False
+  cf.test_initial = False
+  cf.attention = False
+
+  cf.rng_seed = None 
+
+  # usually use %>wandb offline to switch to disable syncing with server
+  cf.with_wandb = True
+  setup_wandb( cf.with_wandb, cf, par_rank, 'train', mode='offline')  
+
+  # cf.file_path = '/p/scratch/atmo-rep/data/era5_1deg/months/era5_y2021_res100_chunk32.zarr'
+  # cf.file_path = '/ec/res4/scratch/nacl/atmorep/era5_y2021_res100_chunk32.zarr'
+  # # # in steps x lat_degrees x lon_degrees
+  # cf.n_size = [36, 1*9*6, 1.*9*12]
+
+  # # # # # cf.file_path = '/p/scratch/atmo-rep/data/era5_1deg/months/era5_y2021_res025_chunk16.zarr'
+  # # # # cf.file_path = '/p/scratch/atmo-rep/data/era5_1deg/months/era5_y2021_res025_chunk32.zarr'
+  # cf.file_path = '/ec/res4/scratch/nacl/atmorep/era5_y2021_res025_chunk32.zarr'
+  # # # 
+  # # # cf.file_path = '/p/scratch/atmo-rep/data/era5_1deg/months/era5_y2021_res025_chunk8.zarr'
+  # # cf.file_path = '/ec/res4/scratch/nacl/atmorep/era5_y2021_res025_chunk8_lat180_lon180.zarr'
+  # # # cf.file_path = '/ec/res4/scratch/nacl/atmorep/era5_y2021_res025_chunk16.zarr'
+  cf.file_path = '/gpfs/scratch/ehpc03/era5_y1979_2021_res025_chunk8.zarr/'
+  # # # in steps x lat_degrees x lon_degrees
+  cf.n_size = [36, 0.25*9*6, 0.25*9*12]
+
+  # cf.file_path = '/ec/res4/scratch/nacl/atmorep/era5_y2021_res100_chunk16.zarr'
+  #cf.file_path = '/p/scratch/atmo-rep/data/era5_1deg/months/era5_y2021_res100_chunk16.zarr'
+  #cf.n_size = [36, 1*9*6, 1.*9*12]
+
+  if cf.with_wandb and 0 == cf.par_rank :
+    cf.write_json( wandb)
+    cf.print()
+
+  trainer = Trainer_BERT( cf, devices).create()
+  trainer.run()
+
+####################################################################################################
+if __name__ == '__main__':
+  
+  try :
+
+    train()
+ 
+    #  wandb_id, epoch, epoch_continue = 'gxfywjzl', 127, 127
+    #  Trainer = Trainer_BERT
+    #  train_continue( wandb_id, epoch, Trainer, epoch_continue)
+
+  except :
+    
+    extype, value, tb = sys.exc_info()
+    traceback.print_exc()
+    pdb.post_mortem(tb)
+

vendor/atmorep-official/atmorep/core/train_multi.py

@@ -0,0 +1,277 @@
+####################################################################################################
+#
+#  Copyright (C) 2022
+#
+####################################################################################################
+#
+#  project     : atmorep
+#
+#  author      : atmorep collaboration
+# 
+#  description :
+#
+#  license     :
+#
+####################################################################################################
+
+import torch
+import os
+import sys
+import traceback
+import pdb
+import wandb
+
+import atmorep.config.config as config
+from atmorep.core.trainer import Trainer_BERT
+from atmorep.utils.utils import Config
+from atmorep.utils.utils import setup_ddp
+from atmorep.utils.utils import setup_wandb
+from atmorep.utils.utils import init_torch
+
+
+####################################################################################################
+def train_continue( wandb_id, epoch, Trainer, epoch_continue = -1) :
+
+  devices = init_torch()
+  with_ddp = True
+  par_rank, par_size = setup_ddp( with_ddp)
+
+  cf = Config().load_json( wandb_id)
+
+  cf.num_accs_per_task = len(devices)   # number of GPUs / accelerators per task
+  cf.with_ddp = with_ddp
+  cf.par_rank = par_rank
+  cf.par_size = par_size
+  cf.optimizer_zero = False
+  cf.attention = False
+  # name has changed but ensure backward compatibility
+  if hasattr( cf, 'loader_num_workers') :
+    cf.num_loader_workers = cf.loader_num_workers
+  if not hasattr( cf, 'n_size'):
+    cf.n_size = [36, 0.25*9*6, 0.25*9*12] 
+  if not hasattr(cf, 'num_samples_per_epoch'):
+    cf.num_samples_per_epoch = 1024
+  if not hasattr(cf, 'num_samples_validate'):
+    cf.num_samples_validate = 128
+  if not hasattr(cf, 'with_mixed_precision'):
+    cf.with_mixed_precision = True
+
+  if not hasattr(cf, 'years_val'):
+    cf.years_val = cf.years_test
+  
+  #cf.with_mixed_precision = False
+  # any parameter in cf can be overwritten when training is continued, e.g. we can increase the 
+  # masking rate 
+  # cf.fields = [ [ 'specific_humidity', [ 1, 2048, [ ], 0 ], 
+  #                               [ 96, 105, 114, 123, 137 ], 
+  #                               [12, 6, 12], [3, 9, 9], [0.5, 0.9, 0.1, 0.05] ] ]
+
+  setup_wandb( cf.with_wandb, cf, par_rank, project_name='train', mode='offline')  
+  # resuming a run requires online mode, which is not available everywhere
+  #setup_wandb( cf.with_wandb, cf, par_rank, wandb_id = wandb_id)  
+  
+  if cf.with_wandb and 0 == cf.par_rank :
+    cf.write_json( wandb)
+    cf.print()
+
+  if -1 == epoch_continue :
+    epoch_continue = epoch
+
+  # run
+  trainer = Trainer.load( cf, wandb_id, epoch, devices)
+  print( 'Loaded run \'{}\' at epoch {}.'.format( wandb_id, epoch))
+  trainer.run( epoch_continue)
+
+####################################################################################################
+def train() :
+
+  devices = init_torch()
+  with_ddp = True
+  par_rank, par_size = setup_ddp( with_ddp)
+
+  torch.backends.cuda.matmul.allow_tf32 = True
+
+  cf = Config()
+  # parallelization
+  cf.with_ddp = with_ddp
+  cf.num_accs_per_task = len(devices)   # number of GPUs / accelerators per task
+  cf.par_rank = par_rank
+  cf.par_size = par_size
+  
+  # format: list of fields where for each field the list is 
+  # [ name , 
+  #   [ dynamic or static field { 1, 0 }, embedding dimension, , device id ],
+  #   [ vertical levels ],
+  #   [ num_tokens],
+  #   [ token size],
+  #   [ total masking rate, rate masking, rate noising, rate for multi-res distortion]
+  # ]
+
+  # cf.fields = [ 
+  #               [ 'velocity_u', [ 1, 1024, ['velocity_v', 'temperature'], 0 ], 
+  #                               [ 96, 105, 114, 123, 137 ], 
+  #                               [12, 3, 6], [3, 18, 18], [0.5, 0.9, 0.2, 0.05] ],
+  #               [ 'velocity_v', [ 1, 1024, ['velocity_u', 'temperature'], 1 ], 
+  #                               [ 96, 105, 114, 123, 137 ], 
+  #                               [12, 3, 6], [3, 18, 18], [0.5, 0.9, 0.2, 0.05] ], 
+  #               [ 'specific_humidity', [ 1, 1024, ['velocity_u', 'velocity_v', 'temperature'], 2 ], 
+  #                             [ 96, 105, 114, 123, 137 ], 
+  #                             [12, 3, 6], [3, 18, 18], [0.5, 0.9, 0.2, 0.05] ],
+  #               [ 'velocity_z', [ 1, 1024, ['velocity_u', 'velocity_v', 'temperature'], 3 ], 
+  #                             [ 96, 105, 114, 123, 137 ], 
+  #                             [12, 3, 6], [3, 18, 18], [0.5, 0.9, 0.2, 0.05] ],
+  #                [ 'temperature', [ 1, 512, ['velocity_u', 'velocity_v', 'specific_humidity'], 3 ], 
+  #                             [ 96, 105, 114, 123, 137 ], 
+  #                             [12, 2, 4], [3, 27, 27], [0.5, 0.9, 0.2, 0.05], 'local' ],
+  #             ]
+
+  cf.fields = [ 
+                [ 'velocity_u', [ 1, 1024, ['velocity_v', 'temperature'], 0, ['j8dwr5qj', -2] ], 
+                                [ 96, 105, 114, 123, 137 ], 
+                                [12, 3, 6], [3, 18, 18], [0.5, 0.9, 0.2, 0.05] ],
+                [ 'velocity_v', [ 1, 1024, ['velocity_u', 'temperature'], 1, ['0tlnm5up', -2] ], 
+                                [ 96, 105, 114, 123, 137 ], 
+                                [12, 3, 6], [3, 18, 18], [0.5, 0.9, 0.2, 0.05] ], 
+                [ 'specific_humidity', [ 1, 1024, ['velocity_u', 'velocity_v', 'temperature'], 2, ['v63l01zu', -2] ], 
+                              [ 96, 105, 114, 123, 137 ], 
+                              [12, 3, 6], [3, 18, 18], [0.5, 0.9, 0.2, 0.05] ],
+                [ 'velocity_z', [ 1, 1024, ['velocity_u', 'velocity_v', 'temperature'], 3, ['9l1errbo', -2] ], 
+                              [ 96, 105, 114, 123, 137 ], 
+                              [12, 3, 6], [3, 18, 18], [0.5, 0.9, 0.2, 0.05] ],
+                 [ 'temperature', [ 1, 1024, ['velocity_u', 'velocity_v', 'specific_humidity'], 3, ['7ojls62c', -2] ], 
+                              [ 96, 105, 114, 123, 137 ], 
+                              [12, 2, 4], [3, 27, 27], [0.5, 0.9, 0.2, 0.05], 'local' ],
+                # ['total_precip', [1, 1536, ['velocity_u', 'velocity_v', 'velocity_z', 'specific_humidity'], 3, ['3kdutwqb', 900]], 
+                #               [0], 
+                #               [12, 6, 12], [3, 9, 9], [0.25, 0.9, 0.1, 0.05]] 
+              ]
+
+  cf.fields_prediction = [
+                          ['velocity_u', 0.225], ['velocity_v', 0.225], 
+                          ['specific_humidity', 0.15], ['velocity_z', 0.1], ['temperature', 0.2]
+             #             ['total_precip', 0.1] 
+              ]
+
+  cf.fields_targets = []
+
+  cf.years_train = list( range( 1979, 2021))
+  cf.years_val = [2021]  #[2018] 
+  cf.month = None
+  cf.geo_range_sampling = [[ -90., 90.], [ 0., 360.]]
+  cf.time_sampling = 1   # sampling rate for time steps
+  # random seeds
+  cf.torch_seed = torch.initial_seed()
+  # training params
+  cf.batch_size_validation = 1 #64
+
+  cf.batch_size = 96
+  cf.num_epochs = 128
+  cf.num_samples_per_epoch = 4096*12
+  cf.num_samples_validate = 128*12
+  cf.num_loader_workers = 5
+  
+  # additional infos
+  cf.size_token_info = 8
+  cf.size_token_info_net = 16
+  cf.grad_checkpointing = True
+  cf.with_cls = False
+  # network config
+  cf.with_mixed_precision = True 
+  cf.with_layernorm = True
+  cf.coupling_num_heads_per_field = 1
+  cf.dropout_rate = 0.05
+  cf.with_qk_lnorm = True
+  # encoder
+  cf.encoder_num_layers = 6
+  cf.encoder_num_heads = 16
+  cf.encoder_num_mlp_layers = 2
+  cf.encoder_att_type = 'dense'
+  # decoder
+  cf.decoder_num_layers = 6
+  cf.decoder_num_heads = 16
+  cf.decoder_num_mlp_layers = 2
+  cf.decoder_self_att = False
+  cf.decoder_cross_att_ratio = 0.5
+  cf.decoder_cross_att_rate = 1.0
+  cf.decoder_att_type = 'dense'
+  # tail net
+  cf.net_tail_num_nets = 16
+  cf.net_tail_num_layers = 0
+  # loss
+  cf.losses = ['mse_ensemble', 'stats']  # mse, mse_ensemble, stats, crps, weighted_mse
+  # training
+  cf.optimizer_zero = False
+  
+  cf.lr_start = 0.00001 #5. * 10e-7
+  cf.lr_max = 0.00002
+  cf.lr_min = 0.00001
+  cf.weight_decay = 0.025 #0.1
+  cf.lr_decay_rate = 1.025
+  cf.lr_start_epochs = 3
+  cf.model_log_frequency = 256 #save checkpoint every X batches
+
+  # BERT
+  # strategies: 'BERT', 'forecast', 'temporal_interpolation'
+  cf.BERT_strategy = 'BERT' #'BERT'
+  cf.forecast_num_tokens = 2      # only needed / used for BERT_strategy 'forecast
+
+  cf.BERT_fields_synced = False   # apply synchronized / identical masking to all fields 
+                                  # (fields need to have same BERT params for this to have effect)
+  cf.BERT_mr_max = 2              # maximum reduction rate for resolution
+  
+  # debug / output
+  cf.log_test_num_ranks = 0
+  cf.save_grads = False
+  cf.profile = False
+  cf.test_initial = False
+  cf.attention = False
+
+  cf.rng_seed = None 
+
+  # usually use %>wandb offline to switch to disable syncing with server
+  cf.with_wandb = True
+  setup_wandb( cf.with_wandb, cf, par_rank, 'train', mode='offline')  
+
+  # cf.file_path = '/p/scratch/atmo-rep/data/era5_1deg/months/era5_y2021_res100_chunk32.zarr'
+  # cf.file_path = '/ec/res4/scratch/nacl/atmorep/era5_y2021_res100_chunk32.zarr'
+  # # # in steps x lat_degrees x lon_degrees
+  # cf.n_size = [36, 1*9*6, 1.*9*12]
+
+  # # # # # cf.file_path = '/p/scratch/atmo-rep/data/era5_1deg/months/era5_y2021_res025_chunk16.zarr'
+  # # # # cf.file_path = '/p/scratch/atmo-rep/data/era5_1deg/months/era5_y2021_res025_chunk32.zarr'
+  # cf.file_path = '/ec/res4/scratch/nacl/atmorep/era5_y2021_res025_chunk32.zarr'
+  # # # 
+  # # # cf.file_path = '/p/scratch/atmo-rep/data/era5_1deg/months/era5_y2021_res025_chunk8.zarr'
+  # # cf.file_path = '/ec/res4/scratch/nacl/atmorep/era5_y2021_res025_chunk8_lat180_lon180.zarr'
+  # # # cf.file_path = '/ec/res4/scratch/nacl/atmorep/era5_y2021_res025_chunk16.zarr'
+  cf.file_path = '/gpfs/scratch/ehpc03/era5_y1979_2021_res025_chunk8.zarr/'
+  # # # in steps x lat_degrees x lon_degrees
+  cf.n_size = [36, 0.25*9*6, 0.25*9*12]
+
+  # cf.file_path = '/ec/res4/scratch/nacl/atmorep/era5_y2021_res100_chunk16.zarr'
+  #cf.file_path = '/p/scratch/atmo-rep/data/era5_1deg/months/era5_y2021_res100_chunk16.zarr'
+  #cf.n_size = [36, 1*9*6, 1.*9*12]
+
+  if cf.with_wandb and 0 == cf.par_rank :
+    cf.write_json( wandb)
+    cf.print()
+
+  trainer = Trainer_BERT( cf, devices).create()
+  trainer.run()
+
+####################################################################################################
+if __name__ == '__main__':
+  
+  try :
+    
+    train()
+
+    #  wandb_id, epoch, epoch_continue = 'uvrdtc0a', 95, 95 #multiformer from scratch
+    #  Trainer = Trainer_BERT
+    #  train_continue( wandb_id, epoch, Trainer, epoch_continue)
+
+  except :
+    
+    extype, value, tb = sys.exc_info()
+    traceback.print_exc()
+    pdb.post_mortem(tb)

vendor/atmorep-official/atmorep/core/trainer.py

@@ -0,0 +1,863 @@
+####################################################################################################
+#
+#  Copyright (C) 2022
+#
+####################################################################################################
+#
+#  project     : atmorep
+#
+#  author      : atmorep collaboration
+# 
+#  description :
+#
+#  license     :
+#
+####################################################################################################
+
+import torch
+import torchinfo
+import numpy as np
+import time
+import code
+
+from pathlib import Path
+import os
+import datetime
+import functools
+
+import wandb
+# import horovod.torch as hvd
+import torch.distributed as dist
+from torch.distributed.optim import ZeroRedundancyOptimizer
+import torch.utils.data.distributed
+
+import atmorep.config.config as config
+
+from atmorep.core.atmorep_model import AtmoRep
+from atmorep.core.atmorep_model import AtmoRepData
+from atmorep.training.bert import prepare_batch_BERT_multifield
+from atmorep.transformer.transformer_base import positional_encoding_harmonic
+
+import atmorep.utils.token_infos_transformations as token_infos_transformations
+
+from atmorep.utils.utils import Gaussian, CRPS, kernel_crps, weighted_mse, NetMode, tokenize, detokenize
+from atmorep.datasets.data_writer import write_forecast, write_BERT, write_attention
+from atmorep.datasets.normalizer import denormalize
+
+####################################################################################################
+class Trainer_Base() :
+
+  def __init__( self, cf, devices ) :
+
+    self.cf = cf
+    self.devices = devices
+    self.device_in = devices[0]
+    self.device_out = devices[-1]
+
+    self.fields_prediction_idx = []
+    self.loss_weights = torch.zeros( len(cf.fields_prediction) )
+    for ifield, field in enumerate(cf.fields_prediction) :
+      self.loss_weights[ifield] = self.cf.fields_prediction[ifield][1]
+      for idx, field_info in enumerate(cf.fields) :
+        if field_info[0] == field[0] :
+          self.fields_prediction_idx.append( idx)
+          break
+    self.loss_weights = self.loss_weights.to( self.device_out)
+
+    self.MSELoss = torch.nn.MSELoss()
+
+    # transformation for token infos
+    if hasattr( cf, 'token_infos_transformation') :
+      self.tok_infos_trans = getattr( token_infos_transformations, cf.token_infos_transformation)
+    else :
+      self.tok_infos_trans = getattr( token_infos_transformations, 'identity')
+
+    if 0 == cf.par_rank :
+      directory = Path( config.path_results, 'id{}'.format( cf.wandb_id))
+      if not os.path.exists(directory):
+        os.makedirs( directory)
+      directory = Path( config.path_models, 'id{}'.format( cf.wandb_id))
+      if not os.path.exists(directory):
+        os.makedirs( directory)
+
+  ###################################################
+  def create( self, load_embeds=True) :
+    net = AtmoRep( self.cf) 
+    self.model = AtmoRepData( net)
+
+    self.model.create( self.pre_batch, self.devices, load_embeds)
+    
+    # TODO: pass the properly to model / net
+    self.model.net.encoder_to_decoder = self.encoder_to_decoder
+    self.model.net.decoder_to_tail = self.decoder_to_tail
+    return self
+
+  ###################################################
+  @classmethod
+  def load( Typename, cf, model_id, epoch, devices) :
+    trainer = Typename( cf, devices).create( load_embeds=False)
+    trainer.model.net = trainer.model.net.load( model_id, devices, cf, epoch)
+
+    # TODO: pass the properly to model / net
+    trainer.model.net.encoder_to_decoder = trainer.encoder_to_decoder
+    trainer.model.net.decoder_to_tail = trainer.decoder_to_tail
+
+    str = 'Loaded model id = {}{}.'.format( model_id, f' at epoch = {epoch}' if epoch> -2 else '')
+    print( str)
+    return trainer
+
+  ###################################################
+  def save( self, epoch) :
+    self.model.net.save( epoch)
+
+  ###################################################
+  def get_learn_rates( self) :
+
+    cf = self.cf
+    size_padding = 5
+    learn_rates = np.zeros( cf.num_epochs + size_padding) 
+  
+    learn_rates[:cf.lr_start_epochs] = np.linspace( cf.lr_start, cf.lr_max, num = cf.lr_start_epochs)
+    lr = learn_rates[cf.lr_start_epochs-1]
+    ic = 0
+    for epoch in range( cf.lr_start_epochs, cf.num_epochs + size_padding) :
+      lr = max( lr / cf.lr_decay_rate, cf.lr_min)
+      learn_rates[epoch] = lr 
+      if ic > 9999 :  # sanity check
+        assert "Maximum number of epochs exceeded."
+
+    return learn_rates
+
+  ###################################################
+  def run( self, epoch = -1) :
+
+    cf = self.cf
+    model = self.model
+
+    learn_rates = self.get_learn_rates()
+    
+    if cf.with_ddp :
+      self.model_ddp = torch.nn.parallel.DistributedDataParallel( model, static_graph=True)
+      if not cf.optimizer_zero :
+        self.optimizer = torch.optim.AdamW( self.model_ddp.parameters(), lr=cf.lr_start,
+                                            weight_decay=cf.weight_decay)
+      else :
+        self.optimizer = ZeroRedundancyOptimizer(self.model_ddp.parameters(),
+                                                              optimizer_class=torch.optim.AdamW,
+                                                              lr=cf.lr_start )
+    else :
+      self.optimizer = torch.optim.AdamW( self.model.parameters(), lr=cf.lr_start,
+                                          weight_decay=cf.weight_decay)
+    
+    self.grad_scaler = torch.cuda.amp.GradScaler(enabled=cf.with_mixed_precision)
+
+    if 0 == cf.par_rank :
+      # print( self.model.net)
+      model_parameters = filter(lambda p: p.requires_grad, self.model_ddp.parameters())
+      num_params = sum([np.prod(p.size()) for p in model_parameters])
+      print( f'Number of trainable parameters: {num_params:,}')
+  
+    if cf.test_initial :
+      cur_test_loss = self.validate( epoch, cf.BERT_strategy).cpu().numpy()
+      test_loss = np.array( [cur_test_loss])
+    else :
+      # generic value based on data normalization
+      test_loss = np.array( [1.0]) 
+    epoch += 1
+   
+    if cf.profile :
+      lr = learn_rates[epoch]
+      for g in self.optimizer.param_groups:
+        g['lr'] = lr
+      self.profile()
+
+    # training loop
+    while True :
+
+      if epoch >= cf.num_epochs :
+        break
+        
+      lr = learn_rates[epoch]
+      for g in self.optimizer.param_groups:
+        g['lr'] = lr
+
+      tstr = datetime.datetime.now().strftime("%H:%M:%S")
+      print( '{} : {} :: batch_size = {}, lr = {}'.format( epoch, tstr, cf.batch_size, lr) )
+
+      self.train( epoch)
+      
+      if cf.with_wandb and 0 == cf.par_rank :
+        self.save( epoch)
+
+      cur_test_loss = self.validate( epoch, cf.BERT_strategy).cpu().numpy()
+      # self.validate( epoch, 'forecast')
+
+      # save model 
+      if cur_test_loss < test_loss.min() :
+        self.save( -2)
+      test_loss = np.append( test_loss, [cur_test_loss])
+
+      epoch += 1
+
+    tstr = datetime.datetime.now().strftime("%H:%M:%S")
+    print( 'Finished training at {} with test loss = {}.'.format( tstr, test_loss[-1]) )
+
+    # save final network
+    if cf.with_wandb and 0 == cf.par_rank :
+      self.save( -2)
+    
+  ###################################################
+  def train( self, epoch):
+
+    model = self.model
+    cf = self.cf
+    
+    model.mode( NetMode.train)
+    self.optimizer.zero_grad()
+
+    loss_total = [[] for i in range(len(cf.losses)) ]
+    std_dev_total = [[] for i in range(len(self.fields_prediction_idx)) ]
+    mse_loss_total = []
+    grad_loss_total = []
+    ctr = 0
+
+    self.optimizer.zero_grad()
+    time_start = time.time()
+    
+    for batch_idx in range( model.len( NetMode.train)) :
+      
+      batch_data = self.model.next()
+      _, _, _, tmksd_list, weight_list = batch_data[0]
+      with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=cf.with_mixed_precision):
+        batch_data = self.prepare_batch( batch_data)
+        preds, _ = self.model_ddp( batch_data)
+        loss, mse_loss, losses = self.loss( preds, batch_idx, tmksd_list, weight_list)
+      
+      self.grad_scaler.scale(loss).backward()
+      self.grad_scaler.step(self.optimizer)
+      self.grad_scaler.update()
+
+      self.optimizer.zero_grad()
+
+      [loss_total[idx].append( losses[key]) for idx, key in enumerate(losses)]
+      mse_loss_total.append( mse_loss.detach().cpu() )
+      grad_loss_total.append( loss.detach().cpu() )
+      [std_dev_total[idx].append( pred[1].detach().cpu()) for idx, pred in enumerate(preds)]
+
+      # logging
+
+      if int((batch_idx * cf.batch_size) / 8) > ctr :
+        
+        # wandb logging
+        if cf.with_wandb and (0 == cf.par_rank) :
+          loss_dict = { "training loss": torch.mean( torch.tensor( mse_loss_total)), 
+                        "gradient loss": torch.mean( torch.tensor( grad_loss_total)) }
+          # log individual loss terms for individual fields
+          for idx, cur_loss in enumerate(loss_total) :
+            loss_name = self.cf.losses[idx]
+            lt = torch.tensor(cur_loss)
+            for i, field in enumerate(cf.fields_prediction) :
+              idx_name = loss_name + ', ' + field[0]
+              idx_std_name = 'stddev, ' + field[0]
+              loss_dict[idx_name] = torch.mean( lt[:,i]).cpu().detach()
+              loss_dict[idx_std_name] = torch.mean(torch.cat(std_dev_total[i],0)).cpu().detach()
+          wandb.log( loss_dict )
+      
+          # console output
+          samples_sec = cf.batch_size / (time.time() - time_start)
+          str = 'epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:1.5f} : {:1.5f} :: {:1.5f} ({:2.2f} s/sec)'
+          print( str.format( epoch, batch_idx, model.len( NetMode.train),
+                            100. * batch_idx/model.len(NetMode.train), 
+                            torch.mean( torch.tensor( grad_loss_total)), 
+                            torch.mean(torch.tensor(mse_loss_total)),
+                            torch.mean( preds[0][1]), samples_sec ), flush=True)
+    
+          # save model (use -2 as epoch to indicate latest, stored without epoch specification)
+          if batch_idx % cf.model_log_frequency == 0 :
+            self.save( -2)
+
+        # reset
+        loss_total = [[] for i in range(len(cf.losses)) ]
+        mse_loss_total = []
+        grad_loss_total = []
+        std_dev_total = [[] for i in range(len(self.fields_prediction_idx)) ]
+        
+        ctr += 1
+        time_start = time.time()
+
+    # save gradients
+    if cf.save_grads and cf.with_wandb and (0 == cf.par_rank) :
+    
+      dir_name = './grads/id{}'.format( cf.wandb_id)
+      if not os.path.exists(dir_name):
+        os.makedirs(dir_name)
+
+      rmsprop_ws = []
+      for k in range( len(self.optimizer.state_dict()['state']) ) : 
+        rmsprop_ws.append(self.optimizer.state_dict()['state'][k]['exp_avg_sq'].mean().unsqueeze(0))
+      rmsprop_ws = torch.cat( rmsprop_ws)
+      fname = '{}/{}_epoch{}_rmsprop.npy'.format( dir_name, cf.wandb_id, epoch)
+      np.save( fname, rmsprop_ws.cpu().detach().numpy() ) 
+     
+      idx = 0
+      for name, param in self.model.named_parameters():
+        if param.requires_grad : 
+          fname = '{}/{}_epoch{}_{:05d}_{}_grad.npy'.format( dir_name, cf.wandb_id, epoch, idx,name)
+          np.save( fname, param.grad.cpu().detach().numpy() ) 
+          idx += 1
+
+    # clean memory
+    self.optimizer.zero_grad()
+    del batch_data, loss, loss_total, mse_loss_total, grad_loss_total, std_dev_total
+      
+  ###################################################
+  def profile( self):
+
+    model = self.model
+    cf = self.cf
+
+    model.mode( NetMode.train)
+    self.optimizer.zero_grad()
+
+    # See https://pytorch.org/tutorials/intermediate/tensorboard_profiler_tutorial.html
+    # for details on how to load and analyse report
+    # https://pytorch.org/blog/trace-analysis-for-masses/
+
+    # do for all par_ranks to avoid that they run out of sync
+    print( '---------------------------------')
+    print( 'Profiling:')
+    pname = './logs/profile_par_rank' + str(cf.par_rank) + '_' + cf.wandb_id + '/profile'
+    with torch.profiler.profile( activities=[torch.profiler.ProfilerActivity.CPU,
+                                            torch.profiler.ProfilerActivity.CUDA],
+                      schedule=torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=2),
+                      on_trace_ready=torch.profiler.tensorboard_trace_handler(pname),
+                      profile_memory=True, record_shapes=True, with_stack=True) as prof:
+      for batch_idx in range( 2 * (1+1+3) ) :
+
+        batch_data = self.model.next()
+
+        with torch.autocast(device_type='cuda',dtype=torch.float16, enabled=cf.with_mixed_precision):
+          batch_data = self.prepare_batch( batch_data)
+          preds, _ = self.model_ddp( batch_data)
+          loss, mse_loss, losses = self.loss( preds, batch_idx)
+
+        self.grad_scaler.scale(loss).backward()
+        self.grad_scaler.step(self.optimizer)
+        self.grad_scaler.update()
+        self.optimizer.zero_grad()
+
+        prof.step()
+
+    print( 'Profiling finished.')
+    print( '---------------------------------')
+
+  ###################################################
+  def validate( self, epoch, BERT_test_strategy = 'BERT'):
+
+    cf = self.cf
+    BERT_strategy_train = cf.BERT_strategy
+    cf.BERT_strategy = BERT_test_strategy
+    self.model.mode( NetMode.test)
+    total_loss = 0.
+    total_losses = torch.zeros( len(self.fields_prediction_idx) )
+    test_len = 0
+
+    self.mode_test = True
+      
+    # run test set evaluation
+    with torch.no_grad() : 
+      for it in range( self.model.len( NetMode.test)) :
+        batch_data = self.model.next()
+        if cf.par_rank < cf.log_test_num_ranks :
+          # keep on cpu since it will otherwise clog up GPU memory
+          (sources, _ , targets, tmis_list, _) = batch_data[0]
+          log_sources = ( [source.detach().clone().cpu() for source in sources ],
+                          [target.detach().clone().cpu() for target in targets ],
+                           tmis_list)
+        
+        with torch.autocast(device_type='cuda',dtype=torch.float16,enabled=cf.with_mixed_precision):
+          batch_data = self.prepare_batch( batch_data)
+          preds, atts = self.model( batch_data)
+        loss = torch.tensor( 0.)
+        ifield = 0
+        for pred, idx in zip( preds, self.fields_prediction_idx) :
+          
+          target = self.targets[idx]
+          # hook for custom test loss
+          self.test_loss( pred, target)
+          # base line loss
+          cur_loss = self.MSELoss( pred[0], target = target ).cpu().item()
+    
+          loss += cur_loss 
+          total_losses[ifield] += cur_loss
+          ifield += 1
+  
+        total_loss += loss
+        test_len += 1
+        
+        # store detailed results on current test set for book keeping
+        if cf.par_rank < cf.log_test_num_ranks :
+          log_preds = [[p.detach().clone().cpu() for p in pred] for pred in preds]
+          self.log_validate( epoch, it, log_sources, log_preds)
+          if cf.attention:
+            self.log_attention( epoch, it, atts)
+                             
+    # average over all nodes
+    total_loss /= test_len * len(self.cf.fields_prediction)
+    total_losses /= test_len
+
+    if cf.with_ddp :
+      total_loss_cuda = total_loss.cuda()
+      total_losses_cuda = total_losses.cuda()
+      dist.all_reduce( total_loss_cuda, op=torch.distributed.ReduceOp.AVG )
+      dist.all_reduce( total_losses_cuda, op=torch.distributed.ReduceOp.AVG )
+      total_loss = total_loss_cuda.cpu()
+      total_losses = total_losses_cuda.cpu()
+
+    if 0 == cf.par_rank :
+      print( 'validation loss for strategy={} at epoch {} : {}'.format( BERT_test_strategy,
+                                                                  epoch, total_loss),
+                                                                  flush=True)
+    if cf.with_wandb and (0 == cf.par_rank) :
+      loss_dict = {"val. loss {}".format(BERT_test_strategy) : total_loss}
+      total_losses = total_losses.cpu().detach()
+      for i, field in enumerate(cf.fields_prediction) :
+        idx_name = 'val., {}, '.format(BERT_test_strategy) + field[0]
+        loss_dict[idx_name] = total_losses[i]
+        print( 'validation loss for {} : {}'.format( field[0], total_losses[i] ))
+      wandb.log( loss_dict)
+    batch_data = []
+    torch.cuda.empty_cache()
+
+    cf.BERT_strategy = BERT_strategy_train
+    self.mode_test = False
+    
+    return total_loss
+
+  ###################################################
+  def test_loss( self, pred, target) :
+    '''Hook for custom test loss'''
+    pass
+
+  ###################################################
+  def loss( self, preds, batch_idx = 0, tmidx_list = None, weights_list = None) :
+
+    # TODO: move implementations to individual files
+
+    cf = self.cf
+    mse_loss_total = torch.tensor( 0.,)
+    losses = dict(zip(cf.losses,[[] for loss in cf.losses ]))
+    
+    for pred, idx in zip( preds, self.fields_prediction_idx) :
+      target = self.targets[idx]
+
+      mse_loss = self.MSELoss( pred[0], target = target) 
+      mse_loss_total += mse_loss.cpu().detach()
+
+      # MSE loss
+      if 'mse' in self.cf.losses :
+        losses['mse'].append( mse_loss)
+
+      # MSE loss
+      if 'mse_ensemble' in self.cf.losses :
+        loss_en = torch.tensor( 0., device=target.device)
+        for en in torch.transpose( pred[2], 1, 0) :
+          loss_en += self.MSELoss( en, target = target) 
+        losses['mse_ensemble'].append( loss_en / pred[2].shape[1])
+
+      if 'weighted_mse' in self.cf.losses :
+        loss_en = torch.tensor( 0., device=target.device)
+        field_info = cf.fields[idx]
+        token_size = field_info[4]
+
+        weights = torch.Tensor(np.array([w for batch in weights_list[idx] for w in batch]))
+        weights = weights.view(*weights.shape, 1, 1).repeat(1, 1, token_size[0], token_size[2]).swapaxes(1, 2)
+        weights = weights.reshape([weights.shape[0], -1]).to(target.get_device())
+       
+        for en in torch.transpose( pred[2], 1, 0) :
+          loss_en += weighted_mse( en, target, weights)          
+
+        losses['weighted_mse'].append( loss_en / pred[2].shape[1]) 
+
+      # Generalized cross entroy loss for continuous distributions
+      if 'stats' in self.cf.losses :
+        stats_loss = Gaussian( target, pred[0], pred[1])  
+        diff = (stats_loss-1.) 
+        # stats_loss = 0.01 * torch.mean( diff * diff) + torch.mean( torch.sqrt(torch.abs( pred[1])) )
+        stats_loss = torch.mean( diff * diff) + torch.mean( torch.sqrt( torch.abs( pred[1])) )
+        losses['stats'].append( stats_loss) 
+      
+      # Generalized cross entroy loss for continuous distributions
+      if 'stats_area' in self.cf.losses :
+        diff = torch.abs( torch.special.erf( (target - pred[0]) / (pred[1] * pred[1])) )
+        stats_area = 0.2 * torch.mean( diff * diff) + torch.mean( torch.sqrt(torch.abs( pred[1])) )
+        losses['stats_area'].append( stats_area)
+
+      # CRPS score
+      if 'crps' in self.cf.losses :
+        crps_loss = torch.mean( CRPS( target, pred[0], pred[1]))
+        losses['crps'].append( crps_loss)
+
+      if 'kernel_crps' in self.cf.losses :
+        kcrps_loss = torch.mean( kernel_crps( target,torch.transpose( pred[2], 1, 0)))
+        losses['kernel_crps'].append( kcrps_loss)
+    
+    #TODO: uncomment it and add it when running in debug mode
+    # field_losses = ""
+    # for ifield, field in enumerate(cf.fields):
+    #   ifield_loss = 0
+    #   for key in losses :  
+    #     ifield_loss += losses[key][ifield].to(self.device_out)
+    #   ifield_loss /= len(losses.keys())
+    #   field_losses +=  f"{field[0]}: {ifield_loss}; "
+    # print(field_losses, flush = True)
+          
+    loss = torch.tensor( 0., device=self.device_out)
+    tot_weight = torch.tensor( 0., device=self.device_out)
+    for key in losses :
+      #print( 'LOSS : {} :: {}'.format( key, losses[key]))
+      for ifield, val in enumerate(losses[key]) :
+        loss += self.loss_weights[ifield] * val.to( self.device_out)
+        tot_weight += self.loss_weights[ifield] 
+    loss /= tot_weight
+    mse_loss = mse_loss_total / len(self.cf.fields_prediction)
+
+    return loss, mse_loss, losses
+
+####################################################################################################
+class Trainer_BERT( Trainer_Base) :
+
+  ###################################################
+  def __init__( self, cf, devices) :
+    
+    Trainer_Base.__init__( self, cf, devices)
+
+    self.rng_seed = cf.rng_seed
+    if not self.rng_seed :
+      self.rng_seed = int(torch.randint( 100000000, (1,))) 
+    # TODO: generate only rngs that are needed
+    ll = len(cf.fields) * 8 #len(cf.vertical_levels)
+    if cf.BERT_fields_synced :
+      self.rngs = [np.random.default_rng(self.rng_seed) for _ in range(ll)]
+    else : 
+      self.rngs = [np.random.default_rng(self.rng_seed+i) for i in range(ll)]
+
+    # batch preprocessing to be done in loader (mainly for performance reasons since it's 
+    # parallelized there)
+    self.pre_batch = functools.partial( prepare_batch_BERT_multifield, self.cf, self.rngs, 
+                                                          self.cf.fields, self.cf.BERT_strategy )
+
+  ###################################################
+  def prepare_batch( self, xin) :
+    '''Move data to device and some additional final preprocessing before model eval'''
+
+    cf = self.cf
+    devs = self.devices
+
+    # unpack loader output
+    # xin[0] since BERT does not have targets
+    (sources, token_infos, targets, fields_tokens_masked_idx_list, _) = xin[0]
+    (self.sources_idxs, self.sources_info) = xin[2]
+    
+    # network input
+    batch_data = [ ( sources[i].to( devs[ cf.fields[i][1][3] ], non_blocking=True), 
+                    self.tok_infos_trans(token_infos[i]).to( self.devices[0], non_blocking=True)) 
+                                                                  for i in range(len(sources)) ]
+
+    # store token number since BERT selects sub-cube (optionally)
+    self.num_tokens = []
+    for field_idx in range(len(batch_data)) :
+      self.num_tokens.append( list(batch_data[field_idx][0].shape[2:5]))
+
+    # target
+    self.targets = []
+    for ifield in self.fields_prediction_idx :
+      self.targets.append( targets[ifield].to( devs[cf.fields[ifield][1][3]], non_blocking=True ))
+
+    # idxs of masked tokens
+    tmi_out = [ ]
+    for i,tmi in enumerate(fields_tokens_masked_idx_list) :
+      cdev = devs[cf.fields[i][1][3]]
+      tmi_out += [ [torch.cat(tmi_l).to( cdev, non_blocking=True) for tmi_l in tmi] ]
+    self.tokens_masked_idx = tmi_out
+  
+    return batch_data
+
+  ###################################################
+  def encoder_to_decoder( self, embeds_layers) :
+    return ([embeds_layers[i][-1] for i in range(len(embeds_layers))] , embeds_layers )
+
+  ###################################################
+  def decoder_to_tail( self, idx_pred, pred) :
+    '''Positional encoding of masked tokens for tail network evaluation'''
+
+    field_idx = self.fields_prediction_idx[idx_pred]
+    dev = self.devices[ self.cf.fields[field_idx][1][3] ]
+    target_idx = self.tokens_masked_idx[field_idx]
+    assert len(target_idx) > 0, 'no masked tokens but target variable'
+    
+    # select "fixed" masked tokens for loss computation
+    
+    # flatten token dimensions: remove space-time separation
+    pred = torch.flatten( pred, 2, 3).to( dev)
+    # extract masked token level by level
+    pred_masked = []
+    for lidx, level in enumerate(self.cf.fields[field_idx][2]) :
+      # select masked tokens, flattened along batch dimension for easier indexing and processing
+      pred_l = torch.flatten( pred[:,lidx], 0, 1)
+      pred_masked.append( pred_l[ target_idx[lidx] ])
+    
+    # flatten along level dimension, for loss evaluation we effectively have level, batch, ...
+    # as ordering of dimensions
+    pred_masked = torch.cat( pred_masked, 0)
+
+    return pred_masked
+
+  ###################################################
+  def log_validate( self, epoch, bidx, log_sources, log_preds) :
+    '''Hook for logging: output associated with concrete training strategy.'''
+
+    if not hasattr( self.cf, 'wandb_id') :
+      return
+
+    if 'forecast' in self.cf.BERT_strategy :
+      self.log_validate_forecast( epoch, bidx, log_sources, log_preds)
+    elif 'BERT' in self.cf.BERT_strategy or 'temporal_interpolation' == self.cf.BERT_strategy :
+      self.log_validate_BERT( epoch, bidx, log_sources, log_preds)
+    else :
+      assert False
+  
+  ###################################################
+  def log_validate_forecast( self, epoch, batch_idx, log_sources, log_preds) :
+    '''Logging for BERT_strategy=forecast.'''
+
+    cf = self.cf
+
+    # save source: remains identical so just save ones
+    (sources, targets, _) = log_sources
+
+    sources_out, targets_out, preds_out, ensembles_out = [ ], [ ], [ ], [ ] 
+    batch_size = len(self.sources_info)  
+    # reconstruct geo-coords (identical for all fields)
+    forecast_num_tokens = 1
+    if hasattr( cf, 'forecast_num_tokens') :
+      forecast_num_tokens = cf.forecast_num_tokens
+  
+    coords = []
+    for fidx, field_info in enumerate(cf.fields) :
+      # reshape from tokens to contiguous physical field
+      num_levels = len(field_info[2])
+      source = detokenize( sources[fidx].cpu().detach().numpy())
+      # recover tokenized shape
+      target = detokenize( targets[fidx].cpu().detach().numpy().reshape( [ num_levels, -1, 
+                           forecast_num_tokens, *field_info[3][1:], *field_info[4] ]).swapaxes(0,1))
+     
+      coords_b = []
+  
+      for bidx in range(batch_size):
+        dates   = self.sources_info[bidx][0]
+        lats    = self.sources_info[bidx][1]
+        lons    = self.sources_info[bidx][2]
+        dates_t = self.sources_info[bidx][0][ -forecast_num_tokens*field_info[4][0] : ]
+        
+        lats_idx = self.sources_idxs[bidx][1]
+        lons_idx = self.sources_idxs[bidx][2]
+
+        for vidx, _ in enumerate(field_info[2]) :
+          normalizer, year_base = self.model.normalizer( fidx, vidx, lats_idx, lons_idx) 
+          source[bidx,vidx] = denormalize( source[bidx,vidx], normalizer, dates, year_base)
+          target[bidx,vidx] = denormalize( target[bidx,vidx], normalizer, dates_t, year_base)
+
+        coords_b += [[dates, 90.-lats, lons, dates_t]]
+
+      # append
+      sources_out.append( [field_info[0], source])
+      targets_out.append( [field_info[0], target])
+      coords.append(coords_b)
+
+    # process predicted fields
+    for fidx, fn in enumerate(cf.fields_prediction) :
+      field_info = cf.fields[ self.fields_prediction_idx[fidx] ]
+      num_levels = len(field_info[2])
+      # predictions
+      pred = log_preds[fidx][0].cpu().detach().numpy()
+      pred = detokenize( pred.reshape( [ num_levels, -1, 
+                                    forecast_num_tokens, *field_info[3][1:], *field_info[4] ]).swapaxes(0,1))
+      # ensemble
+      ensemble = log_preds[fidx][2].cpu().detach().numpy().swapaxes(0,1)
+      ensemble = detokenize( ensemble.reshape( [ cf.net_tail_num_nets, num_levels, -1, 
+                                            forecast_num_tokens, *field_info[3][1:], *field_info[4] ]).swapaxes(1, 2)).swapaxes(0,1)
+      
+      # denormalize
+      for bidx in range(batch_size) : 
+        lats  = self.sources_info[bidx][1]
+        lons  = self.sources_info[bidx][2]
+        dates_t = self.sources_info[bidx][0][ -forecast_num_tokens*field_info[4][0] : ]
+       
+        for vidx, vl in enumerate(field_info[2]) :
+          normalizer, year_base = self.model.normalizer( self.fields_prediction_idx[fidx], vidx, lats_idx, lons_idx)
+          pred[bidx,vidx] = denormalize( pred[bidx,vidx], normalizer, dates_t, year_base)
+          ensemble[bidx,:,vidx] = denormalize(ensemble[bidx,:,vidx], normalizer, dates_t, year_base)
+
+      # append
+      preds_out.append( [fn[0], pred])
+      ensembles_out.append( [fn[0], ensemble])
+
+    levels = np.array(cf.fields[0][2])
+    
+    write_forecast( cf.wandb_id, epoch, batch_idx,
+                                 levels, sources_out,
+                                 targets_out, preds_out,
+                                 ensembles_out, coords)
+  
+  ###################################################
+  
+  def split_data(self, data, idx_list, token_size) :
+    lens_batches = [[len(t) for t in tt] for tt in idx_list]
+    lens_levels = [torch.tensor( tt).sum() for tt in lens_batches]
+    data_b = torch.split( data, lens_levels)    
+    # split according to batch
+    return [torch.split( data_b[vidx], lens) for vidx,lens in enumerate(lens_batches)]
+
+  def get_masked_data(self, field_info, data, idx_list, ensemble = False):
+  
+    cf = self.cf
+    batch_size = len(self.sources_info)  
+    num_levels = len(field_info[2])
+    num_tokens = field_info[3]
+    token_size = field_info[4]
+    data_b  =  self.split_data(data, idx_list, token_size)
+   
+    # recover token shape
+    if ensemble:
+      return [[data_b[vidx][bidx].reshape([-1, cf.net_tail_num_nets, *token_size])
+                                                            for bidx in range(batch_size)]
+                                                            for vidx in range(num_levels)]
+    else:
+      return [[data_b[vidx][bidx].reshape([-1, *token_size]) for bidx in range(batch_size)]
+                                                             for vidx in range(num_levels)]
+
+  ###################################################
+  def log_validate_BERT( self, epoch, batch_idx, log_sources, log_preds) :
+    '''Logging for BERT_strategy=BERT.'''
+
+    cf = self.cf
+    batch_size = len(self.sources_info) 
+
+    # save source: remains identical so just save ones
+    (sources, targets, tokens_masked_idx_list) = log_sources
+
+    sources_out, targets_out, preds_out, ensembles_out = [ ], [ ], [ ], [ ]
+    coords = []
+
+    for fidx, field_info in enumerate(cf.fields) : 
+
+      # reconstruct coordinates
+      is_predicted = fidx in self.fields_prediction_idx
+      num_levels = len(field_info[2])
+      num_tokens = field_info[3]
+      token_size = field_info[4]
+      sources_b = detokenize( sources[fidx].numpy())
+     
+      if is_predicted :
+        targets_b   = self.get_masked_data(field_info, targets[fidx], tokens_masked_idx_list[fidx])
+        preds_mu_b  = self.get_masked_data(field_info, log_preds[fidx][0], tokens_masked_idx_list[fidx])
+        preds_ens_b = self.get_masked_data(field_info, log_preds[fidx][2], tokens_masked_idx_list[fidx], ensemble = True)
+
+      # for all batch items
+      coords_b = []
+      for bidx in range(batch_size):
+        dates = self.sources_info[bidx][0]
+        lats  = self.sources_info[bidx][1]
+        lons  = self.sources_info[bidx][2]
+
+        lats_idx = self.sources_idxs[bidx][1]
+        lons_idx = self.sources_idxs[bidx][2]
+
+        # target etc are aliasing targets_b which simplifies bookkeeping below
+        if is_predicted :
+          target   = [targets_b[vidx][bidx] for vidx in range(num_levels)]
+          pred_mu  = [preds_mu_b[vidx][bidx] for vidx in range(num_levels)]
+          pred_ens = [preds_ens_b[vidx][bidx] for vidx in range(num_levels)]
+
+        coords_mskd_l = []
+        for vidx, _ in enumerate(field_info[2]) :
+
+          normalizer, year_base = self.model.normalizer( fidx, vidx, lats_idx, lons_idx)
+          sources_b[bidx,vidx] = denormalize(sources_b[bidx,vidx], normalizer, dates, year_base = 2021) 
+
+          if is_predicted :
+            idx = tokens_masked_idx_list[fidx][vidx][bidx]
+            grid = np.flip(np.array( np.meshgrid( lons, lats)), axis = 0) #flip to have lat on pos 0 and lon on pos 1
+            grid_idx = np.flip(np.array( np.meshgrid( lons_idx, lats_idx)), axis = 0) #flip to have lat on pos 0 and lon on pos 1
+       
+            # recover time dimension since idx assumes the full space-time cube
+            grid = torch.from_numpy( np.array( np.broadcast_to( grid,
+                                shape = [token_size[0]*num_tokens[0], *grid.shape])).swapaxes(0,1))
+            grid_lats_toked = tokenize( grid[0], token_size).flatten( 0, 2)
+            grid_lons_toked = tokenize( grid[1], token_size).flatten( 0, 2)
+          
+            idx_loc = idx - np.prod(num_tokens) * bidx
+            #save only useful info for each bidx. shape e.g. [n_bidx, lat_token_size*lat_num_tokens]
+            lats_mskd = np.array([np.unique(t) for t in grid_lats_toked[ idx_loc ].numpy()])
+            lons_mskd = np.array([np.unique(t) for t in grid_lons_toked[ idx_loc ].numpy()])
+
+            #time: idx ranges from 0->863 12x6x12 
+            t_idx = (idx_loc // (num_tokens[1]*num_tokens[2])) * token_size[0]
+            #create range from t_idx-2 to t_idx
+            t_idx = np.array([np.arange(t, t + token_size[0]) for t in t_idx])
+            dates_mskd = dates[t_idx]
+            
+            for ii,(t,p,e,da,la,lo) in enumerate(zip( target[vidx], pred_mu[vidx], pred_ens[vidx],
+                                                    dates_mskd, lats_mskd, lons_mskd)) :
+              normalizer_ii = normalizer
+              if len(normalizer.shape) > 2: #local normalization                                     
+                lats_mskd_idx = np.where(np.isin(lats,la))[0]
+                lons_mskd_idx = np.where(np.isin(lons,lo))[0]
+                #normalizer_ii = normalizer[:, :, lats_mskd_idx, lons_mskd_idx] problems in python 3.9
+                normalizer_ii = normalizer[:, :, lats_mskd_idx[0]:lats_mskd_idx[-1]+1, lons_mskd_idx[0]:lons_mskd_idx[-1]+1]
+             
+              targets_b[vidx][bidx][ii]   = denormalize(t, normalizer_ii, da, year_base)  
+              preds_mu_b[vidx][bidx][ii]  = denormalize(p, normalizer_ii, da, year_base) 
+              preds_ens_b[vidx][bidx][ii] = denormalize(e, normalizer_ii, da, year_base)
+            
+            coords_mskd_l += [[dates_mskd, 90.-lats_mskd, lons_mskd] ]
+       
+        coords_b += [ [dates, 90. - lats, lons] + coords_mskd_l ]
+      
+      coords += [ coords_b ]
+      fn = field_info[0]
+      sources_out.append( [fn, sources_b])
+
+      targets_out.append([fn, [[t.numpy(force=True) for t in t_v] for t_v in targets_b]] if is_predicted else [fn, []])
+      preds_out.append( [fn, [[p.numpy(force=True) for p in p_v] for p_v in preds_mu_b]] if is_predicted else [fn, []] )
+      ensembles_out.append( [fn, [[p.numpy(force=True) for p in p_v] for p_v in preds_ens_b]] if is_predicted else [fn, []] )
+
+    levels = [[np.array(l) for l in field[2]] for field in cf.fields]
+    write_BERT( cf.wandb_id, epoch, batch_idx, 
+                levels, sources_out, targets_out,
+                preds_out, ensembles_out, coords )
+
+######################################################
+
+  def log_attention( self, epoch, bidx, attention) : 
+    '''Hook for logging: output attention maps.'''
+    cf = self.cf
+
+    attn_out = []
+    for fidx, field_info in enumerate(cf.fields) : 
+     
+      # coordinates 
+      coords_b = []
+      for bidx in range(batch_size):
+        dates = self.sources_info[bidx][0]
+        lats  = 90. - self.sources_info[bidx][1]
+        lons  = self.sources_info[bidx][2]
+        coords_b += [ [dates, lats, lons] ]
+
+      is_predicted = fidx in self.fields_prediction_idx
+      attn_out.append([field_info[0], attention[fidx]] if is_predicted else [fn, []])
+      
+    levels = [[np.array(l) for l in field[2]] for field in cf.fields]
+    write_attention(cf.wandb_id, epoch,
+                    bidx, levels, attn_out,  coords_b )

vendor/atmorep-official/atmorep/datasets/data_writer.py

@@ -0,0 +1,181 @@
+####################################################################################################
+#
+#  Copyright (C) 2022
+#
+####################################################################################################
+#
+#  project     : atmorep
+#
+#  author      : atmorep collaboration
+# 
+#  description :
+#
+#  license     :
+#
+####################################################################################################
+
+import numpy as np
+import xarray as xr
+import zarr
+import atmorep.config.config as config
+
+def write_item(ds_field, name_idx, data, levels, coords, name  = 'sample' ):
+  ds_batch_item = ds_field.create_group( f'{name}={name_idx:05d}' )
+  ds_batch_item.create_dataset( 'data', data=data)
+  ds_batch_item.create_dataset( 'ml', data=levels)
+  ds_batch_item.create_dataset( 'datetime', data=coords[0].astype('datetime64[ns]'))
+  ds_batch_item.create_dataset( 'lat', data=np.array(coords[1]).astype(np.float32))
+  ds_batch_item.create_dataset( 'lon', data=np.array(coords[2]).astype(np.float32))
+  return ds_batch_item
+
+####################################################################################################
+def write_forecast( model_id, epoch, batch_idx, levels, sources, 
+                    targets, preds, ensembles, coords,  
+                    zarr_store_type = 'ZipStore' ) :
+  ''' 
+    sources : num_fields x [field name , data]
+    targets :
+    preds, ensemble share coords with targets
+  '''
+  sources_coords = [[c[:3] for c in coord_field ] for coord_field in coords]
+  targets_coords = [[[c[-1], c[1], c[2]] for c in coord_field ] for coord_field in coords]
+  fname =  f'{config.path_results}/id{model_id}/results_id{model_id}_epoch{epoch:05d}' + '_{}.zarr'
+
+  zarr_store = getattr( zarr, zarr_store_type)
+
+  store_source = zarr_store( fname.format( 'source'))
+  exp_source = zarr.group(store=store_source)
+ 
+  for fidx, field in enumerate(sources) :
+    ds_field = exp_source.require_group( f'{field[0]}')
+    batch_size = field[1].shape[0]
+    for bidx in range( field[1].shape[0]) :
+      sample = batch_idx * batch_size + bidx
+      write_item(ds_field, sample, field[1][bidx], levels, sources_coords[fidx][bidx])
+  store_source.close()
+
+  store_target = zarr_store( fname.format( 'target'))
+  exp_target = zarr.group(store=store_target)
+  for fidx, field in enumerate(targets) :
+    ds_field = exp_target.require_group( f'{field[0]}')
+    batch_size = field[1].shape[0]
+    for bidx in range( field[1].shape[0]) :
+      sample = batch_idx * batch_size + bidx
+      write_item(ds_field, sample, field[1][bidx], levels, targets_coords[fidx][bidx])
+  store_target.close()
+
+  store_pred = zarr_store( fname.format( 'pred'))
+  exp_pred = zarr.group(store=store_pred)
+  for fidx, field in enumerate(preds) :
+    ds_field = exp_pred.require_group( f'{field[0]}')
+    batch_size = field[1].shape[0]
+    for bidx in range( field[1].shape[0]) :
+      sample = batch_idx * batch_size + bidx
+      write_item(ds_field, sample, field[1][bidx], levels, targets_coords[fidx][bidx]) 
+  store_pred.close()
+
+  store_ens = zarr_store( fname.format( 'ens'))
+  exp_ens = zarr.group(store=store_ens)
+  for fidx, field in enumerate(ensembles) :
+    ds_field = exp_ens.require_group( f'{field[0]}')
+    batch_size = field[1].shape[0]
+    for bidx in range( field[1].shape[0]) :
+      sample = batch_idx * batch_size + bidx
+      write_item(ds_field, sample, field[1][bidx], levels, targets_coords[fidx][bidx])
+  store_ens.close()
+
+####################################################################################################
+def write_BERT( model_id, epoch, batch_idx, levels, sources, 
+                targets, preds, ensembles, coords, 
+                zarr_store_type = 'ZipStore' ) : 
+                                            
+  '''
+    sources : num_fields x [field name , data]
+    targets :
+    preds, ensemble share coords with targets
+  '''
+
+  sources_coords = [[c[:3] for c in coord_field ] for coord_field in coords]
+  targets_coords = [[c[3:] for c in coord_field ] for coord_field in coords]
+  
+  fname =  f'{config.path_results}/id{model_id}/results_id{model_id}_epoch{epoch:05d}' + '_{}.zarr'
+
+  zarr_store = getattr( zarr, zarr_store_type)
+
+  store_source = zarr_store( fname.format( 'source'))
+  exp_source = zarr.group(store=store_source)
+  for fidx, field in enumerate(sources) :
+    ds_field = exp_source.require_group( f'{field[0]}')
+    batch_size = field[1].shape[0]
+    for bidx in range( field[1].shape[0]) :
+      sample = batch_idx * batch_size + bidx
+      write_item(ds_field, sample, field[1][bidx], levels[fidx], sources_coords[fidx][bidx] )
+  store_source.close()
+
+  store_target = zarr_store( fname.format( 'target'))
+  exp_target = zarr.group(store=store_target)
+  for fidx, field in enumerate(targets) :
+    if 0 == len(field[1]) :  # skip fields that were not predicted
+        continue
+    batch_size = len(field[1][0])
+    ds_field = exp_target.require_group( f'{field[0]}')
+    for bidx in range( len(field[1][0])) :
+      sample = batch_idx * batch_size + bidx
+      ds_target_b = ds_field.create_group( f'sample={sample:05d}')
+      for vidx in range(len(levels[fidx])) :
+        write_item(ds_target_b, levels[fidx][vidx], field[1][vidx][bidx], levels[fidx][vidx], targets_coords[fidx][bidx][vidx], name = 'ml' )
+  store_target.close()
+
+  store_pred = zarr_store( fname.format( 'pred'))
+  exp_pred = zarr.group(store=store_pred)
+  for fidx, field in enumerate(preds) :
+    if 0 == len(field[1]) :  # skip fields that were not predicted
+      continue
+    batch_size = len(field[1][0])
+    ds_pred = exp_pred.require_group( f'{field[0]}')
+    for bidx in range( len(field[1][0])) :
+      sample = batch_idx * batch_size + bidx
+      ds_pred_b = ds_pred.create_group( f'sample={sample:05d}')
+      for vidx in range(len(levels[fidx])) :
+        write_item(ds_pred_b, levels[fidx][vidx], field[1][vidx][bidx], levels[fidx][vidx], 
+                                  targets_coords[fidx][bidx][vidx], name = 'ml' )
+  store_pred.close()
+
+  store_ens = zarr_store( fname.format( 'ens'))
+  exp_ens = zarr.group(store=store_ens)
+  for fidx, field in enumerate(ensembles) :
+    if 0 == len(field[1]) :  # skip fields that were not predicted
+      continue
+    batch_size = len(field[1][0])
+    ds_ens = exp_ens.require_group( f'{field[0]}')
+    for bidx in range( len(field[1][0])) :
+      sample = batch_idx * batch_size + bidx
+      ds_ens_b = ds_ens.create_group( f'sample={sample:05d}')
+      for vidx in range(len(levels[fidx])) :
+        write_item(ds_ens_b, levels[fidx][vidx], field[1][vidx][bidx], levels[fidx][vidx], 
+                                targets_coords[fidx][bidx][vidx], name = 'ml' )
+  store_ens.close()
+
+####################################################################################################
+def write_attention(model_id, epoch, batch_idx, levels, attn, coords, zarr_store_type = 'ZipStore' ) :
+
+  fname =  f'{config.path_results}/id{model_id}/results_id{model_id}_epoch{epoch:05d}' + '_{}.zarr'
+  zarr_store = getattr( zarr, zarr_store_type)
+
+  store_attn = zarr_store( fname.format( 'attention'))
+  exp_attn = zarr.group(store=store_attn)
+
+  for fidx, atts_f in enumerate(attn) :
+    ds_field = exp_attn.require_group( f'{atts_f[0]}')
+    ds_field_b = ds_field.require_group( f'batch={batch_idx:05d}')
+    for lidx, atts_f_l in enumerate(atts_f[1]) :  # layer in the network
+      ds_f_l = ds_field_b.require_group( f'layer={lidx:05d}')
+      ds_f_l.create_dataset( 'ml', data=levels[fidx])
+      ds_f_l.create_dataset( 'datetime', data=coords[0][fidx])
+      ds_f_l.create_dataset( 'lat', data=coords[1][fidx])
+      ds_f_l.create_dataset( 'lon', data=coords[2][fidx])
+      ds_f_l_h = ds_f_l.require_group('heads')
+      for hidx, atts_f_l_head in enumerate(atts_f_l) :  # number of attention head
+        if atts_f_l_head != None :
+          ds_f_l_h.create_dataset(f'{hidx}', data=atts_f_l_head.numpy() )
+  store_attn.close()

vendor/atmorep-official/atmorep/datasets/multifield_data_sampler.py

@@ -0,0 +1,371 @@
+####################################################################################################
+#
+#  Copyright (C) 2022
+#
+####################################################################################################
+#
+#  project     : atmorep
+#
+#  author      : atmorep collaboration
+# 
+#  description :
+#
+#  license     :
+#
+####################################################################################################
+
+import dask.config as dc
+import dask.array as da
+import torch
+import numpy as np
+import zarr
+import pandas as pd
+from datetime import datetime
+import time
+import os
+
+from atmorep.datasets.normalizer import normalize
+from atmorep.utils.utils import tokenize, get_weights
+
+class MultifieldDataSampler( torch.utils.data.IterableDataset):
+    
+  ###################################################
+  def __init__( self, file_path, fields, years, batch_size, pre_batch, n_size,
+                num_samples, with_shuffle = False, time_sampling = 1, with_source_idxs = False, compute_weights = False, 
+                fields_targets = None, pre_batch_targets = None ) :
+    '''
+      Data set for single dynamic field at an arbitrary number of vertical levels
+
+      nsize : neighborhood in (tsteps, deg_lat, deg_lon)
+    '''
+    super( MultifieldDataSampler).__init__()
+
+    self.fields = fields
+    self.batch_size = batch_size
+    self.n_size = n_size
+    self.num_samples = num_samples
+    self.with_source_idxs = with_source_idxs
+    self.compute_weights = compute_weights
+    self.with_shuffle = with_shuffle
+    self.pre_batch = pre_batch
+    
+    assert os.path.exists(file_path), f"File path {file_path} does not exist"
+    self.ds = zarr.open( file_path)
+   
+    self.dask_array_data = da.from_zarr(self.ds['data'])
+    self.dask_array_sfc  = da.from_zarr(self.ds['data_sfc'])
+
+    self.ds_global = self.ds.attrs['is_global']
+
+    self.lats = np.array( self.ds['lats'])
+    self.lons = np.array( self.ds['lons'])
+    
+    sh = self.ds['data'].shape
+    st = self.ds['time'].shape
+    self.ds_len = st[0] 
+    print( f'self.ds[\'data\'] : {sh} :: {st}')
+    print( f'self.lats : {self.lats.shape}', flush=True)
+    print( f'self.lons : {self.lons.shape}', flush=True)
+    self.fields_idxs = []
+
+    self.time_sampling = time_sampling
+    self.range_lat = np.array( self.lats[ [0,-1] ])
+    self.range_lon = np.array( self.lons[ [0,-1] ])
+    self.res = np.array(self.ds.attrs['res'])
+    self.year_base = self.ds['time'][0].astype(datetime).year
+
+    # ensure neighborhood does not exceed domain (either at pole or for finite domains)
+    self.range_lat += np.array([n_size[1] / 2., -n_size[1] / 2.])
+    # lon: no change for periodic case
+    if self.ds_global < 1.:
+      self.range_lon += np.array([n_size[2]/2., -n_size[2]/2.])
+    
+    # data normalizers
+    self.normalizers = []
+    for ifield, field_info in enumerate(fields) :
+      corr_type = 'global' if len(field_info) <= 6 else field_info[6]
+      nf_name = 'global_norm' if corr_type == 'global' else 'norm'
+      self.normalizers.append( [] )
+      for vl in field_info[2]: 
+        if vl == 0:
+          field_idx = self.ds.attrs['fields_sfc'].index( field_info[0])
+          n_name = f'normalization/{nf_name}_sfc'
+          self.normalizers[ifield] += [self.ds[n_name].oindex[ :, :, field_idx]] 
+        else:
+          vl_idx = self.ds.attrs['levels'].index(vl)
+          field_idx = self.ds.attrs['fields'].index( field_info[0])
+          n_name = f'normalization/{nf_name}'
+          self.normalizers[ifield] += [self.ds[n_name].oindex[ :, :, field_idx, vl_idx]] 
+    
+    # extract indices for selected years
+    self.times = pd.DatetimeIndex( self.ds['time'])
+    idxs_years = self.times.year == years[0]
+    for year in years[1:] :
+      idxs_years = np.logical_or( idxs_years, self.times.year == year)
+    self.idxs_years = np.where( idxs_years)[0]
+
+    self.num_samples = min( self.num_samples, self.idxs_years.shape[0])
+
+  ###################################################
+  def shuffle( self) :
+
+    worker_info = torch.utils.data.get_worker_info()
+    rng_seed = None
+    if worker_info is not None :
+      rng_seed = int(time.time()) // (worker_info.id+1) + worker_info.id
+
+    rng = np.random.default_rng( rng_seed)
+    self.idxs_perm_t = rng.permutation( self.idxs_years)[ : self.num_samples // self.batch_size]
+    
+    lats = rng.random(self.num_samples) * (self.range_lat[1] - self.range_lat[0]) +self.range_lat[0]
+    lons = rng.random(self.num_samples) * (self.range_lon[1] - self.range_lon[0]) +self.range_lon[0]
+
+    # align with grid
+    res_inv = 1.0 / self.res * 1.00001
+    lats = self.res[0] * np.round( lats * res_inv[0])
+    lons = self.res[1] * np.round( lons * res_inv[1])
+
+    self.idxs_perm = np.stack( [lats, lons], axis=1)
+
+  ###################################################
+  def __iter__(self):
+
+    if self.with_shuffle :
+      self.shuffle()
+
+    lats, lons = self.lats, self.lons
+    ts, n_size = self.time_sampling, self.n_size
+    ns_2 = np.array(self.n_size) / 2.
+    res = self.res
+
+    iter_start, iter_end = self.worker_workset()
+
+    for bidx in range( iter_start, iter_end) :
+
+      sources, token_infos = [[] for _ in self.fields], [[] for _ in self.fields]
+      sources_infos, source_idxs = [], []
+  
+      i_bidx = self.idxs_perm_t[bidx]
+      idxs_t = list(np.arange( i_bidx - n_size[0]*ts, i_bidx, ts, dtype=np.int64))
+      # data_tt_sfc = self.ds['data_sfc'].oindex[idxs_t]
+      # data_tt = self.ds['data'].oindex[idxs_t]
+      with dc.set(**{'array.slicing.split_large_chunks': True}):
+        data_tt_sfc = self.dask_array_sfc[idxs_t].compute()
+        data_tt     = self.dask_array_data[idxs_t].compute()
+
+
+      for sidx in range(self.batch_size) :
+        
+        idx = self.idxs_perm[bidx*self.batch_size+sidx]
+        # slight asymetry with offset by res/2 is required to match desired token count
+        lat_ran = np.where(np.logical_and(lats>idx[0]-ns_2[1]-res[0]/2.,lats<idx[0]+ns_2[1]))[0]
+        # handle periodicity of lon
+        assert not ((idx[1]-ns_2[2]) < 0. and (idx[1]+ns_2[2]) > 360.)
+        il, ir = (idx[1]-ns_2[2]-res[1]/2., idx[1]+ns_2[2])
+        if il < 0. :
+          lon_ran = np.concatenate( [np.where( lons > il+360)[0], np.where(lons < ir)[0]], 0)
+        elif ir > 360. :
+          lon_ran = np.concatenate( [np.where( lons > il)[0], np.where(lons < ir-360)[0]], 0)
+        else : 
+          lon_ran = np.where(np.logical_and( lons > il, lons < ir))[0]
+        
+        sources_infos += [ [ self.ds['time'][ idxs_t ].astype(datetime), 
+                             self.lats[lat_ran], self.lons[lon_ran], self.res ] ]
+
+        if self.with_source_idxs :
+          source_idxs += [ (idxs_t, lat_ran, lon_ran) ]
+
+        # extract data
+        for ifield, field_info in enumerate(self.fields):  
+          source_lvl, tok_info_lvl  = [], []
+          tok_size  = field_info[4]
+          num_tokens = field_info[3]
+          corr_type = 'global' if len(field_info) <= 6 else field_info[6]
+        
+          for ilevel, vl in enumerate(field_info[2]):
+            if vl == 0 : #surface level
+              field_idx = self.ds.attrs['fields_sfc'].index( field_info[0])
+              data_t = data_tt_sfc[ :, field_idx ]
+            else :
+              field_idx = self.ds.attrs['fields'].index( field_info[0])
+              vl_idx = self.ds.attrs['levels'].index(vl)
+              data_t = data_tt[ :, field_idx, vl_idx ]
+          
+            source_data, tok_info = [], []
+            # extract data, normalize and tokenize
+            cdata = data_t[ ... , lat_ran[:,np.newaxis], lon_ran[np.newaxis,:]]
+            
+            normalizer = self.normalizers[ifield][ilevel]
+
+            if corr_type != 'global': 
+              #normalizer = normalizer[ ... , lat_ran[:,np.newaxis], lon_ran[np.newaxis,:]]
+              if lat_ran[0] < lat_ran[-1] and lon_ran[0] < lon_ran[-1]:
+                lat_max, lat_min = max(lat_ran), min(lat_ran)
+                lon_max, lon_min = max(lon_ran), min(lon_ran)
+                normalizer = normalizer[:,:,lat_min:lat_max+1,lon_min:lon_max+1]
+                #normalizer_vu = normalizer[:,:,lat_min:lat_max+1,lon_min:lon_max+1]
+                #cdata = normalize(cdata, normalizer_vu, sources_infos[-1][0], year_base = self.year_base) 
+              else:
+                normalizer = normalizer[ ... , lat_ran[:,np.newaxis], lon_ran[np.newaxis,:]]
+                #cdata = normalize(cdata, normalizer, sources_infos[-1][0], year_base = self.year_base)
+            #else:
+            cdata = normalize(cdata, normalizer, sources_infos[-1][0], year_base = self.year_base)
+            
+            source_data = tokenize( torch.from_numpy( cdata), tok_size )    
+            # token_infos uses center of the token: *last* datetime and center in space
+            dates = self.ds['time'][ idxs_t ].astype(datetime)
+            cdates = dates[tok_size[0]-1::tok_size[0]]
+            # use -1 is to start days from 0
+            dates = [(d.year, d.timetuple().tm_yday-1, d.hour) for d in cdates] 
+            lats_sidx = self.lats[lat_ran][ tok_size[1]//2 :: tok_size[1] ]
+            lons_sidx = self.lons[lon_ran][ tok_size[2]//2 :: tok_size[2] ]
+            # tensor product for token_infos
+            tok_info += [[[[[ year, day, hour, vl, lat, lon, vl, self.res[0]] for lon in lons_sidx]
+                                                                              for lat in lats_sidx]
+                                                                  for (year, day, hour) in dates]]
+
+            source_lvl += [ source_data ]
+            tok_info_lvl += [ torch.tensor(tok_info, dtype=torch.float32).flatten( 1, -2)]      
+          sources[ifield] += [ torch.stack(source_lvl, 0) ]
+          token_infos[ifield] += [ torch.stack(tok_info_lvl, 0) ]
+      
+      # concatenate batches   
+      sources = [torch.stack(sources_field).transpose(1,0) for sources_field in sources]
+      token_infos = [torch.stack(tis_field).transpose(1,0) for tis_field in token_infos]
+      sources = self.pre_batch( sources, token_infos )
+
+      tmidx_list = sources[-1]
+      weights_idx_list = []
+      if self.compute_weights:
+        for ifield, field_info in enumerate(self.fields): 
+          weights = []
+          for ilevel, vl in enumerate(field_info[2]):
+            for ibatch in range(self.batch_size):
+              
+              lats_idx = source_idxs[ibatch][1] 
+              lons_idx = source_idxs[ibatch][2]
+
+              idx_base = tmidx_list[ifield][ilevel][ibatch]
+              idx_loc = idx_base - np.prod(num_tokens) * ibatch
+              
+              grid = np.flip(np.array( np.meshgrid( lons_idx, lats_idx)), axis = 0) #flip to have lat on pos 0 and lon on pos 1
+              grid = torch.from_numpy( np.array( np.broadcast_to( grid,
+                                      shape = [tok_size[0]*num_tokens[0], *grid.shape])).swapaxes(0,1))
+
+              grid_lats_toked = tokenize( grid[0], tok_size).flatten( 0, 2)  
+
+              lats_mskd_b = np.array([np.unique(t) for t in grid_lats_toked[ idx_loc ].numpy()])
+
+              weights.append([get_weights(la) for la in lats_mskd_b])
+
+          weights_idx_list.append(weights)
+      sources = (*sources, weights_idx_list)
+
+      # TODO: implement (only required when prediction target comes from different data stream)
+      targets, target_info = None, None
+      target_idxs = None
+     
+      yield ( sources, targets, (source_idxs, sources_infos), (target_idxs, target_info))
+
+  ###################################################
+  def set_data( self, times_pos, batch_size = None) :
+    '''
+      times_pos = np.array( [ [year, month, day, hour, lat, lon], ...]  )
+        - lat \in [90,-90] = [90N, 90S]
+        - lon \in [0,360]
+        - (year,month) pairs should be a limited number since all data for these is loaded
+    '''
+    # generate all the data
+    self.idxs_perm = np.zeros( (len(times_pos), 2))
+    self.idxs_perm_t = []
+    self.num_samples = len(times_pos)
+    for idx, item in enumerate( times_pos) :
+
+      assert item[2] >= 1 and item[2] <= 31
+      assert item[3] >= 0 and item[3] < int(24 / self.time_sampling)
+      assert item[4] >= -90. and item[4] <= 90.
+
+      tstamp = pd.to_datetime( f'{item[0]}-{item[1]}-{item[2]}-{item[3]}', format='%Y-%m-%d-%H')
+      
+      self.idxs_perm_t += [ np.where( self.times == tstamp)[0]+1 ] #The +1 assures that tsamp is included in the range
+
+      # work with mathematical lat coordinates from here on
+      self.idxs_perm[idx] = np.array( [90. - item[4], item[5]])
+   
+    self.idxs_perm_t = np.array(self.idxs_perm_t).squeeze()
+
+  ###################################################
+  def set_global( self, times, batch_size = None, token_overlap = [0, 0]) :
+    ''' generate patch/token positions for global grid '''
+    token_overlap = np.array( token_overlap).astype(np.int64)
+
+    # assumed that sanity checking that field data is consistent has been done 
+    ifield = 0
+    field = self.fields[ifield]
+
+    res = self.res
+    side_len = np.array( [field[3][1] * field[4][1]*res[0], field[3][2] * field[4][2]*res[1]] )
+    overlap = np.array([token_overlap[0]*field[4][1]*res[0],token_overlap[1]*field[4][2]*res[1]])
+    side_len_2 = side_len / 2.
+    assert all( overlap <= side_len_2), 'token_overlap too large for #tokens, reduce if possible'
+
+    # generate tiles
+    times_pos = []
+    for ctime in times :
+
+      lat = side_len_2[0].item()
+      num_tiles_lat = 0
+      while (lat + side_len_2[0].item()) < 180. :
+        num_tiles_lat += 1
+        lon = side_len_2[1].item() - overlap[1].item()/2.
+        num_tiles_lon = 0
+        while (lon - side_len_2[1]) < 360. :
+          times_pos += [[*ctime, -lat + 90., np.mod(lon,360.) ]]
+          lon += side_len[1].item() - overlap[1].item()
+          num_tiles_lon += 1
+        lat += side_len[0].item() - overlap[0].item()
+
+      # add one additional row if no perfect tiling (sphere is toric in longitude so no special
+      # handling necessary but not in latitude)
+      # the added row is such that it goes exaclty down to the South pole and the offset North-wards
+      # is computed based on this
+      lat -= side_len[0] - overlap[0]
+      if lat - side_len_2[0] < 180. :
+        num_tiles_lat += 1
+        lat = 180. - side_len_2[0].item() + res[0]
+        lon = side_len_2[1].item() - overlap[1].item()/2.
+        while (lon - side_len_2[1]) < 360. :
+          times_pos += [[*ctime, -lat + 90., np.mod(lon,360.) ]]
+          lon += side_len[1].item() - overlap[1].item()
+
+    # adjust batch size if necessary so that the evaluations split up across batches of equal size
+    batch_size = len(times_pos) #num_tiles_lon
+    
+    print( 'Number of batches per global forecast: {}'.format( num_tiles_lat) )
+
+    self.set_data( times_pos, batch_size)
+
+  ###################################################
+  def __len__(self):
+      return self.num_samples // self.batch_size
+
+  ###################################################
+  def worker_workset( self) :
+
+    worker_info = torch.utils.data.get_worker_info()
+
+    if worker_info is None: 
+      iter_start = 0
+      iter_end = self.num_samples
+   
+    else:  
+      # split workload
+      per_worker = len(self) // worker_info.num_workers
+      worker_id = worker_info.id
+      iter_start = int(worker_id * per_worker)
+      iter_end = int(iter_start + per_worker)
+      if worker_info.id+1 == worker_info.num_workers :
+        iter_end = len(self)
+
+    return iter_start, iter_end
+

vendor/atmorep-official/atmorep/datasets/normalizer.py

@@ -0,0 +1,82 @@
+####################################################################################################
+#
+#  Copyright (C) 2022
+#
+####################################################################################################
+#
+#  project     : atmorep
+#
+#  author      : atmorep collaboration
+# 
+#  description :
+#
+#  license     :
+#
+####################################################################################################
+
+import code
+import numpy as np
+import xarray as xr
+import atmorep.config.config as config
+
+######################################################
+#                   Normalize                        #
+######################################################
+
+def normalize( data, norm, dates, year_base = 1979) :
+  corr_data = np.array([norm[12*(dt.year-year_base) + dt.month-1] for dt in dates])
+  mean, var = corr_data[:, 0], corr_data[:, 1]
+  if (var == 0.).all() :
+    print( f'Warning: var == 0') 
+    assert False
+  if len(norm.shape) > 2 : #global norm
+    return normalize_local(data, mean, var)
+  else:
+    return normalize_global( data, mean, var)
+  
+######################################################
+def normalize_local( data, mean, var) :
+  data = (data - mean) / var
+  return data
+
+######################################################
+def normalize_global( data, mean, var) :
+  for i in range( data.shape[0]) :
+    data[i] = (data[i] - mean[i]) / var[i]
+  return data
+
+
+######################################################
+#                  Denormalize                       #
+######################################################
+def denormalize(data, norm, dates, year_base = 1979) :
+  corr_data = np.array([norm[12*(dt.year-year_base) + dt.month-1] for dt in dates])
+  mean, var = corr_data[:, 0], corr_data[:, 1]
+  if len(norm.shape) > 2 :
+    return denormalize_local(data, mean, var)
+  else:
+    return denormalize_global(data, mean, var)  
+
+######################################################
+
+def denormalize_local(data, mean, var) :
+  if len(data.shape) > 3: #ensemble
+    for i in range( data.shape[0]) :
+      data[i] = (data[i] * var) + mean
+  else:
+      data = (data * var) + mean
+  return data
+
+######################################################
+
+def denormalize_global(data, mean, var) :
+  if len(data.shape) > 3: #ensemble
+    data = data.swapaxes(0,1)
+    for i in range( data.shape[0]) :
+      data[i] = ((data[i] * var[i]) + mean[i])
+    data = data.swapaxes(0,1)
+  else:
+    for i in range( data.shape[0]) :
+      data[i] = (data[i] * var[i]) + mean[i]
+    
+  return data

vendor/atmorep-official/atmorep/tests/__init__.py

@@ -0,0 +1 @@
+ 

vendor/atmorep-official/atmorep/tests/conftest.py

@@ -0,0 +1,8 @@
+def pytest_addoption(parser):
+    parser.addoption("--field", action="store", help="field to run the test on")
+    parser.addoption("--model_id", action="store", help="wandb ID of the atmorep model")
+    parser.addoption("--epoch", action="store", help="field to run the test on", default = "0")
+    parser.addoption("--strategy", action="store", help="BERT or forecast")
+
+
+

vendor/atmorep-official/atmorep/tests/test_utils.py

@@ -0,0 +1,77 @@
+import numpy as np
+import pandas as pd 
+
+def era5_fname():
+    return "/gpfs/scratch/ehpc03/data/{}/ml{}/era5_{}_y{}_m{}_ml{}.grib"
+
+def atmorep_pred():
+    return "./results/id{}/results_id{}_epoch{}_pred.zarr"
+
+def atmorep_target():
+    return "./results/id{}/results_id{}_epoch{}_target.zarr"
+
+def grib_index(field):
+    grib_idxs = {"velocity_u": "u",
+                 "temperature": "t", 
+                 "total_precip": "tp", 
+                 "velocity_v": "v", 
+                 "velocity_z": "z", 
+                 "vorticity" : "vo", 
+                 "divergence" : "d", 
+                 "specific_humidity": "q"}
+
+    return grib_idxs[field]
+
+##################################################################
+
+def get_BERT(atmorep, field, sample, level):
+    atmorep_sample = atmorep[f"{field}/sample={sample:05d}/ml={level:05d}"] 
+    data = atmorep_sample.data[0,0] 
+    datetime = pd.Timestamp(atmorep_sample.datetime[0,0])
+    lats = atmorep_sample.lat[0]
+    lons = atmorep_sample.lon[0]
+    return data, datetime, lats, lons
+
+def get_forecast(atmorep, field, sample,level_idx):
+    atmorep_sample = atmorep[f"{field}/sample={sample:05d}"]
+    data = atmorep_sample.data[level_idx, 0]
+    datetime = pd.Timestamp(atmorep_sample.datetime[0])
+    lats = atmorep_sample.lat
+    lons = atmorep_sample.lon
+    return data, datetime, lats, lons
+
+######################################
+
+def check_lats(lats_pred, lats_target):
+    assert (lats_pred[:] == lats_target[:]).all(), "Mismatch between latitudes"
+    assert (lats_pred[:] <= 90.).all(), f"latitudes are between {np.amin(lats_pred)}- {np.amax(lats_pred)}"
+    assert (lats_pred[:] >= -90.).all(), f"latitudes are between {np.amin(lats_pred)}- {np.amax(lats_pred)}"
+
+def check_lons(lons_pred, lons_target):
+    assert (lons_pred[:] == lons_target[:]).all(), "Mismatch between longitudes"
+    assert (lons_pred[:] >= 0.).all(), "longitudes are between {np.amin(lons_pred)}- {np.amax(lons_pred)}"
+    assert (lons_pred[:] <= 360.).all(), "longitudes are between {np.amin(lons_pred)}- {np.amax(lons_pred)}"
+
+def check_datetimes(datetimes_pred, datetimes_target):
+    assert (datetimes_pred == datetimes_target), "Mismatch between datetimes"
+
+######################################
+
+#calculate RMSE
+def compute_RMSE(pred, target):
+    return np.sqrt(np.mean((pred-target)**2))
+
+
+def get_max_RMSE(field):
+    #TODO: optimize thresholds
+    values = {"temperature" : 3, 
+              "velocity_u" : 0.2, #????
+              "velocity_v": 0.2,  #????
+              "velocity_z": 0.2,  #????
+              "vorticity" : 0.2,    #????
+              "divergence": 0.2,    #????
+              "specific_humidity": 0.2,  #????
+              "total_precip": 1, #?????
+            }
+    
+    return values[field]

vendor/atmorep-official/atmorep/tests/validation_test.py

@@ -0,0 +1,126 @@
+import pytest
+import zarr
+import cfgrib 
+import xarray as xr
+import numpy as np 
+import random as rnd
+import warnings
+import os
+
+from atmorep.tests.test_utils import *
+
+# run it with e.g. pytest -s atmorep/tests/validation_test.py --field temperature --model_id ztsut0mr --strategy BERT
+
+@pytest.fixture
+def field(request):
+    return request.config.getoption("field")
+
+@pytest.fixture
+def model_id(request):
+    return request.config.getoption("model_id")
+
+@pytest.fixture
+def epoch(request):
+    request.config.getoption("epoch")
+
+@pytest.fixture(autouse = True) 
+def BERT(request):
+    strategy = request.config.getoption("strategy")
+    return (strategy == 'BERT' or strategy == 'temporal_interpolation')
+
+@pytest.fixture(autouse = True) 
+def strategy(request):
+    return request.config.getoption("strategy")
+
+#TODO: add test for global_forecast vs ERA5
+
+def test_datetime(field, model_id, BERT, epoch = 0):
+
+    """
+    Check against ERA5 timestamps.
+    Loop over all levels individually. 50 random samples for each level.
+    """
+
+    store = zarr.ZipStore(atmorep_target().format(model_id, model_id, str(epoch).zfill(5)))
+    atmorep = zarr.group(store)
+
+    nsamples = min(len(atmorep[field]), 50)
+    samples = rnd.sample(range(len(atmorep[field])), nsamples)
+    levels = [int(f.split("=")[1]) for f in atmorep[f"{field}/sample=00000"]] if BERT else atmorep[f"{field}/sample=00000"].ml[:]
+   
+    get_data = get_BERT if BERT else get_forecast
+   
+    for level in levels:
+        #TODO: make it more elegant
+        level_idx = level if BERT else np.where(levels == level)[0].tolist()[0]
+
+        for s in samples:
+            data, datetime, lats, lons = get_data(atmorep, field, s, level_idx)
+            year, month = datetime.year, str(datetime.month).zfill(2)
+
+            era5_path = era5_fname().format(field, level, field, year, month, level)
+            if not os.path.isfile(era5_path):
+                warnings.warn(UserWarning((f"Timestamp {datetime} not found in ERA5. Skipping")))
+                continue
+            era5 = xr.open_dataset(era5_path, engine = "cfgrib")[grib_index(field)].sel(time = datetime, latitude = lats, longitude = lons)
+
+            #assert (data[0] == era5.values[0]).all(), "Mismatch between ERA5 and AtmoRep Timestamps"
+            assert np.isclose(data[0], era5.values[0],rtol=1e-04, atol=1e-07).all(), "Mismatch between ERA5 and AtmoRep Timestamps"
+
+#############################################################################
+
+def test_coordinates(field, model_id, BERT, epoch = 0):
+    """
+    Check that coordinates match between target and prediction. 
+    Check also that latitude and longitudes are in geographical coordinates
+    50 random samples.
+    """
+   
+    store_t = zarr.ZipStore(atmorep_target().format(model_id, model_id, str(epoch).zfill(5)))
+    target = zarr.group(store_t)
+
+    store_p = zarr.ZipStore(atmorep_pred().format(model_id, model_id, str(epoch).zfill(5)))
+    pred = zarr.group(store_p)
+
+    nsamples = min(len(target[field]), 50)
+    samples = rnd.sample(range(len(target[field])), nsamples)
+    levels = [int(f.split("=")[1]) for f in target[f"{field}/sample=00000"]] if BERT else target[f"{field}/sample=00000"].ml[:]
+   
+    get_data = get_BERT if BERT else get_forecast
+
+    for level in levels:
+        level_idx = level if BERT else np.where(levels == level)[0].tolist()[0]
+        for s in samples:
+            _, datetime_target, lats_target, lons_target = get_data(target,field, s, level_idx)
+            _, datetime_pred, lats_pred, lons_pred = get_data(pred, field, s, level_idx)
+
+            check_lats(lats_pred, lats_target)
+            check_lons(lons_pred, lons_target)
+            check_datetimes(datetime_pred, datetime_target)
+
+#########################################################################
+
+def test_rmse(field, model_id, BERT, epoch = 0):
+    """
+    Test that for each field the RMSE does not exceed a certain value. 
+    50 random samples.
+    """
+    store_t = zarr.ZipStore(atmorep_target().format(model_id, model_id, str(epoch).zfill(5)))
+    target = zarr.group(store_t)
+
+    store_p = zarr.ZipStore(atmorep_pred().format(model_id, model_id, str(epoch).zfill(5)))
+    pred = zarr.group(store_p)
+    
+    nsamples = min(len(target[field]), 50)
+    samples = rnd.sample(range(len(target[field])), nsamples)
+    levels = [int(f.split("=")[1]) for f in target[f"{field}/sample=00000"]] if BERT else target[f"{field}/sample=00000"].ml[:]
+    
+    get_data = get_BERT if BERT else get_forecast
+    
+    for level in levels:
+        level_idx = level if BERT else np.where(levels == level)[0].tolist()[0]
+        for s in samples:
+            sample_target, _, _, _ = get_data(target,field, s, level_idx)
+            sample_pred, _, _, _ = get_data(pred,field, s, level_idx)
+
+            assert compute_RMSE(sample_target, sample_pred).mean() < get_max_RMSE(field)

vendor/atmorep-official/atmorep/training/bert.py

@@ -0,0 +1,228 @@
+####################################################################################################
+#
+#  Copyright (C) 2022
+#
+####################################################################################################
+#
+#  project     : atmorep
+#
+#  author      : atmorep collaboration
+# 
+#  description :
+#
+#  license     :
+#
+####################################################################################################
+
+import torch
+import numpy as np
+from functools import partial
+import code
+
+####################################################################################################
+def prepare_batch_BERT_multifield( cf, rngs, fields, BERT_strategy, fields_data, fields_infos) :
+  
+  fields_tokens_masked_idx_list = [[] for _ in fields_data]
+  fields_targets = [[] for _ in fields_data]
+  sources = [[] for _ in fields_data]
+  token_infos = [[] for _ in fields_data]
+
+  if not BERT_strategy :
+    BERT_strategy = cf.BERT_strategy
+
+  if BERT_strategy == 'BERT' :
+    bert_f = prepare_batch_BERT_field
+  elif BERT_strategy == 'global_forecast' :
+    bert_f = prepare_batch_BERT_forecast_field
+  elif BERT_strategy == 'forecast' :
+    bert_f = prepare_batch_BERT_forecast_field
+  elif BERT_strategy == 'temporal_interpolation' :
+    bert_f = prepare_batch_BERT_temporal_field
+  else :
+    assert False
+
+  rng_idx = 1
+  for ifield, (field, infos) in enumerate(zip(fields_data, fields_infos)) :
+    for ilevel, (field_data, token_info) in enumerate(zip(field, infos)) :
+     
+      # no masking for static fields or if masking rate = 0
+      if fields[ifield][1][0] > 0 and fields[ifield][5][0] > 0. :
+
+        ret = bert_f( cf, ifield, field_data, token_info, rngs[rng_idx])
+        (field_data, token_info, target, tokens_masked_idx_list) = ret
+        
+        if target is not None :
+          fields_targets[ifield].append( target)
+        fields_tokens_masked_idx_list[ifield].append( tokens_masked_idx_list)
+
+      rng_idx += 1
+
+      sources[ifield].append( field_data.unsqueeze(1) )
+      token_infos[ifield].append( token_info )
+
+    # merge along vertical level
+    sources[ifield] = torch.cat( sources[ifield], 1)
+    token_infos[ifield] = torch.cat( token_infos[ifield], 1)
+    # merge along vertical level, for target we have level, batch, ... ordering 
+    fields_targets[ifield] = torch.cat( fields_targets[ifield],0) \
+                                if len(fields_targets[ifield]) > 0 else fields_targets[ifield]
+
+  return (sources, token_infos, fields_targets, fields_tokens_masked_idx_list)
+
+####################################################################################################
+def prepare_batch_BERT_field( cf, ifield, source, token_info, rng) :
+
+  # shortcuts
+  mr = partial( torch.nn.functional.interpolate, mode='trilinear')
+  fl = torch.flatten
+  tr = torch.transpose
+  sq = torch.squeeze
+  usq = torch.unsqueeze 
+  cnt_nz = torch.count_nonzero
+
+  # collapse token dimensions 
+  source_shape0 = source.shape
+  source = torch.flatten( torch.flatten( source, 1, 3), 2, 4)
+  
+  # select random token in the selected space-time cube to be masked/deleted
+  BERT_frac = cf.fields[ifield][5][0]
+  BERT_frac_mask = cf.fields[ifield][5][1]
+  BERT_frac_rndm = cf.fields[ifield][5][2]
+  BERT_frac_mr = cf.fields[ifield][5][3]
+  BERT_mr_max = 2
+  token_size = cf.fields[ifield][4]
+  batch_dim = source.shape[0]
+  num_tokens = source.shape[1] 
+ 
+  masking_ratios = rng.random( batch_dim) * BERT_frac
+  # number of tokens masked per batch entry
+  nums_masked = np.ceil( num_tokens * masking_ratios).astype(np.int64)
+  tokens_masked_idx_list = [ torch.tensor(rng.permutation(num_tokens)[:nms]) for nms in nums_masked]
+  
+  # linear indices for masking
+  tokens_masked_idx_list = [tokens_masked_idx_list[i] + num_tokens * i for i in range(batch_dim)]
+  idx = torch.cat( tokens_masked_idx_list)
+
+  # flatten along first two dimension to simplify linear indexing (which then requires an
+  # easily computable row offset)
+  source_shape = source.shape
+  source = torch.flatten( source, 0, 1)
+ 
+  # keep masked tokens for loss computation
+  target = source[idx].clone()
+
+  # climatological mean of normalized data
+  global_mean = 0. * torch.mean(source, 0)
+  global_std  = torch.std(source, 0)
+
+  # Conditional masking (all are sampled independently so the fractions are only amortized)
+  # BERT_frac_mask = fraction masked       (80%)
+  # BERT_frac_rndm = fraction noisy tokens (10%)
+  # face_mr_BERT   = fraction of multi-res "masked" tokens (downsampled instead of fully masked)
+  # remainder      = untouched             (10%)
+  # 
+  # conditional idx for tokens to be masked with mean or random values
+  idx_mask_cond = torch.tensor( rng.random( idx.shape[0])) < BERT_frac_mask
+  idx_rndm_cond = torch.tensor( rng.random( idx.shape[0])) < BERT_frac_rndm
+  idx_mr_cond   = torch.tensor( rng.random( idx.shape[0])) < BERT_frac_mr
+  
+  # set mask 
+  source[ idx[idx_mask_cond] ] = global_mean
+  
+  # set noisy tokens (noise is field independent)
+  # TODO: fudge factor of 0.1
+  cnt_nz_rndm = cnt_nz(idx_rndm_cond)
+  dim_embed = source.shape[1]
+  rnd_w = tr( torch.tensor( rng.random( cnt_nz_rndm, dtype=np.float32)).repeat((dim_embed,1)), 0,1 )
+  source[ idx[idx_rndm_cond] ] = rnd_w * global_mean + (1.-rnd_w) * source[ idx[idx_rndm_cond] ] + \
+            0.1 * global_std * torch.randn( (cnt_nz_rndm, source.shape[1]))
+  
+  # randomly coarsened tokens
+  if BERT_frac_mr > 0. and idx_mr_cond.shape[0] > 0 and idx_mr_cond.any().item() :
+    # version with uniform multi-res downsampling per batch for computational efficiency 
+    # (in particular on GPU)
+    ts = torch.tensor( [token_size[1], token_size[2]], dtype=torch.int)
+    mrs = ((rng.random() * (ts - BERT_mr_max*ts)) + BERT_mr_max*ts).int()
+    mrs = (token_size[0], mrs[0], mrs[1])
+    ts = token_size 
+    # interpolate to smaller size and then interpolate up -> coarsened version with same size
+    # unsqueeze(usq()) is required since channel dimension is expected
+    temp = mr( mr( usq( source[ idx[idx_mr_cond] ].reshape( (-1,ts[0],ts[1],ts[2])), 1), mrs), ts) 
+    source[ idx[idx_mr_cond] ] = sq( fl( temp, -3, -1))
+
+  # recover batch dimension which was flattend for easier indexing and also token dimensions
+  source = torch.reshape( torch.reshape( source, source_shape), source_shape0)
+  
+  return (source, token_info, target, tokens_masked_idx_list)
+
+####################################################################################################
+def prepare_batch_BERT_forecast_field( cf, ifield, source, token_info, rng) :
+ 
+  nt = cf.forecast_num_tokens
+  num_tokens = source.shape[-6:-3]
+  num_tokens_space = num_tokens[1] * num_tokens[2] 
+  idxs = (num_tokens[0]-nt) * num_tokens_space + torch.arange(nt * num_tokens_space)
+  
+  # collapse token dimensions 
+  source_shape0 = source.shape
+  source = torch.flatten( torch.flatten( source, 1, 3), 2, 4)
+
+  # linear indices for masking
+  num_tokens = source.shape[1]
+  tokens_masked_idx_list = [idxs + num_tokens * i for i in range( source.shape[0] )]
+  idx = torch.cat( tokens_masked_idx_list)
+
+  source_shape = source.shape
+  # flatten along first two dimension to simplify linear indexing (which then requires an
+  # easily computable row offset)
+  source = torch.flatten( source, 0, 1)
+ 
+  # keep masked tokens for loss computation
+  target = source[idx].clone()
+
+  # masking
+  global_mean = 0. * torch.mean(source, 0)
+  source[ idx ] = global_mean
+ 
+  # recover batch dimension which was flattend for easier indexing
+  source = torch.reshape( torch.reshape( source, source_shape), source_shape0)
+ 
+  return (source, token_info, target, tokens_masked_idx_list)
+
+####################################################################################################
+def prepare_batch_BERT_temporal_field( cf, ifield, source, token_info, rng) :
+ 
+  num_tokens = source.shape[-6:-3]
+  num_tokens_space = num_tokens[1] * num_tokens[2]
+  
+  #backward compatibility: mask only middle token
+  if not hasattr( cf, 'idx_time_mask'):
+    idx_time_mask = int( np.floor(num_tokens[0] / 2.)) 
+    idxs = idx_time_mask * num_tokens_space + torch.arange(num_tokens_space)
+  else: #list of idx_time_mask
+    idxs = torch.concat([i*num_tokens_space + torch.arange(num_tokens_space) for i in cf.idx_time_mask])
+  
+  # collapse token dimensions 
+  source_shape0 = source.shape
+  source = torch.flatten( torch.flatten( source, 1, 3), 2, 4)
+
+  # linear indices for masking
+  num_tokens = source.shape[1]
+  idx = torch.cat( [idxs + num_tokens * i for i in range( source.shape[0] )] )
+  tokens_masked_idx_list = [idxs + num_tokens * i for i in range( source.shape[0] )]
+  source_shape = source.shape
+  # flatten along first two dimension to simplify linear indexing (which then requires an
+  # easily computable row offset)
+  source = torch.flatten( source, 0, 1)
+ 
+  # keep masked tokens for loss computation
+  target = source[idx].clone()
+
+  # masking
+  global_mean = 0. * torch.mean(source, 0)
+  source[ idx ] = global_mean
+ 
+  # recover batch dimension which was flattend for easier indexing
+  source = torch.reshape( torch.reshape( source, source_shape), source_shape0)
+ 
+  return (source, token_info, target, tokens_masked_idx_list)

vendor/atmorep-official/atmorep/transformer/axial_attention.py

@@ -0,0 +1,329 @@
+####################################################################################################
+#
+#  Copyright (C) 2022
+#
+####################################################################################################
+#
+#  project     : atmorep
+#
+#  author      : atmorep collaboration / lucidrains@github.com
+# 
+#  description : Based on:
+# https://github.com/lucidrains/axial-attention/blob/master/axial_attention/axial_attention.py
+#
+#  license     :
+#
+####################################################################################################
+
+# 
+import torch
+from torch import nn
+from operator import itemgetter
+import code
+# code.interact(local=locals())
+
+####################################################################################################
+
+# helper functions
+
+def exists(val):
+    return val is not None
+
+def map_el_ind(arr, ind):
+    return list(map(itemgetter(ind), arr))
+
+def sort_and_return_indices(arr):
+    indices = [ind for ind in range(len(arr))]
+    arr = zip(arr, indices)
+    arr = sorted(arr)
+    return map_el_ind(arr, 0), map_el_ind(arr, 1)
+
+# calculates the permutation to bring the input tensor to something attend-able
+# also calculates the inverse permutation to bring the tensor back to its original shape
+
+def calculate_permutations(num_dimensions, emb_dim, axial_dims = None) :
+    
+  total_dimensions = num_dimensions + 2
+  emb_dim = emb_dim if emb_dim > 0 else (emb_dim + total_dimensions)
+  if not axial_dims :
+    axial_dims = [ind for ind in range(1, total_dimensions) if ind != emb_dim]
+
+  permutations = []
+
+  for axial_dim in axial_dims:
+    last_two_dims = [axial_dim, emb_dim]
+    dims_rest = set(range(0, total_dimensions)) - set(last_two_dims)
+    permutation = [*dims_rest, *last_two_dims]
+    permutations.append(permutation)
+      
+  return permutations
+
+####################################################################################################
+
+# helper classes
+
+class PermuteToFrom(nn.Module):
+    
+  def __init__(self, permutation, fn):
+      
+    super().__init__()
+    self.fn = fn
+    _, inv_permutation = sort_and_return_indices(permutation)
+    self.permutation = permutation
+    self.inv_permutation = inv_permutation
+
+  def forward(self, x, kv = None, **kwargs):
+
+    if None == kv :
+
+      axial = x.permute(*self.permutation).contiguous()
+
+      shape = axial.shape
+      *_, t, d = shape
+
+      # merge all but axial dimension
+      axial = axial.reshape(-1, t, d)
+
+      # attention
+      axial = self.fn(axial, **kwargs)
+
+      # restore to original shape and permutation
+      axial = axial.reshape(*shape)
+      axial = axial.permute(*self.inv_permutation).contiguous()
+      
+    else :
+                    
+      axial = x.permute(*self.permutation)
+      axial_kv = kv.permute(*self.permutation)
+
+      shape = list(axial.shape)
+      shape_kv = axial_kv.shape
+      *_, t, d = shape
+      *_, t_kv, d_kv = shape_kv
+
+      # adjust token number along non-axial dimensions for efficient axial attention
+      # physically consistent by inserting the missing tokens for a lower res version
+      # axial dimension does not have to be matched; repeat_interleave is noop if ratio = 1
+      ratio1, ratio2 = shape[1] / shape_kv[1], shape[2] / shape_kv[2]
+      if int(ratio1) != 1 :
+        axial_kv = torch.repeat_interleave( axial_kv, max(1,int(ratio1)), dim=1 ).contiguous()
+        axial    = torch.repeat_interleave( axial, max(1,int(1./ratio1)), dim=1 ).contiguous()
+      if int(ratio2) != 1 :
+        axial_kv = torch.repeat_interleave( axial_kv, max(1,int(ratio2)), dim=2 ).contiguous()
+        axial    = torch.repeat_interleave( axial, max(1,int(1./ratio2)), dim=2 ).contiguous()
+      assert axial.shape[1:3] == axial_kv.shape[1:3], 'axial attn requires matching token numbers'
+
+
+      # merge all but axial dimension
+      axial_shape_perm = list(axial.shape)
+      axial = axial.reshape(-1, t, d)
+      axial_kv = axial_kv.reshape(-1, t_kv, d_kv)
+
+      # attention
+      axial = self.fn( axial, axial_kv, **kwargs)
+
+      # embedding dimension does not match input since not un-embedded in head
+      axial_shape_perm[-1] = -1
+      shape[-1] = -1
+      
+      # restore original, permuted shape modulo embedding dimension
+      axial = axial.reshape( *axial_shape_perm)
+      
+      # recover original token numbers for q if necessary by taking mean over inserted 
+      if shape[1] < shape_kv[1] :
+        axial = axial.reshape( shape[:2] + [int(1./ratio1)] + shape[2:] ).mean( dim=2)
+      if shape[2] < shape_kv[2] :
+        axial = axial.reshape( shape[:3] + [int(1./ratio2)] + shape[3:] ).mean( dim=3)
+
+      # restore to original shape modulo embedding dimension
+      axial = axial.reshape( *shape)
+      axial = axial.permute( *self.inv_permutation).contiguous()
+
+    return axial
+
+####################################################################################################
+# self attention
+
+class SelfAttention(nn.Module):
+
+  #############################################
+  def __init__(self, dim, heads, dim_heads = None):
+      
+    super().__init__()
+    
+    self.dim_heads = (dim // heads) if dim_heads is None else dim_heads
+    dim_hidden = self.dim_heads * heads
+
+    self.heads = heads
+    self.to_q = nn.Linear(dim, dim_hidden, bias = False)
+    self.to_kv = nn.Linear(dim, 2 * dim_hidden, bias = False)
+    self.to_out = nn.Linear(dim_hidden, dim)
+
+  #############################################
+  def forward(self, x, kv = None):
+      
+    kv = x if kv is None else kv
+    q, k, v = (self.to_q(x), *self.to_kv(kv).chunk(2, dim=-1))
+
+    b, t, d, h, e = *q.shape, self.heads, self.dim_heads
+
+    merge_heads = lambda x: x.reshape(b, -1, h, e).transpose(1, 2).reshape(b * h, -1, e)
+    q, k, v = map(merge_heads, (q, k, v))
+
+    dots = torch.einsum('bie,bje->bij', q, k) * (e ** -0.5)
+    dots = dots.softmax(dim=-1)
+    out = torch.einsum('bij,bje->bie', dots, v)
+
+    out = out.reshape(b, h, -1, e).transpose(1, 2).reshape(b, -1, d)
+    out = self.to_out(out)
+
+    return out
+
+####################################################################################################
+# axial attention class
+
+class AxialAttention( nn.Module):
+
+  #############################################
+  def __init__(self, dim, num_dimensions = 2, heads = 8, dim_heads = None, dim_index = -1, 
+                        sum_axial_out = True):
+
+    assert (dim % heads) == 0, 'hidden dimension must be divisible by number of heads'
+    super().__init__()
+    
+    self.dim = dim
+    self.total_dimensions = num_dimensions + 2
+    self.dim_index = dim_index if dim_index > 0 else (dim_index + self.total_dimensions)
+
+    attentions = []
+    for permutation in calculate_permutations(num_dimensions, dim_index):
+      attentions.append(PermuteToFrom(permutation, SelfAttention(dim, heads,dim_heads)))
+
+    self.axial_attentions = nn.ModuleList(attentions)
+    self.sum_axial_out = sum_axial_out
+
+  #############################################
+  def forward(self, x, kv = None):
+
+    assert len(x.shape) == self.total_dimensions, 'input does not have correct number of dimensions'
+    assert x.shape[self.dim_index] == self.dim, 'input does not have correct input dimension'
+
+    if self.sum_axial_out:
+      return sum(map(lambda axial_attn: axial_attn(x, kv), self.axial_attentions))
+
+    out = x
+    for axial_attn in self.axial_attentions:
+      out = axial_attn(out)
+
+    return out
+
+####################################################################################################
+
+class CrossAttention(nn.Module):
+
+  #############################################
+  def __init__(self, dims_embed, num_heads, dim_heads):
+    
+    super().__init__()
+    
+    assert 2 == len(dims_embed)
+    self.dim_heads = dim_heads
+    dim_hidden = self.dim_heads * num_heads
+
+    self.num_heads = num_heads
+    self.to_q = nn.Linear( dims_embed[0], dim_hidden, bias = False) 
+    self.to_kv = nn.Linear( dims_embed[1], 2 * dim_hidden, bias = False) 
+    
+  #############################################
+  def forward(self, x, kv):
+    
+    # per head projection
+    q, k, v = (self.to_q(x), *self.to_kv(kv).chunk(2, dim=-1))
+    
+    # extract require shape / dimension parameters
+    b, t, d, h, e = *q.shape, self.num_heads, self.dim_heads
+
+    # reshape and transpose so that axial dimension is second but last
+    mh = lambda x: x.reshape(b, -1, h, e).transpose(1, 2).reshape(b * h, -1, e)
+    q, k, v = map( mh, (q, k, v))
+
+    # compute attention 
+    dots = torch.einsum('bie,bje->bij', q, k) * (e ** -0.5)
+    dots = dots.softmax(dim=-1)
+    out = torch.einsum('bij,bje->bie', dots, v)
+
+    # recover input shape
+    out = out.reshape(b, h, -1, e).transpose(1, 2).reshape(b, -1, d)
+   
+    return out
+
+####################################################################################################
+
+class MultiFieldAxialAttention( nn.Module):
+
+  #############################################
+  def __init__(self, att_dims, dims_embed, num_heads_self, num_heads_cross, 
+                     sum_axial_out = False, dropout_rate = 0.0):
+
+    super().__init__()
+
+    assert 0 == (dims_embed[0] % (num_heads_self + num_heads_cross))
+    dim_embed_heads = int(dims_embed[0] / (num_heads_self + num_heads_cross))
+
+    self.num_fields = len(dims_embed)
+    self.dims_embed = dims_embed
+    self.att_dims = att_dims
+    num_dimensions = 3   # vertical, time, space (folded together) 
+    idx_embed = -1
+
+    self.sum_axial_out = sum_axial_out
+
+    self.lnorms = torch.nn.ModuleList()
+    for ifield in range( self.num_fields) :
+      self.lnorms.append( nn.LayerNorm( dims_embed[ifield], elementwise_affine=False) )
+    
+    # for all axes, create self and cross attention heads
+    self.axial_attentions = nn.ModuleList( [nn.ModuleList() for _ in att_dims] )
+    for ip, permutation in enumerate(calculate_permutations( num_dimensions, idx_embed, att_dims)) :
+      
+      # self-attention 
+      self.axial_attentions[ip].append( PermuteToFrom( permutation, CrossAttention( [dims_embed[0],dims_embed[0]], num_heads_self, dim_embed_heads)))
+      
+      # cross attention
+      for n in range( 1, len(dims_embed)) :
+        hs = PermuteToFrom( permutation, CrossAttention( [dims_embed[0], dims_embed[n]], 
+                                                          num_heads_cross, dim_embed_heads))
+        self.axial_attentions[ip].append( hs)
+
+    dim_hidden = dim_embed_heads * (num_heads_self + (num_heads_cross * (self.num_fields-1)) )
+    to_outs = [nn.Linear( dim_hidden, dims_embed[0]) for _ in att_dims]
+    self.to_outs = torch.nn.ModuleList( to_outs)
+
+    self.dropout = torch.nn.Dropout( p=dropout_rate)
+
+  #############################################
+  def forward(self, *args):
+
+    pass_through = args[0]
+    fields_lnormed = [self.lnorms[ifield](field) for ifield, field in enumerate( args)]
+    
+    out = []
+    for i_ax, axial_attns in enumerate( self.axial_attentions):
+      outs_axis = []
+      for i, attn_heads in enumerate(axial_attns) :
+        outs_axis.append( attn_heads( fields_lnormed[0], fields_lnormed[i]) )
+      ax_out = self.to_outs[i_ax]( torch.cat( outs_axis, -1) )
+      # parallel processing with sum_axial_out and then summation or serial otherwise
+      if self.sum_axial_out :
+        out.append( ax_out)
+      else : 
+        fields_lnormed[0] = ax_out 
+    
+    if self.sum_axial_out :
+      out = sum( out)
+    else :
+      out = fields_lnormed[0]
+
+    out = self.dropout( out)
+
+    return pass_through + out

vendor/atmorep-official/atmorep/transformer/decoder.py

@@ -0,0 +1,82 @@
+####################################################################################################
+#
+#  Copyright (C) 2022
+#
+####################################################################################################
+#
+#  project     : atmorep
+#
+#  author      : atmorep collaboration
+# 
+#  description :
+#
+#  license     :
+#
+####################################################################################################
+
+import torch
+import numpy as np
+import math
+
+from atmorep.transformer.transformer_base import MLP
+from atmorep.transformer.transformer_attention import MultiSelfAttentionHead, MultiCrossAttentionHead
+
+
+class Decoder(torch.nn.Module) :
+
+  ###################################
+  def __init__(self, num_layers, dim_embed = 2048, 
+                     num_heads = 8, num_mlp_layers = 2,
+                     self_att = False, cross_att_ratio = 0.5 ):
+    '''
+    Vaswani transformer corresponds to self_att = True and cross_att_ratio = 1.
+    '''
+    super( Decoder, self).__init__()
+    
+    self.num_layers = num_layers
+    self.dim_embed = dim_embed 
+
+    self.len_block = 2
+    if self_att :
+      self.len_block = 3
+
+    num_heads_other = int(num_heads * cross_att_ratio)
+    num_heads_self = num_heads - num_heads_other
+
+    self.blocks = torch.nn.ModuleList()
+    for _ in range( self.num_layers) :
+      # attention sub-block
+      if self_att :
+        self.blocks.append( MultiSelfAttentionHead( dim_embed, num_heads_self))
+      # cross attention between encoder and decoder
+      self.blocks.append( MultiCrossAttentionHead( dim_embed, num_heads_self, num_heads_other))
+      # feature space mapping sub-block
+      self.blocks.append( MLP( dim_embed, num_mlp_layers))
+
+  ###################################
+  def forward(self, token_seq_embed, encoder_out):
+    '''Evaluate decoder'''
+
+    for il in range(self.num_layers) : 
+      token_seq_embed = self.blocks[2*il]( token_seq_embed, encoder_out[il] )
+      token_seq_embed = self.blocks[2*il+1]( token_seq_embed, encoder_out[il] )
+
+    return token_seq_embed
+
+  ###################################
+  def get_attention( self, xin, iblock) :
+    ''' 
+    Get attention and projected values from specific layer and her head
+    '''
+
+    assert False
+
+    # # embedding
+    # token_seq_embed = prepare_token( self, xin, self.size_token_info)
+
+    # # attention heads + feature space mappings
+    # for idx in range(2*iblock) : 
+    #   token_seq_embed = self.blocks[idx]( token_seq_embed)
+    # atts, vsh = self.blocks[2*iblock].get_attention( token_seq_embed)
+
+    # return (atts, vsh)

vendor/atmorep-official/atmorep/transformer/interformer.py

@@ -0,0 +1,52 @@
+####################################################################################################
+#
+#  Copyright (C) 2022
+#
+####################################################################################################
+#
+#  project     : atmorep
+#
+#  author      : atmorep collaboration
+# 
+#  description :
+#
+#  license     :
+#
+####################################################################################################
+
+import torch
+import numpy as np
+import math
+
+from atmorep.transformer.transformer_base import MLP
+from atmorep.transformer.transformer_attention import MultiInterAttentionHead
+
+
+class Interformer(torch.nn.Module) :
+
+  def __init__(self, num_layers, dim_input, dims_embed,
+                      num_heads_self = 2, num_heads_coupling = 2, num_mlp_layers = 2,
+                      num_tokens = [], size_token_info = 6):
+    '''
+    
+    '''
+    
+    super(Interformer, self).__init__()
+    self.num_layers = num_layers
+    self.num_tokens = num_tokens
+    self.size_token_info = size_token_info
+    
+    # learnable linear embedding
+    self.token_embed = torch.nn.Linear( dim_input, dims_embed[0] - size_token_info)
+
+    self.blocks = torch.nn.ModuleList()
+    for _ in range( num_layers) :
+      # attention sub-block
+      self.blocks.append( MultiInterAttentionHead( dims_embed, num_heads_self, num_heads_coupling, 
+                                                   num_tokens ))
+      # feature space mapping sub-block
+      self.blocks.append( MLP( dims_embed[0], num_mlp_layers))
+
+  def forward(self, xin):
+    # never called directly but only through multiformer
+    assert False