Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,6 @@ 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
+weights/trained_model/decoder/tas/lt_1/epoch_18 filter=lfs diff=lfs merge=lfs -text
+weights/trained_model/encoder/epoch_96 filter=lfs diff=lfs merge=lfs -text
+weights/trained_model/processor/forecast_1/epoch_0 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,204 @@
+---
+license: other
+language:
+- en
+- zh
+tags:
+- OneScience
+- Earth Science
+- Weather Forecasting
+- Medium-Range Weather Forecasting
+frameworks: PyTorch
+datasets:
+- OneScience/Aardvark-Weather
+---
+
+<p align="center">
+  <strong><span style="font-size: 30px;">Aardvark Weather</span></strong>
+</p>
+
+# Model Overview
+
+Aardvark Weather is an end-to-end multimodal weather forecasting model that generates global gridded forecasts and station-level predictions through an observation encoder, a global forecast processor, and a station decoder.
+
+Paper: *End-to-end data-driven weather prediction*
+
+https://www.nature.com/articles/s41586-025-08897-0
+
+# Model Description
+
+This model package reuses the official code and weights to provide the following run pipeline:
+
+```text
+Official Multimodal Sample
+-> Encoder
+-> Day-1 Processor
+-> TAS Decoder
+-> 1-day global forecast and station 2-meter temperature
+```
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Official Model Verification | Inspect official samples, configuration, and checkpoints. |
+| Global Weather Forecasting | Output a global 1.5° gridded state of 24 variables. |
+| Station Temperature Forecasting | Output 2-meter temperature at 8,719 stations. |
+
+# 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**
+
+- Inference with the official weights requires an NVIDIA GPU.
+- CPU can be used for resource and checkpoint inspection; running full inference on CPU is not recommended.
+
+### Download the Model Package
+
+```bash
+hf download --model OneScience-Group/Aardvark-Weather --local-dir ./Aardvark-Weather
+cd Aardvark-Weather
+```
+
+### 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 the environment lacks Aardvark dependencies, refer to `official-src/environment.yml` for additional installation. The current adaptation is compatible with the `Block` parameter of the newer `timm` package.
+
+### Data & Weights
+
+The model package already includes the resources required for 1-day temperature inference:
+
+```text
+weights/sample_data/sample_data_final.pkl
+weights/trained_model/encoder/epoch_96
+weights/trained_model/processor/forecast_1/epoch_0
+weights/trained_model/decoder/tas/lt_1/epoch_18
+official-src/data/grid_lon_lat/
+official-src/data/norm_factors/
+```
+
+To re-download from source:
+
+```text
+Official Code: https://github.com/anna-allen/aardvark-weather-public
+Official Weights: https://huggingface.co/datasets/av555/aardvark-weather
+```
+
+### Training
+
+The training entry point provides a complete pipeline with epochs, validation, early stopping, learning rate scheduling, best/latest checkpointing, and resumption. By default, the Encoder and Day-1 Processor are frozen and only the TAS Decoder is trained:
+
+```bash
+python scripts/train.py
+```
+
+Paper-style end-to-end joint fine-tuning:
+
+```bash
+python scripts/train.py --train-modules all --epochs 10 --train-steps 100
+```
+
+Resuming training:
+
+```bash
+python scripts/train.py --resume weight/training/last.pth
+```
+
+The default configuration resides in `conf/config.yaml`. `--data` can point to an official-schema pickle file or a directory containing multiple `.pkl` files; multiple files are partitioned into train and validation sets in a deterministic manner. `--batch-size` concatenates multiple tasks along the existing batch dimension of the official task. When only the single official sample is included, training and validation reuse the same task — this allows end-to-end validation of the training software pipeline but does not constitute an independent validation set, nor can it provide the data diversity required to reproduce paper-level accuracy.
+
+Training artifacts:
+
+```text
+weight/training/best.pth
+weight/training/last.pth
+weight/training/history.json
+weight/training/train.json
+```
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Using weights obtained from training:
+
+```bash
+python scripts/inference.py --checkpoint weight/training/best.pth
+```
+
+Inference loads the official sample and the Day-1 `tas` weights by default. Results are saved to:
+
+```text
+result/inference_one_day.json
+result/prediction.pt
+result/target.pt
+```
+
+### Result Inspection
+
+```bash
+python scripts/result.py
+```
+
+Verified output shapes:
+
+```text
+initial_state:   [1, 121, 240, 24]
+global_forecast: [1, 121, 240, 24]
+station_tas:     [1, 8719]
+```
+
+Current results constitute a connectivity/end-to-end verification and do not reproduce the paper's RMSE/MAE metrics.
+
+The result script additionally produces `result/metrics.json` and `result/comparison.png`. The `normalized_mae` and `normalized_rmse` are computed in the normalized space of the official sample and should not be directly compared to the paper's physical-unit metrics.
+
+### Paper vs. Current Implementation I/O
+
+| Item | Paper | Current Package |
+| --- | --- | --- |
+| Input | Multimodal satellite, station, ship, and radiosonde observations | Bundled official sample pickle with a field structure consistent with the official Encoder |
+| Global State | `24 × 121 × 240`, 1.5° | Day-1 supported; output `[1,121,240,24]` |
+| Station Output | 2 m temperature and 10 m wind, up to Day-10 | Day-1 TAS only; `[1,8719]` |
+| Training | Staged pre-training followed by ~25,000 steps of end-to-end fine-tuning | Configurable full training loop; supports Decoder-only or full-model joint fine-tuning |
+| Evaluation | Grid-point weighted RMSE and station MAE in physical units | MAE/RMSE in normalized sample space |
+
+All commands should be run from the project root; `scripts/inference.py --root` converts to an absolute path. The official model internally depends on CUDA, so CPU is currently unavailable for inference. Training data must adhere to the official multimodal task dictionary schema; the current directory does not synthesize satellite or station observations — the bundled official sample serves as the default training-pipeline validator. The model package retains only official resources under `weights/` and does not carry local training artifacts under `weight/` or generated outputs under `result/`. Paper-level training still requires preparing observation data spanning the full date range and converting it into the same `.pkl` task contract.
+
+### Real Data
+
+Using real dates requires preparing ASCAT, AMSU-A/B, HIRS, IASI, GridSat, HadISD, ICOADS, IGRA, ERA5, topography, climatology, and their corresponding normalization statistics.
+
+# 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: https://github.com/anna-allen/aardvark-weather-public
+- This directory is an independent adaptation of the official Aardvark Weather model.
+- Code, weights, and data are subject to their respective official licenses and data terms.

conf/config.yaml

@@ -0,0 +1,14 @@
+training:
+  data: weights/sample_data/sample_data_final.pkl
+  output_dir: weight/training
+  epochs: 10
+  train_steps: 8
+  validation_steps: 2
+  validation_fraction: 0.2
+  batch_size: 1
+  learning_rate: 1.0e-5
+  weight_decay: 1.0e-6
+  gradient_clip: 1.0
+  patience: 3
+  seed: 2026
+  train_modules: decoder

configuration.json

@@ -0,0 +1 @@
+{"framework":"Pytorch","task":"weather-forecasting","implementation":"official-one-day-tas-adapter"}

model/__init__.py

@@ -0,0 +1 @@
+"""Aardvark-local adapters around the pinned official implementation."""

model/aardvark_adapter.py

@@ -0,0 +1,153 @@
+"""Thin, model-local wrapper for the official Aardvark Weather model."""
+
+from __future__ import annotations
+
+import hashlib
+import importlib
+import os
+import pickle
+import sys
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+
+
+TOP_LEVEL_KEYS = {"assimilation", "forecast", "downscaling", "y_target"}
+REQUIRED_ASSIMILATION_KEYS = {
+    "x_context_hadisd_current", "y_context_hadisd_current", "climatology_current",
+    "sat_x_current", "sat_current", "icoads_x_current", "icoads_current",
+    "igra_x_current", "igra_current", "amsua_current", "amsua_x_current",
+    "amsub_current", "amsub_x_current", "iasi_current", "iasi_x_current",
+    "ascat_current", "ascat_x_current", "hirs_current", "hirs_x_current",
+    "y_target_current", "era5_x_current",
+    "era5_elev_current", "era5_lonlat_current", "aux_time_current", "lt",
+    "y_target",
+}
+REQUIRED_FORECAST_KEYS = {"y_context", "y_target", "lt"}
+REQUIRED_DOWNSCALING_KEYS = {
+    "x_target", "alt_target", "y_target", "y_context", "x_context", "aux_time", "lt",
+}
+
+
+def sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def validate_sample(sample_path: Path) -> dict[str, Any]:
+    with sample_path.open("rb") as handle:
+        sample = pickle.load(handle)
+    if not isinstance(sample, dict) or set(sample) != TOP_LEVEL_KEYS:
+        raise ValueError(f"sample top-level keys mismatch: {list(sample) if isinstance(sample, dict) else type(sample)}")
+    expected = {
+        "assimilation": REQUIRED_ASSIMILATION_KEYS,
+        "forecast": REQUIRED_FORECAST_KEYS,
+        "downscaling": REQUIRED_DOWNSCALING_KEYS,
+    }
+    for name, keys in expected.items():
+        if not isinstance(sample[name], dict) or set(sample[name]) != keys:
+            raise ValueError(f"sample {name} keys mismatch: {list(sample[name])}")
+    if not isinstance(sample["y_target"], torch.Tensor) or sample["y_target"].ndim != 2:
+        raise ValueError("sample y_target must be a rank-2 torch.Tensor")
+    return {
+        "top_level_keys": sorted(sample),
+        "nested_keys": {name: sorted(value) for name, value in expected.items()},
+        "y_target_shape": list(sample["y_target"].shape),
+        "nan_counts": {
+            name: int(value.isnan().sum())
+            for name, value in sample["downscaling"].items()
+            if isinstance(value, torch.Tensor) and value.is_floating_point()
+        },
+    }
+
+
+def validate_checkpoint(path: Path) -> dict[str, Any]:
+    checkpoint = torch.load(path, map_location="cpu")
+    if not isinstance(checkpoint, dict) or "model_state_dict" not in checkpoint:
+        raise ValueError(f"checkpoint contract mismatch: {path}")
+    state = checkpoint["model_state_dict"]
+    if not isinstance(state, dict) or not state:
+        raise ValueError(f"empty model_state_dict: {path}")
+    return {"path": str(path), "key_count": len(state), "has_model_state_dict": True}
+
+
+def load_sample(sample_path: Path) -> dict[str, Any]:
+    with sample_path.open("rb") as handle:
+        return pickle.load(handle)
+
+
+def build_one_day_model(weights_root: Path, official_root: Path, device: str):
+    encoder = weights_root / "trained_model/encoder"
+    processor = weights_root / "trained_model/processor"
+    decoder = weights_root / "trained_model/decoder/tas"
+    sys.path.insert(0, str(official_root / "aardvark"))
+    _install_timm_compatibility()
+    official_e2e = importlib.import_module("e2e_model")
+    caller_dir = Path.cwd()
+    try:
+        os.chdir(official_root / "aardvark")
+        model = official_e2e.ConvCNPWeatherE2E(
+            device=device,
+            lead_time=1,
+            se_model_path=str(encoder),
+            forecast_model_path=str(processor),
+            sf_model_path=str(decoder) + "/",
+            return_gridded=True,
+            aux_data_path=str(official_root / "data") + "/",
+        )
+    finally:
+        os.chdir(caller_dir)
+    return model
+
+
+def run_one_day(sample_path: Path, weights_root: Path, official_root: Path, device: str) -> dict[str, Any]:
+    sample_report = validate_sample(sample_path)
+    encoder = weights_root / "trained_model/encoder"
+    processor = weights_root / "trained_model/processor"
+    decoder = weights_root / "trained_model/decoder/tas"
+    checkpoint_report = [
+        validate_checkpoint(encoder / "epoch_96"),
+        validate_checkpoint(processor / "forecast_1/epoch_0"),
+        validate_checkpoint(decoder / "lt_1/epoch_18"),
+    ]
+    sample = load_sample(sample_path)
+    model = build_one_day_model(weights_root, official_root, device)
+    model.eval()
+    with torch.inference_mode():
+        station, global_forecast, initial_state = model(sample)
+    for name, tensor in (("station_tas", station), ("global_forecast", global_forecast), ("initial_state", initial_state)):
+        if not isinstance(tensor, torch.Tensor) or not bool(torch.isfinite(tensor).all()):
+            raise ValueError(f"{name} contains non-finite values")
+    return {
+        "sample": sample_report,
+        "checkpoints": checkpoint_report,
+        "device": device,
+        "lead_time_days": 1,
+        "station_tas_shape": list(station.shape),
+        "global_forecast_shape": list(global_forecast.shape),
+        "initial_state_shape": list(initial_state.shape),
+        "finite_outputs": True,
+    }
+
+
+def _install_timm_compatibility() -> None:
+    """Bridge the old timm 0.6 Block constructor used by the official code."""
+    import timm.models.vision_transformer as vision_transformer
+
+    original = vision_transformer.Block
+    if getattr(original, "_aardvark_compat", False):
+        return
+
+    class AardvarkBlock(original):
+        _aardvark_compat = True
+
+        def __init__(self, *args: Any, drop: float = 0.0, **kwargs: Any) -> None:
+            kwargs.setdefault("proj_drop", drop)
+            super().__init__(*args, **kwargs)
+
+    vision_transformer.Block = AardvarkBlock

model/sample_dataset.py

@@ -0,0 +1,76 @@
+"""Datasets for training Aardvark from official-schema pickle tasks."""
+
+from __future__ import annotations
+
+import copy
+import pickle
+import random
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+
+from .aardvark_adapter import validate_sample
+
+
+def discover_samples(path: Path) -> list[Path]:
+    path = path.resolve()
+    samples = sorted(path.glob("*.pkl")) if path.is_dir() else [path]
+    if not samples or any(not sample.is_file() for sample in samples):
+        raise FileNotFoundError(f"No Aardvark sample pickle found at {path}")
+    for sample in samples:
+        validate_sample(sample)
+    return samples
+
+
+def split_samples(samples: list[Path], validation_fraction: float, seed: int) -> tuple[list[Path], list[Path]]:
+    if not 0.0 < validation_fraction < 1.0:
+        raise ValueError("validation_fraction must be between 0 and 1")
+    if len(samples) == 1:
+        return samples, samples
+    shuffled = samples.copy()
+    random.Random(seed).shuffle(shuffled)
+    validation_count = max(1, round(len(shuffled) * validation_fraction))
+    validation_count = min(validation_count, len(shuffled) - 1)
+    return shuffled[validation_count:], shuffled[:validation_count]
+
+
+class AardvarkTaskDataset(Dataset):
+    """Repeat one or more already-batched official tasks for a fixed number of steps."""
+
+    def __init__(self, samples: list[Path], steps: int) -> None:
+        if steps < 1:
+            raise ValueError("steps must be at least 1")
+        self.tasks = [self._load(path) for path in samples]
+        self.steps = steps
+
+    @staticmethod
+    def _load(path: Path) -> dict[str, Any]:
+        with path.open("rb") as handle:
+            return pickle.load(handle)
+
+    def __len__(self) -> int:
+        return self.steps
+
+    def __getitem__(self, index: int) -> dict[str, Any]:
+        return copy.deepcopy(self.tasks[index % len(self.tasks)])
+
+
+def collate_tasks(items: list[Any]) -> Any:
+    """Concatenate the batch dimension already present in official task tensors."""
+    first = items[0]
+    if isinstance(first, torch.Tensor):
+        return torch.cat(items, dim=0)
+    if isinstance(first, np.ndarray):
+        return np.concatenate(items, axis=0)
+    if isinstance(first, dict):
+        return {key: collate_tasks([item[key] for item in items]) for key in first}
+    if isinstance(first, tuple):
+        return tuple(collate_tasks(list(values)) for values in zip(*items))
+    if isinstance(first, list):
+        return [collate_tasks(list(values)) for values in zip(*items)]
+    if all(item == first for item in items):
+        return first
+    raise TypeError(f"Cannot collate Aardvark values of type {type(first).__name__}")

official-src/LICENSE

