GraphCast_1deg.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# **Notebook Description**\n",
        "\n",
        "7/22/24\n",
        "\n",
        "Authors:<br>\n",
        "Jacob Radford<br>\n",
        "Jebb Stewart\n",
        "<br><br>\n",
        "This notebook is designed for running a coarse (1°) version of GraphCast using free Google Colaboratory resources and plotting the output. You will need to log in with a Google acocunt to access these resources. Note that lines beginning with an exclamation point are calls to the command line rather than native Python code."
      ],
      "metadata": {
        "id": "YFFgnD8eLI3-"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **Step 0: Connect to a runtime**\n",
        "\n",
        "At the top right of this page (next to \"Connect\") click on the dropdown arrow, then \"change runtime type.\" Select \"Python 3\" and \"T4 GPU\". The T4 GPU runtime may or may not be available at any given time depending on demand (unless you purchase colaboratory compute units).\n"
      ],
      "metadata": {
        "id": "GQ8wJI0BKtUP"
      }
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nwtunkYzq0r_"
      },
      "source": [
        "# **Step 1: Install all of the required packages**\n",
        "\n",
        "*   ai-models-gfs: Extension of ai-models, a package for easily running AIWP models\n",
        "*   ai-models-graphcast-gfs: Extension of ai-models-graphcast, the GraphCast plug-in for ai-models\n",
        "*   basemap: Package for plotting data on maps\n",
        "*   git+https://github.com/deepmind/graphcast.git: The GraphCast repository\n",
        "*   jax and jaxlib: Machine learning framework for running GraphCast with a GPU. If using TPU v2 instead of the T4 GPU, run this instead:\n",
        "\n",
        "    !pip install --upgrade \"jax[tpu]<0.4.24\" -f  https://storage.googleapis.com/jax-releases/libtpu_releases.html\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "XLa4eusxDFxh"
      },
      "outputs": [],
      "source": [
        "#@title Install packages\n",
        "!pip install ai-models-gfs==0.0.10 ai-models-graphcast-gfs==0.0.12\n",
        "!pip install basemap basemap-data-hires\n",
        "!pip install git+https://github.com/deepmind/graphcast.git\n",
        "!pip install jax==0.4.23 jaxlib==0.4.23+cuda12.cudnn89 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "S0ZVLJhnrDzp"
      },
      "source": [
        "# **Step 2: Run GraphCast for Hurricane Beryl**\n",
        "\n",
        "Here is where we run GraphCast. Again note that this is just a call to the command line with the following arguments:\n",
        "\n",
        "\n",
        "*   --input (cds, gfs, gdas): Input data source\n",
        "*   --download-assets: Needed first time to download model weights\n",
        "*   --date: Date of initialization\n",
        "*   --time: Time of initialization (6 hour increments)\n",
        "*   --lead-time: How long to run the model forward in hours\n",
        "*   --onedeg: If included, run 1° version, else 0.25° version of GraphCast\n",
        "*   --nc-or-grib: Whether to output as grib, netcdf4, or both (g, n, ng)\n",
        "*   --path: Output file path\n",
        "*   Model name (graphcast, fourcastnetv2-small, panguweather)\n",
        "\n",
        "Assuming this runs without error, a netCDF or grib file will be produced.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "corkOdnxEKdu"
      },
      "outputs": [],
      "source": [
        "#@title Run Graphcast with command line\n",
        "!ai-models-gfs --input gfs --download-assets --date 20240702 --time 1200 --lead-time 240 --onedeg --nc-or-grib n --path 20240702_12_graphcast graphcast"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **Step 3: Plot the data**\n",
        "\n",
        "Now that the model has completed we can plot the output. We will do so using a nice function that the GraphCast team has provided."
      ],
      "metadata": {
        "id": "5nbERfHmM4Ui"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CtdVUkdtE8cN"
      },
      "outputs": [],
      "source": [
        "#@title Load the netCDF file data that we just produced into xarray dataset\n",
        "import xarray\n",
        "graphcast_20240702_12 = xarray.open_dataset('20240702_12_graphcast.nc')\n",
        "graphcast_20240702_12"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "riNUPBBlg8fN"
      },
      "outputs": [],
      "source": [
        "# @title Define plotting functions courtesy of Google DeepMind team\n",
        "from typing import Optional\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "import ipywidgets as widgets\n",
        "import numpy as np\n",
        "import math\n",
        "import datetime\n",
        "from IPython.display import HTML\n",
        "from matplotlib import animation\n",
        "from mpl_toolkits.basemap import Basemap, shiftgrid\n",
        "\n",
        "def select(\n",
        "    data: xarray.Dataset,\n",
        "    variable: str,\n",
        "    level: Optional[int] = None,\n",
        "    max_steps: Optional[int] = None\n",
        "    ) -> 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 < 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",
        "    lat_bounds: Optional[tuple[float, float]] = None,\n",
        "    lon_bounds: Optional[tuple[float, float]] = None,\n",
        "    vminpercent: float = 5,\n",
        "    vmaxpercent: float = 95\n",
        "    ) -> tuple[xarray.Dataset, matplotlib.colors.Normalize, str]:\n",
        "\n",
        "    if lat_bounds and lon_bounds:\n",
        "        data = data.sel(latitude=slice(lat_bounds[0], lat_bounds[1]),\n",
        "                        longitude=slice(lon_bounds[0], lon_bounds[1]))\n",
        "\n",
        "    vmin = np.nanpercentile(data, (vminpercent if robust else 0))\n",
        "    vmax = np.nanpercentile(data, (vmaxpercent 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 convert_longitudes(lon):\n",
        "    lon = np.asarray(lon)\n",
        "    lon = ((lon + 180) % 360) - 180\n",
        "    return lon\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",
        "    lat_bounds: tuple[float, float] = (-90, 90),\n",
        "    lon_bounds: tuple[float, float] = (-180, 180),\n",
        "    nlevels=21\n",
        "    ) -> 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, plot_size * rows))\n",
        "    figure.suptitle(fig_title, fontsize=16)\n",
        "    figure.subplots_adjust(wspace=0.3, hspace=0.3)  # Adjust these values to control spacing\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",
        "\n",
        "        # Convert longitudes\n",
        "        lon = plot_data.coords['longitude'].values\n",
        "        lat = plot_data.coords['latitude'].values\n",
        "        plot_data_shifted, lon_shifted = shiftgrid(180, plot_data.values, lon, start=False)\n",
        "\n",
        "        # Create basemap with specified bounds\n",
        "        m = Basemap(projection='cyl', resolution='c', ax=ax,\n",
        "                    llcrnrlat=lat_bounds[0], urcrnrlat=lat_bounds[1],\n",
        "                    llcrnrlon=lon_bounds[0], urcrnrlon=lon_bounds[1])\n",
        "        m.drawcoastlines()\n",
        "        m.drawcountries()\n",
        "\n",
        "        lon_shifted, lat_shifted = np.meshgrid(lon_shifted, lat)\n",
        "        x, y = m(lon_shifted, lat_shifted)\n",
        "\n",
        "        # Define levels and BoundaryNorm\n",
        "        levels = np.linspace(norm.vmin, norm.vmax, nlevels)\n",
        "        norm = matplotlib.colors.BoundaryNorm(levels, ncolors=256)\n",
        "\n",
        "        im = m.pcolormesh(x, y, plot_data_shifted[0], norm=norm, 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.datetime.utcfromtimestamp(first_data[\"time\"][frame].item() / 1000000000).strftime('%Y-%m-%d %H:%M:%S')\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_array(shiftgrid(180, plot_data.isel(time=frame, missing_dims=\"ignore\").values, plot_data.coords['longitude'].values, start=False)[0])\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": "code",
      "execution_count": null,
      "metadata": {
        "id": "eCn6fXdzg_CV"
      },
      "outputs": [],
      "source": [
        "# @title Choose data to plot\n",
        "\n",
        "plot_example_variable = widgets.Dropdown(\n",
        "    options=graphcast_20240702_12.data_vars.keys(),\n",
        "    value=\"t\",\n",
        "    description=\"Variable\")\n",
        "plot_example_level = widgets.Dropdown(\n",
        "    options=graphcast_20240702_12.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=graphcast_20240702_12.dims[\"time\"], value=graphcast_20240702_12.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": {
        "id": "GjcbQuFC-_SY"
      },
      "outputs": [],
      "source": [
        "#@title Make the plot\n",
        "plot_size = 6\n",
        "\n",
        "#These are the variables that need diverging color map\n",
        "if plot_example_variable.value in ['u10','v10','u','v','w']:\n",
        "  center = 1\n",
        "else:\n",
        "  center = None\n",
        "data = {\n",
        "  \" \": scale(select(graphcast_20240702_12, plot_example_variable.value, plot_example_level.value, plot_example_max_steps.value),\n",
        "              robust=plot_example_robust.value,center=center,vminpercent=5,vmaxpercent=95),\n",
        "}\n",
        "fig_title = plot_example_variable.value\n",
        "if \"level\" in graphcast_20240702_12[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,lat_bounds=(0, 40),lon_bounds=(-110,-45),nlevels=20)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **Bonus Application 1: Large ensembles**\n",
        "\n",
        "Now that we've run GraphCast, let's see what else we can do with it. A commonly discussed application is developing large ensembles. In the interest of time and compute resources we won't make a *large* ensemble, but we'll at least try a small one."
      ],
      "metadata": {
        "id": "wCwZDYorQvRc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Run GraphCast 4 times with random perturbations\n",
        "for i in range(0,4):\n",
        "  command = f\"ai-models-gfs --input gfs --download-assets --date 20240702 --time 1200 --lead-time 120 --onedeg --nc-or-grib n --path 20240702_12_mem{str(i).zfill(2)} --ensemble 0.0050 0.20 graphcast\"\n",
        "  !{command}"
      ],
      "metadata": {
        "id": "ry0jrUlGRgc1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Load the netCDF file data that we just produced into xarray dataset\n",
        "import xarray\n",
        "file_paths = ['20240702_12_mem00.nc', '20240702_12_mem01.nc', '20240702_12_mem02.nc', '20240702_12_mem03.nc']\n",
        "example_batch = xarray.open_dataset(file_paths[0])"
      ],
      "metadata": {
        "id": "nH0A8I3D3rLi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Define paneled plotting functions courtesy of Google DeepMind team\n",
        "from typing import Optional\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "import ipywidgets as widgets\n",
        "import numpy as np\n",
        "import math\n",
        "import datetime\n",
        "from IPython.display import HTML\n",
        "from matplotlib import animation\n",
        "from mpl_toolkits.basemap import Basemap, shiftgrid\n",
        "import xarray as xr\n",
        "\n",
        "def select_ensemble(\n",
        "    data: xr.Dataset,\n",
        "    variable: str,\n",
        "    level: Optional[int] = None,\n",
        "    max_steps: Optional[int] = None\n",
        "    ) -> xr.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 < 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_ensemble(\n",
        "    data: xr.Dataset,\n",
        "    center: Optional[float] = None,\n",
        "    robust: bool = False,\n",
        "    lat_bounds: Optional[tuple[float, float]] = None,\n",
        "    lon_bounds: Optional[tuple[float, float]] = None,\n",
        "    vminpercent: float = 5,\n",
        "    vmaxpercent: float = 95\n",
        "    ) -> tuple[xr.Dataset, float, float, matplotlib.colors.Normalize, str]:\n",
        "\n",
        "    if lat_bounds and lon_bounds:\n",
        "        data = data.sel(latitude=slice(lat_bounds[0], lat_bounds[1]),\n",
        "                        longitude=slice(lon_bounds[0], lon_bounds[1]))\n",
        "\n",
        "    vmin = np.nanpercentile(data, (vminpercent if robust else 0))\n",
        "    vmax = np.nanpercentile(data, (vmaxpercent 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, vmin, vmax, matplotlib.colors.Normalize(vmin, vmax),\n",
        "            (\"RdBu_r\" if center is not None else \"viridis\"))\n",
        "\n",
        "def convert_longitudes(lon):\n",
        "    lon = np.asarray(lon)\n",
        "    lon = ((lon + 180) % 360) - 180\n",
        "    return lon\n",
        "\n",
        "def plot_data_ensemble(\n",
        "    data: dict[str, tuple[xr.Dataset, float, float, matplotlib.colors.Normalize, str]],\n",
        "    fig_title: str,\n",
        "    plot_size: float = 5,\n",
        "    robust: bool = False,\n",
        "    lat_bounds: tuple[float, float] = (-90, 90),\n",
        "    lon_bounds: tuple[float, float] = (-180, 180),\n",
        "    nlevels=21\n",
        "    ) -> tuple[xr.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[0].sizes.get(\"time\", 1) for d in data.values())\n",
        "\n",
        "    cols = 2\n",
        "    rows = 2\n",
        "    figure = plt.figure(figsize=(plot_size * 2 * cols, plot_size * rows))\n",
        "    figure.suptitle(fig_title, fontsize=16)\n",
        "    figure.subplots_adjust(wspace=0, hspace=0)\n",
        "    figure.tight_layout()\n",
        "\n",
        "    # Determine common vmin and vmax\n",
        "    vmin = min(d[1] for d in data.values())\n",
        "    vmax = max(d[2] for d in data.values())\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",
        "\n",
        "        # Convert longitudes\n",
        "        lon = plot_data.coords['longitude'].values\n",
        "        lat = plot_data.coords['latitude'].values\n",
        "        plot_data_shifted, lon_shifted = shiftgrid(180, plot_data.values, lon, start=False)\n",
        "\n",
        "        # Create basemap with specified bounds\n",
        "        m = Basemap(projection='cyl', resolution='c', ax=ax,\n",
        "                    llcrnrlat=lat_bounds[0], urcrnrlat=lat_bounds[1],\n",
        "                    llcrnrlon=lon_bounds[0], urcrnrlon=lon_bounds[1])\n",
        "        m.drawcoastlines()\n",
        "        m.drawcountries()\n",
        "\n",
        "        lon_shifted, lat_shifted = np.meshgrid(lon_shifted, lat)\n",
        "        x, y = m(lon_shifted, lat_shifted)\n",
        "\n",
        "        # Define 20 levels and BoundaryNorm\n",
        "        levels = np.linspace(vmin, vmax, nlevels)\n",
        "        norm = matplotlib.colors.BoundaryNorm(levels, ncolors=256)\n",
        "\n",
        "        im = m.pcolormesh(x, y, plot_data_shifted[0], norm=norm, 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.datetime.utcfromtimestamp(first_data[\"time\"][frame].item() / 1000000000).strftime('%Y-%m-%d %H:%M:%S')\n",
        "\n",
        "            figure.suptitle(f\"{fig_title}, {td}\", fontsize=16)\n",
        "        else:\n",
        "            figure.suptitle(fig_title, fontsize=16)\n",
        "        for im, (plot_data, _, _, _, _) in zip(images, data.values()):\n",
        "            im.set_array(shiftgrid(180, plot_data.isel(time=frame, missing_dims=\"ignore\").values, plot_data.coords['longitude'].values, start=False)[0])\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())"
      ],
      "metadata": {
        "id": "ycG_ePgwRVzF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Choose data to plot\n",
        "plot_example_variable = widgets.Dropdown(\n",
        "    options=example_batch.data_vars.keys(),\n",
        "    value=\"t\",\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",
        "])"
      ],
      "metadata": {
        "id": "quUYzJbKawRa"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Make the plot\n",
        "variable_name = plot_example_variable.value\n",
        "level = plot_example_level.value\n",
        "robust = plot_example_robust.value\n",
        "max_steps = plot_example_max_steps.value\n",
        "\n",
        "data = {}\n",
        "for i, file_path in enumerate(file_paths):\n",
        "    ds = xr.open_dataset(file_path)\n",
        "    scaled_data = scale_ensemble(select_ensemble(ds, variable_name, level, max_steps), robust=robust,vminpercent=5,vmaxpercent=95)\n",
        "    data[f\"File {i+1}\"] = scaled_data\n",
        "\n",
        "fig_title = f\"{variable_name} at {level} hPa\" if level else variable_name\n",
        "plot_data_ensemble(data, fig_title=fig_title, plot_size=4, robust=robust, lat_bounds=(5, 35), lon_bounds=(-105, -50), nlevels=20)\n"
      ],
      "metadata": {
        "id": "RxcgKrr9bO29"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **Bonus Application 2: The What-If Machine**\n",
        "\n",
        "![Alt Text](https://static0.gamerantimages.com/wordpress/wp-content/uploads/2024/03/futurama-professor-farnsworth.jpg?q=50&fit=crop&w=1500&dpr=1.5)\n",
        "\n"
      ],
      "metadata": {
        "id": "rKD2V4z07Nyj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title What if we made the gulf 10 kelvin warmer?\n",
        "!ai-models-gfs --input gfs --download-assets --date 20240702 --time 1200 --lead-time 240 --onedeg --nc-or-grib n --path 20240702_12_perturb --perturbation 2m_temperature 10 25 -90 600 0 graphcast"
      ],
      "metadata": {
        "id": "0Zf0jn7q76yp"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tN1zgA9f_lf4"
      },
      "outputs": [],
      "source": [
        "#@title Load the netCDF file data that we just produced into xarray dataset\n",
        "import xarray\n",
        "graphcast_20240702_12_perturb = xarray.open_dataset('20240702_12_perturb.nc')\n",
        "graphcast_20240702_12_perturb"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tNxNRcDh-zd-"
      },
      "outputs": [],
      "source": [
        "# @title Choose data to plot\n",
        "\n",
        "plot_example_variable = widgets.Dropdown(\n",
        "    options=graphcast_20240702_12_perturb.data_vars.keys(),\n",
        "    value=\"t\",\n",
        "    description=\"Variable\")\n",
        "plot_example_level = widgets.Dropdown(\n",
        "    options=graphcast_20240702_12_perturb.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=graphcast_20240702_12_perturb.dims[\"time\"], value=graphcast_20240702_12_perturb.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": {
        "id": "bZyitJlPhIEw"
      },
      "outputs": [],
      "source": [
        "#@title Make the plot\n",
        "plot_size = 6\n",
        "\n",
        "#These are the variables that need diverging color map\n",
        "if plot_example_variable.value in ['u10','v10','u','v','w']:\n",
        "  center = 1\n",
        "else:\n",
        "  center = None\n",
        "data = {\n",
        "  \" \": scale(select(graphcast_20240702_12_perturb, plot_example_variable.value, plot_example_level.value, plot_example_max_steps.value),\n",
        "              robust=plot_example_robust.value,center=center,vminpercent=5,vmaxpercent=95),\n",
        "}\n",
        "fig_title = plot_example_variable.value\n",
        "if \"level\" in graphcast_20240702_12_perturb[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,lat_bounds=(0, 40),lon_bounds=(-110,-45),nlevels=20)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **Bonus Application 3: Other AIWP models**\n",
        "\n",
        "With the ai-models package, we can other AIWP models like FourCastNetv2-small and PanguWeather. We'll give FourcastNetv2 a try next.\n",
        "\n",
        "**Because this is a 0.25° model it needs more memory and so changing the runtime type to TPU v2 is recommended. TPU v2 may or may not be available at any given time (unless you purchase colaboratory compute units).**\n",
        "\n",
        "PanguWeather (and 0.25° GraphCast) requires a more powerful GPU and so we won't be able to run them in the colab notebook."
      ],
      "metadata": {
        "id": "KrsRiF8ThXuc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Install FourCastNetv2 (and ai-models-gfs again because we switched runtime type)\n",
        "%%capture\n",
        "!pip install basemap basemap-data-hires ai-models-gfs ai-models-fourcastnetv2-gfs ai-models-panguweather-gfs"
      ],
      "metadata": {
        "id": "LVHklxy_ZzZv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Run FourCastNetv2-small for Hurricane Beryl\n",
        "!ai-models-gfs --input gfs --download-assets --date 20240702 --time 1200 --lead-time 120 --nc-or-grib n --path 20240702_12_fourcastnet fourcastnetv2-small"
      ],
      "metadata": {
        "id": "cIYqfP2aHx5r"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Load the netCDF file data that we just produced into xarray dataset\n",
        "import xarray\n",
        "fourcastnet_20240702_12 = xarray.open_dataset('20240702_12_fourcastnet.nc')"
      ],
      "metadata": {
        "id": "XmcIHM1MZszy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Choose data to plot\n",
        "\n",
        "plot_example_variable = widgets.Dropdown(\n",
        "    options=fourcastnet_20240702_12.data_vars.keys(),\n",
        "    value=\"t\",\n",
        "    description=\"Variable\")\n",
        "plot_example_level = widgets.Dropdown(\n",
        "    options=fourcastnet_20240702_12.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=fourcastnet_20240702_12.dims[\"time\"], value=fourcastnet_20240702_12.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",
        "])"
      ],
      "metadata": {
        "id": "P7Zt19zl1heR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Make the plot\n",
        "plot_size = 6\n",
        "\n",
        "#These are the variables that need diverging color map\n",
        "if plot_example_variable.value in ['u10','v10','u','v','w']:\n",
        "  center = 1\n",
        "else:\n",
        "  center = None\n",
        "data = {\n",
        "  \" \": scale(select(fourcastnet_20240702_12, plot_example_variable.value, plot_example_level.value, plot_example_max_steps.value),\n",
        "              robust=plot_example_robust.value,center=center,vminpercent=5,vmaxpercent=95),\n",
        "}\n",
        "fig_title = plot_example_variable.value\n",
        "if \"level\" in fourcastnet_20240702_12[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,lat_bounds=(0, 40),lon_bounds=(-110,-45),nlevels=100)"
      ],
      "metadata": {
        "id": "sgFQeJKl1is2"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **Bonus Application 4: Choose your own GraphCast adventure**"
      ],
      "metadata": {
        "id": "hUXJaRE5Yo2-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Install packages again for T4 runtime\n",
        "!pip install ai-models-gfs==0.0.10 ai-models-graphcast-gfs==0.0.12\n",
        "!pip install basemap basemap-data-hires\n",
        "!pip install git+https://github.com/deepmind/graphcast.git\n",
        "!pip install jax==0.4.23 jaxlib==0.4.23+cuda12.cudnn89 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html"
      ],
      "metadata": {
        "id": "PVC-F6fWqThd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Run Graphcast with custom configuration (connect to T4)\n",
        "# Variable options:  temperature,\n",
        "#                    geopotential,\n",
        "#                    u_component_of_wind,\n",
        "#                    v_component_of_wind,\n",
        "#                    vertical_velocity,\n",
        "#                    specific_humidity,\n",
        "#                    2m_temperature,\n",
        "#                    mean_sea_level_pressure,\n",
        "#                    10m_v_component_of_wind,\n",
        "#                    10m_u_component_of_wind,\n",
        "\n",
        "date = \"20240616\"\n",
        "time = \"0600\"\n",
        "perturb = True\n",
        "perturbation_variable = \"2m_temperature\"\n",
        "perturbation_magnitude = \"10\"\n",
        "perturbation_latitude = \"25\"\n",
        "perturbation_longitude = \"-90\"\n",
        "perturbation_radius = \"600\"\n",
        "perturbation_level = \"0\"\n",
        "\n",
        "if perturb:\n",
        "    command_parts = [\n",
        "        \"ai-models-gfs\",\n",
        "        \"--input gfs\",\n",
        "        \"--download-assets\",\n",
        "        f\"--date {date}\",\n",
        "        f\"--time {time}\",\n",
        "        \"--lead-time 240\",\n",
        "        \"--onedeg\",\n",
        "        \"--nc-or-grib n\",\n",
        "        f\"--perturbation {perturbation_variable} {perturbation_magnitude} {perturbation_latitude} {perturbation_longitude} {perturbation_radius} {perturbation_level}\",\n",
        "        f\"--path {date}_{time}_graphcast_cyoa\",\n",
        "        \"graphcast\"\n",
        "    ]\n",
        "else:\n",
        "    command_parts = [\n",
        "        \"ai-models-gfs\",\n",
        "        \"--input gfs\",\n",
        "        \"--download-assets\",\n",
        "        f\"--date {date}\",\n",
        "        f\"--time {time}\",\n",
        "        \"--lead-time 240\",\n",
        "        \"--onedeg\",\n",
        "        \"--nc-or-grib n\",\n",
        "        f\"--path {date}_{time}_graphcast_cyoa\",\n",
        "        \"graphcast\"\n",
        "    ]\n",
        "\n",
        "command = \" \".join(command_parts)\n",
        "!{command}"
      ],
      "metadata": {
        "id": "Z3ckAsqnYwVr"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4fvx1oLRqlVR"
      },
      "outputs": [],
      "source": [
        "#@title Load the netCDF file data that we just produced into xarray dataset\n",
        "import xarray\n",
        "graphcast_cyoa = xarray.open_dataset(f'{date}_{time}_graphcast_cyoa.nc')\n",
        "graphcast_cyoa"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tWDCdtDRqlVS"
      },
      "outputs": [],
      "source": [
        "# @title Choose data to plot\n",
        "\n",
        "plot_example_variable = widgets.Dropdown(\n",
        "    options=graphcast_cyoa.data_vars.keys(),\n",
        "    value=\"t\",\n",
        "    description=\"Variable\")\n",
        "plot_example_level = widgets.Dropdown(\n",
        "    options=graphcast_cyoa.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=graphcast_cyoa.dims[\"time\"], value=graphcast_cyoa.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": {
        "id": "c0sxeBjTqlVS"
      },
      "outputs": [],
      "source": [
        "#@title Make the plot\n",
        "plot_size = 6\n",
        "\n",
        "#These are the variables that need diverging color map\n",
        "if plot_example_variable.value in ['u10','v10','u','v','w']:\n",
        "  center = 1\n",
        "else:\n",
        "  center = None\n",
        "data = {\n",
        "  \" \": scale(select(graphcast_cyoa, plot_example_variable.value, plot_example_level.value, plot_example_max_steps.value),\n",
        "              robust=plot_example_robust.value,center=center,vminpercent=.1,vmaxpercent=99.9),\n",
        "}\n",
        "fig_title = plot_example_variable.value\n",
        "if \"level\" in graphcast_cyoa[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,lat_bounds=(-90, 90),lon_bounds=(-180,180),nlevels=100)"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# **Bonus Application 5: Choose your own FourCastNetv2 adventure**"
      ],
      "metadata": {
        "id": "mZDqMO1WRcSR"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Install FourCastNetv2 (and ai-models-gfs again because we switched runtime type)\n",
        "!pip install basemap basemap-data-hires ai-models-gfs ai-models-fourcastnetv2-gfs ai-models-panguweather-gfs"
      ],
      "metadata": {
        "id": "EXzPvrQlR6l0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Run FourCastNetv2-small with custom configuration (connect to TPUv2)\n",
        "date = \"20240616\"\n",
        "time = \"0600\"\n",
        "command_parts = [\n",
        "    \"ai-models-gfs\",\n",
        "    \"--input gfs\",\n",
        "    \"--download-assets\",\n",
        "    f\"--date {date}\",\n",
        "    f\"--time {time}\",\n",
        "    \"--lead-time 120\",\n",
        "    \"--nc-or-grib n\",\n",
        "    f\"--path {date}_{time}_fourcastnet_cyoa\",\n",
        "    \"fourcastnetv2-small\"\n",
        "]\n",
        "command = \" \".join(command_parts)\n",
        "print(command)\n",
        "!{command}"
      ],
      "metadata": {
        "id": "2G_sfbBISONc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Load the netCDF file data that we just produced into xarray dataset\n",
        "import xarray\n",
        "fourcastnet_cyoa = xarray.open_dataset(f'{date}_{time}_fourcastnet_cyoa.nc')\n",
        "fourcastnet_cyoa"
      ],
      "metadata": {
        "id": "rw1usyMTS1oR"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Choose data to plot\n",
        "\n",
        "plot_example_variable = widgets.Dropdown(\n",
        "    options=fourcastnet_cyoa.data_vars.keys(),\n",
        "    value=\"t\",\n",
        "    description=\"Variable\")\n",
        "plot_example_level = widgets.Dropdown(\n",
        "    options=fourcastnet_cyoa.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=fourcastnet_cyoa.dims[\"time\"], value=fourcastnet_cyoa.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",
        "])"
      ],
      "metadata": {
        "id": "_1aN9SZnS-eH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Make the plot\n",
        "plot_size = 6\n",
        "\n",
        "#These are the variables that need diverging color map\n",
        "if plot_example_variable.value in ['u10','v10','u','v','w']:\n",
        "  center = 1\n",
        "else:\n",
        "  center = None\n",
        "data = {\n",
        "  \" \": scale(select(fourcastnet_cyoa, plot_example_variable.value, plot_example_level.value, plot_example_max_steps.value),\n",
        "              robust=plot_example_robust.value,center=center,vminpercent=.1,vmaxpercent=99.9),\n",
        "}\n",
        "fig_title = plot_example_variable.value\n",
        "if \"level\" in fourcastnet_cyoa[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,lat_bounds=(-90, 90),lon_bounds=(-180,180),nlevels=100)"
      ],
      "metadata": {
        "id": "HfVQa-PvTAYQ"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "TPU",
    "colab": {
      "gpuType": "V28",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}