{ "cells": [ { "id": "cell0", "cell_type": "markdown", "source": [ "> Copyright 2026 Google LLC.\n", ">\n", "> Licensed under the Apache License, Version 2.0 (the \"License\");\n", "> you may not use this file except in compliance with the License.\n", "> You may obtain a copy of the License at\n", ">\n", "> http://www.apache.org/licenses/LICENSE-2.0\n", ">\n", "> Unless required by applicable law or agreed to in writing, software\n", "> distributed under the License is distributed on an \"AS-IS\" BASIS,\n", "> WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", "> See the License for the specific language governing permissions and\n", "> limitations under the License." ], "metadata": {} }, { "id": "cell1", "cell_type": "markdown", "source": [ "# WeatherNext 2 Demo" ], "metadata": {} }, { "id": "u7Q-uwC15Nnv", "cell_type": "markdown", "source": [ "## Installation and Initialization" ], "metadata": {} }, { "id": "AMjmnev04-Y7", "cell_type": "code", "source": [ "# @title Upgrade packages (kernel needs to be restarted after running this cell).\n", "%pip install -U importlib_metadata" ], "metadata": { "cellView": "form" }, "execution_count": null, "outputs": [] }, { "id": "2UAdIgq15Ilg", "cell_type": "code", "source": [ "# @title Pip install repo and dependencies and reconfigure jax if running on TPU.\n", "%pip install --upgrade https://github.com/google-deepmind/weathernext/archive/master.zip --ignore-requires-python\n", "# This is required due to outdated jax and libtpu versions in Colab TPU images.\n", "%pip uninstall -y libtpu libtpu-nightly\n", "%pip install -U \"jax[tpu]\" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html" ], "metadata": { "cellView": "form" }, "execution_count": null, "outputs": [] }, { "id": "cell3", "cell_type": "code", "source": [ "# @title Imports\n", "\n", "import dataclasses\n", "import datetime\n", "import math\n", "from google.colab import auth\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", "import pandas as pd\n", "\n", "import xarray_jax\n", "from weathernext.utils import rollout\n", "from weathernext.utils import checkpoint\n", "from weathernext.utils import data_utils\n", "from weathernext.utils import xarray_tree\n", "from weathernext.utils import fiddle_config_io\n", "from weathernext.weathernext2 import fgn\n", "from weathernext.cyclones import constants as cyclone_constants\n", "from weathernext.cyclones import direct_tracker_6h_v1_config\n", "from weathernext.cyclones import ibtracs_netcdf_to_csv\n" ], "metadata": { "cellView": "form" }, "execution_count": null, "outputs": [] }, { "id": "cell4", "cell_type": "code", "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", " ) -> 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", " ) -> 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", " cyclone_dfs: Optional[dict[str, pd.DataFrame]] = None,\n", " init_time: Optional[np.datetime64] = None,\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", " # Derive per-frame valid times from the lead-time coordinate + init_time.\n", " if cyclone_dfs:\n", " lead_times = first_data.coords[\"time\"].values[:max_steps]\n", " valid_times = pd.Timestamp(init_time) + pd.to_timedelta(lead_times)\n", "\n", " # Build lat/lon to pixel coordinate mappings from the first data array.\n", " lat_coords = first_data.coords[\"lat\"].values\n", " lon_coords = first_data.coords[\"lon\"].values\n", " lat_min, lat_max = lat_coords.min(), lat_coords.max()\n", " lon_min, lon_max = lon_coords.min(), lon_coords.max()\n", " n_lat = len(lat_coords)\n", " n_lon = len(lon_coords)\n", "\n", " def latlon_to_pixel(lats, lons):\n", " \"\"\"Convert lat/lon arrays to pixel coordinates for imshow.\"\"\"\n", " px = (lons - lon_min) / (lon_max - lon_min) * (n_lon - 1)\n", " py = (lats - lat_min) / (lat_max - lat_min) * (n_lat - 1)\n", " return px, py\n", "\n", " def get_cyclone_positions(df, valid_time):\n", " \"\"\"Filter a cyclone DataFrame for a single valid time, return (lats, lons).\"\"\"\n", " frame_df = df[df[\"valid_time\"] == valid_time]\n", " lons = frame_df[\"lon\"].values % 360 # Wrap to [0, 360) to match grid.\n", " return frame_df[\"lat\"].values, lons\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", " scatters = []\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", " if cyclone_dfs and title in cyclone_dfs:\n", " lats_0, lons_0 = get_cyclone_positions(\n", " cyclone_dfs[title], valid_times[0])\n", " px, py = latlon_to_pixel(lats_0, lons_0)\n", " sc = ax.scatter(\n", " px, py, marker=\"x\", c=\"red\", s=20, linewidths=0.8, zorder=5)\n", " scatters.append((title, sc))\n", " else:\n", " scatters.append((title, None))\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", " # Update cyclone scatter positions.\n", " for title, sc in scatters:\n", " if sc is not None:\n", " lats_f, lons_f = get_cyclone_positions(\n", " cyclone_dfs[title], valid_times[frame])\n", " offsets = (np.column_stack(latlon_to_pixel(lats_f, lons_f))\n", " if len(lats_f) > 0 else np.empty((0, 2)))\n", " sc.set_offsets(offsets)\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" ], "metadata": { "cellView": "form" }, "execution_count": null, "outputs": [] }, { "id": "cell5", "cell_type": "markdown", "source": [ "# Load the Data and initialize the model" ], "metadata": {} }, { "id": "cell6", "cell_type": "code", "source": [ "# @title Load the weights, params and data.\n", "\n", "# These can be updated as per info in the README (but will require accelerators\n", "# not freely available in Colab).\n", "model_name = \"WeatherNextCyclones_Mini\"\n", "split = \"2024\"\n", "data_resolution = \"1.0\"\n", "# Steps must be one of \"01\", \"04\", \"12\", \"20\", or \"30\"\n", "steps = \"20\"\n", "\n", "config_name = f\"weathernext2/configs/{model_name}\"\n", "weights_path = f\"weathernext2/params/{model_name}_<{split}.npz\"\n", "data_path = f\"weathernext2/dataset/source-hres_forecast_init-2024-10-07 00:00:00_res-{data_resolution}_levels-13_steps-{steps}.nc\"\n", "\n", "gcs_client = storage.Client.create_anonymous_client()\n", "gcs_bucket = gcs_client.get_bucket(\"dm_graphcast\")\n", "\n", "with gcs_bucket.blob(data_path).open(\"rb\") as f:\n", " example_batch = xarray.load_dataset(f).compute()\n", "config = fiddle_config_io.get_fiddle_config_by_name(config_name)\n", "with gcs_bucket.blob(weights_path).open(\"rb\") as f:\n", " ckpt = checkpoint.load(f, fgn.CheckPoint)\n", "\n", "try:\n", " ! wget https://www.ncei.noaa.gov/data/international-best-track-archive-for-climate-stewardship-ibtracs/v04r01/access/netcdf/IBTrACS.ALL.v04r01.nc\n", " ds_ibtracs = xarray.open_dataset(\"IBTrACS.ALL.v04r01.nc\", engine=\"h5netcdf\").load()\n", "except:\n", " print(f\"Failed to download IBTrACS data. Cells requiring IBTrACS will fail. See https://www.ncei.noaa.gov/news/cloud-migration for more details on recent IBTrACS data access issues.\")\n", " ds_ibtracs = None\n" ], "metadata": { "cellView": "form" }, "execution_count": null, "outputs": [] }, { "id": "cell7", "cell_type": "code", "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.sizes[\"time\"], value=example_batch.sizes[\"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": { "cellView": "form" }, "execution_count": null, "outputs": [] }, { "id": "cell8", "cell_type": "code", "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" ], "metadata": { "cellView": "form" }, "execution_count": null, "outputs": [] }, { "id": "cell9", "cell_type": "code", "source": [ "# @title Extract training and eval data\n", "task_config = config.task\n", "\n", "train_inputs, train_targets, train_forcings = data_utils.extract_inputs_targets_forcings(\n", " example_batch, target_lead_times=slice(\"6h\", \"6h\"), # 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(\"6h\", f\"{(example_batch.sizes['time']-2)*6}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" ], "metadata": { "cellView": "form" }, "execution_count": null, "outputs": [] }, { "id": "cell10", "cell_type": "code", "source": [ "# @title Build jitted functions.\n", "\n", "backend = jax.default_backend()\n", "transformer_kwargs = config.predictor_kwargs['noisy_function_kwargs'][\n", " 'mesh_model_ctor'\n", "].keywords['transformer_kwargs']\n", "# If running on GPU, use alternative attention implementation. Note that this\n", "# is slower and significantly more memory hungry. Inference with WN-mini\n", "# can fit on a P100 GPU, but note that gradient computation will fail OOM.\n", "if backend == 'gpu':\n", " transformer_kwargs['attention_type'] = 'triblockdiag_mha'\n", "# Overriding the following block sizes enable running the model on a TPU v5e-1.\n", "# Note that in general, when running on different TPU versions, optimal block\n", "# sizes are likely be different, and thus some tuning is recommended.\n", "elif backend == 'tpu':\n", " transformer_kwargs.update({\n", " 'block_q': 128,\n", " 'block_kv': 128,\n", " 'block_kv_compute': 128,\n", " 'block_q_dkv': 128,\n", " 'block_kv_dkv': 128,\n", " 'block_kv_dkv_compute': 128,\n", " })\n", "\n", "config_inference = fgn.PredictorConfig(\n", " task=config.task,\n", " predictor_constructor=config.predictor_constructor,\n", " predictor_kwargs=config.predictor_kwargs,\n", " predictor_wrappers=config.predictor_wrappers[:-1], # Remove ensemble wrapper\n", ")\n", "\n", "\n", "@hk.transform\n", "def run_forward(inputs, targets_template, forcings):\n", " predictor = fgn.construct_predictor(config_inference)\n", " return predictor(inputs, targets_template=targets_template, forcings=forcings)\n", "\n", "\n", "# With ensemble wrapper for training on multiple samples.\n", "config_train = config\n", "\n", "\n", "@hk.transform\n", "def loss_fn(inputs, targets, forcings):\n", " predictor = fgn.construct_predictor(config_train)\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", "def grads_fn(params, inputs, targets, forcings):\n", " def _aux(params, inputs, targets, forcings):\n", " loss, diagnostics = loss_fn.apply(\n", " params, jax.random.PRNGKey(0), inputs, targets, forcings\n", " )\n", " return loss, diagnostics\n", " (loss, diagnostics), grads = jax.value_and_grad(\n", " _aux, has_aux=True)(params, inputs, targets, forcings)\n", " return loss, diagnostics, grads\n", "\n", "loss_fn_jitted = jax.jit(loss_fn.apply)\n", "grads_fn_jitted = jax.jit(grads_fn)\n", "run_forward_jitted = jax.jit(\n", " lambda rng, i, t, f: run_forward.apply(ckpt.params, rng, i, t, f)\n", ")\n", "# We also produce a pmapped version for running inference in parallel.\n", "run_forward_pmap = xarray_jax.pmap(run_forward_jitted, dim=\"sample\")" ], "metadata": {}, "execution_count": null, "outputs": [] }, { "id": "cell11", "cell_type": "code", "source": [ "# @title Randomly initialize network parameters\n", "# This cell demonstrates how to initialize random network weights.\n", "\n", "rng = jax.random.PRNGKey(0)\n", "init_params = jax.jit(run_forward.init)(\n", " rng,\n", " train_inputs,\n", " train_targets,\n", " train_forcings,\n", ")\n", "\n", "num_params = sum(x.size for x in jax.tree_util.tree_leaves(init_params))\n", "print(f\"Initialized {num_params:,} parameters:\")\n", "for module_name in init_params:\n", " num_module_params = sum(\n", " x.size for x in jax.tree_util.tree_leaves(init_params[module_name]))\n", " print(f\" {module_name}: {num_module_params:,} parameters\")" ], "metadata": { "cellView": "form" }, "execution_count": null, "outputs": [] }, { "id": "cell12", "cell_type": "markdown", "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." ], "metadata": {} }, { "id": "cell13", "cell_type": "code", "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())}\")" ], "metadata": { "cellView": "form" }, "execution_count": null, "outputs": [] }, { "id": "cell14", "cell_type": "code", "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)" ], "metadata": { "cellView": "form" }, "execution_count": null, "outputs": [] }, { "id": "cell15", "cell_type": "code", "source": [ "# @title Process IBTrACS for tracking\n", "\n", "if ds_ibtracs is None:\n", " raise ValueError(\"IBTrACS data not available. See `Load the weights, params and data` cell.\")\n", "\n", "init_time = example_batch.isel(batch=0, time=1).datetime.values\n", "ds_ibtracs[\"time\"] = ds_ibtracs[\"time\"].dt.round(\"1s\")\n", "all_storms_df, initial_storms_df = ibtracs_netcdf_to_csv.prepare_ibtracs_storms_dfs(\n", " ibtracs_ds=ds_ibtracs,\n", " init_time=init_time,\n", ")\n" ], "metadata": { "cellView": "form" }, "execution_count": null, "outputs": [] }, { "id": "cell16", "cell_type": "code", "source": [ "# @title Run the tracker\n", "tracker_config = direct_tracker_6h_v1_config.get_config()\n", "tracker = tracker_config.tracker_constructor(**tracker_config.tracker_kwargs)\n", "\n", "predicted_ensemble_tracks = []\n", "for sample_idx in predictions.sample:\n", " data_to_track = predictions.isel(batch=0).sel(sample=sample_idx)\n", " data_to_track = tracker.preprocess_gridded_ds(data_to_track)\n", " data_to_track = data_to_track.expand_dims(forecast_datetime=[init_time])\n", " data_to_track = data_to_track.assign_coords(\n", " lead_time_secs=data_to_track.time.astype(\"timedelta64[s]\").astype(int)\n", " )\n", " data_to_track = data_to_track.assign_coords(\n", " date_time=data_to_track.forecast_datetime.data + data_to_track.time\n", " )\n", " data_to_track = data_to_track.as_numpy()\n", "\n", " predicted_tracks = tracker(\n", " gridded_ds=data_to_track,\n", " initial_storms_df=initial_storms_df,\n", " do_cyclogenesis=True,\n", " )\n", "\n", " predicted_tracks[cyclone_constants.LON] = (\n", " predicted_tracks[cyclone_constants.LON] % 360\n", " )\n", " predicted_ensemble_tracks.append(predicted_tracks)" ], "metadata": { "cellView": "form" }, "execution_count": null, "outputs": [] }, { "id": "cell17", "cell_type": "code", "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.sizes[\"time\"],\n", " value=predictions.sizes[\"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", "])" ], "metadata": { "cellView": "form" }, "execution_count": null, "outputs": [] }, { "id": "cell18", "cell_type": "code", "source": [ "# @title Plot prediction samples and diffs\n", "\n", "plot_size = 5\n", "plot_max_steps = min(predictions.sizes[\"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", " cyclone_dfs = {\n", " \"Targets\": all_storms_df,\n", " \"Predictions\": predicted_ensemble_tracks[sample_idx],\n", " }\n", " display.display(\n", " plot_data(data, fig_title + f\", Sample {sample_idx}\", plot_size,\n", " plot_pred_robust.value, cyclone_dfs=cyclone_dfs,\n", " init_time=init_time))\n" ], "metadata": { "cellView": "form" }, "execution_count": null, "outputs": [] }, { "id": "cell19", "cell_type": "code", "source": [ "# @title Plot ensemble mean and CRPS\n", "\n", "def crps(targets, predictions, bias_corrected = True):\n", " if predictions.sizes.get(\"sample\", 1) < 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.sizes[\"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))" ], "metadata": { "cellView": "form" }, "execution_count": null, "outputs": [] }, { "id": "cell20", "cell_type": "markdown", "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." ], "metadata": {} }, { "id": "cell21", "cell_type": "code", "source": [ "# @title Loss computation\n", "loss, diagnostics = loss_fn_jitted(\n", " ckpt.params,\n", " jax.random.PRNGKey(0),\n", " train_inputs,\n", " train_targets,\n", " train_forcings)\n", "print(\"Loss:\", float(loss))\n", "print(diagnostics)" ], "metadata": { "cellView": "form" }, "execution_count": null, "outputs": [] }, { "id": "cell22", "cell_type": "code", "source": [ "# @title Gradient computation (may OOM depending on runtime)\n", "loss, diagnostics, grads = grads_fn_jitted(\n", " ckpt.params,\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}\")\n" ], "metadata": { "cellView": "form" }, "execution_count": null, "outputs": [] } ], "metadata": { "colab": { "private_outputs": true, "provenance": [] } }, "nbformat_minor": 5, "nbformat": 4 }