@@ -0,0 +1,121 @@
+Creative Commons Legal Code
+
+CC0 1.0 Universal
+
+    CREATIVE COMMONS CORPORATION IS NOT A LAW FIRM AND DOES NOT PROVIDE
+    LEGAL SERVICES. DISTRIBUTION OF THIS DOCUMENT DOES NOT CREATE AN
+    ATTORNEY-CLIENT RELATIONSHIP. CREATIVE COMMONS PROVIDES THIS
+    INFORMATION ON AN "AS-IS" BASIS. CREATIVE COMMONS MAKES NO WARRANTIES
+    REGARDING THE USE OF THIS DOCUMENT OR THE INFORMATION OR WORKS
+    PROVIDED HEREUNDER, AND DISCLAIMS LIABILITY FOR DAMAGES RESULTING FROM
+    THE USE OF THIS DOCUMENT OR THE INFORMATION OR WORKS PROVIDED
+    HEREUNDER.
+
+Statement of Purpose
+
+The laws of most jurisdictions throughout the world automatically confer
+exclusive Copyright and Related Rights (defined below) upon the creator
+and subsequent owner(s) (each and all, an "owner") of an original work of
+authorship and/or a database (each, a "Work").
+
+Certain owners wish to permanently relinquish those rights to a Work for
+the purpose of contributing to a commons of creative, cultural and
+scientific works ("Commons") that the public can reliably and without fear
+of later claims of infringement build upon, modify, incorporate in other
+works, reuse and redistribute as freely as possible in any form whatsoever
+and for any purposes, including without limitation commercial purposes.
+These owners may contribute to the Commons to promote the ideal of a free
+culture and the further production of creative, cultural and scientific
+works, or to gain reputation or greater distribution for their Work in
+part through the use and efforts of others.
+
+For these and/or other purposes and motivations, and without any
+expectation of additional consideration or compensation, the person
+associating CC0 with a Work (the "Affirmer"), to the extent that he or she
+is an owner of Copyright and Related Rights in the Work, voluntarily
+elects to apply CC0 to the Work and publicly distribute the Work under its
+terms, with knowledge of his or her Copyright and Related Rights in the
+Work and the meaning and intended legal effect of CC0 on those rights.
+
+1. Copyright and Related Rights. A Work made available under CC0 may be
+protected by copyright and related or neighboring rights ("Copyright and
+Related Rights"). Copyright and Related Rights include, but are not
+limited to, the following:
+
+  i. the right to reproduce, adapt, distribute, perform, display,
+     communicate, and translate a Work;
+ ii. moral rights retained by the original author(s) and/or performer(s);
+iii. publicity and privacy rights pertaining to a person's image or
+     likeness depicted in a Work;
+ iv. rights protecting against unfair competition in regards to a Work,
+     subject to the limitations in paragraph 4(a), below;
+  v. rights protecting the extraction, dissemination, use and reuse of data
+     in a Work;
+ vi. database rights (such as those arising under Directive 96/9/EC of the
+     European Parliament and of the Council of 11 March 1996 on the legal
+     protection of databases, and under any national implementation
+     thereof, including any amended or successor version of such
+     directive); and
+vii. other similar, equivalent or corresponding rights throughout the
+     world based on applicable law or treaty, and any national
+     implementations thereof.
+
+2. Waiver. To the greatest extent permitted by, but not in contravention
+of, applicable law, Affirmer hereby overtly, fully, permanently,
+irrevocably and unconditionally waives, abandons, and surrenders all of
+Affirmer's Copyright and Related Rights and associated claims and causes
+of action, whether now known or unknown (including existing as well as
+future claims and causes of action), in the Work (i) in all territories
+worldwide, (ii) for the maximum duration provided by applicable law or
+treaty (including future time extensions), (iii) in any current or future
+medium and for any number of copies, and (iv) for any purpose whatsoever,
+including without limitation commercial, advertising or promotional
+purposes (the "Waiver"). Affirmer makes the Waiver for the benefit of each
+member of the public at large and to the detriment of Affirmer's heirs and
+successors, fully intending that such Waiver shall not be subject to
+revocation, rescission, cancellation, termination, or any other legal or
+equitable action to disrupt the quiet enjoyment of the Work by the public
+as contemplated by Affirmer's express Statement of Purpose.
+
+3. Public License Fallback. Should any part of the Waiver for any reason
+be judged legally invalid or ineffective under applicable law, then the
+Waiver shall be preserved to the maximum extent permitted taking into
+account Affirmer's express Statement of Purpose. In addition, to the
+extent the Waiver is so judged Affirmer hereby grants to each affected
+person a royalty-free, non transferable, non sublicensable, non exclusive,
+irrevocable and unconditional license to exercise Affirmer's Copyright and
+Related Rights in the Work (i) in all territories worldwide, (ii) for the
+maximum duration provided by applicable law or treaty (including future
+time extensions), (iii) in any current or future medium and for any number
+of copies, and (iv) for any purpose whatsoever, including without
+limitation commercial, advertising or promotional purposes (the
+"License"). The License shall be deemed effective as of the date CC0 was
+applied by Affirmer to the Work. Should any part of the License for any
+reason be judged legally invalid or ineffective under applicable law, such
+partial invalidity or ineffectiveness shall not invalidate the remainder
+of the License, and in such case Affirmer hereby affirms that he or she
+will not (i) exercise any of his or her remaining Copyright and Related
+Rights in the Work or (ii) assert any associated claims and causes of
+action with respect to the Work, in either case contrary to Affirmer's
+express Statement of Purpose.
+
+4. Limitations and Disclaimers.
+
+ a. No trademark or patent rights held by Affirmer are waived, abandoned,
+    surrendered, licensed or otherwise affected by this document.
+ b. Affirmer offers the Work as-is and makes no representations or
+    warranties of any kind concerning the Work, express, implied,
+    statutory or otherwise, including without limitation warranties of
+    title, merchantability, fitness for a particular purpose, non
+    infringement, or the absence of latent or other defects, accuracy, or
+    the present or absence of errors, whether or not discoverable, all to
+    the greatest extent permissible under applicable law.
+ c. Affirmer disclaims responsibility for clearing rights of other persons
+    that may apply to the Work or any use thereof, including without
+    limitation any person's Copyright and Related Rights in the Work.
+    Further, Affirmer disclaims responsibility for obtaining any necessary
+    consents, permissions or other rights required for any use of the
+    Work.
+ d. Affirmer understands and acknowledges that Creative Commons is not a
+    party to this document and has no duty or obligation with respect to
+    this CC0 or use of the Work.

official-src/README.md

@@ -0,0 +1,71 @@
+# Aardvark Weather
+
+This repo contains code and weights used to run the Aardvark Weather model (https://www.nature.com/articles/s41586-025-08897-0). We provide sample data, model weights and code to demonstrate how the trained model is run to produce forecasts. 
+
+## Model weights
+Weights for the trained model presented in the paper are provided at `https://huggingface.co/datasets/av555/aardvark-weather/tree/main/trained_model`.
+Within this folder there are weights for the encoder, decoder and processor and end to end finetuned weights for one day lead time.
+For the decoder and end-to-end, we provide weights for both temperature and wind speed. 
+
+## Data
+Sample data is provided in `data/sample_data_final.pkl`. In addition we provide normalisation factors used to generate plots of predictions. For a visualisation of what is included in each timeslice of data passed to the model, see notebooks/data_demo.ipynb. For those interested in training their own end-to-end models, we have additionally prepared a machine learning ready dataset for the data sources utilised in the paper available at https://huggingface.co/datasets/av555/aardvark-weather. 
+
+## Forecast notebooks
+Notebooks demonstrating producing a forecast using the trained models are included in the `notebooks/` folder.
+Aardvark produces multiple modalities of forecasts.
+The notebook `forecast_demo.ipynb` demonstrates loading the complete Aardvark Weather model, generating predictions from the sample data and provides visualisations of the output global gridded and station forecasts.
+The notebook `e2e_finetune_demo.ipynb` provides a demonstration of generating optimised station forecasts from the sample data using the end to end finetuned model.
+
+## Training
+Aardvark weather is trained in multiple stages.
+Scripts to train the model are included in `training/`.
+Training of the three main modules is handled in `train_module.py`, processor finetuning in `finetune.py` and end-to-end finetuning in `e2e_train.py`.
+
+__Please note:__ the commands and related scripts under the "Encoder," "Processor," and "Decoder" sections below cannot be executed as they depend on local data loading pipelines, setup for the specific training compute infrastructure. The purpose of these training and finetuning commands is to illustrate the salient points of the training process for the purposes of transperancy and for those interested in the details of the training process and not to provide an executable version.
+
+### Encoder
+To train the encoder, run
+```
+bash train_encoder.sh
+```
+
+### Processor
+To pre-train the processor, we use the code in
+```
+bash train_processor.sh
+```
+
+To finetune the processor, we use the code in
+```
+bash finetune.sh
+```
+
+### Decoder
+To train the decoder, we use the code in
+```
+bash train_decoder.sh
+```
+
+### End-to-end
+To tune the model end-to-end, we use the code in
+```
+bash train_e2e.sh
+```
+## FAQ
+
+#### Can I access the data the model was trained on?
+We provide a dataset with observational data from 2007-2019 at 24 hour resolution at https://huggingface.co/datasets/av555/aardvark-weather. We hope that this will allow others to develop their own end-to-end weather models and explore the vast design space of ML architectures for this task. 
+
+#### Can I run the model in real time? 
+Unfortunately the datasets used in this initial prototype are not available in real time. We are however in the process of building a fully operational system, Aardvark 2.0. Updates to this will be provided here.  
+
+#### This model is at lower resolution than other AI models. Do you have any plans to improve on this?
+We are currently working on a new version of our model which runs at 0.25 degrees for a wider range of pressure levels. 
+
+#### Will further versions of the model and dataset be provided?
+
+Yes, this project is ongoing and we aim to continue updating this dataset and model. Updates currently in progress include the development of a fully operational system with real time data feed, extension of current historical dataset to 2025 and inclusion of further instruments. If you would like to receive updates as new products become available please email av555@cam.ac.uk with subject line 'Aardvark updates'.
+
+#### Can this dataset be used commercially?
+
+The dataset is released under a non-commercial no-derivatives license. This is due to the lisencing inherited from the source datasets.

official-src/aardvark/architectures.py

@@ -0,0 +1,37 @@
+import torch
+import torch.nn as nn
+
+from utils import *
+
+
+class MLP(nn.Module):
+    """
+    Multi-layer perceptron
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        h_channels=64,
+        h_layers=4,
+    ):
+
+        super().__init__()
+
+        def hidden_block(h_channels):
+            h = nn.Sequential(
+                nn.Linear(h_channels, h_channels),
+                nn.ReLU(),
+            )
+            return h
+
+        self.mlp = nn.Sequential(
+            nn.Linear(in_channels, h_channels),
+            nn.ReLU(),
+            *[hidden_block(h_channels) for _ in range(h_layers)],
+            nn.Linear(h_channels, out_channels)
+        )
+
+    def forward(self, x):
+        return self.mlp(x)

official-src/aardvark/data_shapes.py

@@ -0,0 +1,31 @@
+"""
+Data is stored in memmaps for access speed. Here we give the shapes required to load these files
+"""
+
+CLIMATOLOGY_SHAPE = (4, 366, 24, 240, 121)
+
+ICOADS_Y_SHAPE = (33601, 5, 12000)
+ICOADS_X_SHAPE = (33601, 2, 12000)
+
+IGRA_Y_SHAPE = (33604, 24, 1375)
+IGRA_X_SHAPE = (1375, 2)
+
+AMSUA_Y_SHAPE = (21916, 180, 360, 13)
+AMSUB_Y_SHAPE = (21916, 360, 181, 12)
+ASCAT_Y_SHAPE = (21913, 360, 181, 17)
+HIRS_Y_SHAPE = (21913, 360, 181, 26)
+GRIDSAT_Y_SHAPE = (48211, 2, 514, 200)
+IASI_Y_SHAPE = (23373, 360, 181, 52)
+
+
+def get_hadisd_shape(mode):
+    """
+    Return the shape of the HadISD array depending on variable
+    """
+
+    if mode != "train":
+        dim_1 = 415
+    else:
+        var_dict = {"tas": 8719, "tds": 8617, "psl": 8016, "u": 8721, "v": 8721}
+        dim_1 = var_dict[var]
+    return (106652, dim_1)

official-src/aardvark/e2e_model.py

@@ -0,0 +1,241 @@
+import pickle
+
+import torch
+import torch.nn as nn
+import numpy as np
+
+from models import ConvCNPWeather
+from misc_downscaling_functionality import ConvCNPWeatherOnToOff
+
+
+class ConvCNPWeatherE2E(nn.Module):
+    """
+    Complete Aardvark weather model. This chains together the trained encoder,
+    processor and decoder modules to create complete forecasts. It can be finetuned
+    end to end to optimise predictions for a specific variable and location.
+    """
+
+    def __init__(
+        self,
+        device,
+        lead_time,
+        se_model_path,
+        forecast_model_path,
+        sf_model_path,
+        return_gridded=False,
+        aux_data_path=None,
+    ):
+
+        super().__init__()
+
+        self.device = device
+        self.lead_time = lead_time
+        self.return_gridded = return_gridded
+
+        # Load encoder
+        self.se_model = self.load_se_model(se_model_path)
+
+        # Load processor
+        self.forecast_model = nn.ModuleList(
+            [
+                self.load_forecast_model(forecast_model_path, l + 1)
+                for l in range(lead_time)
+            ]
+        )
+
+        # Load decoder
+        self.sf_model = self.load_sf_model(sf_model_path, lead_time)
+        self.decoder = None
+
+        # Setup normalisation factors
+        self.forecast_input_means = (
+            self.to_tensor(np.load(aux_data_path + "norm_factors/mean_4u_1.npy"))
+            .unsqueeze(0)
+            .unsqueeze(0)
+            .unsqueeze(0)
+        )
+        self.forecast_input_stds = (
+            self.to_tensor(np.load(aux_data_path + "norm_factors/std_4u_1.npy"))
+            .unsqueeze(0)
+            .unsqueeze(0)
+            .unsqueeze(0)
+        )
+
+        self.forecast_pred_diff_means = (
+            self.to_tensor(np.load(aux_data_path + "norm_factors/mean_diff_4u_1.npy"))
+            .unsqueeze(0)
+            .unsqueeze(0)
+            .unsqueeze(0)
+        )
+        self.forecast_pred_diff_stds = (
+            self.to_tensor(np.load(aux_data_path + "norm_factors/std_diff_4u_1.npy"))
+            .unsqueeze(0)
+            .unsqueeze(0)
+            .unsqueeze(0)
+        )
+
+    def to_tensor(self, arr):
+        return torch.from_numpy(arr).float().to(self.device)
+
+    def load_se_model(self, se_model_path):
+        """
+        Load the trained encoder module
+        """
+
+        with open(se_model_path + "/config.pkl", "rb") as handle:
+            forecast_config = pickle.load(handle)
+
+        model = ConvCNPWeather(
+            in_channels=forecast_config["in_channels"],
+            out_channels=forecast_config["out_channels"],
+            int_channels=forecast_config["int_channels"],
+            device="cuda",
+            res=forecast_config["res"],
+            gnp=bool(0),
+            decoder=forecast_config["decoder"],
+            mode=forecast_config["mode"],
+            film=bool(0),
+        )
+
+        best_epoch = np.argmin(np.load("{}/losses_0.npy".format(se_model_path)))
+        state_dict = torch.load(
+            "{}/epoch_{}".format(se_model_path, best_epoch),
+            map_location="cuda",
+        )["model_state_dict"]
+        state_dict = {k[7:]: v for k, v in zip(state_dict.keys(), state_dict.values())}
+        model.load_state_dict(state_dict)
+        model = model.to("cuda")
+        return model
+
+    def load_forecast_model(self, forecast_model_path, lead_time):
+        """
+        Load the trained processor module
+        """
+
+        with open(forecast_model_path + "/config.pkl", "rb") as handle:
+            forecast_config = pickle.load(handle)
+
+        model = ConvCNPWeather(
+            in_channels=forecast_config["in_channels"],
+            out_channels=forecast_config["out_channels"],
+            int_channels=forecast_config["int_channels"],
+            device="cuda",
+            res=forecast_config["res"],
+            gnp=bool(0),
+            decoder=forecast_config["decoder"],
+            mode=forecast_config["mode"],
+            film=False,
+        )
+        state_dict = torch.load(
+            f"{forecast_model_path}/forecast_{lead_time}/epoch_0",
+            map_location="cuda",
+        )["model_state_dict"]
+        state_dict = {k[7:]: v for k, v in zip(state_dict.keys(), state_dict.values())}
+        model.load_state_dict(state_dict)
+        model = model.to("cuda")
+        return model
+
+    def load_sf_model(self, sf_model_path, lead_time):
+        """
+        Load the trained decoder module
+        """
+
+        with open(sf_model_path + "config.pkl", "rb") as handle:
+            config = pickle.load(handle)
+
+        model = ConvCNPWeatherOnToOff(
+            in_channels=config["in_channels"],
+            out_channels=config["out_channels"],
+            int_channels=config["int_channels"],
+            device="cuda",
+            res=config["res"],
+            decoder=config["decoder"],
+            mode=config["mode"],
+            film=False,
+        )
+
+        best_epoch = np.argmin(
+            np.load("{}/lt_{}/losses_0.npy".format(sf_model_path, lead_time))
+        )
+        full_state_dict = torch.load(
+            sf_model_path + f"/lt_{lead_time}/epoch_{best_epoch}", map_location="cuda"
+        )
+        state_dict = full_state_dict["model_state_dict"]
+        state_dict = {k[7:]: v for k, v in zip(state_dict.keys(), state_dict.values())}
+        model.load_state_dict(state_dict)
+        model = model.to("cuda")
+        model.eval()
+
+        return model
+
+    def process_se_output(self, task, x):
+        """
+        Reshape and normalise encoder output for input to processor
+        """
+
+        task["forecast"]["y_context"][:, :24, ...] = x.permute(0, 3, 2, 1)
+        if self.return_gridded:
+            return task, x.permute(0, 3, 2, 1)
+        return task
+
+    def process_forecast_output(self, task, x, last=False):
+        """
+        Reshape and normalise processor output for input to decoder
+        """
+
+        base_context = task["forecast"]["y_context"][:, :-11, ...].permute(0, 2, 3, 1)
+
+        base_context = (
+            base_context * self.forecast_input_stds + self.forecast_input_means
+        ).permute(0, 3, 2, 1)
+
+        x = self.forecast_pred_diff_means + x * self.forecast_pred_diff_stds
+
+        unnorm_x = x + base_context.permute(0, 2, 3, 1)
+
+        forecast = unnorm_x
+
+        x = (unnorm_x - self.forecast_input_means) / self.forecast_input_stds
+
+        task["downscaling"]["y_context"][:, :24, ...] = x.permute(0, 3, 2, 1)
+        task["forecast"]["y_context"] = torch.cat(
+            [
+                x.permute(0, 3, 2, 1),
+                task["forecast"]["y_context"][:, 24:, ...],
+            ],
+            axis=1,
+        )
+        if self.return_gridded:
+            return task, forecast
+        return task
+
+    def forward(self, task, film_index=None):
+        """
+        Produce a forecast
+        """
+
+        # Generate initial state
+        x = self.se_model(task["assimilation"], film_index=None)
+        if self.return_gridded:
+            task, initial_state = self.process_se_output(task, x)
+        else:
+            task = self.process_se_output(task, x)
+
+        # Generate forecast
+        for lt in range(self.lead_time):
+            x = self.forecast_model[lt](task["forecast"], film_index=None)
+            if self.return_gridded:
+                task, forecast = self.process_forecast_output(task, x)
+            else:
+                task = self.process_forecast_output(task, x)
+
+        # Generate station forecast
+        x = self.sf_model(task["downscaling"], film_index=None)
+
+        if self.return_gridded:
+            initial_state = (
+                initial_state.permute(0, 3, 2, 1) * self.forecast_input_stds
+                + self.forecast_input_means
+            )
+            return x, forecast, initial_state
+        return x

official-src/aardvark/e2e_train.py

@@ -0,0 +1,242 @@
+import os
+import sys
+import pickle
+import argparse
+
+import torch
+from torch.utils.data import DataLoader
+import torch.distributed as dist
+import torch.utils.data.distributed
+import torch.multiprocessing as mp
+from torch.utils.data.distributed import DistributedSampler
+from torch.distributed import init_process_group, destroy_process_group
+
+from misc_downscaling_functionality import DownscalingRmseLoss
+from e2e_model import *
+from loader import *
+from models import *
+from unet_wrap_padding import *
+from trainer import DDPTrainerE2E
+from loss_functions import WeightedRmseLoss, PressureWeightedRmseLoss, RmseLoss
+
+torch.set_float32_matmul_precision("medium")
+
+
+def ddp_setup(rank, world_size, master_port):
+    """
+    Setup DDP
+    """
+
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = master_port
+    init_process_group(
+        backend="nccl",
+        rank=rank,
+        world_size=world_size,
+    )
+    torch.cuda.set_device(rank)
+
+
+def start_date(name):
+    """
+    Set split start dates
+    """
+
+    if name == "train":
+        return "2007-01-02"
+    elif name == "val":
+        return "2019-01-01"
+    elif name == "test":
+        return "2018-01-01"
+    else:
+        raise Exception(f"Unrecognised split name {name}")
+
+
+def end_date(name):
+    """
+    Set split end dates
+    """
+
+    if name == "train":
+        return "2017-12-31"
+    elif name == "val":
+        return "2019-11-01"
+    elif name == "test":
+        return "2018-12-21"
+    else:
+        raise Exception(f"Unrecognised split name {name}")
+
+
+def main(rank, world_size, output_dir, args):
+    """
+    Run end to end training
+    """
+
+    master_port = args.master_port
+    lead_time = args.lead_time
+    era5_mode = args.era5_mode
+    ddp_setup(rank, world_size, master_port)
+
+    # Setup loss function
+    if args.loss == "lw_rmse":
+        lf = WeightedRmseLoss(
+            args.res,
+            start_ind=0,
+            end_ind=24,
+            weight_per_variable=False,
+        )
+    elif args.loss == "lw_rmse_pressure_weighted":
+        lf = PressureWeightedRmseLoss(args.res, era5_mode)
+    elif args.loss == "rmse":
+        lf = RmseLoss()
+    elif args.loss == "downscaling_rmse":
+        lf = DownscalingRmseLoss()
+
+    # Instantiate model
+    model = ConvCNPWeatherE2E(
+        device="cuda",
+        lead_time=lead_time,
+        se_model_path=args.se_model_path,
+        forecast_model_path=args.forecast_model_path,
+        sf_model_path=args.sf_model_path,
+    )
+    dist.barrier()
+
+    # Instantiate loaders
+    train_dataset = WeatherDatasetE2E(
+        device="cuda",
+        hadisd_mode="train",
+        start_date="2007-01-02",
+        end_date="2017-12-31",
+        lead_time=lead_time,
+        era5_mode="4u",
+        mode="train",
+        res=args.res,
+        var_start=0,
+        var_end=24,
+        diff=bool(0),
+        two_frames=bool(0),
+        region=args.region,
+        hadisd_var=args.var,
+        max_steps_per_epoch=args.max_steps_per_epoch,
+    )
+
+    val_dataset = WeatherDatasetE2E(
+        device="cuda",
+        hadisd_mode="train",
+        start_date="2019-01-01",
+        end_date="2019-12-21",
+        lead_time=lead_time,
+        era5_mode="4u",
+        mode="train",
+        res=args.res,
+        var_start=0,
+        var_end=24,
+        diff=bool(0),
+        two_frames=bool(0),
+        region=args.region,
+        hadisd_var=args.var,
+    )
+
+    test_dataset = WeatherDatasetE2E(
+        device="cuda",
+        hadisd_mode="train",
+        start_date="2018-01-01",
+        end_date="2018-12-21",
+        lead_time=lead_time,
+        era5_mode="4u",
+        mode="train",
+        res=args.res,
+        var_start=0,
+        var_end=24,
+        diff=bool(0),
+        two_frames=bool(0),
+        region=args.region,
+        hadisd_var=args.var,
+    )
+
+    train_sampler = DistributedSampler(train_dataset)
+    val_sampler = DistributedSampler(val_dataset)
+    test_sampler = DistributedSampler(test_dataset)
+
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        sampler=train_sampler,
+    )
+
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        sampler=val_sampler,
+    )
+
+    test_loader = DataLoader(
+        test_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        sampler=test_sampler,
+    )
+
+    # Instantiate trainer
+    trainer = DDPTrainerE2E(
+        model,
+        rank,
+        train_loader,
+        val_loader,
+        lf,
+        output_dir,
+        args.lr,
+        train_sampler,
+        hadisd_variable_name=args.var,
+        weight_decay=args.weight_decay,
+        weights_path=None,
+        tune_film=0,
+        test_loader=test_loader,
+    )
+
+    # Train model
+    trainer.train(n_epochs=args.epoch)
+    destroy_process_group()
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output_dir")
+    parser.add_argument("--loss")
+    parser.add_argument("--batch_size", type=int, default=3)
+    parser.add_argument("--epoch", type=int, default=10)
+    parser.add_argument("--master_port", default="12345")
+    parser.add_argument("--lr", type=float, default=5e-5)
+    parser.add_argument("--lead_time", type=int)
+    parser.add_argument("--era5_mode", default="4u")
+    parser.add_argument("--sf_model_path")
+    parser.add_argument("--se_model_path")
+    parser.add_argument("--forecast_model_path")
+
+    parser.add_argument("--weight_decay", type=float, default=1e-6)
+    parser.add_argument("--max_steps_per_epoch", type=int)
+    parser.add_argument("--res", type=int, default=1)
+    parser.add_argument("--frequency", type=int, default=6)
+    parser.add_argument("--region", default="global")
+    parser.add_argument("--var", choices=["tas", "u", "v", "psl", "ws"])
+    args = parser.parse_args()
+
+    torch.device("cuda")
+
+    # Make results directory
+    output_dir = args.output_dir
+    try:
+        os.mkdir(output_dir)
+    except:
+        pass
+
+    # Save config
+    with open(output_dir + "/config.pkl", "wb") as f:
+        pickle.dump(vars(args), f)
+
+    world_size = torch.cuda.device_count()
+    mp.spawn(main, args=[world_size, output_dir, args], nprocs=world_size)

official-src/aardvark/finetune.py

@@ -0,0 +1,431 @@
+"""
+NB: this script is for illustration purposes only and is not runnable as our
+full dataset is not provided as part of the submission due to size constraints.
+Many of the relevant paths to the data have been thus replaced by dummy paths.
+"""
+
+import os
+import argparse
+import pickle
+import subprocess
+from collections import OrderedDict
+
+import numpy as np
+from tqdm import tqdm
+import torch
+import torch.multiprocessing as mp
+import torch.utils.data.distributed
+from torch.utils.data import DataLoader, DistributedSampler
+from torch.distributed import init_process_group, destroy_process_group
+
+from loss_functions import WeightedRmseLoss, PressureWeightedRmseLoss, RmseLoss
+from trainer import DDPTrainer
+from loader import *
+from models import *
+from unet_wrap_padding import *
+
+torch.set_float32_matmul_precision("medium")
+
+weights = np.load(
+    "path_to_lat_weights/weights_lat_1.npy",
+).T[np.newaxis, ..., np.newaxis]
+
+
+def weighted_rmse_loss(target, output):
+    return np.sqrt(np.nanmean(((target - output) ** 2) * weights, axis=(0, 1, 2)))
+
+
+def unnorm_era5(x, mean, std):
+    x = x * std + mean
+    return x
+
+
+def norm_era5(x, mean, std):
+    x = (x - mean) / std
+    return x
+
+
+def ddp_setup(rank, world_size, master_port):
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = master_port
+    init_process_group(backend="nccl", rank=rank, world_size=world_size)
+    torch.cuda.set_device(rank)
+
+
+def main(rank, world_size, output_dir, args):
+    """
+    Finetune the processor module for each leadtime
+    """
+
+    # Setup
+    master_port = args.master_port
+    lead_time = args.lead_time
+    quicklook_id = args.output_dir
+    era5_mode = args.era5_mode
+    ddp_setup(rank, world_size, master_port)
+
+    # Instantiate loss function
+    if args.loss == "lw_rmse":
+        lf = WeightedRmseLoss(args.res, weight_per_variable=False)
+    elif args.loss == "lw_rmse_pressure_weighted":
+        lf = PressureWeightedRmseLoss(args.res, args.era5_mode)
+    elif args.loss == "rmse":
+        lf = RmseLoss()
+
+    # Load the pre-trained model
+    with open(args.forecast_model_path + "/config.pkl", "rb") as handle:
+        forecast_config = pickle.load(handle)
+
+    model = ConvCNPWeather(
+        in_channels=forecast_config["in_channels"],
+        out_channels=forecast_config["out_channels"],
+        int_channels=forecast_config["int_channels"],
+        device="cuda",
+        res=forecast_config["res"],
+        gnp=bool(0),
+        decoder=forecast_config["decoder"],
+        mode=forecast_config["mode"],
+        film=bool(forecast_config["film"]),
+    )
+
+    try:
+        best_epoch = np.argmin(np.load(forecast_config["output_dir"] + "/losses_0.npy"))
+        checkpoint = torch.load(
+            forecast_config["output_dir"] + "/epoch_{}".format(best_epoch)
+        )
+
+    except:
+        best_epoch = np.argmin(
+            np.load("../results/" + forecast_config["output_dir"] + "/losses_0.npy")
+        )
+        checkpoint = torch.load(
+            "../results/"
+            + forecast_config["output_dir"]
+            + "/epoch_{}".format(best_epoch)
+        )
+
+    new_state_dict = OrderedDict()
+    for k, v in checkpoint["model_state_dict"].items():
+        name = k[7:]  # remove `module.`
+        new_state_dict[name] = v
+
+    # Set the path to the prediction from the encoder
+    path_to_context = args.assimilation_model_path
+
+    # Finetune for each leadtime
+    for lead_time in range(1, 11):
+        print(f"Training lead time {lead_time}")
+
+        # Setup the model
+        model = ConvCNPWeather(
+            in_channels=forecast_config["in_channels"],
+            out_channels=forecast_config["out_channels"],
+            int_channels=forecast_config["int_channels"],
+            device="cuda",
+            res=forecast_config["res"],
+            gnp=bool(0),
+            decoder=forecast_config["decoder"],
+            mode=forecast_config["mode"],
+            film=bool(forecast_config["film"]),
+        )
+
+        model.load_state_dict(new_state_dict)
+
+        # Setup the loaders
+        if lead_time == 1:
+            path_to_context = args.assimilation_model_path
+        else:
+            path_to_context = f"{args.output_dir}/"
+
+        print(f"Loading context data from {path_to_context}...")
+        train_dataset = ForecastLoader(
+            device="cuda",
+            mode="train",
+            lead_time=lead_time,
+            era5_mode=args.era5_mode,
+            res=args.res,
+            frequency=args.frequency,
+            diff=bool(args.diff),
+            ic_path=path_to_context,
+            finetune_step=lead_time,
+            random_lt=False,
+            finetune_eval_every=250,
+            eval_steps=False,
+        )
+
+        val_dataset = ForecastLoader(
+            device="cuda",
+            mode="val",
+            lead_time=lead_time,
+            era5_mode=args.era5_mode,
+            res=args.res,
+            frequency=args.frequency,
+            diff=bool(args.diff),
+            ic_path=path_to_context,
+            finetune_step=lead_time,
+        )
+
+        test_dataset = ForecastLoader(
+            device="cuda",
+            mode="test",
+            lead_time=lead_time,
+            era5_mode=args.era5_mode,
+            res=args.res,
+            frequency=args.frequency,
+            diff=bool(args.diff),
+            ic_path=path_to_context,
+            finetune_step=lead_time,
+        )
+
+        train_sampler = DistributedSampler(train_dataset)
+        val_sampler = DistributedSampler(val_dataset)
+        test_sampler = DistributedSampler(test_dataset)
+
+        train_loader = DataLoader(
+            train_dataset,
+            batch_size=args.batch_size,
+            shuffle=False,
+            sampler=train_sampler,
+        )
+
+        val_loader = DataLoader(
+            val_dataset, batch_size=args.batch_size, shuffle=False, sampler=val_sampler
+        )
+
+        test_loader = DataLoader(
+            test_dataset,
+            batch_size=args.batch_size,
+            shuffle=False,
+            sampler=test_sampler,
+        )
+
+        # Setup output dir
+        try:
+            output_dir = f"{args.output_dir}/forecast_{lead_time}/"
+            os.mkdir(output_dir)
+        except:
+            pass
+
+        # Run the finetuning
+        trainer = DDPTrainer(
+            model,
+            rank,
+            train_loader,
+            val_loader,
+            lf,
+            output_dir,
+            args.lr,
+            train_sampler,
+            weight_decay=args.weight_decay,
+            tune_film=False,
+        )
+
+        # Finetune
+        n_epochs = args.finetune_epochs
+        trainer.train(n_epochs=n_epochs)
+
+        torch.distributed.barrier()
+
+        # Load the model just trained
+
+        best_epoch = np.argmin(
+            np.load(f"{args.output_dir}/forecast_{lead_time}/losses_0.npy")
+        )
+        checkpoint = torch.load(
+            f"{args.output_dir}/forecast_{lead_time}/epoch_{best_epoch}"
+        )
+
+        new_state_dict = OrderedDict()
+        for k, v in checkpoint["model_state_dict"].items():
+            name = k[7:]
+            new_state_dict[name] = v
+
+        model.load_state_dict(new_state_dict)
+        model.eval()
+
+        # Now need to make predictions using this model to create data to train on for the next leadtime
+
+        # Setup the output arrays
+        if rank == 0:
+            val_ic = np.memmap(
+                f"{args.output_dir}/ic_val_{lead_time}.mmap",
+                dtype="float32",
+                mode="w+",
+                shape=(len(val_dataset), 121, 240, 24),
+            )
+
+            test_ic = np.memmap(
+                f"{args.output_dir}/ic_test_{lead_time}.mmap",
+                dtype="float32",
+                mode="w+",
+                shape=(len(test_dataset), 121, 240, 24),
+            )
+
+            test_ic_target = np.memmap(
+                f"{args.output_dir}/ic_test_target_{lead_time}.mmap",
+                dtype="float32",
+                mode="w+",
+                shape=(len(test_dataset), 121, 240, 24),
+            )
+
+            train_ic = np.memmap(
+                f"{args.output_dir}/ic_train_{lead_time}.mmap",
+                dtype="float32",
+                mode="w+",
+                shape=(len(train_dataset), 121, 240, 24),
+            )
+
+        torch.distributed.barrier()
+        val_ic = np.memmap(
+            f"{args.output_dir}/ic_val_{lead_time}.mmap",
+            dtype="float32",
+            mode="r+",
+            shape=(len(val_dataset), 121, 240, 24),
+        )
+
+        train_ic = np.memmap(
+            f"{args.output_dir}/ic_train_{lead_time}.mmap",
+            dtype="float32",
+            mode="r+",
+            shape=(len(train_dataset), 121, 240, 24),
+        )
+
+        test_ic = np.memmap(
+            f"{args.output_dir}/ic_test_{lead_time}.mmap",
+            dtype="float32",
+            mode="r+",
+            shape=(len(test_dataset), 121, 240, 24),
+        )
+
+        test_ic_target = np.memmap(
+            f"{args.output_dir}/ic_test_target_{lead_time}.mmap",
+            dtype="float32",
+            mode="r+",
+            shape=(len(test_dataset), 121, 240, 24),
+        )
+
+        print(f"{args.output_dir}/ic_train_{lead_time}.mmap")
+        if lead_time > 1:
+            if rank == 0:
+                subprocess.run(
+                    [
+                        "rm",
+                        f"{args.output_dir}/ic_train_{lead_time-1}.mmap",
+                    ]
+                )
+
+        torch.distributed.barrier()
+
+        # Generate predictions on the test set
+        with tqdm(test_loader, unit="batch") as tepoch:
+            for count, batch in enumerate(tepoch):
+                inds = batch["target_index"].detach().cpu().numpy().astype(int)[:, 0]
+                out = model(batch, film_index=batch["lt"])
+
+                base_context = test_loader.dataset.unnorm_base_context(
+                    batch["y_context"][:, :-11, ...]
+                ).permute(0, 3, 2, 1)
+                unnorm_pred = test_loader.dataset.unnorm_pred(out)
+                unnorm_target = test_loader.dataset.unnorm_pred(batch["y_target"])
+
+                unnorm_pred = unnorm_pred + base_context
+                unnorm_target = unnorm_target + base_context
+
+                test_ic[inds, ...] = unnorm_pred.detach().cpu()
+                test_ic_target[inds, ...] = unnorm_target.detach().cpu()
+
+        torch.distributed.barrier()
+
+        # Generate predictions on the validation set
+        with tqdm(val_loader, unit="batch") as tepoch:
+            for count, batch in enumerate(tepoch):
+                inds = batch["target_index"].detach().cpu().numpy().astype(int)[:, 0]
+                out = model(batch, film_index=batch["lt"])
+
+                base_context = val_loader.dataset.unnorm_base_context(
+                    batch["y_context"][:, :-11, ...]
+                ).permute(0, 3, 2, 1)
+                unnorm_pred = val_loader.dataset.unnorm_pred(out)
+                unnorm_pred = unnorm_pred + base_context
+
+                val_ic[inds, ...] = unnorm_pred.detach().cpu()
+
+        torch.distributed.barrier()
+
+        # Generate predictions on the train set
+        with tqdm(train_loader, unit="batch") as tepoch:
+            for count, batch in enumerate(tepoch):
+                inds = batch["target_index"].detach().cpu().numpy().astype(int)[:, 0]
+                out = model(batch, film_index=batch["lt"])
+                base_context = val_loader.dataset.unnorm_base_context(
+                    batch["y_context"][:, :-11, ...]
+                ).permute(0, 3, 2, 1)
+                unnorm_pred = val_loader.dataset.unnorm_pred(out)
+                unnorm_pred = unnorm_pred + base_context
+
+                train_ic[inds, ...] = unnorm_pred.detach().cpu()
+
+        torch.distributed.barrier()
+
+        print("Setting up next lead time...")
+        path_to_context = f"{args.output_dir}/"
+
+        # Set the initial model weights for the next leadtime to be the current trained weights
+        if lead_time > 1:
+            best_epoch = np.argmin(
+                np.load(f"{args.output_dir}/forecast_{lead_time}/losses_0.npy")
+            )
+            checkpoint = torch.load(
+                f"{args.output_dir}/forecast_{lead_time}/epoch_{best_epoch}"
+            )
+            new_state_dict = OrderedDict()
+            for k, v in checkpoint["model_state_dict"].items():
+                name = k[7:]
+                new_state_dict[name] = v
+
+        torch.cuda.empty_cache()
+
+    destroy_process_group()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--output_dir")
+    parser.add_argument("--assimilation_model_path")
+    parser.add_argument("--forecast_model_path")
+
+    parser.add_argument("--loss", default="lw_rmse_pressure_weighted")
+    parser.add_argument("--batch_size", type=int, default=12)
+    parser.add_argument("--master_port", default="12345")
+    parser.add_argument("--lr", type=float, default=1e-4)
+    parser.add_argument("--lead_time", type=int, default=1)
+    parser.add_argument("--era5_mode", default="4u")
+    parser.add_argument("--weight_decay", type=float, default=1e-6)
+    parser.add_argument("--res", type=int, default=1)
+    parser.add_argument("--frequency", type=int, default=6)
+
+    parser.add_argument("--diff", type=int, default=1)
+    parser.add_argument("--start_ind", type=int, default=0)
+    parser.add_argument("--end_ind", type=int, default=24)
+    parser.add_argument("--finetune_epochs", type=int, default=5)
+
+    args = parser.parse_args()
+
+    device = torch.device("cuda")
+
+    # Setup results directory
+    output_dir = args.output_dir
+    try:
+        os.mkdir(output_dir)
+    except:
+        pass
+
+    # Save config
+    config = vars(args)
+
+    with open(output_dir + "/config.pkl", "wb") as f:
+        pickle.dump(config, f)
+
+    world_size = torch.cuda.device_count()
+    mp.spawn(main, args=[world_size, output_dir, args], nprocs=world_size)

official-src/aardvark/generate_initial_condition_single.py

@@ -0,0 +1,148 @@
+"""
+NB: this script is for illustration purposes only and is not runnable as our
+full dataset is not provided as part of the submission due to size constraints.
+Many of the relevant paths to the data have been thus replaced by dummy paths.
+"""
+
+import argparse
+import pickle
+
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+import torch
+import torch.nn as nn
+import torch.utils.data.distributed
+from torch.utils.data import DataLoader
+
+from loader import WeatherDatasetAssimilation
+from models import *
+
+torch.set_float32_matmul_precision("medium")
+
+
+def unnorm(x, mean, std, diff=False, av_2019=None):
+
+    x = x * std + mean
+    if diff:
+        return x + av_2019.transpose(0, 3, 2, 1)
+    return x
+
+
+if __name__ == "__main__":
+
+    """
+    Generate encoder predictions to be used as finetuning data for the processor module
+    """
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--encoder_model_path")
+    args = parser.parse_args()
+
+    # Load experiment config
+    with open(args.encoder_model_path + "/config.pkl", "rb") as handle:
+        forecast_config = pickle.load(handle)
+
+    device = "cuda"
+
+    # Setup normalisation factors
+    era5_mean_spatial = None
+    means = np.load(
+        "aux_data_path/norm_factors/mean_{}_{}.npy".format(
+            forecast_config["era5_mode"], forecast_config["res"]
+        )
+    )[np.newaxis, np.newaxis, np.newaxis, :]
+    stds = np.load(
+        "aux_data_path/norm_factors/std_{}_{}.npy".format(
+            forecast_config["era5_mode"], forecast_config["res"]
+        )
+    )[np.newaxis, np.newaxis, np.newaxis, :]
+
+    # Specify dates to generate predictions for
+    labels = ["train", "test", "val"]
+    dates = [
+        ["2007-01-02", "2017-12-31"],
+        ["2018-01-01", "2018-12-31"],
+        ["2019-01-01", "2019-12-31"],
+    ]
+
+    # Iterate over loaders
+    for label, date in zip(labels, dates):
+
+        n_times = pd.date_range(date[0], date[1], freq="6H")
+
+        # Setup a memmap to write out to
+        ic = np.memmap(
+            "{}/ic_{}.mmap".format(args.encoder_model_path, label),
+            dtype="float32",
+            mode="w+",
+            shape=(len(n_times), 121, 240, 24),
+        )
+
+        var_group_preds = []
+        var_group_targets = []
+
+        # Setup loader
+        dataset = WeatherDatasetAssimilation(
+            device="cuda",
+            hadisd_mode="train",
+            start_date=date[0],
+            end_date=date[1],
+            lead_time=0,
+            era5_mode="4u",
+            res=1,
+            var_start=0,
+            var_end=24,
+            diff=False,
+        )
+
+        loader = DataLoader(dataset, batch_size=64, shuffle=False)  # ,
+
+        # Instantiate and load model
+        model = ConvCNPWeather(
+            in_channels=forecast_config["in_channels"],
+            out_channels=forecast_config["out_channels"],
+            int_channels=forecast_config["int_channels"],
+            device="cuda",
+            res=forecast_config["res"],
+            gnp=bool(0),
+            decoder=forecast_config["decoder"],
+            mode=forecast_config["mode"],
+            film=bool(0),
+        )
+
+        best_epoch = np.argmin(
+            np.load("{}/losses_0.npy".format(args.encoder_model_path))
+        )
+        state_dict = torch.load(
+            "{}/epoch_{}".format(args.encoder_model_path, best_epoch),
+            map_location=device,
+        )["model_state_dict"]
+        state_dict = {k[7:]: v for k, v in zip(state_dict.keys(), state_dict.values())}
+        model.load_state_dict(state_dict)
+        model = nn.DataParallel(model)
+        model = model.cuda()
+
+        model.eval()
+
+        # Generate predictions
+        total = []
+        target = []
+
+        sum_count = 0
+        with torch.no_grad():
+            with tqdm(loader, unit="batch") as tepoch:
+                for count, batch in enumerate(tepoch):
+
+                    out = model(batch, film_index=batch["lt"]).detach().cpu().numpy()
+
+                    out_unnorm = unnorm(
+                        out,
+                        means,
+                        stds,
+                        diff=False,
+                        av_2019=era5_mean_spatial,
+                    )
+
+                    ic[sum_count : sum_count + out.shape[0], ...] = out_unnorm
+                    sum_count += out.shape[0]

official-src/aardvark/loader.py

@@ -0,0 +1,1825 @@
+import time as timelib
+from time import time
+
+import numpy as np
+import pandas as pd
+import torch
+from torch.utils.data import Dataset
+
+from loader_utils_new import *
+from data_shapes import *
+
+
+class WeatherDataset(Dataset):
+    """
+    Base weather dataset class
+    """
+
+    def __init__(
+        self,
+        device,
+        hadisd_mode,
+        start_date,
+        end_date,
+        lead_time,
+        era5_mode="train",
+        res=1,
+        filter_dates=None,
+        diff=None,
+    ):
+
+        super().__init__()
+
+        # Setup
+        self.device = device
+        self.mode = hadisd_mode
+        self.data_path = "path_to_data/"
+        self.aux_data_path = "path_to_auxiliary_data/"
+        self.start_date = start_date
+        self.end_date = end_date
+        self.lead_time = lead_time
+        self.era5_mode = era5_mode
+        self.res = res
+        self.filter_dates = filter_dates
+        self.diff = diff
+
+        # Date indexing
+        self.dates = pd.date_range(start_date, end_date, freq="6H")
+        if self.filter_dates == "start":
+            self.index = np.array([i for i, d in enumerate(self.dates) if d.month < 7])
+        elif self.filter_dates == "end":
+            self.index = np.array([i for i, d in enumerate(self.dates) if d.month >= 7])
+        else:
+            self.index = np.array(range(len(self.dates)))
+
+        # Load the input modalities
+        print("Loading IGRA")
+        self.load_igra()
+
+        print("Loading AMSU-A")
+        self.load_amsua()
+
+        print("Loading AMSU-B")
+        self.load_amsub()
+
+        print("Loading ICOADS")
+        self.load_icoads()
+
+        print("Loading IASI")
+        self.load_iasi()
+
+        print("Loading GEO")
+        self.load_sat_data()
+
+        print("Loading HADISD")
+        self.load_hadisd(self.mode)
+
+        print("Loading ASCAT")
+        self.load_ascat_data()
+        self.load_hirs_data()
+
+        # Load the ground truth data for training
+        print("Loading ERA5")
+        self.era5_sfc = [
+            self.load_era5(year)
+            for year in range(int(start_date[:4]), int(end_date[:4]) + 1)
+        ]
+
+        # Internal grid to longitude latitude correspondence
+        self.era5_x = [
+            self.to_tensor(
+                np.load(self.data_path + "era5/era5_x_{}.npy".format(self.res))
+            )
+            / LATLON_SCALE_FACTOR,
+            self.to_tensor(
+                np.load(self.data_path + "era5/era5_y_{}.npy".format(self.res))
+            )
+            / LATLON_SCALE_FACTOR,
+        ]
+
+        # Orography
+        self.era5_elev = self.to_tensor(
+            np.load(self.data_path + "era5/elev_vars_{}.npy".format(self.res))
+        )
+        self.era5_elev = torch.flip(self.era5_elev.permute(0, 2, 1), [-1])
+        xx, yy = torch.meshgrid(self.era5_x[0], self.era5_x[1])
+        self.era5_lonlat = torch.stack([xx, yy])
+
+        # Climatology
+        self.climatology = np.memmap(
+            self.data_path + "climatology_data.mmap",
+            dtype="float32",
+            mode="r",
+            shape=CLIMATOLOGY_SHAPE,
+        )
+
+        # Setup normalisation factors
+        if self.diff:
+            self.era5_mean_spatial = np.load(
+                self.aux_data_path + "era5_spatial_means.npy"
+            )[0, ...]
+            self.means = np.load(self.aux_data_path + "era5_avdiff_means.npy")[
+                :, np.newaxis, np.newaxis, ...
+            ]
+            self.stds = np.load(self.aux_data_path + "era5_avdiff_stds.npy")[
+                :, np.newaxis, np.newaxis, ...
+            ]
+        else:
+            self.means = np.load(
+                self.aux_data_path
+                + "norm_factors/mean_{}_{}.npy".format(self.era5_mode, self.res)
+            )[:, np.newaxis, np.newaxis, ...]
+            self.stds = np.load(
+                self.aux_data_path
+                + "norm_factors/std_{}_{}.npy".format(self.era5_mode, self.res)
+            )[:, np.newaxis, np.newaxis, ...]
+
+    def load_icoads(self):
+        """
+        Load the ICOADS data
+        """
+
+        self.icoads_y = np.memmap(
+            self.data_path + "icoads/1999_2021_icoads_y.mmap",
+            dtype="float32",
+            mode="r",
+            shape=ICOADS_Y_SHAPE,
+        )
+
+        self.icoads_x = (
+            np.memmap(
+                self.data_path + "icoads/1999_2021_icoads_x.mmap",
+                dtype="float32",
+                mode="r",
+                shape=ICOADS_X_SHAPE,
+            )
+            / LATLON_SCALE_FACTOR
+        )
+        self.icoads_means = self.to_tensor(
+            np.load(self.aux_data_path + "norm_factors/mean_icoads.npy")
+        )
+        self.icoads_stds = self.to_tensor(
+            np.load(self.aux_data_path + "norm_factors/std_icoads.npy")
+        )
+        self.icoads_means = self.to_tensor(
+            np.nanmean(self.icoads_y[-365 * 4 :, ...], axis=(0, 2))[:, np.newaxis]
+        )
+        self.icoads_stds = self.to_tensor(
+            np.nanstd(self.icoads_y[-365 * 4 :, ...], axis=(0, 2))[:, np.newaxis]
+        )
+        self.icoads_index_offset = ICOADS_OFFSETS[self.start_date]
+        return
+
+    def load_igra(self):
+        """
+        Load the IGRA data
+        """
+
+        self.igra_y = np.memmap(
+            self.data_path + "igra/1999_2021_igra_y.mmap",
+            dtype="float32",
+            mode="r",
+            shape=IGRA_Y_SHAPE,
+        )
+
+        self.igra_x = np.copy(
+            np.memmap(
+                self.data_path + "igra/1999_2021_igra_x.mmap",
+                dtype="float32",
+                mode="r",
+                shape=IGRA_X_SHAPE,
+            )
+        )
+        self.igra_x = self.igra_x / LATLON_SCALE_FACTOR
+
+        self.igra_means = self.to_tensor(
+            np.load(self.aux_data_path + "norm_factors/mean_igra.npy")
+        )
+        self.igra_stds = self.to_tensor(
+            np.load(self.aux_data_path + "norm_factors/std_igra.npy")
+        )
+
+        self.igra_index_offset = IGRA_OFFSETS[self.start_date]
+
+        return
+
+    def load_amsua(self):
+        """
+        Load the AMSU-A data
+        """
+
+        self.amsua_y = np.memmap(
+            self.data_path + "amsua/2007_2021_amsua.mmap",
+            dtype="float32",
+            mode="r",
+            shape=AMSUA_Y_SHAPE,
+        )
+        self.amsua_index_offset = AMSUA_OFFSETS[self.start_date]
+
+        xx = np.linspace(-180, 179, 360, dtype=np.float32)
+        xx = ((xx + 360) % 360) / LATLON_SCALE_FACTOR
+        yy = np.linspace(90, -90, 180, dtype=np.float32) / LATLON_SCALE_FACTOR
+        self.amsua_x = [xx, yy]
+
+        self.amsua_means = self.to_tensor(
+            np.load(self.aux_data_path + "norm_factors/mean_amsua.npy")
+        )
+        self.amsua_stds = self.to_tensor(
+            np.load(self.aux_data_path + "norm_factors/std_amsua.npy")
+        )
+
+        return
+
+    def load_amsub(self):
+        """
+        Load the AMSU-B data
+        """
+
+        self.amsub_y = np.memmap(
+            self.data_path + "amsub_mhs/2007_2021_amsub.mmap",
+            dtype="float32",
+            mode="r",
+            shape=AMSUB_Y_SHAPE,
+        )
+        self.amsub_index_offset = AMSUB_OFFSETS[self.start_date]
+
+        xx = np.linspace(0, 359, 360, dtype=np.float32)
+        xx = ((xx + 360) % 360) / LATLON_SCALE_FACTOR
+        yy = np.linspace(90, -90, 181, dtype=np.float32) / LATLON_SCALE_FACTOR
+        self.amsub_x = [xx, yy]
+
+        self.amsub_means = self.to_tensor(
+            np.load(self.aux_data_path + "norm_factors/mean_amsub.npy")
+        )
+        self.amsub_stds = self.to_tensor(
+            np.load(self.aux_data_path + "norm_factors/std_amsub.npy")
+        )
+
+        return
+
+    def load_ascat_data(self):
+        """
+        Load the ASCAT data
+        """
+
+        self.ascat_y = np.memmap(
+            self.data_path + "ascat/2007_2021_ascat.mmap",
+            dtype="float32",
+            mode="r",
+            shape=ASCAT_Y_SHAPE,
+        )
+        self.ascat_index_offset = ASCAT_OFFSETS[self.start_date]
+
+        xx = np.linspace(0, 359, 360, dtype=np.float32)
+        xx = ((xx + 360) % 360) / LATLON_SCALE_FACTOR
+        yy = np.linspace(-90, 90, 181, dtype=np.float32) / LATLON_SCALE_FACTOR
+        self.ascat_x = [xx, np.copy(yy[::-1])]
+
+        self.ascat_means = self.to_tensor(
+            np.load(self.aux_data_path + "norm_factors/mean_ascat.npy")
+        )
+        self.ascat_stds = self.to_tensor(
+            np.load(self.aux_data_path + "norm_factors/std_ascat.npy")
+        )
+
+        return
+
+    def load_hirs_data(self):
+        """
+        Load the HIRS data
+        """
+
+        self.hirs_y = np.memmap(
+            self.data_path + "hirs/2007_2021_hirs.mmap",
+            dtype="float32",
+            mode="r",
+            shape=HIRS_Y_SHAPE,
+        )
+        self.hirs_index_offset = ASCAT_OFFSETS[self.start_date]
+
+        xx = np.linspace(0, 359, 360, dtype=np.float32)
+        xx = ((xx + 360) % 360) / LATLON_SCALE_FACTOR
+        yy = np.linspace(-90, 90, 181, dtype=np.float32) / LATLON_SCALE_FACTOR
+        self.hirs_x = [xx, np.copy(yy[::-1])]
+
+        self.hirs_means = self.to_tensor(
+            np.load(self.aux_data_path + "norm_factors/hirs_means.npy")
+        )
+        self.hirs_stds = self.to_tensor(
+            np.load(self.aux_data_path + "norm_factors/hirs_stds.npy")
+        )
+
+        return
+
+    def load_sat_data(self):
+        """
+        Load the GRIDSAT data
+        """
+
+        self.sat_y = np.memmap(
+            self.data_path + "gridsat/gridsat_data.mmap",
+            dtype="float32",
+            mode="r",
+            shape=GRIDSAT_Y_SHAPE,
+        )
+
+        xx = np.load(self.data_path + "gridsat/sat_x.npy") / LATLON_SCALE_FACTOR
+        yy = np.load(self.data_path + "gridsat/sat_y.npy") / LATLON_SCALE_FACTOR
+        self.sat_x = [xx, yy]
+        self.sat_index_offset = SAT_OFFSETS[self.start_date]
+
+        self.sat_means = self.to_tensor(
+            np.load(self.aux_data_path + "norm_factors/mean_sat.npy")
+        )
+        self.sat_stds = self.to_tensor(
+            np.load(self.aux_data_path + "norm_factors/std_sat.npy")
+        )
+
+        return
+
+    def load_iasi(self):
+        """
+        Load the IASI data
+        """
+
+        self.iasi = np.memmap(
+            self.data_path + "2007_2021_iasi_subset.mmap",
+            dtype="float32",
+            mode="r",
+            shape=IASI_Y_SHAPE,
+        )
+        self.iasi_index_offset = ASCAT_OFFSETS[self.start_date]
+
+        xx = np.linspace(0, 359, 360, dtype=np.float32)
+        xx = ((xx + 360) % 360) / LATLON_SCALE_FACTOR
+        yy = np.linspace(-90, 90, 181, dtype=np.float32) / LATLON_SCALE_FACTOR
+        self.iasi_x = [xx, np.copy(yy[::-1])]
+
+        self.iasi_means = self.to_tensor(
+            np.load(self.aux_data_path + "norm_factors/mean_iasi.npy")
+        )
+        self.iasi_stds = self.to_tensor(
+            np.load(self.aux_data_path + "norm_factors/std_iasi.npy")
+        )
+
+        return
+
+    def load_hadisd(self, mode):
+        """
+        Load the HADISD data
+        """
+
+        self.hadisd_x = []
+        self.hadisd_alt = []
+        self.hadisd_y = []
+        hadisd_vars = ["tas", "tds", "psl", "u", "v"]
+        for var in hadisd_vars:
+            lon = lon_to_0_360(
+                np.load(
+                    self.data_path + "hadisd_processed/{}_lon_{}.npy".format(var, mode)
+                )
+            )
+            lat = np.load(
+                self.data_path + "hadisd_processed/{}_lat_{}.npy".format(var, mode)
+            )
+            alt = np.load(
+                self.data_path + "hadisd_processed/{}_alt_{}.npy".format(var, mode)
+            )
+
+            vals = np.memmap(
+                self.data_path
+                + "hadisd_processed/{}_vals_{}.memmap".format(var, self.mode),
+                dtype="float32",
+                mode="r",
+                shape=get_hadisd_shape(mode),
+            )
+
+            self.hadisd_x.append(np.stack([lon, lat], axis=-1) / LATLON_SCALE_FACTOR)
+            self.hadisd_alt.append(alt)
+            self.hadisd_y.append(vals)
+
+        self.hadisd_index_offset = HADISD_OFFSETS[self.start_date]
+
+        self.hadisd_means = [
+            self.to_tensor(
+                np.load(
+                    self.aux_data_path + "norm_factors/mean_hadisd_{}.npy".format(var)
+                )
+            )
+            for var in hadisd_vars
+        ]
+        self.hadisd_stds = [
+            self.to_tensor(
+                np.load(
+                    self.aux_data_path + "norm_factors/std_hadisd_{}.npy".format(var)
+                )
+            )
+            for var in hadisd_vars
+        ]
+
+        return
+
+    def load_era5(self, year):
+        """
+        Load the ERA5 training data
+        """
+
+        if year % 4 == 0:
+            d = 366 * 4
+        else:
+            d = 365 * 4
+
+        if self.era5_mode == "sfc":
+            levels = 4
+        elif self.era5_mode == "13u":
+            levels = 69
+        else:
+            levels = 24
+
+        if self.res == 1:
+            x = 240
+            y = 121
+        elif self.res == 5:
+            x = 64
+            y = 32
+        mmap = np.memmap(
+            self.data_path
+            + "/era5/era5_{}_{}_6_{}.memmap".format(self.era5_mode, self.res, year),
+            dtype="float32",
+            mode="r",
+            shape=(d, levels, x, y),
+        )
+        return mmap
+
+    def norm_era5(self, x):
+
+        x = (x - self.means) / self.stds
+        return x
+
+    def unnorm_era5(self, x):
+
+        x = x * self.stds + self.means
+        return x
+
+    def norm_data(self, x, means, stds):
+        return (x - means) / stds
+
+    def norm_hadisd(self, x):
+        for i in range(5):
+            x[i] = (x[i] - self.hadisd_means[i]) / self.hadisd_stds[i]
+        return x
+
+    def __len__(self):
+        return self.index.shape[0] - 1 - 1
+
+    def to_tensor(self, arr):
+        return torch.from_numpy(arr).float().to(self.device)
+
+    def get_time_aux(self, current_date):
+        """
+        Return the auxiliary temporal channels given a date
+        """
+
+        doy = current_date.dayofyear
+        year = (current_date.year - 2007) / 15
+        time_of_day = current_date.hour
+        return np.array(
+            [
+                np.cos(np.pi * 2 * doy / DAYS_IN_YEAR),
+                np.sin(np.pi * 2 * doy / DAYS_IN_YEAR),
+                np.cos(np.pi * 2 * time_of_day / 24),
+                np.sin(np.pi * 2 * time_of_day / 24),
+                year,
+            ]
+        )
+
+
+class WeatherDatasetAssimilation(WeatherDataset):
+    """
+    Encoder training loader
+    """
+
+    def __init__(
+        self,
+        device,
+        hadisd_mode,
+        start_date,
+        end_date,
+        lead_time,
+        era5_mode="sfc",
+        res=1,
+        filter_dates=None,
+        var_start=0,
+        var_end=24,
+        diff=False,
+        two_frames=False,
+    ):
+
+        super().__init__(
+            device,
+            hadisd_mode,
+            start_date,
+            end_date,
+            lead_time,
+            era5_mode,
+            res=res,
+            filter_dates=filter_dates,
+            diff=diff,
+        )
+
+        # Setup
+
+        self.var_start = var_start
+        self.var_end = var_end
+        self.diff = diff
+        self.two_frames = two_frames
+
+    def load_era5_time(self, index):
+        """
+        ERA5 ground truth data loading
+        """
+
+        date = self.dates[index]
+        year = date.year
+        hour = date.hour
+        doy = (date.dayofyear - 1) * 4 + (hour // 6)
+
+        era5 = self.era5_sfc[year - int(self.start_date[:4])][doy, ...]
+        era5 = np.copy(era5)
+        if self.diff:
+            era5 = era5 - self.era5_mean_spatial
+        era5 = self.norm_era5(era5[np.newaxis, ...])[0, ...]
+        return era5
+
+    def load_year_end(self, year, doy):
+        data_1 = self.era5_sfc[year - int(self.start_date[:4])][doy : doy + 1, ...]
+        missing = self.lead_time - data_1.shape[0] + 1
+        data_2 = self.era5_sfc[year - int(self.start_date[:4]) + 1][:missing, ...]
+        data = np.concatenate([data_1, data_2])
+        return data
+
+    def load_era5_slice(self, index):
+        """
+        ERA5 ground truth data loading
+        """
+
+        date = self.dates[index]
+        year = date.year
+        doy = (date.dayofyear - 1) * 4
+
+        next_date = self.dates[index + 1]
+        next_year = next_date.year
+
+        if next_year != year:
+            era5 = self.load_year_end(year, doy)
+        else:
+            era5 = self.era5_sfc[year - int(self.start_date[:4])][doy : doy + 1, ...]
+
+        era5 = self.norm_era5(np.copy(era5))
+        return era5
+
+    def __getitem__(self, index):
+
+        if self.two_frames:
+            # Case 1: loading t=0 and t=-1
+            index = index + 1
+            current = self.get_index(index, "current")
+            prev = self.get_index(index - 1, "prev")
+            current["y_target"] = current["y_target_current"]
+
+            return {**current, **prev}
+        else:
+            # Case 2: loading t=0
+            current = self.get_index(index, "current")
+            current["y_target"] = current["y_target_current"]
+
+            return {**current}
+
+    def unnorm_pred(self, x):
+        dev = x.device
+        x = x.detach().cpu().numpy()
+
+        x = (
+            x
+            * self.stds[np.newaxis, ...].transpose(0, 2, 3, 1)[
+                ..., self.var_start : self.var_end
+            ]
+            + self.means[np.newaxis, ...].transpose(0, 2, 3, 1)[
+                ..., self.var_start : self.var_end
+            ]
+        )
+        if bool(self.diff):
+            x = (
+                x
+                + self.era5_mean_spatial[np.newaxis, ...].transpose(0, 3, 2, 1)[
+                    ..., self.var_start : self.var_end
+                ]
+            )
+        return torch.from_numpy(x).float().to(dev)
+
+    def get_index(self, index, prefix):
+        """
+        Load data for the relevant index respecting different offsets depending on the modality
+        """
+
+        index = self.index[index]
+        date = self.dates[index]
+
+        # ICOADS
+        icoads_x = self.icoads_x[index + self.icoads_index_offset, ...]
+        icoads_y = self.icoads_y[index + self.icoads_index_offset, ...]
+        icoads_x = [icoads_x[0, :], icoads_x[1, :]]
+        icoads_x = [self.to_tensor(i) for i in icoads_x]
+        icoads_y = self.to_tensor(icoads_y)
+        icoads_y = self.norm_data(icoads_y, self.icoads_means, self.icoads_stds)
+
+        # GRIDSAT
+        sat_y = self.sat_y[index + self.sat_index_offset, ...]
+        sat_x = [self.to_tensor(i) for i in self.sat_x]
+        sat_y = self.to_tensor(sat_y)
+        sat_y = self.norm_data(sat_y, self.sat_means, self.sat_stds)
+
+        # AMSU-A
+        amsua_y = self.to_tensor(self.amsua_y[index + self.amsua_index_offset, ...])
+        amsua_y[amsua_y < -998] = torch.nan
+        amsua_x = [self.to_tensor(i) for i in self.amsua_x]
+        amsua_y[amsua_y < -998] = np.nan
+        amsua_y = self.norm_data(amsua_y, self.amsua_means, self.amsua_stds)
+
+        # AMSU-B
+        amsub_y = self.to_tensor(self.amsub_y[index + self.amsub_index_offset, ...])
+        amsub_y[amsub_y < -998] = torch.nan
+        amsub_x = [self.to_tensor(i) for i in self.amsub_x]
+        amsub_y[amsub_y < -998] = np.nan
+        amsub_y = self.norm_data(amsub_y, self.amsub_means, self.amsub_stds)
+
+        # IASI
+        iasi_y = self.to_tensor(self.iasi[index + self.iasi_index_offset, ...])
+        iasi_x = [self.to_tensor(i) for i in self.iasi_x]
+        iasi_y = self.norm_data(iasi_y, self.iasi_means, self.iasi_stds)
+
+        # IGRA
+        igra_y = self.to_tensor(self.igra_y[index + self.igra_index_offset, ...])
+        igra_x = [self.igra_x[:, 0], self.igra_x[:, 1]]
+        igra_x = [self.to_tensor(i) for i in igra_x]
+        igra_y = self.norm_data(igra_y, self.igra_means, self.igra_stds)
+
+        # ASCAT
+        ascat_y = self.to_tensor(self.ascat_y[index + self.ascat_index_offset, ...])
+        ascat_x = [self.to_tensor(i) for i in self.ascat_x]
+        ascat_y[..., 4][ascat_y[..., 4] < -990] = np.nan
+        ascat_y = self.norm_data(ascat_y, self.ascat_means, self.ascat_stds)
+
+        # HIRS
+        hirs_y = self.to_tensor(self.hirs_y[index + self.hirs_index_offset, ...])
+        hirs_y[hirs_y < -998] = np.nan
+        hirs_x = [self.to_tensor(i) for i in self.hirs_x]
+        hirs_y = self.norm_data(hirs_y, self.hirs_means, self.hirs_stds)
+
+        # HadISD
+        x_context_hadisd = self.hadisd_x
+        y_context_hadisd = [
+            i[index + self.hadisd_index_offset, :] for i in self.hadisd_y
+        ]
+        x_context_hadisd = [self.to_tensor(i).permute(1, 0) for i in x_context_hadisd]
+        y_context_hadisd = [self.to_tensor(i) for i in y_context_hadisd]
+        y_context_hadisd = self.norm_hadisd(y_context_hadisd)
+
+        # ERA5
+        era5 = self.to_tensor(self.load_era5_time(index))
+        era5_target = era5.permute(2, 1, 0)
+        era5_x = self.era5_x
+
+        # AUxiliary variables
+        aux_time = self.to_tensor(self.get_time_aux(date))
+        climatology = self.climatology[date.hour // 6, date.dayofyear - 1, ...]
+
+        task = {
+            "x_context_hadisd_{}".format(prefix): x_context_hadisd,
+            "y_context_hadisd_{}".format(prefix): y_context_hadisd,
+            "climatology_{}".format(prefix): self.to_tensor(climatology),
+            "sat_x_{}".format(prefix): sat_x,
+            "sat_{}".format(prefix): sat_y,
+            "icoads_x_{}".format(prefix): icoads_x,
+            "icoads_{}".format(prefix): icoads_y,
+            "igra_x_{}".format(prefix): igra_x,
+            "igra_{}".format(prefix): igra_y,
+            "amsua_{}".format(prefix): amsua_y,
+            "amsua_x_{}".format(prefix): amsua_x,
+            "amsub_{}".format(prefix): amsub_y,
+            "amsub_x_{}".format(prefix): amsub_x,
+            "iasi_{}".format(prefix): iasi_y,
+            "iasi_x_{}".format(prefix): iasi_x,
+            "ascat_{}".format(prefix): ascat_y,
+            "ascat_x_{}".format(prefix): ascat_x,
+            "hirs_{}".format(prefix): hirs_y,
+            "hirs_x_{}".format(prefix): hirs_x,
+            "y_target_{}".format(prefix): era5_target[
+                ..., self.var_start : self.var_end
+            ],
+            "era5_x_{}".format(prefix): era5_x,
+            "era5_elev_{}".format(prefix): self.era5_elev,
+            "era5_lonlat_{}".format(prefix): self.era5_lonlat,
+            "aux_time_{}".format(prefix): aux_time,
+            "lt": torch.Tensor([self.var_start]),
+        }
+
+        return task
+
+
+class HadISDDataset(Dataset):
+    """
+    HadISD dataset for decoder training
+    """
+
+    def __init__(self, var, mode, device, start_date, end_date):
+        super().__init__()
+
+        # Setup
+        if not mode in ["train", "val", "test"]:
+            raise Exception(f"mode is {mode}. Must be train, val, or test.")
+
+        self.var = var
+        self.mode = mode
+        self.start_date = start_date
+        self.device = device
+        dates = pd.date_range(start_date, end_date, freq="6H")
+        self.index = np.array(range(len(dates)))
+
+        # Load the hadISD data
+        self.load_hadisd()
+
+    def load_hadisd(self):
+        """
+        Load the raw HadISD data
+        """
+
+        data_path = "path_to_data/"
+        aux_data_path = "path_to_auxiliary_data/"
+        var = self.var
+        mode = self.mode
+
+        vals = np.memmap(
+            data_path + f"hadisd_processed/{var}_vals_{mode}.memmap",
+            dtype="float32",
+            mode="r",
+            shape=get_hadisd_shape(mode),
+        )
+
+        lon = lon_to_0_360(
+            np.load(data_path + f"hadisd_processed/{var}_lon_{mode}.npy")
+        )
+        lat = np.load(data_path + f"hadisd_processed/{var}_lat_{mode}.npy")
+        self.hadisd_x = np.stack([lon, lat], axis=-1) / LATLON_SCALE_FACTOR
+        self.hadisd_alt = np.load(
+            data_path + f"hadisd_processed/{var}_alt_{mode}_final.npy"
+        )
+        self.hadisd_y = vals
+
+        self.hadisd_index_offset = HADISD_OFFSETS[self.start_date]
+        self.hadisd_means = self.to_tensor(
+            np.load(aux_data_path + f"norm_factors/mean_hadisd_{var}.npy")
+        )
+        self.hadisd_stds = self.to_tensor(
+            np.load(aux_data_path + f"norm_factors/std_hadisd_{var}.npy")
+        )
+        return
+
+    def norm_hadisd(self, x):
+        return (x - self.hadisd_means) / self.hadisd_stds
+
+    def unnorm_pred(self, x):
+        return self.hadisd_means + self.hadisd_stds * x
+
+    def __len__(self):
+        return self.index.shape[0] - 2
+
+    def to_tensor(self, arr):
+        return torch.from_numpy(np.array(arr)).float().to(self.device)
+
+    def __getitem__(self, index):
+        index = self.index[index]
+
+        # Get longitude-latitude locations
+        x_target = self.to_tensor(self.hadisd_x).permute(1, 0)
+
+        # Get altitude and normalise
+        m_alt = np.expand_dims(np.load("path_to_mean_alt.npy"), 1)
+        s_alt = np.expand_dims(np.load("path_to_std_alt.npy"), 1)
+        alt_target = self.to_tensor((self.hadisd_alt - m_alt) / s_alt)[:, :]
+
+        # Get observations
+        y_target = self.norm_hadisd(
+            self.to_tensor(self.hadisd_y[index + self.hadisd_index_offset, :])
+        )
+
+        assert x_target.shape[0] == 2
+        n_stations = x_target.shape[1]
+        assert alt_target.shape[1] == n_stations
+        assert y_target.shape[0] == n_stations
+
+        return {"x": x_target, "altitude": alt_target, "y": y_target}
+
+
+class AardvarkICDataset(Dataset):
+    """
+    Helper dataset to handle initial condition loading for decoder training
+    """
+
+    def __init__(self, device, start_date, end_date, lead_time=0):
+        super().__init__()
+
+        # Setup
+
+        if lead_time == 0:
+            # If leadtime is 0 load the output of the encoder
+            if start_date == "2007-01-02" and end_date == "2017-12-31":
+                ic_fname = "ic_train.mmap"
+            elif start_date == "2019-01-01" and end_date == "2019-12-01":
+                ic_fname = "ic_val.mmap"
+            elif start_date == "2018-01-01" and end_date == "2018-12-31":
+                ic_fname = "ic_test.mmap"
+            else:
+                print((start_date, end_date))
+                raise Exception("Invalid start and end date")
+
+            dates = pd.date_range(start_date, end_date, freq="6H")
+
+            self.data = np.memmap(
+                "path_to_encoder_predictions/" + ic_fname,
+                dtype="float32",
+                mode="r",
+                shape=(len(dates), 121, 240, 24),  # shape of the output
+            )
+        else:
+            # if leadtime >0 load the forecast generated from the encoder prediction
+            if start_date == "2007-01-02" and end_date == "2017-12-31":
+                ic_fname = f"ic_train_{lead_time}.mmap"
+
+            elif start_date == "2019-01-01" and end_date == "2019-12-01":
+                ic_fname = f"ic_val_{lead_time}.mmap"
+            elif start_date == "2018-01-01" and end_date == "2018-12-31":
+                ic_fname = f"ic_test_{lead_time}.mmap"
+            else:
+                print((start_date, end_date))
+                raise Exception("Invalid start and end date.")
+
+            dates = pd.date_range(start_date, end_date, freq="6H")[(lead_time) * 4 :]
+            ic_shape = (len(dates), 121, 240, 24)
+
+            self.data = np.memmap(
+                self.data_path + "forecast_finetune/" + ic_fname,
+                dtype="float32",
+                mode="r",
+                shape=ic_shape,
+            )
+
+        self.device = device
+
+        # Normalisation
+        aux_data_path = "path_to_auxiliary_data/"
+        mean_factors_path = aux_data_path + f"norm_factors/mean_4u_1.npy"
+        std_factors_path = aux_data_path + f"norm_factors/std_4u_1.npy"
+        self.means = np.load(mean_factors_path)[:, np.newaxis, np.newaxis, ...]
+        self.stds = np.load(std_factors_path)[:, np.newaxis, np.newaxis, ...]
+
+    def __getitem__(self, index):
+        # Load Aardvark prediction and normalise
+        data_raw = np.transpose(np.copy(self.data[index, :, :, :]), (2, 1, 0))
+        data = (data_raw - self.means) / self.stds
+        return torch.from_numpy(data).to(self.device)
+
+
+class WeatherDatasetDownscaling(Dataset):
+    """
+    Main decoder training dataset. Uses AardvarkICDataset and HadISDDataset to
+    handle processor output and station data
+    """
+
+    def __init__(
+        self,
+        device,
+        hadisd_mode,
+        start_date,
+        end_date,
+        context_mode,
+        era5_mode="sfc",
+        res=1,
+        hadisd_var="tas",
+        lead_time=1,
+    ):
+        # The context mode determines whether we make use of ERA5 or our own ICs.
+        if not context_mode in ["era5", "aardvark"]:
+            raise Exception(
+                f"context_mode must be era5 or aardvark, got {context_mode}"
+            )
+
+        super().__init__()
+
+        # Setup
+        self.lead_time = lead_time
+
+        self.device = device
+        self.data_path = "path_to_data/"
+        self.aux_data_path = "path_to_auxiliary_data/"
+        self.start_date = start_date
+        self.end_date = end_date
+        self.era5_mode = era5_mode
+        self.res = res
+        self.context_mode = context_mode
+
+        self.dates = pd.date_range(start_date, end_date, freq="6H")
+        self.index = np.array(range(len(self.dates)))
+
+        # Load ERA5 data for pre-training
+        self.era5_sfc = [
+            self.load_era5(year)
+            for year in range(int(start_date[:4]), int(end_date[:4]) + 1)
+        ]
+
+        raw_era5_lon = np.load(self.data_path + f"era5/era5_x_{res}.npy")
+        raw_era5_lat = np.load(self.data_path + f"era5/era5_y_{res}.npy")
+        self.era5_x = [
+            self.to_tensor(raw_era5_lon) / LATLON_SCALE_FACTOR,
+            self.to_tensor(raw_era5_lat) / LATLON_SCALE_FACTOR,
+        ]
+
+        # Load orography
+        elev_path = self.data_path + f"era5/elev_vars_{res}.npy"
+        self.era5_elev = self.to_tensor(np.load(elev_path)).permute(0, 2, 1)
+
+        # Normalisation
+        mean_factors_path = (
+            self.aux_data_path + f"norm_factors/mean_{era5_mode}_{res}.npy"
+        )
+        std_factors_path = (
+            self.aux_data_path + f"norm_factors/std_{era5_mode}_{res}.npy"
+        )
+        self.means = np.load(mean_factors_path)[:, np.newaxis, np.newaxis, ...]
+        self.stds = np.load(std_factors_path)[:, np.newaxis, np.newaxis, ...]
+
+        # HadISD data
+        self.hadisd_data = HadISDDataset(
+            var=hadisd_var,
+            mode=hadisd_mode,
+            device=device,
+            start_date=start_date,
+            end_date=end_date,
+        )
+
+        if context_mode == "aardvark":
+            # Load the Aardvark encoder predictions
+            self.aardvark_data = AardvarkICDataset(
+                device, start_date, end_date, lead_time
+            )
+
+    def load_era5(self, year):
+        """
+        Load the raw ERA5 data
+        """
+
+        if year % 4 == 0:
+            d = 366 * 4
+        else:
+            d = 365 * 4
+
+        if self.era5_mode == "sfc":
+            levels = 4
+        elif self.era5_mode == "13u":
+            levels = 69
+        else:
+            levels = 24
+
+        if self.res == 1:
+            x = 240
+            y = 121
+        elif self.res == 5:
+            x = 64
+            y = 32
+        mmap = np.memmap(
+            self.data_path
+            + "era5/era5_{}_{}_6_{}.memmap".format(self.era5_mode, self.res, year),
+            dtype="float32",
+            mode="r",
+            shape=(d, levels, x, y),
+        )
+        return mmap
+
+    def norm_era5(self, x):
+        x = (x - self.means) / self.stds
+        return x
+
+    def unnorm_era5(self, x):
+        x = x * self.stds + self.means
+        return x
+
+    def unnorm_pred(self, x):
+        return self.hadisd_data.unnorm_pred(x)
+
+    def norm_data(self, x, means, stds):
+        return (x - means) / stds
+
+    def __len__(self):
+        return self.index.shape[0] - (self.lead_time) * 4
+
+    def to_tensor(self, arr):
+        return torch.from_numpy(np.array(arr)).float().to(self.device)
+
+    def get_time_aux(self, current_date):
+        """
+        Get auxiliary time variables for a given date
+        """
+
+        doy = current_date.dayofyear
+        year = (current_date.year - 2007) / 15
+        time_of_day = current_date.hour
+        return np.array(
+            [
+                np.cos(np.pi * 2 * doy / DAYS_IN_YEAR),
+                np.sin(np.pi * 2 * doy / DAYS_IN_YEAR),
+                np.cos(np.pi * 2 * time_of_day / 24),
+                np.sin(np.pi * 2 * time_of_day / 24),
+                year,
+            ]
+        )
+
+    def load_era5_time(self, index):
+        """
+        Load ERA5 training data
+        """
+
+        date = self.dates[index]
+        year = date.year
+        hour = date.hour
+        doy = (date.dayofyear - 1) * 4 + (hour // 6)
+
+        era5 = self.era5_sfc[year - int(self.start_date[:4])][doy, ...]
+        era5 = np.copy(era5)
+        era5 = self.norm_era5(era5[np.newaxis, ...])[0, ...]
+        return era5
+
+    def load_year_end(self, year, doy):
+        data_1 = self.era5_sfc[year - int(self.start_date[:4])][doy : doy + 1, ...]
+        missing = self.lead_time - data_1.shape[0] + 1
+        data_2 = self.era5_sfc[year - int(self.start_date[:4]) + 1][:missing, ...]
+        data = np.concatenate([data_1, data_2])
+        return data
+
+    def __getitem__(self, index):
+
+        index = self.index[index]
+        date = self.dates[index + 4 * self.lead_time]
+
+        # Get HadISD data
+        hadisd_slice = self.hadisd_data[index + 4 * self.lead_time]
+
+        # Get lon-lat
+        x_context = self.era5_x
+        n_lon = x_context[0].shape[0]
+        n_lat = x_context[1].shape[0]
+
+        # Get auxiliary time
+        aux_time = torch.reshape(self.to_tensor(self.get_time_aux(date)), (-1, 1, 1))
+
+        # Load the context (either aardvark or ERA5 for use in pre-training)
+        if self.context_mode == "era5":
+            y_context_obs = self.to_tensor(
+                self.load_era5_time(index + 4 * self.lead_time)
+            )
+
+        elif self.context_mode == "aardvark":
+            y_context_obs = self.aardvark_data[index]
+
+        else:
+            raise Exception
+
+        y_context = torch.cat(
+            [
+                y_context_obs,
+                self.era5_elev.permute(0, 2, 1),
+                aux_time.repeat(1, n_lon, n_lat),
+            ]
+        )
+
+        assert y_context.shape[1] == n_lon
+        assert y_context.shape[2] == n_lat
+
+        x = hadisd_slice["x"]
+        alt = hadisd_slice["altitude"]
+        y = hadisd_slice["y"]
+
+        return {
+            "x_target": x,
+            "alt_target": alt,
+            "y_target": y,
+            "y_context": y_context,
+            "x_context": x_context,
+            "aux_time": aux_time,
+            "lt": torch.Tensor([0]),
+        }
+
+
+class ForecasterDatasetDownscaling(Dataset):
+    """
+    Dataset to generate decoder predictions from pre-saved Aardvark forecasts
+    """
+
+    def __init__(
+        self,
+        start_date,
+        end_date,
+        lead_time,
+        hadisd_var,
+        mode,
+        device,
+        forecast_path,
+        region="global",
+    ):
+        super().__init__()
+
+        # Setup
+
+        if not mode in ["train", "val", "test"]:
+            raise Exception(f"Mode is {mode}. Must be either train, val, or test")
+
+        self.device = device
+        self.start_date = start_date
+        self.end_date = end_date
+        self.lead_time = lead_time
+        self.mode = mode
+        self.offset = np.timedelta64(lead_time, "D").astype("timedelta64[ns]")
+
+        self.dates = pd.date_range(start_date, end_date, freq="6H")[:-30]
+
+        # Normalisation
+        aux_data_path = "auxiliary_data_path/"
+        self.means = np.load(aux_data_path + "norm_factors/mean_4u_1.npy")
+        self.stds = np.load(aux_data_path + "norm_factors/std_4u_1.npy")
+
+        # Load auxiliary data
+        self.load_npy_file()
+        data_path = "data_path/"
+        res = "1"
+        raw_era5_lon = np.load(data_path + f"era5/era5_x_{res}.npy")
+        raw_era5_lat = np.load(data_path + f"era5/era5_y_{res}.npy")
+        self.era5_x = [
+            self.to_tensor(raw_era5_lon) / LATLON_SCALE_FACTOR,
+            self.to_tensor(raw_era5_lat) / LATLON_SCALE_FACTOR,
+        ]
+        elev_path = data_path + f"era5/elev_vars_{res}.npy"
+        self.era5_elev = self.to_tensor(np.load(elev_path)).permute(0, 2, 1)
+
+        # Load hadISD
+        self.hadisd_data = HadISDDataset(
+            var=hadisd_var,
+            mode="train",
+            device=device,
+            start_date=start_date,
+            end_date=end_date,
+        )
+
+        # Subset to region
+        self.region = region
+        if self.region != "global":
+            self.mask = np.load(
+                self.data_path + f"hadisd_processed/tas_mask_train_{region}.npy"
+            )
+
+    def date_range(self):
+        return np.arange(
+            start=np.datetime64(self.start_date).astype("datetime64[ns]"),
+            stop=np.datetime64(self.end_date).astype("datetime64[ns]"),
+            step=np.timedelta64(1, "D").astype("timedelta64[ns]"),
+        )
+
+    def load_npy_file(self):
+        """
+        Load the pre-saved Aardvark forecasts
+        """
+
+        dates = pd.date_range(self.start_date, self.end_date, freq="6H")
+
+        if self.mode == "train":
+            dates = dates[:-40]  # Need 10 day offset at end of year
+
+        self.Y_context = np.memmap(
+            "path_to_forecasts/forecast_{}.mmap".format(self.mode),
+            dtype="float32",
+            mode="r",
+            shape=(len(dates), 121, 240, 24, 11),
+        )
+
+        return
+
+    def norm_era5(self, x):
+        return (x - self.means) / self.stds
+
+    def norm_hadisd(self, x):
+        return self.hadisd_data.norm_hadisd(x)
+
+    def unnorm_pred(self, x):
+        return self.hadisd_data.unnorm_pred(x)
+
+    def __len__(self):
+        return len(self.dates) - 40  # Need 10 day offset at end of year
+
+    def to_tensor(self, arr):
+        return torch.from_numpy(np.array(arr)).float().to(self.device)
+
+    def get_time_aux(self, index):
+        """
+        Get the auxiliary time variables
+        """
+
+        current_date = (self.dates + self.offset)[index]
+        doy = current_date.dayofyear
+        year = (current_date.year - 2007) / 15
+        time_of_day = current_date.hour
+        return np.array(
+            [
+                np.cos(np.pi * 2 * doy / DAYS_IN_YEAR),
+                np.sin(np.pi * 2 * doy / DAYS_IN_YEAR),
+                np.cos(np.pi * 2 * time_of_day / 24),
+                np.sin(np.pi * 2 * time_of_day / 24),
+                year,
+            ]
+        )
+
+    def __getitem__(self, index):
+
+        # Load target data
+        hadisd_slice = self.hadisd_data[index + 4 * self.lead_time]
+
+        x_context = self.era5_x
+        n_lon = x_context[0].shape[0]
+        n_lat = x_context[1].shape[0]
+
+        # Load auxiliary time
+        aux_time = torch.reshape(self.to_tensor(self.get_time_aux(index)), (-1, 1, 1))
+
+        # Load input
+        y_context = self.norm_era5(self.Y_context[index, ..., self.lead_time])
+        y_context = torch.cat(
+            [
+                self.to_tensor(y_context).permute(2, 1, 0),
+                self.era5_elev.permute(0, 2, 1),
+                aux_time.repeat(1, n_lon, n_lat),
+            ]
+        )
+
+        assert y_context.shape[1] == n_lon
+        assert y_context.shape[2] == n_lat
+
+        # Handle region masking
+        if self.region != "global":
+            hadisd_slice["y"][self.mask] = np.nan
+
+        return {
+            "x_target": hadisd_slice["x"],
+            "alt_target": hadisd_slice["altitude"],
+            "y_target": hadisd_slice["y"],
+            "y_context": y_context,
+            "x_context": x_context,
+            "aux_time": aux_time,
+            "lt": torch.Tensor([0]),
+        }
+
+
+class ForecastLoader(Dataset):
+    """
+    Loader for finetuning the processor module
+    """
+
+    def __init__(
+        self,
+        device,
+        mode,
+        lead_time,
+        era5_mode="sfc",
+        res=5,
+        frequency=24,
+        norm=True,
+        diff=False,
+        rollout=False,
+        random_lt=False,
+        u_only=False,
+        ic_path=None,
+        finetune_step=None,
+        finetune_eval_every=100,
+        eval_steps=False,
+    ):
+
+        super().__init__()
+
+        # Setup
+        self.device = device
+        self.mode = mode
+        self.data_path = "data_path/"
+
+        self.lead_time = lead_time
+        self.era5_mode = era5_mode
+        self.res = res
+        self.frequency = frequency
+        self.norm = norm
+        self.diff = diff
+        self.rollout = rollout
+        self.random_lt = random_lt
+        self.u_only = u_only
+        self.ic_path = ic_path
+
+        self.finetune_step = finetune_step
+        self.finetune_eval_every = finetune_eval_every
+        self.eval_steps = eval_steps
+
+        if self.frequency == 6:
+            self.lead_time = self.lead_time * 4
+            freq = "6H"
+
+        else:
+            freq = "1D"
+
+        if self.mode == "train":
+            self.dates = pd.date_range("1979-01-01", "2017-12-31", freq=freq)
+        elif self.mode == "tune":
+            self.dates = pd.date_range("2018-01-01", "2018-12-31", freq=freq)
+        elif self.mode == "test":
+            self.dates = pd.date_range("2018-01-01", "2018-12-31", freq=freq)
+        elif self.mode == "val":
+            self.dates = pd.date_range("2019-01-01", "2019-12-31", freq=freq)
+
+        # Load the predictions from the previous leadtime to be the new context set
+        if self.finetune_step is not None:
+
+            if self.mode == "train":
+                self.dates = pd.date_range("2007-01-02", "2017-12-31", freq=freq)
+                ic_shape = (
+                    len(self.dates) - max(0, (self.finetune_step - 1) * 4),
+                    121,
+                    240,
+                    24,
+                )
+            elif self.mode == "val":
+                self.dates = pd.date_range("2019-01-01", "2019-12-31", freq=freq)
+                ic_shape = (
+                    len(self.dates) - max(0, (self.finetune_step - 1) * 4),
+                    121,
+                    240,
+                    24,
+                )
+            elif self.mode == "test":
+                self.dates = pd.date_range("2018-01-01", "2018-12-31", freq=freq)
+                ic_shape = (
+                    len(self.dates) - max(0, (self.finetune_step - 1) * 4),
+                    121,
+                    240,
+                    24,
+                )
+
+            if self.finetune_step > 1:
+                print(ic_shape)
+                self.ic = np.memmap(
+                    self.ic_path
+                    + "ic_{}_{}.mmap".format(self.mode, self.finetune_step - 1),
+                    dtype="float32",
+                    mode="r",
+                    shape=ic_shape,
+                )
+            elif self.ic_path is not None:
+
+                self.ic = np.memmap(
+                    self.ic_path + "ic_{}.mmap".format(self.mode),
+                    dtype="float32",
+                    mode="r",
+                    shape=ic_shape,
+                )
+
+        elif self.ic_path is not None:
+            if self.mode == "train":
+                self.dates = pd.date_range("2007-01-02", "2017-12-31", freq=freq)
+            ic_shape = (len(self.dates), 121, 240, 24)
+
+            self.ic = np.memmap(
+                self.ic_path + "/ic_{}.mmap".format(self.mode),
+                dtype="float32",
+                mode="r",
+                shape=ic_shape,
+            )
+
+        # Orography
+        self.era5_elev = np.float32(
+            np.load(self.data_path + "era5/elev_vars_{}.npy".format(res))
+        )
+        elev_mean = self.era5_elev.mean(axis=(1, 2))[:, np.newaxis, np.newaxis]
+        elev_std = self.era5_elev.std(axis=(1, 2))[:, np.newaxis, np.newaxis]
+        self.era5_elev = (self.era5_elev - elev_mean) / elev_std
+
+        # ERA5 ground truth data for training
+        self.era5_sfc = [
+            self.load_era5(year)
+            for year in range(int(self.dates[0].year), int(self.dates[-1].year) + 1)
+        ]
+
+        # Noramalisation factors
+        self.means = (
+            self.to_tensor(
+                np.load(
+                    self.data_path
+                    + "norm_factors/mean_{}_{}.npy".format(self.era5_mode, self.res)
+                )
+            )
+            .unsqueeze(1)
+            .unsqueeze(1)
+        )
+        self.stds = (
+            self.to_tensor(
+                np.load(
+                    self.data_path
+                    + "norm_factors/std_{}_{}.npy".format(self.era5_mode, self.res)
+                )
+            )
+            .unsqueeze(1)
+            .unsqueeze(1)
+        )
+        self.diff_means = (
+            self.to_tensor(
+                np.load(
+                    self.data_path
+                    + "norm_factors/mean_diff_{}_{}.npy".format(
+                        self.era5_mode, self.res
+                    )
+                )
+            )
+            .unsqueeze(0)
+            .unsqueeze(0)
+        )
+        self.diff_stds = (
+            self.to_tensor(
+                np.load(
+                    self.data_path
+                    + "norm_factors/std_diff_{}_{}.npy".format(self.era5_mode, self.res)
+                )
+            )
+            .unsqueeze(0)
+            .unsqueeze(0)
+        )
+
+        self.diff_means_1 = (
+            self.to_tensor(
+                np.load(
+                    self.data_path
+                    + "norm_factors/mean_diff_{}_{}_6h.npy".format(
+                        self.era5_mode, self.res
+                    )
+                )
+            )
+            .unsqueeze(0)
+            .unsqueeze(0)
+        )
+        self.diff_stds_1 = (
+            self.to_tensor(
+                np.load(
+                    self.data_path
+                    + "norm_factors/std_diff_{}_{}_6h.npy".format(
+                        self.era5_mode, self.res
+                    )
+                )
+            )
+            .unsqueeze(0)
+            .unsqueeze(0)
+        )
+
+        self.diff_means_2 = (
+            self.to_tensor(
+                np.load(
+                    self.data_path
+                    + "norm_factors/mean_diff_{}_{}_12h.npy".format(
+                        self.era5_mode, self.res
+                    )
+                )
+            )
+            .unsqueeze(0)
+            .unsqueeze(0)
+        )
+        self.diff_stds_2 = (
+            self.to_tensor(
+                np.load(
+                    self.data_path
+                    + "norm_factors/std_diff_{}_{}_12h.npy".format(
+                        self.era5_mode, self.res
+                    )
+                )
+            )
+            .unsqueeze(0)
+            .unsqueeze(0)
+        )
+
+        self.means_dict = {
+            0: self.diff_means,
+            2: self.diff_means_2,
+            3: self.diff_means_1,
+        }
+
+        self.stds_dict = {0: self.diff_stds, 2: self.diff_stds_2, 3: self.diff_stds_1}
+
+    def __len__(self):
+        if np.logical_and(self.eval_steps, self.mode == "train"):
+            return self.finetune_eval_every * 12 * 4
+
+        return self.dates.shape[0] - self.lead_time
+
+    def to_tensor(self, arr):
+
+        return torch.from_numpy(arr).float().to(self.device)
+
+    def norm_era5(self, x):
+        x = (x - self.means) / self.stds
+        return x
+
+    def norm_era5_tendency(self, x, lt_offset):
+
+        x = (x - self.means_dict[lt_offset]) / self.stds_dict[lt_offset]
+        return x
+
+    def unnorm_pred(self, x):
+        x = x * self.diff_stds.unsqueeze(0) + self.diff_means.unsqueeze(0)
+        return x
+
+    def unnorm_base_context(self, x):
+        x = x * self.stds.unsqueeze(0) + self.means.unsqueeze(0)
+        return x
+
+    def load_era5(self, year):
+        """
+        Load ERA5 data for training
+        """
+
+        if year % 4 == 0:
+            d = 366
+        else:
+            d = 365
+
+        if self.frequency == 6:
+            d = d * 4
+
+        if self.era5_mode == "sfc":
+            levels = 4
+        elif self.era5_mode == "13u":
+            levels = 69
+        else:
+            levels = 24
+
+        if self.res == 1:
+            x = 240
+            y = 121
+        elif self.res == 5:
+            x = 64
+            y = 32
+
+        mmap = np.memmap(
+            self.data_path
+            + "era5/era5_{}_{}_{}_{}.memmap".format(
+                self.era5_mode, self.res, self.frequency, year
+            ),
+            dtype="float32",
+            mode="r",
+            shape=(d, levels, x, y),
+        )
+        return mmap
+
+    def load_era5_time(self, index):
+        """
+        Load ERA5 data for training
+        """
+        date = self.dates[index]
+        year = date.year
+        doy = date.dayofyear - 1
+        hour = date.hour
+        if self.frequency == 6:
+            era5 = self.era5_sfc[year - int(self.dates[0].year)][
+                doy * 4 + hour // 6, ...
+            ]
+        else:
+            era5 = self.era5_sfc[year - int(self.dates[0].year)][doy, ...]
+
+        return np.copy(era5)
+
+    def make_time_channels(self, index, x, y):
+        """
+        Make auxiliary time channels
+        """
+
+        date = self.dates[index]
+        hour = date.hour
+        doy = date.dayofyear - 1
+        if date.year % 4 == 0:
+            n_days = 366
+        else:
+            n_days = 365
+        hour_sin = np.sin(hour * np.pi / 12) * np.float32(np.ones((1, x, y)))
+        hour_cos = np.cos(hour * np.pi / 12) * np.float32(np.ones((1, x, y)))
+        doy_sin = np.sin(doy * 2 * np.pi / n_days) * np.float32(np.ones((1, x, y)))
+        doy_cos = np.cos(doy * 2 * np.pi / n_days) * np.float32(np.ones((1, x, y)))
+
+        return np.concatenate([hour_sin, hour_cos, doy_sin, doy_cos])
+
+    def __getitem__(self, index):
+
+        # Option to offset to random leadtime
+        lt_offset = 0
+        if self.random_lt:
+            lt_offset = np.random.choice([0, 2, 3])
+
+        # Load ground truth data
+        y_target = self.to_tensor(
+            self.load_era5_time(index + self.lead_time - lt_offset)
+        )
+
+        # Load either initial condition or ERA5 depending on task
+        if self.ic_path is not None:
+            era5_ts0 = self.ic[index].copy().transpose(2, 1, 0)
+
+        else:
+            era5_ts0 = self.load_era5_time(index)
+
+        # Auxiliary time
+        time = self.make_time_channels(index, era5_ts0.shape[1], era5_ts0.shape[2])
+        era5_ts0 = self.to_tensor(
+            np.concatenate([era5_ts0, self.era5_elev, time], axis=0)
+        )
+        y_context = era5_ts0.permute(0, 2, 1)[:, ...]
+
+        # Normalisation
+        if self.diff:
+            y_target = (y_target - era5_ts0[:24, ...]).permute(2, 1, 0)
+            y_target = self.norm_era5_tendency(y_target, lt_offset)
+            y_context[:24, ...] = self.norm_era5(y_context[:24, ...])
+
+        else:
+            if self.norm:
+                y_context[:24, ...] = self.norm_era5(y_context[:24, ...], lt_offset)
+                y_target = self.norm_era5(y_target, lt_offset)
+            y_target = y_target.permute(2, 1, 0)
+
+        if self.rollout:
+            # Option to return entire timeseries of target data
+            targets = []
+            for t in range(self.lead_time + 1):
+                t = self.to_tensor(self.load_era5_time(index + t))
+                targets.append(t.permute(2, 1, 0))
+            targets = torch.stack(targets, dim=-1)[..., ::4]
+
+            return {
+                "y_context": y_context.permute(0, 2, 1),
+                "y_target": y_target,
+                "targets": targets,
+                "lt": self.to_tensor(np.array([lt_offset])),
+            }
+
+        else:
+            return {
+                "y_context": y_context.permute(0, 2, 1),
+                "y_target": y_target[..., :],
+                "lt": self.to_tensor(np.array([lt_offset])),
+                "target_index": self.to_tensor(np.array([index])),
+            }
+
+
+class WeatherDatasetE2E(WeatherDataset):
+    """
+    Dataset for running Aardvark end-to-end
+    """
+
+    def __init__(
+        self,
+        device,
+        hadisd_mode,
+        start_date,
+        end_date,
+        lead_time,
+        mode,
+        hadisd_var,
+        max_steps_per_epoch=None,
+        era5_mode="sfc",
+        res=1,
+        filter_dates=None,
+        var_start=0,
+        var_end=24,
+        diff=False,
+        two_frames=False,
+        region="global",
+    ):
+
+        super().__init__(
+            device,
+            hadisd_mode,
+            start_date,
+            end_date,
+            lead_time,
+            era5_mode,
+            res=res,
+            filter_dates=filter_dates,
+            diff=diff,
+        )
+
+        # Setup
+        self.var_start = var_start
+        self.var_end = var_end
+        self.diff = diff
+        self.two_frames = two_frames
+        self.region = region
+        self.lead_time = lead_time
+        self.mode = mode
+        self.max_steps_per_epoch = max_steps_per_epoch
+
+        # Initialise encoder dataset
+        self.assimilation_dataset = WeatherDatasetAssimilation(
+            device="cuda",
+            hadisd_mode="train",
+            start_date=start_date,
+            end_date=end_date,
+            lead_time=0,
+            era5_mode="4u",
+            res=1,
+            var_start=0,
+            var_end=24,
+            diff=False,
+            two_frames=False,
+        )
+
+        # Initialise forecast dataset
+        self.forecast_dataset = ForecastLoader(
+            device="cuda",
+            mode=mode,
+            lead_time=lead_time,
+            era5_mode=era5_mode,
+            res=1,
+            frequency=6,
+            diff=True,
+            u_only=False,
+            random_lt=False,
+        )
+
+        # Initialise downscaling dataset
+        self.downscaling_dataset = ForecasterDatasetDownscaling(
+            start_date=start_date,
+            end_date=end_date,
+            lead_time=lead_time,
+            hadisd_var=hadisd_var,
+            mode=mode,
+            device=device,
+            forecast_path=None,
+            region=region,
+        )
+
+    def __len__(self):
+        if self.max_steps_per_epoch:
+            return self.max_steps_per_epoch
+        return len(self.downscaling_dataset) - 40  # Need 10 day offset at end of year
+
+    def __getitem__(self, index):
+
+        if self.max_steps_per_epoch:
+            index = np.random.choice(
+                np.arange(len(self.downscaling_dataset) - 40)
+            )  # Need 10 day offset at end of year
+
+        # Get data for each of the three datasets
+        assimilation = self.assimilation_dataset.__getitem__(index)
+        forecast = self.forecast_dataset.__getitem__(index)
+        downscaling = self.downscaling_dataset.__getitem__(index)
+
+        # Create task
+        task = {
+            "assimilation": assimilation,
+            "forecast": forecast,
+            "downscaling": downscaling,
+            "index": torch.tensor(index),
+        }
+
+        # Add y target to allow for end to end finetuning if needed
+        task["y_target"] = task["downscaling"]["y_target"]
+
+        return task
+
+    def unnorm_pred(self, x):
+
+        dev = x.device
+        x = x.detach().cpu().numpy()
+
+        x = (
+            x
+            * self.stds[np.newaxis, ...].transpose(0, 2, 3, 1)[
+                ..., self.var_start : self.var_end
+            ]
+            + self.means[np.newaxis, ...].transpose(0, 2, 3, 1)[
+                ..., self.var_start : self.var_end
+            ]
+        )
+        if bool(self.diff):
+            x = (
+                x
+                + self.era5_mean_spatial[np.newaxis, ...].transpose(0, 3, 2, 1)[
+                    ..., self.var_start : self.var_end
+                ]
+            )
+        return torch.from_numpy(x).float().to(dev)

official-src/aardvark/loader_utils_new.py

@@ -0,0 +1,88 @@
+import numpy as np
+import torch
+import pandas as pd
+
+LATLON_SCALE_FACTOR = 360
+DAYS_IN_YEAR = 366
+
+DAILY_SCALE_FACTOR = {
+    "ERA5": 4,
+    "HADISD": 4,
+    "IR": 4,
+    "SOUNDER": 1,
+    "ICOADS": 1,
+}
+
+date_list = [
+    "1979-01-01",
+    "1999-01-01",
+    "1999-01-02",
+    "2002-01-02",
+    "2007-01-01",
+    "2007-01-02",
+    "2007-01-03",
+    "2013-01-02",
+    "2014-01-01",
+    "2017-01-02",
+    "2018-01-01",
+    "2019-01-01",
+    "2020-01-02",
+    "2020-01-01",
+    "2021-01-02",
+    "2021-01-01",
+]
+
+
+def generate_offsets(date_list, dates):
+    offsets = {}
+    for d in date_list:
+        try:
+            offsets[d] = np.where(dates == d)[0][0]
+        except:
+            offsets[d] = -1
+    return offsets
+
+
+IC_OFFSETS = generate_offsets(
+    date_list, pd.date_range("1999-01-02", "2021-12-31 18:00", freq="6H")
+)
+
+AMSUA_OFFSETS = generate_offsets(
+    date_list, pd.date_range("2007-01-01", "2021-12-31 18:00", freq="6H")
+)
+
+AMSUB_OFFSETS = generate_offsets(
+    date_list, pd.date_range("2007-01-01", "2021-12-31 18:00", freq="6H")
+)
+
+ASCAT_OFFSETS = generate_offsets(
+    date_list, pd.date_range("2007-01-01", "2021-12-31", freq="6H")
+)
+
+ATMS_OFFSETS = generate_offsets(
+    date_list, pd.date_range("2013-01-02", "2021-12-31", freq="1D")
+)
+
+ICOADS_OFFSETS = generate_offsets(
+    date_list, pd.date_range("1999-01-01 06:00", "2021-12-31", freq="6H")
+)
+
+IGRA_OFFSETS = generate_offsets(
+    date_list, pd.date_range("1999-01-01 00:00", "2021-12-31 18:00", freq="6H")
+)
+
+SAT_OFFSETS = generate_offsets(
+    date_list, pd.date_range("1990-01-01 00:00", "2021-12-31 18:00", freq="6H")
+)
+
+HADISD_OFFSETS = generate_offsets(
+    date_list, pd.date_range("1950-01-01 00:00", "2021-12-31 18:00", freq="6H")
+)
+
+
+def lon_to_0_360(x):
+    return (x + 360) % 360
+
+
+def lat_to_m90_90(x):
+    return torch.flip(x, [-1])

official-src/aardvark/loss_functions.py

@@ -0,0 +1,166 @@
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+class RmseLoss(nn.Module):
+    """
+    RMSE loss
+    """
+
+    def __init__(self, start_ind=0, end_ind=24):
+
+        super().__init__()
+        self.start_ind = start_ind
+        self.end_ind = end_ind
+
+    def forward(
+        self,
+        target,
+        output,
+        prev_step_output,
+        fix_sigma=False,
+        unwrap=False,
+        expand=False,
+    ):
+
+        squared_diff = ((target.to(output.device) - output) ** 2)[
+            ..., self.start_ind : self.end_ind
+        ]
+        return torch.mean(torch.sqrt(torch.nanmean(squared_diff, dim=(1, 2, 3))))
+
+
+class PressureWeightedRmseLoss(nn.Module):
+    """
+    Latitude weighted pressure weighted RMSE loss used in training the processor
+    """
+
+    def __init__(
+        self,
+        res,
+        era5_mode,
+        data_dir,
+        aux_data_dir,
+        weight_per_variable=False,
+    ):
+        super().__init__()
+
+        self.weights = torch.from_numpy(
+            np.load(aux_data_dir + "lat_weights/weights_lat_{}.npy".format(res)).T[
+                np.newaxis, ..., np.newaxis
+            ]
+        ).float()
+
+        self.weight_per_variable = weight_per_variable
+        self.variable_weights = torch.from_numpy(
+            np.load(aux_data_dir + "loss_weights.npy")[
+                np.newaxis, np.newaxis, np.newaxis, :
+            ]
+        ).float()
+
+        self.pressure_levels = (
+            torch.from_numpy(
+                np.load(
+                    data_dir + "era5/era5_pressure_levels_{}.npy".format(era5_mode)
+                )[np.newaxis, np.newaxis, np.newaxis, :]
+            ).float()
+            / 1000
+        )
+
+    def forward(
+        self,
+        target,
+        output,
+        prev_step_output,
+        fix_sigma=False,
+        unwrap=False,
+        expand=False,
+    ):
+
+        squared_diff = (target.to(output.device) - output) ** 2
+
+        if not expand:
+            weighted_sqared_diff = (
+                squared_diff
+                * self.weights.to(target.device)
+                * self.pressure_levels.to(target.device)
+            )
+            return torch.mean(
+                torch.sqrt(torch.nanmean(weighted_sqared_diff, dim=(1, 2, 3)))
+            )
+
+        weighted_sqared_diff = squared_diff * self.weights.to(target.device)
+
+        if self.weight_per_variable:
+            weighted_sqared_diff = weighted_sqared_diff * self.variable_weights.to(
+                weighted_sqared_diff.device
+            )
+
+        return torch.mean(
+            torch.sqrt(torch.nanmean(weighted_sqared_diff, dim=(1, 2))), dim=0
+        )
+
+
+class WeightedRmseLoss(nn.Module):
+    """
+    Latitude weighted RMSE loss
+    """
+
+    def __init__(
+        self,
+        res,
+        data_dir,
+        aux_data_dir,
+        weight_per_variable=False,
+        start_ind=0,
+        end_ind=24,
+    ):
+
+        super().__init__()
+        self.start_ind = start_ind
+        self.end_ind = end_ind
+
+        self.weights = torch.from_numpy(
+            np.load(aux_data_dir + "lat_weights/weights_lat_{}.npy".format(res)).T[
+                np.newaxis, ..., np.newaxis
+            ]
+        ).float()
+
+        self.weight_per_variable = weight_per_variable
+        self.variable_weights = torch.from_numpy(
+            np.load(aux_data_dir + "loss_weights.npy")[
+                np.newaxis, np.newaxis, np.newaxis, start_ind:end_ind
+            ]
+        ).float()
+
+    def forward(
+        self,
+        target,
+        output,
+        prev_step_output,
+        fix_sigma=False,
+        unwrap=False,
+        expand=False,
+    ):
+        squared_diff = (target.to(output.device) - output) ** 2
+
+        if not expand:
+            weighted_sqared_diff = squared_diff * self.weights.to(target.device)
+
+            if self.weight_per_variable:
+                weighted_sqared_diff = weighted_sqared_diff * self.variable_weights.to(
+                    weighted_sqared_diff.device
+                )
+
+            x = torch.nanmean(
+                torch.sqrt(torch.nanmean(weighted_sqared_diff, dim=(1, 2, 3)))
+            )
+
+            return x
+
+        weighted_sqared_diff = squared_diff * self.weights.to(target.device)
+        x = torch.mean(
+            torch.sqrt(torch.nanmean(weighted_sqared_diff, dim=(1, 2))), dim=0
+        )
+
+        return x

official-src/aardvark/misc_downscaling_functionality.py

@@ -0,0 +1,273 @@
+import pickle
+
+import torch
+import numpy as np
+import torch.nn as nn
+
+from set_convs import convDeepSet
+from unet_wrap_padding import Unet
+from vit import *
+from models import *
+
+hadisd_publisher_shifts = {
+    "tas": 273.15,
+    "u": 0.0,
+    "v": 0.0,
+    "psl": 0.0,
+    "ws": 0.0,
+}
+
+hadisd_publisher_scales = {
+    "tas": 10,
+    "u": 10,
+    "v": 10,
+    "psl": 100,
+    "ws": 10.0,
+}
+
+
+def hadisd_normalisation_factors(var: str):
+    path = "/home/azureuser/aux_data/norm_factors/"
+    return {
+        "mean": np.load(path + f"mean_hadisd_{var}.npy"),
+        "std": np.load(path + f"std_hadisd_{var}.npy"),
+    }
+
+
+def unnormalise_hadisd_var(x, var):
+    factors = hadisd_normalisation_factors(var)
+    hadisd_shift = hadisd_publisher_shifts[var]
+    hadisd_scale = hadisd_publisher_scales[var]
+
+    return hadisd_shift + hadisd_scale * (factors["mean"] + factors["std"] * x)
+
+
+class DownscalingRmseLoss(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, target, output, prev_step, fix_sigma=None, expand=True):
+
+        target = torch.flatten(target.to(output.device))
+        output = torch.flatten(output)
+
+        tmp = torch.isnan(target)
+        clean_target = target[~tmp]
+        clean_output = output[~tmp]
+
+        return torch.mean((clean_target - clean_output) ** 2)
+
+
+class ConvCNPWeatherOnToOff(nn.Module):
+    """
+    ConvCNP for decoder
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        int_channels,
+        device,
+        res,
+        data_path="../data/",
+        mode="end_to_end",
+        decoder=None,
+        film=False,
+    ):
+
+        super().__init__()
+
+        # Setup
+        self.device = device
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.int_channels = int_channels
+        self.decoder = decoder
+        self.int_x = 256
+        self.int_y = 128
+        self.mode = mode
+        self.film = film
+
+        # Load lon-lat of internal discretisation
+        self.era5_x = (
+            torch.from_numpy(
+                np.load(data_path + "grid_lon_lat/era5_x_{}.npy".format(res))
+            ).float()
+            / 360
+        )
+        self.era5_y = (
+            torch.from_numpy(
+                np.load(data_path + "grid_lon_lat/era5_y_{}.npy".format(res))
+            ).float()
+            / 360
+        )
+
+        # Setup setconv
+        self.sc_out = convDeepSet(
+            0.001, "OnToOff", density_channel=False, device=self.device
+        )
+
+        if self.mode not in ["downscaling", "end_to_end"]:
+            unet_out_channels = out_channels
+        else:
+            unet_out_channels = int_channels
+
+        # UNet backbone
+        if self.decoder == "base":
+            self.decoder_lr = Unet(
+                in_channels=in_channels,
+                out_channels=unet_out_channels,
+                div_factor=1,
+                film=film,
+            )
+
+        else:
+            raise Exception(f"Expected to use base decoder, but got {self.decoder}")
+
+        # Postprocessing MLP
+        self.mlp = DownscalingMLP(
+            in_channels=24 + 9,
+            out_channels=1,
+            h_channels=64,
+            h_layers=2,
+        )
+
+    def forward(self, task, film_index):
+
+        x = task["y_context"]
+        batch_size = x.shape[0]
+
+        # UNet backbone
+        x = self.decoder_lr(x, film_index=task["lt"])
+
+        # Transform to station predictions with setconv
+        num_channels = x.shape[3]
+        x = x.permute(0, 3, 1, 2)
+        assert list(x.shape) == [batch_size, num_channels, 240, 121]
+        x_target = task["x_target"]
+        num_stations = x_target.shape[2]
+
+        x = self.sc_out(
+            x_in=task["x_context"],
+            wt=x,
+            x_out=[x_target[:, 0, :], x_target[:, 1, :]],
+        )
+        assert x.shape[0] == batch_size
+        assert x.shape[2] == num_stations
+
+        # Concatenate auxiliary data at stations
+        alt_target = task["alt_target"]
+        assert torch.isnan(alt_target).sum() == 0
+        assert alt_target.shape[0] == batch_size
+        assert alt_target.shape[2] == num_stations
+
+        aux_time = task["aux_time"].squeeze(-1).repeat(1, 1, num_stations)
+        assert aux_time.shape[0] == batch_size
+        assert aux_time.shape[2] == num_stations
+
+        x = torch.cat([x, alt_target, x_target, aux_time], dim=1).permute(0, 2, 1)
+        assert x.shape[0] == batch_size
+        assert x.shape[1] == num_stations
+
+        tmp = self.mlp(x)
+        assert list(tmp.shape) == [batch_size, num_stations, 1]
+        y_hat = tmp.squeeze(-1)
+        assert list(y_hat.shape) == [batch_size, num_stations]
+        return y_hat
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, n_channels):
+        super().__init__()
+        self.block = nn.Sequential(nn.Linear(n_channels, n_channels), nn.ReLU())
+
+    def forward(self, x):
+        return self.block(x) + x
+
+
+class DownscalingMLP(nn.Module):
+    """
+    MLP for handling auxiliary data at station locations
+    """
+
+    def __init__(self, in_channels, out_channels, h_channels, h_layers):
+        super().__init__()
+
+        self.mlp = nn.Sequential(
+            nn.Linear(in_channels, h_channels),
+            *[ResidualBlock(h_channels) for _ in range(h_layers)],
+            nn.Linear(h_channels, out_channels),
+        )
+
+    def forward(self, x):
+        return self.mlp(x)
+
+
+def find_epoch(results_dir: str) -> int:
+    losses = np.load(results_dir + "losses_0.npy")
+    epoch = np.argmin(losses)
+
+    return epoch
+
+
+def load_config(results_dir: str) -> dict:
+    with open(results_dir + "config.pkl", "rb") as f:
+        return pickle.load(f)
+
+
+def load_model(results_dir, device, config_dir=None, epoch=None):
+    """
+    Load the ConvCNPWeatherOnToOff to generate forecasts
+    """
+
+    config_dir = config_dir or results_dir
+    config = load_config(config_dir)
+    epoch = epoch or find_epoch(results_dir)
+
+    model = ConvCNPWeatherOnToOff(
+        in_channels=config["in_channels"],
+        out_channels=config["out_channels"],
+        int_channels=config["int_channels"],
+        device=device,
+        res=config["res"],
+        decoder=config["decoder"],
+        mode=config["mode"],
+        film=False,
+    )
+
+    full_state_dict = torch.load(results_dir + f"epoch_{epoch}", map_location=device)
+    state_dict = full_state_dict["model_state_dict"]
+    state_dict = {k[7:]: v for k, v in zip(state_dict.keys(), state_dict.values())}
+    model.load_state_dict(state_dict)
+    model = model.to(device)
+    model.eval()
+
+    return model
+
+
+def load_model_e2e(results_dir, lead_time, device):
+    """
+    Load the full E2E model to generate forecasts
+    """
+
+    config = load_config(results_dir)
+    epoch = find_epoch(results_dir)
+
+    model = ConvCNPWeatherE2E(
+        device="cuda",
+        lead_time=lead_time,
+        se_model_path="../results/assimilation/all_v4_vit_ps3",
+        forecast_model_path="../results/forecast/vit_fnl_randomlt_mlp",
+        sf_model_path="../results/downscaling_NEW",
+    )
+
+    full_state_dict = torch.load(results_dir + f"epoch_{epoch}", map_location=device)
+    state_dict = full_state_dict["model_state_dict"]
+    state_dict = {k[7:]: v for k, v in zip(state_dict.keys(), state_dict.values())}
+    model.load_state_dict(state_dict)
+    model = model.to(device)
+    model.eval()
+
+    return model

