gencast_mini_demo.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9KHpdDlzYNDI"
      },
      "source": [
        "\u003e Copyright 2024 DeepMind Technologies Limited.\n",
        "\u003e\n",
        "\u003e Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "\u003e you may not use this file except in compliance with the License.\n",
        "\u003e You may obtain a copy of the License at\n",
        "\u003e\n",
        "\u003e      http://www.apache.org/licenses/LICENSE-2.0\n",
        "\u003e\n",
        "\u003e Unless required by applicable law or agreed to in writing, software\n",
        "\u003e distributed under the License is distributed on an \"AS-IS\" BASIS,\n",
        "\u003e WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "\u003e See the License for the specific language governing permissions and\n",
        "\u003e limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GnMHywhUhlgJ"
      },
      "source": [
        "# GenCast Mini Demo\n",
        "\n",
        "This notebook demonstrates running `GenCast 1p0deg Mini \u003c2019`.\n",
        "\n",
        "`GenCast 1p0deg Mini \u003c2019` is a GenCast model at 1deg resolution, with 13 pressure levels and a 4 times refined icosahedral mesh. It is trained on ERA5 data from 1979 to 2018, and can be causally evaluated on 2019 and later years.\n",
        "\n",
        "While other GenCast models are [available](https://github.com/google-deepmind/graphcast/blob/main/README.md), this model has the smallest memory footprint of those provided and is the only one runnable with the freely provided TPUv2-8 configuration in Colab. You can select this configuration in `Runtime\u003eChange Runtime Type`.\n",
        "\n",
        "**N.B.** The performance of `GenCast 1p0deg Mini \u003c2019` is reasonable but is not representative of the performance of the other GenCast models described in the [README](https://github.com/google-deepmind/graphcast/blob/main/README.md).\n",
        "\n",
        "To run the other models using Google Cloud Compute, refer to [gencast_demo_cloud_vm.ipynb](https://colab.research.google.com/github/deepmind/graphcast/blob/master/gencast_demo_cloud_vm.ipynb)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yMbbXFl4msJw"
      },
      "source": [
        "# Installation and Initialization"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-gAH79SRwp9G"
      },
      "outputs": [],
      "source": [
        "# @title Upgrade packages (kernel needs to be restarted after running this cell).\n",
        "\n",
        "%pip install -U importlib_metadata"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "233zaiZYqCnc"
      },
      "outputs": [],
      "source": [
        "# @title Pip install repo and dependencies\n",
        "\n",
        "%pip install --upgrade https://github.com/deepmind/graphcast/archive/master.zip"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Z_j8ej4Pyg1L"
      },
      "outputs": [],
      "source": [
        "# @title Imports\n",
        "\n",
        "import dataclasses\n",
        "import datetime\n",
        "import math\n",
        "from google.cloud import storage\n",
        "from typing import Optional\n",
        "import haiku as hk\n",
        "from IPython.display import HTML\n",
        "from IPython import display\n",
        "import ipywidgets as widgets\n",
        "import jax\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "from matplotlib import animation\n",
        "import numpy as np\n",
        "import xarray\n",
        "\n",
        "from graphcast import rollout\n",
        "from graphcast import xarray_jax\n",
        "from graphcast import normalization\n",
        "from graphcast import checkpoint\n",
        "from graphcast import data_utils\n",
        "from graphcast import xarray_tree\n",
        "from graphcast import gencast\n",
        "from graphcast import denoiser\n",
        "from graphcast import nan_cleaning\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "OzYgQ0QN-kn8"
      },
      "outputs": [],
      "source": [
        "# @title Plotting functions\n",
        "\n",
        "def select(\n",
        "    data: xarray.Dataset,\n",
        "    variable: str,\n",
        "    level: Optional[int] = None,\n",
        "    max_steps: Optional[int] = None\n",
        "    ) -\u003e xarray.Dataset:\n",
        "  data = data[variable]\n",
        "  if \"batch\" in data.dims:\n",
        "    data = data.isel(batch=0)\n",
        "  if max_steps is not None and \"time\" in data.sizes and max_steps \u003c data.sizes[\"time\"]:\n",
        "    data = data.isel(time=range(0, max_steps))\n",
        "  if level is not None and \"level\" in data.coords:\n",
        "    data = data.sel(level=level)\n",
        "  return data\n",
        "\n",
        "def scale(\n",
        "    data: xarray.Dataset,\n",
        "    center: Optional[float] = None,\n",
        "    robust: bool = False,\n",
        "    ) -\u003e tuple[xarray.Dataset, matplotlib.colors.Normalize, str]:\n",
        "  vmin = np.nanpercentile(data, (2 if robust else 0))\n",
        "  vmax = np.nanpercentile(data, (98 if robust else 100))\n",
        "  if center is not None:\n",
        "    diff = max(vmax - center, center - vmin)\n",
        "    vmin = center - diff\n",
        "    vmax = center + diff\n",
        "  return (data, matplotlib.colors.Normalize(vmin, vmax),\n",
        "          (\"RdBu_r\" if center is not None else \"viridis\"))\n",
        "\n",
        "def plot_data(\n",
        "    data: dict[str, xarray.Dataset],\n",
        "    fig_title: str,\n",
        "    plot_size: float = 5,\n",
        "    robust: bool = False,\n",
        "    cols: int = 4\n",
        "    ) -\u003e tuple[xarray.Dataset, matplotlib.colors.Normalize, str]:\n",
        "\n",
        "  first_data = next(iter(data.values()))[0]\n",
        "  max_steps = first_data.sizes.get(\"time\", 1)\n",
        "  assert all(max_steps == d.sizes.get(\"time\", 1) for d, _, _ in data.values())\n",
        "\n",
        "  cols = min(cols, len(data))\n",
        "  rows = math.ceil(len(data) / cols)\n",
        "  figure = plt.figure(figsize=(plot_size * 2 * cols,\n",
        "                               plot_size * rows))\n",
        "  figure.suptitle(fig_title, fontsize=16)\n",
        "  figure.subplots_adjust(wspace=0, hspace=0)\n",
        "  figure.tight_layout()\n",
        "\n",
        "  images = []\n",
        "  for i, (title, (plot_data, norm, cmap)) in enumerate(data.items()):\n",
        "    ax = figure.add_subplot(rows, cols, i+1)\n",
        "    ax.set_xticks([])\n",
        "    ax.set_yticks([])\n",
        "    ax.set_title(title)\n",
        "    im = ax.imshow(\n",
        "        plot_data.isel(time=0, missing_dims=\"ignore\"), norm=norm,\n",
        "        origin=\"lower\", cmap=cmap)\n",
        "    plt.colorbar(\n",
        "        mappable=im,\n",
        "        ax=ax,\n",
        "        orientation=\"vertical\",\n",
        "        pad=0.02,\n",
        "        aspect=16,\n",
        "        shrink=0.75,\n",
        "        cmap=cmap,\n",
        "        extend=(\"both\" if robust else \"neither\"))\n",
        "    images.append(im)\n",
        "\n",
        "  def update(frame):\n",
        "    if \"time\" in first_data.dims:\n",
        "      td = datetime.timedelta(microseconds=first_data[\"time\"][frame].item() / 1000)\n",
        "      figure.suptitle(f\"{fig_title}, {td}\", fontsize=16)\n",
        "    else:\n",
        "      figure.suptitle(fig_title, fontsize=16)\n",
        "    for im, (plot_data, norm, cmap) in zip(images, data.values()):\n",
        "      im.set_data(plot_data.isel(time=frame, missing_dims=\"ignore\"))\n",
        "\n",
        "  ani = animation.FuncAnimation(\n",
        "      fig=figure, func=update, frames=max_steps, interval=250)\n",
        "  plt.close(figure.number)\n",
        "  return HTML(ani.to_jshtml())\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rQWk0RRuCjDN"
      },
      "source": [
        "# Load the Data and initialize the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ttMHeiGCppjB"
      },
      "outputs": [],
      "source": [
        "# @title Authenticate with Google Cloud Storage\n",
        "\n",
        "# Gives you an authenticated client, in case you want to use a private bucket.\n",
        "gcs_client = storage.Client.create_anonymous_client()\n",
        "gcs_bucket = gcs_client.get_bucket(\"dm_graphcast\")\n",
        "dir_prefix = \"gencast/\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ty5WSDRjhDBF"
      },
      "source": [
        "## Load the model params\n",
        "\n",
        "Choose one of the two ways of getting model params:\n",
        "- **checkpoint**: You'll get sensible predictions, but are limited to the model architecture that it was trained with, which may not fit on your device.\n",
        "- **random**: You'll get random predictions, but you can change the model architecture and data resolution which may run faster or fit on your device.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PMoFXuZXs-xg"
      },
      "outputs": [],
      "source": [
        "# @title Choose the model\n",
        "\n",
        "params_file_options = [\n",
        "    name for blob in gcs_bucket.list_blobs(prefix=(dir_prefix+\"params/\"))\n",
        "    if (name := blob.name.removeprefix(dir_prefix+\"params/\"))]  # Drop empty string.\n",
        "\n",
        "latent_value_options = [int(2**i) for i in range(4, 10)]\n",
        "random_latent_size = widgets.Dropdown(\n",
        "    options=latent_value_options, value=512,description=\"Latent size:\")\n",
        "random_attention_type = widgets.Dropdown(\n",
        "    options=[\"splash_mha\", \"triblockdiag_mha\", \"mha\"], value=\"splash_mha\", description=\"Attention:\")\n",
        "random_mesh_size = widgets.IntSlider(\n",
        "    value=4, min=4, max=6, description=\"Mesh size:\")\n",
        "random_num_heads = widgets.Dropdown(\n",
        "    options=[int(2**i) for i in range(0, 3)], value=4,description=\"Num heads:\")\n",
        "random_attention_k_hop = widgets.Dropdown(\n",
        "    options=[int(2**i) for i in range(2, 5)], value=16,description=\"Attn k hop:\")\n",
        "random_resolution = widgets.Dropdown(\n",
        "    options=[\"1p0\", \"0p25\"], value=\"1p0\", description=\"Resolution:\")\n",
        "\n",
        "def update_latent_options(*args):\n",
        "  def _latent_valid_for_attn(attn, latent, heads):\n",
        "    head_dim, rem = divmod(latent, heads)\n",
        "    if rem != 0:\n",
        "      return False\n",
        "    # Required for splash attn.\n",
        "    if head_dim % 128 != 0:\n",
        "      return attn != \"splash_mha\"\n",
        "    return True\n",
        "  attn = random_attention_type.value\n",
        "  heads = random_num_heads.value\n",
        "  random_latent_size.options = [\n",
        "      latent for latent in latent_value_options\n",
        "      if _latent_valid_for_attn(attn, latent, heads)]\n",
        "\n",
        "# Observe changes to only allow for valid combinations.\n",
        "random_attention_type.observe(update_latent_options, \"value\")\n",
        "random_latent_size.observe(update_latent_options, \"value\")\n",
        "random_num_heads.observe(update_latent_options, \"value\")\n",
        "\n",
        "params_file = widgets.Dropdown(\n",
        "    options=[f for f in params_file_options if \"Mini\" in f],\n",
        "    description=\"Params file:\",\n",
        "    layout={\"width\": \"max-content\"})\n",
        "\n",
        "source_tab = widgets.Tab([\n",
        "    params_file,\n",
        "    widgets.VBox([\n",
        "        random_attention_type,\n",
        "        random_mesh_size,\n",
        "        random_num_heads,\n",
        "        random_latent_size,\n",
        "        random_attention_k_hop,\n",
        "        random_resolution\n",
        "    ]),\n",
        "])\n",
        "source_tab.set_title(0, \"Checkpoint\")\n",
        "source_tab.set_title(1, \"Random\")\n",
        "widgets.VBox([\n",
        "    source_tab,\n",
        "    widgets.Label(value=\"Run the next cell to load the model. Rerunning this cell clears your selection.\")\n",
        "])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "cgfYjE1YhALA"
      },
      "outputs": [],
      "source": [
        "# @title Load the model\n",
        "\n",
        "source = source_tab.get_title(source_tab.selected_index)\n",
        "\n",
        "if source == \"Random\":\n",
        "  params = None  # Filled in below\n",
        "  state = {}\n",
        "  task_config = gencast.TASK\n",
        "  # Use default values.\n",
        "  sampler_config = gencast.SamplerConfig()\n",
        "  noise_config = gencast.NoiseConfig()\n",
        "  noise_encoder_config = denoiser.NoiseEncoderConfig()\n",
        "  # Configure, otherwise use default values.\n",
        "  denoiser_architecture_config = denoiser.DenoiserArchitectureConfig(\n",
        "    sparse_transformer_config = denoiser.SparseTransformerConfig(\n",
        "        attention_k_hop=random_attention_k_hop.value,\n",
        "        attention_type=random_attention_type.value,\n",
        "        d_model=random_latent_size.value,\n",
        "        num_heads=random_num_heads.value\n",
        "        ),\n",
        "    mesh_size=random_mesh_size.value,\n",
        "    latent_size=random_latent_size.value,\n",
        "  )\n",
        "else:\n",
        "  assert source == \"Checkpoint\"\n",
        "  with gcs_bucket.blob(dir_prefix + f\"params/{params_file.value}\").open(\"rb\") as f:\n",
        "    ckpt = checkpoint.load(f, gencast.CheckPoint)\n",
        "  params = ckpt.params\n",
        "  state = {}\n",
        "\n",
        "  task_config = ckpt.task_config\n",
        "  sampler_config = ckpt.sampler_config\n",
        "  noise_config = ckpt.noise_config\n",
        "  noise_encoder_config = ckpt.noise_encoder_config\n",
        "  denoiser_architecture_config = ckpt.denoiser_architecture_config\n",
        "  print(\"Model description:\\n\", ckpt.description, \"\\n\")\n",
        "  print(\"Model license:\\n\", ckpt.license, \"\\n\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z2AqgxUgiALy"
      },
      "source": [
        "## Load the example data\n",
        "\n",
        "Example ERA5 datasets are available at 0.25 degree and 1 degree resolution.\n",
        "\n",
        "Example HRES-fc0 datasets are available at 0.25 degree resolution.\n",
        "\n",
        "Some transformations were done from the base datasets:\n",
        "- We accumulated precipitation over 12 hours instead of the default 1 hour.\n",
        "- For HRES-fc0 sea surface temperature, we assigned NaNs to grid cells in which sea surface temperature was NaN in the ERA5 dataset (this remains fixed at all times).\n",
        "\n",
        "The data resolution must match the model that is loaded. Since we are running GenCast Mini, this will be 1 degree.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5XGzOww0y_BC"
      },
      "outputs": [],
      "source": [
        "# @title Get and filter the list of available example datasets\n",
        "\n",
        "dataset_file_options = [\n",
        "    name for blob in gcs_bucket.list_blobs(prefix=(dir_prefix + \"dataset/\"))\n",
        "    if (name := blob.name.removeprefix(dir_prefix+\"dataset/\"))]  # Drop empty string.\n",
        "\n",
        "def parse_file_parts(file_name):\n",
        "  return dict(part.split(\"-\", 1) for part in file_name.split(\"_\"))\n",
        "\n",
        "\n",
        "def data_valid_for_model(file_name: str, params_file_name: str):\n",
        "  \"\"\"Check data type and resolution matches.\"\"\"\n",
        "  data_file_parts = parse_file_parts(file_name.removesuffix(\".nc\"))\n",
        "  data_res = data_file_parts[\"res\"].replace(\".\", \"p\")\n",
        "  if source == \"Random\":\n",
        "    return random_resolution.value == data_res\n",
        "  res_matches = data_res in params_file_name.lower()\n",
        "  source_matches = \"Operational\" in params_file_name\n",
        "  if data_file_parts[\"source\"] == \"era5\":\n",
        "    source_matches = not source_matches\n",
        "  return res_matches and source_matches\n",
        "\n",
        "dataset_file = widgets.Dropdown(\n",
        "    options=[\n",
        "        (\", \".join([f\"{k}: {v}\" for k, v in parse_file_parts(option.removesuffix(\".nc\")).items()]), option)\n",
        "        for option in dataset_file_options\n",
        "        if data_valid_for_model(option, params_file.value)\n",
        "    ],\n",
        "    description=\"Dataset file:\",\n",
        "    layout={\"width\": \"max-content\"})\n",
        "widgets.VBox([\n",
        "    dataset_file,\n",
        "    widgets.Label(value=\"Run the next cell to load the dataset. Rerunning this cell clears your selection and refilters the datasets that match your model.\")\n",
        "])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "Yz-ekISoJxeZ"
      },
      "outputs": [],
      "source": [
        "# @title Load weather data\n",
        "\n",
        "with gcs_bucket.blob(dir_prefix+f\"dataset/{dataset_file.value}\").open(\"rb\") as f:\n",
        "  example_batch = xarray.load_dataset(f).compute()\n",
        "\n",
        "assert example_batch.dims[\"time\"] \u003e= 3  # 2 for input, \u003e=1 for targets\n",
        "\n",
        "print(\", \".join([f\"{k}: {v}\" for k, v in parse_file_parts(dataset_file.value.removesuffix(\".nc\")).items()]))\n",
        "\n",
        "example_batch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "lXjFvdE6qStr"
      },
      "outputs": [],
      "source": [
        "# @title Choose data to plot\n",
        "\n",
        "plot_example_variable = widgets.Dropdown(\n",
        "    options=example_batch.data_vars.keys(),\n",
        "    value=\"2m_temperature\",\n",
        "    description=\"Variable\")\n",
        "plot_example_level = widgets.Dropdown(\n",
        "    options=example_batch.coords[\"level\"].values,\n",
        "    value=500,\n",
        "    description=\"Level\")\n",
        "plot_example_robust = widgets.Checkbox(value=True, description=\"Robust\")\n",
        "plot_example_max_steps = widgets.IntSlider(\n",
        "    min=1, max=example_batch.dims[\"time\"], value=example_batch.dims[\"time\"],\n",
        "    description=\"Max steps\")\n",
        "\n",
        "widgets.VBox([\n",
        "    plot_example_variable,\n",
        "    plot_example_level,\n",
        "    plot_example_robust,\n",
        "    plot_example_max_steps,\n",
        "    widgets.Label(value=\"Run the next cell to plot the data. Rerunning this cell clears your selection.\")\n",
        "])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "iqzXVpn9_b15"
      },
      "outputs": [],
      "source": [
        "# @title Plot example data\n",
        "\n",
        "plot_size = 7\n",
        "\n",
        "data = {\n",
        "    \" \": scale(select(example_batch, plot_example_variable.value, plot_example_level.value, plot_example_max_steps.value),\n",
        "              robust=plot_example_robust.value),\n",
        "}\n",
        "fig_title = plot_example_variable.value\n",
        "if \"level\" in example_batch[plot_example_variable.value].coords:\n",
        "  fig_title += f\" at {plot_example_level.value} hPa\"\n",
        "\n",
        "plot_data(data, fig_title, plot_size, plot_example_robust.value)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "njD4jsPTPKvJ"
      },
      "outputs": [],
      "source": [
        "# @title Extract training and eval data\n",
        "\n",
        "train_inputs, train_targets, train_forcings = data_utils.extract_inputs_targets_forcings(\n",
        "    example_batch, target_lead_times=slice(\"12h\", \"12h\"), # Only 1AR training.\n",
        "    **dataclasses.asdict(task_config))\n",
        "\n",
        "eval_inputs, eval_targets, eval_forcings = data_utils.extract_inputs_targets_forcings(\n",
        "    example_batch, target_lead_times=slice(\"12h\", f\"{(example_batch.dims['time']-2)*12}h\"), # All but 2 input frames.\n",
        "    **dataclasses.asdict(task_config))\n",
        "\n",
        "print(\"All Examples:  \", example_batch.dims.mapping)\n",
        "print(\"Train Inputs:  \", train_inputs.dims.mapping)\n",
        "print(\"Train Targets: \", train_targets.dims.mapping)\n",
        "print(\"Train Forcings:\", train_forcings.dims.mapping)\n",
        "print(\"Eval Inputs:   \", eval_inputs.dims.mapping)\n",
        "print(\"Eval Targets:  \", eval_targets.dims.mapping)\n",
        "print(\"Eval Forcings: \", eval_forcings.dims.mapping)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-DJzie5me2-H"
      },
      "outputs": [],
      "source": [
        "# @title Load normalization data\n",
        "\n",
        "with gcs_bucket.blob(dir_prefix+\"stats/diffs_stddev_by_level.nc\").open(\"rb\") as f:\n",
        "  diffs_stddev_by_level = xarray.load_dataset(f).compute()\n",
        "with gcs_bucket.blob(dir_prefix+\"stats/mean_by_level.nc\").open(\"rb\") as f:\n",
        "  mean_by_level = xarray.load_dataset(f).compute()\n",
        "with gcs_bucket.blob(dir_prefix+\"stats/stddev_by_level.nc\").open(\"rb\") as f:\n",
        "  stddev_by_level = xarray.load_dataset(f).compute()\n",
        "with gcs_bucket.blob(dir_prefix+\"stats/min_by_level.nc\").open(\"rb\") as f:\n",
        "  min_by_level = xarray.load_dataset(f).compute()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ke2zQyuT_sMA"
      },
      "outputs": [],
      "source": [
        "# @title Build jitted functions, and possibly initialize random weights\n",
        "\n",
        "\n",
        "def construct_wrapped_gencast():\n",
        "  \"\"\"Constructs and wraps the GenCast Predictor.\"\"\"\n",
        "  predictor = gencast.GenCast(\n",
        "      sampler_config=sampler_config,\n",
        "      task_config=task_config,\n",
        "      denoiser_architecture_config=denoiser_architecture_config,\n",
        "      noise_config=noise_config,\n",
        "      noise_encoder_config=noise_encoder_config,\n",
        "  )\n",
        "\n",
        "  predictor = normalization.InputsAndResiduals(\n",
        "      predictor,\n",
        "      diffs_stddev_by_level=diffs_stddev_by_level,\n",
        "      mean_by_level=mean_by_level,\n",
        "      stddev_by_level=stddev_by_level,\n",
        "  )\n",
        "\n",
        "  predictor = nan_cleaning.NaNCleaner(\n",
        "      predictor=predictor,\n",
        "      reintroduce_nans=True,\n",
        "      fill_value=min_by_level,\n",
        "      var_to_clean='sea_surface_temperature',\n",
        "  )\n",
        "\n",
        "  return predictor\n",
        "\n",
        "\n",
        "@hk.transform_with_state\n",
        "def run_forward(inputs, targets_template, forcings):\n",
        "  predictor = construct_wrapped_gencast()\n",
        "  return predictor(inputs, targets_template=targets_template, forcings=forcings)\n",
        "\n",
        "\n",
        "@hk.transform_with_state\n",
        "def loss_fn(inputs, targets, forcings):\n",
        "  predictor = construct_wrapped_gencast()\n",
        "  loss, diagnostics = predictor.loss(inputs, targets, forcings)\n",
        "  return xarray_tree.map_structure(\n",
        "      lambda x: xarray_jax.unwrap_data(x.mean(), require_jax=True),\n",
        "      (loss, diagnostics),\n",
        "  )\n",
        "\n",
        "\n",
        "def grads_fn(params, state, inputs, targets, forcings):\n",
        "  def _aux(params, state, i, t, f):\n",
        "    (loss, diagnostics), next_state = loss_fn.apply(\n",
        "        params, state, jax.random.PRNGKey(0), i, t, f\n",
        "    )\n",
        "    return loss, (diagnostics, next_state)\n",
        "\n",
        "  (loss, (diagnostics, next_state)), grads = jax.value_and_grad(\n",
        "      _aux, has_aux=True\n",
        "  )(params, state, inputs, targets, forcings)\n",
        "  return loss, diagnostics, next_state, grads\n",
        "\n",
        "\n",
        "if params is None:\n",
        "  init_jitted = jax.jit(loss_fn.init)\n",
        "  params, state = init_jitted(\n",
        "      rng=jax.random.PRNGKey(0),\n",
        "      inputs=train_inputs,\n",
        "      targets=train_targets,\n",
        "      forcings=train_forcings,\n",
        "  )\n",
        "\n",
        "\n",
        "loss_fn_jitted = jax.jit(\n",
        "    lambda rng, i, t, f: loss_fn.apply(params, state, rng, i, t, f)[0]\n",
        ")\n",
        "grads_fn_jitted = jax.jit(grads_fn)\n",
        "run_forward_jitted = jax.jit(\n",
        "    lambda rng, i, t, f: run_forward.apply(params, state, rng, i, t, f)[0]\n",
        ")\n",
        "# We also produce a pmapped version for running in parallel.\n",
        "run_forward_pmap = xarray_jax.pmap(run_forward_jitted, dim=\"sample\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VBNutliiCyqA"
      },
      "source": [
        "# Run the model\n",
        "\n",
        "The `chunked_prediction_generator_multiple_runs` iterates over forecast steps, where the 1 step forecast is jitted and samples are pmapped across the chips.\n",
        "This allows us to make efficient use of all devices and parallelise generating an ensemble across them. We then combine the chunks at the end to form our final forecast.\n",
        "\n",
        "Note that the `Autoregressive rollout` cell will take longer than the standard inference time to run when executed for the first time, as this will include code compilation time. This cost does not increase with the number of devices, it is a fixed-cost one time operation whose result can be reused across any number of devices."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ewPPSpxr33hN"
      },
      "outputs": [],
      "source": [
        "# The number of ensemble members should be a multiple of the number of devices.\n",
        "print(f\"Number of local devices {len(jax.local_devices())}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7obeY9i9oTtD"
      },
      "outputs": [],
      "source": [
        "# @title Autoregressive rollout (loop in python)\n",
        "\n",
        "print(\"Inputs:  \", eval_inputs.dims.mapping)\n",
        "print(\"Targets: \", eval_targets.dims.mapping)\n",
        "print(\"Forcings:\", eval_forcings.dims.mapping)\n",
        "\n",
        "num_ensemble_members = 8 # @param int\n",
        "rng = jax.random.PRNGKey(0)\n",
        "# We fold-in the ensemble member, this way the first N members should always\n",
        "# match across different runs which use take the same inputs, regardless of\n",
        "# total ensemble size.\n",
        "rngs = np.stack(\n",
        "    [jax.random.fold_in(rng, i) for i in range(num_ensemble_members)], axis=0)\n",
        "\n",
        "chunks = []\n",
        "for chunk in rollout.chunked_prediction_generator_multiple_runs(\n",
        "    # Use pmapped version to parallelise across devices.\n",
        "    predictor_fn=run_forward_pmap,\n",
        "    rngs=rngs,\n",
        "    inputs=eval_inputs,\n",
        "    targets_template=eval_targets * np.nan,\n",
        "    forcings=eval_forcings,\n",
        "    num_steps_per_chunk = 1,\n",
        "    num_samples = num_ensemble_members,\n",
        "    pmap_devices=jax.local_devices()\n",
        "    ):\n",
        "    chunks.append(chunk)\n",
        "predictions = xarray.combine_by_coords(chunks)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "nIoUfBxBAwqm"
      },
      "outputs": [],
      "source": [
        "# @title Choose predictions to plot\n",
        "\n",
        "plot_pred_variable = widgets.Dropdown(\n",
        "    options=predictions.data_vars.keys(),\n",
        "    value=\"2m_temperature\",\n",
        "    description=\"Variable\")\n",
        "plot_pred_level = widgets.Dropdown(\n",
        "    options=predictions.coords[\"level\"].values,\n",
        "    value=500,\n",
        "    description=\"Level\")\n",
        "plot_pred_robust = widgets.Checkbox(value=True, description=\"Robust\")\n",
        "plot_pred_max_steps = widgets.IntSlider(\n",
        "    min=1,\n",
        "    max=predictions.dims[\"time\"],\n",
        "    value=predictions.dims[\"time\"],\n",
        "    description=\"Max steps\")\n",
        "plot_pred_samples = widgets.IntSlider(\n",
        "    min=1,\n",
        "    max=num_ensemble_members,\n",
        "    value=num_ensemble_members,\n",
        "    description=\"Samples\")\n",
        "\n",
        "widgets.VBox([\n",
        "    plot_pred_variable,\n",
        "    plot_pred_level,\n",
        "    plot_pred_robust,\n",
        "    plot_pred_max_steps,\n",
        "    plot_pred_samples,\n",
        "    widgets.Label(value=\"Run the next cell to plot the predictions. Rerunning this cell clears your selection.\")\n",
        "])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "wn7dccXO5R7C"
      },
      "outputs": [],
      "source": [
        "# @title Plot prediction samples and diffs\n",
        "\n",
        "plot_size = 5\n",
        "plot_max_steps = min(predictions.dims[\"time\"], plot_pred_max_steps.value)\n",
        "\n",
        "fig_title = plot_pred_variable.value\n",
        "if \"level\" in predictions[plot_pred_variable.value].coords:\n",
        "  fig_title += f\" at {plot_pred_level.value} hPa\"\n",
        "\n",
        "for sample_idx in range(plot_pred_samples.value):\n",
        "  data = {\n",
        "      \"Targets\": scale(select(eval_targets, plot_pred_variable.value, plot_pred_level.value, plot_max_steps), robust=plot_pred_robust.value),\n",
        "      \"Predictions\": scale(select(predictions.isel(sample=sample_idx), plot_pred_variable.value, plot_pred_level.value, plot_max_steps), robust=plot_pred_robust.value),\n",
        "      \"Diff\": scale((select(eval_targets, plot_pred_variable.value, plot_pred_level.value, plot_max_steps) -\n",
        "                          select(predictions.isel(sample=sample_idx), plot_pred_variable.value, plot_pred_level.value, plot_max_steps)),\n",
        "                        robust=plot_pred_robust.value, center=0),\n",
        "  }\n",
        "  display.display(plot_data(data, fig_title + f\", Sample {sample_idx}\", plot_size, plot_pred_robust.value))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "X3m9lW5fN4oL"
      },
      "outputs": [],
      "source": [
        "# @title Plot ensemble mean and CRPS\n",
        "\n",
        "def crps(targets, predictions, bias_corrected = True):\n",
        "  if predictions.sizes.get(\"sample\", 1) \u003c 2:\n",
        "    raise ValueError(\n",
        "        \"predictions must have dim 'sample' with size at least 2.\")\n",
        "  sum_dims = [\"sample\", \"sample2\"]\n",
        "  preds2 = predictions.rename({\"sample\": \"sample2\"})\n",
        "  num_samps = predictions.sizes[\"sample\"]\n",
        "  num_samps2 = (num_samps - 1) if bias_corrected else num_samps\n",
        "  mean_abs_diff = np.abs(\n",
        "      predictions - preds2).sum(\n",
        "          dim=sum_dims, skipna=False) / (num_samps * num_samps2)\n",
        "  mean_abs_err = np.abs(targets - predictions).sum(dim=\"sample\", skipna=False) / num_samps\n",
        "  return mean_abs_err - 0.5 * mean_abs_diff\n",
        "\n",
        "\n",
        "plot_size = 5\n",
        "plot_max_steps = min(predictions.dims[\"time\"], plot_pred_max_steps.value)\n",
        "\n",
        "fig_title = plot_pred_variable.value\n",
        "if \"level\" in predictions[plot_pred_variable.value].coords:\n",
        "  fig_title += f\" at {plot_pred_level.value} hPa\"\n",
        "\n",
        "data = {\n",
        "    \"Targets\": scale(select(eval_targets, plot_pred_variable.value, plot_pred_level.value, plot_max_steps), robust=plot_pred_robust.value),\n",
        "    \"Ensemble Mean\": scale(select(predictions.mean(dim=[\"sample\"]), plot_pred_variable.value, plot_pred_level.value, plot_max_steps), robust=plot_pred_robust.value),\n",
        "    \"Ensemble CRPS\": scale(crps((select(eval_targets, plot_pred_variable.value, plot_pred_level.value, plot_max_steps)),\n",
        "                        select(predictions, plot_pred_variable.value, plot_pred_level.value, plot_max_steps)),\n",
        "                      robust=plot_pred_robust.value, center=0),\n",
        "}\n",
        "display.display(plot_data(data, fig_title, plot_size, plot_pred_robust.value))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O6ZhRFBPD0kq"
      },
      "source": [
        "# Train the model\n",
        "\n",
        "The following operations requires larger amounts of memory than running inference.\n",
        "\n",
        "The first time executing the cell takes more time, as it includes the time to jit the function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nv-u3dAP7IRZ"
      },
      "outputs": [],
      "source": [
        "# @title Loss computation\n",
        "loss, diagnostics = loss_fn_jitted(\n",
        "    jax.random.PRNGKey(0),\n",
        "    train_inputs,\n",
        "    train_targets,\n",
        "    train_forcings)\n",
        "print(\"Loss:\", float(loss))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mBNFq1IGZNLz"
      },
      "outputs": [],
      "source": [
        "# @title Gradient computation\n",
        "loss, diagnostics, next_state, grads = grads_fn_jitted(\n",
        "    params=params,\n",
        "    state=state,\n",
        "    inputs=train_inputs,\n",
        "    targets=train_targets,\n",
        "    forcings=train_forcings)\n",
        "mean_grad = np.mean(jax.tree_util.tree_flatten(jax.tree_util.tree_map(lambda x: np.abs(x).mean(), grads))[0])\n",
        "print(f\"Loss: {loss:.4f}, Mean |grad|: {mean_grad:.6f}\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "GenCast Mini Demo",
      "private_outputs": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}