Upload folder using huggingface_hub

README.md

@@ -0,0 +1,151 @@
+---
+frameworks: JAX
+language:
+- en
+license: apache-2.0
+tags:
+- OneScience
+- Earth Science
+- Weather Forecasting
+- Ensemble Forecasting
+- ERA5
+tasks: []
+datasets:
+  - OneScience/ERA5
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">GenCast</span>
+  </strong>
+</p>
+
+# Model Overview
+
+GenCast is a probabilistic global weather forecasting model developed by Google DeepMind. Its paper appeared as the cover article of the leading scientific journal *Nature* on December 4, 2024.
+
+Paper: *GenCast: Diffusion-Based Ensemble Forecasting for Medium-Range Weather*
+
+https://arxiv.org/abs/2312.15796
+
+# Model Description
+
+GenCast is an ensemble forecasting model built with graph neural networks and diffusion models. Across a comprehensive set of evaluations, it outperformed ENS, the European Centre for Medium-Range Weather Forecasts' (ECMWF) leading ensemble forecasting system.
+
+
+# Use Cases
+
+| Use Case | Description |
+| :---: | :--- |
+| Weather forecasting training | Train the model on ERA5 data in HDF5 format that conforms to the GenCast data protocol. |
+| Quick local validation | Use synthetic data to validate data loading, model training and inference, and visualization of inference results. |
+| ModelScope/OneCode execution | Download the standalone model package, install its dependencies, and run the included scripts directly. |
+| Multi-GPU training | Use JAX `pmap` for data-parallel training across multiple GPUs or accelerators on a single host. |
+
+
+# Usage
+
+## 1. Using OneCode
+
+Use the OneCode online environment for an intelligent, one-click AI4S development experience:
+
+[Try one-click AI4S development with OneCode](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Setup
+
+**Hardware Requirements**
+
+- A GPU or DCU is recommended.
+- A CPU can be used for import checks and connectivity validation with a minimal configuration, but full training and inference will be slow.
+- DCU users must install DTK in advance. DTK 25.04.2 or later is recommended; alternatively, use the OneScience-recommended version compatible with your cluster.
+
+### Download the Model Package
+
+```bash
+hf download --model OneScience-Group/GenCast --local-dir ./GenCast
+cd GenCast
+```
+
+### Set Up the Runtime Environment
+
+**DCU Environment**
+
+```bash
+# Activate DTK and conda first.
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# Installation with uv is also supported.
+pip install onescience[earth-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+**GPU Environment**
+
+```bash
+# Activate conda first.
+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
+# Installation with uv is also supported.
+pip install onescience[earth-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+### Training Data
+
+The OneScience community provides ERA5 data for training. Because of file-size constraints, the repository currently contains a self-contained data slice. Download the data with the following command and ensure that the data path in `conf/config.yaml` is configured correctly:
+
+```bash
+hf download --dataset OneScience-Group/ERA5 --local-dir ./data
+```
+
+
+### Training
+
+Single GPU:
+
+```bash
+# If real data is unavailable, first run `python scripts/fake_data.py` to generate synthetic data.
+python scripts/train.py
+```
+
+Multiple GPUs:
+
+```bash
+CUDA_VISIBLE_DEVICES=0,1 python scripts/train.py --config conf/config.yaml --parallel-mode pmap --num-devices 2 --global-batch-size 2
+# CUDA_VISIBLE_DEVICES specifies the GPU indices to expose.
+# --num-devices specifies the number of GPUs to use.
+# --global-batch-size specifies the batch size and must be divisible by the number of GPUs.
+```
+
+After training, the weights are saved to `data/checkpoints/model_bak.npz`.
+
+
+### Pre-trained Weights
+
+This repository will provide weights trained on ERA5 reanalysis data in the `weights/` directory. The weight files are being prepared and will be uploaded soon.
+
+
+### Inference
+
+By default, inference loads `data/checkpoints/model_bak.npz`:
+
+```bash
+python scripts/inference.py
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+
+# Official OneScience Resources
+
+| 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 and License
+
+- This repository is a reproduction of the original GenCast paper.

conf/config.yaml

@@ -0,0 +1,76 @@
+runtime:
+  platform: "auto"
+
+data:
+  data_dir: "./data"
+  static_dir: "./data/static"
+  stats_dir: "./data/stats"
+  train_years: [2000, 2001]
+  test_years: [2003]
+  train_stride: 1
+  test_stride: 1
+  reintroduce_sst_nans: true
+  # Source total_precipitation must be accumulated over each source interval,
+  # in the same units as the official GenCast statistics.
+  precipitation_interval_hours: 6
+
+# Full-scale random initialization. For released weights, inference ignores
+# this section and loads every architecture field from the official checkpoint.
+model:
+  mesh_size: 4  # GenCast Mini (2562 nodes); 6=full (40962 nodes, TPU-scale)
+  latent_size: 512
+  hidden_layers: 1
+  radius_query_fraction_edge_length: 0.6
+  attention_k_hop: 16
+  attention_type: "triblockdiag_mha"  # GPU/CPU; use "splash_mha" only for TPU
+  mask_type: "lazy"
+  num_layers: 16
+  num_heads: 4
+  ffw_hidden: 2048
+
+sampler:
+  max_noise_level: 80.0
+  min_noise_level: 0.03
+  num_noise_levels: 20
+  rho: 7.0
+  stochastic_churn_rate: 2.5
+  churn_min_noise_level: 0.75
+  churn_max_noise_level: .inf
+  noise_level_inflation_factor: 1.05
+
+training:
+  max_steps: 10
+  learning_rate: 0.0001
+  betas: [0.9, 0.999]
+  epsilon: 1.0e-8
+  seed: 42
+  save_interval: 1000
+
+parallel:
+  mode: "pmap"
+  num_devices: 1
+  global_batch_size: 1
+  axis_name: "devices"
+
+inference:
+  official_checkpoint: null
+  # Officially documented GPU substitution for TPU splash attention.
+  attention_type_override: null
+  prediction_steps: 30
+  num_members: 4
+  seed: 42
+  # Keep full-resolution ensemble memory bounded by writing each lead/member.
+  stream_chunks: true
+
+checkpoint:
+  trainer: "./data/checkpoints/model_bak.npz"
+  resume: null
+
+output:
+  prediction: "./result/prediction.nc"
+  plot: "./result/gencast_forecast.png"
+
+fake_data:
+  height: 9
+  width: 16
+  timesteps: 8

configuration.json

@@ -0,0 +1 @@
+{"framework":"Pytorch","task":"other"}

model/common.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import xarray
+import yaml
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+
+
+def load_config(path: str | Path) -> dict[str, Any]:
+    with Path(path).open(encoding="utf-8") as source:
+        return yaml.safe_load(source)
+
+
+def resolve_path(path: str | Path) -> Path:
+    candidate = Path(path).expanduser()
+    return candidate if candidate.is_absolute() else PROJECT_ROOT / candidate
+
+
+def configure_jax(platform: str) -> None:
+    if platform != "auto" and "JAX_PLATFORM_NAME" not in os.environ:
+        os.environ["JAX_PLATFORM_NAME"] = platform
+
+
+def load_stats(stats_dir: str | Path) -> dict[str, xarray.Dataset]:
+    directory = resolve_path(stats_dir)
+    names = (
+        "diffs_stddev_by_level",
+        "mean_by_level",
+        "stddev_by_level",
+        "min_by_level",
+    )
+    stats = {}
+    for name in names:
+        path = directory / f"{name}.nc"
+        if not path.exists():
+            raise FileNotFoundError(f"Missing GenCast statistic: {path}")
+        stats[name] = xarray.load_dataset(path).compute()
+    from model.graphcast import gencast, graphcast
+
+    inputs = set(gencast.TASK.input_variables) - set(graphcast.GENERATED_FORCING_VARS)
+    targets = set(gencast.TASK.target_variables)
+    required_by_stat = {
+        "mean_by_level": inputs | (targets - inputs),
+        "stddev_by_level": inputs | (targets - inputs),
+        "diffs_stddev_by_level": targets & inputs,
+        "min_by_level": {"sea_surface_temperature"},
+    }
+    for stat_name, dataset in stats.items():
+        missing = sorted(required_by_stat[stat_name] - set(dataset.data_vars))
+        if missing:
+            raise ValueError(f"{stat_name} is missing GenCast variables: {missing}")
+        for name, values in dataset.data_vars.items():
+            array = np.asarray(values)
+            if not np.all(np.isfinite(array)):
+                raise ValueError(f"{stat_name}.{name} contains non-finite values")
+            if "level" in values.dims and tuple(values.level.values) != tuple(
+                gencast.TASK.pressure_levels
+            ):
+                raise ValueError(f"{stat_name}.{name} does not use GenCast WB13 order")
+            if stat_name in ("stddev_by_level", "diffs_stddev_by_level") and np.any(array <= 0):
+                raise ValueError(f"{stat_name}.{name} must be strictly positive")
+    return stats
+
+
+def save_trainer_checkpoint(
+    path: str | Path,
+    *,
+    params: Any,
+    state: Any,
+    optimizer_state: Any,
+    step: int,
+    config: dict[str, Any],
+) -> None:
+    import jax
+
+    destination = resolve_path(path)
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    leaves, treedef = jax.tree_util.tree_flatten(
+        {"params": params, "state": state, "optimizer_state": optimizer_state}
+    )
+    arrays = {f"leaf_{i}": np.asarray(value) for i, value in enumerate(leaves)}
+    arrays["treedef"] = np.array([treedef], dtype=object)
+    arrays["step"] = np.asarray(step, dtype=np.int64)
+    arrays["config_json"] = np.asarray(json.dumps(config, sort_keys=True))
+    temporary = destination.with_suffix(destination.suffix + ".tmp")
+    with temporary.open("wb") as output:
+        np.savez(output, **arrays)
+    os.replace(temporary, destination)
+
+
+def load_trainer_checkpoint(
+    path: str | Path,
+) -> tuple[Any, Any, Any, int, dict[str, Any]]:
+    import jax
+
+    source_path = resolve_path(path)
+    with np.load(source_path, allow_pickle=True) as source:
+        treedef = source["treedef"].item()
+        leaves = [source[f"leaf_{i}"] for i in range(len(source.files) - 3)]
+        tree = jax.tree_util.tree_unflatten(treedef, leaves)
+        saved_config = json.loads(str(source["config_json"]))
+        return (
+            tree["params"], tree["state"], tree["optimizer_state"],
+            int(source["step"]), saved_config,
+        )
+
+
+def validate_checkpoint_config(
+    current: dict[str, Any],
+    saved: dict[str, Any],
+    *,
+    scope: str = "resume",
+) -> None:
+    """Validate checkpoint compatibility for training resume or inference."""
+    if scope not in ("resume", "inference"):
+        raise ValueError("scope must be 'resume' or 'inference'")
+
+    inference_paths = (
+        ("model",), ("sampler",), ("data", "stats_dir"),
+        ("data", "static_dir"), ("data", "precipitation_interval_hours"),
+    )
+    resume_only_paths = (
+        ("training", "learning_rate"),
+        ("training", "betas"), ("training", "epsilon"),
+        ("training", "seed"), ("data", "data_dir"),
+        ("data", "train_years"), ("data", "train_stride"),
+        ("parallel", "mode"), ("parallel", "num_devices"),
+        ("parallel", "global_batch_size"), ("parallel", "axis_name"),
+    )
+    if scope == "resume":
+        if "parallel" not in saved:
+            saved = dict(saved)
+            saved["parallel"] = {
+                "mode": "single",
+                "num_devices": 1,
+                "global_batch_size": 1,
+                "axis_name": "devices",
+            }
+        paths = inference_paths + resume_only_paths
+    else:
+        paths = inference_paths
+
+    for path in paths:
+        current_value: Any = current
+        saved_value: Any = saved
+        for key in path:
+            current_value = current_value[key]
+            saved_value = saved_value[key]
+        if current_value != saved_value:
+            name = ".".join(path)
+            raise ValueError(
+                f"Trainer checkpoint configuration mismatch for {name} "
+                f"during {scope}: "
+                f"saved={saved_value!r}, current={current_value!r}"
+            )

model/data_loader.py

@@ -0,0 +1,359 @@
+"""将 ERA5 HDF5 严格适配为官方 GenCast xarray 数据协议。"""
+
+from __future__ import annotations
+
+import bisect
+import datetime as dt
+from pathlib import Path
+from typing import Any, Iterator
+
+import h5py
+import numpy as np
+import xarray
+
+try:
+    from onescience.datapipes.climate import ERA5Dataset as _ERA5Dataset
+except ModuleNotFoundError as error:
+    if error.name not in ("torch", "onescience"):
+        raise
+
+    class _ERA5Dataset:
+        """Minimal discovery fallback for JAX-only OneScience environments."""
+
+        def __init__(self, dataset_dir, used_years, used_variables, **_):
+            self.dataset_dir = dataset_dir
+            self.used_years = used_years
+            self.used_variables = used_variables
+            self._init_avail_samples()
+            self._init_normalized_files()
+
+        def _init_avail_samples(self):
+            data_dir = Path(self.dataset_dir) / "data"
+            available = {int(path.stem): path for path in data_dir.glob("*.h5")}
+            missing_years = sorted(set(self.used_years) - set(available))
+            if missing_years:
+                raise ValueError(f"Years not found in dataset: {missing_years}")
+            first = available[self.used_years[0]]
+            with h5py.File(first, "r") as source:
+                fields = source["fields"]
+                variables = [
+                    value.decode() if isinstance(value, bytes) else str(value)
+                    for value in fields.attrs["variables"]
+                ]
+                self.T, self.C, self.H, self.W = fields.shape
+                self.time_step = int(fields.attrs["time_step"])
+            missing_variables = sorted(set(self.used_variables) - set(variables))
+            if missing_variables:
+                raise ValueError(f"Variables not found in dataset: {missing_variables}")
+            self.file_map = {year: str(available[year]) for year in self.used_years}
+
+        def _init_normalized_files(self):
+            pass
+
+from model.graphcast import data_utils
+from model.graphcast import gencast
+from model.graphcast import graphcast
+
+
+PRESSURE_LEVELS = tuple(graphcast.PRESSURE_LEVELS_WEATHERBENCH_13)
+SURFACE_VARIABLES = tuple(gencast.TARGET_SURFACE_NO_PRECIP_VARS)
+ATMOSPHERIC_VARIABLES = tuple(graphcast.TARGET_ATMOSPHERIC_VARS)
+STATIC_VARIABLES = tuple(graphcast.STATIC_VARS)
+RAW_PRECIPITATION = "total_precipitation"
+TARGET_PRECIPITATION = "total_precipitation_12hr"
+MODEL_TARGET_CHANNELS = 6 + 6 * len(PRESSURE_LEVELS)
+ERA5_VARIABLES = (
+    *SURFACE_VARIABLES,
+    RAW_PRECIPITATION,
+    *(f"{name}_{level}" for name in ATMOSPHERIC_VARIABLES for level in PRESSURE_LEVELS),
+)
+
+
+def expected_target_channel_names() -> tuple[str, ...]:
+    channels: list[str] = []
+    atmospheric = set(ATMOSPHERIC_VARIABLES)
+    for name in sorted(gencast.TASK.target_variables):
+        if name in atmospheric:
+            channels.extend(f"{name}_{level}" for level in PRESSURE_LEVELS)
+        else:
+            channels.append(name)
+    return tuple(channels)
+
+
+class GenCastERA5Dataset(_ERA5Dataset):
+    """Reuse ERA5Dataset discovery while enforcing GenCast's named protocol."""
+
+    def __init__(
+        self,
+        dataset_dir: str | Path,
+        used_years: list[int],
+        *,
+        static_dir: str | Path | None = None,
+        prediction_steps: int = 1,
+        stride: int = 1,
+        task_config: Any = gencast.TASK,
+        precipitation_interval_hours: int = 6,
+        load_future_targets: bool = True,
+    ) -> None:
+        super().__init__(
+            dataset_dir=str(dataset_dir),
+            used_years=used_years,
+            used_variables=list(ERA5_VARIABLES),
+            input_steps=1,
+            output_steps=1,
+            normalize=False,
+        )
+        self.static_dir = Path(static_dir or Path(dataset_dir) / "static")
+        self.prediction_steps = int(prediction_steps)
+        self.stride = int(stride)
+        self.task_config = task_config
+        self.precipitation_interval_hours = int(precipitation_interval_hours)
+        self.load_future_targets = bool(load_future_targets)
+        self._validate_task_config()
+        if self.prediction_steps < 1 or self.stride < 1:
+            raise ValueError("prediction_steps and stride must be positive")
+        self._inspect_years()
+
+    def _init_normalized_files(self) -> None:
+        # GenCast uses named by-level NetCDF statistics in the model wrapper.
+        self.mu = self.sd = None
+
+    def _inspect_years(self) -> None:
+        self._year_meta: list[dict[str, Any]] = []
+        self._cumulative: list[int] = []
+        total = 0
+        for year in self.used_years:
+            path = Path(self.file_map[year])
+            with h5py.File(path, "r") as source:
+                fields = source["fields"]
+                variables = [
+                    value.decode() if isinstance(value, bytes) else str(value)
+                    for value in fields.attrs["variables"]
+                ]
+                time_step = int(fields.attrs["time_step"])
+                shape = tuple(fields.shape)
+            if time_step not in (6, 12):
+                raise ValueError(f"{path}: GenCast requires 6h or 12h ERA5, got {time_step}h")
+            if self.precipitation_interval_hours != time_step:
+                raise ValueError(
+                    f"{path}: total_precipitation must be an accumulation over each "
+                    f"{time_step}h source interval; configured "
+                    f"{self.precipitation_interval_hours}h"
+                )
+            missing = sorted(set(ERA5_VARIABLES) - set(variables))
+            if missing:
+                raise ValueError(f"{path}: missing GenCast ERA5 variables: {missing}")
+            frame_stride = 12 // time_step
+            # The -12h input also needs a complete 12h precipitation window.
+            first_reference = 2 * frame_stride - 1
+            last_reference = (
+                shape[0] - frame_stride * self.prediction_steps - 1
+                if self.load_future_targets else shape[0] - 1
+            )
+            references = list(range(first_reference, last_reference + 1, self.stride))
+            meta = {
+                "year": year,
+                "path": path,
+                "shape": shape,
+                "time_step": time_step,
+                "frame_stride": frame_stride,
+                "variables": variables,
+                "references": references,
+            }
+            self._year_meta.append(meta)
+            total += len(references)
+            self._cumulative.append(total)
+        self.total_samples = total
+        if not total:
+            raise ValueError("No complete GenCast samples are available")
+
+    def __len__(self) -> int:
+        return self.total_samples
+
+    def __getitem__(self, index: int):
+        if index < 0:
+            index += len(self)
+        if index < 0 or index >= len(self):
+            raise IndexError(index)
+        year_index = bisect.bisect_right(self._cumulative, index)
+        start = 0 if year_index == 0 else self._cumulative[year_index - 1]
+        meta = self._year_meta[year_index]
+        reference_index = meta["references"][index - start]
+        dataset = self._read_dataset(meta, reference_index)
+        return data_utils.extract_inputs_targets_forcings(
+            dataset,
+            target_lead_times=slice("12h", f"{12 * self.prediction_steps}h"),
+            input_variables=self.task_config.input_variables,
+            target_variables=self.task_config.target_variables,
+            forcing_variables=self.task_config.forcing_variables,
+            pressure_levels=self.task_config.pressure_levels,
+            input_duration=self.task_config.input_duration,
+        )
+
+    def _read_dataset(self, meta: dict[str, Any], reference_index: int) -> xarray.Dataset:
+        frame_stride = meta["frame_stride"]
+        frame_indices = [
+            reference_index - frame_stride,
+            reference_index,
+            *(reference_index + frame_stride * step for step in range(1, self.prediction_steps + 1)),
+        ]
+        variable_index = {name: i for i, name in enumerate(meta["variables"])}
+        selected_names = list(SURFACE_VARIABLES) + [
+            f"{name}_{level}"
+            for name in ATMOSPHERIC_VARIABLES
+            for level in PRESSURE_LEVELS
+        ]
+        selected_indices = [variable_index[name] for name in selected_names]
+        order = np.argsort(selected_indices)
+        inverse = np.empty(len(order), dtype=np.int64)
+        inverse[order] = np.arange(len(order))
+        read_count = len(frame_indices) if self.load_future_targets else 2
+        with h5py.File(meta["path"], "r") as source:
+            fields = source["fields"]
+            loaded = np.stack([
+                fields[t, np.asarray(selected_indices)[order], :, :][inverse]
+                for t in frame_indices[:read_count]
+            ]).astype(np.float32)
+            values = np.full(
+                (len(frame_indices), *loaded.shape[1:]), np.nan, dtype=np.float32
+            )
+            values[:read_count] = loaded
+            precipitation = np.full(
+                (len(frame_indices), *loaded.shape[-2:]), np.nan, dtype=np.float32
+            )
+            if self.load_future_targets:
+                precipitation[:] = np.stack([
+                    self._precipitation_12h(
+                        fields, variable_index[RAW_PRECIPITATION], t, frame_stride
+                    )
+                    for t in frame_indices
+                ]).astype(np.float32)
+
+        # OneScience ERA5 uses north-to-south storage; GenCast spherical noise requires ascending lat.
+        values = values[..., ::-1, :]
+        precipitation = precipitation[..., ::-1, :]
+        height, width = values.shape[-2:]
+        lat = np.linspace(-90.0, 90.0, height, dtype=np.float32)
+        lon = np.linspace(0.0, 360.0, width, endpoint=False, dtype=np.float32)
+        reference_time = dt.datetime(meta["year"], 1, 1) + dt.timedelta(
+            hours=reference_index * meta["time_step"]
+        )
+        datetimes = np.asarray([
+            np.datetime64(reference_time + dt.timedelta(hours=(t - reference_index) * meta["time_step"]))
+            for t in frame_indices
+        ], dtype="datetime64[ns]")
+        times = np.asarray([
+            np.timedelta64((t - reference_index) * meta["time_step"], "h")
+            for t in frame_indices
+        ], dtype="timedelta64[ns]")
+
+        data_vars: dict[str, Any] = {}
+        cursor = 0
+        for name in SURFACE_VARIABLES:
+            data_vars[name] = (("batch", "time", "lat", "lon"), values[:, cursor][None])
+            cursor += 1
+        for name in ATMOSPHERIC_VARIABLES:
+            data_vars[name] = (
+                ("batch", "time", "level", "lat", "lon"),
+                values[:, cursor:cursor + len(PRESSURE_LEVELS)][None],
+            )
+            cursor += len(PRESSURE_LEVELS)
+        data_vars[TARGET_PRECIPITATION] = (
+            ("batch", "time", "lat", "lon"), precipitation[None]
+        )
+        data_vars.update(self._load_static(height, width))
+        dataset = xarray.Dataset(
+            data_vars=data_vars,
+            coords={
+                "batch": np.arange(1),
+                "time": times,
+                "datetime": (("batch", "time"), datetimes[None]),
+                "level": np.asarray(PRESSURE_LEVELS, dtype=np.int32),
+                "lat": lat,
+                "lon": lon,
+            },
+        )
+        dataset.attrs["forecast_reference_time"] = np.datetime_as_string(
+            np.datetime64(reference_time), unit="h"
+        )
+        self.validate_dataset(dataset)
+        return dataset
+
+    @staticmethod
+    def _precipitation_12h(fields, channel: int, end: int, frame_stride: int):
+        start = end - frame_stride + 1
+        if start < 0:
+            raise IndexError("Insufficient precipitation history for 12h accumulation")
+        return np.sum(fields[start:end + 1, channel], axis=0)
+
+    def _load_static(self, height: int, width: int) -> dict[str, Any]:
+        paths = {
+            "geopotential_at_surface": self.static_dir / "geopotential_at_surface.npy",
+            "land_sea_mask": self.static_dir / "land_mask.npy",
+        }
+        result = {}
+        for name, path in paths.items():
+            if not path.exists():
+                raise FileNotFoundError(f"Missing GenCast static field: {path}")
+            values = np.load(path).astype(np.float32)
+            if values.shape != (height, width):
+                raise ValueError(f"{path}: expected {(height, width)}, got {values.shape}")
+            result[name] = (("lat", "lon"), values[::-1])
+        return result
+
+    def _validate_task_config(self) -> None:
+        expected = gencast.TASK
+        for field in (
+            "input_variables", "target_variables", "forcing_variables",
+            "pressure_levels", "input_duration",
+        ):
+            if getattr(self.task_config, field) != getattr(expected, field):
+                raise ValueError(
+                    "This ERA5 adapter supports the official WB13 GenCast task "
+                    f"only; checkpoint field {field} differs"
+                )
+
+    @staticmethod
+    def validate_dataset(dataset: xarray.Dataset) -> None:
+        missing = sorted(
+            set(gencast.TASK.input_variables + gencast.TASK.target_variables)
+            - set(dataset.data_vars)
+            - set(graphcast.GENERATED_FORCING_VARS)
+        )
+        if missing:
+            raise ValueError(f"Missing GenCast variables: {missing}")
+        if tuple(int(level) for level in dataset.level.values) != PRESSURE_LEVELS:
+            raise ValueError("GenCast WB13 pressure-level order changed")
+        if not np.all(np.diff(dataset.lat.values) > 0):
+            raise ValueError("GenCast latitude must be strictly ascending")
+        height, width = dataset.sizes["lat"], dataset.sizes["lon"]
+        if width != 2 * (height - 1):
+            raise ValueError(
+                "GenCast equiangular grids with poles require lon=2*(lat-1), "
+                f"got lat={height}, lon={width}"
+            )
+        if len(expected_target_channel_names()) != MODEL_TARGET_CHANNELS:
+            raise AssertionError("The official GenCast target contract must contain 84 channels")
+
+
+def batch_iterator(
+    dataset: GenCastERA5Dataset,
+    *,
+    shuffle: bool,
+    seed: int,
+    batch_size: int = 1,
+) -> Iterator:
+    if batch_size < 1:
+        raise ValueError("batch_size must be positive")
+    indices = np.arange(len(dataset))
+    if shuffle:
+        np.random.default_rng(seed).shuffle(indices)
+    for start in range(0, len(indices) - batch_size + 1, batch_size):
+        samples = [dataset[int(index)] for index in indices[start:start + batch_size]]
+        if batch_size == 1:
+            yield samples[0]
+            continue
+        yield tuple(
+            xarray.concat(values, dim="batch", data_vars="minimal", coords="minimal")
+            for values in zip(*samples)
+        )

model/gencast.py

@@ -0,0 +1,194 @@
+"""项目内官方等价 GenCast JAX/Haiku 实现封装。"""
+
+from __future__ import annotations
+
+import dataclasses
+from pathlib import Path
+from typing import Any
+
+import haiku as hk
+import jax
+import xarray
+
+from model.graphcast import checkpoint
+from model.graphcast import denoiser
+from model.graphcast import gencast
+from model.graphcast import nan_cleaning
+from model.graphcast import normalization
+from model.graphcast import xarray_jax
+from model.graphcast import xarray_tree
+
+
+def build_model_config(config: dict[str, Any]) -> tuple[
+    Any, denoiser.DenoiserArchitectureConfig, gencast.SamplerConfig,
+    gencast.NoiseConfig, denoiser.NoiseEncoderConfig
+]:
+    """Build a random-weight configuration without changing GenCast semantics."""
+    model_cfg = config["model"]
+    sampler_cfg = config["sampler"]
+    transformer = denoiser.SparseTransformerConfig(
+        attention_k_hop=int(model_cfg["attention_k_hop"]),
+        d_model=int(model_cfg["latent_size"]),
+        num_layers=int(model_cfg["num_layers"]),
+        num_heads=int(model_cfg["num_heads"]),
+        attention_type=str(model_cfg["attention_type"]),
+        mask_type=str(model_cfg.get("mask_type", "full")),
+        ffw_hidden=int(model_cfg["ffw_hidden"]),
+    )
+    architecture = denoiser.DenoiserArchitectureConfig(
+        sparse_transformer_config=transformer,
+        mesh_size=int(model_cfg["mesh_size"]),
+        latent_size=int(model_cfg["latent_size"]),
+        hidden_layers=int(model_cfg.get("hidden_layers", 1)),
+        radius_query_fraction_edge_length=float(
+            model_cfg.get("radius_query_fraction_edge_length", 0.6)
+        ),
+    )
+    sampler = gencast.SamplerConfig(**sampler_cfg)
+    return (
+        gencast.TASK,
+        architecture,
+        sampler,
+        gencast.NoiseConfig(),
+        denoiser.NoiseEncoderConfig(),
+    )
+
+
+def load_model_checkpoint(path: str | Path) -> gencast.CheckPoint:
+    """Load the typed official GenCast NPZ checkpoint."""
+    with Path(path).open("rb") as source:
+        return checkpoint.load(source, gencast.CheckPoint)
+
+
+class GenCastModel:
+    """Owns official-equivalent GenCast loss and sampling Haiku transforms."""
+
+    def __init__(
+        self,
+        *,
+        task_config: Any,
+        architecture_config: denoiser.DenoiserArchitectureConfig,
+        sampler_config: gencast.SamplerConfig,
+        noise_config: gencast.NoiseConfig,
+        noise_encoder_config: denoiser.NoiseEncoderConfig,
+        diffs_stddev_by_level: xarray.Dataset,
+        mean_by_level: xarray.Dataset,
+        stddev_by_level: xarray.Dataset,
+        min_by_level: xarray.Dataset,
+        reintroduce_nans: bool = True,
+    ) -> None:
+        self.task_config = task_config
+        self.architecture_config = architecture_config
+        self.sampler_config = sampler_config
+        self.noise_config = noise_config
+        self.noise_encoder_config = noise_encoder_config
+        self.diffs_stddev_by_level = diffs_stddev_by_level
+        self.mean_by_level = mean_by_level
+        self.stddev_by_level = stddev_by_level
+        self.min_by_level = min_by_level
+        self.reintroduce_nans = reintroduce_nans
+
+        def construct() -> Any:
+            predictor = gencast.GenCast(
+                task_config=self.task_config,
+                denoiser_architecture_config=self.architecture_config,
+                sampler_config=self.sampler_config,
+                noise_config=self.noise_config,
+                noise_encoder_config=self.noise_encoder_config,
+            )
+            predictor = normalization.InputsAndResiduals(
+                predictor,
+                diffs_stddev_by_level=self.diffs_stddev_by_level,
+                mean_by_level=self.mean_by_level,
+                stddev_by_level=self.stddev_by_level,
+            )
+            return nan_cleaning.NaNCleaner(
+                predictor,
+                var_to_clean="sea_surface_temperature",
+                fill_value=self.min_by_level,
+                reintroduce_nans=self.reintroduce_nans,
+            )
+
+        @hk.transform_with_state
+        def loss_fn(inputs, targets, forcings):
+            loss, diagnostics = construct().loss(inputs, targets, forcings)
+            return xarray_tree.map_structure(
+                lambda value: xarray_jax.unwrap_data(
+                    value.mean(), require_jax=True
+                ),
+                (loss, diagnostics),
+            )
+
+        @hk.transform_with_state
+        def forward_fn(inputs, targets_template, forcings):
+            return construct()(
+                inputs,
+                targets_template=targets_template,
+                forcings=forcings,
+            )
+
+        self.loss_fn = loss_fn
+        self.forward_fn = forward_fn
+
+    @classmethod
+    def from_config_and_stats(
+        cls, config: dict[str, Any], stats: dict[str, xarray.Dataset]
+    ) -> "GenCastModel":
+        configs = build_model_config(config)
+        return cls(
+            task_config=configs[0],
+            architecture_config=configs[1],
+            sampler_config=configs[2],
+            noise_config=configs[3],
+            noise_encoder_config=configs[4],
+            diffs_stddev_by_level=stats["diffs_stddev_by_level"],
+            mean_by_level=stats["mean_by_level"],
+            stddev_by_level=stats["stddev_by_level"],
+            min_by_level=stats["min_by_level"],
+            reintroduce_nans=bool(config.get("data", {}).get("reintroduce_sst_nans", True)),
+        )
+
+    @classmethod
+    def from_checkpoint_and_stats(
+        cls,
+        model_checkpoint: gencast.CheckPoint,
+        stats: dict[str, xarray.Dataset],
+        *,
+        attention_type: str | None = None,
+    ) -> "GenCastModel":
+        architecture = model_checkpoint.denoiser_architecture_config
+        if attention_type is not None:
+            architecture = dataclasses.replace(
+                architecture,
+                sparse_transformer_config=dataclasses.replace(
+                    architecture.sparse_transformer_config,
+                    attention_type=attention_type,
+                    mask_type="full",
+                ),
+            )
+        return cls(
+            task_config=model_checkpoint.task_config,
+            architecture_config=architecture,
+            sampler_config=model_checkpoint.sampler_config,
+            noise_config=model_checkpoint.noise_config,
+            noise_encoder_config=model_checkpoint.noise_encoder_config,
+            diffs_stddev_by_level=stats["diffs_stddev_by_level"],
+            mean_by_level=stats["mean_by_level"],
+            stddev_by_level=stats["stddev_by_level"],
+            min_by_level=stats["min_by_level"],
+        )
+
+    def init(self, rng, inputs, targets, forcings):
+        return self.loss_fn.init(rng, inputs, targets, forcings)
+
+    def loss(self, params, state, rng, inputs, targets, forcings):
+        return self.loss_fn.apply(params, state, rng, inputs, targets, forcings)
+
+    def predict(self, params, state, rng, inputs, targets_template, forcings):
+        return self.forward_fn.apply(
+            params, state, rng, inputs, targets_template, forcings
+        )
+
+
+def parameter_count(params: Any) -> int:
+    return sum(int(value.size) for value in jax.tree_util.tree_leaves(params))

model/graphcast/autoregressive.py

@@ -0,0 +1,312 @@
+# Copyright 2023 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""A Predictor wrapping a one-step Predictor to make autoregressive predictions.
+"""
+
+from typing import Optional, cast
+
+from absl import logging
+from . import predictor_base
+from . import xarray_jax
+from . import xarray_tree
+import haiku as hk
+import jax
+import xarray
+
+
+def _unflatten_and_expand_time(flat_variables, tree_def, time_coords):
+  variables = jax.tree_util.tree_unflatten(tree_def, flat_variables)
+  return variables.expand_dims(time=time_coords, axis=0)
+
+
+def _get_flat_arrays_and_single_timestep_treedef(variables):
+  flat_arrays = jax.tree_util.tree_leaves(variables.transpose('time', ...))
+  _, treedef = jax.tree_util.tree_flatten(variables.isel(time=0, drop=True))
+  return flat_arrays, treedef
+
+
+class Predictor(predictor_base.Predictor):
+  """Wraps a one-step Predictor to make multi-step predictions autoregressively.
+
+  The wrapped Predictor will be used to predict a single timestep conditional
+  on the inputs passed to the outer Predictor. Its predictions are then
+  passed back in as inputs at the next timestep, for as many timesteps as are
+  requested in the targets_template. (When multiple timesteps of input are
+  used, a rolling window of inputs is maintained with new predictions
+  concatenated onto the end).
+
+  You may ask for additional variables to be predicted as targets which aren't
+  used as inputs. These will be predicted as output variables only and not fed
+  back in autoregressively. All target variables must be time-dependent however.
+
+  You may also specify static (non-time-dependent) inputs which will be passed
+  in at each timestep but are not predicted.
+
+  At present, any time-dependent inputs must also be present as targets so they
+  can be passed in autoregressively.
+
+  The loss of the wrapped one-step Predictor is averaged over all timesteps to
+  give a loss for the autoregressive Predictor.
+  """
+
+  def __init__(
+      self,
+      predictor: predictor_base.Predictor,
+      noise_level: Optional[float] = None,
+      gradient_checkpointing: bool = False,
+      ):
+    """Initializes an autoregressive predictor wrapper.
+
+    Args:
+      predictor: A predictor to wrap in an auto-regressive way.
+      noise_level: Optional value that multiplies the standard normal noise
+        added to the time-dependent variables of the predictor inputs. In
+        particular, no noise is added to the predictions that are fed back
+        auto-regressively. Defaults to not adding noise.
+      gradient_checkpointing: If True, gradient checkpointing will be
+        used at each step of the computation to save on memory. Roughtly this
+        should make the backwards pass two times more expensive, and the time
+        per step counting the forward pass, should only increase by about 50%.
+        Note this parameter will be ignored with a warning if the scan sequence
+        length is 1.
+    """
+    self._predictor = predictor
+    self._noise_level = noise_level
+    self._gradient_checkpointing = gradient_checkpointing
+
+  def _get_and_validate_constant_inputs(self, inputs, targets, forcings):
+    constant_inputs = inputs.drop_vars(targets.keys(), errors='ignore')
+    constant_inputs = constant_inputs.drop_vars(
+        forcings.keys(), errors='ignore')
+    for name, var in constant_inputs.items():
+      if 'time' in var.dims:
+        raise ValueError(
+            f'Time-dependent input variable {name} must either be a forcing '
+            'variable, or a target variable to allow for auto-regressive '
+            'feedback.')
+    return constant_inputs
+
+  def _validate_targets_and_forcings(self, targets, forcings):
+    for name, var in targets.items():
+      if 'time' not in var.dims:
+        raise ValueError(f'Target variable {name} must be time-dependent.')
+
+    for name, var in forcings.items():
+      if 'time' not in var.dims:
+        raise ValueError(f'Forcing variable {name} must be time-dependent.')
+
+    overlap = forcings.keys() & targets.keys()
+    if overlap:
+      raise ValueError('The following were specified as both targets and '
+                       f'forcings, which isn\'t allowed: {overlap}')
+
+  def _update_inputs(self, inputs, next_frame):
+    num_inputs = inputs.dims['time']
+
+    predicted_or_forced_inputs = next_frame[list(inputs.keys())]
+
+    # Combining datasets with inputs and target time stamps aligns them.
+    # Only keep the num_inputs trailing frames for use as next inputs.
+    return (xarray.concat([inputs, predicted_or_forced_inputs], dim='time')
+            .tail(time=num_inputs)
+            # Update the time coordinate to reset the lead times for
+            # next AR iteration.
+            .assign_coords(time=inputs.coords['time']))
+
+  def __call__(self,
+               inputs: xarray.Dataset,
+               targets_template: xarray.Dataset,
+               forcings: xarray.Dataset,
+               **kwargs) -> xarray.Dataset:
+    """Calls the Predictor.
+
+    Args:
+      inputs: input variable used to make predictions. Inputs can include both
+        time-dependent and time independent variables. Any time-dependent
+        input variables must also be present in the targets_template or the
+        forcings.
+      targets_template: A target template containing informations about which
+        variables should be predicted and the time alignment of the predictions.
+        All target variables must be time-dependent.
+        The number of time frames is used to set the number of unroll of the AR
+        predictor (e.g. multiple unroll of the inner predictor for one time step
+        in the targets is not supported yet).
+      forcings: Variables that will be fed to the model. The variables
+        should not overlap with the target ones. The time coordinates of the
+        forcing variables should match the target ones.
+        Forcing variables which are also present in the inputs, will be used to
+        supply ground-truth values for those inputs when they are passed to the
+        underlying predictor at timesteps beyond the first timestep.
+      **kwargs: Additional arguments passed along to the inner Predictor.
+
+    Returns:
+      predictions: the model predictions matching the target template.
+
+    Raise:
+      ValueError: if the time coordinates of the inputs and targets are not
+        different by a constant time step.
+    """
+
+    constant_inputs = self._get_and_validate_constant_inputs(
+        inputs, targets_template, forcings)
+    self._validate_targets_and_forcings(targets_template, forcings)
+
+    # After the above checks, the remaining inputs must be time-dependent:
+    inputs = inputs.drop_vars(constant_inputs.keys())
+
+    # A predictions template only including the next time to predict.
+    target_template = targets_template.isel(time=[0])
+
+    flat_forcings, forcings_treedef = (
+        _get_flat_arrays_and_single_timestep_treedef(forcings))
+    scan_variables = flat_forcings
+
+    def one_step_prediction(inputs, scan_variables):
+
+      flat_forcings = scan_variables
+      forcings = _unflatten_and_expand_time(flat_forcings, forcings_treedef,
+                                            target_template.coords['time'])
+
+      # Add constant inputs:
+      all_inputs = xarray.merge([constant_inputs, inputs])
+      predictions: xarray.Dataset = self._predictor(
+          all_inputs, target_template,
+          forcings=forcings,
+          **kwargs)
+
+      next_frame = xarray.merge([predictions, forcings])
+      next_inputs = self._update_inputs(inputs, next_frame)
+
+      # Drop the length-1 time dimension, since scan will concat all the outputs
+      # for different times along a new leading time dimension:
+      predictions = predictions.squeeze('time', drop=True)
+      # We return the prediction flattened into plain jax arrays, because the
+      # extra leading dimension added by scan prevents the tree_util
+      # registrations in xarray_jax from unflattening them back into an
+      # xarray.Dataset automatically:
+      flat_pred = jax.tree_util.tree_leaves(predictions)
+      return next_inputs, flat_pred
+
+    if self._gradient_checkpointing:
+      scan_length = targets_template.dims['time']
+      if scan_length <= 1:
+        logging.warning(
+            'Skipping gradient checkpointing for sequence length of 1')
+      else:
+        # Just in case we take gradients (e.g. for control), although
+        # in most cases this will just be for a forward pass.
+        one_step_prediction = hk.remat(one_step_prediction)
+
+    # Loop (without unroll) with hk states in cell (jax.lax.scan won't do).
+    _, flat_preds = hk.scan(one_step_prediction, inputs, scan_variables)
+
+    # The result of scan will have an extra leading axis on all arrays,
+    # corresponding to the target times in this case. We need to be prepared for
+    # it when unflattening the arrays back into a Dataset:
+    scan_result_template = (
+        target_template.squeeze('time', drop=True)
+        .expand_dims(time=targets_template.coords['time'], axis=0))
+    _, scan_result_treedef = jax.tree_util.tree_flatten(scan_result_template)
+    predictions = jax.tree_util.tree_unflatten(scan_result_treedef, flat_preds)
+    return predictions
+
+  def loss(self,
+           inputs: xarray.Dataset,
+           targets: xarray.Dataset,
+           forcings: xarray.Dataset,
+           **kwargs
+           ) -> predictor_base.LossAndDiagnostics:
+    """The mean of the per-timestep losses of the underlying predictor."""
+    if targets.sizes['time'] == 1:
+      # If there is only a single target timestep then we don't need any
+      # autoregressive feedback and can delegate the loss directly to the
+      # underlying single-step predictor. This means the underlying predictor
+      # doesn't need to implement .loss_and_predictions.
+      return self._predictor.loss(inputs, targets, forcings, **kwargs)
+
+    constant_inputs = self._get_and_validate_constant_inputs(
+        inputs, targets, forcings)
+    self._validate_targets_and_forcings(targets, forcings)
+    # After the above checks, the remaining inputs must be time-dependent:
+    inputs = inputs.drop_vars(constant_inputs.keys())
+
+    if self._noise_level:
+      def add_noise(x):
+        return x + self._noise_level * jax.random.normal(
+            hk.next_rng_key(), shape=x.shape)
+      # Add noise to time-dependent variables of the inputs.
+      inputs = jax.tree.map(add_noise, inputs)
+
+    # The per-timestep targets passed by scan to one_step_loss below will have
+    # no leading time axis. We need a treedef without the time axis to use
+    # inside one_step_loss to unflatten it back into a dataset:
+    flat_targets, target_treedef = _get_flat_arrays_and_single_timestep_treedef(
+        targets)
+    scan_variables = flat_targets
+
+    flat_forcings, forcings_treedef = (
+        _get_flat_arrays_and_single_timestep_treedef(forcings))
+    scan_variables = (flat_targets, flat_forcings)
+
+    def one_step_loss(inputs, scan_variables):
+      flat_target, flat_forcings = scan_variables
+      forcings = _unflatten_and_expand_time(flat_forcings, forcings_treedef,
+                                            targets.coords['time'][:1])
+
+      target = _unflatten_and_expand_time(flat_target, target_treedef,
+                                          targets.coords['time'][:1])
+
+      # Add constant inputs:
+      all_inputs = xarray.merge([constant_inputs, inputs])
+
+      (loss, diagnostics), predictions = self._predictor.loss_and_predictions(
+          all_inputs,
+          target,
+          forcings=forcings,
+          **kwargs)
+
+      # Unwrap to jax arrays shape (batch,):
+      loss, diagnostics = xarray_tree.map_structure(
+          xarray_jax.unwrap_data, (loss, diagnostics))
+
+      predictions = cast(xarray.Dataset, predictions)  # Keeps pytype happy.
+      next_frame = xarray.merge([predictions, forcings])
+      next_inputs = self._update_inputs(inputs, next_frame)
+
+      return next_inputs, (loss, diagnostics)
+
+    if self._gradient_checkpointing:
+      scan_length = targets.dims['time']
+      if scan_length <= 1:
+        logging.warning(
+            'Skipping gradient checkpointing for sequence length of 1')
+      else:
+        one_step_loss = hk.remat(one_step_loss)
+
+    # We can pass inputs (the initial state of the loop) in directly as a
+    # Dataset because the shape we pass in to scan is the same as the shape scan
+    # passes to the inner function. But, for scan_variables, we must flatten the
+    # targets (and unflatten them inside the inner function) because they are
+    # passed to the inner function per-timestep without the original time axis.
+    # The same apply to the optional forcing.
+    _, (per_timestep_losses, per_timestep_diagnostics) = hk.scan(
+        one_step_loss, inputs, scan_variables)
+
+    # Re-wrap loss and diagnostics as DataArray and average them over time:
+    (loss, diagnostics) = jax.tree_util.tree_map(
+        lambda x: xarray_jax.DataArray(x, dims=('time', 'batch')).mean(  # pylint: disable=g-long-lambda
+            'time', skipna=False),
+        (per_timestep_losses, per_timestep_diagnostics))
+
+    return loss, diagnostics

model/graphcast/casting.py

@@ -0,0 +1,205 @@
+# Copyright 2023 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Wrappers that take care of casting."""
+
+import contextlib
+from typing import Any, Mapping, Tuple
+
+import chex
+from . import predictor_base
+import haiku as hk
+import jax
+import jax.numpy as jnp
+import numpy as np
+import xarray
+
+
+PyTree = Any
+
+
+class Bfloat16Cast(predictor_base.Predictor):
+  """Wrapper that casts all inputs to bfloat16 and outputs to targets dtype."""
+
+  def __init__(self, predictor: predictor_base.Predictor, enabled: bool = True):
+    """Inits the wrapper.
+
+    Args:
+      predictor: predictor being wrapped.
+      enabled: disables the wrapper if False, for simpler hyperparameter scans.
+
+    """
+    self._enabled = enabled
+    self._predictor = predictor
+
+  def __call__(self,
+               inputs: xarray.Dataset,
+               targets_template: xarray.Dataset,
+               forcings: xarray.Dataset,
+               **kwargs
+               ) -> xarray.Dataset:
+    if not self._enabled:
+      return self._predictor(inputs, targets_template, forcings, **kwargs)
+
+    with bfloat16_variable_view():
+      predictions = self._predictor(
+          *_all_inputs_to_bfloat16(inputs, targets_template, forcings),
+          **kwargs,)
+
+    predictions_dtype = infer_floating_dtype(predictions)  # pytype: disable=wrong-arg-types
+    if predictions_dtype != jnp.bfloat16:
+      raise ValueError(f'Expected bfloat16 output, got {predictions_dtype}')
+
+    targets_dtype = infer_floating_dtype(targets_template)  # pytype: disable=wrong-arg-types
+    return tree_map_cast(
+        predictions, input_dtype=jnp.bfloat16, output_dtype=targets_dtype)
+
+  def loss(self,
+           inputs: xarray.Dataset,
+           targets: xarray.Dataset,
+           forcings: xarray.Dataset,
+           **kwargs,
+           ) -> predictor_base.LossAndDiagnostics:
+    if not self._enabled:
+      return self._predictor.loss(inputs, targets, forcings, **kwargs)
+
+    with bfloat16_variable_view():
+      loss, scalars = self._predictor.loss(
+          *_all_inputs_to_bfloat16(inputs, targets, forcings), **kwargs)
+
+    if loss.dtype != jnp.bfloat16:
+      raise ValueError(f'Expected bfloat16 loss, got {loss.dtype}')
+
+    targets_dtype = infer_floating_dtype(targets)  # pytype: disable=wrong-arg-types
+
+    # Note that casting back the loss to e.g. float32 should not affect data
+    # types of the backwards pass, because the first thing the backwards pass
+    # should do is to go backwards the casting op and cast back to bfloat16
+    # (and xprofs seem to confirm this).
+    return tree_map_cast((loss, scalars),
+                         input_dtype=jnp.bfloat16, output_dtype=targets_dtype)
+
+  def loss_and_predictions(  # pytype: disable=signature-mismatch  # jax-ndarray
+      self,
+      inputs: xarray.Dataset,
+      targets: xarray.Dataset,
+      forcings: xarray.Dataset,
+      **kwargs,
+      ) -> Tuple[predictor_base.LossAndDiagnostics,
+                 xarray.Dataset]:
+    if not self._enabled:
+      return self._predictor.loss_and_predictions(inputs, targets, forcings,  # pytype: disable=bad-return-type  # jax-ndarray
+                                                  **kwargs)
+
+    with bfloat16_variable_view():
+      (loss, scalars), predictions = self._predictor.loss_and_predictions(
+          *_all_inputs_to_bfloat16(inputs, targets, forcings), **kwargs)
+
+    if loss.dtype != jnp.bfloat16:
+      raise ValueError(f'Expected bfloat16 loss, got {loss.dtype}')
+
+    predictions_dtype = infer_floating_dtype(predictions)  # pytype: disable=wrong-arg-types
+    if predictions_dtype != jnp.bfloat16:
+      raise ValueError(f'Expected bfloat16 output, got {predictions_dtype}')
+
+    targets_dtype = infer_floating_dtype(targets)  # pytype: disable=wrong-arg-types
+    return tree_map_cast(((loss, scalars), predictions),
+                         input_dtype=jnp.bfloat16, output_dtype=targets_dtype)
+
+
+def infer_floating_dtype(data_vars: Mapping[str, chex.Array]) -> np.dtype:
+  """Infers a floating dtype from an input mapping of data."""
+  dtypes = {
+      v.dtype
+      for k, v in data_vars.items() if jnp.issubdtype(v.dtype, np.floating)}
+  if len(dtypes) != 1:
+    dtypes_and_shapes = {
+        k: (v.dtype, v.shape)
+        for k, v in data_vars.items() if jnp.issubdtype(v.dtype, np.floating)}
+    raise ValueError(
+        f'Did not found exactly one floating dtype {dtypes} in input variables:'
+        f'{dtypes_and_shapes}')
+  return list(dtypes)[0]
+
+
+def _all_inputs_to_bfloat16(
+    inputs: xarray.Dataset,
+    targets: xarray.Dataset,
+    forcings: xarray.Dataset,
+    ) -> Tuple[xarray.Dataset,
+               xarray.Dataset,
+               xarray.Dataset]:
+  return (inputs.astype(jnp.bfloat16),
+          jax.tree.map(lambda x: x.astype(jnp.bfloat16), targets),
+          forcings.astype(jnp.bfloat16))
+
+
+def tree_map_cast(inputs: PyTree, input_dtype: np.dtype, output_dtype: np.dtype,
+                  ) -> PyTree:
+  def cast_fn(x):
+    if x.dtype == input_dtype:
+      return x.astype(output_dtype)
+  return jax.tree.map(cast_fn, inputs)
+
+
+@contextlib.contextmanager
+def bfloat16_variable_view(enabled: bool = True):
+  """Context for Haiku modules with float32 params, but bfloat16 activations.
+
+  It works as follows:
+  * Every time a variable is requested to be created/set as np.bfloat16,
+    it will create an underlying float32 variable, instead.
+  * Every time a variable a variable is requested as bfloat16, it will check the
+    variable is of float32 type, and cast the variable to bfloat16.
+
+  Note the gradients are still computed and accumulated as float32, because
+  the params returned by init are float32, so the gradient function with
+  respect to the params will already include an implicit casting to float32.
+
+  Args:
+    enabled: Only enables bfloat16 behavior if True.
+
+  Yields:
+    None
+  """
+
+  if enabled:
+    with hk.custom_creator(
+        _bfloat16_creator, state=True), hk.custom_getter(
+            _bfloat16_getter, state=True), hk.custom_setter(
+                _bfloat16_setter):
+      yield
+  else:
+    yield
+
+
+def _bfloat16_creator(next_creator, shape, dtype, init, context):
+  """Creates float32 variables when bfloat16 is requested."""
+  if context.original_dtype == jnp.bfloat16:
+    dtype = jnp.float32
+  return next_creator(shape, dtype, init)
+
+
+def _bfloat16_getter(next_getter, value, context):
+  """Casts float32 to bfloat16 when bfloat16 was originally requested."""
+  if context.original_dtype == jnp.bfloat16:
+    assert value.dtype == jnp.float32
+    value = value.astype(jnp.bfloat16)
+  return next_getter(value)
+
+
+def _bfloat16_setter(next_setter, value, context):
+  """Casts bfloat16 to float32 when bfloat16 was originally set."""
+  if context.original_dtype == jnp.bfloat16:
+    value = value.astype(jnp.float32)
+  return next_setter(value)

model/graphcast/checkpoint.py

@@ -0,0 +1,170 @@
+# Copyright 2023 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Serialize and deserialize trees."""
+
+import dataclasses
+import io
+import types
+from typing import Any, BinaryIO, Optional, TypeVar
+
+import numpy as np
+
+_T = TypeVar("_T")
+
+
+def dump(dest: BinaryIO, value: Any) -> None:
+  """Dump a tree of dicts/dataclasses to a file object.
+
+  Args:
+    dest: a file object to write to.
+    value: A tree of dicts, lists, tuples and dataclasses of numpy arrays and
+      other basic types. Unions are not supported, other than Optional/None
+      which is only supported in dataclasses, not in dicts, lists or tuples.
+      All leaves must be coercible to a numpy array, and recoverable as a single
+      arg to a type.
+  """
+  buffer = io.BytesIO()  # In case the destination doesn't support seeking.
+  np.savez(buffer, **_flatten(value))
+  dest.write(buffer.getvalue())
+
+
+def load(source: BinaryIO, typ: type[_T]) -> _T:
+  """Load from a file object and convert it to the specified type.
+
+  Args:
+    source: a file object to read from.
+    typ: a type object that acts as a schema for deserialization. It must match
+      what was serialized. If a type is Any, it will be returned however numpy
+      serialized it, which is what you want for a tree of numpy arrays.
+
+  Returns:
+    the deserialized value as the specified type.
+  """
+  return _convert_types(typ, _unflatten(np.load(source)))
+
+
+_SEP = ":"
+
+
+def _flatten(tree: Any) -> dict[str, Any]:
+  """Flatten a tree of dicts/dataclasses/lists/tuples to a single dict."""
+  if dataclasses.is_dataclass(tree):
+    # Don't use dataclasses.asdict as it is recursive so skips dropping None.
+    tree = {f.name: v for f in dataclasses.fields(tree)
+            if (v := getattr(tree, f.name)) is not None}
+  elif isinstance(tree, (list, tuple)):
+    tree = dict(enumerate(tree))
+
+  assert isinstance(tree, dict)
+
+  flat = {}
+  for k, v in tree.items():
+    k = str(k)
+    assert _SEP not in k
+    if dataclasses.is_dataclass(v) or isinstance(v, (dict, list, tuple)):
+      for a, b in _flatten(v).items():
+        flat[f"{k}{_SEP}{a}"] = b
+    else:
+      assert v is not None
+      flat[k] = v
+  return flat
+
+
+def _unflatten(flat: dict[str, Any]) -> dict[str, Any]:
+  """Unflatten a dict to a tree of dicts."""
+  tree = {}
+  for flat_key, v in flat.items():
+    node = tree
+    keys = flat_key.split(_SEP)
+    for k in keys[:-1]:
+      if k not in node:
+        node[k] = {}
+      node = node[k]
+    node[keys[-1]] = v
+  return tree
+
+
+def _convert_types(typ: type[_T], value: Any) -> _T:
+  """Convert some structure into the given type. The structures must match."""
+  if typ in (Any, ...):
+    return value
+
+  if typ in (int, float, str, bool):
+    return typ(value)
+
+  if typ is np.ndarray:
+    assert isinstance(value, np.ndarray)
+    return value
+
+  if dataclasses.is_dataclass(typ):
+    kwargs = {}
+    for f in dataclasses.fields(typ):
+      # Only support Optional for dataclasses, as numpy can't serialize it
+      # directly (without pickle), and dataclasses are the only case where we
+      # can know the full set of values and types and therefore know the
+      # non-existence must mean None.
+      if isinstance(f.type, (types.UnionType, type(Optional[int]))):
+        constructors = [t for t in f.type.__args__ if t is not types.NoneType]
+        if len(constructors) != 1:
+          raise TypeError(
+              "Optional works, Union with anything except None doesn't")
+        if f.name not in value:
+          kwargs[f.name] = None
+          continue
+        constructor = constructors[0]
+      else:
+        constructor = f.type
+
+      if f.name in value:
+        kwargs[f.name] = _convert_types(constructor, value[f.name])
+      else:
+        raise ValueError(f"Missing value: {f.name}")
+    return typ(**kwargs)
+
+  base_type = getattr(typ, "__origin__", None)
+
+  if base_type is dict:
+    assert len(typ.__args__) == 2
+    key_type, value_type = typ.__args__
+    return {_convert_types(key_type, k): _convert_types(value_type, v)
+            for k, v in value.items()}
+
+  if base_type is list:
+    assert len(typ.__args__) == 1
+    value_type = typ.__args__[0]
+    return [_convert_types(value_type, v)
+            for _, v in sorted(value.items(), key=lambda x: int(x[0]))]
+
+  if base_type is tuple:
+    if len(typ.__args__) == 2 and typ.__args__[1] == ...:
+      # An arbitrary length tuple of a single type, eg: tuple[int, ...]
+      value_type = typ.__args__[0]
+      return tuple(_convert_types(value_type, v)
+                   for _, v in sorted(value.items(), key=lambda x: int(x[0])))
+    else:
+      # A fixed length tuple of arbitrary types, eg: tuple[int, str, float]
+      assert len(typ.__args__) == len(value)
+      return tuple(
+          _convert_types(t, v)
+          for t, (_, v) in zip(
+              typ.__args__, sorted(value.items(), key=lambda x: int(x[0]))))
+
+  # This is probably unreachable with reasonable serializable inputs.
+  try:
+    return typ(value)
+  except TypeError as e:
+    raise TypeError(
+        "_convert_types expects the type argument to be a dataclass defined "
+        "with types that are valid constructors (eg tuple is fine, Tuple "
+        "isn't), and accept a numpy array as the sole argument.") from e

model/graphcast/data_utils.py

@@ -0,0 +1,362 @@
+# Copyright 2023 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Dataset utilities."""
+
+from typing import Any, Mapping, Sequence, Tuple, Union
+
+from . import solar_radiation
+import numpy as np
+import pandas as pd
+import xarray
+
+TimedeltaLike = Any  # Something convertible to pd.Timedelta.
+TimedeltaStr = str  # A string convertible to pd.Timedelta.
+
+TargetLeadTimes = Union[
+    TimedeltaLike,
+    Sequence[TimedeltaLike],
+    slice  # with TimedeltaLike as its start and stop.
+]
+
+_SEC_PER_HOUR = 3600
+_HOUR_PER_DAY = 24
+SEC_PER_DAY = _SEC_PER_HOUR * _HOUR_PER_DAY
+_AVG_DAY_PER_YEAR = 365.24219
+AVG_SEC_PER_YEAR = SEC_PER_DAY * _AVG_DAY_PER_YEAR
+
+DAY_PROGRESS = "day_progress"
+YEAR_PROGRESS = "year_progress"
+_DERIVED_VARS = {
+    DAY_PROGRESS,
+    f"{DAY_PROGRESS}_sin",
+    f"{DAY_PROGRESS}_cos",
+    YEAR_PROGRESS,
+    f"{YEAR_PROGRESS}_sin",
+    f"{YEAR_PROGRESS}_cos",
+}
+TISR = "toa_incident_solar_radiation"
+
+
+def get_year_progress(seconds_since_epoch: np.ndarray) -> np.ndarray:
+  """Computes year progress for times in seconds.
+
+  Args:
+    seconds_since_epoch: Times in seconds since the "epoch" (the point at which
+      UNIX time starts).
+
+  Returns:
+    Year progress normalized to be in the [0, 1) interval for each time point.
+  """
+
+  # Start with the pure integer division, and then float at the very end.
+  # We will try to keep as much precision as possible.
+  years_since_epoch = (
+      seconds_since_epoch / SEC_PER_DAY / np.float64(_AVG_DAY_PER_YEAR)
+  )
+  # Note depending on how these ops are down, we may end up with a "weak_type"
+  # which can cause issues in subtle ways, and hard to track here.
+  # In any case, casting to float32 should get rid of the weak type.
+  # [0, 1.) Interval.
+  return np.mod(years_since_epoch, 1.0).astype(np.float32)
+
+
+def get_day_progress(
+    seconds_since_epoch: np.ndarray,
+    longitude: np.ndarray,
+) -> np.ndarray:
+  """Computes day progress for times in seconds at each longitude.
+
+  Args:
+    seconds_since_epoch: 1D array of times in seconds since the 'epoch' (the
+      point at which UNIX time starts).
+    longitude: 1D array of longitudes at which day progress is computed.
+
+  Returns:
+    2D array of day progress values normalized to be in the [0, 1) inverval
+      for each time point at each longitude.
+  """
+
+  # [0.0, 1.0) Interval.
+  day_progress_greenwich = (
+      np.mod(seconds_since_epoch, SEC_PER_DAY) / SEC_PER_DAY
+  )
+
+  # Offset the day progress to the longitude of each point on Earth.
+  longitude_offsets = np.deg2rad(longitude) / (2 * np.pi)
+  day_progress = np.mod(
+      day_progress_greenwich[..., np.newaxis] + longitude_offsets, 1.0
+  )
+  return day_progress.astype(np.float32)
+
+
+def featurize_progress(
+    name: str, dims: Sequence[str], progress: np.ndarray
+) -> Mapping[str, xarray.Variable]:
+  """Derives features used by ML models from the `progress` variable.
+
+  Args:
+    name: Base variable name from which features are derived.
+    dims: List of the output feature dimensions, e.g. ("day", "lon").
+    progress: Progress variable values.
+
+  Returns:
+    Dictionary of xarray variables derived from the `progress` values. It
+    includes the original `progress` variable along with its sin and cos
+    transformations.
+
+  Raises:
+    ValueError if the number of feature dimensions is not equal to the number
+      of data dimensions.
+  """
+  if len(dims) != progress.ndim:
+    raise ValueError(
+        f"Number of feature dimensions ({len(dims)}) must be equal to the"
+        f" number of data dimensions: {progress.ndim}."
+    )
+  progress_phase = progress * (2 * np.pi)
+  return {
+      name: xarray.Variable(dims, progress),
+      name + "_sin": xarray.Variable(dims, np.sin(progress_phase)),
+      name + "_cos": xarray.Variable(dims, np.cos(progress_phase)),
+  }
+
+
+def get_seconds_since_epoch(datetime_sequence: xarray.DataArray) -> np.ndarray:
+  """Computes seconds since epoch from `data` in place if missing."""
+  # Note `datetime_sequence.astype("datetime64[s]").astype(np.int64)`
+  # does not work as xarrays always cast dates into nanoseconds!
+  return datetime_sequence.data.astype("datetime64[s]").astype(np.int64)
+
+
+def add_derived_vars(data: xarray.Dataset) -> None:
+  """Adds year and day progress features to `data` in place if missing.
+
+  Args:
+    data: Xarray dataset to which derived features will be added.
+
+  Raises:
+    ValueError if `datetime` or `lon` are not in `data` coordinates.
+  """
+
+  for coord in ("datetime", "lon"):
+    if coord not in data.coords:
+      raise ValueError(f"'{coord}' must be in `data` coordinates.")
+
+  # Compute seconds since epoch.
+  seconds_since_epoch = get_seconds_since_epoch(data.coords["datetime"])
+  batch_dim = ("batch",) if "batch" in data.dims else ()
+
+  # Add year progress features if missing.
+  if YEAR_PROGRESS not in data.data_vars:
+    year_progress = get_year_progress(seconds_since_epoch)
+    data.update(
+        featurize_progress(
+            name=YEAR_PROGRESS,
+            dims=batch_dim + ("time",),
+            progress=year_progress,
+        )
+    )
+
+  # Add day progress features if missing.
+  if DAY_PROGRESS not in data.data_vars:
+    longitude_coord = data.coords["lon"]
+    day_progress = get_day_progress(seconds_since_epoch, longitude_coord.data)
+    data.update(
+        featurize_progress(
+            name=DAY_PROGRESS,
+            dims=batch_dim + ("time",) + longitude_coord.dims,
+            progress=day_progress,
+        )
+    )
+
+
+def add_tisr_var(data: xarray.Dataset) -> None:
+  """Adds TISR feature to `data` in place if missing.
+
+  Args:
+    data: Xarray dataset to which TISR feature will be added.
+
+  Raises:
+    ValueError if `datetime`, 'lat', or `lon` are not in `data` coordinates.
+  """
+
+  if TISR in data.data_vars:
+    return
+
+  for coord in ("datetime", "lat", "lon"):
+    if coord not in data.coords:
+      raise ValueError(f"'{coord}' must be in `data` coordinates.")
+
+  # Remove `batch` dimension of size one if present. An error will be raised if
+  # the `batch` dimension exists and has size greater than one.
+  data_no_batch = data.squeeze("batch") if "batch" in data.dims else data
+
+  tisr = solar_radiation.get_toa_incident_solar_radiation_for_xarray(
+      data_no_batch, use_jit=True
+  )
+
+  if "batch" in data.dims:
+    tisr = tisr.expand_dims("batch", axis=0)
+
+  data.update({TISR: tisr})
+
+
+def extract_input_target_times(
+    dataset: xarray.Dataset,
+    input_duration: TimedeltaLike,
+    target_lead_times: TargetLeadTimes,
+    ) -> Tuple[xarray.Dataset, xarray.Dataset]:
+  """Extracts inputs and targets for prediction, from a Dataset with a time dim.
+
+  The input period is assumed to be contiguous (specified by a duration), but
+  the targets can be a list of arbitrary lead times.
+
+  Examples:
+
+    # Use 18 hours of data as inputs, and two specific lead times as targets:
+    # 3 days and 5 days after the final input.
+    extract_inputs_targets(
+        dataset,
+        input_duration='18h',
+        target_lead_times=('3d', '5d')
+    )
+
+    # Use 1 day of data as input, and all lead times between 6 hours and
+    # 24 hours inclusive as targets. Demonstrates a friendlier supported string
+    # syntax.
+    extract_inputs_targets(
+        dataset,
+        input_duration='1 day',
+        target_lead_times=slice('6 hours', '24 hours')
+    )
+
+    # Just use a single target lead time of 3 days:
+    extract_inputs_targets(
+        dataset,
+        input_duration='24h',
+        target_lead_times='3d'
+    )
+
+  Args:
+    dataset: An xarray.Dataset with a 'time' dimension whose coordinates are
+      timedeltas. It's assumed that the time coordinates have a fixed offset /
+      time resolution, and that the input_duration and target_lead_times are
+      multiples of this.
+    input_duration: pandas.Timedelta or something convertible to it (e.g. a
+      shorthand string like '6h' or '5d12h').
+    target_lead_times: Either a single lead time, a slice with start and stop
+      (inclusive) lead times, or a sequence of lead times. Lead times should be
+      Timedeltas (or something convertible to). They are given relative to the
+      final input timestep, and should be positive.
+
+  Returns:
+    inputs:
+    targets:
+      Two datasets with the same shape as the input dataset except that a
+      selection has been made from the time axis, and the origin of the
+      time coordinate will be shifted to refer to lead times relative to the
+      final input timestep. So for inputs the times will end at lead time 0,
+      for targets the time coordinates will refer to the lead times requested.
+  """
+
+  (target_lead_times, target_duration
+   ) = _process_target_lead_times_and_get_duration(target_lead_times)
+
+  # Shift the coordinates for the time axis so that a timedelta of zero
+  # corresponds to the forecast reference time. That is, the final timestep
+  # that's available as input to the forecast, with all following timesteps
+  # forming the target period which needs to be predicted.
+  # This means the time coordinates are now forecast lead times.
+  time = dataset.coords["time"]
+  dataset = dataset.assign_coords(time=time + target_duration - time[-1])
+
+  # Slice out targets:
+  targets = dataset.sel({"time": target_lead_times})
+
+  input_duration = pd.Timedelta(input_duration)
+  # Both endpoints are inclusive with label-based slicing, so we offset by a
+  # small epsilon to make one of the endpoints non-inclusive:
+  zero = pd.Timedelta(0)
+  epsilon = pd.Timedelta(1, "ns")
+  inputs = dataset.sel({"time": slice(-input_duration + epsilon, zero)})
+  return inputs, targets
+
+
+def _process_target_lead_times_and_get_duration(
+    target_lead_times: TargetLeadTimes) -> TimedeltaLike:
+  """Returns the minimum duration for the target lead times."""
+  if isinstance(target_lead_times, slice):
+    # A slice of lead times. xarray already accepts timedelta-like values for
+    # the begin/end/step of the slice.
+    if target_lead_times.start is None:
+      # If the start isn't specified, we assume it starts at the next timestep
+      # after lead time 0 (lead time 0 is the final input timestep):
+      target_lead_times = slice(
+          pd.Timedelta(1, "ns"), target_lead_times.stop, target_lead_times.step
+      )
+    target_duration = pd.Timedelta(target_lead_times.stop)
+  else:
+    if not isinstance(target_lead_times, (list, tuple, set)):
+      # A single lead time, which we wrap as a length-1 array to ensure there
+      # still remains a time dimension (here of length 1) for consistency.
+      target_lead_times = [target_lead_times]
+
+    # A list of multiple (not necessarily contiguous) lead times:
+    target_lead_times = [pd.Timedelta(x) for x in target_lead_times]
+    target_lead_times.sort()
+    target_duration = target_lead_times[-1]
+  return target_lead_times, target_duration
+
+
+def extract_inputs_targets_forcings(
+    dataset: xarray.Dataset,
+    *,
+    input_variables: Tuple[str, ...],
+    target_variables: Tuple[str, ...],
+    forcing_variables: Tuple[str, ...],
+    pressure_levels: Tuple[int, ...],
+    input_duration: TimedeltaLike,
+    target_lead_times: TargetLeadTimes,
+    ) -> Tuple[xarray.Dataset, xarray.Dataset, xarray.Dataset]:
+  """Extracts inputs, targets and forcings according to requirements."""
+  dataset = dataset.sel(level=list(pressure_levels))
+
+  # "Forcings" include derived variables that do not exist in the original ERA5
+  # or HRES datasets, as well as other variables (e.g. tisr) that need to be
+  # computed manually for the target lead times. Compute the requested ones.
+  if set(forcing_variables) & _DERIVED_VARS:
+    add_derived_vars(dataset)
+  if set(forcing_variables) & {TISR}:
+    add_tisr_var(dataset)
+
+  # `datetime` is needed by add_derived_vars but breaks autoregressive rollouts.
+  dataset = dataset.drop_vars("datetime")
+
+  inputs, targets = extract_input_target_times(
+      dataset,
+      input_duration=input_duration,
+      target_lead_times=target_lead_times)
+
+  if set(forcing_variables) & set(target_variables):
+    raise ValueError(
+        f"Forcing variables {forcing_variables} should not "
+        f"overlap with target variables {target_variables}."
+    )
+
+  inputs = inputs[list(input_variables)]
+  # The forcing uses the same time coordinates as the target.
+  forcings = targets[list(forcing_variables)]
+  targets = targets[list(target_variables)]
+
+  return inputs, targets, forcings

model/graphcast/deep_typed_graph_net.py

@@ -0,0 +1,461 @@
+# Copyright 2023 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""JAX implementation of Graph Networks Simulator.
+
+Generalization to TypedGraphs of the deep Graph Neural Network from:
+
+@inproceedings{pfaff2021learning,
+  title={Learning Mesh-Based Simulation with Graph Networks},
+  author={Pfaff, Tobias and Fortunato, Meire and Sanchez-Gonzalez, Alvaro and
+      Battaglia, Peter},
+  booktitle={International Conference on Learning Representations},
+  year={2021}
+}
+
+@inproceedings{sanchez2020learning,
+  title={Learning to simulate complex physics with graph networks},
+  author={Sanchez-Gonzalez, Alvaro and Godwin, Jonathan and Pfaff, Tobias and
+      Ying, Rex and Leskovec, Jure and Battaglia, Peter},
+  booktitle={International conference on machine learning},
+  pages={8459--8468},
+  year={2020},
+  organization={PMLR}
+}
+"""
+
+import functools
+from typing import Callable, List, Mapping, Optional, Tuple
+
+import chex
+from . import mlp as mlp_builder
+from . import typed_graph
+from . import typed_graph_net
+import haiku as hk
+import jax
+import jax.numpy as jnp
+import jraph
+
+
+GraphToGraphNetwork = Callable[[typed_graph.TypedGraph], typed_graph.TypedGraph]
+
+
+class DeepTypedGraphNet(hk.Module):
+  """Deep Graph Neural Network.
+
+  It works with TypedGraphs with typed nodes and edges. It runs message
+  passing on all of the node sets and all of the edge sets in the graph. For
+  each message passing step a `typed_graph_net.InteractionNetwork` is used to
+  update the full TypedGraph by using different MLPs for each of the node sets
+  and each of the edge sets.
+
+  If embed_{nodes,edges} is specified the node/edge features will be embedded
+  into a fixed dimensionality before running the first step of message passing.
+
+  If {node,edge}_output_size the final node/edge features will be embedded into
+  the specified output size.
+
+  This class may be used for shared or unshared message passing:
+  * num_message_passing_steps = N, num_processor_repetitions = 1, gives
+    N layers of message passing with fully unshared weights:
+    [W_1, W_2, ... , W_M] (default)
+  * num_message_passing_steps = 1, num_processor_repetitions = M, gives
+    N layers of message passing with fully shared weights:
+    [W_1] * M
+  * num_message_passing_steps = N, num_processor_repetitions = M, gives
+    M*N layers of message passing with both shared and unshared message passing
+    such that the weights used at each iteration are:
+    [W_1, W_2, ... , W_N] * M
+
+  """
+
+  def __init__(self,
+               *,
+               node_latent_size: Mapping[str, int],
+               edge_latent_size: Mapping[str, int],
+               mlp_hidden_size: int,
+               mlp_num_hidden_layers: int,
+               num_message_passing_steps: int,
+               num_processor_repetitions: int = 1,
+               embed_nodes: bool = True,
+               embed_edges: bool = True,
+               node_output_size: Optional[Mapping[str, int]] = None,
+               edge_output_size: Optional[Mapping[str, int]] = None,
+               include_sent_messages_in_node_update: bool = False,
+               use_layer_norm: bool = True,
+               use_norm_conditioning: bool = False,
+               activation: str = "relu",
+               f32_aggregation: bool = False,
+               aggregate_edges_for_nodes_fn: str = "segment_sum",
+               aggregate_normalization: Optional[float] = None,
+               name: str = "DeepTypedGraphNet"):
+    """Inits the model.
+
+    Args:
+      node_latent_size: Size of the node latent representations.
+      edge_latent_size: Size of the edge latent representations.
+      mlp_hidden_size: Hidden layer size for all MLPs.
+      mlp_num_hidden_layers: Number of hidden layers in all MLPs.
+      num_message_passing_steps: Number of unshared message passing steps
+         in the processor steps.
+      num_processor_repetitions: Number of times that the same processor is
+         applied sequencially.
+      embed_nodes: If False, the node embedder will be omitted.
+      embed_edges: If False, the edge embedder will be omitted.
+      node_output_size: Size of the output node representations for
+         each node type. For node types not specified here, the latent node
+         representation from the output of the processor will be returned.
+      edge_output_size: Size of the output edge representations for
+         each edge type. For edge types not specified here, the latent edge
+         representation from the output of the processor will be returned.
+      include_sent_messages_in_node_update: Whether to include pooled sent
+          messages from each node in the node update.
+      use_layer_norm: Whether it uses layer norm or not.
+      use_norm_conditioning: If True, the latent feaures outputted by the
+        activation normalization that follows the MLPs (e.g. LayerNorm), rather
+        than being scaled/offset by learned  parameters of the normalization
+        module, will be scaled/offset by offsets/biases produced by a linear
+        layer (with different weights for each MLP), which takes an extra
+        argument "global_norm_conditioning". This argument is used to condition
+        the normalization of all nodes and all edges (hence global), and would
+        usually only have a batch and feature axis. This is typically used to
+        condition diffusion models on the "diffusion time". Will raise an error
+        if this is set to True but the "global_norm_conditioning" is not passed
+        to the __call__ method, as well as if this is set to False, but
+        "global_norm_conditioning" is passed to the call method.
+      activation: name of activation function.
+      f32_aggregation: Use float32 in the edge aggregation.
+      aggregate_edges_for_nodes_fn: function used to aggregate messages to each
+        node.
+      aggregate_normalization: An optional constant that normalizes the output
+        of aggregate_edges_for_nodes_fn. For context, this can be used to
+        reduce the shock the model undergoes when switching resolution, which
+        increase the number of edges connected to a node. In particular, this is
+        useful when using segment_sum, but should not be combined with
+        segment_mean.
+      name: Name of the model.
+    """
+
+    super().__init__(name=name)
+
+    self._node_latent_size = node_latent_size
+    self._edge_latent_size = edge_latent_size
+    self._mlp_hidden_size = mlp_hidden_size
+    self._mlp_num_hidden_layers = mlp_num_hidden_layers
+    self._num_message_passing_steps = num_message_passing_steps
+    self._num_processor_repetitions = num_processor_repetitions
+    self._embed_nodes = embed_nodes
+    self._embed_edges = embed_edges
+    self._node_output_size = node_output_size
+    self._edge_output_size = edge_output_size
+    self._include_sent_messages_in_node_update = (
+        include_sent_messages_in_node_update)
+    if use_norm_conditioning and not use_layer_norm:
+      raise ValueError(
+          "`norm_conditioning` can only be used when "
+          "`use_layer_norm` is true."
+      )
+    self._use_layer_norm = use_layer_norm
+    self._use_norm_conditioning = use_norm_conditioning
+    self._activation = _get_activation_fn(activation)
+    self._f32_aggregation = f32_aggregation
+    self._aggregate_edges_for_nodes_fn = _get_aggregate_edges_for_nodes_fn(
+        aggregate_edges_for_nodes_fn)
+    self._aggregate_normalization = aggregate_normalization
+
+    if aggregate_normalization:
+      # using aggregate_normalization only makes sense with segment_sum.
+      assert aggregate_edges_for_nodes_fn == "segment_sum"
+
+  def __call__(self,
+               input_graph: typed_graph.TypedGraph,
+               global_norm_conditioning: Optional[chex.Array] = None
+               ) -> typed_graph.TypedGraph:
+    """Forward pass of the learnable dynamics model."""
+    embedder_network, processor_networks, decoder_network = (
+        self._networks_builder(input_graph, global_norm_conditioning)
+    )
+
+    # Embed input features (if applicable).
+    latent_graph_0 = self._embed(input_graph, embedder_network)
+
+    # Do `m` message passing steps in the latent graphs.
+    latent_graph_m = self._process(latent_graph_0, processor_networks)
+
+    # Compute outputs from the last latent graph (if applicable).
+    return self._output(latent_graph_m, decoder_network)
+
+  def _networks_builder(
+      self,
+      graph_template: typed_graph.TypedGraph,
+      global_norm_conditioning: Optional[chex.Array] = None,
+  ) -> Tuple[
+      GraphToGraphNetwork, List[GraphToGraphNetwork], GraphToGraphNetwork
+  ]:
+    # TODO(aelkadi): move to mlp_builder.
+    def build_mlp(name, output_size):
+      mlp = hk.nets.MLP(
+          output_sizes=[self._mlp_hidden_size] * self._mlp_num_hidden_layers + [
+              output_size], name=name + "_mlp", activation=self._activation)
+      return jraph.concatenated_args(mlp)
+
+    def build_mlp_with_maybe_layer_norm(name, output_size):
+      network = build_mlp(name, output_size)
+      stages = [network]
+      if self._use_norm_conditioning:
+        if global_norm_conditioning is None:
+          raise ValueError(
+              "When using norm conditioning, `global_norm_conditioning` must"
+              "be passed to the call method.")
+        # If using norm conditioning, it is no longer the responsibility of the
+        # LayerNorm module itself to learn its scale and offset. These will be
+        # learned for the module by the norm conditioning layer instead.
+        create_scale = create_offset = False
+      else:
+        if global_norm_conditioning is not None:
+          raise ValueError(
+              "`globa_norm_conditioning` was passed, but `norm_conditioning`"
+              " is not enabled.")
+        create_scale = create_offset = True
+
+      if self._use_layer_norm:
+        layer_norm = hk.LayerNorm(
+            axis=-1, create_scale=create_scale, create_offset=create_offset,
+            name=name + "_layer_norm")
+        stages.append(layer_norm)
+
+      if self._use_norm_conditioning:
+        norm_conditioning_layer = mlp_builder.LinearNormConditioning(
+            name=name + "_norm_conditioning")
+        norm_conditioning_layer = functools.partial(
+            norm_conditioning_layer,
+            # Broadcast to the node/edge axis.
+            norm_conditioning=global_norm_conditioning[None],
+        )
+        stages.append(norm_conditioning_layer)
+
+      network = hk.Sequential(stages)
+      return jraph.concatenated_args(network)
+
+    # The embedder graph network independently embeds edge and node features.
+    if self._embed_edges:
+      embed_edge_fn = _build_update_fns_for_edge_types(
+          build_mlp_with_maybe_layer_norm,
+          graph_template,
+          "encoder_edges_",
+          output_sizes=self._edge_latent_size)
+    else:
+      embed_edge_fn = None
+    if self._embed_nodes:
+      embed_node_fn = _build_update_fns_for_node_types(
+          build_mlp_with_maybe_layer_norm,
+          graph_template,
+          "encoder_nodes_",
+          output_sizes=self._node_latent_size)
+    else:
+      embed_node_fn = None
+    embedder_kwargs = dict(
+        embed_edge_fn=embed_edge_fn,
+        embed_node_fn=embed_node_fn,
+    )
+    embedder_network = typed_graph_net.GraphMapFeatures(
+        **embedder_kwargs)
+
+    if self._f32_aggregation:
+      def aggregate_fn(data, *args, **kwargs):
+        dtype = data.dtype
+        data = data.astype(jnp.float32)
+        output = self._aggregate_edges_for_nodes_fn(data, *args, **kwargs)
+        if self._aggregate_normalization:
+          output = output / self._aggregate_normalization
+        output = output.astype(dtype)
+        return output
+
+    else:
+      def aggregate_fn(data, *args, **kwargs):
+        output = self._aggregate_edges_for_nodes_fn(data, *args, **kwargs)
+        if self._aggregate_normalization:
+          output = output / self._aggregate_normalization
+        return output
+
+    # Create `num_message_passing_steps` graph networks with unshared parameters
+    # that update the node and edge latent features.
+    # Note that we can use `modules.InteractionNetwork` because
+    # it also outputs the messages as updated edge latent features.
+    processor_networks = []
+    for step_i in range(self._num_message_passing_steps):
+      processor_networks.append(
+          typed_graph_net.InteractionNetwork(
+              update_edge_fn=_build_update_fns_for_edge_types(
+                  build_mlp_with_maybe_layer_norm,
+                  graph_template,
+                  f"processor_edges_{step_i}_",
+                  output_sizes=self._edge_latent_size),
+              update_node_fn=_build_update_fns_for_node_types(
+                  build_mlp_with_maybe_layer_norm,
+                  graph_template,
+                  f"processor_nodes_{step_i}_",
+                  output_sizes=self._node_latent_size),
+              aggregate_edges_for_nodes_fn=aggregate_fn,
+              include_sent_messages_in_node_update=(
+                  self._include_sent_messages_in_node_update),
+              ))
+
+    # The output MLPs converts edge/node latent features into the output sizes.
+    output_kwargs = dict(
+        embed_edge_fn=_build_update_fns_for_edge_types(
+            build_mlp, graph_template, "decoder_edges_", self._edge_output_size)
+        if self._edge_output_size else None,
+        embed_node_fn=_build_update_fns_for_node_types(
+            build_mlp, graph_template, "decoder_nodes_", self._node_output_size)
+        if self._node_output_size else None,)
+    output_network = typed_graph_net.GraphMapFeatures(
+        **output_kwargs)
+    return embedder_network, processor_networks, output_network
+
+  def _embed(
+      self,
+      input_graph: typed_graph.TypedGraph,
+      embedder_network: GraphToGraphNetwork,
+  ) -> typed_graph.TypedGraph:
+    """Embeds the input graph features into a latent graph."""
+
+    # Copy the context to all of the node types, if applicable.
+    context_features = input_graph.context.features
+    if jax.tree_util.tree_leaves(context_features):
+      # This code assumes a single input feature array for the context and for
+      # each node type.
+      assert len(jax.tree_util.tree_leaves(context_features)) == 1
+      new_nodes = {}
+      for node_set_name, node_set in input_graph.nodes.items():
+        node_features = node_set.features
+        broadcasted_context = jnp.repeat(
+            context_features, node_set.n_node, axis=0,
+            total_repeat_length=node_features.shape[0])
+        new_nodes[node_set_name] = node_set._replace(
+            features=jnp.concatenate(
+                [node_features, broadcasted_context], axis=-1))
+      input_graph = input_graph._replace(
+          nodes=new_nodes,
+          context=input_graph.context._replace(features=()))
+
+    # Embeds the node and edge features.
+    latent_graph_0 = embedder_network(input_graph)
+    return latent_graph_0
+
+  def _process(
+      self,
+      latent_graph_0: typed_graph.TypedGraph,
+      processor_networks: List[GraphToGraphNetwork],
+  ) -> typed_graph.TypedGraph:
+    """Processes the latent graph with several steps of message passing."""
+
+    # Do `num_message_passing_steps` with each of the `self._processor_networks`
+    # with unshared weights, and repeat that `self._num_processor_repetitions`
+    # times.
+    latent_graph = latent_graph_0
+    for unused_repetition_i in range(self._num_processor_repetitions):
+      for processor_network in processor_networks:
+        latent_graph = self._process_step(processor_network, latent_graph)
+
+    return latent_graph
+
+  def _process_step(
+      self, processor_network_k,
+      latent_graph_prev_k: typed_graph.TypedGraph) -> typed_graph.TypedGraph:
+    """Single step of message passing with node/edge residual connections."""
+
+    # One step of message passing.
+    latent_graph_k = processor_network_k(latent_graph_prev_k)
+
+    # Add residuals.
+    nodes_with_residuals = {}
+    for k, prev_set in latent_graph_prev_k.nodes.items():
+      nodes_with_residuals[k] = prev_set._replace(
+          features=prev_set.features + latent_graph_k.nodes[k].features)
+
+    edges_with_residuals = {}
+    for k, prev_set in latent_graph_prev_k.edges.items():
+      edges_with_residuals[k] = prev_set._replace(
+          features=prev_set.features + latent_graph_k.edges[k].features)
+
+    latent_graph_k = latent_graph_k._replace(
+        nodes=nodes_with_residuals, edges=edges_with_residuals)
+    return latent_graph_k
+
+  def _output(
+      self,
+      latent_graph: typed_graph.TypedGraph,
+      output_network: GraphToGraphNetwork,
+  ) -> typed_graph.TypedGraph:
+    """Produces the output from the latent graph."""
+    return output_network(latent_graph)
+
+
+def _build_update_fns_for_node_types(
+    builder_fn, graph_template, prefix, output_sizes=None):
+  """Builds an update function for all node types or a subset of them."""
+
+  output_fns = {}
+  for node_set_name in graph_template.nodes.keys():
+    if output_sizes is None:
+      # Use the default output size for all types.
+      output_size = None
+    else:
+      # Otherwise, ignore any type that does not have an explicit output size.
+      if node_set_name in output_sizes:
+        output_size = output_sizes[node_set_name]
+      else:
+        continue
+    output_fns[node_set_name] = builder_fn(
+        f"{prefix}{node_set_name}", output_size)
+  return output_fns
+
+
+def _build_update_fns_for_edge_types(
+    builder_fn, graph_template, prefix, output_sizes=None):
+  """Builds an edge function for all node types or a subset of them."""
+  output_fns = {}
+  for edge_set_key in graph_template.edges.keys():
+    edge_set_name = edge_set_key.name
+    if output_sizes is None:
+      # Use the default output size for all types.
+      output_size = None
+    else:
+      # Otherwise, ignore any type that does not have an explicit output size.
+      if edge_set_name in output_sizes:
+        output_size = output_sizes[edge_set_name]
+      else:
+        continue
+    output_fns[edge_set_name] = builder_fn(
+        f"{prefix}{edge_set_name}", output_size)
+  return output_fns
+
+
+def _get_activation_fn(name):
+  """Return activation function corresponding to function_name."""
+  if name == "identity":
+    return lambda x: x
+  if hasattr(jax.nn, name):
+    return getattr(jax.nn, name)
+  if hasattr(jnp, name):
+    return getattr(jnp, name)
+  raise ValueError(f"Unknown activation function {name} specified.")
+
+
+def _get_aggregate_edges_for_nodes_fn(name):
+  """Return aggregate_edges_for_nodes_fn corresponding to function_name."""
+  if hasattr(jraph, name):
+    return getattr(jraph, name)
+  raise ValueError(
+      f"Unknown aggregate_edges_for_nodes_fn function {name} specified.")

model/graphcast/denoiser.py

@@ -0,0 +1,852 @@
+# Copyright 2024 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Support for wrapping a general Predictor to act as a Denoiser."""
+
+import dataclasses
+from typing import Any, Callable, Mapping, Optional, Sequence, Tuple
+
+import chex
+from . import deep_typed_graph_net
+from . import denoisers_base as base
+from . import grid_mesh_connectivity
+from . import icosahedral_mesh
+from . import model_utils
+from . import sparse_transformer
+from . import transformer
+from . import typed_graph
+from . import xarray_jax
+import haiku as hk
+import jax
+import jax.numpy as jnp
+import numpy as np
+from scipy import sparse
+import xarray
+
+
+Kwargs = Mapping[str, Any]
+NoiseLevelEncoder = Callable[[jnp.ndarray], jnp.ndarray]
+
+
+class FourierFeaturesMLP(hk.Module):
+  """A simple MLP applied to Fourier features of values or their logarithms."""
+
+  def __init__(
+      self,
+      base_period: float,
+      num_frequencies: int,
+      output_sizes: Sequence[int],
+      apply_log_first: bool = False,
+      w_init=None,
+      activation=jax.nn.gelu,
+      **mlp_kwargs
+  ):
+    """Initializes the module.
+
+    Args:
+      base_period:
+        See model_utils.fourier_features. Note this would apply to log inputs if
+        apply_log_first is used.
+      num_frequencies:
+        See model_utils.fourier_features.
+      output_sizes:
+        Layer sizes for the MLP.
+      apply_log_first:
+        Whether to take the log of the inputs before computing Fourier features.
+      w_init:
+        Weights initializer for the MLP, default setting aims to produce
+        approx unit-variance outputs given the input sin/cos features.
+      activation:
+      **mlp_kwargs:
+        Further settings for the MLP.
+    """
+    super().__init__()
+    self._base_period = base_period
+    self._num_frequencies = num_frequencies
+    self._apply_log_first = apply_log_first
+    if w_init is None:
+      # Scale of 2 is appropriate for input layer as sin/cos fourier features
+      # have variance 0.5 for random inputs. Also reasonable to use for later
+      # layers as relu activation cuts variance in half for inputs to later
+      # layers and gelu something close enough too.
+      w_init = hk.initializers.VarianceScaling(
+          2.0, mode="fan_in", distribution="uniform"
+      )
+    self._mlp = hk.nets.MLP(
+        output_sizes=output_sizes,
+        w_init=w_init,
+        activation=activation,
+        **mlp_kwargs)
+
+  def __call__(self, values: jnp.ndarray) -> jnp.ndarray:
+    if self._apply_log_first:
+      values = jnp.log(values)
+
+    features = model_utils.fourier_features(
+        values, self._base_period, self._num_frequencies)
+
+    return self._mlp(features)
+
+
+@chex.dataclass(frozen=True, eq=True)
+class NoiseEncoderConfig:
+  """Configures the noise level encoding.
+
+  Properties:
+    apply_log_first: Whether to take the log of the inputs before computing
+      Fourier features.
+    base_period: The base period to use. This should be greater or equal to the
+      range of the values, or to the period if the values have periodic
+      semantics (e.g. 2pi if they represent angles). Frequencies used will be
+      integer multiples of 1/base_period.
+    num_frequencies: The number of frequencies to use, we will use integer
+      multiples of 1/base_period from 1 up to num_frequencies inclusive. (We
+      don't include a zero frequency as this would just give constant features
+      which are redundant if a bias term is present).
+    output_sizes: Layer sizes for the MLP.
+  """
+  apply_log_first: bool = True
+  base_period: float = 16.0
+  num_frequencies: int = 32
+  # 2-layer MLP applied to Fourier features
+  output_sizes: tuple[int, int] = (32, 16)
+
+
+@chex.dataclass(eq=True)
+class SparseTransformerConfig:
+  """Sparse Transformer config."""
+  # Neighbours to attend to.
+  attention_k_hop: int
+  # Primary width, the number of channels on the carrier path.
+  d_model: int
+  # Depth, or num transformer blocks. One 'layer' is attn + ffw.
+  num_layers: int = 16
+  # Number of heads for self-attention.
+  num_heads: int = 4
+  # Attention type.
+  attention_type: str = "splash_mha"
+  # mask type if splash attention being used.
+  mask_type: str = "lazy"
+  block_q: int = 1024
+  block_kv: int = 512
+  block_kv_compute: int = 256
+  block_q_dkv: int = 512
+  block_kv_dkv: int = 1024
+  block_kv_dkv_compute: int = 1024
+  # Init scale for final ffw layer (divided by depth)
+  ffw_winit_final_mult: float = 0.0
+  # Init scale for mha w (divided by depth).
+  attn_winit_final_mult: float = 0.0
+  # Number of hidden units in the MLP blocks. Defaults to 4 * d_model.
+  ffw_hidden: int = 2048
+  # Name for haiku module.
+  name: Optional[str] = None
+
+
+@chex.dataclass(eq=True)
+class DenoiserArchitectureConfig:
+  """Defines the GenCast architecture.
+
+  Properties:
+    sparse_transformer_config: Config for the mesh transformer.
+    mesh_size: How many refinements to do on the multi-mesh.
+    latent_size: How many latent features to include in the various MLPs.
+    hidden_layers: How many hidden layers for each MLP.
+    radius_query_fraction_edge_length: Scalar that will be multiplied by the
+      length of the longest edge of the finest mesh to define the radius of
+      connectivity to use in the Grid2Mesh graph. Reasonable values are
+      between 0.6 and 1. 0.6 reduces the number of grid points feeding into
+      multiple mesh nodes and therefore reduces edge count and memory use, but
+      1 gives better predictions.
+    norm_conditioning_features: List of feature names which will be used to
+      condition the GNN via norm_conditioning, rather than as regular
+      features. If this is provided, the GNN has to support the
+      `global_norm_conditioning` argument. For now it only supports global
+      norm conditioning (e.g. the same vector conditions all edges and nodes
+      normalization), which means features passed here must not have "lat" or
+      "lon" axes. In the future it may support node level norm conditioning
+      too.
+    grid2mesh_aggregate_normalization: Optional constant to normalize the output
+      of aggregate_edges_for_nodes_fn in the mesh2grid GNN. This can be used to
+        reduce the shock the model undergoes when switching resolution, which
+        increases the number of edges connected to a node.
+    node_output_size: Size of the output node representations for
+        each node type. For node types not specified here, the latent node
+        representation from the output of the processor will be returned.
+  """
+
+  sparse_transformer_config: SparseTransformerConfig
+  mesh_size: int
+  latent_size: int = 512
+  hidden_layers: int = 1
+  radius_query_fraction_edge_length: float = 0.6
+  norm_conditioning_features: tuple[str, ...] = ("noise_level_encodings",)
+  grid2mesh_aggregate_normalization: Optional[float] = None
+  node_output_size: Optional[int] = None
+
+
+class Denoiser(base.Denoiser):
+  """Wraps a general deterministic Predictor to act as a Denoiser.
+
+  This passes an encoding of the noise level as an additional input to the
+  Predictor as an additional input 'noise_level_encodings' with shape
+  ('batch', 'noise_level_encoding_channels'). It passes the noisy_targets as
+  additional forcings (since they are also per-target-timestep data that the
+  predictor needs to condition on) with the same names as the original target
+  variables.
+  """
+
+  def __init__(
+      self,
+      noise_encoder_config: Optional[NoiseEncoderConfig],
+      denoiser_architecture_config: DenoiserArchitectureConfig,
+  ):
+    self._predictor = _DenoiserArchitecture(
+        denoiser_architecture_config=denoiser_architecture_config,
+    )
+    # Use default values if not specified.
+    if noise_encoder_config is None:
+      noise_encoder_config = NoiseEncoderConfig()
+    self._noise_level_encoder = FourierFeaturesMLP(**noise_encoder_config)
+
+  def __call__(
+      self,
+      inputs: xarray.Dataset,
+      noisy_targets: xarray.Dataset,
+      noise_levels: xarray.DataArray,
+      forcings: Optional[xarray.Dataset] = None,
+      **kwargs) -> xarray.Dataset:
+    if forcings is None: forcings = xarray.Dataset()
+    forcings = forcings.assign(noisy_targets)
+
+    if noise_levels.dims != ("batch",):
+      raise ValueError("noise_levels expected to be shape (batch,).")
+    noise_level_encodings = self._noise_level_encoder(
+        xarray_jax.unwrap_data(noise_levels)
+    )
+    noise_level_encodings = xarray_jax.Variable(
+        ("batch", "noise_level_encoding_channels"), noise_level_encodings
+    )
+    inputs = inputs.assign(noise_level_encodings=noise_level_encodings)
+
+    return self._predictor(
+        inputs=inputs,
+        targets_template=noisy_targets,
+        forcings=forcings,
+        **kwargs)
+
+
+class _DenoiserArchitecture:
+  """GenCast Predictor.
+
+  The model works on graphs that take into account:
+  * Mesh nodes: nodes for the vertices of the mesh.
+  * Grid nodes: nodes for the points of the grid.
+  * Nodes: When referring to just "nodes", this means the joint set of
+    both mesh nodes, concatenated with grid nodes.
+
+  The model works with 3 graphs:
+  * Grid2Mesh graph: Graph that contains all nodes. This graph is strictly
+    bipartite with edges going from grid nodes to mesh nodes using a
+    fixed radius query. The grid2mesh_gnn will operate in this graph. The output
+    of this stage will be a latent representation for the mesh nodes, and a
+    latent representation for the grid nodes.
+  * Mesh graph: Graph that contains mesh nodes only. The mesh_gnn will
+    operate in this graph. It will update the latent state of the mesh nodes
+    only.
+  * Mesh2Grid graph: Graph that contains all nodes. This graph is strictly
+    bipartite with edges going from mesh nodes to grid nodes such that each grid
+    node is connected to 3 nodes of the mesh triangular face that contains
+    the grid points. The mesh2grid_gnn will operate in this graph. It will
+    process the updated latent state of the mesh nodes, and the latent state
+    of the grid nodes, to produce the final output for the grid nodes.
+
+  The model is built on top of `TypedGraph`s so the different types of nodes and
+  edges can be stored and treated separately.
+  """
+
+  def __init__(
+      self,
+      denoiser_architecture_config: DenoiserArchitectureConfig,
+  ):
+    """Initializes the predictor."""
+    self._spatial_features_kwargs = dict(
+        add_node_positions=False,
+        add_node_latitude=True,
+        add_node_longitude=True,
+        add_relative_positions=True,
+        relative_longitude_local_coordinates=True,
+        relative_latitude_local_coordinates=True,
+    )
+
+    # Construct the mesh.
+    mesh = icosahedral_mesh.get_last_triangular_mesh_for_sphere(
+        splits=denoiser_architecture_config.mesh_size
+    )
+    # Permute the mesh to a banded structure so we can run sparse attention
+    # operations.
+    self._mesh = _permute_mesh_to_banded(mesh=mesh)
+
+    # Encoder, which moves data from the grid to the mesh with a single message
+    # passing step.
+    self._grid2mesh_gnn = (
+        deep_typed_graph_net.DeepTypedGraphNet(
+            activation="swish",
+            aggregate_normalization=(
+                denoiser_architecture_config.grid2mesh_aggregate_normalization
+            ),
+            edge_latent_size=dict(
+                grid2mesh=denoiser_architecture_config.latent_size
+            ),
+            embed_edges=True,
+            embed_nodes=True,
+            f32_aggregation=True,
+            include_sent_messages_in_node_update=False,
+            mlp_hidden_size=denoiser_architecture_config.latent_size,
+            mlp_num_hidden_layers=denoiser_architecture_config.hidden_layers,
+            name="grid2mesh_gnn",
+            node_latent_size=dict(
+                grid_nodes=denoiser_architecture_config.latent_size,
+                mesh_nodes=denoiser_architecture_config.latent_size
+            ),
+            node_output_size=None,
+            num_message_passing_steps=1,
+            use_layer_norm=True,
+            use_norm_conditioning=True,
+        )
+    )
+
+    # Processor - performs multiple rounds of message passing on the mesh.
+    self._mesh_gnn = transformer.MeshTransformer(
+        name="mesh_transformer",
+        transformer_ctor=sparse_transformer.Transformer,
+        transformer_kwargs=dataclasses.asdict(
+            denoiser_architecture_config.sparse_transformer_config
+        ),
+    )
+
+    # Decoder, which moves data from the mesh back into the grid with a single
+    # message passing step.
+    self._mesh2grid_gnn = (
+        deep_typed_graph_net.DeepTypedGraphNet(
+            activation="swish",
+            edge_latent_size=dict(
+                mesh2grid=denoiser_architecture_config.latent_size
+            ),
+            embed_nodes=False,
+            f32_aggregation=False,
+            include_sent_messages_in_node_update=False,
+            mlp_hidden_size=denoiser_architecture_config.latent_size,
+            mlp_num_hidden_layers=denoiser_architecture_config.hidden_layers,
+            name="mesh2grid_gnn",
+            node_latent_size=dict(
+                grid_nodes=denoiser_architecture_config.latent_size,
+                mesh_nodes=denoiser_architecture_config.latent_size,
+            ),
+            node_output_size={
+                "grid_nodes": denoiser_architecture_config.node_output_size
+            },
+            num_message_passing_steps=1,
+            use_layer_norm=True,
+            use_norm_conditioning=True,
+        )
+    )
+
+    self._norm_conditioning_features = (
+        denoiser_architecture_config.norm_conditioning_features
+    )
+    # Obtain the query radius in absolute units for the unit-sphere for the
+    # grid2mesh model, by rescaling the `radius_query_fraction_edge_length`.
+    self._query_radius = (
+        _get_max_edge_distance(self._mesh)
+        * denoiser_architecture_config.radius_query_fraction_edge_length
+    )
+
+    # Other initialization is delayed until the first call (`_maybe_init`)
+    # when we get some sample data so we know the lat/lon values.
+    self._initialized = False
+
+    # A "_init_mesh_properties":
+    # This one could be initialized at init but we delay it for consistency too.
+    self._num_mesh_nodes = None  # num_mesh_nodes
+    self._mesh_nodes_lat = None  # [num_mesh_nodes]
+    self._mesh_nodes_lon = None  # [num_mesh_nodes]
+
+    # A "_init_grid_properties":
+    self._grid_lat = None  # [num_lat_points]
+    self._grid_lon = None  # [num_lon_points]
+    self._num_grid_nodes = None  # num_lat_points * num_lon_points
+    self._grid_nodes_lat = None  # [num_grid_nodes]
+    self._grid_nodes_lon = None  # [num_grid_nodes]
+
+    # A "_init_{grid2mesh,processor,mesh2grid}_graph"
+    self._grid2mesh_graph_structure = None
+    self._mesh_graph_structure = None
+    self._mesh2grid_graph_structure = None
+
+  def __call__(self,
+               inputs: xarray.Dataset,
+               targets_template: xarray.Dataset,
+               forcings: xarray.Dataset,
+               ) -> xarray.Dataset:
+    self._maybe_init(inputs)
+
+    # Convert all input data into flat vectors for each of the grid nodes.
+    # xarray (batch, time, lat, lon, level, multiple vars, forcings)
+    # -> [num_grid_nodes, batch, num_channels]
+    grid_node_features, global_norm_conditioning = (
+        self._inputs_to_grid_node_features_and_norm_conditioning(
+            inputs, forcings
+        )
+    )
+
+    # [num_mesh_nodes, batch, latent_size], [num_grid_nodes, batch, latent_size]
+    (latent_mesh_nodes, latent_grid_nodes) = self._run_grid2mesh_gnn(
+        grid_node_features, global_norm_conditioning
+    )
+
+    # Run message passing in the multimesh.
+    # [num_mesh_nodes, batch, latent_size]
+    updated_latent_mesh_nodes = self._run_mesh_gnn(
+        latent_mesh_nodes, global_norm_conditioning
+    )
+
+    # Transfer data from the mesh to the grid.
+    # [num_grid_nodes, batch, output_size]
+    output_grid_nodes = self._run_mesh2grid_gnn(
+        updated_latent_mesh_nodes, latent_grid_nodes, global_norm_conditioning
+    )
+
+    # Convert output flat vectors for the grid nodes to the format of the
+    # output. [num_grid_nodes, batch, output_size] -> xarray (batch, one time
+    # step, lat, lon, level, multiple vars)
+    return self._grid_node_outputs_to_prediction(
+        output_grid_nodes, targets_template
+    )
+
+  def _maybe_init(self, sample_inputs: xarray.Dataset):
+    """Inits everything that has a dependency on the input coordinates."""
+    if not self._initialized:
+      self._init_mesh_properties()
+      self._init_grid_properties(
+          grid_lat=sample_inputs.lat, grid_lon=sample_inputs.lon)
+      self._grid2mesh_graph_structure = self._init_grid2mesh_graph()
+      self._mesh_graph_structure = self._init_mesh_graph()
+      self._mesh2grid_graph_structure = self._init_mesh2grid_graph()
+
+      self._initialized = True
+
+  def _init_mesh_properties(self):
+    """Inits static properties that have to do with mesh nodes."""
+    self._num_mesh_nodes = self._mesh.vertices.shape[0]
+    mesh_phi, mesh_theta = model_utils.cartesian_to_spherical(
+        self._mesh.vertices[:, 0],
+        self._mesh.vertices[:, 1],
+        self._mesh.vertices[:, 2])
+    (
+        mesh_nodes_lat,
+        mesh_nodes_lon,
+    ) = model_utils.spherical_to_lat_lon(
+        phi=mesh_phi, theta=mesh_theta)
+    # Convert to f32 to ensure the lat/lon features aren't in f64.
+    self._mesh_nodes_lat = mesh_nodes_lat.astype(np.float32)
+    self._mesh_nodes_lon = mesh_nodes_lon.astype(np.float32)
+
+  def _init_grid_properties(self, grid_lat: np.ndarray, grid_lon: np.ndarray):
+    """Inits static properties that have to do with grid nodes."""
+    self._grid_lat = grid_lat.astype(np.float32)
+    self._grid_lon = grid_lon.astype(np.float32)
+    # Initialized the counters.
+    self._num_grid_nodes = grid_lat.shape[0] * grid_lon.shape[0]
+
+    # Initialize lat and lon for the grid.
+    grid_nodes_lon, grid_nodes_lat = np.meshgrid(grid_lon, grid_lat)
+    self._grid_nodes_lon = grid_nodes_lon.reshape([-1]).astype(np.float32)
+    self._grid_nodes_lat = grid_nodes_lat.reshape([-1]).astype(np.float32)
+
+  def _init_grid2mesh_graph(self) -> typed_graph.TypedGraph:
+    """Build Grid2Mesh graph."""
+
+    # Create some edges according to distance between mesh and grid nodes.
+    assert self._grid_lat is not None and self._grid_lon is not None
+    (grid_indices, mesh_indices) = grid_mesh_connectivity.radius_query_indices(
+        grid_latitude=self._grid_lat,
+        grid_longitude=self._grid_lon,
+        mesh=self._mesh,
+        radius=self._query_radius)
+
+    # Edges sending info from grid to mesh.
+    senders = grid_indices
+    receivers = mesh_indices
+
+    # Precompute structural node and edge features according to config options.
+    # Structural features are those that depend on the fixed values of the
+    # latitude and longitudes of the nodes.
+    (senders_node_features, receivers_node_features,
+     edge_features) = model_utils.get_bipartite_graph_spatial_features(
+         senders_node_lat=self._grid_nodes_lat,
+         senders_node_lon=self._grid_nodes_lon,
+         receivers_node_lat=self._mesh_nodes_lat,
+         receivers_node_lon=self._mesh_nodes_lon,
+         senders=senders,
+         receivers=receivers,
+         edge_normalization_factor=None,
+         **self._spatial_features_kwargs,
+     )
+
+    n_grid_node = np.array([self._num_grid_nodes])
+    n_mesh_node = np.array([self._num_mesh_nodes])
+    n_edge = np.array([mesh_indices.shape[0]])
+    grid_node_set = typed_graph.NodeSet(
+        n_node=n_grid_node, features=senders_node_features)
+    mesh_node_set = typed_graph.NodeSet(
+        n_node=n_mesh_node, features=receivers_node_features)
+    edge_set = typed_graph.EdgeSet(
+        n_edge=n_edge,
+        indices=typed_graph.EdgesIndices(senders=senders, receivers=receivers),
+        features=edge_features)
+    nodes = {"grid_nodes": grid_node_set, "mesh_nodes": mesh_node_set}
+    edges = {
+        typed_graph.EdgeSetKey("grid2mesh", ("grid_nodes", "mesh_nodes")):
+            edge_set
+    }
+    grid2mesh_graph = typed_graph.TypedGraph(
+        context=typed_graph.Context(n_graph=np.array([1]), features=()),
+        nodes=nodes,
+        edges=edges)
+    return grid2mesh_graph
+
+  def _init_mesh_graph(self) -> typed_graph.TypedGraph:
+    """Build Mesh graph."""
+    # Work simply on the mesh edges.
+    # N.B.To make sure ordering is preserved, any changes to faces_to_edges here
+    # should be reflected in the other 2 calls to faces_to_edges in this file.
+    senders, receivers = icosahedral_mesh.faces_to_edges(self._mesh.faces)
+
+    # Precompute structural node and edge features according to config options.
+    # Structural features are those that depend on the fixed values of the
+    # latitude and longitudes of the nodes.
+    assert self._mesh_nodes_lat is not None and self._mesh_nodes_lon is not None
+    node_features, edge_features = model_utils.get_graph_spatial_features(
+        node_lat=self._mesh_nodes_lat,
+        node_lon=self._mesh_nodes_lon,
+        senders=senders,
+        receivers=receivers,
+        **self._spatial_features_kwargs,
+    )
+
+    n_mesh_node = np.array([self._num_mesh_nodes])
+    n_edge = np.array([senders.shape[0]])
+    assert n_mesh_node == len(node_features)
+    mesh_node_set = typed_graph.NodeSet(
+        n_node=n_mesh_node, features=node_features)
+    edge_set = typed_graph.EdgeSet(
+        n_edge=n_edge,
+        indices=typed_graph.EdgesIndices(senders=senders, receivers=receivers),
+        features=edge_features)
+    nodes = {"mesh_nodes": mesh_node_set}
+    edges = {
+        typed_graph.EdgeSetKey("mesh", ("mesh_nodes", "mesh_nodes")): edge_set
+    }
+    mesh_graph = typed_graph.TypedGraph(
+        context=typed_graph.Context(n_graph=np.array([1]), features=()),
+        nodes=nodes,
+        edges=edges)
+
+    return mesh_graph
+
+  def _init_mesh2grid_graph(self) -> typed_graph.TypedGraph:
+    """Build Mesh2Grid graph."""
+
+    # Create some edges according to how the grid nodes are contained by
+    # mesh triangles.
+    (grid_indices,
+     mesh_indices) = grid_mesh_connectivity.in_mesh_triangle_indices(
+         grid_latitude=self._grid_lat,
+         grid_longitude=self._grid_lon,
+         mesh=self._mesh)
+
+    # Edges sending info from mesh to grid.
+    senders = mesh_indices
+    receivers = grid_indices
+
+    # Precompute structural node and edge features according to config options.
+    assert self._mesh_nodes_lat is not None and self._mesh_nodes_lon is not None
+    (senders_node_features, receivers_node_features,
+     edge_features) = model_utils.get_bipartite_graph_spatial_features(
+         senders_node_lat=self._mesh_nodes_lat,
+         senders_node_lon=self._mesh_nodes_lon,
+         receivers_node_lat=self._grid_nodes_lat,
+         receivers_node_lon=self._grid_nodes_lon,
+         senders=senders,
+         receivers=receivers,
+         edge_normalization_factor=None,
+         **self._spatial_features_kwargs,
+     )
+
+    n_grid_node = np.array([self._num_grid_nodes])
+    n_mesh_node = np.array([self._num_mesh_nodes])
+    n_edge = np.array([senders.shape[0]])
+    grid_node_set = typed_graph.NodeSet(
+        n_node=n_grid_node, features=receivers_node_features)
+    mesh_node_set = typed_graph.NodeSet(
+        n_node=n_mesh_node, features=senders_node_features)
+    edge_set = typed_graph.EdgeSet(
+        n_edge=n_edge,
+        indices=typed_graph.EdgesIndices(senders=senders, receivers=receivers),
+        features=edge_features)
+    nodes = {"grid_nodes": grid_node_set, "mesh_nodes": mesh_node_set}
+    edges = {
+        typed_graph.EdgeSetKey("mesh2grid", ("mesh_nodes", "grid_nodes")):
+            edge_set
+    }
+    mesh2grid_graph = typed_graph.TypedGraph(
+        context=typed_graph.Context(n_graph=np.array([1]), features=()),
+        nodes=nodes,
+        edges=edges)
+    return mesh2grid_graph
+
+  def _run_grid2mesh_gnn(self, grid_node_features: chex.Array,
+                         global_norm_conditioning: Optional[chex.Array] = None,
+                         ) -> tuple[chex.Array, chex.Array]:
+    """Runs the grid2mesh_gnn, extracting latent mesh and grid nodes."""
+
+    # Concatenate node structural features with input features.
+    batch_size = grid_node_features.shape[1]
+
+    grid2mesh_graph = self._grid2mesh_graph_structure
+    assert grid2mesh_graph is not None
+    grid_nodes = grid2mesh_graph.nodes["grid_nodes"]
+    mesh_nodes = grid2mesh_graph.nodes["mesh_nodes"]
+    new_grid_nodes = grid_nodes._replace(
+        features=jnp.concatenate([
+            grid_node_features,
+            _add_batch_second_axis(
+                grid_nodes.features.astype(grid_node_features.dtype),
+                batch_size)
+        ],
+                                 axis=-1))
+
+    # To make sure capacity of the embedded is identical for the grid nodes and
+    # the mesh nodes, we also append some dummy zero input features for the
+    # mesh nodes.
+    dummy_mesh_node_features = jnp.zeros(
+        (self._num_mesh_nodes,) + grid_node_features.shape[1:],
+        dtype=grid_node_features.dtype)
+    new_mesh_nodes = mesh_nodes._replace(
+        features=jnp.concatenate([
+            dummy_mesh_node_features,
+            _add_batch_second_axis(
+                mesh_nodes.features.astype(dummy_mesh_node_features.dtype),
+                batch_size)
+        ],
+                                 axis=-1))
+
+    # Broadcast edge structural features to the required batch size.
+    grid2mesh_edges_key = grid2mesh_graph.edge_key_by_name("grid2mesh")
+    edges = grid2mesh_graph.edges[grid2mesh_edges_key]
+
+    new_edges = edges._replace(
+        features=_add_batch_second_axis(
+            edges.features.astype(dummy_mesh_node_features.dtype), batch_size))
+
+    input_graph = self._grid2mesh_graph_structure._replace(
+        edges={grid2mesh_edges_key: new_edges},
+        nodes={
+            "grid_nodes": new_grid_nodes,
+            "mesh_nodes": new_mesh_nodes
+        })
+
+    # Run the GNN.
+    grid2mesh_out = self._grid2mesh_gnn(input_graph, global_norm_conditioning)
+    latent_mesh_nodes = grid2mesh_out.nodes["mesh_nodes"].features
+    latent_grid_nodes = grid2mesh_out.nodes["grid_nodes"].features
+    return latent_mesh_nodes, latent_grid_nodes
+
+  def _run_mesh_gnn(self, latent_mesh_nodes: chex.Array,
+                    global_norm_conditioning: Optional[chex.Array] = None
+                    ) -> chex.Array:
+    """Runs the mesh_gnn, extracting updated latent mesh nodes."""
+
+    # Add the structural edge features of this graph. Note we don't need
+    # to add the structural node features, because these are already part of
+    # the latent state, via the original Grid2Mesh gnn, however, we need
+    # the edge ones, because it is the first time we are seeing this particular
+    # set of edges.
+    batch_size = latent_mesh_nodes.shape[1]
+
+    mesh_graph = self._mesh_graph_structure
+    assert mesh_graph is not None
+    mesh_edges_key = mesh_graph.edge_key_by_name("mesh")
+    edges = mesh_graph.edges[mesh_edges_key]
+
+    # We are assuming here that the mesh gnn uses a single set of edge keys
+    # named "mesh" for the edges and that it uses a single set of nodes named
+    # "mesh_nodes"
+    msg = ("The setup currently requires to only have one kind of edge in the"
+           " mesh GNN.")
+    assert len(mesh_graph.edges) == 1, msg
+
+    new_edges = edges._replace(
+        features=_add_batch_second_axis(
+            edges.features.astype(latent_mesh_nodes.dtype), batch_size))
+
+    nodes = mesh_graph.nodes["mesh_nodes"]
+    nodes = nodes._replace(features=latent_mesh_nodes)
+
+    input_graph = mesh_graph._replace(
+        edges={mesh_edges_key: new_edges}, nodes={"mesh_nodes": nodes})
+
+    # Run the GNN.
+    return self._mesh_gnn(input_graph,
+                          global_norm_conditioning=global_norm_conditioning
+                          ).nodes["mesh_nodes"].features
+
+  def _run_mesh2grid_gnn(self,
+                         updated_latent_mesh_nodes: chex.Array,
+                         latent_grid_nodes: chex.Array,
+                         global_norm_conditioning: Optional[chex.Array] = None,
+                         ) -> chex.Array:
+    """Runs the mesh2grid_gnn, extracting the output grid nodes."""
+
+    # Add the structural edge features of this graph. Note we don't need
+    # to add the structural node features, because these are already part of
+    # the latent state, via the original Grid2Mesh gnn, however, we need
+    # the edge ones, because it is the first time we are seeing this particular
+    # set of edges.
+    batch_size = updated_latent_mesh_nodes.shape[1]
+
+    mesh2grid_graph = self._mesh2grid_graph_structure
+    assert mesh2grid_graph is not None
+    mesh_nodes = mesh2grid_graph.nodes["mesh_nodes"]
+    grid_nodes = mesh2grid_graph.nodes["grid_nodes"]
+    new_mesh_nodes = mesh_nodes._replace(features=updated_latent_mesh_nodes)
+    new_grid_nodes = grid_nodes._replace(features=latent_grid_nodes)
+    mesh2grid_key = mesh2grid_graph.edge_key_by_name("mesh2grid")
+    edges = mesh2grid_graph.edges[mesh2grid_key]
+
+    new_edges = edges._replace(
+        features=_add_batch_second_axis(
+            edges.features.astype(latent_grid_nodes.dtype), batch_size))
+
+    input_graph = mesh2grid_graph._replace(
+        edges={mesh2grid_key: new_edges},
+        nodes={
+            "mesh_nodes": new_mesh_nodes,
+            "grid_nodes": new_grid_nodes
+        })
+
+    # Run the GNN.
+    output_graph = self._mesh2grid_gnn(input_graph, global_norm_conditioning)
+    output_grid_nodes = output_graph.nodes["grid_nodes"].features
+
+    return output_grid_nodes
+
+  def _inputs_to_grid_node_features_and_norm_conditioning(
+      self,
+      inputs: xarray.Dataset,
+      forcings: xarray.Dataset,
+      ) -> Tuple[chex.Array, Optional[chex.Array]]:
+    """xarray ->[n_grid_nodes, batch, n_channels], [batch, n_cond channels]."""
+
+    if self._norm_conditioning_features:
+      norm_conditioning_inputs = inputs[list(self._norm_conditioning_features)]
+      inputs = inputs.drop_vars(list(self._norm_conditioning_features))
+
+      if "lat" in norm_conditioning_inputs or "lon" in norm_conditioning_inputs:
+        raise ValueError("Features with lat or lon dims are not currently "
+                         "supported for norm conditioning.")
+      global_norm_conditioning = xarray_jax.unwrap_data(
+          model_utils.dataset_to_stacked(norm_conditioning_inputs,
+                                         preserved_dims=("batch",),
+                                         ).transpose("batch", ...))
+
+    else:
+      global_norm_conditioning = None
+
+    # xarray `Dataset` (batch, time, lat, lon, level, multiple vars)
+    # to xarray `DataArray` (batch, lat, lon, channels)
+    stacked_inputs = model_utils.dataset_to_stacked(inputs)
+    stacked_forcings = model_utils.dataset_to_stacked(forcings)
+    stacked_inputs = xarray.concat(
+        [stacked_inputs, stacked_forcings], dim="channels")
+
+    # xarray `DataArray` (batch, lat, lon, channels)
+    # to single numpy array with shape [lat_lon_node, batch, channels]
+    grid_xarray_lat_lon_leading = model_utils.lat_lon_to_leading_axes(
+        stacked_inputs)
+    # ["node", "batch", "features"]
+    grid_node_features = xarray_jax.unwrap(
+        grid_xarray_lat_lon_leading.data
+    ).reshape((-1,) + grid_xarray_lat_lon_leading.data.shape[2:])
+    return grid_node_features, global_norm_conditioning
+
+  def _grid_node_outputs_to_prediction(
+      self,
+      grid_node_outputs: chex.Array,
+      targets_template: xarray.Dataset,
+  ) -> xarray.Dataset:
+    """[num_grid_nodes, batch, num_outputs] -> xarray."""
+
+    # numpy array with shape [lat_lon_node, batch, channels]
+    assert self._grid_lat is not None and self._grid_lon is not None
+    grid_shape = (self._grid_lat.shape[0], self._grid_lon.shape[0])
+    grid_outputs_lat_lon_leading = grid_node_outputs.reshape(
+        grid_shape + grid_node_outputs.shape[1:])
+    dims = ("lat", "lon", "batch", "channels")
+    grid_xarray_lat_lon_leading = xarray_jax.DataArray(
+        data=grid_outputs_lat_lon_leading,
+        dims=dims)
+    grid_xarray = model_utils.restore_leading_axes(grid_xarray_lat_lon_leading)
+
+    # xarray `DataArray` (batch, lat, lon, channels)
+    # to xarray `Dataset` (batch, one time step, lat, lon, level, multiple vars)
+    return model_utils.stacked_to_dataset(
+        grid_xarray.variable, targets_template)
+
+
+def _add_batch_second_axis(data, batch_size):
+  # data [leading_dim, trailing_dim]
+  assert data.ndim == 2
+  ones = jnp.ones([batch_size, 1], dtype=data.dtype)
+  return data[:, None] * ones  # [leading_dim, batch, trailing_dim]
+
+
+def _get_max_edge_distance(mesh):
+  # N.B.To make sure ordering is preserved, any changes to faces_to_edges here
+  # should be reflected in the other 2 calls to faces_to_edges in this file.
+  senders, receivers = icosahedral_mesh.faces_to_edges(mesh.faces)
+  edge_distances = np.linalg.norm(
+      mesh.vertices[senders] - mesh.vertices[receivers], axis=-1)
+  return edge_distances.max()
+
+
+def _permute_mesh_to_banded(mesh):
+  """Permutes the mesh nodes such that adjacency matrix has banded structure."""
+  # Build adjacency matrix.
+  # N.B.To make sure ordering is preserved, any changes to faces_to_edges here
+  # should be reflected in the other 2 calls to faces_to_edges in this file.
+  senders, receivers = icosahedral_mesh.faces_to_edges(mesh.faces)
+  num_mesh_nodes = mesh.vertices.shape[0]
+  adj_mat = sparse.csr_matrix((num_mesh_nodes, num_mesh_nodes))
+  adj_mat[senders, receivers] = 1
+  # Permutation to banded (this algorithm is deterministic, a given sparse
+  # adjacency matrix will yield the same permutation every time this is run).
+  mesh_permutation = sparse.csgraph.reverse_cuthill_mckee(
+      adj_mat, symmetric_mode=True
+  )
+  vertex_permutation_map = {j: i for i, j in enumerate(mesh_permutation)}
+  permute_func = np.vectorize(lambda x: vertex_permutation_map[x])
+  return icosahedral_mesh.TriangularMesh(
+      vertices=mesh.vertices[mesh_permutation], faces=permute_func(mesh.faces)
+  )

model/graphcast/denoisers_base.py

@@ -0,0 +1,53 @@
+# Copyright 2024 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Base class for Denoisers used in diffusion Predictors.
+
+Denoisers are a bit like deterministic Predictors, except:
+* Their __call__ method also conditions on noisy_targets and the noise_levels
+  of those noisy targets
+* They don't have an overrideable loss function (the loss is assumed to be some
+  form of MSE and is implemented outside the Denoiser itself)
+"""
+
+from typing import Optional, Protocol
+
+import xarray
+
+
+class Denoiser(Protocol):
+  """A denoising model that conditions on inputs as well as noise level."""
+
+  def __call__(
+      self,
+      inputs: xarray.Dataset,
+      noisy_targets: xarray.Dataset,
+      noise_levels: xarray.DataArray,
+      forcings: Optional[xarray.Dataset] = None,
+      **kwargs) -> xarray.Dataset:
+    """Computes denoised targets from noisy targets.
+
+    Args:
+      inputs: Inputs to condition on, as for Predictor.__call__.
+      noisy_targets: Targets which have had i.i.d. zero-mean Gaussian noise
+        added to them (where the noise level used may vary along the 'batch'
+        dimension).
+      noise_levels: A DataArray with dimensions ('batch',) specifying the noise
+        levels that were used for each example in the batch.
+      forcings: Optional additional per-target-timestep forcings to condition
+        on, as for Predictor.__call__.
+      **kwargs: Any additional custom kwargs.
+
+    Returns:
+      Denoised predictions with the same shape as noisy_targets.
+    """

model/graphcast/dpm_solver_plus_plus_2s.py

@@ -0,0 +1,187 @@
+# Copyright 2024 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""DPM-Solver++ 2S sampler from https://arxiv.org/abs/2211.01095."""
+
+from typing import Optional
+
+from . import casting
+from . import denoisers_base
+from . import samplers_base as base
+from . import samplers_utils as utils
+from . import xarray_jax
+import haiku as hk
+import jax.numpy as jnp
+import xarray
+
+
+class Sampler(base.Sampler):
+  """Sampling using DPM-Solver++ 2S from [1].
+
+  This is combined with optional stochastic churn as described in [2].
+
+  The '2S' terminology from [1] means that this is a second-order (2),
+  single-step (S) solver. Here 'single-step' here distinguishes it from
+  'multi-step' methods where the results of function evaluations from previous
+  steps are reused in computing updates for subsequent steps. The solver still
+  uses multiple steps though.
+
+  [1] DPM-Solver++: Fast Solver for Guided Sampling of Diffusion Probabilistic
+  Models, https://arxiv.org/abs/2211.01095
+  [2] Elucidating the Design Space of Diffusion-Based Generative Models,
+  https://arxiv.org/abs/2206.00364
+  """
+
+  def __init__(self,
+               denoiser: denoisers_base.Denoiser,
+               max_noise_level: float,
+               min_noise_level: float,
+               num_noise_levels: int,
+               rho: float,
+               stochastic_churn_rate: float,
+               churn_min_noise_level: float,
+               churn_max_noise_level: float,
+               noise_level_inflation_factor: float
+               ):
+    """Initializes the sampler.
+
+    Args:
+      denoiser: A Denoiser which predicts noise-free targets.
+      max_noise_level: The highest noise level used at the start of the
+        sequence of reverse diffusion steps.
+      min_noise_level: The lowest noise level used at the end of the sequence of
+        reverse diffusion steps.
+      num_noise_levels: Determines the number of noise levels used and hence the
+        number of reverse diffusion steps performed.
+      rho: Parameter affecting the spacing of noise steps. Higher values will
+        concentrate noise steps more around zero.
+      stochastic_churn_rate: S_churn from the paper. This controls the rate
+        at which noise is re-injected/'churned' during the sampling algorithm.
+        If this is set to zero then we are performing deterministic sampling
+        as described in Algorithm 1.
+      churn_min_noise_level: Minimum noise level at which stochastic churn
+        occurs. S_min from the paper. Only used if stochastic_churn_rate > 0.
+      churn_max_noise_level: Maximum noise level at which stochastic churn
+        occurs. S_min from the paper. Only used if stochastic_churn_rate > 0.
+      noise_level_inflation_factor: This can be used to set the actual amount of
+        noise injected higher than what the denoiser is told has been added.
+        The motivation is to compensate for a tendency of L2-trained denoisers
+        to remove slightly too much noise / blur too much. S_noise from the
+        paper. Only used if stochastic_churn_rate > 0.
+    """
+    super().__init__(denoiser)
+    self._noise_levels = utils.noise_schedule(
+        max_noise_level, min_noise_level, num_noise_levels, rho)
+    self._stochastic_churn = stochastic_churn_rate > 0
+    self._per_step_churn_rates = utils.stochastic_churn_rate_schedule(
+        self._noise_levels, stochastic_churn_rate, churn_min_noise_level,
+        churn_max_noise_level)
+    self._noise_level_inflation_factor = noise_level_inflation_factor
+
+  def __call__(
+      self,
+      inputs: xarray.Dataset,
+      targets_template: xarray.Dataset,
+      forcings: Optional[xarray.Dataset] = None,
+      **kwargs) -> xarray.Dataset:
+
+    dtype = casting.infer_floating_dtype(targets_template)  # pytype: disable=wrong-arg-types
+    noise_levels = jnp.array(self._noise_levels).astype(dtype)
+    per_step_churn_rates = jnp.array(self._per_step_churn_rates).astype(dtype)
+
+    def denoiser(noise_level: jnp.ndarray, x: xarray.Dataset) -> xarray.Dataset:
+      """Computes D(x, sigma, y)."""
+      bcast_noise_level = xarray_jax.DataArray(
+          jnp.tile(noise_level, x.sizes['batch']), dims=('batch',))
+      # Estimate the expectation of the fully-denoised target x0, conditional on
+      # inputs/forcings, noisy targets and their noise level:
+      return self._denoiser(
+          inputs=inputs,
+          noisy_targets=x,
+          noise_levels=bcast_noise_level,
+          forcings=forcings)
+
+    def body_fn(i: jnp.ndarray, x: xarray.Dataset) -> xarray.Dataset:
+      """One iteration of the sampling algorithm.
+
+      Args:
+        i: Sampling iteration.
+        x: Noisy targets at iteration i, these will have noise level
+          self._noise_levels[i].
+
+      Returns:
+        Noisy targets at the next lowest noise level self._noise_levels[i+1].
+      """
+      def init_noise(template):
+        return noise_levels[0] * utils.spherical_white_noise_like(template)
+
+      # Initialise the inputs if i == 0.
+      # This is done here to ensure both noise sampler calls can use the same
+      # spherical harmonic basis functions. While there may be a small compute
+      # cost the memory savings can be significant.
+      # TODO(dominicmasters): Figure out if we can merge the two noise sampler
+      # calls into one to avoid this hack.
+      maybe_init_noise = (i == 0).astype(noise_levels[0].dtype)
+      x = x + init_noise(x) * maybe_init_noise
+
+      noise_level = noise_levels[i]
+
+      if self._stochastic_churn:
+        # We increase the noise level of x a bit before taking it down again:
+        x, noise_level = utils.apply_stochastic_churn(
+            x, noise_level,
+            stochastic_churn_rate=per_step_churn_rates[i],
+            noise_level_inflation_factor=self._noise_level_inflation_factor)
+
+      # Apply one step of the ODE solver to take x down to the next lowest
+      # noise level.
+
+      # Note that the Elucidating paper's choice of sigma(t)=t and s(t)=1
+      # (corresponding to alpha(t)=1 in the DPM paper) as well as the standard
+      # choice of r=1/2 (corresponding to a geometric mean for the s_i
+      # midpoints) greatly simplifies the update from the DPM-Solver++ paper.
+      # You need to do a bit of algebraic fiddling to arrive at the below after
+      # substituting these choices into DPMSolver++'s Algorithm 1. The simpler
+      # update we arrive at helps with intuition too.
+
+      next_noise_level = noise_levels[i + 1]
+      # This is s_{i+1} from the paper. They don't explain how the s_i are
+      # chosen, but the default choice seems to be a geometric mean, which is
+      # equivalent to setting all the r_i = 1/2.
+      mid_noise_level = jnp.sqrt(noise_level * next_noise_level)
+
+      mid_over_current = mid_noise_level / noise_level
+      x_denoised = denoiser(noise_level, x)
+      # This turns out to be a convex combination of current and denoised x,
+      # which isn't entirely apparent from the paper formulae:
+      x_mid = mid_over_current * x + (1 - mid_over_current) * x_denoised
+
+      next_over_current = next_noise_level / noise_level
+      x_mid_denoised = denoiser(mid_noise_level, x_mid)  # pytype: disable=wrong-arg-types
+      x_next = next_over_current * x + (1 - next_over_current) * x_mid_denoised
+
+      # For the final step to noise level 0, we do an Euler update which
+      # corresponds to just returning the denoiser's prediction directly.
+      #
+      # In fact the behaviour above when next_noise_level == 0 is almost
+      # equivalent, except that it runs the denoiser a second time to denoise
+      # from noise level 0. The denoiser should just be the identity function in
+      # this case, but it hasn't necessarily been trained at noise level 0 so
+      # we avoid relying on this.
+      return utils.tree_where(next_noise_level == 0, x_denoised, x_next)
+
+    # Init with zeros but apply additional noise at step 0 to initialise the
+    # state.
+    noise_init = xarray.zeros_like(targets_template)
+    return hk.fori_loop(
+        0, len(noise_levels) - 1, body_fun=body_fn, init_val=noise_init)

model/graphcast/gencast.py

@@ -0,0 +1,284 @@
+# Copyright 2024 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Denoising diffusion models based on the framework of [1].
+
+Throughout we will refer to notation and equations from [1].
+
+  [1] Elucidating the Design Space of Diffusion-Based Generative Models
+  Karras, Aittala, Aila and Laine, 2022
+  https://arxiv.org/abs/2206.00364
+"""
+
+from typing import Any, Optional, Tuple
+
+import chex
+from . import casting
+from . import denoiser
+from . import dpm_solver_plus_plus_2s
+from . import graphcast
+from . import losses
+from . import predictor_base
+from . import samplers_utils
+from . import xarray_jax
+import haiku as hk
+import jax
+import xarray
+
+
+TARGET_SURFACE_VARS = (
+    '2m_temperature',
+    'mean_sea_level_pressure',
+    '10m_v_component_of_wind',
+    '10m_u_component_of_wind',  # GenCast predicts in 12hr timesteps.
+    'total_precipitation_12hr',
+    'sea_surface_temperature',
+)
+
+TARGET_SURFACE_NO_PRECIP_VARS = (
+    '2m_temperature',
+    'mean_sea_level_pressure',
+    '10m_v_component_of_wind',
+    '10m_u_component_of_wind',
+    'sea_surface_temperature',
+)
+
+
+TASK = graphcast.TaskConfig(
+    input_variables=(
+        # GenCast doesn't take precipitation as input.
+        TARGET_SURFACE_NO_PRECIP_VARS
+        + graphcast.TARGET_ATMOSPHERIC_VARS
+        + graphcast.GENERATED_FORCING_VARS
+        + graphcast.STATIC_VARS
+    ),
+    target_variables=TARGET_SURFACE_VARS + graphcast.TARGET_ATMOSPHERIC_VARS,
+    # GenCast doesn't take incident solar radiation as a forcing.
+    forcing_variables=graphcast.GENERATED_FORCING_VARS,
+    pressure_levels=graphcast.PRESSURE_LEVELS_WEATHERBENCH_13,
+    # GenCast takes the current frame and the frame 12 hours prior.
+    input_duration='24h',
+)
+
+
+@chex.dataclass(frozen=True, eq=True)
+class SamplerConfig:
+  """Configures the sampler used to draw samples from GenCast.
+
+      max_noise_level: The highest noise level used at the start of the
+        sequence of reverse diffusion steps.
+      min_noise_level: The lowest noise level used at the end of the sequence of
+        reverse diffusion steps.
+      num_noise_levels: Determines the number of noise levels used and hence the
+        number of reverse diffusion steps performed.
+      rho: Parameter affecting the spacing of noise steps. Higher values will
+        concentrate noise steps more around zero.
+      stochastic_churn_rate: S_churn from the paper. This controls the rate
+        at which noise is re-injected/'churned' during the sampling algorithm.
+        If this is set to zero then we are performing deterministic sampling
+        as described in Algorithm 1.
+      churn_max_noise_level: Maximum noise level at which stochastic churn
+        occurs. S_min from the paper. Only used if stochastic_churn_rate > 0.
+      churn_min_noise_level: Minimum noise level at which stochastic churn
+        occurs. S_min from the paper. Only used if stochastic_churn_rate > 0.
+      noise_level_inflation_factor: This can be used to set the actual amount of
+        noise injected higher than what the denoiser is told has been added.
+        The motivation is to compensate for a tendency of L2-trained denoisers
+        to remove slightly too much noise / blur too much. S_noise from the
+        paper. Only used if stochastic_churn_rate > 0.
+  """
+  max_noise_level: float = 80.
+  min_noise_level: float = 0.03
+  num_noise_levels: int = 20
+  rho: float = 7.
+  # Stochastic sampler settings.
+  stochastic_churn_rate: float = 2.5
+  churn_min_noise_level: float = 0.75
+  churn_max_noise_level: float = float('inf')
+  noise_level_inflation_factor: float = 1.05
+
+
+@chex.dataclass(frozen=True, eq=True)
+class NoiseConfig:
+  training_noise_level_rho: float = 7.0
+  training_max_noise_level: float = 88.0
+  training_min_noise_level: float = 0.02
+
+
+@chex.dataclass(frozen=True, eq=True)
+class CheckPoint:
+  description: str
+  license: str
+  params: dict[str, Any]
+  task_config: graphcast.TaskConfig
+  denoiser_architecture_config: denoiser.DenoiserArchitectureConfig
+  sampler_config: SamplerConfig
+  noise_config: NoiseConfig
+  noise_encoder_config: denoiser.NoiseEncoderConfig
+
+
+class GenCast(predictor_base.Predictor):
+  """Predictor for a denoising diffusion model following the framework of [1].
+
+    [1] Elucidating the Design Space of Diffusion-Based Generative Models
+    Karras, Aittala, Aila and Laine, 2022
+    https://arxiv.org/abs/2206.00364
+
+  Unlike the paper, we have a conditional model and our denoising function
+  conditions on previous timesteps.
+
+  As the paper demonstrates, the sampling algorithm can be varied independently
+  of the denoising model and its training procedure, and it is separately
+  configurable here.
+  """
+
+  def __init__(
+      self,
+      task_config: graphcast.TaskConfig,
+      denoiser_architecture_config: denoiser.DenoiserArchitectureConfig,
+      sampler_config: Optional[SamplerConfig] = None,
+      noise_config: Optional[NoiseConfig] = None,
+      noise_encoder_config: Optional[denoiser.NoiseEncoderConfig] = None,
+  ):
+    """Constructs GenCast."""
+    # Output size depends on number of variables being predicted.
+    num_surface_vars = len(
+        set(task_config.target_variables)
+        - set(graphcast.ALL_ATMOSPHERIC_VARS)
+    )
+    num_atmospheric_vars = len(
+        set(task_config.target_variables)
+        & set(graphcast.ALL_ATMOSPHERIC_VARS)
+    )
+    num_outputs = (
+        num_surface_vars
+        + len(task_config.pressure_levels) * num_atmospheric_vars
+    )
+    denoiser_architecture_config.node_output_size = num_outputs
+    self._denoiser = denoiser.Denoiser(
+        noise_encoder_config,
+        denoiser_architecture_config,
+    )
+    self._sampler_config = sampler_config
+    # Singleton to avoid re-initializing the sampler for each inference call.
+    self._sampler = None
+    self._noise_config = noise_config
+
+  def _c_in(self, noise_scale: xarray.DataArray) -> xarray.DataArray:
+    """Scaling applied to the noisy targets input to the underlying network."""
+    return (noise_scale**2 + 1)**-0.5
+
+  def _c_out(self, noise_scale: xarray.DataArray) -> xarray.DataArray:
+    """Scaling applied to the underlying network's raw outputs."""
+    return noise_scale * (noise_scale**2 + 1)**-0.5
+
+  def _c_skip(self, noise_scale: xarray.DataArray) -> xarray.DataArray:
+    """Scaling applied to the skip connection."""
+    return 1 / (noise_scale**2 + 1)
+
+  def _loss_weighting(self, noise_scale: xarray.DataArray) -> xarray.DataArray:
+    r"""The loss weighting \lambda(\sigma) from the paper."""
+    return self._c_out(noise_scale) ** -2
+
+  def _preconditioned_denoiser(
+      self,
+      inputs: xarray.Dataset,
+      noisy_targets: xarray.Dataset,
+      noise_levels: xarray.DataArray,
+      forcings: Optional[xarray.Dataset] = None,
+      **kwargs) -> xarray.Dataset:
+    """The preconditioned denoising function D from the paper (Eqn 7)."""
+    raw_predictions = self._denoiser(
+        inputs=inputs,
+        noisy_targets=noisy_targets * self._c_in(noise_levels),
+        noise_levels=noise_levels,
+        forcings=forcings,
+        **kwargs)
+    return (raw_predictions * self._c_out(noise_levels) +
+            noisy_targets * self._c_skip(noise_levels))
+
+  def loss_and_predictions(
+      self,
+      inputs: xarray.Dataset,
+      targets: xarray.Dataset,
+      forcings: Optional[xarray.Dataset] = None,
+  ) -> Tuple[predictor_base.LossAndDiagnostics, xarray.Dataset]:
+    return self.loss(inputs, targets, forcings), self(inputs, targets, forcings)
+
+  def loss(self,
+           inputs: xarray.Dataset,
+           targets: xarray.Dataset,
+           forcings: Optional[xarray.Dataset] = None,
+           ) -> predictor_base.LossAndDiagnostics:
+
+    if self._noise_config is None:
+      raise ValueError('Noise config must be specified to train GenCast.')
+
+    # Sample noise levels:
+    dtype = casting.infer_floating_dtype(targets)  # pytype: disable=wrong-arg-types
+    key = hk.next_rng_key()
+    batch_size = inputs.sizes['batch']
+    noise_levels = xarray_jax.DataArray(
+        data=samplers_utils.rho_inverse_cdf(
+            min_value=self._noise_config.training_min_noise_level,
+            max_value=self._noise_config.training_max_noise_level,
+            rho=self._noise_config.training_noise_level_rho,
+            cdf=jax.random.uniform(key, shape=(batch_size,), dtype=dtype)),
+        dims=('batch',))
+
+    # Sample noise and apply it to targets:
+    noise = (
+        samplers_utils.spherical_white_noise_like(targets) * noise_levels
+    )
+    noisy_targets = targets + noise
+
+    denoised_predictions = self._preconditioned_denoiser(
+        inputs, noisy_targets, noise_levels, forcings)
+
+    loss, diagnostics = losses.weighted_mse_per_level(
+        denoised_predictions,
+        targets,
+        # Weights are same as we used for GraphCast.
+        per_variable_weights={
+            # Any variables not specified here are weighted as 1.0.
+            # A single-level variable, but an important headline variable
+            # and also one which we have struggled to get good performance
+            # on at short lead times, so leaving it weighted at 1.0, equal
+            # to the multi-level variables:
+            '2m_temperature': 1.0,
+            # New single-level variables, which we don't weight too highly
+            # to avoid hurting performance on other variables.
+            '10m_u_component_of_wind': 0.1,
+            '10m_v_component_of_wind': 0.1,
+            'mean_sea_level_pressure': 0.1,
+            'sea_surface_temperature': 0.1,
+            'total_precipitation_12hr': 0.1
+        },
+    )
+    loss *= self._loss_weighting(noise_levels)
+    return loss, diagnostics
+
+  def __call__(self,
+               inputs: xarray.Dataset,
+               targets_template: xarray.Dataset,
+               forcings: Optional[xarray.Dataset] = None,
+               **kwargs) -> xarray.Dataset:
+    if self._sampler_config is None:
+      raise ValueError(
+          'Sampler config must be specified to run inference on GenCast.'
+      )
+    if self._sampler is None:
+      self._sampler = dpm_solver_plus_plus_2s.Sampler(
+          self._preconditioned_denoiser, **self._sampler_config
+      )
+    return self._sampler(inputs, targets_template, forcings, **kwargs)

model/graphcast/graphcast.py

@@ -0,0 +1,796 @@
+# Copyright 2023 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""A predictor that runs multiple graph neural networks on mesh data.
+
+It learns to interpolate between the grid and the mesh nodes, with the loss
+and the rollouts ultimately computed at the grid level.
+
+It uses ideas similar to those in Keisler (2022):
+
+Reference:
+  https://arxiv.org/pdf/2202.07575.pdf
+
+It assumes data across time and level is stacked, and operates only operates in
+a 2D mesh over latitudes and longitudes.
+"""
+
+from typing import Any, Callable, Mapping, Optional
+
+import chex
+from . import deep_typed_graph_net
+from . import grid_mesh_connectivity
+from . import icosahedral_mesh
+from . import losses
+from . import model_utils
+from . import predictor_base
+from . import typed_graph
+from . import xarray_jax
+import jax.numpy as jnp
+import jraph
+import numpy as np
+import xarray
+
+Kwargs = Mapping[str, Any]
+
+GNN = Callable[[jraph.GraphsTuple], jraph.GraphsTuple]
+
+
+# https://www.ecmwf.int/en/forecasts/dataset/ecmwf-reanalysis-v5
+PRESSURE_LEVELS_ERA5_37 = (
+    1, 2, 3, 5, 7, 10, 20, 30, 50, 70, 100, 125, 150, 175, 200, 225, 250, 300,
+    350, 400, 450, 500, 550, 600, 650, 700, 750, 775, 800, 825, 850, 875, 900,
+    925, 950, 975, 1000)
+
+# https://www.ecmwf.int/en/forecasts/datasets/set-i
+PRESSURE_LEVELS_HRES_25 = (
+    1, 2, 3, 5, 7, 10, 20, 30, 50, 70, 100, 150, 200, 250, 300, 400, 500, 600,
+    700, 800, 850, 900, 925, 950, 1000)
+
+# https://agupubs.onlinelibrary.wiley.com/doi/full/10.1029/2020MS002203
+PRESSURE_LEVELS_WEATHERBENCH_13 = (
+    50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)
+
+PRESSURE_LEVELS = {
+    13: PRESSURE_LEVELS_WEATHERBENCH_13,
+    25: PRESSURE_LEVELS_HRES_25,
+    37: PRESSURE_LEVELS_ERA5_37,
+}
+
+# The list of all possible atmospheric variables. Taken from:
+# https://confluence.ecmwf.int/display/CKB/ERA5%3A+data+documentation#ERA5:datadocumentation-Table9
+ALL_ATMOSPHERIC_VARS = (
+    "potential_vorticity",
+    "specific_rain_water_content",
+    "specific_snow_water_content",
+    "geopotential",
+    "temperature",
+    "u_component_of_wind",
+    "v_component_of_wind",
+    "specific_humidity",
+    "vertical_velocity",
+    "vorticity",
+    "divergence",
+    "relative_humidity",
+    "ozone_mass_mixing_ratio",
+    "specific_cloud_liquid_water_content",
+    "specific_cloud_ice_water_content",
+    "fraction_of_cloud_cover",
+)
+
+TARGET_SURFACE_VARS = (
+    "2m_temperature",
+    "mean_sea_level_pressure",
+    "10m_v_component_of_wind",
+    "10m_u_component_of_wind",
+    "total_precipitation_6hr",
+)
+TARGET_SURFACE_NO_PRECIP_VARS = (
+    "2m_temperature",
+    "mean_sea_level_pressure",
+    "10m_v_component_of_wind",
+    "10m_u_component_of_wind",
+)
+TARGET_ATMOSPHERIC_VARS = (
+    "temperature",
+    "geopotential",
+    "u_component_of_wind",
+    "v_component_of_wind",
+    "vertical_velocity",
+    "specific_humidity",
+)
+TARGET_ATMOSPHERIC_NO_W_VARS = (
+    "temperature",
+    "geopotential",
+    "u_component_of_wind",
+    "v_component_of_wind",
+    "specific_humidity",
+)
+EXTERNAL_FORCING_VARS = (
+    "toa_incident_solar_radiation",
+)
+GENERATED_FORCING_VARS = (
+    "year_progress_sin",
+    "year_progress_cos",
+    "day_progress_sin",
+    "day_progress_cos",
+)
+FORCING_VARS = EXTERNAL_FORCING_VARS + GENERATED_FORCING_VARS
+STATIC_VARS = (
+    "geopotential_at_surface",
+    "land_sea_mask",
+)
+
+
+@chex.dataclass(frozen=True, eq=True)
+class TaskConfig:
+  """Defines inputs and targets on which a model is trained and/or evaluated."""
+  input_variables: tuple[str, ...]
+  # Target variables which the model is expected to predict.
+  target_variables: tuple[str, ...]
+  forcing_variables: tuple[str, ...]
+  pressure_levels: tuple[int, ...]
+  input_duration: str
+
+TASK = TaskConfig(
+    input_variables=(
+        TARGET_SURFACE_VARS + TARGET_ATMOSPHERIC_VARS + FORCING_VARS +
+        STATIC_VARS),
+    target_variables=TARGET_SURFACE_VARS + TARGET_ATMOSPHERIC_VARS,
+    forcing_variables=FORCING_VARS,
+    pressure_levels=PRESSURE_LEVELS_ERA5_37,
+    input_duration="12h",
+)
+TASK_13 = TaskConfig(
+    input_variables=(
+        TARGET_SURFACE_VARS + TARGET_ATMOSPHERIC_VARS + FORCING_VARS +
+        STATIC_VARS),
+    target_variables=TARGET_SURFACE_VARS + TARGET_ATMOSPHERIC_VARS,
+    forcing_variables=FORCING_VARS,
+    pressure_levels=PRESSURE_LEVELS_WEATHERBENCH_13,
+    input_duration="12h",
+)
+TASK_13_PRECIP_OUT = TaskConfig(
+    input_variables=(
+        TARGET_SURFACE_NO_PRECIP_VARS + TARGET_ATMOSPHERIC_VARS + FORCING_VARS +
+        STATIC_VARS),
+    target_variables=TARGET_SURFACE_VARS + TARGET_ATMOSPHERIC_VARS,
+    forcing_variables=FORCING_VARS,
+    pressure_levels=PRESSURE_LEVELS_WEATHERBENCH_13,
+    input_duration="12h",
+)
+
+
+@chex.dataclass(frozen=True, eq=True)
+class ModelConfig:
+  """Defines the architecture of the GraphCast neural network architecture.
+
+  Properties:
+    resolution: The resolution of the data, in degrees (e.g. 0.25 or 1.0).
+    mesh_size: How many refinements to do on the multi-mesh.
+    gnn_msg_steps: How many Graph Network message passing steps to do.
+    latent_size: How many latent features to include in the various MLPs.
+    hidden_layers: How many hidden layers for each MLP.
+    radius_query_fraction_edge_length: Scalar that will be multiplied by the
+        length of the longest edge of the finest mesh to define the radius of
+        connectivity to use in the Grid2Mesh graph. Reasonable values are
+        between 0.6 and 1. 0.6 reduces the number of grid points feeding into
+        multiple mesh nodes and therefore reduces edge count and memory use, but
+        1 gives better predictions.
+    mesh2grid_edge_normalization_factor: Allows explicitly controlling edge
+        normalization for mesh2grid edges. If None, defaults to max edge length.
+        This supports using pre-trained model weights with a different graph
+        structure to what it was trained on.
+  """
+  resolution: float
+  mesh_size: int
+  latent_size: int
+  gnn_msg_steps: int
+  hidden_layers: int
+  radius_query_fraction_edge_length: float
+  mesh2grid_edge_normalization_factor: Optional[float] = None
+
+
+@chex.dataclass(frozen=True, eq=True)
+class CheckPoint:
+  params: dict[str, Any]
+  model_config: ModelConfig
+  task_config: TaskConfig
+  description: str
+  license: str
+
+
+class GraphCast(predictor_base.Predictor):
+  """GraphCast Predictor.
+
+  The model works on graphs that take into account:
+  * Mesh nodes: nodes for the vertices of the mesh.
+  * Grid nodes: nodes for the points of the grid.
+  * Nodes: When referring to just "nodes", this means the joint set of
+    both mesh nodes, concatenated with grid nodes.
+
+  The model works with 3 graphs:
+  * Grid2Mesh graph: Graph that contains all nodes. This graph is strictly
+    bipartite with edges going from grid nodes to mesh nodes using a
+    fixed radius query. The grid2mesh_gnn will operate in this graph. The output
+    of this stage will be a latent representation for the mesh nodes, and a
+    latent representation for the grid nodes.
+  * Mesh graph: Graph that contains mesh nodes only. The mesh_gnn will
+    operate in this graph. It will update the latent state of the mesh nodes
+    only.
+  * Mesh2Grid graph: Graph that contains all nodes. This graph is strictly
+    bipartite with edges going from mesh nodes to grid nodes such that each grid
+    nodes is connected to 3 nodes of the mesh triangular face that contains
+    the grid points. The mesh2grid_gnn will operate in this graph. It will
+    process the updated latent state of the mesh nodes, and the latent state
+    of the grid nodes, to produce the final output for the grid nodes.
+
+  The model is built on top of `TypedGraph`s so the different types of nodes and
+  edges can be stored and treated separately.
+
+  """
+
+  def __init__(self, model_config: ModelConfig, task_config: TaskConfig):
+    """Initializes the predictor."""
+    self._spatial_features_kwargs = dict(
+        add_node_positions=False,
+        add_node_latitude=True,
+        add_node_longitude=True,
+        add_relative_positions=True,
+        relative_longitude_local_coordinates=True,
+        relative_latitude_local_coordinates=True,
+    )
+
+    # Specification of the multimesh.
+    self._meshes = (
+        icosahedral_mesh.get_hierarchy_of_triangular_meshes_for_sphere(
+            splits=model_config.mesh_size))
+
+    # Encoder, which moves data from the grid to the mesh with a single message
+    # passing step.
+    self._grid2mesh_gnn = deep_typed_graph_net.DeepTypedGraphNet(
+        embed_nodes=True,  # Embed raw features of the grid and mesh nodes.
+        embed_edges=True,  # Embed raw features of the grid2mesh edges.
+        edge_latent_size=dict(grid2mesh=model_config.latent_size),
+        node_latent_size=dict(
+            mesh_nodes=model_config.latent_size,
+            grid_nodes=model_config.latent_size),
+        mlp_hidden_size=model_config.latent_size,
+        mlp_num_hidden_layers=model_config.hidden_layers,
+        num_message_passing_steps=1,
+        use_layer_norm=True,
+        include_sent_messages_in_node_update=False,
+        activation="swish",
+        f32_aggregation=True,
+        aggregate_normalization=None,
+        name="grid2mesh_gnn",
+    )
+
+    # Processor, which performs message passing on the multi-mesh.
+    self._mesh_gnn = deep_typed_graph_net.DeepTypedGraphNet(
+        embed_nodes=False,  # Node features already embdded by previous layers.
+        embed_edges=True,  # Embed raw features of the multi-mesh edges.
+        node_latent_size=dict(mesh_nodes=model_config.latent_size),
+        edge_latent_size=dict(mesh=model_config.latent_size),
+        mlp_hidden_size=model_config.latent_size,
+        mlp_num_hidden_layers=model_config.hidden_layers,
+        num_message_passing_steps=model_config.gnn_msg_steps,
+        use_layer_norm=True,
+        include_sent_messages_in_node_update=False,
+        activation="swish",
+        f32_aggregation=False,
+        name="mesh_gnn",
+    )
+
+    num_surface_vars = len(
+        set(task_config.target_variables) - set(ALL_ATMOSPHERIC_VARS))
+    num_atmospheric_vars = len(
+        set(task_config.target_variables) & set(ALL_ATMOSPHERIC_VARS))
+    num_outputs = (num_surface_vars +
+                   len(task_config.pressure_levels) * num_atmospheric_vars)
+
+    # Decoder, which moves data from the mesh back into the grid with a single
+    # message passing step.
+    self._mesh2grid_gnn = deep_typed_graph_net.DeepTypedGraphNet(
+        # Require a specific node dimensionaly for the grid node outputs.
+        node_output_size=dict(grid_nodes=num_outputs),
+        embed_nodes=False,  # Node features already embdded by previous layers.
+        embed_edges=True,  # Embed raw features of the mesh2grid edges.
+        edge_latent_size=dict(mesh2grid=model_config.latent_size),
+        node_latent_size=dict(
+            mesh_nodes=model_config.latent_size,
+            grid_nodes=model_config.latent_size),
+        mlp_hidden_size=model_config.latent_size,
+        mlp_num_hidden_layers=model_config.hidden_layers,
+        num_message_passing_steps=1,
+        use_layer_norm=True,
+        include_sent_messages_in_node_update=False,
+        activation="swish",
+        f32_aggregation=False,
+        name="mesh2grid_gnn",
+    )
+
+    # Obtain the query radius in absolute units for the unit-sphere for the
+    # grid2mesh model, by rescaling the `radius_query_fraction_edge_length`.
+    self._query_radius = (_get_max_edge_distance(self._finest_mesh)
+                          * model_config.radius_query_fraction_edge_length)
+    self._mesh2grid_edge_normalization_factor = (
+        model_config.mesh2grid_edge_normalization_factor
+    )
+
+    # Other initialization is delayed until the first call (`_maybe_init`)
+    # when we get some sample data so we know the lat/lon values.
+    self._initialized = False
+
+    # A "_init_mesh_properties":
+    # This one could be initialized at init but we delay it for consistency too.
+    self._num_mesh_nodes = None  # num_mesh_nodes
+    self._mesh_nodes_lat = None  # [num_mesh_nodes]
+    self._mesh_nodes_lon = None  # [num_mesh_nodes]
+
+    # A "_init_grid_properties":
+    self._grid_lat = None  # [num_lat_points]
+    self._grid_lon = None  # [num_lon_points]
+    self._num_grid_nodes = None  # num_lat_points * num_lon_points
+    self._grid_nodes_lat = None  # [num_grid_nodes]
+    self._grid_nodes_lon = None  # [num_grid_nodes]
+
+    # A "_init_{grid2mesh,processor,mesh2grid}_graph"
+    self._grid2mesh_graph_structure = None
+    self._mesh_graph_structure = None
+    self._mesh2grid_graph_structure = None
+
+  @property
+  def _finest_mesh(self):
+    return self._meshes[-1]
+
+  def __call__(self,
+               inputs: xarray.Dataset,
+               targets_template: xarray.Dataset,
+               forcings: xarray.Dataset,
+               is_training: bool = False,
+               ) -> xarray.Dataset:
+    self._maybe_init(inputs)
+
+    # Convert all input data into flat vectors for each of the grid nodes.
+    # xarray (batch, time, lat, lon, level, multiple vars, forcings)
+    # -> [num_grid_nodes, batch, num_channels]
+    grid_node_features = self._inputs_to_grid_node_features(inputs, forcings)
+
+    # Transfer data for the grid to the mesh,
+    # [num_mesh_nodes, batch, latent_size], [num_grid_nodes, batch, latent_size]
+    (latent_mesh_nodes, latent_grid_nodes
+     ) = self._run_grid2mesh_gnn(grid_node_features)
+
+    # Run message passing in the multimesh.
+    # [num_mesh_nodes, batch, latent_size]
+    updated_latent_mesh_nodes = self._run_mesh_gnn(latent_mesh_nodes)
+
+    # Transfer data frome the mesh to the grid.
+    # [num_grid_nodes, batch, output_size]
+    output_grid_nodes = self._run_mesh2grid_gnn(
+        updated_latent_mesh_nodes, latent_grid_nodes)
+
+    # Conver output flat vectors for the grid nodes to the format of the output.
+    # [num_grid_nodes, batch, output_size] ->
+    # xarray (batch, one time step, lat, lon, level, multiple vars)
+    return self._grid_node_outputs_to_prediction(
+        output_grid_nodes, targets_template)
+
+  def loss_and_predictions(  # pytype: disable=signature-mismatch  # jax-ndarray
+      self,
+      inputs: xarray.Dataset,
+      targets: xarray.Dataset,
+      forcings: xarray.Dataset,
+      ) -> tuple[predictor_base.LossAndDiagnostics, xarray.Dataset]:
+    # Forward pass.
+    predictions = self(
+        inputs, targets_template=targets, forcings=forcings, is_training=True)
+    # Compute loss.
+    loss = losses.weighted_mse_per_level(
+        predictions, targets,
+        per_variable_weights={
+            # Any variables not specified here are weighted as 1.0.
+            # A single-level variable, but an important headline variable
+            # and also one which we have struggled to get good performance
+            # on at short lead times, so leaving it weighted at 1.0, equal
+            # to the multi-level variables:
+            "2m_temperature": 1.0,
+            # New single-level variables, which we don't weight too highly
+            # to avoid hurting performance on other variables.
+            "10m_u_component_of_wind": 0.1,
+            "10m_v_component_of_wind": 0.1,
+            "mean_sea_level_pressure": 0.1,
+            "total_precipitation_6hr": 0.1,
+        })
+    return loss, predictions  # pytype: disable=bad-return-type  # jax-ndarray
+
+  def loss(  # pytype: disable=signature-mismatch  # jax-ndarray
+      self,
+      inputs: xarray.Dataset,
+      targets: xarray.Dataset,
+      forcings: xarray.Dataset,
+      ) -> predictor_base.LossAndDiagnostics:
+    loss, _ = self.loss_and_predictions(inputs, targets, forcings)
+    return loss  # pytype: disable=bad-return-type  # jax-ndarray
+
+  def _maybe_init(self, sample_inputs: xarray.Dataset):
+    """Inits everything that has a dependency on the input coordinates."""
+    if not self._initialized:
+      self._init_mesh_properties()
+      self._init_grid_properties(
+          grid_lat=sample_inputs.lat, grid_lon=sample_inputs.lon)
+      self._grid2mesh_graph_structure = self._init_grid2mesh_graph()
+      self._mesh_graph_structure = self._init_mesh_graph()
+      self._mesh2grid_graph_structure = self._init_mesh2grid_graph()
+
+      self._initialized = True
+
+  def _init_mesh_properties(self):
+    """Inits static properties that have to do with mesh nodes."""
+    self._num_mesh_nodes = self._finest_mesh.vertices.shape[0]
+    mesh_phi, mesh_theta = model_utils.cartesian_to_spherical(
+        self._finest_mesh.vertices[:, 0],
+        self._finest_mesh.vertices[:, 1],
+        self._finest_mesh.vertices[:, 2])
+    (
+        mesh_nodes_lat,
+        mesh_nodes_lon,
+    ) = model_utils.spherical_to_lat_lon(
+        phi=mesh_phi, theta=mesh_theta)
+    # Convert to f32 to ensure the lat/lon features aren't in f64.
+    self._mesh_nodes_lat = mesh_nodes_lat.astype(np.float32)
+    self._mesh_nodes_lon = mesh_nodes_lon.astype(np.float32)
+
+  def _init_grid_properties(self, grid_lat: np.ndarray, grid_lon: np.ndarray):
+    """Inits static properties that have to do with grid nodes."""
+    self._grid_lat = grid_lat.astype(np.float32)
+    self._grid_lon = grid_lon.astype(np.float32)
+    # Initialized the counters.
+    self._num_grid_nodes = grid_lat.shape[0] * grid_lon.shape[0]
+
+    # Initialize lat and lon for the grid.
+    grid_nodes_lon, grid_nodes_lat = np.meshgrid(grid_lon, grid_lat)
+    self._grid_nodes_lon = grid_nodes_lon.reshape([-1]).astype(np.float32)
+    self._grid_nodes_lat = grid_nodes_lat.reshape([-1]).astype(np.float32)
+
+  def _init_grid2mesh_graph(self) -> typed_graph.TypedGraph:
+    """Build Grid2Mesh graph."""
+
+    # Create some edges according to distance between mesh and grid nodes.
+    assert self._grid_lat is not None and self._grid_lon is not None
+    (grid_indices, mesh_indices) = grid_mesh_connectivity.radius_query_indices(
+        grid_latitude=self._grid_lat,
+        grid_longitude=self._grid_lon,
+        mesh=self._finest_mesh,
+        radius=self._query_radius)
+
+    # Edges sending info from grid to mesh.
+    senders = grid_indices
+    receivers = mesh_indices
+
+    # Precompute structural node and edge features according to config options.
+    # Structural features are those that depend on the fixed values of the
+    # latitude and longitudes of the nodes.
+    (senders_node_features, receivers_node_features,
+     edge_features) = model_utils.get_bipartite_graph_spatial_features(
+         senders_node_lat=self._grid_nodes_lat,
+         senders_node_lon=self._grid_nodes_lon,
+         receivers_node_lat=self._mesh_nodes_lat,
+         receivers_node_lon=self._mesh_nodes_lon,
+         senders=senders,
+         receivers=receivers,
+         edge_normalization_factor=None,
+         **self._spatial_features_kwargs,
+     )
+
+    n_grid_node = np.array([self._num_grid_nodes])
+    n_mesh_node = np.array([self._num_mesh_nodes])
+    n_edge = np.array([mesh_indices.shape[0]])
+    grid_node_set = typed_graph.NodeSet(
+        n_node=n_grid_node, features=senders_node_features)
+    mesh_node_set = typed_graph.NodeSet(
+        n_node=n_mesh_node, features=receivers_node_features)
+    edge_set = typed_graph.EdgeSet(
+        n_edge=n_edge,
+        indices=typed_graph.EdgesIndices(senders=senders, receivers=receivers),
+        features=edge_features)
+    nodes = {"grid_nodes": grid_node_set, "mesh_nodes": mesh_node_set}
+    edges = {
+        typed_graph.EdgeSetKey("grid2mesh", ("grid_nodes", "mesh_nodes")):
+            edge_set
+    }
+    grid2mesh_graph = typed_graph.TypedGraph(
+        context=typed_graph.Context(n_graph=np.array([1]), features=()),
+        nodes=nodes,
+        edges=edges)
+    return grid2mesh_graph
+
+  def _init_mesh_graph(self) -> typed_graph.TypedGraph:
+    """Build Mesh graph."""
+    merged_mesh = icosahedral_mesh.merge_meshes(self._meshes)
+
+    # Work simply on the mesh edges.
+    senders, receivers = icosahedral_mesh.faces_to_edges(merged_mesh.faces)
+
+    # Precompute structural node and edge features according to config options.
+    # Structural features are those that depend on the fixed values of the
+    # latitude and longitudes of the nodes.
+    assert self._mesh_nodes_lat is not None and self._mesh_nodes_lon is not None
+    node_features, edge_features = model_utils.get_graph_spatial_features(
+        node_lat=self._mesh_nodes_lat,
+        node_lon=self._mesh_nodes_lon,
+        senders=senders,
+        receivers=receivers,
+        **self._spatial_features_kwargs,
+    )
+
+    n_mesh_node = np.array([self._num_mesh_nodes])
+    n_edge = np.array([senders.shape[0]])
+    assert n_mesh_node == len(node_features)
+    mesh_node_set = typed_graph.NodeSet(
+        n_node=n_mesh_node, features=node_features)
+    edge_set = typed_graph.EdgeSet(
+        n_edge=n_edge,
+        indices=typed_graph.EdgesIndices(senders=senders, receivers=receivers),
+        features=edge_features)
+    nodes = {"mesh_nodes": mesh_node_set}
+    edges = {
+        typed_graph.EdgeSetKey("mesh", ("mesh_nodes", "mesh_nodes")): edge_set
+    }
+    mesh_graph = typed_graph.TypedGraph(
+        context=typed_graph.Context(n_graph=np.array([1]), features=()),
+        nodes=nodes,
+        edges=edges)
+
+    return mesh_graph
+
+  def _init_mesh2grid_graph(self) -> typed_graph.TypedGraph:
+    """Build Mesh2Grid graph."""
+
+    # Create some edges according to how the grid nodes are contained by
+    # mesh triangles.
+    (grid_indices,
+     mesh_indices) = grid_mesh_connectivity.in_mesh_triangle_indices(
+         grid_latitude=self._grid_lat,
+         grid_longitude=self._grid_lon,
+         mesh=self._finest_mesh)
+
+    # Edges sending info from mesh to grid.
+    senders = mesh_indices
+    receivers = grid_indices
+
+    # Precompute structural node and edge features according to config options.
+    assert self._mesh_nodes_lat is not None and self._mesh_nodes_lon is not None
+    (senders_node_features, receivers_node_features,
+     edge_features) = model_utils.get_bipartite_graph_spatial_features(
+         senders_node_lat=self._mesh_nodes_lat,
+         senders_node_lon=self._mesh_nodes_lon,
+         receivers_node_lat=self._grid_nodes_lat,
+         receivers_node_lon=self._grid_nodes_lon,
+         senders=senders,
+         receivers=receivers,
+         edge_normalization_factor=self._mesh2grid_edge_normalization_factor,
+         **self._spatial_features_kwargs,
+     )
+
+    n_grid_node = np.array([self._num_grid_nodes])
+    n_mesh_node = np.array([self._num_mesh_nodes])
+    n_edge = np.array([senders.shape[0]])
+    grid_node_set = typed_graph.NodeSet(
+        n_node=n_grid_node, features=receivers_node_features)
+    mesh_node_set = typed_graph.NodeSet(
+        n_node=n_mesh_node, features=senders_node_features)
+    edge_set = typed_graph.EdgeSet(
+        n_edge=n_edge,
+        indices=typed_graph.EdgesIndices(senders=senders, receivers=receivers),
+        features=edge_features)
+    nodes = {"grid_nodes": grid_node_set, "mesh_nodes": mesh_node_set}
+    edges = {
+        typed_graph.EdgeSetKey("mesh2grid", ("mesh_nodes", "grid_nodes")):
+            edge_set
+    }
+    mesh2grid_graph = typed_graph.TypedGraph(
+        context=typed_graph.Context(n_graph=np.array([1]), features=()),
+        nodes=nodes,
+        edges=edges)
+    return mesh2grid_graph
+
+  def _run_grid2mesh_gnn(self, grid_node_features: chex.Array,
+                         ) -> tuple[chex.Array, chex.Array]:
+    """Runs the grid2mesh_gnn, extracting latent mesh and grid nodes."""
+
+    # Concatenate node structural features with input features.
+    batch_size = grid_node_features.shape[1]
+
+    grid2mesh_graph = self._grid2mesh_graph_structure
+    assert grid2mesh_graph is not None
+    grid_nodes = grid2mesh_graph.nodes["grid_nodes"]
+    mesh_nodes = grid2mesh_graph.nodes["mesh_nodes"]
+    new_grid_nodes = grid_nodes._replace(
+        features=jnp.concatenate([
+            grid_node_features,
+            _add_batch_second_axis(
+                grid_nodes.features.astype(grid_node_features.dtype),
+                batch_size)
+        ],
+                                 axis=-1))
+
+    # To make sure capacity of the embedded is identical for the grid nodes and
+    # the mesh nodes, we also append some dummy zero input features for the
+    # mesh nodes.
+    dummy_mesh_node_features = jnp.zeros(
+        (self._num_mesh_nodes,) + grid_node_features.shape[1:],
+        dtype=grid_node_features.dtype)
+    new_mesh_nodes = mesh_nodes._replace(
+        features=jnp.concatenate([
+            dummy_mesh_node_features,
+            _add_batch_second_axis(
+                mesh_nodes.features.astype(dummy_mesh_node_features.dtype),
+                batch_size)
+        ],
+                                 axis=-1))
+
+    # Broadcast edge structural features to the required batch size.
+    grid2mesh_edges_key = grid2mesh_graph.edge_key_by_name("grid2mesh")
+    edges = grid2mesh_graph.edges[grid2mesh_edges_key]
+
+    new_edges = edges._replace(
+        features=_add_batch_second_axis(
+            edges.features.astype(dummy_mesh_node_features.dtype), batch_size))
+
+    input_graph = self._grid2mesh_graph_structure._replace(
+        edges={grid2mesh_edges_key: new_edges},
+        nodes={
+            "grid_nodes": new_grid_nodes,
+            "mesh_nodes": new_mesh_nodes
+        })
+
+    # Run the GNN.
+    grid2mesh_out = self._grid2mesh_gnn(input_graph)
+    latent_mesh_nodes = grid2mesh_out.nodes["mesh_nodes"].features
+    latent_grid_nodes = grid2mesh_out.nodes["grid_nodes"].features
+    return latent_mesh_nodes, latent_grid_nodes
+
+  def _run_mesh_gnn(self, latent_mesh_nodes: chex.Array) -> chex.Array:
+    """Runs the mesh_gnn, extracting updated latent mesh nodes."""
+
+    # Add the structural edge features of this graph. Note we don't need
+    # to add the structural node features, because these are already part of
+    # the latent state, via the original Grid2Mesh gnn, however, we need
+    # the edge ones, because it is the first time we are seeing this particular
+    # set of edges.
+    batch_size = latent_mesh_nodes.shape[1]
+
+    mesh_graph = self._mesh_graph_structure
+    assert mesh_graph is not None
+    mesh_edges_key = mesh_graph.edge_key_by_name("mesh")
+    edges = mesh_graph.edges[mesh_edges_key]
+
+    # We are assuming here that the mesh gnn uses a single set of edge keys
+    # named "mesh" for the edges and that it uses a single set of nodes named
+    # "mesh_nodes"
+    msg = ("The setup currently requires to only have one kind of edge in the"
+           " mesh GNN.")
+    assert len(mesh_graph.edges) == 1, msg
+
+    new_edges = edges._replace(
+        features=_add_batch_second_axis(
+            edges.features.astype(latent_mesh_nodes.dtype), batch_size))
+
+    nodes = mesh_graph.nodes["mesh_nodes"]
+    nodes = nodes._replace(features=latent_mesh_nodes)
+
+    input_graph = mesh_graph._replace(
+        edges={mesh_edges_key: new_edges}, nodes={"mesh_nodes": nodes})
+
+    # Run the GNN.
+    return self._mesh_gnn(input_graph).nodes["mesh_nodes"].features
+
+  def _run_mesh2grid_gnn(self,
+                         updated_latent_mesh_nodes: chex.Array,
+                         latent_grid_nodes: chex.Array,
+                         ) -> chex.Array:
+    """Runs the mesh2grid_gnn, extracting the output grid nodes."""
+
+    # Add the structural edge features of this graph. Note we don't need
+    # to add the structural node features, because these are already part of
+    # the latent state, via the original Grid2Mesh gnn, however, we need
+    # the edge ones, because it is the first time we are seeing this particular
+    # set of edges.
+    batch_size = updated_latent_mesh_nodes.shape[1]
+
+    mesh2grid_graph = self._mesh2grid_graph_structure
+    assert mesh2grid_graph is not None
+    mesh_nodes = mesh2grid_graph.nodes["mesh_nodes"]
+    grid_nodes = mesh2grid_graph.nodes["grid_nodes"]
+    new_mesh_nodes = mesh_nodes._replace(features=updated_latent_mesh_nodes)
+    new_grid_nodes = grid_nodes._replace(features=latent_grid_nodes)
+    mesh2grid_key = mesh2grid_graph.edge_key_by_name("mesh2grid")
+    edges = mesh2grid_graph.edges[mesh2grid_key]
+
+    new_edges = edges._replace(
+        features=_add_batch_second_axis(
+            edges.features.astype(latent_grid_nodes.dtype), batch_size))
+
+    input_graph = mesh2grid_graph._replace(
+        edges={mesh2grid_key: new_edges},
+        nodes={
+            "mesh_nodes": new_mesh_nodes,
+            "grid_nodes": new_grid_nodes
+        })
+
+    # Run the GNN.
+    output_graph = self._mesh2grid_gnn(input_graph)
+    output_grid_nodes = output_graph.nodes["grid_nodes"].features
+
+    return output_grid_nodes
+
+  def _inputs_to_grid_node_features(
+      self,
+      inputs: xarray.Dataset,
+      forcings: xarray.Dataset,
+      ) -> chex.Array:
+    """xarrays -> [num_grid_nodes, batch, num_channels]."""
+
+    # xarray `Dataset` (batch, time, lat, lon, level, multiple vars)
+    # to xarray `DataArray` (batch, lat, lon, channels)
+    stacked_inputs = model_utils.dataset_to_stacked(inputs)
+    stacked_forcings = model_utils.dataset_to_stacked(forcings)
+    stacked_inputs = xarray.concat(
+        [stacked_inputs, stacked_forcings], dim="channels")
+
+    # xarray `DataArray` (batch, lat, lon, channels)
+    # to single numpy array with shape [lat_lon_node, batch, channels]
+    grid_xarray_lat_lon_leading = model_utils.lat_lon_to_leading_axes(
+        stacked_inputs)
+    return xarray_jax.unwrap(grid_xarray_lat_lon_leading.data).reshape(
+        (-1,) + grid_xarray_lat_lon_leading.data.shape[2:])
+
+  def _grid_node_outputs_to_prediction(
+      self,
+      grid_node_outputs: chex.Array,
+      targets_template: xarray.Dataset,
+      ) -> xarray.Dataset:
+    """[num_grid_nodes, batch, num_outputs] -> xarray."""
+
+    # numpy array with shape [lat_lon_node, batch, channels]
+    # to xarray `DataArray` (batch, lat, lon, channels)
+    assert self._grid_lat is not None and self._grid_lon is not None
+    grid_shape = (self._grid_lat.shape[0], self._grid_lon.shape[0])
+    grid_outputs_lat_lon_leading = grid_node_outputs.reshape(
+        grid_shape + grid_node_outputs.shape[1:])
+    dims = ("lat", "lon", "batch", "channels")
+    grid_xarray_lat_lon_leading = xarray_jax.DataArray(
+        data=grid_outputs_lat_lon_leading,
+        dims=dims)
+    grid_xarray = model_utils.restore_leading_axes(grid_xarray_lat_lon_leading)
+
+    # xarray `DataArray` (batch, lat, lon, channels)
+    # to xarray `Dataset` (batch, one time step, lat, lon, level, multiple vars)
+    return model_utils.stacked_to_dataset(
+        grid_xarray.variable, targets_template)
+
+
+def _add_batch_second_axis(data, batch_size):
+  # data [leading_dim, trailing_dim]
+  assert data.ndim == 2
+  ones = jnp.ones([batch_size, 1], dtype=data.dtype)
+  return data[:, None] * ones  # [leading_dim, batch, trailing_dim]
+
+
+def _get_max_edge_distance(mesh):
+  senders, receivers = icosahedral_mesh.faces_to_edges(mesh.faces)
+  edge_distances = np.linalg.norm(
+      mesh.vertices[senders] - mesh.vertices[receivers], axis=-1)
+  return edge_distances.max()

model/graphcast/grid_mesh_connectivity.py

@@ -0,0 +1,133 @@
+# Copyright 2023 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tools for converting from regular grids on a sphere, to triangular meshes."""
+
+from . import icosahedral_mesh
+import numpy as np
+import scipy
+import trimesh
+
+
+def _grid_lat_lon_to_coordinates(
+    grid_latitude: np.ndarray, grid_longitude: np.ndarray) -> np.ndarray:
+  """Lat [num_lat] lon [num_lon] to 3d coordinates [num_lat, num_lon, 3]."""
+  # Convert to spherical coordinates phi and theta defined in the grid.
+  # Each [num_latitude_points, num_longitude_points]
+  phi_grid, theta_grid = np.meshgrid(
+      np.deg2rad(grid_longitude),
+      np.deg2rad(90 - grid_latitude))
+
+  # [num_latitude_points, num_longitude_points, 3]
+  # Note this assumes unit radius, since for now we model the earth as a
+  # sphere of unit radius, and keep any vertical dimension as a regular grid.
+  return np.stack(
+      [np.cos(phi_grid)*np.sin(theta_grid),
+       np.sin(phi_grid)*np.sin(theta_grid),
+       np.cos(theta_grid)], axis=-1)
+
+
+def radius_query_indices(
+    *,
+    grid_latitude: np.ndarray,
+    grid_longitude: np.ndarray,
+    mesh: icosahedral_mesh.TriangularMesh,
+    radius: float) -> tuple[np.ndarray, np.ndarray]:
+  """Returns mesh-grid edge indices for radius query.
+
+  Args:
+    grid_latitude: Latitude values for the grid [num_lat_points]
+    grid_longitude: Longitude values for the grid [num_lon_points]
+    mesh: Mesh object.
+    radius: Radius of connectivity in R3. for a sphere of unit radius.
+
+  Returns:
+    tuple with `grid_indices` and `mesh_indices` indicating edges between the
+    grid and the mesh such that the distances in a straight line (not geodesic)
+    are smaller than or equal to `radius`.
+    * grid_indices: Indices of shape [num_edges], that index into a
+      [num_lat_points, num_lon_points] grid, after flattening the leading axes.
+    * mesh_indices: Indices of shape [num_edges], that index into mesh.vertices.
+  """
+
+  # [num_grid_points=num_lat_points * num_lon_points, 3]
+  grid_positions = _grid_lat_lon_to_coordinates(
+      grid_latitude, grid_longitude).reshape([-1, 3])
+
+  # [num_mesh_points, 3]
+  mesh_positions = mesh.vertices
+  kd_tree = scipy.spatial.cKDTree(mesh_positions)
+
+  # [num_grid_points, num_mesh_points_per_grid_point]
+  # Note `num_mesh_points_per_grid_point` is not constant, so this is a list
+  # of arrays, rather than a 2d array.
+  query_indices = kd_tree.query_ball_point(x=grid_positions, r=radius)
+
+  grid_edge_indices = []
+  mesh_edge_indices = []
+  for grid_index, mesh_neighbors in enumerate(query_indices):
+    grid_edge_indices.append(np.repeat(grid_index, len(mesh_neighbors)))
+    mesh_edge_indices.append(mesh_neighbors)
+
+  # [num_edges]
+  grid_edge_indices = np.concatenate(grid_edge_indices, axis=0).astype(int)
+  mesh_edge_indices = np.concatenate(mesh_edge_indices, axis=0).astype(int)
+
+  return grid_edge_indices, mesh_edge_indices
+
+
+def in_mesh_triangle_indices(
+    *,
+    grid_latitude: np.ndarray,
+    grid_longitude: np.ndarray,
+    mesh: icosahedral_mesh.TriangularMesh) -> tuple[np.ndarray, np.ndarray]:
+  """Returns mesh-grid edge indices for grid points contained in mesh triangles.
+
+  Args:
+    grid_latitude: Latitude values for the grid [num_lat_points]
+    grid_longitude: Longitude values for the grid [num_lon_points]
+    mesh: Mesh object.
+
+  Returns:
+    tuple with `grid_indices` and `mesh_indices` indicating edges between the
+    grid and the mesh vertices of the triangle that contain each grid point.
+    The number of edges is always num_lat_points * num_lon_points * 3
+    * grid_indices: Indices of shape [num_edges], that index into a
+      [num_lat_points, num_lon_points] grid, after flattening the leading axes.
+    * mesh_indices: Indices of shape [num_edges], that index into mesh.vertices.
+  """
+
+  # [num_grid_points=num_lat_points * num_lon_points, 3]
+  grid_positions = _grid_lat_lon_to_coordinates(
+      grid_latitude, grid_longitude).reshape([-1, 3])
+
+  mesh_trimesh = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces)
+
+  # [num_grid_points] with mesh face indices for each grid point.
+  _, _, query_face_indices = trimesh.proximity.closest_point(
+      mesh_trimesh, grid_positions)
+
+  # [num_grid_points, 3] with mesh node indices for each grid point.
+  mesh_edge_indices = mesh.faces[query_face_indices]
+
+  # [num_grid_points, 3] with grid node indices, where every row simply contains
+  # the row (grid_point) index.
+  grid_indices = np.arange(grid_positions.shape[0])
+  grid_edge_indices = np.tile(grid_indices.reshape([-1, 1]), [1, 3])
+
+  # Flatten to get a regular list.
+  # [num_edges=num_grid_points*3]
+  mesh_edge_indices = mesh_edge_indices.reshape([-1])
+  grid_edge_indices = grid_edge_indices.reshape([-1])
+
+  return grid_edge_indices, mesh_edge_indices

model/graphcast/icosahedral_mesh.py

@@ -0,0 +1,285 @@
+# Copyright 2023 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Utils for creating icosahedral meshes."""
+
+import itertools
+from typing import List, NamedTuple, Sequence, Tuple
+
+import numpy as np
+from scipy.spatial import transform
+
+
+class TriangularMesh(NamedTuple):
+  """Data structure for triangular meshes.
+
+  Attributes:
+    vertices: spatial positions of the vertices of the mesh of shape
+        [num_vertices, num_dims].
+    faces: triangular faces of the mesh of shape [num_faces, 3]. Contains
+        integer indices into `vertices`.
+
+  """
+  vertices: np.ndarray
+  faces: np.ndarray
+
+
+def merge_meshes(
+    mesh_list: Sequence[TriangularMesh]) -> TriangularMesh:
+  """Merges all meshes into one. Assumes the last mesh is the finest.
+
+  Args:
+     mesh_list: Sequence of meshes, from coarse to fine refinement levels. The
+       vertices and faces may contain those from preceding, coarser levels.
+
+  Returns:
+     `TriangularMesh` for which the vertices correspond to the highest
+     resolution mesh in the hierarchy, and the faces are the join set of the
+     faces at all levels of the hierarchy.
+  """
+  for mesh_i, mesh_ip1 in itertools.pairwise(mesh_list):
+    num_nodes_mesh_i = mesh_i.vertices.shape[0]
+    assert np.allclose(mesh_i.vertices, mesh_ip1.vertices[:num_nodes_mesh_i])
+
+  return TriangularMesh(
+      vertices=mesh_list[-1].vertices,
+      faces=np.concatenate([mesh.faces for mesh in mesh_list], axis=0))
+
+
+def get_hierarchy_of_triangular_meshes_for_sphere(
+    splits: int) -> List[TriangularMesh]:
+  """Returns a sequence of meshes, each with triangularization sphere.
+
+  Starting with a regular icosahedron (12 vertices, 20 faces, 30 edges) with
+  circumscribed unit sphere. Then, each triangular face is iteratively
+  subdivided into 4 triangular faces `splits` times. The new vertices are then
+  projected back onto the unit sphere. All resulting meshes are returned in a
+  list, from lowest to highest resolution.
+
+  The vertices in each face are specified in counter-clockwise order as
+  observed from the outside the icosahedron.
+
+  Args:
+     splits: How many times to split each triangle.
+  Returns:
+     Sequence of `TriangularMesh`s of length `splits + 1` each with:
+
+       vertices: [num_vertices, 3] vertex positions in 3D, all with unit norm.
+       faces: [num_faces, 3] with triangular faces joining sets of 3 vertices.
+           Each row contains three indices into the vertices array, indicating
+           the vertices adjacent to the face. Always with positive orientation
+           (counterclock-wise when looking from the outside).
+  """
+  current_mesh = get_icosahedron()
+  output_meshes = [current_mesh]
+  for _ in range(splits):
+    current_mesh = _two_split_unit_sphere_triangle_faces(current_mesh)
+    output_meshes.append(current_mesh)
+  return output_meshes
+
+
+def get_icosahedron() -> TriangularMesh:
+  """Returns a regular icosahedral mesh with circumscribed unit sphere.
+
+  See https://en.wikipedia.org/wiki/Regular_icosahedron#Cartesian_coordinates
+  for details on the construction of the regular icosahedron.
+
+  The vertices in each face are specified in counter-clockwise order as observed
+  from the outside of the icosahedron.
+
+  Returns:
+     TriangularMesh with:
+
+     vertices: [num_vertices=12, 3] vertex positions in 3D, all with unit norm.
+     faces: [num_faces=20, 3] with triangular faces joining sets of 3 vertices.
+         Each row contains three indices into the vertices array, indicating
+         the vertices adjacent to the face. Always with positive orientation (
+         counterclock-wise when looking from the outside).
+
+  """
+  phi = (1 + np.sqrt(5)) / 2
+  vertices = []
+  for c1 in [1., -1.]:
+    for c2 in [phi, -phi]:
+      vertices.append((c1, c2, 0.))
+      vertices.append((0., c1, c2))
+      vertices.append((c2, 0., c1))
+
+  vertices = np.array(vertices, dtype=np.float32)
+  vertices /= np.linalg.norm([1., phi])
+
+  # I did this manually, checking the orientation one by one.
+  faces = [(0, 1, 2),
+           (0, 6, 1),
+           (8, 0, 2),
+           (8, 4, 0),
+           (3, 8, 2),
+           (3, 2, 7),
+           (7, 2, 1),
+           (0, 4, 6),
+           (4, 11, 6),
+           (6, 11, 5),
+           (1, 5, 7),
+           (4, 10, 11),
+           (4, 8, 10),
+           (10, 8, 3),
+           (10, 3, 9),
+           (11, 10, 9),
+           (11, 9, 5),
+           (5, 9, 7),
+           (9, 3, 7),
+           (1, 6, 5),
+           ]
+
+  # By default the top is an aris parallel to the Y axis.
+  # Need to rotate around the y axis by half the supplementary to the
+  # angle between faces divided by two to get the desired orientation.
+  #                          /O\  (top arist)
+  #                     /          \                           Z
+  # (adjacent face)/                    \  (adjacent face)     ^
+  #           /     angle_between_faces      \                 |
+  #      /                                        \            |
+  #  /                                                 \      YO-----> X
+  # This results in:
+  #  (adjacent faceis now top plane)
+  #  ----------------------O\  (top arist)
+  #                           \
+  #                             \
+  #                               \     (adjacent face)
+  #                                 \
+  #                                   \
+  #                                     \
+
+  angle_between_faces = 2 * np.arcsin(phi / np.sqrt(3))
+  rotation_angle = (np.pi - angle_between_faces) / 2
+  rotation = transform.Rotation.from_euler(seq="y", angles=rotation_angle)
+  rotation_matrix = rotation.as_matrix()
+  vertices = np.dot(vertices, rotation_matrix)
+
+  return TriangularMesh(vertices=vertices.astype(np.float32),
+                        faces=np.array(faces, dtype=np.int32))
+
+
+def _two_split_unit_sphere_triangle_faces(
+    triangular_mesh: TriangularMesh) -> TriangularMesh:
+  """Splits each triangular face into 4 triangles keeping the orientation."""
+
+  # Every time we split a triangle into 4 we will be adding 3 extra vertices,
+  # located at the edge centres.
+  # This class handles the positioning of the new vertices, and avoids creating
+  # duplicates.
+  new_vertices_builder = _ChildVerticesBuilder(triangular_mesh.vertices)
+
+  new_faces = []
+  for ind1, ind2, ind3 in triangular_mesh.faces:
+    # Transform each triangular face into 4 triangles,
+    # preserving the orientation.
+    #                    ind3
+    #                   /    \
+    #                /          \
+    #              /      #3       \
+    #            /                  \
+    #         ind31 -------------- ind23
+    #         /   \                /   \
+    #       /       \     #4     /      \
+    #     /    #1     \        /    #2    \
+    #   /               \    /              \
+    # ind1 ------------ ind12 ------------ ind2
+    ind12 = new_vertices_builder.get_new_child_vertex_index((ind1, ind2))
+    ind23 = new_vertices_builder.get_new_child_vertex_index((ind2, ind3))
+    ind31 = new_vertices_builder.get_new_child_vertex_index((ind3, ind1))
+    # Note how each of the 4 triangular new faces specifies the order of the
+    # vertices to preserve the orientation of the original face. As the input
+    # face should always be counter-clockwise as specified in the diagram,
+    # this means child faces should also be counter-clockwise.
+    new_faces.extend([[ind1, ind12, ind31],  # 1
+                      [ind12, ind2, ind23],  # 2
+                      [ind31, ind23, ind3],  # 3
+                      [ind12, ind23, ind31],  # 4
+                      ])
+  return TriangularMesh(vertices=new_vertices_builder.get_all_vertices(),
+                        faces=np.array(new_faces, dtype=np.int32))
+
+
+class _ChildVerticesBuilder(object):
+  """Bookkeeping of new child vertices added to an existing set of vertices."""
+
+  def __init__(self, parent_vertices):
+
+    # Because the same new vertex will be required when splitting adjacent
+    # triangles (which share an edge) we keep them in a hash table indexed by
+    # sorted indices of the vertices adjacent to the edge, to avoid creating
+    # duplicated child vertices.
+    self._child_vertices_index_mapping = {}
+    self._parent_vertices = parent_vertices
+    # We start with all previous vertices.
+    self._all_vertices_list = list(parent_vertices)
+
+  def _get_child_vertex_key(self, parent_vertex_indices):
+    return tuple(sorted(parent_vertex_indices))
+
+  def _create_child_vertex(self, parent_vertex_indices):
+    """Creates a new vertex."""
+    # Position for new vertex is the middle point, between the parent points,
+    # projected to unit sphere.
+    child_vertex_position = self._parent_vertices[
+        list(parent_vertex_indices)].mean(0)
+    child_vertex_position /= np.linalg.norm(child_vertex_position)
+
+    # Add the vertex to the output list. The index for this new vertex will
+    # match the length of the list before adding it.
+    child_vertex_key = self._get_child_vertex_key(parent_vertex_indices)
+    self._child_vertices_index_mapping[child_vertex_key] = len(
+        self._all_vertices_list)
+    self._all_vertices_list.append(child_vertex_position)
+
+  def get_new_child_vertex_index(self, parent_vertex_indices):
+    """Returns index for a child vertex, creating it if necessary."""
+    # Get the key to see if we already have a new vertex in the middle.
+    child_vertex_key = self._get_child_vertex_key(parent_vertex_indices)
+    if child_vertex_key not in self._child_vertices_index_mapping:
+      self._create_child_vertex(parent_vertex_indices)
+    return self._child_vertices_index_mapping[child_vertex_key]
+
+  def get_all_vertices(self):
+    """Returns an array with old vertices."""
+    return np.array(self._all_vertices_list)
+
+
+def faces_to_edges(faces: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+  """Transforms polygonal faces to sender and receiver indices.
+
+  It does so by transforming every face into N_i edges. Such if the triangular
+  face has indices [0, 1, 2], three edges are added 0->1, 1->2, and 2->0.
+
+  If all faces have consistent orientation, and the surface represented by the
+  faces is closed, then every edge in a polygon with a certain orientation
+  is also part of another polygon with the opposite orientation. In this
+  situation, the edges returned by the method are always bidirectional.
+
+  Args:
+    faces: Integer array of shape [num_faces, 3]. Contains node indices
+        adjacent to each face.
+  Returns:
+    Tuple with sender/receiver indices, each of shape [num_edges=num_faces*3].
+
+  """
+  assert faces.ndim == 2
+  assert faces.shape[-1] == 3
+  senders = np.concatenate([faces[:, 0], faces[:, 1], faces[:, 2]])
+  receivers = np.concatenate([faces[:, 1], faces[:, 2], faces[:, 0]])
+  return senders, receivers
+
+
+def get_last_triangular_mesh_for_sphere(splits: int) -> TriangularMesh:
+  return get_hierarchy_of_triangular_meshes_for_sphere(splits=splits)[-1]

model/graphcast/losses.py

@@ -0,0 +1,179 @@
+# Copyright 2023 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Loss functions (and terms for use in loss functions) used for weather."""
+
+from typing import Mapping
+
+from . import xarray_tree
+import numpy as np
+from typing_extensions import Protocol
+import xarray
+
+
+LossAndDiagnostics = tuple[xarray.DataArray, xarray.Dataset]
+
+
+class LossFunction(Protocol):
+  """A loss function.
+
+  This is a protocol so it's fine to use a plain function which 'quacks like'
+  this. This is just to document the interface.
+  """
+
+  def __call__(self,
+               predictions: xarray.Dataset,
+               targets: xarray.Dataset,
+               **optional_kwargs) -> LossAndDiagnostics:
+    """Computes a loss function.
+
+    Args:
+      predictions: Dataset of predictions.
+      targets: Dataset of targets.
+      **optional_kwargs: Implementations may support extra optional kwargs.
+
+    Returns:
+      loss: A DataArray with dimensions ('batch',) containing losses for each
+        element of the batch. These will be averaged to give the final
+        loss, locally and across replicas.
+      diagnostics: Mapping of additional quantities to log by name alongside the
+        loss. These will will typically correspond to terms in the loss. They
+        should also have dimensions ('batch',) and will be averaged over the
+        batch before logging.
+    """
+
+
+def weighted_mse_per_level(
+    predictions: xarray.Dataset,
+    targets: xarray.Dataset,
+    per_variable_weights: Mapping[str, float],
+) -> LossAndDiagnostics:
+  """Latitude- and pressure-level-weighted MSE loss."""
+  def loss(prediction, target):
+    loss = (prediction - target)**2
+    loss *= normalized_latitude_weights(target).astype(loss.dtype)
+    if 'level' in target.dims:
+      loss *= normalized_level_weights(target).astype(loss.dtype)
+    return _mean_preserving_batch(loss)
+
+  losses = xarray_tree.map_structure(loss, predictions, targets)
+  return sum_per_variable_losses(losses, per_variable_weights)
+
+
+def _mean_preserving_batch(x: xarray.DataArray) -> xarray.DataArray:
+  return x.mean([d for d in x.dims if d != 'batch'], skipna=False)
+
+
+def sum_per_variable_losses(
+    per_variable_losses: Mapping[str, xarray.DataArray],
+    weights: Mapping[str, float],
+) -> LossAndDiagnostics:
+  """Weighted sum of per-variable losses."""
+  if not set(weights.keys()).issubset(set(per_variable_losses.keys())):
+    raise ValueError(
+        'Passing a weight that does not correspond to any variable '
+        f'{set(weights.keys())-set(per_variable_losses.keys())}')
+
+  weighted_per_variable_losses = {
+      name: loss * weights.get(name, 1)
+      for name, loss in per_variable_losses.items()
+  }
+  total = xarray.concat(
+      weighted_per_variable_losses.values(), dim='variable', join='exact').sum(
+          'variable', skipna=False)
+  return total, per_variable_losses  # pytype: disable=bad-return-type
+
+
+def normalized_level_weights(data: xarray.DataArray) -> xarray.DataArray:
+  """Weights proportional to pressure at each level."""
+  level = data.coords['level']
+  return level / level.mean(skipna=False)
+
+
+def normalized_latitude_weights(data: xarray.DataArray) -> xarray.DataArray:
+  """Weights based on latitude, roughly proportional to grid cell area.
+
+  This method supports two use cases only (both for equispaced values):
+  * Latitude values such that the closest value to the pole is at latitude
+    (90 - d_lat/2), where d_lat is the difference between contiguous latitudes.
+    For example: [-89, -87, -85, ..., 85, 87, 89]) (d_lat = 2)
+    In this case each point with `lat` value represents a sphere slice between
+    `lat - d_lat/2` and `lat + d_lat/2`, and the area of this slice would be
+    proportional to:
+    `sin(lat + d_lat/2) - sin(lat - d_lat/2) = 2 * sin(d_lat/2) * cos(lat)`, and
+    we can simply omit the term `2 * sin(d_lat/2)` which is just a constant
+    that cancels during normalization.
+  * Latitude values that fall exactly at the poles.
+    For example: [-90, -88, -86, ..., 86, 88, 90]) (d_lat = 2)
+    In this case each point with `lat` value also represents
+    a sphere slice between `lat - d_lat/2` and `lat + d_lat/2`,
+    except for the points at the poles, that represent a slice between
+    `90 - d_lat/2` and `90` or, `-90` and  `-90 + d_lat/2`.
+    The areas of the first type of point are still proportional to:
+    * sin(lat + d_lat/2) - sin(lat - d_lat/2) = 2 * sin(d_lat/2) * cos(lat)
+    but for the points at the poles now is:
+    * sin(90) - sin(90 - d_lat/2) = 2 * sin(d_lat/4) ^ 2
+    and we will be using these weights, depending on whether we are looking at
+    pole cells, or non-pole cells (omitting the common factor of 2 which will be
+    absorbed by the normalization).
+
+    It can be shown via a limit, or simple geometry, that in the small angles
+    regime, the proportion of area per pole-point is equal to 1/8th
+    the proportion of area covered by each of the nearest non-pole point, and we
+    test for this in the test.
+
+  Args:
+    data: `DataArray` with latitude coordinates.
+  Returns:
+    Unit mean latitude weights.
+  """
+  latitude = data.coords['lat']
+
+  if np.any(np.isclose(np.abs(latitude), 90.)):
+    weights = _weight_for_latitude_vector_with_poles(latitude)
+  else:
+    weights = _weight_for_latitude_vector_without_poles(latitude)
+
+  return weights / weights.mean(skipna=False)
+
+
+def _weight_for_latitude_vector_without_poles(latitude):
+  """Weights for uniform latitudes of the form [+-90-+d/2, ..., -+90+-d/2]."""
+  delta_latitude = np.abs(_check_uniform_spacing_and_get_delta(latitude))
+  if (not np.isclose(np.max(latitude), 90 - delta_latitude/2) or
+      not np.isclose(np.min(latitude), -90 + delta_latitude/2)):
+    raise ValueError(
+        f'Latitude vector {latitude} does not start/end at '
+        '+- (90 - delta_latitude/2) degrees.')
+  return np.cos(np.deg2rad(latitude))
+
+
+def _weight_for_latitude_vector_with_poles(latitude):
+  """Weights for uniform latitudes of the form [+- 90, ..., -+90]."""
+  delta_latitude = np.abs(_check_uniform_spacing_and_get_delta(latitude))
+  if (not np.isclose(np.max(latitude), 90.) or
+      not np.isclose(np.min(latitude), -90.)):
+    raise ValueError(
+        f'Latitude vector {latitude} does not start/end at +- 90 degrees.')
+  weights = np.cos(np.deg2rad(latitude)) * np.sin(np.deg2rad(delta_latitude/2))
+  # The two checks above enough to guarantee that latitudes are sorted, so
+  # the extremes are the poles
+  weights[[0, -1]] = np.sin(np.deg2rad(delta_latitude/4)) ** 2
+  return weights
+
+
+def _check_uniform_spacing_and_get_delta(vector):
+  diff = np.diff(vector)
+  if not np.all(np.isclose(diff[0], diff)):
+    raise ValueError(f'Vector {diff} is not uniformly spaced.')
+  return diff[0]

model/graphcast/mlp.py

@@ -0,0 +1,45 @@
+# Copyright 2024 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Constructors for MLPs."""
+
+import haiku as hk
+import jax
+import jax.numpy as jnp
+
+
+# TODO(aelkadi): Move the mlp factory here from `deep_typed_graph_net.py`.
+
+
+class LinearNormConditioning(hk.Module):
+  """Module for norm conditioning.
+
+  Conditions the normalization of "inputs" by applying a linear layer to the
+  "norm_conditioning" which produces the scale and variance which are applied to
+  each channel (across the last dim) of "inputs".
+  """
+
+  def __init__(self, name="norm_conditioning"):
+    super().__init__(name=name)
+
+  def __call__(self, inputs: jax.Array, norm_conditioning: jax.Array):
+
+    feature_size = inputs.shape[-1]
+    conditional_linear_layer = hk.Linear(
+        output_size=2 * feature_size,
+        w_init=hk.initializers.TruncatedNormal(stddev=1e-8),
+    )
+    conditional_scale_offset = conditional_linear_layer(norm_conditioning)
+    scale_minus_one, offset = jnp.split(conditional_scale_offset, 2, axis=-1)
+    scale = scale_minus_one + 1.
+    return inputs * scale + offset

model/graphcast/model_utils.py

@@ -0,0 +1,807 @@
+# Copyright 2023 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Utilities for building models."""
+
+from typing import Any, Mapping, Optional, Tuple
+
+import jax.numpy as jnp
+import numpy as np
+from scipy.spatial import transform
+import xarray
+
+NumpyInterface = Any
+TransformInterface = Any
+
+
+def get_graph_spatial_features(
+    *, node_lat: np.ndarray, node_lon: np.ndarray,
+    senders: np.ndarray, receivers: np.ndarray,
+    add_node_positions: bool,
+    add_node_latitude: bool,
+    add_node_longitude: bool,
+    add_relative_positions: bool,
+    edge_normalization_factor: Optional[float] = None,
+    relative_longitude_local_coordinates: bool,
+    relative_latitude_local_coordinates: bool,
+    sine_cosine_encoding: bool = False,
+    encoding_num_freqs: int = 10,
+    encoding_multiplicative_factor: float = 1.2,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+  """Computes spatial features for the nodes.
+
+  Args:
+    node_lat: Latitudes in the [-90, 90] interval of shape [num_nodes]
+    node_lon: Longitudes in the [0, 360] interval of shape [num_nodes]
+    senders: Sender indices of shape [num_edges]
+    receivers: Receiver indices of shape [num_edges]
+    add_node_positions: Add unit norm absolute positions.
+    add_node_latitude: Add a feature for latitude (cos(90 - lat))
+        Note even if this is set to False, the model may be able to infer the
+        longitude from relative features, unless
+        `relative_latitude_local_coordinates` is also True, or if there is any
+        bias on the relative edge sizes for different longitudes.
+    add_node_longitude: Add features for longitude (cos(lon), sin(lon)).
+        Note even if this is set to False, the model may be able to infer the
+        longitude from relative features, unless
+        `relative_longitude_local_coordinates` is also True, or if there is any
+        bias on the relative edge sizes for different longitudes.
+    add_relative_positions: Whether to relative positions in R3 to the edges.
+    edge_normalization_factor: Allows explicitly controlling edge normalization.
+        If None, defaults to max edge length. This supports using pre-trained
+        model weights with a different graph structure to what it was trained.
+    relative_longitude_local_coordinates: If True, relative positions are
+        computed in a local space where the receiver is at 0 longitude.
+    relative_latitude_local_coordinates: If True, relative positions are
+        computed in a local space where the receiver is at 0 latitude.
+    sine_cosine_encoding: If True, we will transform the node/edge features
+        with sine and cosine functions, similar to NERF.
+    encoding_num_freqs: frequency parameter
+    encoding_multiplicative_factor: used for calculating the frequency.
+
+  Returns:
+    Arrays of shape: [num_nodes, num_features] and [num_edges, num_features].
+    with node and edge features.
+
+  """
+
+  num_nodes = node_lat.shape[0]
+  num_edges = senders.shape[0]
+  dtype = node_lat.dtype
+  node_phi, node_theta = lat_lon_deg_to_spherical(node_lat, node_lon)
+
+  # Computing some node features.
+  node_features = []
+  if add_node_positions:
+    # Already in [-1, 1.] range.
+    node_features.extend(spherical_to_cartesian(node_phi, node_theta))
+
+  if add_node_latitude:
+    # Using the cos of theta.
+    # From 1. (north pole) to -1 (south pole).
+    node_features.append(np.cos(node_theta))
+
+  if add_node_longitude:
+    # Using the cos and sin, which is already normalized.
+    node_features.append(np.cos(node_phi))
+    node_features.append(np.sin(node_phi))
+
+  if not node_features:
+    node_features = np.zeros([num_nodes, 0], dtype=dtype)
+  else:
+    node_features = np.stack(node_features, axis=-1)
+
+  # Computing some edge features.
+  edge_features = []
+
+  if add_relative_positions:
+
+    relative_position = get_relative_position_in_receiver_local_coordinates(
+        node_phi=node_phi,
+        node_theta=node_theta,
+        senders=senders,
+        receivers=receivers,
+        latitude_local_coordinates=relative_latitude_local_coordinates,
+        longitude_local_coordinates=relative_longitude_local_coordinates
+        )
+
+    # Note this is L2 distance in 3d space, rather than geodesic distance.
+    relative_edge_distances = np.linalg.norm(
+        relative_position, axis=-1, keepdims=True)
+
+    if edge_normalization_factor is None:
+      # Normalize to the maximum edge distance. Note that we expect to always
+      # have an edge that goes in the opposite direction of any given edge
+      # so the distribution of relative positions should be symmetric around
+      # zero. So by scaling by the maximum length, we expect all relative
+      # positions to fall in the [-1., 1.] interval, and all relative distances
+      # to fall in the [0., 1.] interval.
+      edge_normalization_factor = relative_edge_distances.max()
+    edge_features.append(relative_edge_distances / edge_normalization_factor)
+    edge_features.append(relative_position / edge_normalization_factor)
+
+  if not edge_features:
+    edge_features = np.zeros([num_edges, 0], dtype=dtype)
+  else:
+    edge_features = np.concatenate(edge_features, axis=-1)
+
+  if sine_cosine_encoding:
+    def sine_cosine_transform(x: np.ndarray) -> np.ndarray:
+      freqs = encoding_multiplicative_factor**np.arange(encoding_num_freqs)
+      phases = freqs * x[..., None]
+      x_sin = np.sin(phases)
+      x_cos = np.cos(phases)
+      x_cat = np.concatenate([x_sin, x_cos], axis=-1)
+      return x_cat.reshape([x.shape[0], -1])
+
+    node_features = sine_cosine_transform(node_features)
+    edge_features = sine_cosine_transform(edge_features)
+
+  return node_features, edge_features
+
+
+def lat_lon_to_leading_axes(
+    grid_xarray: xarray.DataArray) -> xarray.DataArray:
+  """Reorders xarray so lat/lon axes come first."""
+  # leading + ["lat", "lon"] + trailing
+  # to
+  # ["lat", "lon"] + leading + trailing
+  return grid_xarray.transpose("lat", "lon", ...)
+
+
+def restore_leading_axes(grid_xarray: xarray.DataArray) -> xarray.DataArray:
+  """Reorders xarray so batch/time/level axes come first (if present)."""
+
+  # ["lat", "lon"] + [(batch,) (time,) (level,)] + trailing
+  # to
+  # [(batch,) (time,) (level,)] + ["lat", "lon"] + trailing
+
+  input_dims = list(grid_xarray.dims)
+  output_dims = list(input_dims)
+  for leading_key in ["level", "time", "batch"]:  # reverse order for insert
+    if leading_key in input_dims:
+      output_dims.remove(leading_key)
+      output_dims.insert(0, leading_key)
+  return grid_xarray.transpose(*output_dims)
+
+
+def lat_lon_deg_to_spherical(node_lat: np.ndarray,
+                             node_lon: np.ndarray,
+                             np_: NumpyInterface = np,
+                            ) -> Tuple[np.ndarray, np.ndarray]:
+  phi = np_.deg2rad(node_lon)
+  theta = np_.deg2rad(90 - node_lat)
+  return phi, theta
+
+
+def spherical_to_lat_lon(phi: np.ndarray,
+                         theta: np.ndarray,
+                         np_: NumpyInterface = np,
+                        ) -> Tuple[np.ndarray, np.ndarray]:
+  lon = np_.mod(np_.rad2deg(phi), 360)
+  lat = 90 - np_.rad2deg(theta)
+  return lat, lon
+
+
+def cartesian_to_spherical(x: np.ndarray,
+                           y: np.ndarray,
+                           z: np.ndarray,
+                           np_: NumpyInterface = np,
+                          ) -> Tuple[np.ndarray, np.ndarray]:
+  phi = np_.arctan2(y, x)
+  with np.errstate(invalid="ignore"):  # circumventing b/253179568
+    theta = np_.arccos(z)  # Assuming unit radius.
+  return phi, theta
+
+
+def spherical_to_cartesian(
+    phi: np.ndarray, theta: np.ndarray,
+    np_: NumpyInterface = np,
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+  # Assuming unit radius.
+  return (np_.cos(phi)*np_.sin(theta),
+          np_.sin(phi)*np_.sin(theta),
+          np_.cos(theta))
+
+
+def lat_lon_to_cartesian(
+    lat: np.ndarray, lon: np.ndarray,
+    np_: NumpyInterface = np,
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+  return spherical_to_cartesian(
+      *lat_lon_deg_to_spherical(lat, lon, np_=np_), np_=np_)
+
+
+def cartesian_to_lat_lon(
+    x: np.ndarray,
+    y: np.ndarray,
+    z: np.ndarray,
+    np_: NumpyInterface = np,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+  return spherical_to_lat_lon(
+      *cartesian_to_spherical(x, y, z, np_=np_), np_=np_)
+
+
+def get_relative_position_in_receiver_local_coordinates(
+    node_phi: np.ndarray,
+    node_theta: np.ndarray,
+    senders: np.ndarray,
+    receivers: np.ndarray,
+    latitude_local_coordinates: bool,
+    longitude_local_coordinates: bool,
+    np_: NumpyInterface = np,
+    transform_: TransformInterface = transform,
+    ) -> np.ndarray:
+  """Returns relative position features for the edges.
+
+  The relative positions will be computed in a rotated space for a local
+  coordinate system as defined by the receiver. The relative positions are
+  simply obtained by subtracting sender position minues receiver position in
+  that local coordinate system after the rotation in R^3.
+
+  Args:
+    node_phi: [num_nodes] with polar angles.
+    node_theta: [num_nodes] with azimuthal angles.
+    senders: [num_edges] with indices.
+    receivers: [num_edges] with indices.
+    latitude_local_coordinates: Whether to rotate edges such that in the
+        positions are computed such that the receiver is always at latitude 0.
+    longitude_local_coordinates: Whether to rotate edges such that in the
+        positions are computed such that the receiver is always at longitude 0.
+    np_: Numpy library interface.
+    transform_: scipy.transform library interface.
+
+  Returns:
+    Array of relative positions in R3 [num_edges, 3]
+  """
+
+  node_pos = np_.stack(
+      spherical_to_cartesian(node_phi, node_theta, np_=np_), axis=-1)
+
+  # No rotation in this case.
+  if not (latitude_local_coordinates or longitude_local_coordinates):
+    return node_pos[senders] - node_pos[receivers]
+
+  # Get rotation matrices for the local space space for every node.
+  rotation_matrices = get_rotation_matrices_to_local_coordinates(
+      reference_phi=node_phi,
+      reference_theta=node_theta,
+      rotate_latitude=latitude_local_coordinates,
+      rotate_longitude=longitude_local_coordinates,
+      np_=np_,
+      transform_=transform_)
+
+  # Each edge will be rotated according to the rotation matrix of its receiver
+  # node.
+  edge_rotation_matrices = rotation_matrices[receivers]
+
+  # Rotate all nodes to the rotated space of the corresponding edge.
+  # Note for receivers we can also do the matmul first and the gather second:
+  # ```
+  # receiver_pos_in_rotated_space = rotate_with_matrices(
+  #    rotation_matrices, node_pos)[receivers]
+  # ```
+  # which is more efficient, however, we do gather first to keep it more
+  # symmetric with the sender computation.
+  receiver_pos_in_rotated_space = rotate_with_matrices(
+      edge_rotation_matrices, node_pos[receivers], np_=np_)
+  sender_pos_in_in_rotated_space = rotate_with_matrices(
+      edge_rotation_matrices, node_pos[senders], np_=np_)
+  # Note, here, that because the rotated space is chosen according to the
+  # receiver, if:
+  # * latitude_local_coordinates = True: latitude for the receivers will be
+  #   0, that is the z coordinate will always be 0.
+  # * longitude_local_coordinates = True: longitude for the receivers will be
+  #   0, that is the y coordinate will be 0.
+
+  # Now we can just subtract.
+  # Note we are rotating to a local coordinate system, where the y-z axes are
+  # parallel to a tangent plane to the sphere, but still remain in a 3d space.
+  # Note that if both `latitude_local_coordinates` and
+  # `longitude_local_coordinates` are True, and edges are short,
+  # then the difference in x coordinate between sender and receiver
+  # should be small, so we could consider dropping the new x coordinate if
+  # we wanted to the tangent plane, however in doing so
+  # we would lose information about the curvature of the mesh, which may be
+  # important for very coarse meshes.
+  return sender_pos_in_in_rotated_space - receiver_pos_in_rotated_space
+
+
+def get_rotation_matrices_to_local_coordinates(
+    reference_phi: np.ndarray,
+    reference_theta: np.ndarray,
+    rotate_latitude: bool,
+    rotate_longitude: bool,
+    np_: NumpyInterface = np,
+    transform_: TransformInterface = transform,
+    ) -> np.ndarray:
+  """Returns a rotation matrix to rotate to a point based on a reference vector.
+
+  The rotation matrix is build such that, a vector in the
+  same coordinate system at the reference point that points towards the pole
+  before the rotation, continues to point towards the pole after the rotation.
+
+  Args:
+    reference_phi: [leading_axis] Polar angles of the reference.
+    reference_theta: [leading_axis] Azimuthal angles of the reference.
+    rotate_latitude: Whether to produce a rotation matrix that would rotate
+        R^3 vectors to zero latitude.
+    rotate_longitude: Whether to produce a rotation matrix that would rotate
+        R^3 vectors to zero longitude.
+    np_: Numpy library interface.
+    transform_: scipy.transform library interface.
+
+  Returns:
+    Matrices of shape [leading_axis] such that when applied to the reference
+        position with `rotate_with_matrices(rotation_matrices, reference_pos)`
+
+        * phi goes to 0. if "rotate_longitude" is True.
+
+        * theta goes to np.pi / 2 if "rotate_latitude" is True.
+
+        The rotation consists of:
+        * rotate_latitude = False, rotate_longitude = True:
+            Latitude preserving rotation.
+        * rotate_latitude = True, rotate_longitude = True:
+            Latitude preserving rotation, followed by longitude preserving
+            rotation.
+        * rotate_latitude = True, rotate_longitude = False:
+            Latitude preserving rotation, followed by longitude preserving
+            rotation, and the inverse of the latitude preserving rotation. Note
+            this is computationally different from rotating the longitude only
+            and is. We do it like this, so the polar geodesic curve, continues
+            to be aligned with one of the axis after the rotation.
+
+  """
+
+  # Azimuthal angle we need to apply to move to zero longitude.
+  azimuthal_rotation = - reference_phi
+
+  # Polar angle we need to apply to move from "theta" to zero latitude.
+  polar_rotation = - reference_theta + np.pi/2
+
+  if rotate_longitude and rotate_latitude:
+    # We first rotate to zero longitude around the z axis, and then, when the
+    # point is at x=0 we can simply apply the polar rotation around the y axis.
+    return transform_.Rotation.from_euler(
+        "zy", np_.stack([azimuthal_rotation, polar_rotation],
+                        axis=1)).as_matrix()
+  elif rotate_longitude:
+    # Just like the previous case, but applying only the azimuthal rotation,
+    # leaving the latitude unchanged. Even though it is a single rotation, we
+    # need a "sequence" axis of size 1 (hence the expand_dims).
+    return transform_.Rotation.from_euler(
+        "z", np_.expand_dims(azimuthal_rotation, axis=1)).as_matrix()
+  elif rotate_latitude:
+    # We want to apply the polar rotation only, but we don't know the rotation
+    # axis to apply a polar rotation. The simplest way to achieve this is to
+    # first rotate all the way to longitude 0, then apply the polar rotation
+    # arond the y axis, and then rotate back to the original longitude.
+    return transform_.Rotation.from_euler(
+        "zyz", np_.stack(
+            [azimuthal_rotation, polar_rotation, -azimuthal_rotation]
+            , axis=1)).as_matrix()
+  else:
+    raise ValueError(
+        "At least one of longitude and latitude should be rotated.")
+
+
+def rotate_with_matrices(rotation_matrices: np.ndarray, positions: np.ndarray,
+                         np_: NumpyInterface = np) -> np.ndarray:
+  return np_.einsum("...ji,...i->...j", rotation_matrices, positions)
+
+
+def get_bipartite_graph_spatial_features(
+    *,
+    senders_node_lat: np.ndarray,
+    senders_node_lon: np.ndarray,
+    senders: np.ndarray,
+    receivers_node_lat: np.ndarray,
+    receivers_node_lon: np.ndarray,
+    receivers: np.ndarray,
+    add_node_positions: bool,
+    add_node_latitude: bool,
+    add_node_longitude: bool,
+    add_relative_positions: bool,
+    edge_normalization_factor: Optional[float] = None,
+    relative_longitude_local_coordinates: bool,
+    relative_latitude_local_coordinates: bool,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+  """Computes spatial features for the nodes.
+
+  This function is almost identical to `get_graph_spatial_features`. The only
+  difference is that sender nodes and receiver nodes can be in different arrays.
+  This is necessary to enable combination with typed Graph.
+
+  Args:
+    senders_node_lat: Latitudes in the [-90, 90] interval of shape
+      [num_sender_nodes]
+    senders_node_lon: Longitudes in the [0, 360] interval of shape
+      [num_sender_nodes]
+    senders: Sender indices of shape [num_edges], indices in [0,
+      num_sender_nodes)
+    receivers_node_lat: Latitudes in the [-90, 90] interval of shape
+      [num_receiver_nodes]
+    receivers_node_lon: Longitudes in the [0, 360] interval of shape
+      [num_receiver_nodes]
+    receivers: Receiver indices of shape [num_edges], indices in [0,
+      num_receiver_nodes)
+    add_node_positions: Add unit norm absolute positions.
+    add_node_latitude: Add a feature for latitude (cos(90 - lat)) Note even if
+      this is set to False, the model may be able to infer the longitude from
+      relative features, unless `relative_latitude_local_coordinates` is also
+      True, or if there is any bias on the relative edge sizes for different
+      longitudes.
+    add_node_longitude: Add features for longitude (cos(lon), sin(lon)). Note
+      even if this is set to False, the model may be able to infer the longitude
+      from relative features, unless `relative_longitude_local_coordinates` is
+      also True, or if there is any bias on the relative edge sizes for
+      different longitudes.
+    add_relative_positions: Whether to relative positions in R3 to the edges.
+    edge_normalization_factor: Allows explicitly controlling edge normalization.
+      If None, defaults to max edge length. This supports using pre-trained
+      model weights with a different graph structure to what it was trained on.
+    relative_longitude_local_coordinates: If True, relative positions are
+      computed in a local space where the receiver is at 0 longitude.
+    relative_latitude_local_coordinates: If True, relative positions are
+      computed in a local space where the receiver is at 0 latitude.
+
+  Returns:
+    Arrays of shape: [num_nodes, num_features] and [num_edges, num_features].
+    with node and edge features.
+
+  """
+
+  num_senders = senders_node_lat.shape[0]
+  num_receivers = receivers_node_lat.shape[0]
+  num_edges = senders.shape[0]
+  dtype = senders_node_lat.dtype
+  assert receivers_node_lat.dtype == dtype
+  senders_node_phi, senders_node_theta = lat_lon_deg_to_spherical(
+      senders_node_lat, senders_node_lon)
+  receivers_node_phi, receivers_node_theta = lat_lon_deg_to_spherical(
+      receivers_node_lat, receivers_node_lon)
+
+  # Computing some node features.
+  senders_node_features = []
+  receivers_node_features = []
+  if add_node_positions:
+    # Already in [-1, 1.] range.
+    senders_node_features.extend(
+        spherical_to_cartesian(senders_node_phi, senders_node_theta))
+    receivers_node_features.extend(
+        spherical_to_cartesian(receivers_node_phi, receivers_node_theta))
+
+  if add_node_latitude:
+    # Using the cos of theta.
+    # From 1. (north pole) to -1 (south pole).
+    senders_node_features.append(np.cos(senders_node_theta))
+    receivers_node_features.append(np.cos(receivers_node_theta))
+
+  if add_node_longitude:
+    # Using the cos and sin, which is already normalized.
+    senders_node_features.append(np.cos(senders_node_phi))
+    senders_node_features.append(np.sin(senders_node_phi))
+
+    receivers_node_features.append(np.cos(receivers_node_phi))
+    receivers_node_features.append(np.sin(receivers_node_phi))
+
+  if not senders_node_features:
+    senders_node_features = np.zeros([num_senders, 0], dtype=dtype)
+    receivers_node_features = np.zeros([num_receivers, 0], dtype=dtype)
+  else:
+    senders_node_features = np.stack(senders_node_features, axis=-1)
+    receivers_node_features = np.stack(receivers_node_features, axis=-1)
+
+  # Computing some edge features.
+  edge_features = []
+
+  if add_relative_positions:
+
+    relative_position = get_bipartite_relative_position_in_receiver_local_coordinates(  # pylint: disable=line-too-long
+        senders_node_phi=senders_node_phi,
+        senders_node_theta=senders_node_theta,
+        receivers_node_phi=receivers_node_phi,
+        receivers_node_theta=receivers_node_theta,
+        senders=senders,
+        receivers=receivers,
+        latitude_local_coordinates=relative_latitude_local_coordinates,
+        longitude_local_coordinates=relative_longitude_local_coordinates)
+
+    # Note this is L2 distance in 3d space, rather than geodesic distance.
+    relative_edge_distances = np.linalg.norm(
+        relative_position, axis=-1, keepdims=True)
+
+    if edge_normalization_factor is None:
+      # Normalize to the maximum edge distance. Note that we expect to always
+      # have an edge that goes in the opposite direction of any given edge
+      # so the distribution of relative positions should be symmetric around
+      # zero. So by scaling by the maximum length, we expect all relative
+      # positions to fall in the [-1., 1.] interval, and all relative distances
+      # to fall in the [0., 1.] interval.
+      edge_normalization_factor = relative_edge_distances.max()
+
+    edge_features.append(relative_edge_distances / edge_normalization_factor)
+    edge_features.append(relative_position / edge_normalization_factor)
+
+  if not edge_features:
+    edge_features = np.zeros([num_edges, 0], dtype=dtype)
+  else:
+    edge_features = np.concatenate(edge_features, axis=-1)
+
+  return senders_node_features, receivers_node_features, edge_features
+
+
+def get_bipartite_relative_position_in_receiver_local_coordinates(
+    senders_node_phi: np.ndarray,
+    senders_node_theta: np.ndarray,
+    senders: np.ndarray,
+    receivers_node_phi: np.ndarray,
+    receivers_node_theta: np.ndarray,
+    receivers: np.ndarray,
+    latitude_local_coordinates: bool,
+    longitude_local_coordinates: bool,
+    np_: NumpyInterface = np,
+    transform_: TransformInterface = transform,
+    ) -> np.ndarray:
+  """Returns relative position features for the edges.
+
+  This function is equivalent to
+  `get_relative_position_in_receiver_local_coordinates`, but adapted to work
+  with bipartite typed graphs.
+
+  The relative positions will be computed in a rotated space for a local
+  coordinate system as defined by the receiver. The relative positions are
+  simply obtained by subtracting sender position minues receiver position in
+  that local coordinate system after the rotation in R^3.
+
+  Args:
+    senders_node_phi: [num_sender_nodes] with polar angles.
+    senders_node_theta: [num_sender_nodes] with azimuthal angles.
+    senders: [num_edges] with indices into sender nodes.
+    receivers_node_phi: [num_sender_nodes] with polar angles.
+    receivers_node_theta: [num_sender_nodes] with azimuthal angles.
+    receivers: [num_edges] with indices into receiver nodes.
+    latitude_local_coordinates: Whether to rotate edges such that in the
+      positions are computed such that the receiver is always at latitude 0.
+    longitude_local_coordinates: Whether to rotate edges such that in the
+      positions are computed such that the receiver is always at longitude 0.
+    np_: Numpy library interface.
+    transform_: scipy.transform library interface.
+
+  Returns:
+    Array of relative positions in R3 [num_edges, 3]
+  """
+
+  senders_node_pos = np_.stack(
+      spherical_to_cartesian(
+          senders_node_phi, senders_node_theta, np_=np_), axis=-1)
+
+  receivers_node_pos = np_.stack(
+      spherical_to_cartesian(
+          receivers_node_phi, receivers_node_theta, np_=np_), axis=-1)
+
+  # No rotation in this case.
+  if not (latitude_local_coordinates or longitude_local_coordinates):
+    return senders_node_pos[senders] - receivers_node_pos[receivers]
+
+  # Get rotation matrices for the local space space for every receiver node.
+  receiver_rotation_matrices = get_rotation_matrices_to_local_coordinates(
+      reference_phi=receivers_node_phi,
+      reference_theta=receivers_node_theta,
+      rotate_latitude=latitude_local_coordinates,
+      rotate_longitude=longitude_local_coordinates,
+      np_=np_,
+      transform_=transform_)
+
+  # Each edge will be rotated according to the rotation matrix of its receiver
+  # node.
+  edge_rotation_matrices = receiver_rotation_matrices[receivers]
+
+  # Rotate all nodes to the rotated space of the corresponding edge.
+  # Note for receivers we can also do the matmul first and the gather second:
+  # ```
+  # receiver_pos_in_rotated_space = rotate_with_matrices(
+  #    rotation_matrices, node_pos)[receivers]
+  # ```
+  # which is more efficient, however, we do gather first to keep it more
+  # symmetric with the sender computation.
+  receiver_pos_in_rotated_space = rotate_with_matrices(
+      edge_rotation_matrices, receivers_node_pos[receivers], np_=np_)
+  sender_pos_in_in_rotated_space = rotate_with_matrices(
+      edge_rotation_matrices, senders_node_pos[senders], np_=np_)
+  # Note, here, that because the rotated space is chosen according to the
+  # receiver, if:
+  # * latitude_local_coordinates = True: latitude for the receivers will be
+  #   0, that is the z coordinate will always be 0.
+  # * longitude_local_coordinates = True: longitude for the receivers will be
+  #   0, that is the y coordinate will be 0.
+
+  # Now we can just subtract.
+  # Note we are rotating to a local coordinate system, where the y-z axes are
+  # parallel to a tangent plane to the sphere, but still remain in a 3d space.
+  # Note that if both `latitude_local_coordinates` and
+  # `longitude_local_coordinates` are True, and edges are short,
+  # then the difference in x coordinate between sender and receiver
+  # should be small, so we could consider dropping the new x coordinate if
+  # we wanted to the tangent plane, however in doing so
+  # we would lose information about the curvature of the mesh, which may be
+  # important for very coarse meshes.
+  return sender_pos_in_in_rotated_space - receiver_pos_in_rotated_space
+
+
+def variable_to_stacked(
+    variable: xarray.Variable,
+    sizes: Mapping[str, int],
+    preserved_dims: Tuple[str, ...] = ("batch", "lat", "lon"),
+    ) -> xarray.Variable:
+  """Converts an xarray.Variable to preserved_dims + ("channels",).
+
+  Any dimensions other than those included in preserved_dims get stacked into a
+  final "channels" dimension. If any of the preserved_dims are missing then they
+  are added, with the data broadcast/tiled to match the sizes specified in
+  `sizes`.
+
+  Args:
+    variable: An xarray.Variable.
+    sizes: Mapping including sizes for any dimensions which are not present in
+      `variable` but are needed for the output. This may be needed for example
+      for a static variable with only ("lat", "lon") dims, or if you want to
+      encode just the latitude coordinates (a variable with dims ("lat",)).
+    preserved_dims: dimensions of variable to not be folded in channels.
+
+  Returns:
+    An xarray.Variable with dimensions preserved_dims + ("channels",).
+  """
+  stack_to_channels_dims = [
+      d for d in variable.dims if d not in preserved_dims]
+  if stack_to_channels_dims:
+    variable = variable.stack(channels=stack_to_channels_dims)
+  dims = {dim: variable.sizes.get(dim) or sizes[dim] for dim in preserved_dims}
+  dims["channels"] = variable.sizes.get("channels", 1)
+  return variable.set_dims(dims)
+
+
+def dataset_to_stacked(
+    dataset: xarray.Dataset,
+    sizes: Optional[Mapping[str, int]] = None,
+    preserved_dims: Tuple[str, ...] = ("batch", "lat", "lon"),
+) -> xarray.DataArray:
+  """Converts an xarray.Dataset to a single stacked array.
+
+  This takes each consistuent data_var, converts it into BHWC layout
+  using `variable_to_stacked`, then concats them all along the channels axis.
+
+  Args:
+    dataset: An xarray.Dataset.
+    sizes: Mapping including sizes for any dimensions which are not present in
+      the `dataset` but are needed for the output. See variable_to_stacked.
+    preserved_dims: dimensions from the dataset that should not be folded in
+      the predictions channels.
+
+  Returns:
+    An xarray.DataArray with dimensions preserved_dims + ("channels",).
+    Existing coordinates for preserved_dims axes will be preserved, however
+    there will be no coordinates for "channels".
+  """
+  data_vars = [
+      variable_to_stacked(dataset.variables[name], sizes or dataset.sizes,
+                          preserved_dims)
+      for name in sorted(dataset.data_vars.keys())
+  ]
+  coords = {
+      dim: coord
+      for dim, coord in dataset.coords.items()
+      if dim in preserved_dims
+  }
+  return xarray.DataArray(
+      data=xarray.Variable.concat(data_vars, dim="channels"), coords=coords)
+
+
+def stacked_to_dataset(
+    stacked_array: xarray.Variable,
+    template_dataset: xarray.Dataset,
+    preserved_dims: Tuple[str, ...] = ("batch", "lat", "lon"),
+    ) -> xarray.Dataset:
+  """The inverse of dataset_to_stacked.
+
+  Requires a template dataset to demonstrate the variables/shapes/coordinates
+  required.
+  All variables must have preserved_dims dimensions.
+
+  Args:
+    stacked_array: Data in BHWC layout, encoded the same as dataset_to_stacked
+      would if it was asked to encode `template_dataset`.
+    template_dataset: A template Dataset (or other mapping of DataArrays)
+      demonstrating the shape of output required (variables, shapes,
+      coordinates etc).
+    preserved_dims: dimensions from the target_template that were not folded in
+      the predictions channels. The preserved_dims need to be a subset of the
+      dims of all the variables of template_dataset.
+
+  Returns:
+    An xarray.Dataset (or other mapping of DataArrays) with the same shape and
+    type as template_dataset.
+  """
+  unstack_from_channels_sizes = {}
+  var_names = sorted(template_dataset.keys())
+  for name in var_names:
+    template_var = template_dataset[name]
+    if not all(dim in template_var.dims for dim in preserved_dims):
+      raise ValueError(
+          f"stacked_to_dataset requires all Variables to have {preserved_dims} "
+          f"dimensions, but found only {template_var.dims}.")
+    unstack_from_channels_sizes[name] = {
+        dim: size for dim, size in template_var.sizes.items()
+        if dim not in preserved_dims}
+
+  channels = {name: np.prod(list(unstack_sizes.values()), dtype=np.int64)
+              for name, unstack_sizes in unstack_from_channels_sizes.items()}
+  total_expected_channels = sum(channels.values())
+  found_channels = stacked_array.sizes["channels"]
+  if total_expected_channels != found_channels:
+    raise ValueError(
+        f"Expected {total_expected_channels} channels but found "
+        f"{found_channels}, when trying to convert a stacked array of shape "
+        f"{stacked_array.sizes} to a dataset of shape {template_dataset}.")
+
+  data_vars = {}
+  index = 0
+  for name in var_names:
+    template_var = template_dataset[name]
+    var = stacked_array.isel({"channels": slice(index, index + channels[name])})
+    index += channels[name]
+    var = var.unstack({"channels": unstack_from_channels_sizes[name]})
+    var = var.transpose(*template_var.dims)
+    data_vars[name] = xarray.DataArray(
+        data=var,
+        coords=template_var.coords,
+        # This might not always be the same as the name it's keyed under; it
+        # will refer to the original variable name, whereas the key might be
+        # some alias e.g. temperature_850 under which it should be logged:
+        name=template_var.name,
+    )
+  return type(template_dataset)(data_vars)  # pytype:disable=not-callable,wrong-arg-count
+
+
+def fourier_features(
+    values: jnp.ndarray,
+    base_period: float,
+    num_frequencies: int,
+    ) -> jnp.ndarray:
+  """Maps values to sin/cos features for a range of frequencies.
+
+  Args:
+    values: Values to compute Fourier features for.
+    base_period: The base period to use. This should be greater or equal to the
+      range of the values, or to the period if the values have periodic
+      semantics (e.g. 2pi if they represent angles). Frequencies used will be
+      integer multiples of 1/base_period.
+    num_frequencies: The number of frequencies to use, we will use integer
+      multiples of 1/base_period from 1 up to num_frequencies inclusive. (We
+      don't include a zero frequency as this would just give constant features
+      which are redundant if a bias term is present).
+
+  Returns:
+    Array with same shape as values except with an extra trailing dimension
+    of size 2*num_frequencies, which contains a sin and a cos feature for each
+    frequency.
+  """
+  frequencies = np.arange(1, num_frequencies + 1) / base_period
+  angular_frequencies = jnp.array(2 * np.pi * frequencies, dtype=values.dtype)
+  values_times_angular_freqs = values[..., None] * angular_frequencies
+  return jnp.concatenate(
+      [jnp.cos(values_times_angular_freqs),
+       jnp.sin(values_times_angular_freqs)],
+      axis=-1)
+

model/graphcast/nan_cleaning.py

@@ -0,0 +1,125 @@
+# Copyright 2024 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Wrappers for Predictors which allow them to work with data cleaned of NaNs.
+
+The Predictor which is wrapped sees inputs and targets without NaNs, and makes
+NaNless predictions.
+"""
+
+from typing import Optional, Tuple
+
+from . import predictor_base as base
+import numpy as np
+import xarray
+
+
+class NaNCleaner(base.Predictor):
+  """A predictor wrapper than removes NaNs from ingested data.
+
+  The Predictor which is wrapped sees inputs and targets without NaNs.
+  """
+
+  def __init__(
+      self,
+      predictor: base.Predictor,
+      var_to_clean: str,
+      fill_value: xarray.Dataset,
+      reintroduce_nans: bool = False,
+  ):
+    """Initializes the NaNCleaner."""
+    self._predictor = predictor
+    self._fill_value = fill_value[var_to_clean]
+    self._var_to_clean = var_to_clean
+    self._reintroduce_nans = reintroduce_nans
+
+  def _clean(self, dataset: xarray.Dataset) -> xarray.Dataset:
+    """Cleans the dataset of NaNs."""
+    data_array = dataset[self._var_to_clean]
+    dataset = dataset.assign(
+        {self._var_to_clean: data_array.fillna(self._fill_value)}
+    )
+    return dataset
+
+  def _maybe_reintroduce_nans(
+      self, stale_inputs: xarray.Dataset, predictions: xarray.Dataset
+  ) -> xarray.Dataset:
+    # NaN positions don't change between input frames, if they do then
+    # we should be more careful about re-introducing them.
+    if self._var_to_clean in predictions.keys():
+      nan_mask = np.isnan(stale_inputs[self._var_to_clean]).any(dim='time')
+      with_nan_values = predictions[self._var_to_clean].where(~nan_mask, np.nan)
+      predictions = predictions.assign({self._var_to_clean: with_nan_values})
+    return predictions
+
+  def __call__(
+      self,
+      inputs: xarray.Dataset,
+      targets_template: xarray.Dataset,
+      forcings: Optional[xarray.Dataset] = None,
+      **kwargs,
+  ) -> xarray.Dataset:
+    if self._reintroduce_nans:
+      # Copy inputs before cleaning so that we can reintroduce NaNs later.
+      original_inputs = inputs.copy()
+    if self._var_to_clean in inputs.keys():
+      inputs = self._clean(inputs)
+    if forcings and self._var_to_clean in forcings.keys():
+      forcings = self._clean(forcings)
+    predictions = self._predictor(
+        inputs, targets_template, forcings, **kwargs
+    )
+    if self._reintroduce_nans:
+      predictions = self._maybe_reintroduce_nans(original_inputs, predictions)
+    return predictions
+
+  def loss(
+      self,
+      inputs: xarray.Dataset,
+      targets: xarray.Dataset,
+      forcings: Optional[xarray.Dataset] = None,
+      **kwargs,
+  ) -> base.LossAndDiagnostics:
+    if self._var_to_clean in inputs.keys():
+      inputs = self._clean(inputs)
+    if self._var_to_clean in targets.keys():
+      targets = self._clean(targets)
+    if forcings and self._var_to_clean in forcings.keys():
+      forcings = self._clean(forcings)
+    return self._predictor.loss(
+        inputs, targets, forcings, **kwargs
+    )
+
+  def loss_and_predictions(
+      self,
+      inputs: xarray.Dataset,
+      targets: xarray.Dataset,
+      forcings: Optional[xarray.Dataset] = None,
+      **kwargs,
+  ) -> Tuple[base.LossAndDiagnostics, xarray.Dataset]:
+    if self._reintroduce_nans:
+      # Copy inputs before cleaning so that we can reintroduce NaNs later.
+      original_inputs = inputs.copy()
+    if self._var_to_clean in inputs.keys():
+      inputs = self._clean(inputs)
+    if self._var_to_clean in targets.keys():
+      targets = self._clean(targets)
+    if forcings and self._var_to_clean in forcings.keys():
+      forcings = self._clean(forcings)
+
+    loss, predictions = self._predictor.loss_and_predictions(
+        inputs, targets, forcings, **kwargs
+    )
+    if self._reintroduce_nans:
+      predictions = self._maybe_reintroduce_nans(original_inputs, predictions)
+    return loss, predictions

model/graphcast/normalization.py

@@ -0,0 +1,196 @@
+# Copyright 2023 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Wrappers for Predictors which allow them to work with normalized data.
+
+The Predictor which is wrapped sees normalized inputs and targets, and makes
+normalized predictions. The wrapper handles translating the predictions back
+to the original domain.
+"""
+
+import logging
+from typing import Optional, Tuple
+
+from . import predictor_base
+from . import xarray_tree
+import xarray
+
+
+def normalize(values: xarray.Dataset,
+              scales: xarray.Dataset,
+              locations: Optional[xarray.Dataset],
+              ) -> xarray.Dataset:
+  """Normalize variables using the given scales and (optionally) locations."""
+  def normalize_array(array):
+    if array.name is None:
+      raise ValueError(
+          "Can't look up normalization constants because array has no name.")
+    if locations is not None:
+      if array.name in locations:
+        array = array - locations[array.name].astype(array.dtype)
+      else:
+        logging.warning('No normalization location found for %s', array.name)
+    if array.name in scales:
+      array = array / scales[array.name].astype(array.dtype)
+    else:
+      logging.warning('No normalization scale found for %s', array.name)
+    return array
+  return xarray_tree.map_structure(normalize_array, values)
+
+
+def unnormalize(values: xarray.Dataset,
+                scales: xarray.Dataset,
+                locations: Optional[xarray.Dataset],
+                ) -> xarray.Dataset:
+  """Unnormalize variables using the given scales and (optionally) locations."""
+  def unnormalize_array(array):
+    if array.name is None:
+      raise ValueError(
+          "Can't look up normalization constants because array has no name.")
+    if array.name in scales:
+      array = array * scales[array.name].astype(array.dtype)
+    else:
+      logging.warning('No normalization scale found for %s', array.name)
+    if locations is not None:
+      if array.name in locations:
+        array = array + locations[array.name].astype(array.dtype)
+      else:
+        logging.warning('No normalization location found for %s', array.name)
+    return array
+  return xarray_tree.map_structure(unnormalize_array, values)
+
+
+class InputsAndResiduals(predictor_base.Predictor):
+  """Wraps with a residual connection, normalizing inputs and target residuals.
+
+  The inner predictor is given inputs that are normalized using `locations`
+  and `scales` to roughly zero-mean unit variance.
+
+  For target variables that are present in the inputs, the inner predictor is
+  trained to predict residuals (target - last_frame_of_input) that have been
+  normalized using `residual_scales` (and optionally `residual_locations`) to
+  roughly unit variance / zero mean.
+
+  This replaces `residual.Predictor` in the case where you want normalization
+  that's based on the scales of the residuals.
+
+  Since we return the underlying predictor's loss on the normalized residuals,
+  if the underlying predictor is a sum of per-variable losses, the normalization
+  will affect the relative weighting of the per-variable loss terms (hopefully
+  in a good way).
+
+  For target variables *not* present in the inputs, the inner predictor is
+  trained to predict targets directly, that have been normalized in the same
+  way as the inputs.
+
+  The transforms applied to the targets (the residual connection and the
+  normalization) are applied in reverse to the predictions before returning
+  them.
+  """
+
+  def __init__(
+      self,
+      predictor: predictor_base.Predictor,
+      stddev_by_level: xarray.Dataset,
+      mean_by_level: xarray.Dataset,
+      diffs_stddev_by_level: xarray.Dataset):
+    self._predictor = predictor
+    self._scales = stddev_by_level
+    self._locations = mean_by_level
+    self._residual_scales = diffs_stddev_by_level
+    self._residual_locations = None
+
+  def _unnormalize_prediction_and_add_input(self, inputs, norm_prediction):
+    if norm_prediction.sizes.get('time') != 1:
+      raise ValueError(
+          'normalization.InputsAndResiduals only supports predicting a '
+          'single timestep.')
+    if norm_prediction.name in inputs:
+      # Residuals are assumed to be predicted as normalized (unit variance),
+      # but the scale and location they need mapping to is that of the residuals
+      # not of the values themselves.
+      prediction = unnormalize(
+          norm_prediction, self._residual_scales, self._residual_locations)
+      # A prediction for which we have a corresponding input -- we are
+      # predicting the residual:
+      last_input = inputs[norm_prediction.name].isel(time=-1)
+      prediction = prediction + last_input
+      return prediction
+    else:
+      # A predicted variable which is not an input variable. We are predicting
+      # it directly, so unnormalize it directly to the target scale/location:
+      return unnormalize(norm_prediction, self._scales, self._locations)
+
+  def _subtract_input_and_normalize_target(self, inputs, target):
+    if target.sizes.get('time') != 1:
+      raise ValueError(
+          'normalization.InputsAndResiduals only supports wrapping predictors'
+          'that predict a single timestep.')
+    if target.name in inputs:
+      target_residual = target
+      last_input = inputs[target.name].isel(time=-1)
+      target_residual = target_residual - last_input
+      return normalize(
+          target_residual, self._residual_scales, self._residual_locations)
+    else:
+      return normalize(target, self._scales, self._locations)
+
+  def __call__(self,
+               inputs: xarray.Dataset,
+               targets_template: xarray.Dataset,
+               forcings: xarray.Dataset,
+               **kwargs
+               ) -> xarray.Dataset:
+    norm_inputs = normalize(inputs, self._scales, self._locations)
+    norm_forcings = normalize(forcings, self._scales, self._locations)
+    norm_predictions = self._predictor(
+        norm_inputs, targets_template, forcings=norm_forcings, **kwargs)
+    return xarray_tree.map_structure(
+        lambda pred: self._unnormalize_prediction_and_add_input(inputs, pred),
+        norm_predictions)
+
+  def loss(self,
+           inputs: xarray.Dataset,
+           targets: xarray.Dataset,
+           forcings: xarray.Dataset,
+           **kwargs,
+           ) -> predictor_base.LossAndDiagnostics:
+    """Returns the loss computed on normalized inputs and targets."""
+    norm_inputs = normalize(inputs, self._scales, self._locations)
+    norm_forcings = normalize(forcings, self._scales, self._locations)
+    norm_target_residuals = xarray_tree.map_structure(
+        lambda t: self._subtract_input_and_normalize_target(inputs, t),
+        targets)
+    return self._predictor.loss(
+        norm_inputs, norm_target_residuals, forcings=norm_forcings, **kwargs)
+
+  def loss_and_predictions(  # pytype: disable=signature-mismatch  # jax-ndarray
+      self,
+      inputs: xarray.Dataset,
+      targets: xarray.Dataset,
+      forcings: xarray.Dataset,
+      **kwargs,
+      ) -> Tuple[predictor_base.LossAndDiagnostics,
+                 xarray.Dataset]:
+    """The loss computed on normalized data, with unnormalized predictions."""
+    norm_inputs = normalize(inputs, self._scales, self._locations)
+    norm_forcings = normalize(forcings, self._scales, self._locations)
+    norm_target_residuals = xarray_tree.map_structure(
+        lambda t: self._subtract_input_and_normalize_target(inputs, t),
+        targets)
+    (loss, scalars), norm_predictions = self._predictor.loss_and_predictions(
+        norm_inputs, norm_target_residuals, forcings=norm_forcings, **kwargs)
+    predictions = xarray_tree.map_structure(
+        lambda pred: self._unnormalize_prediction_and_add_input(inputs, pred),
+        norm_predictions)
+    return (loss, scalars), predictions

model/graphcast/predictor_base.py

@@ -0,0 +1,170 @@
+# Copyright 2023 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Abstract base classes for an xarray-based Predictor API."""
+
+import abc
+
+from typing import Tuple
+
+from . import losses
+from . import xarray_jax
+import jax.numpy as jnp
+import xarray
+
+LossAndDiagnostics = losses.LossAndDiagnostics
+
+
+class Predictor(abc.ABC):
+  """A possibly-trainable predictor of weather, exposing an xarray-based API.
+
+  Typically wraps an underlying JAX model and handles translating the xarray
+  Dataset values to and from plain JAX arrays that are convenient for input to
+  (and output from) the underlying model.
+
+  Different subclasses may exist to wrap different kinds of underlying model,
+  e.g. models taking stacked inputs/outputs, models taking separate 2D and 3D
+  inputs/outputs, autoregressive models.
+
+  You can also implement a specific model directly as a Predictor if you want,
+  for example if it has quite specific/unique requirements for its input/output
+  or loss function, or if it's convenient to implement directly using xarray.
+  """
+
+  @abc.abstractmethod
+  def __call__(self,
+               inputs: xarray.Dataset,
+               targets_template: xarray.Dataset,
+               forcings: xarray.Dataset,
+               **optional_kwargs
+               ) -> xarray.Dataset:
+    """Makes predictions.
+
+    This is only used by the Experiment for inference / evaluation, with
+    training going via the .loss method. So it should default to making
+    predictions for evaluation, although you can also support making predictions
+    for use in the loss via an is_training argument -- see
+    LossFunctionPredictor which helps with that.
+
+    Args:
+      inputs: An xarray.Dataset of inputs.
+      targets_template: An xarray.Dataset or other mapping of xarray.DataArrays,
+        with the same shape as the targets, to demonstrate what kind of
+        predictions are required. You can use this to determine which variables,
+        levels and lead times must be predicted.
+        You are free to raise an error if you don't support predicting what is
+        requested.
+      forcings: An xarray.Dataset of forcings terms. Forcings are variables
+        that can be fed to the model, but do not need to be predicted. This is
+        often because this variable can be computed analytically (e.g. the toa
+        radiation of the sun is mostly a function of geometry) or are considered
+        to be controlled for the experiment (e.g., impose a scenario of C02
+        emission into the atmosphere). Unlike `inputs`, the `forcings` can
+        include information "from the future", that is, information at target
+        times specified in the `targets_template`.
+      **optional_kwargs: Implementations may support extra optional kwargs,
+        provided they set appropriate defaults for them.
+
+    Returns:
+      Predictions, as an xarray.Dataset or other mapping of DataArrays which
+      is capable of being evaluated against targets with shape given by
+      targets_template.
+      For probabilistic predictors which can return multiple samples from a
+      predictive distribution, these should (by convention) be returned along
+      an additional 'sample' dimension.
+    """
+
+  def loss(self,
+           inputs: xarray.Dataset,
+           targets: xarray.Dataset,
+           forcings: xarray.Dataset,
+           **optional_kwargs,
+           ) -> LossAndDiagnostics:
+    """Computes a training loss, for predictors that are trainable.
+
+    Why make this the Predictor's responsibility, rather than letting callers
+    compute their own loss function using predictions obtained from
+    Predictor.__call__?
+
+    Doing it this way gives Predictors more control over their training setup.
+    For example, some predictors may wish to train using different targets to
+    the ones they predict at evaluation time -- perhaps different lead times and
+    variables, perhaps training to predict transformed versions of targets
+    where the transform needs to be inverted at evaluation time, etc.
+
+    It's also necessary for generative models (VAEs, GANs, ...) where the
+    training loss is more complex and isn't expressible as a parameter-free
+    function of predictions and targets.
+
+    Args:
+      inputs: An xarray.Dataset.
+      targets: An xarray.Dataset or other mapping of xarray.DataArrays. See
+        docs on __call__ for an explanation about the targets.
+      forcings: xarray.Dataset of forcing terms.
+      **optional_kwargs: Implementations may support extra optional kwargs,
+        provided they set appropriate defaults for them.
+
+    Returns:
+      loss: A DataArray with dimensions ('batch',) containing losses for each
+        element of the batch. These will be averaged to give the final
+        loss, locally and across replicas.
+      diagnostics: Mapping of additional quantities to log by name alongside the
+        loss. These will will typically correspond to terms in the loss. They
+        should also have dimensions ('batch',) and will be averaged over the
+        batch before logging.
+        You need not include the loss itself in this dict; it will be added for
+        you.
+    """
+    del targets, forcings, optional_kwargs
+    batch_size = inputs.sizes['batch']
+    dummy_loss = xarray_jax.DataArray(jnp.zeros(batch_size), dims=('batch',))
+    return dummy_loss, {}  # pytype: disable=bad-return-type
+
+  def loss_and_predictions(
+      self,
+      inputs: xarray.Dataset,
+      targets: xarray.Dataset,
+      forcings: xarray.Dataset,
+      **optional_kwargs,
+      ) -> Tuple[LossAndDiagnostics, xarray.Dataset]:
+    """Like .loss but also returns corresponding predictions.
+
+    Implementing this is optional as it's not used directly by the Experiment,
+    but it is required by autoregressive.Predictor when applying an inner
+    Predictor autoregressively at training time; we need a loss at each step but
+    also predictions to feed back in for the next step.
+
+    Note the loss itself may not be directly regressing the predictions towards
+    targets, the loss may be computed in terms of transformed predictions and
+    targets (or in some other way). For this reason we can't always cleanly
+    separate this into step 1: get predictions, step 2: compute loss from them,
+    hence the need for this combined method.
+
+    Args:
+      inputs:
+      targets:
+      forcings:
+      **optional_kwargs:
+        As for self.loss.
+
+    Returns:
+      (loss, diagnostics)
+        As for self.loss
+      predictions:
+        The predictions which the loss relates to. These should be of the same
+        shape as what you would get from
+        `self.__call__(inputs, targets_template=targets)`, and should be in the
+        same 'domain' as the inputs (i.e. they shouldn't be transformed
+        differently to how the predictor expects its inputs).
+    """
+    raise NotImplementedError

model/graphcast/rollout.py

@@ -0,0 +1,498 @@
+# Copyright 2023 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Utils for rolling out models."""
+
+from typing import Iterator, Optional, Sequence
+
+from absl import logging
+import chex
+import dask.array
+from . import xarray_jax
+from . import xarray_tree
+import jax
+import jax.numpy as jnp
+import numpy as np
+import typing_extensions
+import xarray
+
+
+def _device_put_sharded(data_list, devices, axis_name):
+  """Stack data and put on devices with consistent sharding.
+
+  Creates a mesh with axis_name to ensure JIT cache consistency with pmap.
+
+  Args:
+    data_list: List of data to stack and put on devices.
+    devices: List of devices to put the data on.
+    axis_name: Name of the axis to use for sharding.
+
+  Returns:
+    Data put on devices with consistent sharding.
+  """
+
+  mesh = jax.sharding.Mesh(np.array(devices), (axis_name,))
+  sharding = jax.NamedSharding(mesh, jax.P(axis_name))
+  stack_fn = (
+      jnp.stack
+      if all(isinstance(x, jax.Array) for x in data_list)
+      else np.stack
+  )
+  stacked = stack_fn(data_list, axis=0)
+  return jax.device_put(stacked, sharding)
+
+
+class PredictorFn(typing_extensions.Protocol):
+  """Functional version of base.Predictor.__call__ with explicit rng."""
+
+  def __call__(
+      self, rng: chex.PRNGKey, inputs: xarray.Dataset,
+      targets_template: xarray.Dataset,
+      forcings: xarray.Dataset,
+      **optional_kwargs,
+      ) -> xarray.Dataset:
+    ...
+
+
+def _replicate_dataset(
+    data: xarray.Dataset, replica_dim: str,
+    replicate_to_device: bool,
+    devices: Sequence[jax.Device],
+    ) -> xarray.Dataset:
+  """Used to prepare for xarray_jax.pmap."""
+
+  def replicate_variable(variable: xarray.Variable) -> xarray.Variable:
+    if replica_dim in variable.dims:
+      # TODO(pricei): call device_put_replicated when replicate_to_device==True
+      return variable.transpose(replica_dim, ...)
+    else:
+      data = len(devices) * [variable.data]
+      if replicate_to_device:
+        assert devices is not None
+        data = _device_put_sharded(data, devices, replica_dim)
+      else:
+        data = np.stack(data, axis=0)
+      return xarray_jax.Variable(
+          data=data, dims=(replica_dim,) + variable.dims, attrs=variable.attrs
+      )
+
+  def replicate_dataset(dataset: xarray.Dataset) -> xarray.Dataset:
+    if dataset is None:
+      return None
+    data_variables = {
+        name: replicate_variable(var)
+        for name, var in dataset.data_vars.variables.items()
+    }
+    coords = {name: coord.variable for name, coord in dataset.coords.items()}
+    return xarray.Dataset(data_variables, coords=coords, attrs=dataset.attrs)
+
+  return replicate_dataset(data)
+
+
+def chunked_prediction_generator_multiple_runs(
+    predictor_fn: PredictorFn,
+    rngs: chex.PRNGKey,
+    inputs: xarray.Dataset,
+    targets_template: xarray.Dataset,
+    forcings: Optional[xarray.Dataset],
+    num_samples: Optional[int],
+    pmap_devices: Optional[Sequence[jax.Device]] = None,
+    **chunked_prediction_kwargs,
+) -> Iterator[xarray.Dataset]:
+  """Outputs a trajectory of multiple samples by yielding chunked predictions.
+
+  Args:
+    predictor_fn: Function to use to make predictions for each chunk.
+    rngs: RNG sequence to be used for each ensemble member.
+    inputs: Inputs for the model.
+    targets_template: Template for the target prediction, requires targets
+        equispaced in time.
+    forcings: Optional forcing for the model.
+    num_samples: The number of runs / samples to rollout.
+    pmap_devices: List of devices over which predictor_fn is pmapped, or None if
+      it is not pmapped.
+    **chunked_prediction_kwargs:
+      See chunked_prediction, some of these are required arguments.
+
+  Yields:
+    The predictions for each chunked step of the chunked rollout, such that
+    if all predictions are concatenated in time and sample dimension squeezed,
+    this would match the targets template in structure.
+
+  """
+  if pmap_devices is not None:
+    assert (
+        num_samples % len(pmap_devices) == 0
+    ), "num_samples must be a multiple of len(pmap_devices)"
+
+    def predictor_fn_pmap_named_args(rng, inputs, targets_template, forcings):
+      targets_template = _replicate_dataset(
+          targets_template,
+          replica_dim="sample",
+          replicate_to_device=True,
+          devices=pmap_devices,
+      )
+      return predictor_fn(rng, inputs, targets_template, forcings)
+
+    for i in range(0, num_samples, len(pmap_devices)):
+      sample_idx = slice(i, i + len(pmap_devices))
+      logging.info("Samples %s out of %s", sample_idx, num_samples)
+      logging.flush()
+      sample_group_rngs = _device_put_sharded(
+          rngs[sample_idx], pmap_devices, "sample")
+
+      if "sample" not in inputs.dims:
+        sample_inputs = inputs
+      else:
+        sample_inputs = inputs.isel(sample=sample_idx, drop=True)
+
+      sample_inputs = _replicate_dataset(
+          sample_inputs,
+          replica_dim="sample",
+          replicate_to_device=True,
+          devices=pmap_devices,
+      )
+
+      if forcings is not None:
+        if "sample" not in forcings.dims:
+          sample_forcings = forcings
+        else:
+          sample_forcings = forcings.isel(sample=sample_idx, drop=True)
+
+        # TODO(pricei): We are replicating the full forcings for all rollout
+        # timesteps here, rather than inside `predictor_fn_pmap_named_args` like
+        # the targets_template above, because the forcings are concatenated with
+        # the inputs which will already be replicated. We should refactor this
+        # so that chunked prediction is aware of whether it is being run with
+        # pmap, and if so do the replication and device_put only of the
+        # necessary timesteps, as part of the chunked prediction function.
+        sample_forcings = _replicate_dataset(
+            sample_forcings,
+            replica_dim="sample",
+            replicate_to_device=False,
+            devices=pmap_devices,
+        )
+      else:
+        sample_forcings = None
+
+      for prediction_chunk in chunked_prediction_generator(
+          predictor_fn=predictor_fn_pmap_named_args,
+          rng=sample_group_rngs,
+          inputs=sample_inputs,
+          targets_template=targets_template,
+          forcings=sample_forcings,
+          pmap_devices=pmap_devices,
+          replica_axis="sample",
+          **chunked_prediction_kwargs,
+      ):
+        prediction_chunk.coords["sample"] = np.arange(
+            sample_idx.start, sample_idx.stop, sample_idx.step
+        )
+        yield prediction_chunk
+        del prediction_chunk
+  else:
+    for i in range(num_samples):
+      logging.info("Sample %d/%d", i, num_samples)
+      logging.flush()
+      this_sample_rng = rngs[i]
+
+      if "sample" in inputs.dims:
+        sample_inputs = inputs.isel(sample=i, drop=True)
+      else:
+        sample_inputs = inputs
+
+      sample_forcings = forcings
+      if sample_forcings is not None:
+        if "sample" in sample_forcings.dims:
+          sample_forcings = sample_forcings.isel(sample=i, drop=True)
+
+      for prediction_chunk in chunked_prediction_generator(
+          predictor_fn=predictor_fn,
+          rng=this_sample_rng,
+          inputs=sample_inputs,
+          targets_template=targets_template,
+          forcings=sample_forcings,
+          **chunked_prediction_kwargs):
+        prediction_chunk.coords["sample"] = i
+        yield prediction_chunk
+        del prediction_chunk
+
+
+def chunked_prediction(
+    predictor_fn: PredictorFn,
+    rng: chex.PRNGKey,
+    inputs: xarray.Dataset,
+    targets_template: xarray.Dataset,
+    forcings: xarray.Dataset,
+    num_steps_per_chunk: int = 1,
+    verbose: bool = False,
+) -> xarray.Dataset:
+  """Outputs a long trajectory by iteratively concatenating chunked predictions.
+
+  Args:
+    predictor_fn: Function to use to make predictions for each chunk.
+    rng: Random key.
+    inputs: Inputs for the model.
+    targets_template: Template for the target prediction, requires targets
+        equispaced in time.
+    forcings: Optional forcing for the model.
+    num_steps_per_chunk: How many of the steps in `targets_template` to predict
+        at each call of `predictor_fn`. It must evenly divide the number of
+        steps in `targets_template`.
+    verbose: Whether to log the current chunk being predicted.
+
+  Returns:
+    Predictions for the targets template.
+
+  """
+  chunks_list = []
+  for prediction_chunk in chunked_prediction_generator(
+      predictor_fn=predictor_fn,
+      rng=rng,
+      inputs=inputs,
+      targets_template=targets_template,
+      forcings=forcings,
+      num_steps_per_chunk=num_steps_per_chunk,
+      verbose=verbose,
+  ):
+    chunks_list.append(jax.device_get(prediction_chunk))
+  return xarray.concat(chunks_list, dim="time")
+
+
+def chunked_prediction_generator(
+    predictor_fn: PredictorFn,
+    rng: chex.PRNGKey,
+    inputs: xarray.Dataset,
+    targets_template: xarray.Dataset,
+    forcings: xarray.Dataset,
+    num_steps_per_chunk: int = 1,
+    verbose: bool = False,
+    pmap_devices: Sequence[jax.Device] | None = None,
+    replica_axis: str | None = None,
+) -> Iterator[xarray.Dataset]:
+  """Outputs a long trajectory by yielding chunked predictions.
+
+  Args:
+    predictor_fn: Function to use to make predictions for each chunk.
+    rng: Random key.
+    inputs: Inputs for the model.
+    targets_template: Template for the target prediction, requires targets
+        equispaced in time.
+    forcings: Optional forcing for the model.
+    num_steps_per_chunk: How many of the steps in `targets_template` to predict
+        at each call of `predictor_fn`. It must evenly divide the number of
+        steps in `targets_template`.
+    verbose: Whether to log the current chunk being predicted.
+    pmap_devices: List of devices over which predictor_fn is pmapped, or None if
+      it is not pmapped.
+    replica_axis: Dimension name to use for the replicas.
+
+  Yields:
+    The predictions for each chunked step of the chunked rollout, such as
+    if all predictions are concatenated in time this would match the targets
+    template in structure.
+
+  """
+
+  if pmap_devices is not None and replica_axis is None:
+    raise ValueError("Must provide replica_axis when pmap_devices is provided.")
+
+  # Create copies to avoid mutating inputs.
+  inputs = inputs.copy()
+  targets_template = targets_template.copy()
+  forcings = forcings.copy()
+
+  if "datetime" in inputs.coords:
+    del inputs.coords["datetime"]
+
+  if "datetime" in targets_template.coords:
+    output_datetime = targets_template.coords["datetime"]
+    del targets_template.coords["datetime"]
+  else:
+    output_datetime = None
+
+  if "datetime" in forcings.coords:
+    del forcings.coords["datetime"]
+
+  num_target_steps = targets_template.dims["time"]
+  num_chunks, remainder = divmod(num_target_steps, num_steps_per_chunk)
+  if remainder != 0:
+    raise ValueError(
+        f"The number of steps per chunk {num_steps_per_chunk} must "
+        f"evenly divide the number of target steps {num_target_steps} ")
+
+  if len(np.unique(np.diff(targets_template.coords["time"].data))) > 1:
+    raise ValueError("The targets time coordinates must be evenly spaced")
+
+  # Our template targets will always have a time axis corresponding for the
+  # timedeltas for the first chunk.
+  targets_chunk_time = targets_template.time.isel(
+      time=slice(0, num_steps_per_chunk))
+
+  current_inputs = inputs
+
+  def split_rng_fn(rng):
+    # Note, this is *not* equivalent to `return jax.random.split(rng)`, because
+    # by assigning to a tuple, the single numpy array returned by
+    # `jax.random.split` actually gets split into two arrays, so when calling
+    # the function with pmap the output is Tuple[Array, Array], where the
+    # leading axis of each array is `num devices`.
+    rng1, rng2 = jax.random.split(rng)
+    return rng1, rng2
+
+  if pmap_devices is not None:
+    split_rng_fn = jax.pmap(
+        split_rng_fn, devices=pmap_devices, axis_name=replica_axis
+    )
+
+  for chunk_index in range(num_chunks):
+    if verbose:
+      logging.info("Chunk %d/%d", chunk_index, num_chunks)
+      logging.flush()
+
+    # Select targets for the time period that we are predicting for this chunk.
+    target_offset = num_steps_per_chunk * chunk_index
+    target_slice = slice(target_offset, target_offset + num_steps_per_chunk)
+    current_targets_template = targets_template.isel(time=target_slice)
+
+    # Replace the timedelta, by the one corresponding to the first chunk, so we
+    # don't recompile at every iteration, keeping the
+    actual_target_time = current_targets_template.coords["time"]
+    current_targets_template = current_targets_template.assign_coords(
+        time=targets_chunk_time).compute()
+
+    current_forcings = forcings.isel(time=target_slice)
+    current_forcings = current_forcings.assign_coords(time=targets_chunk_time)
+    current_forcings = current_forcings.compute()
+    # Make predictions for the chunk.
+    rng, this_rng = split_rng_fn(rng)
+    predictions = predictor_fn(
+        rng=this_rng,
+        inputs=current_inputs,
+        targets_template=current_targets_template,
+        forcings=current_forcings)
+
+    # In the pmapped case, profiling reveals that the predictions, forcings and
+    # inputs are all copied onto a single TPU, causing OOM. To avoid this
+    # we pull all of the input/output data off the devices. This will have
+    # some performance impact, but maximise the memory efficiency.
+    # TODO(aelkadi): Pmap `_get_next_inputs` when running under pmap, and
+    # remove the device_get.
+    if pmap_devices is not None:
+      predictions = jax.device_get(predictions)
+      current_forcings = jax.device_get(current_forcings)
+      current_inputs = jax.device_get(current_inputs)
+
+    if chunk_index == num_chunks - 1:
+      # No need to call `_get_next_inputs` on the last iteration.
+      current_inputs = None
+    else:
+      next_frame = xarray.merge([predictions, current_forcings])
+      next_inputs = _get_next_inputs(current_inputs, next_frame)
+      # Shift timedelta coordinates, so we don't recompile at every iteration.
+      next_inputs = next_inputs.assign_coords(
+          time=current_inputs.coords["time"])
+      current_inputs = next_inputs
+
+    # At this point we can assign the actual targets time coordinates.
+    predictions = predictions.assign_coords(time=actual_target_time)
+    if output_datetime is not None:
+      predictions.coords["datetime"] = output_datetime.isel(
+          time=target_slice)
+    yield predictions
+    del predictions
+
+
+def _get_next_inputs(
+    prev_inputs: xarray.Dataset, next_frame: xarray.Dataset,
+    ) -> xarray.Dataset:
+  """Computes next inputs, from previous inputs and predictions."""
+
+  # Make sure are are predicting all inputs with a time axis.
+  non_predicted_or_forced_inputs = list(
+      set(prev_inputs.keys()) - set(next_frame.keys()))
+  if "time" in prev_inputs[non_predicted_or_forced_inputs].dims:
+    raise ValueError(
+        "Found an input with a time index that is not predicted or forced.")
+
+  # Keys we need to copy from predictions to inputs.
+  next_inputs_keys = list(
+      set(next_frame.keys()).intersection(set(prev_inputs.keys())))
+  next_inputs = next_frame[next_inputs_keys]
+
+  # Apply concatenate next frame with inputs, crop what we don't need.
+  num_inputs = prev_inputs.dims["time"]
+  return (
+      xarray.concat(
+          [prev_inputs, next_inputs], dim="time", data_vars="different")
+      .tail(time=num_inputs))
+
+
+def extend_targets_template(
+    targets_template: xarray.Dataset,
+    required_num_steps: int) -> xarray.Dataset:
+  """Extends `targets_template` to `required_num_steps` with lazy arrays.
+
+  It uses lazy dask arrays of zeros, so it does not require instantiating the
+  array in memory.
+
+  Args:
+    targets_template: Input template to extend.
+    required_num_steps: Number of steps required in the returned template.
+
+  Returns:
+    `xarray.Dataset` identical in variables and timestep to `targets_template`
+    full of `dask.array.zeros` such that the time axis has `required_num_steps`.
+
+  """
+
+  # Extend the "time" and "datetime" coordinates
+  time = targets_template.coords["time"]
+
+  # Assert the first target time corresponds to the timestep.
+  timestep = time[0].data
+  if time.shape[0] > 1:
+    assert np.all(timestep == time[1:] - time[:-1])
+
+  extended_time = (np.arange(required_num_steps) + 1) * timestep
+
+  if "datetime" in targets_template.coords:
+    datetime = targets_template.coords["datetime"]
+    extended_datetime = (datetime[0].data - timestep) + extended_time
+  else:
+    extended_datetime = None
+
+  # Replace the values with empty dask arrays extending the time coordinates.
+  datetime = targets_template.coords["time"]
+
+  def extend_time(data_array: xarray.DataArray) -> xarray.DataArray:
+    dims = data_array.dims
+    shape = list(data_array.shape)
+    shape[dims.index("time")] = required_num_steps
+    dask_data = dask.array.zeros(
+        shape=tuple(shape),
+        chunks=-1,  # Will give chunk info directly to `ChunksToZarr``.
+        dtype=data_array.dtype)
+
+    coords = dict(data_array.coords)
+    coords["time"] = extended_time
+
+    if extended_datetime is not None:
+      coords["datetime"] = ("time", extended_datetime)
+
+    return xarray.DataArray(
+        dims=dims,
+        data=dask_data,
+        coords=coords)
+
+  return xarray_tree.map_structure(extend_time, targets_template)

model/graphcast/samplers_base.py

@@ -0,0 +1,47 @@
+# Copyright 2024 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Base class for diffusion samplers."""
+
+import abc
+from typing import Optional
+
+from . import denoisers_base
+import xarray
+
+
+class Sampler(abc.ABC):
+  """A sampling algorithm for a denoising diffusion model.
+
+  This is constructed with a denoising function, and uses it to draw samples.
+  """
+
+  _denoiser: denoisers_base.Denoiser
+
+  def __init__(self, denoiser: denoisers_base.Denoiser):
+    """Constructs Sampler.
+
+    Args:
+      denoiser: A Denoiser which has been trained with an MSE loss to predict
+        the noise-free targets.
+    """
+    self._denoiser = denoiser
+
+  @abc.abstractmethod
+  def __call__(
+      self,
+      inputs: xarray.Dataset,
+      targets_template: xarray.Dataset,
+      forcings: Optional[xarray.Dataset] = None,
+      **kwargs) -> xarray.Dataset:
+    """Draws a sample using self._denoiser. Contract like Predictor.__call__."""

model/graphcast/samplers_utils.py

@@ -0,0 +1,436 @@
+# Copyright 2024 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Utils for diffusion samplers. Makes use of dinosaur.spherical_harmonic."""
+
+import dataclasses
+import functools
+from typing import Any, cast, Optional, Tuple
+
+import chex
+from dinosaur import spherical_harmonic
+from . import xarray_jax
+from . import xarray_tree
+import haiku as hk
+import jax
+import jax.numpy as jnp
+import numpy as np
+import xarray
+
+# Some useful constants useful when dealing with Earth's geometry.
+# The earth isn't really a sphere so these are only approximate, this is the
+# average radius according to https://en.wikipedia.org/wiki/Earth_radius,
+# with the actual value varying from 6378 to 6357km.
+EARTH_RADIUS_KM = 6371.
+# And this is also approximate, but we've chosen to make it consistent with the
+# radius above when modelling the earth as a sphere. This gives a value of
+# around 40030; the actual value varies from 40008 to 40075.
+EARTH_CIRCUMFERENCE_KM = EARTH_RADIUS_KM * 2 * np.pi
+
+
+@dataclasses.dataclass(frozen=True)
+class _ArrayGrid:
+  """A class that performs operations and transformations in the spectral basis.
+
+  Attributes:
+    longitude_wavenumbers: num of longitude wavenumbers in the spectral basis.
+    total_wavenumbers: number of total wavenumbers in the spectral basis.
+    longitude_nodes: number of quadrature nodes along the lon direction.
+    latitude_nodes: number of quadrature nodes along the lat direction.
+    latitude_spacing: either 'gauss' or 'equiangular'. This determines the
+      spacing of nodal grid points in the latitudinal (north-south) direction.
+  """
+  longitude_wavenumbers: int
+  total_wavenumbers: int
+  longitude_nodes: int
+  latitude_nodes: int
+  latitude_spacing: str
+
+  @classmethod
+  def with_lat_lon(
+      cls,
+      lat: np.ndarray,
+      lon: np.ndarray,
+      ) -> '_ArrayGrid':
+    """_ArrayGrid for use with data in specified lat/lon grid (in degrees)."""
+
+    latitude_nodes = lat.shape[0]
+    longitude_nodes = lon.shape[0]
+    latitude_spacing = _infer_latitude_spacing(lat)
+    if latitude_spacing in ['equiangular', 'gauss']:
+      if longitude_nodes != 2 * latitude_nodes:
+        # Technically not a requirement but useful to ensure `max_wavenumber`
+        # makes sense.
+        raise ValueError(
+            'Unexpected number of longitude nodes. '
+            f'Expected {2 * latitude_nodes}, got {longitude_nodes}')
+    elif latitude_spacing == 'equiangular_with_poles':
+      if longitude_nodes != 2 * (latitude_nodes - 1):
+        # Technically not a requirement but useful to ensure `max_wavenumber`
+        # makes s
+        raise ValueError(
+            'Unexpected number of longitude nodes. '
+            f'Expected {2 * (latitude_nodes - 1)}, got {longitude_nodes}')
+    else:
+      raise ValueError(f'Unexpected latitude_spacing={latitude_spacing}')
+    max_wavenumber = int(longitude_nodes // 2) - 1
+    grid = cls(
+        longitude_wavenumbers=max_wavenumber+1,
+        # total_wavenumbers should be one larger than max_wavenumber as the
+        # wavenumbers go from 0 to max_wavenumber inclusive.
+        total_wavenumbers=max_wavenumber+1,
+        longitude_nodes=longitude_nodes,
+        latitude_nodes=latitude_nodes,
+        latitude_spacing=latitude_spacing,
+    )
+    _verify_nodal_axes(lat, lon, grid.nodal_axes)
+    return grid
+
+  @functools.cached_property
+  def _grid(self) -> spherical_harmonic.Grid:
+    return spherical_harmonic.Grid(
+        spherical_harmonics_impl=spherical_harmonic.RealSphericalHarmonics,
+        **dataclasses.asdict(self),
+    )
+
+  @functools.cached_property
+  def nodal_axes(self) -> Tuple[np.ndarray, np.ndarray]:
+    """Longitude and sin(latitude) coordinates of the nodal basis."""
+    return self._grid.nodal_axes
+
+  @functools.cached_property
+  def modal_axes(self) -> Tuple[np.ndarray, np.ndarray]:
+    """Longitudinal and total wavenumbers (m, l) of the modal basis."""
+    return self._grid.modal_axes
+
+  def to_nodal(self, x: chex.Array) -> chex.Array:
+    """Maps `x` from a modal to nodal representation."""
+    return self._grid.to_nodal(x)
+
+
+def _infer_latitude_spacing(lat: np.ndarray) -> str:
+  """Infers the type of latitude spacing given the latitude."""
+  if not np.all(np.diff(lat) > 0.):
+    raise ValueError('Latitude values are expected to be sorted.')
+
+  if np.allclose(np.diff(lat), lat[1] - lat[0]):
+    if np.isclose(max(lat), 90.):
+      spacing = 'equiangular_with_poles'
+    else:
+      spacing = 'equiangular'
+  else:
+    spacing = 'gauss'
+  return spacing
+
+
+def _verify_nodal_axes(lat_coords: np.ndarray, lon_coords: np.ndarray,
+                       nodal_axes: Tuple[np.ndarray, np.ndarray]):
+  nodal_axes_lon, nodal_axes_sin_lat = nodal_axes
+  if not np.allclose(nodal_axes_sin_lat, np.sin(np.deg2rad(lat_coords))):
+    raise ValueError(
+        "Latitude coords don't match those used by "
+        "spherical_harmonic.SphericalHarmonicBasis.")
+  if not np.allclose(nodal_axes_lon, np.deg2rad(lon_coords)):
+    raise ValueError(
+        "Longitude coords don't match those used by "
+        "spherical_harmonic.SphericalHarmonicBasis.")
+
+
+class Grid:
+  """xarray wrapper around _ArrayGrid."""
+
+  @classmethod
+  def for_nodal_data(
+      cls,
+      nodal_data: xarray.DataArray,
+      ) -> 'Grid':
+    """A Grid for use with a given shape of nodal (lat/lon grid) data.
+
+    This uses the maximum number of spherical harmonics that the grid is able
+    to resolve.
+
+    This class supports data arrays with latitude spacings as defined by
+    "dinosaur.spherical_harmonic". In summary:
+    * 'equiangular': equally spaced (by `d_lat`) values between -90 + d_lat and
+      90 - d_lat / 2. In our case, longitude must also be spaced by `d_lat`.
+    * 'equiangular_with_poles': equally spaced (by `d_lat`) values between -90
+      and 90. In our case, longitude must also be spaced by `d_lat`.
+    * 'gauss': Gauss-Legendre nodes.
+
+    Args:
+      nodal_data: An xarray with 'lat' and 'lon' dimensions and coordinates in
+        degrees.
+
+    Returns:
+      A grid with the specified latitude_nodes, with
+      longitude_nodes=2*latitude_nodes and max_wavenumber=latitude_nodes-1.
+    """
+
+    grid = _ArrayGrid.with_lat_lon(
+        nodal_data.coords['lat'].data,
+        nodal_data.coords['lon'].data)
+    return cls(grid,
+               nodal_data.coords['lat'].data,
+               nodal_data.coords['lon'].data)
+
+  def __init__(self,
+               grid: _ArrayGrid,
+               lat_coords: np.ndarray,
+               lon_coords: np.ndarray):
+    _verify_nodal_axes(lat_coords, lon_coords, grid.nodal_axes)
+    self._underlying = grid
+    # Record the exact original lat/lon coords so we can return them exactly
+    # from an inverse transform, avoiding any xarray merge issues if coordinates
+    # are off by a rounding error.
+    self._lat_coords = lat_coords
+    self._lon_coords = lon_coords
+    self._longitude_wavenumber_coords, self._total_wavenumber_coords = (
+        grid.modal_axes)
+
+  @property
+  def total_wavenumber_coords(self) -> xarray.DataArray:
+    """Coords that must be used for 'total_wavenumber' dimension."""
+    return xarray.DataArray(
+        data=self._total_wavenumber_coords,
+        dims=('total_wavenumber',),
+        coords={'total_wavenumber': self._total_wavenumber_coords})
+
+  @property
+  def longitude_wavenumber_coords(self) -> xarray.DataArray:
+    """Coords that must be used for 'longitude_wavenumber' dimension."""
+    return xarray.DataArray(
+        data=self._longitude_wavenumber_coords,
+        dims=('longitude_wavenumber',),
+        coords={'longitude_wavenumber': self._longitude_wavenumber_coords})
+
+  def to_nodal(
+            self, modal_data: xarray.DataArray) -> xarray.DataArray:
+    """Applies the inverse spherical harmonic transform.
+
+    Args:
+      modal_data: A tree of xarray.DataArray with 'longitude_wavenumber' and
+        'total_wavenumber' dimensions with coords
+        `self.longitude_wavenumber_coords` and `self.total_wavenumber_coords`
+        respectively, and with the same sparsity pattern described under
+        `to_modal`.
+
+    Returns:
+      Corresponding tree where the 'longitude_wavenumber' and
+      'total_wavenumber' dimensions are replaced by 'lat', 'lon' dimensions.
+    """
+    def inverse_transform(modal: xarray.DataArray) -> xarray.DataArray:
+      if (not np.all(modal.coords['longitude_wavenumber'] ==
+                     self._longitude_wavenumber_coords) or
+          not np.all(modal.coords['total_wavenumber'] ==
+                     self._total_wavenumber_coords)):
+        raise ValueError('Wavenumber coords don\'t follow required convention.')
+
+      return xarray_jax.apply_ufunc(
+          self._underlying.to_nodal, modal,
+          input_core_dims=[['longitude_wavenumber', 'total_wavenumber']],
+          output_core_dims=[['lon', 'lat']],
+      ).assign_coords(
+          lon=self._lon_coords,
+          lat=self._lat_coords,
+      )
+
+    return xarray_tree.map_structure(inverse_transform, modal_data)
+
+
+def sample(
+    key: jnp.ndarray,
+    power_spectrum: xarray.DataArray,
+    template: xarray.DataArray,
+    grid: Optional[Grid] = None,
+    ) -> xarray.DataArray:
+  """Samples Gaussian Process noise on a sphere, with a given power spectrum.
+
+  This means the noise will have the given power spectrum *in expectation*; the
+  power spectrum of individual samples may vary.
+
+  The noise will be isotropic, meaning the distribution is invariant to
+  rotations of the sphere.
+
+  The marginal variance of the returned values will be equal to the total power,
+  i.e. the sum of power_spectrum. So if you want unit marginal variance, just
+  make sure to normalize the power_spectrum to sum to 1.
+
+  Args:
+    key: JAX rng key.
+    power_spectrum: An array with shape (total_wavenumber,) giving the power
+      which is desired at each total wavenumber (corresponding to a wavelength
+      EARTH_CIRCUMFERENCE/total_wavenumber) for total wavenumbers 0 up to some
+      maximum. This is in squared units of the quantity being sampled.
+    template: An array with the shape that you want the samples in, containing
+      'lat' and 'lon' dimensions. If other dimensions are present, we draw
+      multiple independent samples along these other dimensions.
+    grid: spherical_harmonic.Grid on which to sample the noise. If not specified
+      a grid will be created based on `template`, however note you may save some
+      RAM and compute by re-using a single Grid instance across multiple calls.
+
+  Returns:
+    DataArray with the same shape as template.
+  """
+  if grid is None:
+    grid = Grid.for_nodal_data(template)
+  dims = [d for d in template.dims if d not in ('lat', 'lon')]
+  shape = [template.sizes[d] for d in dims]
+  coords = {name: coord for name, coord in template.coords.items()
+            if name not in ('lat', 'lon')}
+  dims.extend(('total_wavenumber', 'longitude_wavenumber'))
+  shape.extend((len(grid.total_wavenumber_coords),
+                len(grid.longitude_wavenumber_coords)))
+  coords.update({'total_wavenumber': grid.total_wavenumber_coords,
+                 'longitude_wavenumber': grid.longitude_wavenumber_coords})
+  coeffs = xarray_jax.DataArray(
+      data=jax.random.normal(key, shape), dims=dims, coords=coords)
+  # Mask out coefficients which are out of range. This broadcasts to a
+  # triangular mask with shape (total_wavenumber, longitude_wavenumber):
+  mask = (
+      abs(coeffs.longitude_wavenumber) <= coeffs.total_wavenumber
+      ).astype(np.float32)
+  # For total_wavenumber t, there will be 2t+1 non-zero coefficients at
+  # different longitude_wavenumbers. We must normalize the coefficients so that
+  # summing their squares at each total_wavenumber, sums to the corresponding
+  # value in the power spectrum:
+  multiplier = mask * np.sqrt(power_spectrum / mask.sum(
+      'longitude_wavenumber', skipna=False))
+  # And a standard normalization factor used in this implementation of the
+  # spherical harmonic transform:
+  multiplier *= np.sqrt(4 * np.pi)
+  # Only finally multiply by coeffs to avoid too many broadcasting
+  # multiplications:
+  coeffs *= multiplier
+  result = cast(xarray.DataArray, grid.to_nodal(coeffs))
+  result = result.astype(template.dtype)
+  return result.transpose(*template.dims)
+
+
+def spherical_white_noise_like(template: xarray.Dataset) -> xarray.Dataset:
+  """Samples isotropic mean 0 variance 1 white noise on the sphere."""
+  def spherical_white_noise_like_dataarray(data_array: xarray.DataArray
+                                           ) -> xarray.DataArray:
+    num_wavenumbers = data_array.lon.shape[0] // 2
+    key = hk.next_rng_key()
+    return sample(
+        key=key,
+        power_spectrum=xarray_jax.DataArray(
+            data=np.array([1/num_wavenumbers for _ in range(num_wavenumbers)]),
+            dims=['total_wavenumber']),
+        template=data_array)
+  return template.map(spherical_white_noise_like_dataarray)
+
+
+def rho_inverse_cdf(
+    min_value: float,
+    max_value: float,
+    rho: float,
+    cdf: Any) -> Any:
+  """Quantiles of rho distribution used for noise levels at sampling time.
+
+  This is parameterised by rho as in Eqn 5 from the Elucidating paper
+  (but with max/min flipped so that quantiles are given in ascending not
+  descending order). It's equivalent to a Beta[rho, 1] distribution rescaled to
+  [min_value, max_value].
+
+  At sampling time we use noise levels at fixed quantiles of this distribution.
+  Unlike in the paper, we also use the same distribution for noise levels at
+  training time (albeit potentially with different parameters, and sampling from
+  it at random).
+
+  Args:
+    min_value:
+    max_value:
+      Define the support of the distribution.
+    rho:
+      Shape parameter.
+    cdf:
+      Value or values between 0 and 1 indicating which quantile you want. Can
+      be a numpy or jax array.
+
+  Returns:
+    Quantiles of the distribution, with same shape/type as `cdf`.
+  """
+  return (
+      min_value**(1 / rho) + cdf *
+      (max_value**(1 / rho) - min_value**(1 / rho))
+  )**rho
+
+
+def tree_where(
+    cond: jnp.ndarray,
+    xs: Any,
+    ys: Any
+    ) -> Any:
+  """Like jnp.where but works with trees for xs and ys (but not for cond)."""
+  return jax.tree_util.tree_map(lambda x, y: jnp.where(cond, x, y), xs, ys)
+
+
+def noise_schedule(
+    max_noise_level: float = 80.,
+    min_noise_level: float = 0.002,
+    num_noise_levels: int = 30,
+    rho: float = 7.,
+) -> np.ndarray:
+  """Computes a descending noise schedule for sampling, ending with zero."""
+  noise_levels = rho_inverse_cdf(
+      min_value=min_noise_level,
+      max_value=max_noise_level,
+      rho=rho,
+      # We want the noise levels in descending order, so ask for quantiles
+      # 1 down to 0:
+      cdf=np.linspace(1, 0, num_noise_levels))
+  # The final zero noise level is somewhat special-cased. We don't actually
+  # denoise from this noise level but appending it here is convenient for
+  # sampling loop implementations.
+  return np.append(noise_levels, 0.)
+
+
+def stochastic_churn_rate_schedule(
+    noise_levels: np.ndarray,
+    stochastic_churn_rate: float = 0.,
+    churn_min_noise_level: float = 0.05,
+    churn_max_noise_level: float = 50.0,
+) -> np.ndarray:
+  """Computes a stochastic churn rate for each noise level."""
+  num_noise_levels = len(noise_levels)-1  # Exclude final zero noise level.
+  # As in the Elucidated Diffusion paper, clamp this so it doesn't increase the
+  # variance by a factor of more than 2, no matter how few noise levels are
+  # used:
+  per_step_churn_rate = min(stochastic_churn_rate / num_noise_levels,
+                            np.sqrt(2) - 1)
+  return (
+      (churn_min_noise_level <= noise_levels[:-1]) &
+      (noise_levels[:-1] <= churn_max_noise_level)
+  ) * per_step_churn_rate
+
+
+def apply_stochastic_churn(
+    x: Any,
+    noise_level: jax.typing.ArrayLike,
+    stochastic_churn_rate: jax.typing.ArrayLike,
+    noise_level_inflation_factor: jax.typing.ArrayLike,
+) -> tuple[Any, jax.typing.ArrayLike]:
+  """Returns x at higher noise level, and the higher noise level itself."""
+  # We increase the noise level of x a bit before taking it down again:
+  new_noise_level = noise_level * (1.0 + stochastic_churn_rate)
+  noise_diff = new_noise_level**2 - noise_level**2
+  # stochastic_churn_rate == 0 => new_noise_level == noise_level
+  # => noise_diff == 0. This can resolve to a negative value because of
+  # floating point rounding errors. To avoid this we clamp noise_diff to zero if
+  # it's negative.
+  noise_diff = jnp.maximum(noise_diff, 0)
+  extra_noise_stddev = jnp.sqrt(noise_diff)* noise_level_inflation_factor
+  updated_x = x + spherical_white_noise_like(x) * extra_noise_stddev
+  return updated_x, new_noise_level
+

model/graphcast/solar_radiation.py

@@ -0,0 +1,605 @@
+# Copyright 2023 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Computes TOA incident solar radiation compatible with ERA5.
+
+The Top-Of-the-Atmosphere (TOA) incident solar radiation is available in the
+ERA5 dataset as the parameter `toa_incident_solar_radiation` (or `tisr`). This
+represents the TOA solar radiation flux integrated over a period of one hour
+ending at the timestamp given by the `datetime` coordinate. See
+https://confluence.ecmwf.int/display/CKB/ERA5%3A+data+documentation and
+https://codes.ecmwf.int/grib/param-db/?id=212.
+"""
+
+from collections.abc import Callable, Sequence
+import dataclasses
+import functools
+
+import chex
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pandas as pd
+import xarray as xa
+
+
+# Default value of the `integration_period` argument to be compatible with ERA5.
+_DEFAULT_INTEGRATION_PERIOD = pd.Timedelta(hours=1)
+
+# Default value for the `num_integration_bins` argument. This provides a good
+# approximation of the solar radiation in ERA5.
+_DEFAULT_NUM_INTEGRATION_BINS = 360
+
+# The length of a Julian year in days.
+# https://en.wikipedia.org/wiki/Julian_year_(astronomy)
+_JULIAN_YEAR_LENGTH_IN_DAYS = 365.25
+
+# Julian Date for the J2000 epoch, a standard reference used in astronomy.
+# https://en.wikipedia.org/wiki/Epoch_(astronomy)#Julian_years_and_J2000
+_J2000_EPOCH = 2451545.0
+
+# Number of seconds in a day.
+_SECONDS_PER_DAY = 60 * 60 * 24
+
+
+_TimestampLike = str | pd.Timestamp | np.datetime64
+_TimedeltaLike = str | pd.Timedelta | np.timedelta64
+
+
+# Interface for loading Total Solar Irradiance (TSI) data.
+# Returns a xa.DataArray containing yearly average TSI values with a `time`
+# coordinate in units of years since 0000-1-1. E.g. 2023.5 corresponds to
+# the middle of the year 2023.
+TsiDataLoader = Callable[[], xa.DataArray]
+
+
+# Total Solar Irradiance (TSI): Energy input to the top of the Earth's
+# atmosphere in W⋅m⁻². TSI varies with time. This is the reference TSI value
+# that can be used when more accurate data is not available.
+# https://www.ncei.noaa.gov/products/climate-data-records/total-solar-irradiance
+# https://github.com/ecmwf-ifs/ecrad/blob/6db82f929fb75028cc20606a04da87c0abe9b642/radiation/radiation_ecckd.F90#L296
+_REFERENCE_TSI = 1361.0
+
+
+def reference_tsi_data() -> xa.DataArray:
+  """A TsiDataProvider that returns a single reference TSI value."""
+  return xa.DataArray(
+      np.array([_REFERENCE_TSI]),
+      dims=["time"],
+      coords={"time": np.array([0.0])},
+  )
+
+
+def era5_tsi_data() -> xa.DataArray:
+  """A TsiDataProvider that returns ERA5 compatible TSI data."""
+  # ECMWF provided the data used for ERA5, which was hardcoded in the IFS (cycle
+  # 41r2). The values were scaled down to agree better with more recent
+  # observations of the sun.
+  time = np.arange(1951.5, 2035.5, 1.0)
+  tsi = 0.9965 * np.array([
+      # fmt: off
+      # 1951-1995 (non-repeating sequence)
+      1365.7765, 1365.7676, 1365.6284, 1365.6564, 1365.7773,
+      1366.3109, 1366.6681, 1366.6328, 1366.3828, 1366.2767,
+      1365.9199, 1365.7484, 1365.6963, 1365.6976, 1365.7341,
+      1365.9178, 1366.1143, 1366.1644, 1366.2476, 1366.2426,
+      1365.9580, 1366.0525, 1365.7991, 1365.7271, 1365.5345,
+      1365.6453, 1365.8331, 1366.2747, 1366.6348, 1366.6482,
+      1366.6951, 1366.2859, 1366.1992, 1365.8103, 1365.6416,
+      1365.6379, 1365.7899, 1366.0826, 1366.6479, 1366.5533,
+      1366.4457, 1366.3021, 1366.0286, 1365.7971, 1365.6996,
+      # 1996-2008 (13 year cycle, repeated below)
+      1365.6121, 1365.7399, 1366.1021, 1366.3851, 1366.6836,
+      1366.6022, 1366.6807, 1366.2300, 1366.0480, 1365.8545,
+      1365.8107, 1365.7240, 1365.6918,
+      # 2009-2021
+      1365.6121, 1365.7399, 1366.1021, 1366.3851, 1366.6836,
+      1366.6022, 1366.6807, 1366.2300, 1366.0480, 1365.8545,
+      1365.8107, 1365.7240, 1365.6918,
+      # 2022-2034
+      1365.6121, 1365.7399, 1366.1021, 1366.3851, 1366.6836,
+      1366.6022, 1366.6807, 1366.2300, 1366.0480, 1365.8545,
+      1365.8107, 1365.7240, 1365.6918,
+      # fmt: on
+  ])
+  return xa.DataArray(tsi, dims=["time"], coords={"time": time})
+
+
+# HRES compatible TSI data is from IFS cycle 47r1. The dataset can be obtained
+# from the ECRAD package: https://confluence.ecmwf.int/display/ECRAD.
+# The example code below can load this dataset from a local file.
+
+# def hres_tsi_data() -> xa.DataArray:
+#   with open("total_solar_irradiance_CMIP6_47r1.nc", "rb") as f:
+#     with xa.load_dataset(f, decode_times=False) as ds:
+#       return ds["tsi"]
+
+
+_DEFAULT_TSI_DATA_LOADER: TsiDataLoader = era5_tsi_data
+
+
+def get_tsi(
+    timestamps: Sequence[_TimestampLike], tsi_data: xa.DataArray
+) -> chex.Array:
+  """Returns TSI values for the given timestamps.
+
+  TSI values are interpolated from the provided yearly TSI data.
+
+  Args:
+    timestamps: Timestamps for which to compute TSI values.
+    tsi_data: A DataArray with a single dimension `time` that has coordinates in
+      units of years since 0000-1-1. E.g. 2023.5 corresponds to the middle of
+      the year 2023.
+
+  Returns:
+    An Array containing interpolated TSI data.
+  """
+  timestamps = pd.DatetimeIndex(timestamps)
+  timestamps_date = pd.DatetimeIndex(timestamps.date)
+  day_fraction = (timestamps - timestamps_date) / pd.Timedelta(days=1)
+  year_length = 365 + timestamps.is_leap_year
+  year_fraction = (timestamps.dayofyear - 1 + day_fraction) / year_length
+  fractional_year = timestamps.year + year_fraction
+  return np.interp(fractional_year, tsi_data.coords["time"].data, tsi_data.data)
+
+
+@dataclasses.dataclass(frozen=True)
+class _OrbitalParameters:
+  """Parameters characterising Earth's position relative to the Sun.
+
+  The parameters characterize the position of the Earth in its orbit around the
+  Sun for specific points in time. Each attribute is an N-dimensional array
+  to represent orbital parameters for multiple points in time.
+
+  Attributes:
+    theta: The number of Julian years since the Julian epoch J2000.0.
+    rotational_phase: The phase of the Earth's rotation along its axis as a
+      ratio with 0 representing the phase at Julian epoch J2000.0 at exactly
+      12:00 Terrestrial Time (TT). Multiplying this value by `2*pi` yields the
+        phase in radians.
+    sin_declination: Sine of the declination of the Sun as seen from the Earth.
+    cos_declination: Cosine of the declination of the Sun as seen from the
+      Earth.
+    eq_of_time_seconds: The value of the equation of time, in seconds.
+    solar_distance_au: Earth-Sun distance in astronomical units.
+  """
+
+  theta: chex.Array
+  rotational_phase: chex.Array
+  sin_declination: chex.Array
+  cos_declination: chex.Array
+  eq_of_time_seconds: chex.Array
+  solar_distance_au: chex.Array
+
+
+def _get_j2000_days(timestamp: pd.Timestamp) -> float:
+  """Returns the number of days since the J2000 epoch.
+
+  Args:
+    timestamp: A timestamp for which to compute the J2000 days.
+
+  Returns:
+    The J2000 days corresponding to the input timestamp.
+  """
+  return timestamp.to_julian_date() - _J2000_EPOCH
+
+
+def _get_orbital_parameters(j2000_days: chex.Array) -> _OrbitalParameters:
+  """Computes the orbital parameters for the given J2000 days.
+
+  Args:
+    j2000_days: Timestamps represented as the number of days since the J2000
+      epoch.
+
+  Returns:
+    Orbital parameters for the given timestamps. Each attribute of the return
+    value is an array containing the same dimensions as the input.
+  """
+  # Orbital parameters are computed based on the formulas in this code, which
+  # were determined empirically to produce radiation values similar to ERA5:
+  # https://github.com/ECCC-ASTD-MRD/gem/blob/1d711f7b89971cd7b1e10afc7508d1135b51397d/src/rpnphy/src/base/sucst.F90
+  # https://github.com/ECCC-ASTD-MRD/gem/blob/1d711f7b89971cd7b1e10afc7508d1135b51397d/src/rpnphy/src/base/fctast.cdk
+  # https://github.com/ECCC-ASTD-MRD/gem/blob/1d711f7b89971cd7b1e10afc7508d1135b51397d/src/rpnphy/src/base/fcttim.cdk
+  # There are many variations to these formulas, but since the goal is to match
+  # the values in ERA5, the formulas were implemented as is. Comments reference
+  # the notation used in those sources. Here are some additional references
+  # related to the quantities being computed here:
+  # https://aa.usno.navy.mil/faq/sun_approx
+  # https://en.wikipedia.org/wiki/Position_of_the_Sun
+  # https://en.wikipedia.org/wiki/Equation_of_time
+
+  # Number of Julian years since the J2000 epoch (including fractional years).
+  theta = j2000_days / _JULIAN_YEAR_LENGTH_IN_DAYS
+  # The phase of the Earth's rotation along its axis as a ratio. 0 represents
+  # Julian epoch J2000.0 at exactly 12:00 Terrestrial Time (TT).
+  rotational_phase = j2000_days % 1.0
+
+  # REL(PTETA).
+  rel = 1.7535 + 6.283076 * theta
+  # REM(PTETA).
+  rem = 6.240041 + 6.283020 * theta
+  # RLLS(PTETA).
+  rlls = 4.8951 + 6.283076 * theta
+
+  # Variables used in the three polynomials below.
+  one = jnp.ones_like(theta)
+  sin_rel = jnp.sin(rel)
+  cos_rel = jnp.cos(rel)
+  sin_two_rel = jnp.sin(2.0 * rel)
+  cos_two_rel = jnp.cos(2.0 * rel)
+  sin_two_rlls = jnp.sin(2.0 * rlls)
+  cos_two_rlls = jnp.cos(2.0 * rlls)
+  sin_four_rlls = jnp.sin(4.0 * rlls)
+  sin_rem = jnp.sin(rem)
+  sin_two_rem = jnp.sin(2.0 * rem)
+
+  # Ecliptic longitude of the Sun - RLLLS(PTETA).
+  rllls = jnp.dot(
+      jnp.stack(
+          [one, theta, sin_rel, cos_rel, sin_two_rel, cos_two_rel], axis=-1
+      ),
+      jnp.array([4.8952, 6.283320, -0.0075, -0.0326, -0.0003, 0.0002]),
+  )
+
+  # Angle in radians between the Earth's rotational axis and its orbital axis.
+  # Equivalent to 23.4393°.
+  repsm = 0.409093
+
+  # Declination of the Sun - RDS(teta).
+  sin_declination = jnp.sin(repsm) * jnp.sin(rllls)
+  cos_declination = jnp.sqrt(1.0 - sin_declination**2)
+
+  # Equation of time in seconds - RET(PTETA).
+  eq_of_time_seconds = jnp.dot(
+      jnp.stack(
+          [
+              sin_two_rlls,
+              sin_rem,
+              sin_rem * cos_two_rlls,
+              sin_four_rlls,
+              sin_two_rem,
+          ],
+          axis=-1,
+      ),
+      jnp.array([591.8, -459.4, 39.5, -12.7, -4.8]),
+  )
+
+  # Earth-Sun distance in astronomical units - RRS(PTETA).
+  solar_distance_au = jnp.dot(
+      jnp.stack([one, sin_rel, cos_rel], axis=-1),
+      jnp.array([1.0001, -0.0163, 0.0037]),
+  )
+
+  return _OrbitalParameters(
+      theta=theta,
+      rotational_phase=rotational_phase,
+      sin_declination=sin_declination,
+      cos_declination=cos_declination,
+      eq_of_time_seconds=eq_of_time_seconds,
+      solar_distance_au=solar_distance_au,
+  )
+
+
+def _get_solar_sin_altitude(
+    op: _OrbitalParameters,
+    sin_latitude: chex.Array,
+    cos_latitude: chex.Array,
+    longitude: chex.Array,
+) -> chex.Array:
+  """Returns the sine of the solar altitude angle.
+
+  All computations are vectorized. Dimensions of all the inputs should be
+  broadcastable using standard NumPy rules. For example, if `op` has shape
+  `(T, 1, 1)`, `latitude` has shape `(1, H, 1)`, and `longitude` has shape
+  `(1, H, W)`, the return value will have shape `(T, H, W)`.
+
+  Args:
+    op: Orbital parameters characterising Earth's position relative to the Sun.
+    sin_latitude: Sine of latitude coordinates.
+    cos_latitude: Cosine of latitude coordinates.
+    longitude: Longitude coordinates in radians.
+
+  Returns:
+    Sine of the solar altitude angle for each set of orbital parameters and each
+    geographical coordinates. The returned array has the shape resulting from
+    broadcasting all the inputs together.
+  """
+  solar_time = op.rotational_phase + op.eq_of_time_seconds / _SECONDS_PER_DAY
+  # https://en.wikipedia.org/wiki/Hour_angle#Solar_hour_angle
+  hour_angle = 2.0 * jnp.pi * solar_time + longitude
+  # https://en.wikipedia.org/wiki/Solar_zenith_angle
+  sin_altitude = (
+      cos_latitude * op.cos_declination * jnp.cos(hour_angle)
+      + sin_latitude * op.sin_declination
+  )
+  return sin_altitude
+
+
+def _get_radiation_flux(
+    j2000_days: chex.Array,
+    sin_latitude: chex.Array,
+    cos_latitude: chex.Array,
+    longitude: chex.Array,
+    tsi: chex.Array,
+) -> chex.Array:
+  """Computes the instantaneous TOA incident solar radiation flux.
+
+  Computes the instantanous Top-Of-the-Atmosphere (TOA) incident radiation flux
+  in W⋅m⁻² for the given timestamps and locations on the surface of the Earth.
+  See https://en.wikipedia.org/wiki/Solar_irradiance.
+
+  All inputs are assumed to be broadcastable together using standard NumPy
+  rules.
+
+  Args:
+    j2000_days: Timestamps represented as the number of days since the J2000
+      epoch.
+    sin_latitude: Sine of latitude coordinates.
+    cos_latitude: Cosine of latitude coordinates.
+    longitude: Longitude coordinates in radians.
+    tsi: Total Solar Irradiance (TSI) in W⋅m⁻². This can be a scalar (default)
+      to use the same TSI value for all the inputs, or an array to allow TSI to
+      depend on the timestamps.
+
+  Returns:
+    The instataneous TOA incident solar radiation flux in W⋅m⁻² for the given
+    timestamps and geographical coordinates. The returned array has the shape
+    resulting from broadcasting all the inputs together.
+  """
+  op = _get_orbital_parameters(j2000_days)
+  # Attenuation of the solar radiation based on the solar distance.
+  solar_factor = (1.0 / op.solar_distance_au) ** 2
+  sin_altitude = _get_solar_sin_altitude(
+      op, sin_latitude, cos_latitude, longitude
+  )
+  return tsi * solar_factor * jnp.maximum(sin_altitude, 0.0)
+
+
+def _get_integrated_radiation(
+    j2000_days: chex.Array,
+    sin_latitude: chex.Array,
+    cos_latitude: chex.Array,
+    longitude: chex.Array,
+    tsi: chex.Array,
+    integration_period: pd.Timedelta,
+    num_integration_bins: int,
+) -> chex.Array:
+  """Returns the TOA solar radiation flux integrated over a time period.
+
+  Integrates the instantaneous TOA solar radiation flux over a time period.
+  The input timestamps represent the end times of each integration period.
+  When the integration period is one hour this approximates the
+  `toa_incident_solar_radiation` (or `tisr`) parameter from the ERA5 dataset.
+  See https://confluence.ecmwf.int/display/CKB/ERA5%3A+data+documentation and
+  https://codes.ecmwf.int/grib/param-db/?id=212.
+
+  All inputs are assumed to be broadcastable together using standard NumPy
+  rules. To approximate the integral, the instantaneous radiation is computed
+  at `num_integration_bins+1` time steps using `_get_radiation_flux` and
+  integrated using the trapezoidal rule. A dimension is appended at the end
+  of all inputs to compute the instantaneous radiation, which is then integrated
+  over to compute the final result.
+
+  Args:
+    j2000_days: Timestamps represented as the number of days since the J2000
+      epoch. These correspond to the end times of each integration period.
+    sin_latitude: Sine of latitude coordinates.
+    cos_latitude: Cosine of latitude coordinates.
+    longitude: Longitude in radians.
+    tsi: Total Solar Irradiance (TSI) in W⋅m⁻².
+    integration_period: Integration period.
+    num_integration_bins: Number of bins to divide the `integration_period` to
+      approximate the integral using the trapezoidal rule.
+
+  Returns:
+    The TOA solar radiation flux integrated over the requested time period for
+    the given timestamps and geographical coordinates. Unit is J⋅m⁻² .
+  """
+  # Offsets for the integration time steps.
+  offsets = (
+      pd.timedelta_range(
+          start=-integration_period,
+          end=pd.Timedelta(0),
+          periods=num_integration_bins + 1,
+      )
+      / pd.Timedelta(days=1)
+  ).to_numpy()
+
+  # Integration happens over the time dimension. Compute the instantaneous
+  # radiation flux for all the integration time steps by appending a dimension
+  # to all the inputs and adding `offsets` to `j2000_days` (will be broadcast
+  # over all the other dimensions).
+  fluxes = _get_radiation_flux(
+      j2000_days=jnp.expand_dims(j2000_days, axis=-1) + offsets,
+      sin_latitude=jnp.expand_dims(sin_latitude, axis=-1),
+      cos_latitude=jnp.expand_dims(cos_latitude, axis=-1),
+      longitude=jnp.expand_dims(longitude, axis=-1),
+      tsi=jnp.expand_dims(tsi, axis=-1),
+  )
+
+  # Size of each bin in seconds. The instantaneous solar radiation flux is
+  # returned in units of W⋅m⁻². Integrating over time expressed in seconds
+  # yields a result in units of J⋅m⁻².
+  dx = (integration_period / num_integration_bins) / pd.Timedelta(seconds=1)
+  return jax.scipy.integrate.trapezoid(fluxes, dx=dx)
+
+
+_get_integrated_radiation_jitted = jax.jit(
+    _get_integrated_radiation,
+    static_argnames=["integration_period", "num_integration_bins"],
+)
+
+
+def get_toa_incident_solar_radiation(
+    timestamps: Sequence[_TimestampLike],
+    latitude: chex.Array,
+    longitude: chex.Array,
+    tsi_data: xa.DataArray | None = None,
+    integration_period: _TimedeltaLike = _DEFAULT_INTEGRATION_PERIOD,
+    num_integration_bins: int = _DEFAULT_NUM_INTEGRATION_BINS,
+    use_jit: bool = False,
+) -> chex.Array:
+  """Computes the solar radiation incident at the top of the atmosphere.
+
+  The solar radiation is computed for each element in `timestamps` for all the
+  locations on the grid determined by the `latitude` and `longitude` parameters.
+
+  To approximate the `toa_incident_solar_radiation` (or `tisr`) parameter from
+  the ERA5 dataset, set `integration_period` to one hour (default). See
+  https://confluence.ecmwf.int/display/CKB/ERA5%3A+data+documentation and
+  https://codes.ecmwf.int/grib/param-db/?id=212.
+
+  Args:
+    timestamps: Timestamps for which to compute the solar radiation.
+    latitude: The latitude coordinates in degrees of the grid for which to
+      compute the solar radiation.
+    longitude: The longitude coordinates in degrees of the grid for which to
+      compute the solar radiation.
+    tsi_data: A DataArray containing yearly TSI data as returned by a
+      `TsiDataLoader`. The default is to use ERA5 compatible TSI data.
+    integration_period: Timedelta to use to integrate the radiation, e.g. if
+      producing radiation for 1989-11-08 21:00:00, and `integration_period` is
+      "1h", radiation will be integrated from 1989-11-08 20:00:00 to 1989-11-08
+      21:00:00. The default value ("1h") matches ERA5.
+    num_integration_bins: Number of equally spaced bins to divide the
+      `integration_period` in when approximating the integral using the
+      trapezoidal rule. Performance and peak memory usage are affected by this
+      value. The default (360) provides a good approximation, but lower values
+      may work to improve performance and reduce memory usage.
+    use_jit: Set to True to use the jitted implementation, or False (default) to
+      use the non-jitted one.
+
+  Returns:
+    An 3D array with dimensions (time, lat, lon) containing the total
+    top of atmosphere solar radiation integrated for the `integration_period`
+    up to each timestamp.
+  """
+  # Add a trailing dimension to latitude to get dimensions (lat, lon).
+  lat = jnp.radians(latitude).reshape((-1, 1))
+  lon = jnp.radians(longitude)
+  sin_lat = jnp.sin(lat)
+  cos_lat = jnp.cos(lat)
+  integration_period = pd.Timedelta(integration_period)
+  if tsi_data is None:
+    tsi_data = _DEFAULT_TSI_DATA_LOADER()
+  tsi = get_tsi(timestamps, tsi_data)
+  fn = (
+      _get_integrated_radiation_jitted if use_jit else _get_integrated_radiation
+  )
+
+  # Compute integral for each timestamp individually. Although this could be
+  # done in one step, peak memory usage would be proportional to
+  # `len(timestamps) * num_integration_bins`. Computing each timestamp
+  # individually reduces this to `max(len(timestamps), num_integration_bins)`.
+  # E.g. memory usage for a single timestamp, with a full 0.25° grid and 360
+  # integration bins is about 1.5 GB (1440 * 721 * 361 * 4 bytes); computing
+  # forcings for 40 prediction steps would require 60 GB.
+  results = []
+  for idx, timestamp in enumerate(timestamps):
+    results.append(
+        fn(
+            j2000_days=jnp.array(_get_j2000_days(pd.Timestamp(timestamp))),
+            sin_latitude=sin_lat,
+            cos_latitude=cos_lat,
+            longitude=lon,
+            tsi=tsi[idx],
+            integration_period=integration_period,
+            num_integration_bins=num_integration_bins,
+        )
+    )
+  return jnp.stack(results, axis=0)
+
+
+def get_toa_incident_solar_radiation_for_xarray(
+    data_array_like: xa.DataArray | xa.Dataset,
+    tsi_data: xa.DataArray | None = None,
+    integration_period: _TimedeltaLike = _DEFAULT_INTEGRATION_PERIOD,
+    num_integration_bins: int = _DEFAULT_NUM_INTEGRATION_BINS,
+    use_jit: bool = False,
+) -> xa.DataArray:
+  """Computes the solar radiation incident at the top of the atmosphere.
+
+  This method is a wrapper for `get_toa_incident_solar_radiation` using
+  coordinates from an Xarray and returning an Xarray.
+
+  Args:
+    data_array_like: A xa.Dataset or xa.DataArray from which to take the time
+      and spatial coordinates for which to compute the solar radiation. It must
+      contain `lat` and `lon` spatial dimensions with corresponding coordinates.
+      If a `time` dimension is present, the `datetime` coordinate should be a
+      vector associated with that dimension containing timestamps for which to
+      compute the solar radiation. Otherwise, the `datetime` coordinate should
+      be a scalar representing the timestamp for which to compute the solar
+      radiation.
+    tsi_data: A DataArray containing yearly TSI data as returned by a
+      `TsiDataLoader`. The default is to use ERA5 compatible TSI data.
+    integration_period: Timedelta to use to integrate the radiation, e.g. if
+      producing radiation for 1989-11-08 21:00:00, and `integration_period` is
+      "1h", radiation will be integrated from 1989-11-08 20:00:00 to 1989-11-08
+      21:00:00. The default value ("1h") matches ERA5.
+    num_integration_bins: Number of equally spaced bins to divide the
+      `integration_period` in when approximating the integral using the
+      trapezoidal rule. Performance and peak memory usage are affected by this
+      value. The default (360) provides a good approximation, but lower values
+      may work to improve performance and reduce memory usage.
+    use_jit: Set to True to use the jitted implementation, or False to use the
+      non-jitted one.
+
+  Returns:
+    xa.DataArray with dimensions `(time, lat, lon)` if `data_array_like` had
+    a `time` dimension; or dimensions `(lat, lon)` otherwise. The `datetime`
+    coordinates and those for the dimensions are copied to the returned array.
+    The array contains the total top of atmosphere solar radiation integrated
+    for `integration_period` up to the corresponding `datetime`.
+
+  Raises:
+    ValueError: If there are missing coordinates or dimensions.
+  """
+  missing_dims = set(["lat", "lon"]) - set(data_array_like.dims)
+  if missing_dims:
+    raise ValueError(
+        f"'{missing_dims}' dimensions are missing in `data_array_like`."
+    )
+
+  missing_coords = set(["datetime", "lat", "lon"]) - set(data_array_like.coords)
+  if missing_coords:
+    raise ValueError(
+        f"'{missing_coords}' coordinates are missing in `data_array_like`."
+    )
+
+  if "time" in data_array_like.dims:
+    timestamps = data_array_like.coords["datetime"].data
+  else:
+    timestamps = [data_array_like.coords["datetime"].data.item()]
+
+  radiation = get_toa_incident_solar_radiation(
+      timestamps=timestamps,
+      latitude=data_array_like.coords["lat"].data,
+      longitude=data_array_like.coords["lon"].data,
+      tsi_data=tsi_data,
+      integration_period=integration_period,
+      num_integration_bins=num_integration_bins,
+      use_jit=use_jit,
+  )
+
+  if "time" in data_array_like.dims:
+    output = xa.DataArray(radiation, dims=("time", "lat", "lon"))
+  else:
+    output = xa.DataArray(radiation[0], dims=("lat", "lon"))
+
+  # Preserve as many of the original coordinates as possible, so long as the
+  # dimension or the coordinate still exist in the output array.
+  for k, coord in data_array_like.coords.items():
+    if set(coord.dims).issubset(set(output.dims)):
+      output.coords[k] = coord
+  return output

model/graphcast/sparse_transformer.py

@@ -0,0 +1,577 @@
+# Copyright 2024 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Transformer with either dense or sparse attention.
+
+The sparse attention implemented here is for nodes to attend only to themselves
+and their neighbours on the graph). It assumes that the adjacency matrix has a
+banded structure, and is implemented with dense operations computing with only
+the diagonal, super diagonal, and subdiagonal blocks of the tri-block-diagonal
+attention matrix.
+
+The basic model structure of the transformer and some functions were adapted
+from xlm's transformer_simple.py.
+"""
+
+import dataclasses
+import logging
+from typing import Any, Callable, Literal, Optional, Tuple
+
+from . import mlp as mlp_builder
+from . import sparse_transformer_utils as utils
+import haiku as hk
+import jax
+from jax.experimental.pallas.ops.tpu import splash_attention
+import jax.numpy as jnp
+import numpy as np
+import scipy as sp
+
+
+@dataclasses.dataclass
+class _ModelConfig:
+  """Transformer config."""
+  # Depth, or num transformer blocks. One 'layer' is attn + ffw.
+  num_layers: int
+  # Primary width, the number of channels on the carrier path.
+  d_model: int
+  # Number of heads for self-attention.
+  num_heads: int
+  # Mask block size.
+  mask_block_size: int
+  # Attention type - 'mha' or 'triblockdiag_mha'
+  attention_type: str = 'triblockdiag_mha'
+  block_q: Optional[int] = None
+  block_kv: Optional[int] = None
+  block_kv_compute: Optional[int] = None
+  block_q_dkv: Optional[int] = None
+  block_kv_dkv: Optional[int] = None
+  block_kv_dkv_compute: Optional[int] = None
+  # mask type if splash attention being used - 'full' or 'lazy'
+  mask_type: Optional[str] = 'full'
+  # Number of channels per-head for self-attn QK computation.
+  key_size: Optional[int] = None
+  # Number of channels per-head for self-attn V computation.
+  value_size: Optional[int] = None
+  # Activation to use, any in jax.nn.
+  activation: str = 'gelu'
+  # Init scale for ffw layers (divided by num_layers)
+  ffw_winit_mult: float = 2.0
+  # Init scale for final ffw layer (divided by depth)
+  ffw_winit_final_mult: float = 2.0
+  # Init scale for mha proj (divided by depth).
+  attn_winit_mult: float = 2.0
+  # Init scale for mha w (divided by depth).
+  attn_winit_final_mult: float = 2.0
+  # Number of hidden units in the MLP blocks. Defaults to 4 * d_model.
+  ffw_hidden: Optional[int] = None
+
+  def __post_init__(self):
+    if self.ffw_hidden is None:
+      self.ffw_hidden = 4 * self.d_model
+    # Compute key_size and value_size from d_model // num_heads.
+    if self.key_size is None:
+      if self.d_model % self.num_heads != 0:
+        raise ValueError('num_heads has to divide d_model exactly')
+      self.key_size = self.d_model // self.num_heads
+    if self.value_size is None:
+      if self.d_model % self.num_heads != 0:
+        raise ValueError('num_heads has to divide d_model exactly')
+      self.value_size = self.d_model // self.num_heads
+
+
+def get_mask_block_size(mask: sp.sparse.csr_matrix) -> int:
+  """Get blocksize of the adjacency matrix (attn mask) for the permuted mesh."""
+  # sub-diagonal bandwidth
+  lbandwidth = (
+      np.arange(mask.shape[0]) - (mask != 0).argmax(axis=0) + 1).max()
+  # super-diagonal bandwidth
+  ubandwidth = (
+      (mask.shape[0]-1) - np.argmax(mask[::-1,:] != 0, axis=0
+                                    ) - np.arange(mask.shape[0]) + 1).max()
+  block_size = np.maximum(lbandwidth, ubandwidth)
+  return block_size
+
+
+def ffw(x: jnp.ndarray, cfg: _ModelConfig) -> jnp.ndarray:
+  """Feed-forward block."""
+  ffw_winit = hk.initializers.VarianceScaling(cfg.ffw_winit_mult /
+                                              cfg.num_layers)
+  ffw_winit_final = hk.initializers.VarianceScaling(cfg.ffw_winit_final_mult /
+                                                    cfg.num_layers)
+  x = hk.Linear(cfg.ffw_hidden, name='ffw_up', w_init=ffw_winit)(x)
+  x = getattr(jax.nn, cfg.activation)(x)
+  return hk.Linear(cfg.d_model, name='ffw_down', w_init=ffw_winit_final)(x)
+
+
+def triblockdiag_softmax(logits: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]
+                         ) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
+  """Softmax given the diag, upper diag, and lower diag logit blocks."""
+
+  logits_d, logits_u, logits_l = logits
+
+  m = jnp.max(jnp.stack([
+      jax.lax.stop_gradient(logits_d.max(-1, keepdims=True)),
+      jax.lax.stop_gradient(logits_u.max(-1, keepdims=True)),
+      jax.lax.stop_gradient(logits_l.max(-1, keepdims=True))]), axis=0)
+
+  unnormalized_d = jnp.exp(logits_d - m)
+  unnormalized_u = jnp.exp(logits_u - m)
+  unnormalized_l = jnp.exp(logits_l - m)
+
+  denom = (
+      unnormalized_d.sum(-1, keepdims=True)
+      + unnormalized_u.sum(-1, keepdims=True)
+      + unnormalized_l.sum(-1, keepdims=True)
+  )
+
+  logits_d = unnormalized_d / denom
+  logits_u = unnormalized_u / denom
+  logits_l = unnormalized_l / denom
+
+  return (logits_d, logits_u, logits_l)
+
+
+def triblockdiag_mha(q_input: jnp.ndarray, kv_input: jnp.ndarray,
+                     mask: jnp.ndarray, cfg: _ModelConfig,
+                     ) -> jnp.ndarray:
+  """Triblockdiag multihead attention."""
+
+  # q_inputs, kv_input: (batch, num_blocks, block_size, num_heads, d_model)
+  q = multihead_linear(q_input, 'q', cfg)
+  k = multihead_linear(kv_input, 'k', cfg)
+  v = multihead_linear(kv_input, 'v', cfg)
+
+  k = jnp.pad(k, ((0, 0), (1, 1), (0, 0), (0, 0), (0, 0)))
+  v = jnp.pad(v, ((0, 0), (1, 1), (0, 0), (0, 0), (0, 0)))
+
+  def qk_prod(queries, keys):
+    return jnp.einsum('bnqhd,bnkhd->bnhqk', queries, keys)
+
+  # q shape is (batch, num_blocks, block_size, num_heads, qk_dim)
+  # k shape is (batch, num_blocks + 2, block_size, num_heads, qk_dim)
+  logits_d = qk_prod(q, k[:, 1:-1, ...]) * cfg.key_size**-0.5
+  logits_u = qk_prod(q, k[:, 2:, ...]) * cfg.key_size**-0.5
+  logits_l = qk_prod(q, k[:, :-2, ...]) * cfg.key_size**-0.5
+
+  # apply mask
+  logits_d = jnp.where(mask[:, 0, ...], logits_d, -1e30)
+  logits_u = jnp.where(mask[:, 1, ...], logits_u, -1e30)
+  logits_l = jnp.where(mask[:, 2, ...], logits_l, -1e30)
+
+  logits_d, logits_u, logits_l = utils.wrap_fn_for_upcast_downcast(
+      (logits_d, logits_u, logits_l),
+      triblockdiag_softmax
+      )
+
+  def av_prod(attn_weights, values):
+    return jnp.einsum('bnhqk,bnkhd->bnqhd', attn_weights, values)
+
+  out_d = av_prod(logits_d, v[:, 1:-1, ...])
+  out_u = av_prod(logits_u, v[:, 2:, ...])
+  out_l = av_prod(logits_l, v[:, :-2, ...])
+  # x shape is (batch, num_blocks, block_size, num_heads, d_model)
+  x = out_d + out_u + out_l
+
+  x = jnp.reshape(x, x.shape[:-2] + (cfg.num_heads * cfg.value_size,))
+  attn_winit_final = hk.initializers.VarianceScaling(
+      cfg.attn_winit_final_mult / cfg.num_layers)
+  x = hk.Linear(cfg.d_model, name='mha_final', w_init=attn_winit_final)(x)
+  return x
+
+
+def multihead_linear(
+    x: jnp.ndarray, qkv: str, cfg: _ModelConfig
+) -> jnp.ndarray:
+  """Linearly project `x` to have `head_size` dimensions per head."""
+  head_size = cfg.value_size if qkv == 'v' else cfg.key_size
+  attn_winit = hk.initializers.VarianceScaling(cfg.attn_winit_mult /
+                                               cfg.num_layers)
+  out = hk.Linear(
+      cfg.num_heads * head_size,
+      w_init=attn_winit,
+      name='mha_proj_' + qkv,
+      with_bias=False,
+  )(x)
+  shape = out.shape[:-1] + (cfg.num_heads, head_size)
+  return jnp.reshape(out, shape)
+
+
+def mha(q_input: jnp.ndarray, kv_input: jnp.ndarray,
+        mask: jnp.ndarray, cfg: _ModelConfig,
+        normalize_logits: bool = True,
+        ) -> jnp.ndarray:
+  """Multi head attention."""
+
+  q = multihead_linear(q_input, 'q', cfg)
+  k = multihead_linear(kv_input, 'k', cfg)
+  v = multihead_linear(kv_input, 'v', cfg)
+
+  logits = jnp.einsum('bthd, bThd->bhtT', q, k)
+  if normalize_logits:
+    logits *= cfg.key_size**-0.5
+  if mask is not None:
+    def apply_mask(m, l):
+      return jnp.where(m, l, -1e30)
+    logits = jax.vmap(jax.vmap(
+        apply_mask, in_axes=[None, 0]), in_axes=[None, 0])(mask, logits)
+
+  # Wrap softmax weights for upcasting & downcasting in case of BF16 activations
+  weights = utils.wrap_fn_for_upcast_downcast(logits, jax.nn.softmax)
+
+  # Note: our mask never has all 0 rows, since nodes always have self edges,
+  # so no need to account for that possibility explicitly.
+
+  x = jnp.einsum('bhtT,bThd->bthd', weights, v)
+  x = jnp.reshape(x, x.shape[:-2] + (cfg.num_heads * cfg.value_size,))
+
+  attn_winit_final = hk.initializers.VarianceScaling(
+      cfg.attn_winit_final_mult / cfg.num_layers)
+
+  x = hk.Linear(cfg.d_model, name='mha_final', w_init=attn_winit_final)(x)
+  return x
+
+
+def _make_splash_mha(
+    mask,
+    mask_type: str,
+    num_heads: int,
+    block_q: Optional[int] = None,
+    block_kv: Optional[int] = None,
+    block_kv_compute: Optional[int] = None,
+    block_q_dkv: Optional[int] = None,
+    block_kv_dkv: Optional[int] = None,
+    block_kv_dkv_compute: Optional[int] = None,
+    tanh_soft_cap: Optional[float] = None,
+) -> Callable[..., jnp.ndarray]:
+  """Construct attention kernel."""
+  if mask_type == 'full':
+    mask = np.broadcast_to(mask[None],
+                           (num_heads, *mask.shape)).astype(np.bool_)
+
+  block_sizes = splash_attention.BlockSizes(
+      block_q=block_q,
+      block_kv=block_kv,
+      block_kv_compute=block_kv_compute,
+      block_q_dkv=block_q_dkv,
+      block_kv_dkv=block_kv_dkv,
+      block_kv_dkv_compute=block_kv_dkv_compute,
+      use_fused_bwd_kernel=True,
+  )
+  attn = splash_attention.make_splash_mha(mask, block_sizes=block_sizes,
+                                          head_shards=1,
+                                          q_seq_shards=1,
+                                          attn_logits_soft_cap=tanh_soft_cap,
+                                          )
+  return attn
+
+
+def splash_mha(q_input: jnp.ndarray, kv_input: jnp.ndarray,
+               mask: jnp.ndarray | splash_attention.splash_attention_mask.Mask,
+               cfg: _ModelConfig,
+               tanh_soft_cap: Optional[float] = None,
+               normalize_q: bool = True) -> jnp.ndarray:
+  """Splash attention."""
+
+  q = multihead_linear(q_input, 'q', cfg)
+  k = multihead_linear(kv_input, 'k', cfg)
+  v = multihead_linear(kv_input, 'v', cfg)
+
+  _, _, num_heads, head_dim = q.shape
+
+  assert head_dim % 128 == 0  # splash attention kernel requires this
+
+  attn = _make_splash_mha(
+      mask=mask,
+      mask_type=cfg.mask_type,
+      num_heads=num_heads,
+      block_q=cfg.block_q,
+      block_kv=cfg.block_kv,
+      block_kv_compute=cfg.block_kv_compute,
+      block_q_dkv=cfg.block_q_dkv,
+      block_kv_dkv=cfg.block_kv_dkv,
+      block_kv_dkv_compute=cfg.block_kv_dkv_compute,
+      tanh_soft_cap=tanh_soft_cap,
+  )
+  attn = jax.vmap(attn)  # Add batch axis.
+
+  if normalize_q:
+    q *= cfg.key_size**-0.5
+
+  # (batch, nodes, num_heads, head_dim) -> (batch, num_heads, nodes, head_dim)
+  reformat = lambda y: y.transpose(0, 2, 1, 3)
+  x = attn(q=reformat(q), k=reformat(k), v=reformat(v))
+  x = x.transpose(0, 2, 1, 3)
+
+  x = jnp.reshape(x, x.shape[:-2] + (cfg.num_heads * cfg.value_size,))
+
+  attn_winit_final = hk.initializers.VarianceScaling(
+      cfg.attn_winit_final_mult / cfg.num_layers)
+
+  x = hk.Linear(cfg.d_model, name='mha_final', w_init=attn_winit_final)(x)
+  return x
+
+
+def layernorm(
+    x: jnp.ndarray, create_scale: bool, create_offset: bool
+) -> jnp.ndarray:
+  return hk.LayerNorm(
+      axis=-1, create_scale=create_scale, create_offset=create_offset,
+      name='norm')(x)
+
+
+def mask_block_diags(mask: sp.sparse.csr_matrix,
+                     num_padding_nodes: int,
+                     block_size: int) -> jnp.ndarray:
+  """Pad and reshape mask diag, super-siag and sub-diag blocks."""
+  # add zero padding to mask
+  mask_padding_rows = sp.sparse.csr_matrix(
+      (num_padding_nodes, mask.shape[1]), dtype=jnp.int32)
+  mask = sp.sparse.vstack([mask, mask_padding_rows])
+  mask_padding_cols = sp.sparse.csr_matrix(
+      (mask.shape[0], num_padding_nodes), dtype=jnp.int32)
+  mask = sp.sparse.hstack([mask, mask_padding_cols])
+
+  assert (mask.shape[-1] % block_size) == 0
+  mask_daig_blocks = jnp.stack(
+      [jnp.array(mask[i * block_size : (i + 1) * block_size,
+                      i * block_size : (i + 1) * block_size,
+                      ].toarray())
+       for i in range(mask.shape[0] // block_size)])
+  mask_upper_diag_blocks = jnp.stack(
+      [jnp.array(mask[i * block_size : (i + 1) * block_size,
+                      (i + 1) * block_size : (i + 2) * block_size,
+                      ].toarray())
+       for i in range(mask.shape[0] // block_size - 1)]
+      + [jnp.zeros((block_size, block_size), dtype=mask.dtype)])
+  mask_lower_diag_blocks = jnp.stack(
+      [jnp.zeros((block_size, block_size), dtype=mask.dtype)]
+      + [jnp.array(mask[(i + 1) * block_size : (i + 2) * block_size,
+                        i * block_size : (i + 1) * block_size,
+                        ].toarray())
+         for i in range(mask.shape[0] // block_size - 1)])
+  mask = jnp.stack(
+      [mask_daig_blocks, mask_upper_diag_blocks, mask_lower_diag_blocks]
+  )
+  mask = jnp.expand_dims(mask, (0, 3))
+  return mask
+
+
+def _pad_mask(mask, num_padding_nodes: Tuple[int, int]) -> jnp.ndarray:
+  q_padding, kv_padding = num_padding_nodes
+  mask_padding_rows = sp.sparse.csr_matrix(
+      (q_padding, mask.shape[1]), dtype=np.bool_)
+  mask = sp.sparse.vstack([mask, mask_padding_rows])
+  mask_padding_cols = sp.sparse.csr_matrix(
+      (mask.shape[0], kv_padding), dtype=np.bool_)
+  mask = sp.sparse.hstack([mask, mask_padding_cols])
+  return mask
+
+
+class WeatherMeshMask(splash_attention.splash_attention_mask.Mask):
+  """Lazy local mask, prevent attention to embeddings outside window.
+
+  Attributes:
+    mask:
+  """
+
+  _shape: Tuple[int, int]
+  mask: sp.sparse.spmatrix
+
+  def __init__(
+      self,
+      mask: Any
+  ):
+    self._shape = mask.shape
+    self.mask = mask
+
+  @property
+  def shape(self) -> Tuple[int, int]:
+    return self._shape
+
+  def __getitem__(self, idx) -> np.ndarray:
+    if len(idx) != 2:
+      raise NotImplementedError(f'Unsupported slice: {idx}')
+    q_slice, kv_slice = idx
+    if not isinstance(q_slice, slice) or not isinstance(kv_slice, slice):
+      raise NotImplementedError(f'Unsupported slice: {idx}')
+
+    return self.mask[q_slice, kv_slice].toarray()
+
+
+class Block(hk.Module):
+  """Transformer block (mha and ffw)."""
+
+  def __init__(self, cfg, mask, num_nodes, num_padding_nodes, name=None):
+    super().__init__(name=name)
+    self._cfg = cfg
+    self.mask = mask
+    self.num_nodes = num_nodes
+    self.num_padding_nodes = num_padding_nodes
+
+  def __call__(self, x, global_norm_conditioning=jax.Array):
+    # x shape is (batch, num_nodes, feature_dim)
+    def attn(x):
+      if self._cfg.attention_type == 'triblockdiag_mha':
+        # We pad -> reshape -> compute attn -> reshape -> select at each block
+        # so as to avoid complications involved in making the norm layers and
+        # ffw blocks account for the padding. However, this might be decreasing
+        # efficiency.
+
+        # Add padding so that number of nodes is divisible into blocks
+        x = jnp.pad(x, ((0, 0), (0, self.num_padding_nodes), (0, 0)))
+        x = x.reshape(x.shape[0],
+                      x.shape[1]//self._cfg.mask_block_size,
+                      self._cfg.mask_block_size,
+                      x.shape[-1])
+        x = triblockdiag_mha(x, x, mask=self.mask, cfg=self._cfg)
+        x = x.reshape(x.shape[0],
+                      self.num_nodes + self.num_padding_nodes,
+                      x.shape[-1])
+        return x[:,:self.num_nodes, :]
+
+      elif self._cfg.attention_type == 'mha':
+        return mha(x, x, mask=self.mask, cfg=self._cfg)
+
+      elif self._cfg.attention_type == 'splash_mha':
+        # We pad -> reshape -> compute attn -> reshape -> select at each block
+        # so as to avoid complications involved in making the norm layers and
+        # ffw blocks account for the padding. However, this might be decreasing
+        # efficiency.
+
+        # Add padding so that number of nodes is divisible by block sizes.
+        x = jnp.pad(x, ((0, 0), (0, self.num_padding_nodes[0]), (0, 0)))
+        x = splash_mha(x, x, mask=self.mask, cfg=self._cfg)
+        return x[:,:self.num_nodes, :]
+
+      else:
+        raise NotImplementedError()
+
+    def norm_conditioning_layer(x):
+      return mlp_builder.LinearNormConditioning(
+          name=self.name+'_norm_conditioning')(
+              x,
+              norm_conditioning=jnp.expand_dims(global_norm_conditioning, 1)
+              )
+
+    x = x + attn(
+        norm_conditioning_layer(
+            layernorm(x, create_scale=False, create_offset=False)
+        )
+    )
+    x = x + ffw(
+        norm_conditioning_layer(
+            layernorm(x, create_scale=False, create_offset=False)
+        ),
+        self._cfg,
+    )
+    return x
+
+
+class Transformer(hk.Module):
+  """Main transformer module that processes embeddings.
+
+  All but the very first and very last layer of a 'classic' Transformer:
+  Receives already embedded inputs instead of discrete tokens.
+  Outputs an embedding for each 'node'/'position' rather than logits.
+  """
+
+  def __init__(self,
+               adj_mat: sp.sparse.csr_matrix,
+               attention_k_hop: int,
+               attention_type: Literal['splash_mha', 'triblockdiag_mha', 'mha'],
+               mask_type: Literal['full', 'lazy'],
+               num_heads=1,
+               name=None,
+               block_q: Optional[int] = None,
+               block_kv: Optional[int] = None,
+               block_kv_compute: Optional[int] = None,
+               block_q_dkv: Optional[int] = None,
+               block_kv_dkv: Optional[int] = None,
+               block_kv_dkv_compute: Optional[int] = None,
+               **kwargs):
+    super().__init__(name=name)
+
+    # Construct mask and deduce block size.
+    mask = adj_mat ** attention_k_hop
+    mask_block_size = get_mask_block_size(mask)
+    logging.info('mask_block_size: %s.', mask_block_size)
+
+    if attention_type == 'triblockdiag_mha':
+      # we will stack the nodes in blocks of 'block_size' nodes, so we need to
+      # pad the input such that (num_nodes + num_padding_nodes) % block_size = 0
+      self.num_padding_nodes = int(np.ceil(
+          mask.shape[0]/mask_block_size)*mask_block_size
+                                   - mask.shape[0])
+      self.mask = mask_block_diags(
+          mask, self.num_padding_nodes, mask_block_size)
+    elif attention_type == 'splash_mha':
+      max_q_block_size = np.maximum(block_q, block_q_dkv)
+      max_kv_block_size = np.maximum(block_kv, block_kv_dkv)
+      q_padding = int(np.ceil(
+          mask.shape[0]/max_q_block_size)*max_q_block_size - mask.shape[0])
+      kv_padding = int(np.ceil(
+          mask.shape[1]/max_kv_block_size)*max_kv_block_size - mask.shape[1])
+      self.num_padding_nodes = (q_padding, kv_padding)
+      mask = _pad_mask(mask, self.num_padding_nodes)
+      if mask_type == 'lazy':
+        splash_mask = [
+            WeatherMeshMask(mask)
+            for _ in range(num_heads)
+        ]
+        self.mask = splash_attention.splash_attention_mask.MultiHeadMask(
+            splash_mask)
+      elif mask_type == 'full':
+        self.mask = mask.toarray()  # pytype: disable=attribute-error
+    elif attention_type == 'mha':
+      self.mask = jnp.array(mask.toarray())
+      self.num_padding_nodes = 0
+    else:
+      raise ValueError(
+          'Unsupported attention type: %s' % attention_type
+      )
+
+    # Construct config for use within class.
+    self._cfg = _ModelConfig(
+        mask_block_size=mask_block_size,
+        attention_type=attention_type,
+        mask_type=mask_type,
+        num_heads=num_heads,
+        block_q=block_q,
+        block_kv=block_kv,
+        block_kv_compute=block_kv_compute,
+        block_q_dkv=block_q_dkv,
+        block_kv_dkv=block_kv_dkv,
+        block_kv_dkv_compute=block_kv_dkv_compute,
+        **kwargs)
+
+  def __call__(self, node_features, global_norm_conditioning: jax.Array):
+    # node_features expected to have shape (batch, num_nodes, d)
+    x = node_features
+    for i_layer in range(self._cfg.num_layers):
+      x = Block(cfg=self._cfg, mask=self.mask,
+                num_nodes=node_features.shape[1],
+                num_padding_nodes=self.num_padding_nodes,
+                name='block_%02d' % i_layer
+                )(x, global_norm_conditioning=global_norm_conditioning)
+
+    def norm_conditioning_layer(x):
+      return mlp_builder.LinearNormConditioning(
+          name=self.name+'_final_norm_conditioning')(
+              x,
+              norm_conditioning=jnp.expand_dims(global_norm_conditioning, 1)
+              )
+    x = norm_conditioning_layer(
+        layernorm(x, create_scale=False, create_offset=False)
+    )
+
+    return x

model/graphcast/sparse_transformer_utils.py

@@ -0,0 +1,76 @@
+# Copyright 2024 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Utils for training models in low precision."""
+
+import functools
+from typing import Callable, Tuple, Union
+
+import jax
+import jax.numpy as jnp
+
+
+# Wrappers for jax.lax.reduce_precision which is non-differentiable.
+@functools.partial(jax.custom_vjp, nondiff_argnums=(1, 2))
+def reduce_precision(x, exponent_bits, mantissa_bits):
+  return jax.tree_util.tree_map(
+      lambda y: jax.lax.reduce_precision(y, exponent_bits, mantissa_bits), x)
+
+
+def reduce_precision_fwd(x, exponent_bits, mantissa_bits):
+  return reduce_precision(x, exponent_bits, mantissa_bits), None
+
+
+def reduce_precision_bwd(exponent_bits, mantissa_bits, res, dout):
+  del res  # Unused.
+  return reduce_precision(dout, exponent_bits, mantissa_bits),
+
+
+reduce_precision.defvjp(reduce_precision_fwd, reduce_precision_bwd)
+
+
+def wrap_fn_for_upcast_downcast(inputs: Union[jnp.ndarray,
+                                              Tuple[jnp.ndarray, ...]],
+                                fn: Callable[[Union[jnp.ndarray,
+                                                    Tuple[jnp.ndarray, ...]]],
+                                             Union[jnp.ndarray,
+                                                   Tuple[jnp.ndarray, ...]]],
+                                f32_upcast: bool = True,
+                                guard_against_excess_precision: bool = True
+                                ) ->  Union[jnp.ndarray,
+                                            Tuple[jnp.ndarray, ...]]:
+  """Wraps `fn` to  upcast to float32 and then downcast, for use with BF16."""
+  # Do not upcast if the inputs are already in float32.
+  # This removes a no-op `jax.lax.reduce_precision` which is unsupported
+  # in jax2tf at the moment.
+  if isinstance(inputs, Tuple):
+    f32_upcast = f32_upcast and inputs[0].dtype != jnp.float32
+    orig_dtype = inputs[0].dtype
+  else:
+    f32_upcast = f32_upcast and inputs.dtype != jnp.float32
+    orig_dtype = inputs.dtype
+
+  if f32_upcast:
+    inputs = jax.tree_util.tree_map(lambda x: x.astype(jnp.float32), inputs)
+
+    if guard_against_excess_precision:
+      # This is evil magic to guard against differences in precision in the QK
+      # calculation between the forward pass and backwards pass. This is like
+      # --xla_allow_excess_precision=false but scoped here.
+      finfo = jnp.finfo(orig_dtype)  # jnp important!
+      inputs = reduce_precision(inputs, finfo.nexp, finfo.nmant)
+
+  output = fn(inputs)
+  if f32_upcast:
+    output = jax.tree_util.tree_map(lambda x: x.astype(orig_dtype), output)
+  return output

model/graphcast/transformer.py

@@ -0,0 +1,124 @@
+# Copyright 2024 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""A Transformer model for weather predictions.
+
+This model wraps the a transformer model and swaps the leading two axes of the
+nodes in the input graph prior to evaluating the model to make it compatible
+with a [nodes, batch, ...] ordering of the inputs.
+"""
+
+from typing import Any, Mapping, Optional
+
+from . import typed_graph
+import haiku as hk
+import jax
+import jax.numpy as jnp
+import numpy as np
+from scipy import sparse
+
+
+Kwargs = Mapping[str, Any]
+
+
+def _get_adj_matrix_for_edge_set(
+    graph: typed_graph.TypedGraph,
+    edge_set_name: str,
+    add_self_edges: bool,
+):
+  """Returns the adjacency matrix for the given graph and edge set."""
+  # Get nodes and edges of the graph.
+  edge_set_key = graph.edge_key_by_name(edge_set_name)
+  sender_node_set, receiver_node_set = edge_set_key.node_sets
+
+  # Compute number of sender and receiver nodes.
+  sender_n_node = graph.nodes[sender_node_set].n_node[0]
+  receiver_n_node = graph.nodes[receiver_node_set].n_node[0]
+
+  # Build adjacency matrix.
+  adj_mat = sparse.csr_matrix((sender_n_node, receiver_n_node), dtype=np.bool_)
+  edge_set = graph.edges[edge_set_key]
+  s, r = edge_set.indices
+  adj_mat[s, r] = True
+  if add_self_edges:
+    # Should only do this if we are certain the adjacency matrix is square.
+    assert sender_node_set == receiver_node_set
+    adj_mat[np.arange(sender_n_node), np.arange(receiver_n_node)] = True
+  return adj_mat
+
+
+class MeshTransformer(hk.Module):
+  """A Transformer for inputs with ordering [nodes, batch, ...]."""
+
+  def __init__(self,
+               transformer_ctor,
+               transformer_kwargs: Kwargs,
+               name: Optional[str] = None):
+    """Initialises the Transformer model.
+
+    Args:
+      transformer_ctor: Constructor for transformer.
+      transformer_kwargs: Kwargs to pass to the transformer module.
+      name: Optional name for haiku module.
+    """
+    super().__init__(name=name)
+    # We defer the transformer initialisation to the first call to __call__,
+    # where we can build the mask senders and receivers of the TypedGraph
+    self._batch_first_transformer = None
+    self._transformer_ctor = transformer_ctor
+    self._transformer_kwargs = transformer_kwargs
+
+  @hk.name_like('__init__')
+  def _maybe_init_batch_first_transformer(self, x: typed_graph.TypedGraph):
+    if self._batch_first_transformer is not None:
+      return
+    self._batch_first_transformer = self._transformer_ctor(
+        adj_mat=_get_adj_matrix_for_edge_set(
+            graph=x,
+            edge_set_name='mesh',
+            add_self_edges=True,
+        ),
+        **self._transformer_kwargs,
+    )
+
+  def __call__(
+      self, x: typed_graph.TypedGraph,
+      global_norm_conditioning: jax.Array
+  ) -> typed_graph.TypedGraph:
+    """Applies the model to the input graph and returns graph of same shape."""
+
+    if set(x.nodes.keys()) != {'mesh_nodes'}:
+      raise ValueError(
+          f'Expected x.nodes to have key `mesh_nodes`, got {x.nodes.keys()}.'
+      )
+    features = x.nodes['mesh_nodes'].features
+    if features.ndim != 3:  # pytype: disable=attribute-error  # jax-ndarray
+      raise ValueError(
+          'Expected `x.nodes["mesh_nodes"].features` to be 3, got'
+          f' {features.ndim}.'
+      )  # pytype: disable=attribute-error  # jax-ndarray
+
+    # Initialise transformer and mask.
+    self._maybe_init_batch_first_transformer(x)
+
+    y = jnp.transpose(features, axes=[1, 0, 2])
+    y = self._batch_first_transformer(y, global_norm_conditioning)
+    y = jnp.transpose(y, axes=[1, 0, 2])
+    x = x._replace(
+        nodes={
+            'mesh_nodes': x.nodes['mesh_nodes']._replace(
+                features=y.astype(features.dtype)
+            )
+        }
+    )
+    return x

model/graphcast/typed_graph.py

@@ -0,0 +1,97 @@
+# Copyright 2023 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Data-structure for storing graphs with typed edges and nodes."""
+
+from typing import NamedTuple, Any, Union, Tuple, Mapping, TypeVar
+
+ArrayLike = Union[Any]  # np.ndarray, jnp.ndarray, tf.tensor
+ArrayLikeTree = Union[Any, ArrayLike]  # Nest of ArrayLike
+
+_T = TypeVar('_T')
+
+
+# All tensors have a "flat_batch_axis", which is similar to the leading
+# axes of graph_tuples:
+# * In the case of nodes this is simply a shared node and flat batch axis, with
+# size corresponding to the total number of nodes in the flattened batch.
+# * In the case of edges this is simply a shared edge and flat batch axis, with
+# size corresponding to the total number of edges in the flattened batch.
+# * In the case of globals this is simply the number of graphs in the flattened
+# batch.
+
+# All shapes may also have any additional leading shape "batch_shape".
+# Options for building batches are:
+# * Use a provided "flatten" method that takes a leading `batch_shape` and
+#   it into the flat_batch_axis (this will be useful when using `tf.Dataset`
+#   which supports batching into RaggedTensors, with leading batch shape even
+#   if graphs have different numbers of nodes and edges), so the RaggedBatches
+#   can then be converted into something without ragged dimensions that jax can
+#   use.
+# * Directly build a "flat batch" using a provided function for batching a list
+#   of graphs (how it is done in `jraph`).
+
+
+class NodeSet(NamedTuple):
+  """Represents a set of nodes."""
+  n_node: ArrayLike  # [num_flat_graphs]
+  features: ArrayLikeTree  # Prev. `nodes`: [num_flat_nodes] + feature_shape
+
+
+class EdgesIndices(NamedTuple):
+  """Represents indices to nodes adjacent to the edges."""
+  senders: ArrayLike  # [num_flat_edges]
+  receivers: ArrayLike  # [num_flat_edges]
+
+
+class EdgeSet(NamedTuple):
+  """Represents a set of edges."""
+  n_edge: ArrayLike  # [num_flat_graphs]
+  indices: EdgesIndices
+  features: ArrayLikeTree  # Prev. `edges`: [num_flat_edges] + feature_shape
+
+
+class Context(NamedTuple):
+  # `n_graph` always contains ones but it is useful to query the leading shape
+  # in case of graphs without any nodes or edges sets.
+  n_graph: ArrayLike  # [num_flat_graphs]
+  features: ArrayLikeTree  # Prev. `globals`: [num_flat_graphs] + feature_shape
+
+
+class EdgeSetKey(NamedTuple):
+  name: str   # Name of the EdgeSet.
+
+  # Sender node set name and receiver node set name connected by the edge set.
+  node_sets: Tuple[str, str]
+
+
+class TypedGraph(NamedTuple):
+  """A graph with typed nodes and edges.
+
+  A typed graph is made of a context, multiple sets of nodes and multiple
+  sets of edges connecting those nodes (as indicated by the EdgeSetKey).
+  """
+
+  context: Context
+  nodes: Mapping[str, NodeSet]
+  edges: Mapping[EdgeSetKey, EdgeSet]
+
+  def edge_key_by_name(self, name: str) -> EdgeSetKey:
+    found_key = [k for k in self.edges.keys() if k.name == name]
+    if len(found_key) != 1:
+      raise KeyError("invalid edge key '{}'. Available edges: [{}]".format(
+          name, ', '.join(x.name for x in self.edges.keys())))
+    return found_key[0]
+
+  def edge_by_name(self, name: str) -> EdgeSet:
+    return self.edges[self.edge_key_by_name(name)]

model/graphcast/typed_graph_net.py

@@ -0,0 +1,317 @@
+# Copyright 2023 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""A library of typed Graph Neural Networks."""
+
+from typing import Callable, Mapping, Optional, Union
+
+from . import typed_graph
+import jax.numpy as jnp
+import jax.tree_util as tree
+import jraph
+
+
+# All features will be an ArrayTree.
+NodeFeatures = EdgeFeatures = SenderFeatures = ReceiverFeatures = Globals = (
+    jraph.ArrayTree)
+
+# Signature:
+# (node features, outgoing edge features, incoming edge features,
+#  globals) -> updated node features
+GNUpdateNodeFn = Callable[
+    [NodeFeatures, Mapping[str, SenderFeatures], Mapping[str, ReceiverFeatures],
+     Globals],
+    NodeFeatures]
+
+GNUpdateGlobalFn = Callable[
+    [Mapping[str, NodeFeatures], Mapping[str, EdgeFeatures], Globals],
+    Globals]
+
+
+def GraphNetwork(  # pylint: disable=invalid-name
+    update_edge_fn: Mapping[str, jraph.GNUpdateEdgeFn],
+    update_node_fn: Mapping[str, GNUpdateNodeFn],
+    update_global_fn: Optional[GNUpdateGlobalFn] = None,
+    aggregate_edges_for_nodes_fn: jraph.AggregateEdgesToNodesFn = jraph
+    .segment_sum,
+    aggregate_nodes_for_globals_fn: jraph.AggregateNodesToGlobalsFn = jraph
+    .segment_sum,
+    aggregate_edges_for_globals_fn: jraph.AggregateEdgesToGlobalsFn = jraph
+    .segment_sum,
+    ):
+  """Returns a method that applies a configured GraphNetwork.
+
+  This implementation follows Algorithm 1 in https://arxiv.org/abs/1806.01261
+  extended to Typed Graphs with multiple edge sets and node sets and extended to
+  allow aggregating not only edges received by the nodes, but also edges sent by
+  the nodes.
+
+  Example usage::
+
+    gn = GraphNetwork(update_edge_function,
+    update_node_function, **kwargs)
+    # Conduct multiple rounds of message passing with the same parameters:
+    for _ in range(num_message_passing_steps):
+      graph = gn(graph)
+
+  Args:
+    update_edge_fn: mapping of functions used to update a subset of the edge
+      types, indexed by edge type name.
+    update_node_fn: mapping of functions used to update a subset of the node
+      types, indexed by node type name.
+    update_global_fn: function used to update the globals or None to deactivate
+      globals updates.
+    aggregate_edges_for_nodes_fn: function used to aggregate messages to each
+      node.
+    aggregate_nodes_for_globals_fn: function used to aggregate the nodes for the
+      globals.
+    aggregate_edges_for_globals_fn: function used to aggregate the edges for the
+      globals.
+
+  Returns:
+    A method that applies the configured GraphNetwork.
+  """
+
+  def _apply_graph_net(graph: typed_graph.TypedGraph) -> typed_graph.TypedGraph:
+    """Applies a configured GraphNetwork to a graph.
+
+    This implementation follows Algorithm 1 in https://arxiv.org/abs/1806.01261
+    extended to Typed Graphs with multiple edge sets and node sets and extended
+    to allow aggregating not only edges received by the nodes, but also edges
+    sent by the nodes.
+
+    Args:
+      graph: a `TypedGraph` containing the graph.
+
+    Returns:
+      Updated `TypedGraph`.
+    """
+
+    updated_graph = graph
+
+    # Edge update.
+    updated_edges = dict(updated_graph.edges)
+    for edge_set_name, edge_fn in update_edge_fn.items():
+      edge_set_key = graph.edge_key_by_name(edge_set_name)
+      updated_edges[edge_set_key] = _edge_update(
+          updated_graph, edge_fn, edge_set_key)
+    updated_graph = updated_graph._replace(edges=updated_edges)
+
+    # Node update.
+    updated_nodes = dict(updated_graph.nodes)
+    for node_set_key, node_fn in update_node_fn.items():
+      updated_nodes[node_set_key] = _node_update(
+          updated_graph, node_fn, node_set_key, aggregate_edges_for_nodes_fn)
+    updated_graph = updated_graph._replace(nodes=updated_nodes)
+
+    # Global update.
+    if update_global_fn:
+      updated_context = _global_update(
+          updated_graph, update_global_fn,
+          aggregate_edges_for_globals_fn,
+          aggregate_nodes_for_globals_fn)
+      updated_graph = updated_graph._replace(context=updated_context)
+
+    return updated_graph
+
+  return _apply_graph_net
+
+
+def _edge_update(graph, edge_fn, edge_set_key):  # pylint: disable=invalid-name
+  """Updates an edge set of a given key."""
+
+  sender_nodes = graph.nodes[edge_set_key.node_sets[0]]
+  receiver_nodes = graph.nodes[edge_set_key.node_sets[1]]
+  edge_set = graph.edges[edge_set_key]
+  senders = edge_set.indices.senders  # pytype: disable=attribute-error
+  receivers = edge_set.indices.receivers  # pytype: disable=attribute-error
+
+  sent_attributes = tree.tree_map(
+      lambda n: n[senders], sender_nodes.features)
+  received_attributes = tree.tree_map(
+      lambda n: n[receivers], receiver_nodes.features)
+
+  n_edge = edge_set.n_edge
+  sum_n_edge = senders.shape[0]
+  global_features = tree.tree_map(
+      lambda g: jnp.repeat(g, n_edge, axis=0, total_repeat_length=sum_n_edge),
+      graph.context.features)
+  new_features = edge_fn(
+      edge_set.features, sent_attributes, received_attributes,
+      global_features)
+  return edge_set._replace(features=new_features)
+
+
+def _node_update(graph, node_fn, node_set_key, aggregation_fn):  # pylint: disable=invalid-name
+  """Updates an edge set of a given key."""
+  node_set = graph.nodes[node_set_key]
+  sum_n_node = tree.tree_leaves(node_set.features)[0].shape[0]
+
+  sent_features = {}
+  for edge_set_key, edge_set in graph.edges.items():
+    sender_node_set_key = edge_set_key.node_sets[0]
+    if sender_node_set_key == node_set_key:
+      assert isinstance(edge_set.indices, typed_graph.EdgesIndices)
+      senders = edge_set.indices.senders
+      sent_features[edge_set_key.name] = tree.tree_map(
+          lambda e: aggregation_fn(e, senders, sum_n_node), edge_set.features)  # pylint: disable=cell-var-from-loop
+
+  received_features = {}
+  for edge_set_key, edge_set in graph.edges.items():
+    receiver_node_set_key = edge_set_key.node_sets[1]
+    if receiver_node_set_key == node_set_key:
+      assert isinstance(edge_set.indices, typed_graph.EdgesIndices)
+      receivers = edge_set.indices.receivers
+      received_features[edge_set_key.name] = tree.tree_map(
+          lambda e: aggregation_fn(e, receivers, sum_n_node), edge_set.features)  # pylint: disable=cell-var-from-loop
+
+  n_node = node_set.n_node
+  global_features = tree.tree_map(
+      lambda g: jnp.repeat(g, n_node, axis=0, total_repeat_length=sum_n_node),
+      graph.context.features)
+  new_features = node_fn(
+      node_set.features, sent_features, received_features, global_features)
+  return node_set._replace(features=new_features)
+
+
+def _global_update(graph, global_fn, edge_aggregation_fn, node_aggregation_fn):  # pylint: disable=invalid-name
+  """Updates an edge set of a given key."""
+  n_graph = graph.context.n_graph.shape[0]
+  graph_idx = jnp.arange(n_graph)
+
+  edge_features = {}
+  for edge_set_key, edge_set in graph.edges.items():
+    assert isinstance(edge_set.indices, typed_graph.EdgesIndices)
+    sum_n_edge = edge_set.indices.senders.shape[0]
+    edge_gr_idx = jnp.repeat(
+        graph_idx, edge_set.n_edge, axis=0, total_repeat_length=sum_n_edge)
+    edge_features[edge_set_key.name] = tree.tree_map(
+        lambda e: edge_aggregation_fn(e, edge_gr_idx, n_graph),   # pylint: disable=cell-var-from-loop
+        edge_set.features)
+
+  node_features = {}
+  for node_set_key, node_set in graph.nodes.items():
+    sum_n_node = tree.tree_leaves(node_set.features)[0].shape[0]
+    node_gr_idx = jnp.repeat(
+        graph_idx, node_set.n_node, axis=0, total_repeat_length=sum_n_node)
+    node_features[node_set_key] = tree.tree_map(
+        lambda n: node_aggregation_fn(n, node_gr_idx, n_graph),   # pylint: disable=cell-var-from-loop
+        node_set.features)
+
+  new_features = global_fn(node_features, edge_features, graph.context.features)
+  return graph.context._replace(features=new_features)
+
+
+InteractionUpdateNodeFn = Callable[
+    [jraph.NodeFeatures,
+     Mapping[str, SenderFeatures],
+     Mapping[str, ReceiverFeatures]],
+    jraph.NodeFeatures]
+
+
+InteractionUpdateNodeFnNoSentEdges = Callable[
+    [jraph.NodeFeatures,
+     Mapping[str, ReceiverFeatures]],
+    jraph.NodeFeatures]
+
+
+def InteractionNetwork(  # pylint: disable=invalid-name
+    update_edge_fn: Mapping[str, jraph.InteractionUpdateEdgeFn],
+    update_node_fn: Mapping[str, Union[InteractionUpdateNodeFn,
+                                       InteractionUpdateNodeFnNoSentEdges]],
+    aggregate_edges_for_nodes_fn: jraph.AggregateEdgesToNodesFn = jraph
+    .segment_sum,
+    include_sent_messages_in_node_update: bool = False):
+  """Returns a method that applies a configured InteractionNetwork.
+
+  An interaction network computes interactions on the edges based on the
+  previous edges features, and on the features of the nodes sending into those
+  edges. It then updates the nodes based on the incoming updated edges.
+  See https://arxiv.org/abs/1612.00222 for more details.
+
+  This implementation extends the behavior to `TypedGraphs` adding an option
+  to include edge features for which a node is a sender in the arguments to
+  the node update function.
+
+  Args:
+    update_edge_fn: mapping of functions used to update a subset of the edge
+      types, indexed by edge type name.
+    update_node_fn: mapping of functions used to update a subset of the node
+      types, indexed by node type name.
+    aggregate_edges_for_nodes_fn: function used to aggregate messages to each
+      node.
+    include_sent_messages_in_node_update: pass edge features for which a node is
+      a sender to the node update function.
+  """
+  # An InteractionNetwork is a GraphNetwork without globals features,
+  # so we implement the InteractionNetwork as a configured GraphNetwork.
+
+  # An InteractionNetwork edge function does not have global feature inputs,
+  # so we filter the passed global argument in the GraphNetwork.
+  wrapped_update_edge_fn = tree.tree_map(
+      lambda fn: lambda e, s, r, g: fn(e, s, r), update_edge_fn)
+
+  # Similarly, we wrap the update_node_fn to ensure only the expected
+  # arguments are passed to the Interaction net.
+  if include_sent_messages_in_node_update:
+    wrapped_update_node_fn = tree.tree_map(
+        lambda fn: lambda n, s, r, g: fn(n, s, r), update_node_fn)
+  else:
+    wrapped_update_node_fn = tree.tree_map(
+        lambda fn: lambda n, s, r, g: fn(n, r), update_node_fn)
+  return GraphNetwork(
+      update_edge_fn=wrapped_update_edge_fn,
+      update_node_fn=wrapped_update_node_fn,
+      aggregate_edges_for_nodes_fn=aggregate_edges_for_nodes_fn)
+
+
+def GraphMapFeatures(  # pylint: disable=invalid-name
+    embed_edge_fn: Optional[Mapping[str, jraph.EmbedEdgeFn]] = None,
+    embed_node_fn: Optional[Mapping[str, jraph.EmbedNodeFn]] = None,
+    embed_global_fn: Optional[jraph.EmbedGlobalFn] = None):
+  """Returns function which embeds the components of a graph independently.
+
+  Args:
+    embed_edge_fn: mapping of functions used to embed each edge type,
+      indexed by edge type name.
+    embed_node_fn: mapping of functions used to embed each node type,
+      indexed by node type name.
+    embed_global_fn: function used to embed the globals.
+  """
+
+  def _embed(graph: typed_graph.TypedGraph) -> typed_graph.TypedGraph:
+
+    updated_edges = dict(graph.edges)
+    if embed_edge_fn:
+      for edge_set_name, embed_fn in embed_edge_fn.items():
+        edge_set_key = graph.edge_key_by_name(edge_set_name)
+        edge_set = graph.edges[edge_set_key]
+        updated_edges[edge_set_key] = edge_set._replace(
+            features=embed_fn(edge_set.features))
+
+    updated_nodes = dict(graph.nodes)
+    if embed_node_fn:
+      for node_set_key, embed_fn in embed_node_fn.items():
+        node_set = graph.nodes[node_set_key]
+        updated_nodes[node_set_key] = node_set._replace(
+            features=embed_fn(node_set.features))
+
+    updated_context = graph.context
+    if embed_global_fn:
+      updated_context = updated_context._replace(
+          features=embed_global_fn(updated_context.features))
+
+    return graph._replace(edges=updated_edges, nodes=updated_nodes,
+                          context=updated_context)
+
+  return _embed

model/graphcast/xarray_jax.py

@@ -0,0 +1,1080 @@
+# Copyright 2023 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Helpers to use xarray.{Variable,DataArray,Dataset} with JAX.
+
+Allows them to be based on JAX arrays without converting to numpy arrays under
+the hood, so you can start with a JAX array, do some computation with it in
+xarray-land, get a JAX array out the other end and (for example) jax.jit
+through the whole thing. You can even jax.jit a function which accepts and
+returns xarray.Dataset, DataArray and Variable.
+
+## Creating xarray datatypes from jax arrays, and vice-versa.
+
+You can use the xarray_jax.{Variable,DataArray,Dataset} constructors, which have
+the same API as the standard xarray constructors but will accept JAX arrays
+without converting them to numpy.
+
+It does this by wrapping the JAX array in a wrapper before passing it to
+xarray; you can also do this manually by calling xarray_jax.wrap on your JAX
+arrays before passing them to the standard xarray constructors.
+
+To get non-wrapped JAX arrays out the other end, you can use e.g.:
+
+  xarray_jax.jax_vars(dataset)
+  xarray_jax.jax_data(dataset.some_var)
+
+which will complain if the data isn't actually a JAX array. Use this if you need
+to make sure the computation has gone via JAX, e.g. if it's the output of code
+that you want to JIT or compute gradients through. If this is not the case and
+you want to support passing plain numpy arrays through as well as potentially
+JAX arrays, you can use:
+
+  xarray_jax.unwrap_vars(dataset)
+  xarray_jax.unwrap_data(dataset.some_var)
+
+which will unwrap the data if it is a wrapped JAX array, but otherwise pass
+it through to you without complaint.
+
+The wrapped JAX arrays aim to support all the core operations from the numpy
+array API that xarray expects, however there may still be some gaps; if you run
+into any problems around this, you may need to add a few more proxy methods onto
+the wrapper class below.
+
+In future once JAX and xarray support the new  Python array API standard
+(https://data-apis.org/array-api/latest/index.html), we hope to avoid the need
+for wrapping the JAX arrays like this.
+
+## jax.jit and pmap of functions taking and returning xarray datatypes
+
+We register xarray datatypes with jax.tree_util, which allows them to be treated
+as generic containers of jax arrays by various parts of jax including jax.jit.
+
+This allows for, e.g.:
+
+  @jax.jit
+  def foo(input: xarray.Dataset) -> xarray.Dataset:
+    ...
+
+It will not work out-of-the-box with shape-modifying transformations like
+jax.pmap, or e.g. a jax.tree_util.tree_map with some transform that alters array
+shapes or dimension order. That's because we won't know what dimension names
+and/or coordinates to use when unflattening, if the results have a different
+shape to the data that was originally flattened.
+
+You can work around this using xarray_jax.dims_change_on_unflatten, however,
+and in the case of jax.pmap we provide a wrapper xarray_jax.pmap which allows
+it to be used with functions taking and returning xarrays.
+
+## Treatment of coordinates
+
+We don't support passing jax arrays as coordinates when constructing a
+DataArray/Dataset. This is because xarray's advanced indexing and slicing is
+unlikely to work with jax arrays (at least when a Tracer is used during
+jax.jit), and also because some important datatypes used for coordinates, like
+timedelta64 and datetime64, are not supported by jax.
+
+For the purposes of tree_util and jax.jit, coordinates are not treated as leaves
+of the tree (array data 'contained' by a Dataset/DataArray), they are just a
+static part of the structure. That means that if a jit'ed function is called
+twice with Dataset inputs that use different coordinates, it will compile a
+separate version of the function for each. The coordinates are treated like
+static_argnums by jax.jit.
+
+If you want to use dynamic data for coordinates, we recommend making it a
+data_var instead of a coord. You won't be able to do indexing and slicing using
+the coordinate, but that wasn't going to work with a jax array anyway.
+"""
+
+import collections
+import contextlib
+import contextvars
+from typing import Any, Callable, Iterator, Mapping, Optional, Union, Tuple, TypeVar, cast
+from typing import Hashable  # pylint: disable=deprecated-class
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import tree
+import xarray
+
+
+# Types which we wrap with JaxArrayWrapper to allow creating xarray datatypes
+# from them.
+# Note this includes some non-Array types which jax sometimes needs to use as
+# leaves of pytrees, in order to ensure we can still use xarray datatypes as
+# internal pytree nodes in these cases.
+_WRAPPED_TYPES = (
+    jax.Array, jax.ShapeDtypeStruct, jax.stages.ArgInfo)
+
+
+def Variable(dims, data, **kwargs) -> xarray.Variable:  # pylint:disable=invalid-name
+  """Like xarray.Variable, but can wrap JAX arrays."""
+  return xarray.Variable(dims, wrap(data), **kwargs)
+
+
+_JAX_COORD_ATTR_NAME = '_jax_coord'
+
+
+def DataArray(  # pylint:disable=invalid-name
+    data,
+    coords=None,
+    dims=None,
+    name=None,
+    attrs=None,
+    jax_coords=None,
+    ) -> xarray.DataArray:
+  """Like xarray.DataArray, but supports using JAX arrays.
+
+  Args:
+    data: As for xarray.DataArray, except jax arrays are also supported.
+    coords: Coordinates for the array, see xarray.DataArray. These coordinates
+      must be based on plain numpy arrays or something convertible to plain
+      numpy arrays. Their values will form a static part of the data structure
+      from the point of view of jax.tree_util. In particular this means these
+      coordinates will be passed as plain numpy arrays even inside a JIT'd
+      function, and the JIT'd function will be recompiled under the hood if the
+      coordinates of DataArrays passed into it change.
+      If this is not convenient for you, see also jax_coords below.
+    dims: See xarray.DataArray.
+    name: See xarray.DataArray.
+    attrs: See xarray.DataArray.
+    jax_coords: Additional coordinates, which *can* use JAX arrays. These
+      coordinates will be treated as JAX data from the point of view of
+      jax.tree_util, that means when JIT'ing they will be passed as tracers and
+      computation involving them will be JIT'd.
+      Unfortunately a side-effect of this is that they can't be used as index
+      coordinates (because xarray's indexing logic is not JIT-able). If you
+      specify a coordinate with the same name as a dimension here, it will not
+      be set as an index coordinate; this behaviour is different to the default
+      for `coords`, and it means that things like `.sel` based on the jax
+      coordinate will not work.
+      Note we require `jax_coords` to be explicitly specified via a different
+      constructor argument to `coords`, rather than just looking for jax arrays
+      within the `coords` and treating them differently. This is because it
+      affects the way jax.tree_util treats them, which is somewhat orthogonal to
+      whether the value is passed in as numpy or not, and generally needs to be
+      handled consistently so is something we encourage explicit control over.
+
+  Returns:
+    An instance of xarray.DataArray. Where JAX arrays are used as data or
+    coords, they will be wrapped with JaxArrayWrapper and can be unwrapped via
+    `unwrap` and `unwrap_data`.
+  """
+  result = xarray.DataArray(
+      wrap(data), dims=dims, name=name, attrs=attrs or {})
+  return assign_coords(result, coords=coords, jax_coords=jax_coords)
+
+
+def Dataset(  # pylint:disable=invalid-name
+    data_vars=None,
+    coords=None,
+    attrs=None,
+    jax_coords=None,
+    ) -> xarray.Dataset:
+  """Like xarray.Dataset, but can wrap JAX arrays.
+
+  Args:
+    data_vars: As for xarray.Dataset, except jax arrays are also supported.
+    coords: Coordinates for the dataset, see xarray.Dataset. These coordinates
+      must be based on plain numpy arrays or something convertible to plain
+      numpy arrays. Their values will form a static part of the data structure
+      from the point of view of jax.tree_util. In particular this means these
+      coordinates will be passed as plain numpy arrays even inside a JIT'd
+      function, and the JIT'd function will be recompiled under the hood if the
+      coordinates of DataArrays passed into it change.
+      If this is not convenient for you, see also jax_coords below.
+    attrs: See xarray.Dataset.
+    jax_coords: Additional coordinates, which *can* use JAX arrays. These
+      coordinates will be treated as JAX data from the point of view of
+      jax.tree_util, that means when JIT'ing they will be passed as tracers and
+      computation involving them will be JIT'd.
+      Unfortunately a side-effect of this is that they can't be used as index
+      coordinates (because xarray's indexing logic is not JIT-able). If you
+      specify a coordinate with the same name as a dimension here, it will not
+      be set as an index coordinate; this behaviour is different to the default
+      for `coords`, and it means that things like `.sel` based on the jax
+      coordinate will not work.
+      Note we require `jax_coords` to be explicitly specified via a different
+      constructor argument to `coords`, rather than just looking for jax arrays
+      within the `coords` and treating them differently. This is because it
+      affects the way jax.tree_util treats them, which is somewhat orthogonal to
+      whether the value is passed in as numpy or not, and generally needs to be
+      handled consistently so is something we encourage explicit control over.
+
+  Returns:
+    An instance of xarray.Dataset. Where JAX arrays are used as data, they
+    will be wrapped with JaxArrayWrapper.
+  """
+  wrapped_data_vars = {}
+  for name, var_like in (data_vars or {}).items():
+    # xarray.Dataset accepts a few different formats for data_vars:
+    if isinstance(var_like, _WRAPPED_TYPES):
+      wrapped_data_vars[name] = wrap(var_like)
+    elif isinstance(var_like, tuple):
+      # Layout is (dims, data, ...). We wrap data.
+      wrapped_data_vars[name] = (var_like[0], wrap(var_like[1])) + var_like[2:]
+    else:
+      # Could be a plain numpy array or scalar (we don't wrap), or an
+      # xarray.Variable, DataArray etc, which we must assume is already wrapped
+      # if necessary (e.g. if creating using xarray_jax.{Variable,DataArray}).
+      wrapped_data_vars[name] = var_like
+
+  result = xarray.Dataset(
+      data_vars=wrapped_data_vars,
+      attrs=attrs)
+
+  return assign_coords(result, coords=coords, jax_coords=jax_coords)
+
+
+DatasetOrDataArray = TypeVar(
+    'DatasetOrDataArray', xarray.Dataset, xarray.DataArray)
+
+
+def assign_coords(
+    x: DatasetOrDataArray,
+    *,
+    coords: Optional[Mapping[Hashable, Any]] = None,
+    jax_coords: Optional[Mapping[Hashable, Any]] = None,
+    ) -> DatasetOrDataArray:
+  """Replacement for assign_coords which works in presence of jax_coords.
+
+  `jax_coords` allow certain specified coordinates to have their data passed as
+  JAX arrays (including through jax.jit boundaries). The compromise in return is
+  that they are not created as index coordinates and cannot be used for .sel
+  and other coordinate-based indexing operations. See docs for `jax_coords` on
+  xarray_jax.Dataset and xarray_jax.DataArray for more information.
+
+  This function can be used to set jax_coords on an existing DataArray or
+  Dataset, and also to set a mix of jax and non-jax coordinates. It implements
+  some workarounds to prevent xarray trying and failing to create IndexVariables
+  from jax arrays under the hood.
+
+  If you have any jax_coords with the same name as a dimension, you'll need to
+  use this function instead of data_array.assign_coords or dataset.assign_coords
+  in general, to avoid an xarray bug where it tries (and in our case fails) to
+  create indexes for existing jax coords. See
+  https://github.com/pydata/xarray/issues/7885.
+
+  Args:
+    x: An xarray Dataset or DataArray.
+    coords: Dict of (non-JAX) coords, or None if not assigning any.
+    jax_coords: Dict of JAX coords, or None if not assigning any. See docs for
+      xarray_jax.Dataset / DataArray for more information on jax_coords.
+
+  Returns:
+    The Dataset or DataArray with coordinates assigned, similarly to
+    Dataset.assign_coords / DataArray.assign_coords.
+  """
+  coords = {} if coords is None else dict(coords)  # Copy before mutating.
+  jax_coords = {} if jax_coords is None else dict(jax_coords)
+
+  # Any existing JAX coords must be dropped and re-added via the workaround
+  # below, since otherwise .assign_coords will trigger an xarray bug where
+  # it tries to recreate the indexes again for the existing coordinates.
+  # Can remove if/when https://github.com/pydata/xarray/issues/7885 fixed.
+  existing_jax_coords = get_jax_coords(x)
+  jax_coords = existing_jax_coords | jax_coords
+  x = x.drop_vars(existing_jax_coords.keys())
+
+  # We need to ensure that xarray doesn't try to create an index for
+  # coordinates with the same name as a dimension, since this will fail if
+  # given a wrapped JAX tracer.
+  # It appears the only way to avoid this is to name them differently to any
+  # dimension name, then rename them back afterwards.
+  renamed_jax_coords = {}
+  for name, coord in jax_coords.items():
+    if isinstance(coord, xarray.DataArray):
+      coord = coord.variable
+
+    if isinstance(coord, list):
+      coord = np.array(coord)
+
+    if isinstance(coord, xarray.Variable):
+      coord = coord.copy(deep=False)  # Copy before mutating attrs.
+    elif isinstance(coord, tuple):
+      # A tuple represents a pair of (dims, data).
+      dims, data = coord
+      coord = Variable(dims, data)
+    elif jnp.isscalar(coord):
+      # A scalar coord maps to a scalar Variable:
+      coord = Variable(dims=(), data=coord)
+    elif isinstance(coord, jax.typing.ArrayLike) and jnp.ndim(coord) == 1:
+      # A 1D array maps to a 1D Variable whose dimension is the same as the
+      # coordinate name:
+      coord = Variable((name,), coord)
+    else:
+      raise ValueError(f'Unsupported value for coordinate {name}')
+
+    # We set an attr on each jax_coord identifying it as such. These attrs on
+    # the coord Variable gets reflected on the coord DataArray exposed too, and
+    # when set on coordinates they generally get preserved under the default
+    # keep_attrs setting.
+    # These attrs are used by jax.tree_util registered flatten/unflatten to
+    # determine which coords need to be treated as leaves of the flattened
+    # structure vs static data.
+    coord.attrs[_JAX_COORD_ATTR_NAME] = True
+    renamed_jax_coords[f'__NONINDEX_{name}'] = coord
+
+  x = x.assign_coords(coords=coords | renamed_jax_coords)
+
+  rename_back_mapping = {f'__NONINDEX_{name}': name for name in jax_coords}
+  if isinstance(x, xarray.Dataset):
+    # Using 'rename' doesn't work if renaming to the same name as a dimension.
+    return x.rename_vars(rename_back_mapping)
+  else:  # DataArray
+    return x.rename(rename_back_mapping)
+
+
+def get_jax_coords(x: DatasetOrDataArray) -> Mapping[Hashable, Any]:
+  return {
+      name: coord_var
+      for name, coord_var in x.coords.variables.items()
+      if coord_var.attrs.get(_JAX_COORD_ATTR_NAME, False)}
+
+
+def assign_jax_coords(
+    x: DatasetOrDataArray,
+    jax_coords: Optional[Mapping[Hashable, Any]] = None,
+    **jax_coords_kwargs
+    ) -> DatasetOrDataArray:
+  """Assigns only jax_coords, with same API as xarray's assign_coords."""
+  return assign_coords(x, jax_coords=jax_coords or jax_coords_kwargs)
+
+
+def wrap(value):
+  """Wraps JAX arrays for use in xarray, passing through other values."""
+  if isinstance(value, _WRAPPED_TYPES):
+    return JaxArrayWrapper(value)
+  else:
+    return value
+
+
+def unwrap(value, require_jax=False):
+  """Unwraps wrapped JAX arrays used in xarray, passing through other values."""
+  if isinstance(value, JaxArrayWrapper):
+    return value.jax_array
+  elif isinstance(value, jax.Array):
+    return value
+  elif require_jax:
+    raise TypeError(f'Expected JAX array, found {type(value)}.')
+  else:
+    return value
+
+
+def _wrapped(func):
+  """Surrounds a function with JAX array unwrapping/wrapping."""
+  def wrapped_func(*args, **kwargs):
+    args, kwargs = tree.map_structure(unwrap, (args, kwargs))
+    result = func(*args, **kwargs)
+    return tree.map_structure(wrap, result)
+  return wrapped_func
+
+
+def unwrap_data(
+    value: Union[xarray.Variable, xarray.DataArray],
+    require_jax: bool = False
+    ) -> Union[jax.Array, np.ndarray]:
+  """The unwrapped (see unwrap) data of a an xarray.Variable or DataArray."""
+  return unwrap(value.data, require_jax=require_jax)
+
+
+def unwrap_vars(
+    dataset: Mapping[Hashable, xarray.DataArray],
+    require_jax: bool = False
+    ) -> Mapping[str, Union[jax.Array, np.ndarray]]:
+  """The unwrapped data (see unwrap) of the variables in a dataset."""
+  # xarray types variable names as Hashable, but in practice they're invariably
+  # strings and we convert to str to allow for a more useful return type.
+  return {str(name): unwrap_data(var, require_jax=require_jax)
+          for name, var in dataset.items()}
+
+
+def unwrap_coords(
+    dataset: Union[xarray.Dataset, xarray.DataArray],
+    require_jax: bool = False
+    ) -> Mapping[str, Union[jax.Array, np.ndarray]]:
+  """The unwrapped data (see unwrap) of the coords in a Dataset or DataArray."""
+  return {str(name): unwrap_data(var, require_jax=require_jax)
+          for name, var in dataset.coords.items()}
+
+
+def jax_data(value: Union[xarray.Variable, xarray.DataArray]) -> jax.Array:
+  """Like unwrap_data, but will complain if not a jax array."""
+  # Implementing this separately so we can give a more specific return type
+  # for it.
+  return cast(jax.Array, unwrap_data(value, require_jax=True))
+
+
+def jax_vars(
+    dataset: Mapping[Hashable, xarray.DataArray]) -> Mapping[str, jax.Array]:
+  """Like unwrap_vars, but will complain if vars are not all jax arrays."""
+  return cast(Mapping[str, jax.Array], unwrap_vars(dataset, require_jax=True))
+
+
+class JaxArrayWrapper(np.lib.mixins.NDArrayOperatorsMixin):
+  """Wraps a JAX array into a duck-typed array suitable for use with xarray.
+
+  This uses an older duck-typed array protocol based on __array_ufunc__ and
+  __array_function__ which works with numpy and xarray. (In newer versions
+  of xarray it implements xarray.namedarray._typing._array_function.)
+
+  This is in the process of being superseded by the Python array API standard
+  (https://data-apis.org/array-api/latest/index.html), but JAX hasn't
+  implemented it yet. Once they have, we should be able to get rid of
+  this wrapper and use JAX arrays directly with xarray.
+
+  """
+
+  def __init__(self, jax_array):
+    self.jax_array = jax_array
+
+  def __array_ufunc__(self, ufunc, method, *args, **kwargs):
+    for x in args:
+      if not isinstance(x, (jax.typing.ArrayLike, type(self))):
+        return NotImplemented
+    if method != '__call__':
+      return NotImplemented
+    try:
+      # Get the corresponding jax.numpy function to the NumPy ufunc:
+      func = getattr(jnp, ufunc.__name__)
+    except AttributeError:
+      return NotImplemented
+    # There may be an 'out' kwarg requesting an in-place operation, e.g. when
+    # this is called via __iadd__ (+=), __imul__ (*=) etc. JAX doesn't support
+    # in-place operations so we just remove this argument and have the ufunc
+    # return a fresh JAX array instead.
+    kwargs.pop('out', None)
+    return _wrapped(func)(*args, **kwargs)
+
+  def __array_function__(self, func, types, args, kwargs):
+    try:
+      # Get the corresponding jax.np function to the NumPy function:
+      func = getattr(jnp, func.__name__)
+    except AttributeError:
+      return NotImplemented
+    return _wrapped(func)(*args, **kwargs)
+
+  def __repr__(self):
+    return f'xarray_jax.JaxArrayWrapper({repr(self.jax_array)})'
+
+  # NDArrayOperatorsMixin already proxies most __dunder__ operator methods.
+  # We need to proxy through a few more methods in a similar way:
+
+  # Essential array properties:
+
+  @property
+  def shape(self):
+    return self.jax_array.shape
+
+  @property
+  def dtype(self):
+    return self.jax_array.dtype
+
+  @property
+  def ndim(self):
+    return self.jax_array.ndim
+
+  @property
+  def size(self):
+    return self.jax_array.size
+
+  @property
+  def real(self):
+    return self.jax_array.real
+
+  @property
+  def imag(self):
+    return self.jax_array.imag
+
+  # Array methods not covered by NDArrayOperatorsMixin:
+
+  # Allows conversion to numpy array using np.asarray etc. Warning: doing this
+  # will fail in a jax.jit-ed function.
+  def __array__(self, dtype=None, context=None):
+    return np.asarray(self.jax_array, dtype=dtype)
+
+  __getitem__ = _wrapped(lambda array, *args: array.__getitem__(*args))
+  # We drop the kwargs on this as they are not supported by JAX, but xarray
+  # uses at least one of them (the copy arg).
+  astype = _wrapped(lambda array, *args, **kwargs: array.astype(*args))
+
+  # There are many more methods which are more canonically available via (j)np
+  # functions, e.g. .sum() available via jnp.sum, and also mean, max, min,
+  # argmax, argmin etc. We don't attempt to proxy through all of these as
+  # methods, since this doesn't appear to be expected from a duck-typed array
+  # implementation. But there are a few which xarray calls as methods, so we
+  # proxy those:
+  transpose = _wrapped(jnp.transpose)
+  reshape = _wrapped(jnp.reshape)
+  all = _wrapped(jnp.all)
+
+
+def apply_ufunc(func, *args, require_jax=False, **apply_ufunc_kwargs):
+  """Like xarray.apply_ufunc but for jax-specific ufuncs.
+
+  Many numpy ufuncs will work fine out of the box with xarray_jax and
+  JaxArrayWrapper, since JaxArrayWrapper quacks (mostly) like a numpy array and
+  will convert many numpy operations to jax ops under the hood. For these
+  situations, xarray.apply_ufunc should work fine.
+
+  But sometimes you need a jax-specific ufunc which needs to be given a
+  jax array as input or return a jax array as output. In that case you should
+  use this helper as it will remove any JaxArrayWrapper before calling the func,
+  and wrap the result afterwards before handing it back to xarray.
+
+  Args:
+    func: A function that works with jax arrays (e.g. using functions from
+      jax.numpy) but otherwise meets the spec for the func argument to
+      xarray.apply_ufunc.
+    *args: xarray arguments to be mapped to arguments for func
+      (see xarray.apply_ufunc).
+    require_jax: Whether to require that inputs are based on jax arrays or allow
+      those based on plain numpy arrays too.
+    **apply_ufunc_kwargs: See xarray.apply_ufunc.
+
+  Returns:
+    Corresponding xarray results (see xarray.apply_ufunc).
+  """
+  def wrapped_func(*maybe_wrapped_args):
+    unwrapped_args = [unwrap(a, require_jax) for a in maybe_wrapped_args]
+    result = func(*unwrapped_args)
+    # Result can be an array or a tuple of arrays, this handles both:
+    return jax.tree_util.tree_map(wrap, result)
+  return xarray.apply_ufunc(wrapped_func, *args, **apply_ufunc_kwargs)
+
+
+def pmap(
+    fn: Callable[..., Any],
+    dim: str,
+    axis_name: Optional[str] = None,
+    devices=None,
+    backend=None,
+) -> Callable[..., Any]:
+  """Wraps a subset of jax.pmap functionality to handle xarray input/output.
+
+  Constraints:
+    * Any Dataset or DataArray passed to the function must have `dim` as the
+      first dimension. This will be checked. You can ensure this if necessary
+      by calling `.transpose(dim, ...)` beforehand.
+    * All args and return values will be mapped over the first dimension,
+      it will use in_axes=0, out_axes=0.
+    * No support for static_broadcasted_argnums, donate_argnums etc.
+
+  Args:
+    fn: Function to be pmap'd which takes and returns trees which may contain
+      xarray Dataset/DataArray. Any Dataset/DataArrays passed as input must use
+      `dim` as the first dimension on all arrays.
+    dim: The xarray dimension name corresponding to the first dimension that is
+      pmapped over (pmap is called with in_axes=0, out_axes=0).
+    axis_name: Used by jax to identify the mapped axis so that parallel
+      collectives can be applied. Defaults to same as `dim`.
+    devices:
+    backend:
+      See jax.pmap.
+
+  Returns:
+    A pmap'd version of `fn`, which takes and returns Dataset/DataArray with an
+    extra leading dimension `dim` relative to what the original `fn` sees.
+  """
+  return _vmap_or_pmap(
+      fn, dim, axis_name, devices, backend, is_vmap=False,
+  )
+
+
+def vmap(
+    fn: Callable[..., Any],
+    dim: str,
+    axis_name: Optional[str] = None,
+) -> Callable[..., Any]:
+  """Similar to pmap, but for vmap."""
+  return _vmap_or_pmap(
+      fn, dim, axis_name, None, None, is_vmap=True,
+  )
+
+
+def _vmap_or_pmap(
+    fn: Callable[..., Any],
+    dim: str,
+    axis_name: Optional[str] = None,
+    devices=None,
+    backend=None,
+    is_vmap: bool = False,
+) -> Callable[..., Any]:
+  """See pmap documentations."""
+
+  input_treedef = None
+  output_treedef = None
+
+  def fn_passed_to_pmap(*flat_args):
+    assert input_treedef is not None
+    # Inside the pmap the original first dimension will no longer be present:
+    def check_and_remove_leading_dim(dims):
+      try:
+        index = dims.index(dim)
+      except ValueError:
+        index = None
+      if index != 0:
+        raise ValueError(f'Expected dim {dim} at index 0, found at {index}.')
+      return dims[1:]
+    with dims_change_on_unflatten(check_and_remove_leading_dim):
+      args = jax.tree_util.tree_unflatten(input_treedef, flat_args)
+    result = fn(*args)
+    nonlocal output_treedef
+    flat_result, output_treedef = jax.tree_util.tree_flatten(result)
+    return flat_result
+
+  if is_vmap:
+    assert devices is None
+    assert backend is None
+    pmapped_fn = jax.vmap(
+        fn_passed_to_pmap,
+        axis_name=axis_name or dim,
+        in_axes=0,
+        out_axes=0)
+  else:
+    pmapped_fn = jax.pmap(
+        fn_passed_to_pmap,
+        axis_name=axis_name or dim,
+        in_axes=0,
+        out_axes=0,
+        devices=devices,
+        backend=backend)
+
+  def result_fn(*args):
+    nonlocal input_treedef
+    flat_args, input_treedef = jax.tree_util.tree_flatten(args)
+    flat_result = pmapped_fn(*flat_args)
+    assert output_treedef is not None
+    # After the pmap an extra leading axis will be present, we need to add an
+    # xarray dimension for this when unflattening the result:
+    with dims_change_on_unflatten(lambda dims: (dim,) + dims):
+      return jax.tree_util.tree_unflatten(output_treedef, flat_result)
+
+  return result_fn
+
+
+_PyTree = TypeVar('_PyTree')
+
+
+def tree_map_variables(
+    func: Callable[[xarray.Variable], xarray.Variable],
+    tree_data: _PyTree) -> _PyTree:
+  """Like jax.tree.map but operates with Variables as leaves.
+
+  This will work with any jax.tree_util-registered PyTree containing xarray
+  datatypes. All jax data in xarray datatypes is exposed via xarray.Variable
+  nodes by our registered flatten/unflatten functions and hence here too. Note
+  static coordinate data will not be mapped over however.
+
+  This allows you to see the associated dimensions for each leaf, and to change
+  them. If you change them, it's your responsibility to ensure that when
+  unflattened back into DataArray/Dataset/DataTree the result still makes sense.
+  In particular that any updated shapes are consistent with the shapes of any
+  static (non-jax_coord) coordinates, since these will not be mapped over.
+
+  Args:
+    func: Function from xarray.Variable to xarray.Variable.
+    tree_data: PyTree to be mapped over.
+
+  Returns:
+    PyTree with the same structure as `tree_data` but where xarray.Variables
+    within xarray datatypes have been mapped over by `func`. Any leaves outside
+    of xarray datatypes will be unchanged.
+  """
+  return jax.tree.map(
+      lambda leaf: func(leaf) if isinstance(leaf, xarray.Variable) else leaf,
+      tree_data,
+      is_leaf=lambda x: isinstance(x, xarray.Variable))
+
+
+def tree_map_with_dims(
+    func: Callable[[jax.typing.ArrayLike, tuple[str, ...] | None],
+                   jax.typing.ArrayLike],
+    data: _PyTree,
+) -> _PyTree:
+  """Like jax.tree.map but also passes in xarray dimensions where known.
+
+  This is convenient when applying logic to every jax array in some xarray data
+  structure, which wants to be sensitive to the xarray dimension names.
+  Typical examples of this would be jax operations relating to sharding where
+  you may want to map xarray dimension names to sharding axis names.
+
+  This only supports changing array shapes in limited situations (see below).
+
+  Unlike tree_map_variables above, this will also map over plain jax arrays
+  that don't occur within xarray.Variable nodes; these will be passed to func
+  with dims=None.
+
+  Args:
+    func: A function from (jax_array, dims) -> jax_array. dims will correspond
+      to the dimension names of the xarray.Variable containing the jax_array
+      where it occurs within an xarray.Variable, note this includes arrays
+      within xarray.Dataset and xarray.DataArray too. For plain jax arrays that
+      don't occur within an xarray.Variable, dims will be None.
+      The returned jax array should generally be of the same shape as the input.
+      However you can get away with changing the shape of a particular dimension
+      in limited circumstances: when there are no explicit coordinates involving
+      that dimension, or when the only coordinates involving that dimension are
+      jax_coords and you modify their shapes too in a consistent fashion.
+      You are not allowed to change the dimension order, add or remove
+      dimensions.
+    data: Any pytree with jax ArrayLike leaves suitable for use with
+      `jax.tree.map`. Thanks to xarray_jax such pytrees may include xarray
+      datatypes.
+
+  Returns:
+    A pytree of the same structure as data, with the result of applying func
+    to each jax array found.
+  """
+  # All jax arrays within xarray.Dataset, xarray.DataArray (including
+  # jax_coord arrays) will be exposed via xarray.Variable internal nodes by
+  # xarray_jax's pytree registrations. So to find xarray dimension metadata
+  # it's sufficient to stop descending at xarray.Variable nodes:
+  def is_leaf(x):
+    return isinstance(x, xarray.Variable)
+
+  def wrapped_func(x):
+    if isinstance(x, xarray.Variable):
+      array = unwrap(x.data)
+      array = func(array, x.dims)
+      return Variable(dims=x.dims, data=array)
+    else:
+      return func(x, None)
+
+  return jax.tree_util.tree_map(wrapped_func, data, is_leaf=is_leaf)
+
+
+_Carry = TypeVar('_Carry')
+_X = TypeVar('_X')
+_Y = TypeVar('_Y')
+
+
+def scan(f: Callable[[_Carry, _X], tuple[_Carry, _Y]],
+         init: _Carry,
+         dim: str,
+         xs: _X | None = None,
+         length: int | None = None,
+         reverse: bool = False,
+         unroll: int | bool = 1,
+         ) -> tuple[_Carry, _Y]:
+  """Like jax.lax.scan but supports xarray data.
+
+  This can handle a jax.tree containing any mix of xarray and plain jax data.
+  It scans along the dimension `dim` for xarray data, and the leading axis for
+  any non-xarray data. These scanned-along dimensions must all be consistent in
+  size.
+
+  Static coordinates along `dim` in the `xs` will not be present on the `x`
+  argument to `f`, since they would be different on each iteration and we can't
+  pass them through as static data. jax_coords will be passed through correctly
+  however.
+
+  Static coordinates along `dim` on the `xs` will also not be present on the
+  resulting `ys`, you will need to copy these across yourself if desired.
+  (This one we may be able to fix in future.)
+
+  Args:
+    f: Function to apply at each step of the scan. This should map
+      (carry, x) -> (carry, y), where `x` is a slice of the `xs` along the `dim`
+      axis (with the `dim` axis and any coordinates using it dropped), and `y`
+      is a slice of the desired output along the `dim` axis, with no `dim` axis
+      itself.
+      x, y and carry can in general be trees, in which the above applies to
+      each leaf of the tree.
+    init: Initial value of the carry.
+    dim: The xarray dimension name to scan along, for xarray data.
+    xs: The input to be scanned. If not provided, will scan over `length`
+      iterations with None passed as the `x` argument to `f`.
+    length: The length of the scan, if `xs` are not provided.
+    reverse: Whether to scan in reverse order.
+    unroll: How many steps to unroll the scan, see jax.lax.scan.
+
+  Returns:
+    final_carry: The carry returned from the final step of the scan.
+    ys: Data corresponding to the `y` returned from `f` on each step of the
+      scan, concatenated along an extra leading dimension (named `dim` for
+      xarray data).
+  """
+  if xs is not None:
+    # Ensure `dim` is the leading axis on any xarray data in the `xs`. This is
+    # what jax.lax.scan will scan over. (Any non-array data it's on you to put
+    # the relevant axis first).
+    xs = tree_map_variables(lambda v: v.transpose(dim, ...), xs)
+    xs_leaves, xs_treedef = jax.tree.flatten(xs)
+  else:
+    xs_treedef = None
+    xs_leaves = None
+
+  y_treedef = None
+
+  def scan_fn(carry, x_leaves):
+    if x_leaves is None:
+      x = None
+    else:
+      with dims_change_on_unflatten(lambda dims: dims[1:]):
+        x = jax.tree.unflatten(xs_treedef, x_leaves)
+    carry, y = f(carry, x)
+
+    nonlocal y_treedef
+    y_leaves, y_treedef = jax.tree.flatten(y)
+    return carry, y_leaves
+
+  final_carry, ys_leaves = jax.lax.scan(
+      scan_fn,
+      init,
+      xs_leaves,
+      length=length,
+      reverse=reverse,
+      unroll=unroll)
+
+  assert isinstance(y_treedef, jax.tree_util.PyTreeDef)
+
+  with dims_change_on_unflatten(lambda dims: (dim,) + dims):
+    ys = jax.tree.unflatten(y_treedef, ys_leaves)
+
+  return final_carry, ys
+
+
+# Register xarray datatypes with jax.tree_util.
+
+
+DimsChangeFn = Callable[[Tuple[Hashable, ...]], Tuple[Hashable, ...]]
+_DIMS_CHANGE_ON_UNFLATTEN_FN: contextvars.ContextVar[DimsChangeFn] = (
+    contextvars.ContextVar('dims_change_on_unflatten_fn'))
+
+
+@contextlib.contextmanager
+def dims_change_on_unflatten(dims_change_fn: DimsChangeFn):
+  """Can be used to change the dims used when unflattening arrays into xarrays.
+
+  This is useful when some axes were added to / removed from the underlying jax
+  arrays after they were flattened using jax.tree_util.tree_flatten, and you
+  want to unflatten them again afterwards using the original treedef but
+  adjusted for the added/removed dimensions.
+
+  It can also be used with jax.tree_util.tree_map, when it's called with a
+  function that adds/removes axes or otherwise changes the axis order.
+
+  When dimensions are removed, any coordinates using those removed dimensions
+  will also be removed on unflatten.
+
+  This is implemented as a context manager that sets some thread-local state
+  affecting the behaviour of our unflatten functions, because it's not possible
+  to directly modify the treedef to change the dims/coords in it (and with
+  tree_map, the treedef isn't exposed to you anyway).
+
+  Args:
+    dims_change_fn: Maps a tuple of dimension names for the original
+      Variable/DataArray/Dataset that was flattened, to an updated tuple of
+      dimensions which should be used when unflattening.
+
+  Yields:
+    To a context manager in whose scope jax.tree_util.tree_unflatten and
+    jax.tree_util.tree_map will apply the dims_change_fn before reconstructing
+    xarrays from jax arrays.
+  """
+  token = _DIMS_CHANGE_ON_UNFLATTEN_FN.set(dims_change_fn)
+  try:
+    yield
+  finally:
+    _DIMS_CHANGE_ON_UNFLATTEN_FN.reset(token)
+
+
+def _flatten_variable(v: xarray.Variable) -> Tuple[
+    Tuple[jax.typing.ArrayLike], Tuple[Hashable, ...]]:  # pylint: disable=g-one-element-tuple
+  """Flattens a Variable for jax.tree_util."""
+  children = (unwrap_data(v),)
+  aux = v.dims
+  return children, aux
+
+
+def _unflatten_variable(
+    aux: Tuple[Hashable, ...],
+    children: Tuple[jax.typing.ArrayLike]) -> xarray.Variable:  # pylint: disable=g-one-element-tuple
+  """Unflattens a Variable for jax.tree_util."""
+  dims = aux
+  dims_change_fn = _DIMS_CHANGE_ON_UNFLATTEN_FN.get(None)
+  if dims_change_fn: dims = dims_change_fn(dims)
+  return Variable(dims=dims, data=children[0])
+
+
+def _split_static_and_jax_coords(
+    coords: xarray.core.coordinates.Coordinates) -> Tuple[
+        Mapping[Hashable, xarray.Variable], Mapping[Hashable, xarray.Variable]]:
+  static_coord_vars = {}
+  jax_coord_vars = {}
+  for name, coord in coords.items():
+    if coord.attrs.get(_JAX_COORD_ATTR_NAME, False):
+      jax_coord_vars[name] = coord.variable
+    else:
+      assert not isinstance(coord, (jax.Array, JaxArrayWrapper))
+      static_coord_vars[name] = coord.variable
+  return static_coord_vars, jax_coord_vars
+
+
+def _drop_with_none_of_dims(
+    coord_vars: Mapping[Hashable, xarray.Variable],
+    dims: Tuple[Hashable, ...]) -> Mapping[Hashable, xarray.Variable]:
+  return {name: var for name, var in coord_vars.items()
+          if set(var.dims) <= set(dims)}
+
+
+class _HashableCoords(collections.abc.Mapping):
+  """Wraps a dict of xarray Variables as hashable, used for static coordinates.
+
+  This needs to be hashable so that when an xarray.Dataset is passed to a
+  jax.jit'ed function, jax can check whether it's seen an array with the
+  same static coordinates(*) before or whether it needs to recompile the
+  function for the new values of the static coordinates.
+
+  (*) note jax_coords are not included in this; their value can be different
+  on different calls without triggering a recompile.
+  """
+
+  def __init__(self, coord_vars: Mapping[Hashable, xarray.Variable]):
+    self._variables = coord_vars
+
+  def __repr__(self) -> str:
+    return f'_HashableCoords({repr(self._variables)})'
+
+  def __getitem__(self, key: Hashable) -> xarray.Variable:
+    return self._variables[key]
+
+  def __len__(self) -> int:
+    return len(self._variables)
+
+  def __iter__(self) -> Iterator[Hashable]:
+    return iter(self._variables)
+
+  def __hash__(self):
+    if not hasattr(self, '_hash'):
+      self._hash = hash(frozenset((name, var.data.tobytes())
+                                  for name, var in self._variables.items()))
+    return self._hash
+
+  def __eq__(self, other):
+    if self is other:
+      return True
+    elif not isinstance(other, type(self)):
+      return NotImplemented
+    elif self._variables is other._variables:
+      return True
+    else:
+      return self._variables.keys() == other._variables.keys() and all(
+          variable.equals(other._variables[name])
+          for name, variable in self._variables.items())
+
+
+def _flatten_data_array(v: xarray.DataArray) -> Tuple[
+    # Children (data variable, jax_coord_vars):
+    Tuple[xarray.Variable, Mapping[Hashable, xarray.Variable]],
+    # Static auxiliary data (name, static_coord_vars):
+    Tuple[Optional[Hashable], _HashableCoords]]:
+  """Flattens a DataArray for jax.tree_util."""
+  static_coord_vars, jax_coord_vars = _split_static_and_jax_coords(v.coords)
+  children = (v.variable, jax_coord_vars)
+  aux = (v.name, _HashableCoords(static_coord_vars))
+  return children, aux
+
+
+def _unflatten_data_array(
+    aux: Tuple[Optional[Hashable], _HashableCoords],
+    children: Tuple[xarray.Variable, Mapping[Hashable, xarray.Variable]],
+) -> xarray.DataArray:
+  """Unflattens a DataArray for jax.tree_util."""
+  variable, jax_coord_vars = children
+  name, static_coord_vars = aux
+  if _DIMS_CHANGE_ON_UNFLATTEN_FN.get(None):
+    # Drop static coords which have dims not present in any of the data_vars.
+    # These would generally be dims that were dropped by a dims_change_fn, but
+    # because static coordinates don't go through dims_change_fn on unflatten,
+    # we just drop them where this causes a problem.
+    # Since jax_coords go through the dims_change_fn on unflatten we don't need
+    # to do this for jax_coords.
+    static_coord_vars = _drop_with_none_of_dims(
+        static_coord_vars, variable.dims)
+  return DataArray(
+      variable, name=name, coords=static_coord_vars, jax_coords=jax_coord_vars)
+
+
+def _flatten_dataset(dataset: xarray.Dataset) -> Tuple[
+    # Children (data variables, jax_coord_vars):
+    Tuple[Mapping[Hashable, xarray.Variable],
+          Mapping[Hashable, xarray.Variable]],
+    # Static auxiliary data (static_coord_vars):
+    _HashableCoords]:
+  """Flattens a Dataset for jax.tree_util."""
+  variables = {name: data_array.variable
+               for name, data_array in dataset.data_vars.items()}
+  static_coord_vars, jax_coord_vars = _split_static_and_jax_coords(
+      dataset.coords)
+  children = (variables, jax_coord_vars)
+  aux = _HashableCoords(static_coord_vars)
+  return children, aux
+
+
+def _unflatten_dataset(
+    aux: _HashableCoords,
+    children: Tuple[Mapping[Hashable, xarray.Variable],
+                    Mapping[Hashable, xarray.Variable]],
+    ) -> xarray.Dataset:
+  """Unflattens a Dataset for jax.tree_util."""
+  data_vars, jax_coord_vars = children
+  static_coord_vars = aux
+  dataset = xarray.Dataset(data_vars)
+  if _DIMS_CHANGE_ON_UNFLATTEN_FN.get(None):
+    # Drop static coords which have dims not present in any of the data_vars.
+    # See corresponding comment in _unflatten_data_array.
+    static_coord_vars = _drop_with_none_of_dims(
+        static_coord_vars, dataset.dims)  # pytype: disable=wrong-arg-types
+  return assign_coords(
+      dataset, coords=static_coord_vars, jax_coords=jax_coord_vars)
+
+
+def _flatten_datatree(datatree: xarray.DataTree) -> Tuple[
+    Tuple[Mapping[str, xarray.DataTree], xarray.Dataset], str | None]:
+  """Flattens a DataTree for jax.tree_util."""
+  # For simplicity we assume DataTrees will be flattened/unflattened from the
+  # root. If you give it a non-root-node, it will still work but any parents
+  # (and any coordinates inherited from them) will be lost.
+  node_dataset = datatree.to_dataset(inherit=False)
+  children = (dict(datatree.children), node_dataset)
+  aux = datatree.name
+  return children, aux
+
+
+def _unflatten_datatree(
+    aux: str | None,
+    children: Tuple[Mapping[str, xarray.DataTree], xarray.Dataset],
+) -> xarray.DataTree:
+  """Unflattens a DataTree for jax.tree_util."""
+  children_dict, node_dataset = children
+  name = aux
+  return xarray.DataTree(
+      dataset=node_dataset, children=children_dict, name=name)
+
+
+jax.tree_util.register_pytree_node(
+    xarray.Variable, _flatten_variable, _unflatten_variable)
+# This is a subclass of Variable but still needs registering separately.
+# Flatten/unflatten for IndexVariable is a bit of a corner case but we do
+# need to support it.
+jax.tree_util.register_pytree_node(
+    xarray.IndexVariable, _flatten_variable, _unflatten_variable)
+jax.tree_util.register_pytree_node(
+    xarray.DataArray, _flatten_data_array, _unflatten_data_array)
+jax.tree_util.register_pytree_node(
+    xarray.Dataset, _flatten_dataset, _unflatten_dataset)
+jax.tree_util.register_pytree_node(
+    xarray.DataTree, _flatten_datatree, _unflatten_datatree)

model/graphcast/xarray_tree.py

@@ -0,0 +1,70 @@
+# Copyright 2023 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS-IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Utilities for working with trees of xarray.DataArray (including Datasets).
+
+Note that xarray.Dataset doesn't work out-of-the-box with the `tree` library;
+it won't work as a leaf node since it implements Mapping, but also won't work
+as an internal node since tree doesn't know how to re-create it properly.
+
+To fix this, we reimplement a subset of `map_structure`, exposing its
+constituent DataArrays as leaf nodes. This means it can be mapped over as a
+generic container of DataArrays, while still preserving the result as a Dataset
+where possible.
+
+This is useful because in a few places we need to handle a general
+Mapping[str, DataArray] (where the coordinates might not be compatible across
+the constituent DataArrays) but also the special case of a Dataset nicely.
+
+For the result e.g. of a tree.map_structure(fn, dataset), if fn returns None for
+some of the child DataArrays, they will be omitted from the returned dataset. If
+any values other than DataArrays or None are returned, then we don't attempt to
+return a Dataset and just return a plain dict of the results. Similarly if
+DataArrays are returned but with non-matching coordinates, it will just return a
+plain dict of DataArrays.
+
+Note xarray datatypes are registered with `jax.tree_util` by xarray_jax.py,
+but `jax.tree_util.tree_map` is distinct from the `xarray_tree.map_structure`.
+as the former exposes the underlying JAX/numpy arrays as leaf nodes, while the
+latter exposes DataArrays as leaf nodes.
+"""
+
+from typing import Any, Callable
+
+import xarray
+
+
+def map_structure(func: Callable[..., Any], *structures: Any) -> Any:
+  """Maps func through given structures with xarrays. See tree.map_structure."""
+  if not callable(func):
+    raise TypeError(f'func must be callable, got: {func}')
+  if not structures:
+    raise ValueError('Must provide at least one structure')
+
+  first = structures[0]
+  if isinstance(first, xarray.Dataset):
+    data = {k: func(*[s[k] for s in structures]) for k in first.keys()}
+    if all(isinstance(a, (type(None), xarray.DataArray))
+           for a in data.values()):
+      data_arrays = [v.rename(k) for k, v in data.items() if v is not None]
+      try:
+        return xarray.merge(data_arrays, join='exact')
+      except ValueError:  # Exact join not possible.
+        pass
+    return data
+  if isinstance(first, dict):
+    return {k: map_structure(func, *[s[k] for s in structures])
+            for k in first.keys()}
+  if isinstance(first, (list, tuple, set)):
+    return type(first)(map_structure(func, *s) for s in zip(*structures))
+  return func(*structures)

scripts/fake_data.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+"""生成确定性的 ERA5、GenCast 统计量和静态场。"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+import h5py
+import numpy as np
+import xarray
+import yaml
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+
+PRESSURE_LEVELS = (50, 100, 150, 200, 250, 300, 400, 500, 600, 700, 850, 925, 1000)
+SURFACE = (
+    "2m_temperature", "mean_sea_level_pressure", "10m_v_component_of_wind",
+    "10m_u_component_of_wind", "sea_surface_temperature", "total_precipitation",
+)
+ATMOSPHERIC = (
+    "temperature", "geopotential", "u_component_of_wind",
+    "v_component_of_wind", "vertical_velocity", "specific_humidity",
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", default=str(PROJECT_ROOT / "conf/config.yaml"))
+    parser.add_argument("--height", type=int)
+    parser.add_argument("--width", type=int)
+    parser.add_argument("--timesteps", type=int)
+    parser.add_argument("--seed", type=int, default=42)
+    return parser.parse_args()
+
+
+def _stat_dataset(value: float, *, minimum: bool = False) -> xarray.Dataset:
+    data_vars = {}
+    surface_targets = (
+        "2m_temperature", "mean_sea_level_pressure", "10m_v_component_of_wind",
+        "10m_u_component_of_wind", "sea_surface_temperature", "total_precipitation_12hr",
+    )
+    for name in surface_targets:
+        scalar = -3.0 if minimum and name == "sea_surface_temperature" else value
+        data_vars[name] = xarray.DataArray(np.float32(scalar))
+    for name in ATMOSPHERIC:
+        data_vars[name] = xarray.DataArray(
+            np.full(len(PRESSURE_LEVELS), value, dtype=np.float32),
+            dims=("level",),
+            coords={"level": np.asarray(PRESSURE_LEVELS, dtype=np.int32)},
+        )
+    # Generated forcings already in [-1, 1] — include so normalization is silent
+    for name in ("year_progress_sin", "year_progress_cos",
+                 "day_progress_sin", "day_progress_cos"):
+        data_vars[name] = xarray.DataArray(np.float32(value))
+    for name in ("geopotential_at_surface", "land_sea_mask"):
+        data_vars[name] = xarray.DataArray(np.float32(value))
+    return xarray.Dataset(data_vars)
+
+
+def main() -> None:
+    args = parse_args()
+    with Path(args.config).open(encoding="utf-8") as source:
+        config = yaml.safe_load(source)
+    root = PROJECT_ROOT / config["data"]["data_dir"]
+    height = int(args.height or config["fake_data"]["height"])
+    width = int(args.width or config["fake_data"]["width"])
+    timesteps = int(args.timesteps or config["fake_data"]["timesteps"])
+    if width != 2 * (height - 1):
+        raise ValueError("Synthetic GenCast grid must satisfy width=2*(height-1)")
+    variables = list(SURFACE) + [f"{name}_{level}" for name in ATMOSPHERIC for level in PRESSURE_LEVELS]
+    years = sorted(set(config["data"]["train_years"] + config["data"]["test_years"]))
+    rng = np.random.default_rng(args.seed)
+    (root / "data").mkdir(parents=True, exist_ok=True)
+    for year in years:
+        path = root / "data" / f"{year}.h5"
+        time = np.arange(timesteps, dtype=np.float32)[:, None, None, None]
+        channel = np.arange(len(variables), dtype=np.float32)[None, :, None, None]
+        lat = np.linspace(1.0, -1.0, height, dtype=np.float32)[None, None, :, None]
+        lon = np.linspace(0.0, 2.0 * np.pi, width, endpoint=False, dtype=np.float32)[None, None, None, :]
+        values = 0.01 * time + 0.001 * channel + 0.1 * lat + 0.05 * np.sin(lon)
+        values += rng.normal(0.0, 1e-4, values.shape).astype(np.float32)
+        precip_index = variables.index("total_precipitation")
+        values[:, precip_index] = np.maximum(values[:, precip_index], 0.0)
+        sst_index = variables.index("sea_surface_temperature")
+        values[:, sst_index, : height // 4] = np.nan
+        with h5py.File(path, "w") as output:
+            fields = output.create_dataset("fields", data=values, chunks=(1, len(variables), height, width))
+            fields.attrs["variables"] = variables
+            fields.attrs["time_step"] = 6
+        print(f"Generated {path} shape={values.shape}")
+
+    static_dir = root / "static"
+    static_dir.mkdir(parents=True, exist_ok=True)
+    np.save(
+        static_dir / "geopotential_at_surface.npy",
+        np.zeros((height, width), dtype=np.float32),
+    )
+    land = np.zeros((height, width), dtype=np.float32)
+    land[: height // 4] = 1.0
+    np.save(static_dir / "land_mask.npy", land)
+    stats_dir = root / "stats"
+    stats_dir.mkdir(parents=True, exist_ok=True)
+    _stat_dataset(0.0).to_netcdf(stats_dir / "mean_by_level.nc")
+    _stat_dataset(1.0).to_netcdf(stats_dir / "stddev_by_level.nc")
+    _stat_dataset(1.0).to_netcdf(stats_dir / "diffs_stddev_by_level.nc")
+    _stat_dataset(0.0, minimum=True).to_netcdf(stats_dir / "min_by_level.nc")
+    print(f"Generated GenCast statistics and static fields under {root}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,158 @@
+#!/usr/bin/env python3
+"""使用官方 GenCast DPM-Solver++ 执行集合自回归推理。"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+import warnings
+from pathlib import Path
+
+warnings.filterwarnings("ignore", message="Changing the sparsity structure")
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from model.common import configure_jax, load_config, load_stats, resolve_path
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", default=str(PROJECT_ROOT / "conf/config.yaml"))
+    parser.add_argument("--checkpoint")
+    parser.add_argument("--sample-index", type=int, default=0)
+    parser.add_argument("--num-members", type=int)
+    parser.add_argument("--prediction-steps", type=int)
+    parser.add_argument("--output")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    config = load_config(args.config)
+    configure_jax(config["runtime"].get("platform", "auto"))
+
+    import jax
+    import numpy as np
+    import xarray
+
+    from model.graphcast import rollout
+    from model.gencast import GenCastModel, load_model_checkpoint
+    from model.common import (
+        load_trainer_checkpoint, validate_checkpoint_config,
+    )
+    from model.data_loader import GenCastERA5Dataset
+
+    prediction_steps = int(args.prediction_steps or config["inference"]["prediction_steps"])
+    num_members = int(args.num_members or config["inference"]["num_members"])
+    stats = load_stats(config["data"]["stats_dir"])
+    checkpoint_path = args.checkpoint or config["inference"].get("official_checkpoint")
+    if checkpoint_path:
+        official = load_model_checkpoint(resolve_path(checkpoint_path))
+        model = GenCastModel.from_checkpoint_and_stats(
+            official,
+            stats,
+            attention_type=config["inference"].get("attention_type_override"),
+        )
+        params, state = official.params, {}
+        task_config = official.task_config
+    else:
+        model = GenCastModel.from_config_and_stats(config, stats)
+        params, state, _, _, saved_config = load_trainer_checkpoint(
+            config["checkpoint"]["trainer"]
+        )
+        validate_checkpoint_config(config, saved_config, scope="inference")
+        task_config = model.task_config
+
+    dataset = GenCastERA5Dataset(
+        resolve_path(config["data"]["data_dir"]),
+        list(config["data"]["test_years"]),
+        static_dir=resolve_path(config["data"]["static_dir"]),
+        prediction_steps=prediction_steps,
+        stride=int(config["data"].get("test_stride", 1)),
+        task_config=task_config,
+        precipitation_interval_hours=int(
+            config["data"]["precipitation_interval_hours"]
+        ),
+        load_future_targets=False,
+    )
+    inputs, targets, forcings = dataset[args.sample_index]
+
+    def forward(rng, inputs, targets_template, forcings):
+        return model.predict(
+            params, state, rng, inputs, targets_template, forcings
+        )[0]
+
+    forward = jax.jit(forward)
+    seed = int(config["inference"]["seed"])
+    rngs = np.stack([jax.random.fold_in(jax.random.PRNGKey(seed), i) for i in range(num_members)])
+    chunks = rollout.chunked_prediction_generator_multiple_runs(
+        predictor_fn=forward,
+        rngs=rngs,
+        inputs=inputs,
+        targets_template=targets * np.nan,
+        forcings=forcings,
+        num_steps_per_chunk=1,
+        num_samples=num_members,
+        pmap_devices=None,
+    )
+    output = resolve_path(args.output or config["output"]["prediction"])
+    if bool(config["inference"].get("stream_chunks", True)):
+        output_dir = output.with_suffix("")
+        output_dir.mkdir(parents=True, exist_ok=True)
+        for chunk_index, chunk in enumerate(chunks):
+            host_chunk = jax.device_get(chunk)
+            member = int(host_chunk.coords["sample"])
+            lead = int(host_chunk.time.values[0] / np.timedelta64(1, "h"))
+            host_chunk = host_chunk.drop_vars("sample").assign_coords(time=[lead])
+            host_chunk.coords["time"].attrs = {"long_name": "forecast lead time hours"}
+            host_chunk.attrs.update(
+                model="GenCast", target_channel_count=84,
+                forecast_reference_time=inputs.attrs["forecast_reference_time"],
+            )
+            path = output_dir / f"member_{member:03d}_lead_{lead:04d}h.nc"
+            host_chunk.to_netcdf(path)
+            print(f"Saved prediction chunk to {path}")
+        return
+
+    chunks = list(chunks)
+    member_chunks: list[list[xarray.Dataset]] = [[] for _ in range(num_members)]
+    for chunk in chunks:
+        host_chunk = jax.device_get(chunk)
+        member = int(host_chunk.coords["sample"])
+        member_chunks[member].append(host_chunk.drop_vars("sample"))
+    members = [
+        xarray.concat(parts, dim="time").expand_dims(sample=[member])
+        for member, parts in enumerate(member_chunks)
+    ]
+    predictions = xarray.concat(members, dim="sample")
+    predictions.attrs.update(
+        model="GenCast",
+        target_channel_count=84,
+        ensemble_members=num_members,
+        step_hours=12,
+        forecast_reference_time=inputs.attrs["forecast_reference_time"],
+    )
+    # Store lead time as plain hours; xarray_jax's internal dtype attribute is
+    # not valid CF metadata and conflicts with decoding after NetCDF round-trip.
+    lead_hours = (
+        predictions.coords["time"].values / np.timedelta64(1, "h")
+    ).astype(np.int32)
+    predictions = predictions.assign_coords(time=("time", lead_hours))
+    predictions.coords["time"].attrs = {
+        "long_name": "forecast lead time",
+        "units": "hours",
+    }
+    reference_time = np.datetime64(inputs.attrs["forecast_reference_time"])
+    predictions = predictions.assign_coords(
+        valid_time=("time", reference_time + lead_hours.astype("timedelta64[h]"))
+    )
+    output.parent.mkdir(parents=True, exist_ok=True)
+    temporary = output.with_suffix(output.suffix + ".tmp")
+    predictions.to_netcdf(temporary)
+    temporary.replace(output)
+    print(f"Saved predictions to {output}")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,87 @@
+#!/usr/bin/env python3
+"""可视化 GenCast 集合预报结果。"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+import warnings
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import xarray
+
+warnings.filterwarnings("ignore", message="Changing the sparsity structure")
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from model.common import load_config, resolve_path
+
+
+def load_predictions(path: Path) -> xarray.Dataset:
+    """Load predictions from either a single NetCDF or a chunked directory."""
+    # Auto-detect: if path.nc doesn't exist but path/ directory does, use chunks
+    if not path.exists() and path.suffix == ".nc":
+        chunk_dir = path.with_suffix("")
+        if chunk_dir.is_dir():
+            path = chunk_dir
+    if path.is_dir():
+        files = sorted(path.glob("member_*_lead_*.nc"))
+        if not files:
+            raise FileNotFoundError(f"No prediction chunks found in {path}")
+        datasets = [xarray.load_dataset(f) for f in files]
+        members = sorted(set(int(f.stem.split("_")[1]) for f in files))
+        member_parts = []
+        for m in members:
+            member_ds = [ds for ds, f in zip(datasets, files)
+                         if f"member_{m:03d}_lead_" in str(f)]
+            member_parts.append(xarray.concat(member_ds, dim="time"))
+        return xarray.concat(member_parts, dim="sample")
+    else:
+        return xarray.load_dataset(path)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", default=str(PROJECT_ROOT / "conf/config.yaml"))
+    parser.add_argument("--prediction")
+    parser.add_argument("--variable", default="2m_temperature")
+    parser.add_argument("--output")
+    parser.add_argument("--lead", type=int, default=0, help="lead time index (0-based)")
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    config = load_config(args.config)
+    prediction_path = resolve_path(args.prediction or config["output"]["prediction"])
+    output_path = resolve_path(args.output or config["output"]["plot"])
+    prediction = load_predictions(prediction_path)[args.variable]
+
+    sample_dim = "sample" if "sample" in prediction.dims else None
+    mean = prediction.mean(sample_dim) if sample_dim else prediction
+    spread = prediction.std(sample_dim) if sample_dim else xarray.zeros_like(mean)
+
+    # Select lead time index (0=first, -1=last)
+    lead_idx = args.lead
+    field = mean.isel(batch=0, time=lead_idx).values
+    spread_field = spread.isel(batch=0, time=lead_idx).values
+
+    fig, axes = plt.subplots(1, 2, figsize=(12, 4), constrained_layout=True)
+    im0 = axes[0].imshow(field, origin="lower", cmap="viridis", aspect="auto")
+    axes[0].set_title(f"GenCast ensemble mean: {args.variable}")
+    fig.colorbar(im0, ax=axes[0], orientation="horizontal")
+    im1 = axes[1].imshow(spread_field, origin="lower", cmap="magma", aspect="auto")
+    axes[1].set_title("Ensemble spread")
+    fig.colorbar(im1, ax=axes[1], orientation="horizontal")
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    fig.savefig(output_path, dpi=180)
+    plt.close(fig)
+    print(f"Saved plot to {output_path}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,325 @@
+#!/usr/bin/env python3
+"""使用官方单步 EDM 去噪目标训练 GenCast。"""
+
+from __future__ import annotations
+
+import argparse
+import itertools
+import sys
+import warnings
+from pathlib import Path
+
+import xarray
+
+# Mesh adjacency construction triggers one-time scipy CSR restructure warning.
+warnings.filterwarnings("ignore", message="Changing the sparsity structure")
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from model.common import configure_jax, load_config, load_stats, resolve_path
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--config", default=str(PROJECT_ROOT / "conf/config.yaml"))
+    parser.add_argument("--max-steps", type=int)
+    parser.add_argument("--resume")
+    parser.add_argument("--parallel-mode", choices=("single", "pmap"))
+    parser.add_argument("--num-devices", type=int)
+    parser.add_argument("--global-batch-size", type=int)
+    parser.add_argument("--checkpoint")
+    parser.add_argument("--seed", type=int)
+    return parser.parse_args()
+
+
+def _adam_init(params):
+    import jax
+    import jax.numpy as jnp
+
+    zeros = jax.tree_util.tree_map(jnp.zeros_like, params)
+    return {"count": jnp.asarray(0, dtype=jnp.int32), "mu": zeros, "nu": zeros}
+
+
+def _adam_update(params, grads, state, learning_rate, beta1, beta2, eps):
+    import jax
+    import jax.numpy as jnp
+
+    count = state["count"] + 1
+    mu = jax.tree_util.tree_map(
+        lambda old, grad: beta1 * old + (1.0 - beta1) * grad,
+        state["mu"], grads,
+    )
+    nu = jax.tree_util.tree_map(
+        lambda old, grad: beta2 * old + (1.0 - beta2) * jnp.square(grad),
+        state["nu"], grads,
+    )
+    mu_hat = jax.tree_util.tree_map(lambda value: value / (1.0 - beta1**count), mu)
+    nu_hat = jax.tree_util.tree_map(lambda value: value / (1.0 - beta2**count), nu)
+    params = jax.tree_util.tree_map(
+        lambda value, first, second: value - learning_rate * first / (jnp.sqrt(second) + eps),
+        params, mu_hat, nu_hat,
+    )
+    return params, {"count": count, "mu": mu, "nu": nu}
+
+
+def _replicate(tree, devices):
+    import jax
+
+    return jax.device_put_replicated(tree, devices)
+
+
+def _unreplicate(tree):
+    import jax
+
+    return jax.tree_util.tree_map(lambda value: value[0], tree)
+
+
+def _device_batch(batch, device_count):
+    """Add a leading device dimension to each GenCast xarray input."""
+    result = []
+    for value in batch:
+        if not isinstance(value, xarray.Dataset):
+            raise TypeError("GenCast batches must contain xarray.Dataset values")
+        value = value.transpose("batch", ...)
+        if "batch" not in value.dims:
+            value = value.expand_dims("batch")
+        if value.sizes["batch"] % device_count:
+            raise ValueError("Batch size must be divisible by the device count")
+        local_batch = value.sizes["batch"] // device_count
+        shards = [
+            value.isel(batch=slice(index * local_batch, (index + 1) * local_batch))
+            for index in range(device_count)
+        ]
+        result.append(xarray.concat(shards, dim="device"))
+    return tuple(result)
+
+
+def main() -> None:
+    args = parse_args()
+    config = load_config(args.config)
+    parallel = config.setdefault("parallel", {})
+    if args.parallel_mode is not None:
+        parallel["mode"] = args.parallel_mode
+    if args.num_devices is not None:
+        parallel["num_devices"] = args.num_devices
+    if args.global_batch_size is not None:
+        parallel["global_batch_size"] = args.global_batch_size
+    if args.checkpoint is not None:
+        config["checkpoint"]["trainer"] = args.checkpoint
+    if args.seed is not None:
+        config["training"]["seed"] = args.seed
+    configure_jax(config["runtime"].get("platform", "auto"))
+
+    import jax
+    import jax.numpy as jnp
+
+    from model.gencast import GenCastModel, parameter_count
+    from model.common import (
+        load_trainer_checkpoint, save_trainer_checkpoint,
+        validate_checkpoint_config,
+    )
+    from model.data_loader import GenCastERA5Dataset, batch_iterator
+
+    mode = str(parallel.get("mode", "single")).lower()
+    if mode not in ("single", "pmap"):
+        raise ValueError("parallel.mode must be 'single' or 'pmap'")
+    devices = list(jax.local_devices())
+    requested_devices = int(parallel.get("num_devices", 1))
+    if requested_devices < 1:
+        raise ValueError("parallel.num_devices must be positive")
+    if mode == "pmap":
+        if requested_devices > len(devices):
+            raise ValueError(
+                f"Requested {requested_devices} devices, only {len(devices)} available"
+            )
+        devices = devices[:requested_devices]
+    else:
+        requested_devices = 1
+        devices = devices[:1]
+    global_batch_size = int(parallel.get("global_batch_size", requested_devices))
+    if global_batch_size < 1 or global_batch_size % requested_devices:
+        raise ValueError("global_batch_size must be divisible by the device count")
+    stats = load_stats(config["data"]["stats_dir"])
+    model = GenCastModel.from_config_and_stats(config, stats)
+    dataset = GenCastERA5Dataset(
+        resolve_path(config["data"]["data_dir"]),
+        list(config["data"]["train_years"]),
+        static_dir=resolve_path(config["data"]["static_dir"]),
+        prediction_steps=1,
+        stride=int(config["data"].get("train_stride", 1)),
+        precipitation_interval_hours=int(
+            config["data"]["precipitation_interval_hours"]
+        ),
+    )
+    first_batch = dataset[0]
+    seed = int(config["training"]["seed"])
+    start_step = 0
+    resume = args.resume or config["checkpoint"].get("resume")
+    if resume:
+        params, state, optimizer_state, start_step, saved_config = \
+            load_trainer_checkpoint(resume)
+        validate_checkpoint_config(config, saved_config)
+    else:
+        params, state = model.init(
+            jax.random.fold_in(jax.random.PRNGKey(seed), -1), *first_batch
+        )
+        optimizer_state = _adam_init(params)
+
+    learning_rate = float(config["training"]["learning_rate"])
+    beta1, beta2 = (float(value) for value in config["training"]["betas"])
+    epsilon = float(config["training"].get("epsilon", 1e-8))
+
+    def train_step(params, state, optimizer_state, rng, inputs, targets, forcings):
+        def objective(current_params, current_state):
+            (loss, diagnostics), next_state = model.loss(
+                current_params, current_state, rng, inputs, targets, forcings
+            )
+            return loss, (diagnostics, next_state)
+
+        (loss, (diagnostics, next_state)), grads = jax.value_and_grad(
+            objective, has_aux=True
+        )(params, state)
+        finite = jnp.logical_and(
+            jnp.isfinite(loss),
+            jnp.all(jnp.asarray([jnp.all(jnp.isfinite(x)) for x in jax.tree_util.tree_leaves(grads)])),
+        )
+        new_params, new_optimizer_state = _adam_update(
+            params, grads, optimizer_state, learning_rate, beta1, beta2, epsilon
+        )
+        params = jax.tree_util.tree_map(
+            lambda new, old: jnp.where(finite, new, old), new_params, params
+        )
+        next_state = jax.tree_util.tree_map(
+            lambda new, old: jnp.where(finite, new, old), next_state, state
+        )
+        new_optimizer_state = jax.tree_util.tree_map(
+            lambda new, old: jnp.where(finite, new, old),
+            new_optimizer_state,
+            optimizer_state,
+        )
+        return params, next_state, new_optimizer_state, loss, diagnostics, finite
+
+    if mode == "pmap":
+        axis_name = str(parallel.get("axis_name", "devices"))
+
+        def parallel_train_step(
+            params, state, optimizer_state, rng, inputs, targets, forcings
+        ):
+            rng = jax.random.fold_in(rng, jax.lax.axis_index(axis_name))
+
+            def objective(current_params, current_state):
+                (loss, diagnostics), next_state = model.loss(
+                    current_params, current_state, rng, inputs, targets, forcings
+                )
+                return loss, (diagnostics, next_state)
+
+            (loss, (diagnostics, next_state)), grads = jax.value_and_grad(
+                objective, has_aux=True
+            )(params, state)
+            grads = jax.lax.pmean(grads, axis_name)
+            loss = jax.lax.pmean(loss, axis_name)
+            diagnostics = jax.tree_util.tree_map(
+                lambda value: jax.lax.pmean(value, axis_name), diagnostics
+            )
+            next_state = jax.tree_util.tree_map(
+                lambda value: jax.lax.pmean(value, axis_name), next_state
+            )
+            finite = jnp.logical_and(
+                jnp.isfinite(loss),
+                jnp.all(jnp.asarray([
+                    jnp.all(jnp.isfinite(x))
+                    for x in jax.tree_util.tree_leaves(grads)
+                ])),
+            )
+            finite = jax.lax.pmin(finite, axis_name)
+            new_params, new_optimizer_state = _adam_update(
+                params, grads, optimizer_state, learning_rate, beta1, beta2, epsilon
+            )
+            params = jax.tree_util.tree_map(
+                lambda new, old: jnp.where(finite, new, old), new_params, params
+            )
+            next_state = jax.tree_util.tree_map(
+                lambda new, old: jnp.where(finite, new, old), next_state, state
+            )
+            new_optimizer_state = jax.tree_util.tree_map(
+                lambda new, old: jnp.where(finite, new, old),
+                new_optimizer_state,
+                optimizer_state,
+            )
+            return params, next_state, new_optimizer_state, loss, diagnostics, finite
+
+        from model.graphcast import xarray_jax
+
+        train_step = xarray_jax.pmap(
+            parallel_train_step, dim="device", axis_name=axis_name, devices=devices
+        )
+    else:
+        train_step = jax.jit(train_step)
+    max_steps = int(args.max_steps or config["training"]["max_steps"])
+    save_interval = int(config["training"].get("save_interval", max_steps))
+    checkpoint_path = config["checkpoint"]["trainer"]
+    print(f"Training samples: {len(dataset)}; parameters: {parameter_count(params):,}")
+    if mode == "pmap":
+        params = _replicate(params, devices)
+        state = _replicate(state, devices)
+        optimizer_state = _replicate(optimizer_state, devices)
+        print(
+            f"Parallel mode: pmap; devices: {requested_devices}; "
+            f"global batch: {global_batch_size}"
+        )
+
+    step = start_step
+    batches_per_epoch = len(dataset) // global_batch_size
+    if batches_per_epoch < 1:
+        raise ValueError(
+            f"Dataset has {len(dataset)} samples, fewer than global_batch_size "
+            f"{global_batch_size}"
+        )
+    while step < max_steps:
+        epoch = step // batches_per_epoch
+        offset = step % batches_per_epoch
+        epoch_batches = batch_iterator(
+            dataset,
+            shuffle=True,
+            seed=seed + epoch,
+            batch_size=global_batch_size,
+        )
+        for batch in itertools.islice(epoch_batches, offset, None):
+            if step >= max_steps:
+                break
+            step_rng = jax.random.fold_in(jax.random.PRNGKey(seed), step)
+            if mode == "pmap":
+                batch = _device_batch(batch, requested_devices)
+                step_rng = jax.numpy.broadcast_to(
+                    step_rng, (requested_devices, *step_rng.shape)
+                )
+                params, state, optimizer_state, loss, _, finite = train_step(
+                    params, state, optimizer_state, step_rng, *batch
+                )
+                loss, finite = loss[0], finite[0]
+            else:
+                params, state, optimizer_state, loss, _, finite = train_step(
+                    params, state, optimizer_state, step_rng, *batch
+                )
+            step += 1
+            print(f"step={step} loss={float(loss):.8f} finite={bool(finite)}")
+            if not bool(finite):
+                raise FloatingPointError(f"Non-finite GenCast loss at step {step}")
+            if step % save_interval == 0 or step == max_steps:
+                checkpoint_trees = (params, state, optimizer_state)
+                if mode == "pmap":
+                    checkpoint_trees = tuple(map(_unreplicate, checkpoint_trees))
+                save_trainer_checkpoint(
+                    checkpoint_path,
+                    params=checkpoint_trees[0],
+                    state=checkpoint_trees[1],
+                    optimizer_state=checkpoint_trees[2],
+                    step=step,
+                    config=config,
+                )
+                print(f"Saved checkpoint to {resolve_path(checkpoint_path)}")
+
+
+if __name__ == "__main__":
+    main()