official-src/aardvark/models.py

@@ -0,0 +1,390 @@
+import sys
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from architectures import MLP
+from set_convs import convDeepSet
+from unet_wrap_padding import *
+from vit import *
+
+sys.path.append("../")
+
+
+class ConvCNPWeather(nn.Module):
+    """
+    ConvCNP class used for the encoder and processor modules
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        int_channels,
+        device,
+        res,
+        data_path="../data/",
+        gnp=False,
+        mode="assimilation",
+        decoder=None,
+        film=False,
+        two_frames=False,
+    ):
+
+        super().__init__()
+
+        self.device = device
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.int_channels = int_channels
+        self.decoder = decoder
+        self.int_x = 256
+        self.int_y = 128
+        self.data_path = data_path
+        self.mode = mode
+        self.film = film
+        self.two_frames = two_frames
+
+        N_SAT_VARS = 2
+        N_ICOADS_VARS = 5
+        N_HADISD_VARS = 5
+
+        # Load internal grid longitude-latitude locations
+        self.era5_x = (
+            torch.from_numpy(
+                np.load(self.data_path + "grid_lon_lat/era5_x_{}.npy".format(res))
+            ).float()
+            / 360
+        )
+        self.era5_y = (
+            torch.from_numpy(
+                np.load(self.data_path + "grid_lon_lat/era5_y_{}.npy".format(res))
+            ).float()
+            / 360
+        )
+
+        self.int_grid = [
+            (torch.linspace(0, 360, 240) / 360).float().cuda(),
+            (torch.linspace(-90, 90, 121) / 360).float().cuda(),
+        ]
+
+        self.int_grid = [self.int_grid[0].unsqueeze(0), self.int_grid[1].unsqueeze(0)]
+
+        # Create input setconvs for each data modality
+        self.ascat_setconvs = convDeepSet(
+            0.001, "OnToOn", density_channel=True, device=self.device
+        )
+        self.amsua_setconvs = [
+            convDeepSet(0.001, "OnToOn", density_channel=True, device=self.device)
+            for _ in range(13)
+        ]
+        self.amsub_setconvs = [
+            convDeepSet(0.001, "OnToOn", density_channel=True, device=self.device)
+            for _ in range(12)
+        ]
+        self.hirs_setconvs = [
+            convDeepSet(0.001, "OnToOn", density_channel=True, device=self.device)
+            for _ in range(26)
+        ]
+
+        self.sat_setconvs = [
+            convDeepSet(0.001, "OnToOn", density_channel=True, device=self.device)
+            for _ in range(N_SAT_VARS)
+        ]
+        self.hadisd_setconvs = [
+            convDeepSet(0.001, "OffToOn", density_channel=True, device=self.device)
+            for _ in range(N_HADISD_VARS)
+        ]
+        self.icoads_setconvs = [
+            convDeepSet(0.001, "OffToOn", density_channel=True, device=self.device)
+            for _ in range(N_ICOADS_VARS)
+        ]
+        self.igra_setconvs = [
+            convDeepSet(0.001, "OffToOn", density_channel=True, device=self.device)
+            for _ in range(24)
+        ]
+
+        self.sc_out = convDeepSet(
+            0.001, "OnToOff", density_channel=False, device=self.device
+        )
+
+        # Instantiate the decoder. Here decoder refers to decoder in a convCNP (i.e the ViT backbone)
+        if self.decoder == "vit":
+            self.decoder_lr = ViT(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                h_channels=512,
+                depth=16,
+                patch_size=5,
+                per_var_embedding=True,
+                img_size=[240, 121],
+            )
+
+        elif self.decoder == "vit_assimilation":
+            self.decoder_lr = ViT(
+                in_channels=256,
+                out_channels=out_channels,
+                h_channels=512,
+                depth=8,
+                patch_size=3,
+                per_var_embedding=False,
+                img_size=[256, 128],
+            )
+
+        self.mlp = MLP(
+            in_channels=out_channels,
+            out_channels=out_channels,
+            h_channels=128,
+            h_layers=4,
+        )
+        self.break_next = False
+
+    def encoder_hadisd(self, task, prefix):
+        """
+        Data preprocessing for HadISD
+        """
+
+        encodings = []
+        for channel in range(4):
+            encodings.append(
+                self.hadisd_setconvs[channel](
+                    x_in=[
+                        task["x_context_hadisd_{}".format(prefix)][channel][:, 0, :],
+                        task["x_context_hadisd_{}".format(prefix)][channel][:, 1, :],
+                    ],
+                    wt=task["y_context_hadisd_{}".format(prefix)][channel].unsqueeze(1),
+                    x_out=self.int_grid,
+                )
+            )
+        encodings = torch.cat(encodings, dim=1)
+        return encodings
+
+    def encoder_sat(self, task, prefix):
+        """
+        Data preprocessing for Gridsat
+        """
+
+        encodings = []
+        for channel in range(task["sat_{}".format(prefix)].shape[1]):
+            encodings.append(
+                self.sat_setconvs[channel](
+                    x_in=task["sat_x_{}".format(prefix)],
+                    wt=task["sat_{}".format(prefix)][:, channel : channel + 1, ...],
+                    x_out=self.int_grid,
+                )
+            )
+        encodings = torch.cat(encodings, dim=1)
+        return encodings
+
+    def encoder_icoads(self, task, prefix):
+        """
+        Data preprocessing for ICOADS
+        """
+
+        encodings = []
+        for channel in range(5):
+            encodings.append(
+                self.icoads_setconvs[channel](
+                    x_in=task["icoads_x_{}".format(prefix)],
+                    wt=task["icoads_{}".format(prefix)][:, channel, :].unsqueeze(1),
+                    x_out=self.int_grid,
+                )
+            )
+        encodings = torch.cat(encodings, dim=1)
+
+        return encodings
+
+    def encoder_amsua(self, task, prefix):
+        """
+        Data preprocessing for AMSU-A
+        """
+
+        encodings = []
+        task["amsua_{}".format(prefix)][..., -1] = np.nan
+        task["amsua_{}".format(prefix)][task["amsua_{}".format(prefix)] == 0] = np.nan
+        for i in range(13):
+            encodings.append(
+                self.amsua_setconvs[i](
+                    x_in=task["amsua_x_{}".format(prefix)],
+                    wt=task["amsua_{}".format(prefix)].permute(0, 3, 2, 1)[
+                        :, i : i + 1, ...
+                    ],
+                    x_out=self.int_grid,
+                )
+            )
+
+        encodings = torch.cat(encodings, dim=1)
+        return encodings
+
+    def encoder_amsub(self, task, prefix):
+        """
+        Data preprocessing for AMSU-B
+        """
+
+        encodings = []
+        task["amsub_{}".format(prefix)][task["amsub_{}".format(prefix)] == 0] = np.nan
+        for i in range(12):
+            encodings.append(
+                self.amsua_setconvs[i](
+                    x_in=task["amsub_x_{}".format(prefix)],
+                    wt=task["amsub_{}".format(prefix)].permute(0, 3, 1, 2)[
+                        :, i : i + 1, ...
+                    ],
+                    x_out=self.int_grid,
+                )
+            )
+
+        encodings = torch.cat(encodings, dim=1)
+        return encodings
+
+    def encoder_hirs(self, task, prefix):
+        """
+        Data preprocessing for HIRS
+        """
+
+        encodings = []
+
+        task["hirs_{}".format(prefix)][task["hirs_{}".format(prefix)] == 0] = np.nan
+        for i in range(26):
+            encodings.append(
+                self.hirs_setconvs[i](
+                    x_in=task["hirs_x_{}".format(prefix)],
+                    wt=task["hirs_{}".format(prefix)].permute(0, 3, 1, 2)[
+                        :, i : i + 1, ...
+                    ],
+                    x_out=self.int_grid,
+                )
+            )
+
+        encodings = torch.cat(encodings, dim=1)
+        return encodings
+
+    def encoder_igra(self, task, prefix):
+        """
+        Data preprocessing for IGRA
+        """
+
+        encodings = []
+        for channel in range(24):
+            encodings.append(
+                self.igra_setconvs[channel](
+                    x_in=task["igra_x_{}".format(prefix)],
+                    wt=task["igra_{}".format(prefix)][:, channel, :].unsqueeze(1),
+                    x_out=self.int_grid,
+                )
+            )
+        encodings = torch.cat(encodings, dim=1)
+
+        return encodings
+
+    def encoder_ascat(self, task, prefix):
+        """
+        Data preprocessing for ASCAT
+        """
+
+        task["ascat_{}".format(prefix)][
+            torch.isnan(task["ascat_{}".format(prefix)])
+        ] = 0
+        e = nn.functional.interpolate(
+            task["ascat_{}".format(prefix)].permute(0, 3, 1, 2), size=(240, 121)
+        )
+        e = torch.flip(e, dims=[-1])
+        return e
+
+    def encoder_iasi(self, task, prefix):
+        """
+        Data preprocessing for IASI
+        """
+
+        task["iasi_{}".format(prefix)][torch.isnan(task["iasi_{}".format(prefix)])] = 0
+        e = nn.functional.interpolate(
+            task["iasi_{}".format(prefix)].permute(0, 3, 1, 2), size=(240, 121)
+        )
+        e = torch.flip(e, dims=[-1])
+        return e
+
+    def forward(self, task, film_index):
+
+        # Setup input
+        if self.mode == "assimilation":
+
+            self.int_grid = [i.to(task["y_target"].device) for i in self.int_grid]
+            elev = nn.functional.interpolate(
+                torch.flip(task["era5_elev_current"].permute(0, 1, 3, 2), dims=[2]),
+                size=(self.int_grid[0].shape[1], self.int_grid[1].shape[1]),
+            )
+            elev = torch.flip(task["era5_elev_current"].permute(0, 1, 3, 2), dims=[2])
+
+            if not self.two_frames:
+                encodings = [
+                    self.encoder_iasi(task, "current"),
+                    self.encoder_ascat(task, "current"),
+                    self.encoder_hadisd(task, "current"),
+                    self.encoder_icoads(task, "current"),
+                    self.encoder_sat(task, "current"),
+                    self.encoder_amsua(task, "current"),
+                    self.encoder_amsub(task, "current"),
+                    self.encoder_igra(task, "current"),
+                    self.encoder_hirs(task, "current"),
+                    elev,
+                    task["climatology_current"],
+                    torch.ones_like(elev[:, :5, ...])
+                    * task["aux_time_current"].unsqueeze(-1).unsqueeze(-1),
+                ]
+            else:
+                # Option to pass two timesteps (t=-1 and t=0) as input
+                encodings = [
+                    self.encoder_iasi(task, "current"),
+                    self.encoder_ascat(task, "current"),
+                    self.encoder_hadisd(task, "current"),
+                    self.encoder_icoads(task, "current"),
+                    self.encoder_sat(task, "current"),
+                    self.encoder_amsua(task, "current"),
+                    self.encoder_amsub(task, "current"),
+                    self.encoder_igra(task, "current"),
+                    self.encoder_hirs(task, "current"),
+                    self.encoder_iasi(task, "prev"),
+                    self.encoder_ascat(task, "prev"),
+                    self.encoder_hadisd(task, "prev"),
+                    self.encoder_icoads(task, "prev"),
+                    self.encoder_sat(task, "prev"),
+                    self.encoder_amsua(task, "prev"),
+                    self.encoder_amsub(task, "prev"),
+                    self.encoder_igra(task, "prev"),
+                    self.encoder_hirs(task, "prev"),
+                    elev,
+                    task["climatology_current"],
+                    torch.ones_like(elev[:, :5, ...])
+                    * task["aux_time_current"].unsqueeze(-1).unsqueeze(-1),
+                ]
+            x = torch.cat(encodings, dim=1)
+
+        else:
+            x = task["y_context"]
+
+        if x.shape[-1] > x.shape[-2]:
+            x = x.permute(0, 1, 3, 2)
+
+        # Run ViT backbone
+        if self.decoder == "vit":
+            x = self.decoder_lr(x, lead_times=task["lt"])
+            x = x.permute(0, 3, 1, 2)
+        else:
+            x = nn.functional.interpolate(x, size=(256, 128))
+            x = self.decoder_lr(x, film_index=(task["lt"] * 0) + 1)
+
+        # Process outputs
+
+        if np.logical_and(
+            self.mode == "assimilation", self.decoder == "vit_assimilation"
+        ):
+            x = nn.functional.interpolate(x.permute(0, 3, 1, 2), size=(240, 121))
+            return x.permute(0, 3, 2, 1)
+
+        elif self.mode == "forecast":
+            x = nn.functional.interpolate(x, size=(240, 121)).permute(0, 2, 3, 1)
+            return x.permute(0, 2, 1, 3)
+
+        return x

official-src/aardvark/set_convs.py

@@ -0,0 +1,101 @@
+import torch
+import torch.nn as nn
+
+
+class convDeepSet(nn.Module):
+    """
+    ConvDeepSet used to translate ungridded observations to a gridded representation and
+    vice versa.
+    """
+
+    def __init__(
+        self,
+        init_ls,
+        mode,
+        device,
+        density_channel=True,
+        step=0.25,
+        grid=False,
+    ):
+        super().__init__()
+        self.init_ls = torch.nn.Parameter(torch.tensor([init_ls]))
+        self.grid = grid
+        self.step = step
+        self.density_channel = density_channel
+        self.mode = mode
+
+        self.init_ls.requires_grad = True
+        self.device = device
+
+    def compute_weights(self, x1, x2):
+
+        dists2 = self.pw_dists2(x1.unsqueeze(-1), x2.unsqueeze(-1))
+
+        d = torch.exp((-0.5 * dists2) / (self.init_ls.to(x1.device)) ** 2)
+        return d
+
+    def pw_dists2(self, a, b):
+
+        norms_a = torch.sum(a**2, axis=-1)[..., :, None]
+        norms_b = torch.sum(b**2, axis=-1)[..., None, :]
+
+        return norms_a + norms_b - 2 * torch.matmul(a, b.permute(0, 2, 1))
+
+    def forward(self, x_in, wt, x_out):
+
+        # Add a density channel
+        density_channel = torch.ones_like(wt[:, 0:1, ...])
+        density_channel[torch.isnan(wt[:, 0:1, ...])] = 0
+
+        wt = torch.cat([density_channel, wt], dim=1)
+        wt[torch.isnan(wt)] = 0
+
+        if self.mode == "OffToOn":
+            # Case 1: converting off-the-grid data to a gridded representation
+
+            in_lon_mask = ~torch.isnan(x_in[0])
+            in_lat_mask = ~torch.isnan(x_in[1])
+
+            x_in[0][~in_lon_mask] = 0
+            x_in[1][~in_lat_mask] = 0
+
+            ws = [self.compute_weights(xzi, xi) for xzi, xi in zip(x_in, x_out)]
+
+            ws[0] = ws[0] * in_lon_mask.unsqueeze(-1).int()
+            ws[1] = ws[1] * in_lat_mask.unsqueeze(-1).int()
+
+            ee = torch.einsum("...cw,...wx,...wy->...cxy", wt, ws[0], ws[1])
+
+        elif self.mode == "OnToOn":
+            # Case 2: converting between gridded representations
+
+            ws = [self.compute_weights(xzi, xi) for xzi, xi in zip(x_in, x_out)]
+            ee = torch.einsum("...cwh,...wx,...hy->...cxy", wt, ws[0], ws[1])
+
+        elif self.mode == "OnToOff":
+
+            # Case 3: converting a gridded representation to off-the-grid predictions
+            out_lon_mask = ~torch.isnan(x_out[0])
+            out_lat_mask = ~torch.isnan(x_out[1])
+            x_out[0][~out_lon_mask] = 0
+            x_out[1][~out_lat_mask] = 0
+
+            ws = [self.compute_weights(xzi, xi) for xzi, xi in zip(x_in, x_out)]
+
+            ws[0] = ws[0] * out_lon_mask.unsqueeze(-2).int()
+            ws[1] = ws[1] * out_lat_mask.unsqueeze(-2).int()
+
+            ee = torch.einsum("...cwh,...wx,...hx->...cx", wt, ws[0], ws[1])
+        if self.density_channel:
+            ee = torch.cat(
+                [
+                    ee[:, 0:1, ...],
+                    ee[:, 1:, ...] / torch.clamp(ee[:, 0:1, ...], min=1e-6, max=1e5),
+                ],
+                dim=1,
+            )
+
+            return ee
+        else:
+            ee = ee[:, 1:, ...] / torch.clamp(ee[:, 0:1, ...], min=1e-6, max=1e5)
+            return ee

official-src/aardvark/train_module.py

@@ -0,0 +1,315 @@
+"""
+NB: this script is for illustration purposes only and is not runnable as our
+full dataset is not provided as part of the submission, due to size constraints.
+Many of the relevant paths to the data have been thus replaced by dummy paths.
+"""
+
+import os
+import sys
+import pickle
+import argparse
+
+import torch
+import torch.multiprocessing as mp
+from torch.utils.data import DataLoader
+import torch.utils.data.distributed
+from torch.utils.data.distributed import DistributedSampler
+from torch.distributed import init_process_group, destroy_process_group
+
+
+from trainer import DDPTrainer
+from loss_functions import WeightedRmseLoss, PressureWeightedRmseLoss, RmseLoss
+from misc_downscaling_functionality import ConvCNPWeatherOnToOff, DownscalingRmseLoss
+from loader import *
+from models import *
+from unet_wrap_padding import *
+
+
+sys.path.append("../npw/data")
+torch.set_float32_matmul_precision("medium")
+
+
+def ddp_setup(rank, world_size, master_port):
+    """
+    Args:
+        rank: Unique identifier of each process
+        world_size: Total number of processes
+    """
+
+    os.environ["MASTER_ADDR"] = "localhost"
+    os.environ["MASTER_PORT"] = master_port
+    init_process_group(backend="nccl", rank=rank, world_size=world_size)
+    torch.cuda.set_device(rank)
+
+
+def start_date(name):
+    if name == "train":
+        return "2007-01-02"
+    elif name == "val":
+        return "2019-01-01"
+    elif name == "test":
+        return "2018-01-01"
+    else:
+        raise Exception(f"Unrecognised split name {name}")
+
+
+def end_date(name):
+    if name == "train":
+        return "2017-12-31"
+    elif name == "val":
+        return "2019-11-01"
+    elif name == "test":
+        return "2018-12-21"
+    else:
+        raise Exception(f"Unrecognised split name {name}")
+
+
+def main(rank, world_size, output_dir, args):
+    """
+    Primary training script for the encoder, processor and decoder modules.
+    """
+
+    master_port = args.master_port
+    lead_time = args.lead_time
+    era5_mode = args.era5_mode
+    weights_dir = args.weights_dir
+    ddp_setup(rank, world_size, master_port)
+
+    # Instantiate loss function
+    if args.loss == "lw_rmse":
+        lf = WeightedRmseLoss(
+            args.res,
+            start_ind=args.start_ind,
+            end_ind=args.end_ind,
+            weight_per_variable=bool(args.weight_per_variable),
+        )
+    elif args.loss == "lw_rmse_pressure_weighted":
+        lf = PressureWeightedRmseLoss(args.res, era5_mode)
+    elif args.loss == "rmse":
+        lf = RmseLoss()
+    elif args.loss == "downscaling_rmse":
+        lf = DownscalingRmseLoss()
+
+    # Setup datasets
+
+    # Case 1: training encoder
+    if args.mode == "assimilation":
+        train_dataset = WeatherDatasetAssimilation(
+            device="cuda",
+            hadisd_mode="train",
+            start_date="2007-01-02",
+            end_date="2017-12-31",
+            lead_time=0,
+            era5_mode="4u",
+            res=args.res,
+            var_start=args.start_ind,
+            var_end=args.end_ind,
+            diff=bool(args.diff),
+            two_frames=bool(args.two_frames),
+        )
+        val_dataset = WeatherDatasetAssimilation(
+            device="cuda",
+            hadisd_mode="train",
+            start_date="2019-01-01",
+            end_date="2019-12-31",
+            lead_time=0,
+            era5_mode="4u",
+            res=args.res,
+            var_start=args.start_ind,
+            var_end=args.end_ind,
+            diff=bool(args.diff),
+            two_frames=bool(args.two_frames),
+        )
+
+    # Case 2: training processor
+    elif args.mode == "forecast":
+        if args.ic == "aardvark":
+            train_dataset = FineTuneForecastLoaderNew(
+                device="cuda",
+                mode="train",
+                lead_time=lead_time,
+                era5_mode=era5_mode,
+                res=args.res,
+                frequency=args.frequency,
+                diff=bool(args.diff),
+                aardvark_ic_path=args.aardvark_ic_path,
+                random_lt=True,
+            )
+            val_dataset = FineTuneForecastLoaderNew(
+                device="cuda",
+                mode="val",
+                lead_time=lead_time,
+                era5_mode=era5_mode,
+                res=args.res,
+                frequency=args.frequency,
+                diff=bool(args.diff),
+                aardvark_ic_path=args.aardvark_ic_path,
+            )
+        else:
+            train_dataset = ForecastLoader(
+                device="cuda",
+                mode="train",
+                lead_time=lead_time,
+                era5_mode=era5_mode,
+                res=args.res,
+                frequency=args.frequency,
+                diff=bool(args.diff),
+                u_only=False,
+                random_lt=False,
+            )
+            val_dataset = ForecastLoader(
+                device="cuda",
+                mode="val",
+                lead_time=lead_time,
+                era5_mode=era5_mode,
+                res=args.res,
+                frequency=args.frequency,
+                diff=bool(args.diff),
+                u_only=False,
+                random_lt=False,
+            )
+
+    # Case 3: training decoder
+    elif args.mode == "downscaling":
+
+        train_dataset = ForecasterDatasetDownscaling(
+            start_date="2007-01-02",
+            end_date="2017-12-31",
+            lead_time=args.lead_time,
+            hadisd_var=args.var,
+            mode="train",
+            device="cuda",
+            forecast_path=None,
+        )
+
+        val_dataset = ForecasterDatasetDownscaling(
+            start_date="2019-01-01",
+            end_date="2019-12-21",
+            lead_time=args.lead_time,
+            hadisd_var=args.var,
+            mode="train",
+            device="cuda",
+            forecast_path=None,
+        )
+
+        try:
+            os.mkdir(f"{output_dir}lt_{args.lead_time}")
+        except FileExistsError:
+            pass
+
+        output_dir = f"{output_dir}lt_{args.lead_time}/"
+
+    # Instantiate model
+
+    if args.mode == "downscaling":
+        model = ConvCNPWeatherOnToOff(
+            in_channels=args.in_channels,
+            out_channels=args.end_ind - args.start_ind,
+            int_channels=args.int_channels,
+            device="cuda",
+            res=args.res,
+            decoder=args.decoder,
+            mode=args.mode,
+            film=bool(args.film),
+        )
+    else:
+        model = ConvCNPWeather(
+            in_channels=args.in_channels,
+            out_channels=args.end_ind - args.start_ind,
+            int_channels=args.int_channels,
+            device="cuda",
+            res=args.res,
+            gnp=bool(0),
+            decoder=args.decoder,
+            mode=args.mode,
+            film=bool(args.film),
+            two_frames=bool(args.two_frames),
+        )
+
+    # Instantiate loaders
+    train_sampler = DistributedSampler(train_dataset)
+    val_sampler = DistributedSampler(val_dataset)
+
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        sampler=train_sampler,
+    )
+
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        sampler=val_sampler,
+    )
+
+    # Instantiate trainer
+
+    trainer = DDPTrainer(
+        model,
+        rank,
+        train_loader,
+        val_loader,
+        lf,
+        output_dir,
+        args.lr,
+        train_sampler,
+        weight_decay=args.weight_decay,
+        weights_path=weights_dir,
+        tune_film=args.film,
+    )
+
+    # Train model
+
+    trainer.train(n_epochs=args.epoch)
+    destroy_process_group()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--output_dir")
+    parser.add_argument("--mode")
+    parser.add_argument("--weights_dir")
+    parser.add_argument("--in_channels", type=int)
+    parser.add_argument("--out_channels", type=int)
+    parser.add_argument("--int_channels", type=int)
+    parser.add_argument("--loss")
+    parser.add_argument("--ic")
+    parser.add_argument("--decoder")
+    parser.add_argument("--film")
+    parser.add_argument("--aardvark_ic_path")
+    parser.add_argument("--two_frames", type=int, default=0)
+    parser.add_argument("--weight_per_variable", type=int, default=0)
+    parser.add_argument("--batch_size", type=int, default=128)
+    parser.add_argument("--epoch", type=int, default=50)
+    parser.add_argument("--master_port", default="12345")
+    parser.add_argument("--lr", type=float, default=5e-4)
+    parser.add_argument("--lead_time", type=int)
+    parser.add_argument("--era5_mode", default="4u")
+    parser.add_argument("--weight_decay", type=float, default=1e-6)
+    parser.add_argument("--res", type=int, default=1)
+    parser.add_argument("--frequency", type=int, default=6)
+    parser.add_argument("--diff", type=int, default=1)
+    parser.add_argument("--start_ind", type=int, default=0)
+    parser.add_argument("--end_ind", type=int, default=24)
+    parser.add_argument("--downscaling_train_start_date", default="1979-01-01")
+    parser.add_argument("--downscaling_train_end_date", default="2017-12-31")
+    parser.add_argument("--downscaling_context", default="era5")
+    parser.add_argument("--downscaling_lead_time", type=int)
+    parser.add_argument("--var", default=None)
+    args = parser.parse_args()
+
+    torch.device("cuda")
+
+    # Create results directory
+    output_dir = args.output_dir
+    if not os.path.exists(output_dir):
+        os.mkdir(output_dir)
+
+    # Save config
+    with open(output_dir + "/config.pkl", "wb") as f:
+        pickle.dump(vars(args), f)
+
+    world_size = torch.cuda.device_count()
+    mp.spawn(main, args=[world_size, output_dir, args], nprocs=world_size)

official-src/aardvark/trainer.py

@@ -0,0 +1,602 @@
+import sys
+import subprocess
+
+import numpy as np
+import torch
+from torch.nn.parallel import DistributedDataParallel as DDP
+from tqdm import tqdm
+
+from misc_downscaling_functionality import *
+from models import *
+from loss_functions import *
+
+sys.path.append("..")
+
+
+class DDPTrainer:
+    """
+    Main class for training models using DDP
+    """
+
+    def __init__(
+        self,
+        model,
+        rank,
+        train_loader,
+        val_loader,
+        loss_function,
+        save_path,
+        learning_rate,
+        sampler,
+        weight_decay,
+        test_loader=None,
+        weights_path=None,
+        tune_film=False,
+    ):
+        self.rank = rank
+        self.model = model
+        self.sampler = sampler
+        self.tune_film = tune_film
+        self.train_loader = train_loader
+        self.val_loader = val_loader
+        self.save_path = save_path
+        self.loss_function = loss_function
+        self.best_loss = 1000
+        self.test_loader = test_loader
+
+        self.model = self.model.to(rank)
+        self.model = DDP(self.model, device_ids=[rank], find_unused_parameters=True)
+
+        if self.model.module.decoder == "vit":
+            self.opt = torch.optim.AdamW(
+                model.parameters(),
+                lr=learning_rate,
+                betas=(0.9, 0.95),
+                weight_decay=1e-5,
+            )
+            self.scheduler = optim.lr_scheduler.CosineAnnealingLR(self.opt, 891 * 80)
+
+        else:
+            self.opt = torch.optim.Adam(
+                model.parameters(), lr=learning_rate, weight_decay=weight_decay
+            )
+
+        self.losses = []
+        self.train_losses = []
+        self.maes = []
+
+    def _unravel_to_numpy(self, x):
+        return x.view(-1).detach().cpu().numpy()
+
+    def eval_epoch(self, fix_sigma, epoch):
+
+        self.model.eval()
+
+        mae = []
+        lf = []
+        lf_unnorm = []
+
+        with torch.no_grad():
+            for count, task in enumerate(self.val_loader):
+
+                out = self.model(task, film_index=0)
+
+                prev_step = None
+
+                mae.append(
+                    np.nanmean(
+                        np.abs(
+                            task["y_target"].detach().cpu()
+                            - out[..., : task["y_target"].shape[-1]].detach().cpu()
+                        )
+                    )
+                )
+                l = (
+                    self.loss_function(
+                        task["y_target"], out, prev_step, fix_sigma=fix_sigma
+                    )
+                    .detach()
+                    .item()
+                )
+                lf.append(l)
+
+                try:
+                    ic = self.train_loader.dataset.unnorm_base_context(
+                        task["y_context"][:, :-11, ...]
+                    ).permute(0, 3, 2, 1)
+                    unnorm_pred = self.train_loader.dataset.unnorm_pred(out)
+                    unnorm_target = self.train_loader.dataset.unnorm_pred(
+                        task["y_target"]
+                    )
+
+                    unnorm_pred = unnorm_pred + ic
+                    unnorm_target = unnorm_target + ic
+
+                    lu = (
+                        self.loss_function(
+                            unnorm_target,
+                            unnorm_pred,
+                            prev_step,
+                            fix_sigma=fix_sigma,
+                            expand=True,
+                        )
+                        .detach()
+                        .cpu()
+                        .numpy()
+                    )
+
+                    lf_unnorm.append(lu)
+
+                except:
+                    pass
+
+            if self.test_loader is not None:
+                forecasts = []
+                targets = []
+                stations = []
+
+                for count, task in enumerate(self.test_loader):
+                    out = self.model(task, film_index=0)
+                    forecasts.append(unnormalise_hadisd_tas(out.detach().cpu().numpy()))
+                    targets.append(
+                        unnormalise_hadisd_tas(task["y_target"].detach().cpu().numpy())
+                    )
+                    stations.append(
+                        task["downscaling"]["x_target"].detach().cpu().numpy() * 360
+                    )
+
+                # Save the test outputs
+                np.save(
+                    self.save_path + f"forecast_{self.rank}_{self.epoch}.npy",
+                    np.concatenate(forecasts, axis=0),
+                )
+                np.save(
+                    self.save_path + f"targets_{self.rank}_{self.epoch}.npy",
+                    np.concatenate(targets, axis=0),
+                )
+                np.save(
+                    self.save_path + f"stations_{self.rank}_{self.epoch}.npy",
+                    np.concatenate(stations, axis=0),
+                )
+
+        log_loss = np.nanmean(np.array(lf))
+
+        if log_loss < self.best_loss:
+            np.save(
+                self.save_path + "unnorm_preds.npy",
+                self.train_loader.dataset.unnorm_pred(out).detach().cpu().numpy(),
+            )
+            np.save(
+                self.save_path + "unnorm_targets.npy",
+                self.train_loader.dataset.unnorm_pred(task["y_target"])
+                .detach()
+                .cpu()
+                .numpy(),
+            )
+        log_loss_unnorm = np.nanmean(np.stack(lf_unnorm), axis=0)
+
+        if np.logical_and(self.rank == 0, self.epoch % 5 == 0):
+
+            np.save(self.save_path + "preds_eval.npy", out.cpu().numpy())
+            np.save(
+                self.save_path + "y_target_eval.npy", task["y_target"].cpu().numpy()
+            )
+
+        return log_loss, log_loss_unnorm
+
+    def train(self, n_epochs=100):
+
+        torch.cuda.set_device(self.rank)
+        subprocess.run(["cp", "reproduce_training.sh", f"{self.save_path}"])
+
+        train_loss = []
+        ll = []
+
+        fix_sigma = False
+        prev_step = None
+
+        self.epoch = 0
+        epoch_loss, log_loss_unnorm = self.eval_epoch(fix_sigma, 0)
+        train_loss = np.mean(train_loss)
+
+        for epoch in range(n_epochs):
+            self.epoch = epoch
+
+            self.sampler.set_epoch(epoch)
+
+            self.model.train()
+            train_loss = []
+            with tqdm(self.train_loader, unit="batch") as tepoch:
+                for count, task in enumerate(tepoch):
+
+                    out = self.model(task, film_index=0)
+
+                    loss = self.loss_function(
+                        task["y_target"], out, prev_step, fix_sigma=fix_sigma
+                    )
+
+                    loss.backward()
+                    tepoch.set_postfix(loss=loss.item())
+                    prev_step = out
+
+                    self.opt.step()
+                    self.opt.zero_grad()
+                    train_loss.append(loss.item())
+                    try:
+                        if self.model.module.decoder == "vit":
+                            if epoch > 20:
+                                self.scheduler.step()
+                    except:
+                        pass
+
+            epoch_loss, log_loss_unnorm = self.eval_epoch(fix_sigma, epoch)
+            train_loss = np.mean(train_loss)
+            ll.append(log_loss_unnorm)
+
+            self.losses.append(epoch_loss)
+            self.train_losses.append(train_loss)
+            np.save(
+                self.save_path + "losses_{}.npy".format(self.rank),
+                np.array(self.losses),
+            )
+            np.save(
+                self.save_path + "train_losses_{}.npy".format(self.rank),
+                np.array(self.train_losses),
+            )
+            np.save(self.save_path + "rmse_{}.npy".format(self.rank), np.array(ll))
+
+            if epoch_loss < self.best_loss:
+                self.best_loss = epoch_loss
+
+                if self.rank == 0:
+                    if self.model.module.decoder == "vit":
+                        torch.save(
+                            {
+                                "epoch": epoch,
+                                "model_state_dict": self.model.state_dict(),
+                                "optimizer_state_dict": self.opt.state_dict(),
+                                "scheduler_state_dict": self.scheduler.state_dict(),
+                                "loss": epoch_loss,
+                            },
+                            self.save_path + "epoch_{}".format(epoch),
+                        )
+                    else:
+                        torch.save(
+                            {
+                                "epoch": epoch,
+                                "model_state_dict": self.model.state_dict(),
+                                "optimizer_state_dict": self.opt.state_dict(),
+                                "loss": epoch_loss,
+                            },
+                            self.save_path + "epoch_{}".format(epoch),
+                        )
+
+                try:
+                    np.save(
+                        self.save_path + "preds_train.npy".format(epoch),
+                        out.detach().cpu().numpy(),
+                    )
+                    np.save(
+                        self.save_path + "y_target_train.npy".format(epoch),
+                        task["y_target"].detach().cpu().numpy(),
+                    )
+                except:
+                    pass
+
+
+class DDPTrainerE2E:
+    """
+    Training class for E2E
+    """
+
+    def __init__(
+        self,
+        model,
+        rank,
+        train_loader,
+        val_loader,
+        loss_function,
+        save_path,
+        learning_rate,
+        sampler,
+        weight_decay,
+        hadisd_variable_name,
+        test_loader=None,
+        weights_path=None,
+        tune_film=False,
+    ):
+
+        self.rank = rank
+        self.model = model
+        self.sampler = sampler
+        self.tune_film = tune_film
+        self.train_loader = train_loader
+        self.val_loader = val_loader
+        self.save_path = save_path
+        self.loss_function = loss_function
+        self.best_loss = 1000
+        self.test_loader = test_loader
+        self.hadisd_variable_name = hadisd_variable_name
+
+        self.model = self.model.to(rank)
+        self.model = DDP(self.model, device_ids=[rank], find_unused_parameters=True)
+
+        if self.model.module.decoder == "vit":
+            self.opt = torch.optim.AdamW(
+                model.parameters(),
+                lr=learning_rate,
+                betas=(0.9, 0.95),
+                weight_decay=1e-5,
+            )
+            self.scheduler = optim.lr_scheduler.CosineAnnealingLR(self.opt, 891 * 80)
+        else:
+            self.opt = torch.optim.Adam(
+                model.parameters(), lr=learning_rate, weight_decay=weight_decay
+            )
+
+        self.losses = []
+        self.train_losses = []
+
+        self.maes = []
+
+    def _unravel_to_numpy(self, x):
+        return x.view(-1).detach().cpu().numpy()
+
+    def eval_epoch(self, fix_sigma, epoch):
+
+        self.model.eval()
+
+        mae = []
+        lf = []
+        lf_unnorm = []
+
+        with torch.no_grad():
+            forecasts = []
+            targets = []
+            stations = []
+            indices = []
+            for count, task in tqdm(enumerate(self.val_loader)):
+
+                out = self.model(task, film_index=0)
+                forecasts.append(
+                    unnormalise_hadisd_var(
+                        out.detach().cpu().numpy(), self.hadisd_variable_name
+                    )
+                )
+                targets.append(
+                    unnormalise_hadisd_var(
+                        task["y_target"].detach().cpu().numpy(),
+                        self.hadisd_variable_name,
+                    )
+                )
+                stations.append(
+                    task["downscaling"]["x_target"].detach().cpu().numpy() * 360
+                )
+                indices.append(task["index"])
+
+                prev_step = None
+
+                mae.append(
+                    np.nanmean(
+                        np.abs(
+                            task["y_target"].detach().cpu()
+                            - out[..., : task["y_target"].shape[-1]].detach().cpu()
+                        )
+                    )
+                )
+                l = (
+                    self.loss_function(
+                        task["y_target"], out, prev_step, fix_sigma=fix_sigma
+                    )
+                    .detach()
+                    .item()
+                )
+                lf.append(l)
+
+                try:
+                    ic = self.train_loader.dataset.unnorm_base_context(
+                        task["y_context"][:, :-11, ...]
+                    ).permute(0, 3, 2, 1)
+                    unnorm_pred = self.train_loader.dataset.unnorm_pred(out)
+                    unnorm_target = self.train_loader.dataset.unnorm_pred(
+                        task["y_target"]
+                    )
+
+                    unnorm_pred = unnorm_pred + ic
+                    unnorm_target = unnorm_target + ic
+
+                    lu = (
+                        self.loss_function(
+                            unnorm_target,
+                            unnorm_pred,
+                            prev_step,
+                            fix_sigma=fix_sigma,
+                            expand=True,
+                        )
+                        .detach()
+                        .cpu()
+                        .numpy()
+                    )
+
+                    lf_unnorm.append(lu)
+
+                except:
+                    pass
+
+            # Save the test outputs
+            np.save(
+                self.save_path + f"val_forecast_{self.rank}_{self.epoch}.npy",
+                np.concatenate(forecasts, axis=0),
+            )
+            np.save(
+                self.save_path + f"val_targets_{self.rank}_{self.epoch}.npy",
+                np.concatenate(targets, axis=0),
+            )
+            np.save(
+                self.save_path + f"val_stations_{self.rank}_{self.epoch}.npy",
+                np.concatenate(stations, axis=0),
+            )
+            np.save(
+                self.save_path + f"val_indices_{self.rank}_{self.epoch}.npy",
+                np.concatenate(indices, axis=0),
+            )
+
+            if self.test_loader is not None:
+                forecasts = []
+                targets = []
+                stations = []
+                indices = []
+
+                for count, task in enumerate(self.test_loader):
+                    out = self.model(task, film_index=0)
+                    forecasts.append(
+                        unnormalise_hadisd_var(
+                            out.detach().cpu().numpy(), self.hadisd_variable_name
+                        )
+                    )
+                    targets.append(
+                        unnormalise_hadisd_var(
+                            task["y_target"].detach().cpu().numpy(),
+                            self.hadisd_variable_name,
+                        )
+                    )
+                    stations.append(
+                        task["downscaling"]["x_target"].detach().cpu().numpy() * 360
+                    )
+                    indices.append(task["index"])
+
+                np.save(
+                    self.save_path + f"test_forecast_{self.rank}_{self.epoch}.npy",
+                    np.concatenate(forecasts, axis=0),
+                )
+                np.save(
+                    self.save_path + f"test_targets_{self.rank}_{self.epoch}.npy",
+                    np.concatenate(targets, axis=0),
+                )
+                np.save(
+                    self.save_path + f"test_stations_{self.rank}_{self.epoch}.npy",
+                    np.concatenate(stations, axis=0),
+                )
+                np.save(
+                    self.save_path + f"test_indices_{self.rank}_{self.epoch}.npy",
+                    np.concatenate(indices, axis=0),
+                )
+
+        log_loss = np.nanmean(np.array(lf))
+
+        if log_loss < self.best_loss:
+            np.save(
+                self.save_path + "unnorm_preds.npy",
+                self.train_loader.dataset.unnorm_pred(out).detach().cpu().numpy(),
+            )
+            np.save(
+                self.save_path + "unnorm_targets.npy",
+                self.train_loader.dataset.unnorm_pred(task["y_target"])
+                .detach()
+                .cpu()
+                .numpy(),
+            )
+        log_loss_unnorm = np.nanmean(np.stack(lf_unnorm), axis=0)
+
+        if np.logical_and(self.rank == 0, self.epoch % 5 == 0):
+
+            np.save(self.save_path + "preds_eval.npy", out.cpu().numpy())
+            np.save(
+                self.save_path + "y_target_eval.npy", task["y_target"].cpu().numpy()
+            )
+
+        return log_loss, log_loss_unnorm
+
+    def train(self, n_epochs=100):
+
+        torch.cuda.set_device(self.rank)
+        subprocess.run(["cp", "reproduce_training.sh", f"{self.save_path}"])
+
+        train_loss = []
+        ll = []
+
+        fix_sigma = False
+        prev_step = None
+
+        self.epoch = 0
+        epoch_loss, log_loss_unnorm = self.eval_epoch(fix_sigma, 0)
+        train_loss = np.mean(train_loss)
+
+        for epoch in range(n_epochs):
+            self.epoch = epoch
+            self.sampler.set_epoch(epoch)
+
+            epoch_loss, log_loss_unnorm = self.eval_epoch(fix_sigma, epoch)
+            train_loss = np.mean(train_loss)
+            ll.append(log_loss_unnorm)
+
+            self.losses.append(epoch_loss)
+            self.train_losses.append(train_loss)
+            np.save(
+                self.save_path + "losses_{}.npy".format(self.rank),
+                np.array(self.losses),
+            )
+            np.save(
+                self.save_path + "train_losses_{}.npy".format(self.rank),
+                np.array(self.train_losses),
+            )
+            np.save(self.save_path + "rmse_{}.npy".format(self.rank), np.array(ll))
+
+            if epoch_loss < self.best_loss:
+                self.best_loss = epoch_loss
+
+                if self.rank == 0:
+                    if self.model.module.decoder == "vit":
+                        torch.save(
+                            {
+                                "epoch": epoch,
+                                "model_state_dict": self.model.state_dict(),
+                                "optimizer_state_dict": self.opt.state_dict(),
+                                "scheduler_state_dict": self.scheduler.state_dict(),
+                                "loss": epoch_loss,
+                            },
+                            self.save_path + "epoch_{}".format(epoch),
+                        )
+                    else:
+                        torch.save(
+                            {
+                                "epoch": epoch,
+                                "model_state_dict": self.model.state_dict(),
+                                "optimizer_state_dict": self.opt.state_dict(),
+                                "loss": epoch_loss,
+                            },
+                            self.save_path + "epoch_{}".format(epoch),
+                        )
+
+                try:
+                    np.save(
+                        self.save_path + "preds_train.npy".format(epoch),
+                        out.detach().cpu().numpy(),
+                    )
+                    np.save(
+                        self.save_path + "y_target_train.npy".format(epoch),
+                        task["y_target"].detach().cpu().numpy(),
+                    )
+                except:
+                    pass
+
+            self.model.train()
+            train_loss = []
+            with tqdm(self.train_loader, unit="batch") as tepoch:
+                for count, task in tqdm(enumerate(tepoch)):
+                    out = self.model(task, film_index=0)
+
+                    loss = self.loss_function(
+                        task["y_target"], out, prev_step, fix_sigma=fix_sigma
+                    )
+
+                    loss.backward()
+                    tepoch.set_postfix(loss=loss.item())
+                    prev_step = out
+
+                    self.opt.step()
+                    self.opt.zero_grad()
+                    train_loss.append(loss.item())
+                    if self.model.module.decoder == "vit":
+                        if epoch > 20:
+                            self.scheduler.step()

official-src/aardvark/unet_wrap_padding.py

@@ -0,0 +1,310 @@
+import math
+
+import torch
+import torch.nn as nn
+
+
+def cylindrical_conv_pad(x, w_pad):
+    return torch.cat([x[..., -w_pad:], x, x[..., :w_pad]], axis=-1)
+
+
+class CylindricalConv2D(nn.Conv2d):
+    """
+    UNet with cylinderical boundary conditions
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int,
+    ):
+        super().__init__(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+        )
+
+        assert self.kernel_size[0] % 2 == 1
+        assert self.kernel_size[1] % 2 == 1
+
+        self.h_pad = self.kernel_size[0] // 2
+        self.w_pad = self.kernel_size[1] // 2
+
+    def forward(self, x: torch.Tensor):
+        x = nn.functional.pad(x, (0, 0, self.h_pad, self.h_pad))
+        return super().forward(cylindrical_conv_pad(x, self.w_pad))
+
+
+class CylindricalConvTranspose2D(nn.ConvTranspose2d):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int,
+    ):
+        super().__init__(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            kernel_size=kernel_size,
+            stride=stride,
+        )
+
+        assert self.kernel_size[0] % 2 == 1
+        assert self.kernel_size[1] % 2 == 1
+
+        self.sh, self.sw = self.stride
+        self.kh, self.kw = self.kernel_size
+
+        self.h_pad = math.ceil(((self.sh - 1) + 2 * (self.kh // 2)) / self.sh)
+        self.w_pad = math.ceil(((self.sw - 1) + 2 * (self.kw // 2)) / self.sw)
+
+        self.h0 = self.sh * self.h_pad - (self.sh - 1) + (self.kh // 2)
+        self.w0 = self.sw * self.w_pad - (self.sw - 1) + (self.kw // 2)
+
+        self._bias = nn.Parameter(10**-3 * torch.randn(out_channels))
+
+    def forward(self, x: torch.Tensor):
+
+        Nh = x.shape[2] * self.sh
+        Nw = x.shape[3] * self.sw
+
+        x = cylindrical_conv_pad(x, self.w_pad)
+        x = nn.functional.pad(x, (0, 0, self.h_pad, self.h_pad))
+
+        x = super().forward(x)
+
+        return x[:, :, self.h0 : self.h0 + Nh, self.w0 : self.w0 + Nw]
+
+
+class Down(nn.Module):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        p=0,
+        film=False,
+        down=True,
+        attn=False,
+    ):
+
+        super().__init__()
+
+        self.film = film
+        self.attn = attn
+
+        self.conv_1 = CylindricalConv2D(
+            in_channels, out_channels, kernel_size=3, stride=1
+        )
+        if down:
+            self.conv_2 = CylindricalConv2D(
+                out_channels, out_channels, kernel_size=3, stride=2
+            )
+        else:
+            self.conv_2 = CylindricalConv2D(
+                out_channels, out_channels, kernel_size=3, stride=1
+            )
+
+        self.bn_1 = nn.BatchNorm2d(out_channels)
+        self.bn_2 = nn.BatchNorm2d(out_channels)
+
+        self.activation = nn.GELU()
+
+        if film:
+
+            self.gamma_1 = torch.nn.Parameter(
+                torch.ones(10, out_channels, 1, 1),
+            )
+            self.gamma_2 = torch.nn.Parameter(
+                torch.ones(10, out_channels, 1, 1),
+            )
+            self.beta_1 = torch.nn.Parameter(
+                torch.zeros(10, out_channels, 1, 1),
+            )
+            self.beta_2 = torch.nn.Parameter(
+                torch.zeros(10, out_channels, 1, 1),
+            )
+
+        if self.attn:
+            self.mha = AttentionBlock(n_channels=out_channels, n_heads=8)
+
+    def forward(self, xi, film_index=None):
+
+        film_index = film_index[:, 0].int()
+
+        x = self.conv_1(xi)
+        x = self.bn_1(x)
+        if self.film:
+            g1 = torch.index_select(self.gamma_1, 0, film_index)
+            b1 = torch.index_select(self.beta_1, 0, film_index)
+
+            x = g1 * x + b1
+
+        x = self.activation(x)
+        x = self.conv_2(x)
+        x = self.bn_2(x)
+        if self.film:
+            g2 = torch.index_select(self.gamma_2, 0, film_index)
+            b2 = torch.index_select(self.beta_2, 0, film_index)
+            x = g2 * x + b2
+        x = self.activation(x)
+
+        if self.attn:
+            x = self.mha(x)
+
+        return x
+
+
+class Up(nn.Module):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        p,
+        bilinear=False,
+        film=False,
+        stride=2,
+        attn=False,
+    ):
+        super().__init__()
+
+        self.film = film
+
+        if bilinear:
+            self.up = nn.Upsample(
+                scale_factor=2,
+                mode="bilinear",
+                align_corners=True,
+            )
+        else:
+            self.up = CylindricalConvTranspose2D(
+                in_channels, out_channels, kernel_size=3, stride=stride
+            )
+
+        self.conv = Down(
+            out_channels,
+            out_channels,
+            p=0,
+            film=film,
+            down=False,
+            attn=attn,
+        )
+
+    def forward(self, x1, x2, film_index=None):
+        x1 = self.up(x1)
+        x1 = self.conv(x1, film_index=film_index)
+
+        if x1.shape[-1] != x2.shape[-1]:
+            x1 = x1[..., :, :-1]
+        if x1.shape[-2] != x2.shape[-2]:
+            x1 = x1[..., :-1, :]
+
+        return torch.cat([x2, x1], dim=1)
+
+
+class Unet(nn.Module):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        div_factor=1,
+        p=0.0,
+        context=True,
+        film=False,
+        film_base=True,
+    ):
+        super(Unet, self).__init__()
+
+        self.n_channels = in_channels
+        self.bilinear = True
+        self.fp = nn.Softplus()
+        self.variances = nn.Parameter(torch.zeros([out_channels]))
+        self.context = context
+        self.film = film
+
+        m = 1
+
+        self.down1 = Down(
+            self.n_channels,
+            m * 128 // div_factor,
+            p=0,
+            film=film,
+            attn=False,
+        )
+        self.down2 = Down(
+            m * 128 // div_factor,
+            m * 256 // div_factor,
+            p=0,
+            film=film,
+            attn=False,
+        )
+        self.down3 = Down(
+            m * 256 // div_factor,
+            m * 512 // div_factor,
+            p=0,
+            film=film,
+            attn=False,
+        )
+        self.down4 = Down(
+            m * 512 // div_factor,
+            m * 512 // div_factor,
+            p=0,
+            film=film,
+            attn=False,
+        )
+        self.up1 = Up(
+            m * 512 // div_factor,
+            m * 512 // div_factor,
+            p=0,
+            film=film,
+            attn=False,
+        )
+        self.up2 = Up(
+            m * 1024 // div_factor,
+            m * 256 // div_factor,
+            p=0,
+            film=film,
+            attn=False,
+        )
+        self.up3 = Up(
+            m * 512 // div_factor,
+            m * 128 // div_factor,
+            p=0,
+            film=film,
+            attn=False,
+        )
+        self.up4 = Up(
+            m * 256 // div_factor,
+            m * 64 // div_factor,
+            p=0,
+            film=film,
+            attn=False,
+        )
+
+        self.out = nn.Conv2d(
+            m * 64 // div_factor + in_channels,
+            out_channels,
+            kernel_size=1,
+            bias=False,
+        )
+
+    def forward(self, x, film_index=None):
+
+        x1 = x.contiguous()
+        x2 = self.down1(x1, film_index=film_index)
+        x3 = self.down2(x2, film_index=film_index)
+        x4 = self.down3(x3, film_index=film_index)
+        x5 = self.down4(x4, film_index=film_index)
+        x = self.up1(x5, x4, film_index=film_index)
+        x = self.up2(x, x3, film_index=film_index)
+        x = self.up3(x, x2, film_index=film_index)
+        x = self.up4(x, x1, film_index=film_index)
+        x = self.out(x)
+
+        return x.permute(0, 2, 3, 1)

official-src/aardvark/utils.py

@@ -0,0 +1,23 @@
+from collections import defaultdict
+import numpy as np
+from torch.nn.utils.rnn import pad_sequence
+
+
+def channels_to_2nd_dim(x):
+    return x.permute(*([0, x.dim() - 1] + list(range(1, x.dim() - 1))))
+
+
+def channels_to_final_dim(x):
+    return x.permute(*([0] + list(range(2, x.dim())) + [1]))
+
+
+def collate(tensor_list):
+    out_dict = defaultdict()
+    for k in tensor_list[0].keys():
+        out_dict[k] = [t[k] for t in tensor_list]
+        out_dict[k] = pad_sequence(
+            out_dict[k],
+            padding_value=np.nan,
+            batch_first=True,
+        )
+    return out_dict

official-src/aardvark/vit.py

@@ -0,0 +1,298 @@
+"""
+This file is adapted from Nguyen, Tung, et al. "ClimaX: A foundation model
+for weather and climate." arXiv preprint arXiv:2301.10343 (2023).
+Code from this project is available at https://github.com/microsoft/ClimaX
+"""
+
+from functools import lru_cache
+
+import numpy as np
+import torch
+import torch.nn as nn
+from timm.models.vision_transformer import Block, PatchEmbed, trunc_normal_
+
+from architectures import MLP
+
+
+def get_2d_sincos_pos_embed(embed_dim, grid_size_h, grid_size_w, cls_token=False):
+
+    grid_h = np.arange(grid_size_h, dtype=np.float32)
+    grid_w = np.arange(grid_size_w, dtype=np.float32)
+    grid = np.meshgrid(grid_w, grid_h)
+    grid = np.stack(grid, axis=0)
+
+    grid = grid.reshape([2, 1, grid_size_h, grid_size_w])
+    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
+    if cls_token:
+        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
+    return pos_embed
+
+
+def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
+    assert embed_dim % 2 == 0
+
+    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])
+    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])
+
+    emb = np.concatenate([emb_h, emb_w], axis=1)
+    return emb
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float32)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega
+
+    pos = pos.reshape(-1)
+    out = np.einsum("m,d->md", pos, omega)
+
+    emb_sin = np.sin(out)
+    emb_cos = np.cos(out)
+
+    emb = np.concatenate([emb_sin, emb_cos], axis=1)
+    return emb
+
+
+def interpolate_pos_embed(model, checkpoint_model, new_size=(64, 128)):
+    if "net.pos_embed" in checkpoint_model:
+        pos_embed_checkpoint = checkpoint_model["net.pos_embed"]
+        embedding_size = pos_embed_checkpoint.shape[-1]
+        orig_num_patches = pos_embed_checkpoint.shape[-2]
+        patch_size = model.patch_size
+        w_h_ratio = 2
+        orig_h = int((orig_num_patches // w_h_ratio) ** 0.5)
+        orig_w = w_h_ratio * orig_h
+        orig_size = (orig_h, orig_w)
+        new_size = (new_size[0] // patch_size, new_size[1] // patch_size)
+
+        if orig_size[0] != new_size[0]:
+            print(
+                "Interpolate PEs from %dx%d to %dx%d"
+                % (orig_size[0], orig_size[1], new_size[0], new_size[1])
+            )
+            pos_tokens = pos_embed_checkpoint.reshape(
+                -1, orig_size[0], orig_size[1], embedding_size
+            ).permute(0, 3, 1, 2)
+            new_pos_tokens = torch.nn.functional.interpolate(
+                pos_tokens,
+                size=(new_size[0], new_size[1]),
+                mode="bicubic",
+                align_corners=False,
+            )
+            new_pos_tokens = new_pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+            checkpoint_model["net.pos_embed"] = new_pos_tokens
+
+
+def interpolate_channel_embed(checkpoint_model, new_len):
+    if "net.channel_embed" in checkpoint_model:
+        channel_embed_checkpoint = checkpoint_model["net.channel_embed"]
+        old_len = channel_embed_checkpoint.shape[1]
+        if new_len <= old_len:
+            checkpoint_model["net.channel_embed"] = channel_embed_checkpoint[
+                :, :new_len
+            ]
+
+
+class ViT(nn.Module):
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        h_channels,
+        img_size=[256, 128],
+        patch_size=8,
+        depth=24,
+        decoder_depth=4,
+        num_heads=16,
+        mlp_ratio=4.0,
+        drop_path=0.0,
+        drop_rate=0.0,
+        per_var_embedding=True,
+    ):
+        super().__init__()
+
+        self.img_size = img_size
+        self.patch_size = patch_size
+        default_vars = [str(i) for i in range(in_channels)]
+        self.default_vars = default_vars
+        embed_dim = h_channels
+        self.per_var_embedding = per_var_embedding
+
+        if self.per_var_embedding:
+            self.token_embeds = nn.ModuleList(
+                [
+                    PatchEmbed(img_size, patch_size, 1, embed_dim)
+                    for i in range(len(default_vars))
+                ]
+            )
+        else:
+            self.token_embeds = nn.ModuleList(
+                [PatchEmbed(img_size, patch_size, in_channels, embed_dim)]
+            )
+        self.num_patches = self.token_embeds[0].num_patches
+
+        self.var_embed, self.var_map = self.create_var_embedding(embed_dim)
+        self.var_query = nn.Parameter(torch.zeros(1, 1, embed_dim), requires_grad=True)
+        self.var_agg = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
+
+        self.pos_embed = nn.Parameter(
+            torch.zeros(1, self.num_patches, embed_dim), requires_grad=True
+        )
+        self.lead_time_embed = nn.Linear(1, embed_dim)
+
+        self.out_dim = out_channels
+        self.pos_drop = nn.Dropout(p=drop_rate)
+        dpr = [x.item() for x in torch.linspace(0, drop_path, depth)]
+        self.blocks = nn.ModuleList(
+            [
+                Block(
+                    embed_dim,
+                    num_heads,
+                    mlp_ratio,
+                    qkv_bias=True,
+                    drop_path=dpr[i],
+                    norm_layer=nn.LayerNorm,
+                    drop=drop_rate,
+                )
+                for i in range(depth)
+            ]
+        )
+        self.norm = nn.LayerNorm(embed_dim)
+
+        self.head = nn.ModuleList()
+        for _ in range(decoder_depth):
+            self.head.append(nn.Linear(embed_dim, embed_dim))
+            self.head.append(nn.GELU())
+        self.head.append(nn.Linear(embed_dim, self.out_dim * patch_size**2))
+        self.head = nn.Sequential(*self.head)
+
+        self.initialize_weights()
+        if not self.per_var_embedding:
+            self.mlp = MLP(in_channels=277, out_channels=256)
+
+    def initialize_weights(self):
+        pos_embed = get_2d_sincos_pos_embed(
+            self.pos_embed.shape[-1],
+            int(self.img_size[0] / self.patch_size),
+            int(self.img_size[1] / self.patch_size),
+            cls_token=False,
+        )
+        self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+
+        var_embed = get_1d_sincos_pos_embed_from_grid(
+            self.var_embed.shape[-1], np.arange(len(self.default_vars))
+        )
+        self.var_embed.data.copy_(torch.from_numpy(var_embed).float().unsqueeze(0))
+
+        for i in range(len(self.token_embeds)):
+            w = self.token_embeds[i].proj.weight.data
+            trunc_normal_(w.view([w.shape[0], -1]), std=0.02)
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def create_var_embedding(self, dim):
+        var_embed = nn.Parameter(
+            torch.zeros(1, len(self.default_vars), dim), requires_grad=True
+        )
+        var_map = {}
+        idx = 0
+        for var in self.default_vars:
+            var_map[var] = idx
+            idx += 1
+        return var_embed, var_map
+
+    @lru_cache(maxsize=None)
+    def get_var_ids(self, vars, device):
+        ids = np.array([self.var_map[var] for var in vars])
+        return torch.from_numpy(ids).to(device)
+
+    def get_var_emb(self, var_emb, vars):
+        ids = self.get_var_ids(vars, var_emb.device)
+        return var_emb[:, ids, :]
+
+    def unpatchify(self, x: torch.Tensor, h=None, w=None):
+        p = self.patch_size
+        c = self.out_dim
+        h = self.img_size[0] // p if h is None else h // p
+        w = self.img_size[1] // p if w is None else w // p
+        assert h * w == x.shape[1]
+
+        x = x.reshape(shape=(x.shape[0], h, w, p, p, c))
+        x = torch.einsum("nhwpqc->nchpwq", x)
+        imgs = x.reshape(shape=(x.shape[0], c, h * p, w * p))
+        return imgs
+
+    def aggregate_variables(self, x: torch.Tensor):
+        b, _, l, _ = x.shape
+        x = torch.einsum("bvld->blvd", x)
+        x = x.flatten(0, 1)
+
+        var_query = self.var_query.repeat_interleave(x.shape[0], dim=0)
+        x, _ = self.var_agg(var_query, x, x)
+        x = x.squeeze()
+
+        x = x.unflatten(dim=0, sizes=(b, l))
+        return x
+
+    def mlp_embedding(self, x):
+
+        return
+
+    def forward_encoder(self, x, lead_times, variables):
+
+        if isinstance(variables, list):
+            variables = tuple(variables)
+
+        if self.per_var_embedding:
+            embeds = []
+            var_ids = self.get_var_ids(variables, x.device)
+            for i in range(len(var_ids)):
+                id = var_ids[i]
+                embeds.append(self.token_embeds[id](x[:, i : i + 1]))
+            x = torch.stack(embeds, dim=1)
+
+            var_embed = self.get_var_emb(self.var_embed, variables)
+            x = x + var_embed.unsqueeze(2)
+
+            x = self.aggregate_variables(x)
+
+        else:
+            x = self.mlp(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)
+            x = self.token_embeds[0](x)
+
+        x = x + self.pos_embed
+
+        lead_time_emb = self.lead_time_embed(lead_times.unsqueeze(-1))
+        lead_time_emb = lead_time_emb.unsqueeze(1)
+
+        x = x + lead_time_emb
+
+        x = self.pos_drop(x)
+
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+
+        return x
+
+    def forward(self, x, lead_times=None, film_index=None):
+
+        if lead_times is None:
+            lead_times = torch.ones(x.shape[0]).float().cuda().unsqueeze(-1)
+
+        out_transformers = self.forward_encoder(x, lead_times[:, 0], self.default_vars)
+        preds = self.head(out_transformers)
+        preds = self.unpatchify(preds)
+
+        return preds.permute(0, 2, 3, 1)

official-src/data/grid_lon_lat/era5_x_1.npy

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

official-src/data/grid_lon_lat/era5_y_1.npy

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

official-src/data/norm_factors/mean_4u_1.npy

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

official-src/data/norm_factors/mean_diff_4u_1.npy

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

official-src/data/norm_factors/mean_hadisd_tas.npy

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

official-src/data/norm_factors/mean_hadisd_ws.npy

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

official-src/data/norm_factors/std_4u_1.npy

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

official-src/data/norm_factors/std_diff_4u_1.npy

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

official-src/data/norm_factors/std_hadisd_tas.npy

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

official-src/data/norm_factors/std_hadisd_ws.npy

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

official-src/data/sample_data_final.pkl

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

official-src/environment.yml

@@ -0,0 +1,328 @@
+name: npw
+channels:
+  - pytorch
+  - nvidia
+  - conda-forge
+  - defaults
+dependencies:
+  - _libgcc_mutex=0.1=conda_forge
+  - _openmp_mutex=4.5=2_gnu
+  - anyio=4.2.0=py38h06a4308_0
+  - argon2-cffi=21.3.0=pyhd3eb1b0_0
+  - argon2-cffi-bindings=21.2.0=py38h7f8727e_0
+  - asttokens=2.0.5=pyhd3eb1b0_0
+  - async-lru=2.0.4=py38h06a4308_0
+  - attrs=23.1.0=py38h06a4308_0
+  - babel=2.11.0=py38h06a4308_0
+  - backcall=0.2.0=pyhd3eb1b0_0
+  - beautifulsoup4=4.12.2=py38h06a4308_0
+  - blas=1.0=mkl
+  - bleach=4.1.0=pyhd3eb1b0_0
+  - blosc=1.21.5=hc2324a3_1
+  - bottleneck=1.3.7=py38ha9d4c09_0
+  - brotli=1.0.9=h5eee18b_7
+  - brotli-bin=1.0.9=h5eee18b_7
+  - brotli-python=1.0.9=py38h6a678d5_7
+  - bzip2=1.0.8=h5eee18b_5
+  - c-ares=1.19.1=h5eee18b_0
+  - ca-certificates=2024.6.2=hbcca054_0
+  - cartopy=0.21.1=py38h1176785_0
+  - certifi=2024.6.2=pyhd8ed1ab_0
+  - cffi=1.16.0=py38h5eee18b_0
+  - charset-normalizer=2.0.4=pyhd3eb1b0_0
+  - comm=0.2.1=py38h06a4308_0
+  - contourpy=1.0.5=py38hdb19cb5_0
+  - cuda-cudart=12.1.105=0
+  - cuda-cupti=12.1.105=0
+  - cuda-libraries=12.1.0=0
+  - cuda-nvrtc=12.1.105=0
+  - cuda-nvtx=12.1.105=0
+  - cuda-opencl=12.4.99=0
+  - cuda-runtime=12.1.0=0
+  - cycler=0.11.0=pyhd3eb1b0_0
+  - cyrus-sasl=2.1.28=h52b45da_1
+  - dbus=1.13.18=hb2f20db_0
+  - debugpy=1.6.7=py38h6a678d5_0
+  - decorator=5.1.1=pyhd3eb1b0_0
+  - defusedxml=0.7.1=pyhd3eb1b0_0
+  - eccodes=2.35.0=he84ddb8_0
+  - exceptiongroup=1.2.0=py38h06a4308_0
+  - executing=0.8.3=pyhd3eb1b0_0
+  - expat=2.5.0=h6a678d5_0
+  - ffmpeg=4.3=hf484d3e_0
+  - filelock=3.13.1=py38h06a4308_0
+  - findlibs=0.0.5=pyhd8ed1ab_0
+  - fontconfig=2.14.2=h14ed4e7_0
+  - fonttools=4.25.0=pyhd3eb1b0_0
+  - freeglut=3.4.0=h6a678d5_0
+  - freetype=2.12.1=h4a9f257_0
+  - geos=3.8.0=he6710b0_0
+  - glib=2.78.4=h6a678d5_0
+  - glib-tools=2.78.4=h6a678d5_0
+  - gmp=6.2.1=h295c915_3
+  - gmpy2=2.1.2=py38heeb90bb_0
+  - gnutls=3.6.15=he1e5248_0
+  - gst-plugins-base=1.14.1=h6a678d5_1
+  - gstreamer=1.14.1=h5eee18b_1
+  - hdf4=4.2.15=h9772cbc_5
+  - hdf5=1.14.3=nompi_h4f84152_100
+  - icu=73.2=h59595ed_0
+  - idna=3.4=py38h06a4308_0
+  - importlib-metadata=7.0.1=py38h06a4308_0
+  - importlib_metadata=7.0.1=hd3eb1b0_0
+  - importlib_resources=6.1.1=py38h06a4308_1
+  - intel-openmp=2023.1.0=hdb19cb5_46306
+  - ipykernel=6.28.0=py38h06a4308_0
+  - ipython=8.12.2=py38h06a4308_0
+  - jasper=4.2.4=he283c48_0
+  - jedi=0.18.1=py38h06a4308_1
+  - jinja2=3.1.3=py38h06a4308_0
+  - jpeg=9e=h5eee18b_1
+  - json5=0.9.6=pyhd3eb1b0_0
+  - jsonschema=4.19.2=py38h06a4308_0
+  - jsonschema-specifications=2023.7.1=py38h06a4308_0
+  - jupyter-lsp=2.2.0=py38h06a4308_0
+  - jupyter_client=8.6.0=py38h06a4308_0
+  - jupyter_core=5.5.0=py38h06a4308_0
+  - jupyter_events=0.8.0=py38h06a4308_0
+  - jupyter_server=2.10.0=py38h06a4308_0
+  - jupyter_server_terminals=0.4.4=py38h06a4308_1
+  - jupyterlab=4.0.11=py38h06a4308_0
+  - jupyterlab_pygments=0.1.2=py_0
+  - jupyterlab_server=2.25.1=py38h06a4308_0
+  - kiwisolver=1.4.4=py38h6a678d5_0
+  - krb5=1.20.1=h143b758_1
+  - lame=3.100=h7b6447c_0
+  - lcms2=2.12=h3be6417_0
+  - ld_impl_linux-64=2.38=h1181459_1
+  - lerc=3.0=h295c915_0
+  - libaec=1.1.3=h59595ed_0
+  - libbrotlicommon=1.0.9=h5eee18b_7
+  - libbrotlidec=1.0.9=h5eee18b_7
+  - libbrotlienc=1.0.9=h5eee18b_7
+  - libclang=14.0.6=default_hc6dbbc7_1
+  - libclang13=14.0.6=default_he11475f_1
+  - libcublas=12.1.0.26=0
+  - libcufft=11.0.2.4=0
+  - libcufile=1.9.0.20=0
+  - libcups=2.4.2=h2d74bed_1
+  - libcurand=10.3.5.119=0
+  - libcurl=8.5.0=h251f7ec_0
+  - libcusolver=11.4.4.55=0
+  - libcusparse=12.0.2.55=0
+  - libdeflate=1.17=h5eee18b_1
+  - libedit=3.1.20230828=h5eee18b_0
+  - libev=4.33=h7f8727e_1
+  - libffi=3.4.4=h6a678d5_0
+  - libgcc-ng=13.2.0=h77fa898_7
+  - libgfortran-ng=13.2.0=h69a702a_7
+  - libgfortran5=13.2.0=hca663fb_7
+  - libglib=2.78.4=hdc74915_0
+  - libglu=9.0.0=hac7e632_1003
+  - libgomp=13.2.0=h77fa898_7
+  - libiconv=1.17=hd590300_2
+  - libidn2=2.3.4=h5eee18b_0
+  - libjpeg-turbo=2.0.0=h9bf148f_0
+  - libllvm14=14.0.6=hdb19cb5_3
+  - libnetcdf=4.9.2=nompi_h9612171_113
+  - libnghttp2=1.57.0=h2d74bed_0
+  - libnpp=12.0.2.50=0
+  - libnvjitlink=12.1.105=0
+  - libnvjpeg=12.1.1.14=0
+  - libpng=1.6.43=h2797004_0
+  - libpq=12.17=hdbd6064_0
+  - libsodium=1.0.18=h7b6447c_0
+  - libssh2=1.10.0=hdbd6064_2
+  - libstdcxx-ng=13.2.0=hc0a3c3a_7
+  - libtasn1=4.19.0=h5eee18b_0
+  - libtiff=4.5.1=h6a678d5_0
+  - libunistring=0.9.10=h27cfd23_0
+  - libuuid=2.38.1=h0b41bf4_0
+  - libwebp-base=1.3.2=h5eee18b_0
+  - libxcb=1.15=h7f8727e_0
+  - libxkbcommon=1.7.0=h662e7e4_0
+  - libxml2=2.12.7=hc051c1a_1
+  - libzip=1.10.1=h2629f0a_3
+  - libzlib=1.2.13=h4ab18f5_6
+  - llvm-openmp=14.0.6=h9e868ea_0
+  - lz4-c=1.9.4=h6a678d5_0
+  - markupsafe=2.1.3=py38h5eee18b_0
+  - matplotlib=3.7.2=py38h06a4308_0
+  - matplotlib-base=3.7.2=py38h1128e8f_0
+  - matplotlib-inline=0.1.6=py38h06a4308_0
+  - mistune=2.0.4=py38h06a4308_0
+  - mkl=2023.1.0=h213fc3f_46344
+  - mkl-service=2.4.0=py38h5eee18b_1
+  - mkl_fft=1.3.8=py38h5eee18b_0
+  - mkl_random=1.2.4=py38hdb19cb5_0
+  - mpc=1.1.0=h10f8cd9_1
+  - mpfr=4.0.2=hb69a4c5_1
+  - mpmath=1.3.0=py38h06a4308_0
+  - munkres=1.1.4=py_0
+  - mysql=5.7.24=h721c034_2
+  - nbclient=0.8.0=py38h06a4308_0
+  - nbconvert=7.10.0=py38h06a4308_0
+  - nbformat=5.9.2=py38h06a4308_0
+  - ncurses=6.4=h6a678d5_0
+  - nest-asyncio=1.6.0=py38h06a4308_0
+  - nettle=3.7.3=hbbd107a_1
+  - networkx=3.1=py38h06a4308_0
+  - notebook-shim=0.2.3=py38h06a4308_0
+  - numexpr=2.8.4=py38hc78ab66_1
+  - numpy=1.24.3=py38hf6e8229_1
+  - numpy-base=1.24.3=py38h060ed82_1
+  - openh264=2.1.1=h4ff587b_0
+  - openjpeg=2.4.0=h3ad879b_0
+  - openssl=3.3.1=h4ab18f5_0
+  - overrides=7.4.0=py38h06a4308_0
+  - packaging=23.1=py38h06a4308_0
+  - pandas=2.0.3=py38h1128e8f_0
+  - pandocfilters=1.5.0=pyhd3eb1b0_0
+  - parso=0.8.3=pyhd3eb1b0_0
+  - pcre2=10.42=hebb0a14_0
+  - pexpect=4.8.0=pyhd3eb1b0_3
+  - pickleshare=0.7.5=pyhd3eb1b0_1003
+  - pillow=10.2.0=py38h5eee18b_0
+  - pip=23.3.1=py38h06a4308_0
+  - pkgutil-resolve-name=1.3.10=py38h06a4308_1
+  - platformdirs=3.10.0=py38h06a4308_0
+  - ply=3.11=py38_0
+  - pooch=1.7.0=py38h06a4308_0
+  - proj=8.2.1=ha227179_0
+  - prometheus_client=0.14.1=py38h06a4308_0
+  - prompt-toolkit=3.0.43=py38h06a4308_0
+  - psutil=5.9.0=py38h5eee18b_0
+  - ptyprocess=0.7.0=pyhd3eb1b0_2
+  - pure_eval=0.2.2=pyhd3eb1b0_0
+  - pycparser=2.21=pyhd3eb1b0_0
+  - pygments=2.15.1=py38h06a4308_1
+  - pyparsing=3.0.9=py38h06a4308_0
+  - pyproj=3.4.1=py38hb83c990_0
+  - pyqt=5.15.10=py38h6a678d5_0
+  - pyqt5-sip=12.13.0=py38h5eee18b_0
+  - pyshp=2.3.1=py38h06a4308_0
+  - pysocks=1.7.1=py38h06a4308_0
+  - python=3.8.18=h955ad1f_0
+  - python-dateutil=2.8.2=pyhd3eb1b0_0
+  - python-eccodes=1.7.0=py38he82f83a_2
+  - python-fastjsonschema=2.16.2=py38h06a4308_0
+  - python-json-logger=2.0.7=py38h06a4308_0
+  - python-tzdata=2023.3=pyhd3eb1b0_0
+  - python_abi=3.8=2_cp38
+  - pytorch=2.3.0=py3.8_cuda12.1_cudnn8.9.2_0
+  - pytorch-cuda=12.1=ha16c6d3_5
+  - pytorch-mutex=1.0=cuda
+  - pytz=2023.3.post1=py38h06a4308_0
+  - pyyaml=6.0.1=py38h5eee18b_0
+  - pyzmq=25.1.2=py38h6a678d5_0
+  - qt-main=5.15.2=h53bd1ea_10
+  - readline=8.2=h5eee18b_0
+  - referencing=0.30.2=py38h06a4308_0
+  - requests=2.31.0=py38h06a4308_1
+  - rfc3339-validator=0.1.4=py38h06a4308_0
+  - rfc3986-validator=0.1.1=py38h06a4308_0
+  - rpds-py=0.10.6=py38hb02cf49_0
+  - scipy=1.10.1=py38hf6e8229_1
+  - send2trash=1.8.2=py38h06a4308_0
+  - setuptools=68.2.2=py38h06a4308_0
+  - shapely=2.0.1=py38h0fbb895_0
+  - sip=6.7.12=py38h6a678d5_0
+  - six=1.16.0=pyhd3eb1b0_1
+  - snappy=1.2.0=hdb0a2a9_1
+  - sniffio=1.3.0=py38h06a4308_0
+  - soupsieve=2.5=py38h06a4308_0
+  - sqlite=3.41.2=h5eee18b_0
+  - stack_data=0.2.0=pyhd3eb1b0_0
+  - sympy=1.12=py38h06a4308_0
+  - tbb=2021.8.0=hdb19cb5_0
+  - terminado=0.17.1=py38h06a4308_0
+  - tinycss2=1.2.1=py38h06a4308_0
+  - tk=8.6.12=h1ccaba5_0
+  - tomli=2.0.1=py38h06a4308_0
+  - torchaudio=2.3.0=py38_cu121
+  - torchtriton=2.3.0=py38
+  - torchvision=0.18.0=py38_cu121
+  - tornado=6.3.3=py38h5eee18b_0
+  - tqdm=4.65.0=py38hb070fc8_0
+  - traitlets=5.7.1=py38h06a4308_0
+  - typing-extensions=4.9.0=py38h06a4308_1
+  - typing_extensions=4.9.0=py38h06a4308_1
+  - urllib3=2.1.0=py38h06a4308_1
+  - wcwidth=0.2.5=pyhd3eb1b0_0
+  - webencodings=0.5.1=py38_1
+  - websocket-client=0.58.0=py38h06a4308_4
+  - wheel=0.41.2=py38h06a4308_0
+  - xarray=2022.11.0=py38h06a4308_0
+  - xkeyboard-config=2.42=h4ab18f5_0
+  - xorg-kbproto=1.0.7=h7f98852_1002
+  - xorg-libx11=1.8.9=h8ee46fc_0
+  - xorg-libxau=1.0.11=hd590300_0
+  - xorg-libxext=1.3.4=h0b41bf4_2
+  - xorg-xextproto=7.3.0=h0b41bf4_1003
+  - xorg-xproto=7.0.31=h7f98852_1007
+  - xz=5.4.6=h5eee18b_0
+  - yaml=0.2.5=h7b6447c_0
+  - zeromq=4.3.5=h6a678d5_0
+  - zipp=3.17.0=py38h06a4308_0
+  - zlib=1.2.13=h4ab18f5_6
+  - zstd=1.5.5=hc292b87_0
+  - pip:
+      - aiohttp==3.9.3
+      - aiosignal==1.3.1
+      - appdirs==1.4.4
+      - asciitree==0.3.3
+      - async-timeout==4.0.3
+      - black==24.4.2
+      - cachetools==5.3.3
+      - cfgrib==0.9.12.0
+      - cftime==1.6.3
+      - click==8.1.7
+      - click-plugins==1.1.1
+      - cligj==0.7.2
+      - docker-pycreds==0.4.0
+      - eumdac==2.2.2
+      - fasteners==0.19
+      - fiona==1.9.6
+      - frozenlist==1.4.1
+      - fsspec==2024.3.1
+      - gcsfs==2024.3.1
+      - geographiclib==2.0
+      - geopandas==0.13.2
+      - geopy==2.4.1
+      - gitdb==4.0.11
+      - gitpython==3.1.42
+      - google-api-core==2.17.1
+      - google-auth==2.28.2
+      - google-auth-oauthlib==1.2.0
+      - google-cloud-core==2.4.1
+      - google-cloud-storage==2.16.0
+      - google-crc32c==1.5.0
+      - google-resumable-media==2.7.0
+      - googleapis-common-protos==1.63.0
+      - h5py==3.10.0
+      - huggingface-hub==0.22.2
+      - multidict==6.0.5
+      - mypy-extensions==1.0.0
+      - netcdf4==1.6.5
+      - numcodecs==0.12.1
+      - oauthlib==3.2.2
+      - pathspec==0.12.1
+      - plotly==5.20.0
+      - protobuf==4.25.3
+      - pyasn1==0.5.1
+      - pyasn1-modules==0.3.0
+      - ratelimiter==1.2.0.post0
+      - requests-oauthlib==1.4.0
+      - rsa==4.9
+      - safetensors==0.4.3
+      - scienceplots==2.1.1
+      - sentry-sdk==1.42.0
+      - setproctitle==1.3.3
+      - smmap==5.0.1
+      - tenacity==8.2.3
+      - timm==0.6.13
+      - tokenize-rt==5.2.0
+      - wandb==0.16.4
+      - yarl==1.9.4
+      - zarr==2.16.1
+prefix: /home/azureuser/miniconda3/envs/npw

official-src/notebooks/data_demo.ipynb

@@ -0,0 +1,386 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "f329e4f1-6dd9-4f2e-b5b3-0e8454832eee",
+   "metadata": {},
+   "source": [
+    "# Dataset\n",
+    "This notebook provides an example of the data utilised to generate a forecast using Aardvark Weather. We explore a single timeslice containing all the observations required to generate a forecast. This sample data is the output of the loader `WeatherDatasetE2E` in `../aardvark/loaders.py`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "c17595b2-a624-440b-a2be-6af8d3802f87",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import pickle"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "bc8b869d-33d4-45e5-b3ef-15d4e521e3d5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import matplotlib\n",
+    "from matplotlib import pyplot as plt\n",
+    "\n",
+    "matplotlib.rcParams['mathtext.fontset'] = 'stix'\n",
+    "matplotlib.rcParams['font.family'] = 'STIXGeneral'\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6de4a11f-4677-4552-8630-6b395c27012e",
+   "metadata": {},
+   "source": [
+    "## Data\n",
+    "Open a sample of data required to produce a forecast."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "0483c01d-2e08-469c-b914-42532cc3099b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open('../data/sample_data_final.pkl', 'rb') as fp:\n",
+    "    data = pickle.load(fp)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "83b7c966-fab5-47ce-a91c-d5419acaee40",
+   "metadata": {},
+   "source": [
+    "Multiple different datasets are utilised as input to create a forecast, each with multiple channels including observations and metadata. Example channels for each of these are plotted below. The plot_channel variable in each cell can be adjusted to visualise different channels. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "693aa212-c3ab-4eda-b07a-a6b262da627d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Set up longitude and latitude for plotting\n",
+    "lon = np.linspace(0,359,360)\n",
+    "lat = np.linspace(90,-90,181)\n",
+    "\n",
+    "lon_1p5 = np.linspace(0,359,240)\n",
+    "lat_1p5 = np.linspace(90,-90,121)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "42a2524f-108a-4ed7-bc89-2fb91638a7aa",
+   "metadata": {},
+   "source": [
+    "## Satellite\n",
+    "First visualise the satellite data from HIRS, AMSU-A, AMSU-B, IASI and ASCAT"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1243f422-351a-458c-aaa9-4bf72a959678",
+   "metadata": {},
+   "source": [
+    "### HIRS"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0d750dcc-cfef-4ca5-8f14-a7bae45980f4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig = plt.figure()\n",
+    "plot_channel = 11\n",
+    "p = plt.contourf(\n",
+    "    lon,\n",
+    "    lat,\n",
+    "    data[\"assimilation\"][\"hirs_current\"][0,...,plot_channel].cpu().T, \n",
+    "    levels=100,\n",
+    "    cmap=\"magma\")\n",
+    "cbar = fig.colorbar(p)\n",
+    "cbar.set_label('Normalised radiance') \n",
+    "plt.xlabel(\"Longitude\")\n",
+    "plt.ylabel(\"Latitude\")\n",
+    "plt.title(f\"HIRS channel {plot_channel}\")\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "d5c5e48e-3bab-461f-8425-9efa10c415e0",
+   "metadata": {},
+   "source": [
+    "### AMSU-A"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "79deb284-7ed8-47d2-9006-a13a06d2b835",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig = plt.figure()\n",
+    "plot_channel = 8\n",
+    "p = plt.contourf(\n",
+    "    lon, \n",
+    "    lat[:-1],\n",
+    "    data[\"assimilation\"][\"amsua_current\"][0,...,plot_channel].cpu(), \n",
+    "    levels=100,\n",
+    "    cmap=\"magma\")\n",
+    "cbar = fig.colorbar(p)\n",
+    "cbar.set_label('Normalised radiance') \n",
+    "plt.xlabel(\"Longitude\")\n",
+    "plt.ylabel(\"Latitude\")\n",
+    "plt.title(f\"AMSU-A channel {plot_channel}\")\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9b481a5c-97ad-4016-826d-93771531ad2d",
+   "metadata": {},
+   "source": [
+    "### AMSU-B"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "466434e7-7296-4454-b14c-6ab40bfbe0d1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig = plt.figure()\n",
+    "plot_channel = 10\n",
+    "p = plt.contourf(\n",
+    "    lon,\n",
+    "    lat,\n",
+    "    data[\"assimilation\"][\"amsub_current\"][0,...,plot_channel].T.cpu(), \n",
+    "    levels=100,\n",
+    "    cmap=\"magma\")\n",
+    "cbar = fig.colorbar(p)\n",
+    "cbar.set_label('Normalised radiance') \n",
+    "plt.xlabel(\"Longitude\")\n",
+    "plt.ylabel(\"Latitude\")\n",
+    "plt.title(f\"AMSU-A channel {plot_channel}\")\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "44de7c3f-6a25-47eb-bcfe-ab7bc82a1aac",
+   "metadata": {},
+   "source": [
+    "### IASI"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d5a5caf9-ecc4-403d-97ff-6c96dbc18f13",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig = plt.figure()\n",
+    "plot_channel = 10\n",
+    "p = plt.contourf(\n",
+    "    lon,\n",
+    "    lat,\n",
+    "    data[\"assimilation\"][\"iasi_current\"][0,...,plot_channel].T.cpu(), \n",
+    "    levels=100,\n",
+    "    cmap=\"magma\")\n",
+    "cbar = fig.colorbar(p)\n",
+    "cbar.set_label('Normalised radiance') \n",
+    "plt.xlabel(\"Longitude\")\n",
+    "plt.ylabel(\"Latitude\")\n",
+    "plt.title(f\"IASI channel {plot_channel}\")\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ea0ac772-4150-4bfd-96de-004ac7b8b6b4",
+   "metadata": {},
+   "source": [
+    "### ASCAT"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9981b9c6-2d88-47a1-b135-f70f1335df78",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig = plt.figure()\n",
+    "plot_channel = 5\n",
+    "p = plt.contourf(\n",
+    "    lon,\n",
+    "    lat,\n",
+    "    data[\"assimilation\"][\"ascat_current\"][0,...,plot_channel].T.cpu(), \n",
+    "    levels=100,\n",
+    "    cmap=\"magma\")\n",
+    "cbar = fig.colorbar(p)\n",
+    "cbar.set_label('Normalised radiance') \n",
+    "plt.xlabel(\"Longitude\")\n",
+    "plt.ylabel(\"Latitude\")\n",
+    "plt.title(f\"ASCAT channel {plot_channel}\")\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1b6849d1-c6fd-413f-a72e-3de111228a78",
+   "metadata": {},
+   "source": [
+    "## SYNOPS\n",
+    "We next visualise the SYNOPS data from land stations, marine platforms and radiosonde profiles"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c392b951-b446-4fde-aacb-02be7759bc97",
+   "metadata": {},
+   "source": [
+    "### HadISD"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "611eed9e-8c4a-4ac8-82f4-9969a6ef8eb7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig = plt.figure()\n",
+    "plot_channel = 0\n",
+    "p = plt.scatter(\n",
+    "    data[\"assimilation\"][\"x_context_hadisd_current\"][plot_channel][0,0,:].cpu(),\n",
+    "    data[\"assimilation\"][\"x_context_hadisd_current\"][plot_channel][0,1,:].cpu(),\n",
+    "    c = data[\"assimilation\"][\"y_context_hadisd_current\"][plot_channel][0].T.cpu(), \n",
+    "    cmap=\"magma\")\n",
+    "cbar = fig.colorbar(p)\n",
+    "cbar.set_label('Normalised value') \n",
+    "plt.xlabel(\"Longitude\")\n",
+    "plt.ylabel(\"Latitude\")\n",
+    "plt.title(f\"HadISD channel {plot_channel}\")\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b2ae8c40-28bd-49ad-8a2f-b9f7df706834",
+   "metadata": {},
+   "source": [
+    "### ICOADS"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1e425818-5a30-4641-abc7-08cfd4c6b4fe",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plot_channel = 1\n",
+    "fig = plt.figure()\n",
+    "p = plt.scatter(\n",
+    "    data[\"assimilation\"][\"icoads_x_current\"][0][0,:].cpu(),\n",
+    "    data[\"assimilation\"][\"icoads_x_current\"][1][0,:].cpu(),\n",
+    "    c = data[\"assimilation\"][\"icoads_current\"][0,plot_channel,:].T.cpu(), \n",
+    "    cmap=\"magma\")\n",
+    "cbar = fig.colorbar(p)\n",
+    "cbar.set_label('Normalised value') \n",
+    "plt.xlabel(\"Longitude\")\n",
+    "plt.ylabel(\"Latitude\")\n",
+    "plt.title(f\"ICOADS channel {plot_channel}\")\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "14d103a9-053c-4576-b17f-35ebae3fa44a",
+   "metadata": {},
+   "source": [
+    "### IGRA"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "11d5f5de-22ea-4d6e-a590-c1ce0a842a53",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plot_channel = 1\n",
+    "fig = plt.figure()\n",
+    "p = plt.scatter(\n",
+    "    data[\"assimilation\"][\"igra_x_current\"][0][0,:].cpu(),\n",
+    "    data[\"assimilation\"][\"igra_x_current\"][1][0,:].cpu(),\n",
+    "    c = data[\"assimilation\"][\"igra_current\"][0,plot_channel,:].T.cpu(), \n",
+    "    cmap=\"magma\")\n",
+    "cbar = fig.colorbar(p)\n",
+    "cbar.set_label('Normalised value') \n",
+    "plt.xlabel(\"Longitude\")\n",
+    "plt.ylabel(\"Latitude\")\n",
+    "plt.title(f\"IGRA channel {plot_channel}\")\n",
+    "plt.show()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2b893a77-4c30-4f18-a7fb-b007a676d849",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "521c93b6-babf-41fe-8dc6-c9a6eaef5b6f",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.18"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "ff379d8b57ea10bc9abbe40475203c855480f57b5b4821ba344a2e4e1a8a455a"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

official-src/training/downscaling_run_pipeline.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+
+python3 train_module.py \
+    --output_dir $3 \
+    --master_port 22348 \
+    --decoder base \
+    --loss downscaling_rmse \
+    --diff 0 \
+    --in_channels 36 \
+    --out_channels 1 \
+    --int_channels 24 \
+    --mode downscaling \
+    --lr 5e-4 \
+    --batch_size 64 \
+    --start_ind 0 \
+    --end_ind 24 \
+    --epoch 20 \
+    --weight_decay 1e-6 \
+    --downscaling_context aardvark \
+    --downscaling_train_start_date 2007-01-02 \
+    --downscaling_train_end_date 2017-12-31 \
+    --lead_time $1 \
+    --var $2

official-src/training/finetune.sh

@@ -0,0 +1,9 @@
+python3 ../aardvark/generate_initial_condition_single.py \
+    --assimilation_model_path ENCODER_PATH
+
+python3 ../aardvark/finetune.py \
+    --assimilation_model_path ENCODER_PATH \
+    --forecast_model_path FORECAST_PATH \
+    --output_dir FINETUNE_PATH \
+    --lr 5e-5 \
+    --finetune_epochs 1

official-src/training/test/config.pkl

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

official-src/training/train_decoder.sh

@@ -0,0 +1,43 @@
+python3 ../aardvark/train_module.py \
+    --output_dir DECODER/TAS \
+    --master_port 22348 \
+    --decoder base \
+    --loss downscaling_rmse \
+    --diff 0 \
+    --in_channels 36 \
+    --out_channels 1 \
+    --int_channels 24 \
+    --mode downscaling \
+    --lr 5e-4 \
+    --batch_size 64 \
+    --start_ind 0 \
+    --end_ind 24 \
+    --epoch 20 \
+    --weight_decay 1e-6 \
+    --downscaling_context aardvark \
+    --downscaling_train_start_date 2007-01-02 \
+    --downscaling_train_end_date 2017-12-31 \
+    --lead_time 4 \
+    --var tas
+
+python3 ../aardvark/train_module.py \
+    --output_dir DECODER/WS \
+    --master_port 22348 \
+    --decoder base \
+    --loss downscaling_rmse \
+    --diff 0 \
+    --in_channels 36 \
+    --out_channels 1 \
+    --int_channels 24 \
+    --mode downscaling \
+    --lr 5e-4 \
+    --batch_size 64 \
+    --start_ind 0 \
+    --end_ind 24 \
+    --epoch 20 \
+    --weight_decay 1e-6 \
+    --downscaling_context aardvark \
+    --downscaling_train_start_date 2007-01-02 \
+    --downscaling_train_end_date 2017-12-31 \
+    --lead_time 4 \
+    --var ws

official-src/training/train_e2e.sh

@@ -0,0 +1,21 @@
+# Command for training with wind speed variable (ws)
+python3 ../aardvark/e2e_train.py \
+    --output_dir E2E/WS \
+    --loss downscaling_rmse \
+    --region global \
+    --lead_time 1 \
+    --sf_model_path DECODER_PATH \
+    --se_model_path ENCODER_PATH \
+    --forecast_model_path PROCESSOR_PATH \
+    --var ws
+
+# Command for training with temperature variable (tas)
+python3 ../aardvark/e2e_train.py \
+    --output_dir E2E/TAS \
+    --loss downscaling_rmse \
+    --region global \
+    --lead_time 1 \
+    --sf_model_path DECODER_PATH \
+    --se_model_path ENCODER_PATH \
+    --forecast_model_path PROCESSOR_PATH \
+    --var tas

official-src/training/train_encoder.sh

@@ -0,0 +1,16 @@
+python3 ../aardvark/train_module.py \
+    --output_dir ENCODER_PATH \
+    --master_port 12348 \
+    --decoder vit_assimilation \
+    --loss lw_rmse \
+    --diff 0 \
+    --in_channels 277 \
+    --out_channels 24 \
+    --int_channels 24 \
+    --mode assimilation \
+    --lr 5e-4 \
+    --batch_size 6 \
+    --start_ind 0 \
+    --end_ind 24 \
+    --epoch 100 \
+    --weight_per_variable 1

official-src/training/train_processor.sh

@@ -0,0 +1,17 @@
+python3 ../aardvark/train_module.py \
+    --output_dir FORECAST_PATH/ \
+    --lead_time 1 \
+    --loss lw_rmse_pressure_weighted \
+    --diff 1 \
+    --mode forecast \
+    --in_channels 35 \
+    --out_channels 24 \
+    --int_channels 24 \
+    --decoder vit \
+    --ic era5 \
+    --film 0 \
+    --epoch 200 \
+    --weight_per_variable 0 \
+    --batch_size 24 \
+    --lr 5e-4 \
+    --era5_mode 4u

scripts/inference.py

@@ -0,0 +1,57 @@
+#!/usr/bin/env python
+"""Run the one-day tas forecast through the official Aardvark modules."""
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+import copy
+import torch
+
+from model.aardvark_adapter import build_one_day_model, load_sample
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--root", type=Path, default=ROOT)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--checkpoint", type=Path, help="Optional checkpoint produced by scripts/train.py")
+    parser.add_argument("--output", type=Path, default=ROOT / "result" / "inference_one_day.json")
+    args = parser.parse_args()
+    root = args.root.resolve()
+    sample_path = root / "weights/sample_data/sample_data_final.pkl"
+    sample = load_sample(sample_path)
+    model = build_one_day_model(root / "weights", root / "official-src", args.device)
+    if args.checkpoint:
+        tuned = args.checkpoint.resolve()
+        payload = torch.load(tuned, map_location=args.device, weights_only=False)
+        target_model = model if payload["train_modules"] == "all" else model.sf_model
+        target_model.load_state_dict(payload["model"])
+    model.eval()
+    target = sample["y_target"].cpu()
+    with torch.inference_mode():
+        station, global_forecast, initial_state = model(copy.deepcopy(sample))
+    result_dir = root / "result"
+    result_dir.mkdir(parents=True, exist_ok=True)
+    torch.save(station.cpu(), result_dir / "prediction.pt")
+    torch.save(target, result_dir / "target.pt")
+    report = {
+        "device": args.device,
+        "lead_time_days": 1,
+        "station_tas_shape": list(station.shape),
+        "global_forecast_shape": list(global_forecast.shape),
+        "initial_state_shape": list(initial_state.shape),
+        "finite_outputs": bool(torch.isfinite(station).all()),
+        "weights": str(tuned) if args.checkpoint else "official",
+    }
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n")
+    print(json.dumps(report, indent=2, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()