gif_animation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9ba8ea1d-e90a-4c78-89a2-fa891b08c0ed",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"><span style=\"color: #008000; text-decoration-color: #008000; font-weight: bold\">Successfully refreshed tokens!</span> <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">Token stored at </span><span style=\"color: #bf7fbf; text-decoration-color: #bf7fbf\">/Users/juliusbusecke/.arraylake/</span><span style=\"color: #ff7fff; text-decoration-color: #ff7fff\">token.json</span>\n",
       "</pre>\n"
      ],
      "text/plain": [
       "\u001b[1;32mSuccessfully refreshed tokens!\u001b[0m \u001b[2mToken stored at \u001b[0m\u001b[2;35m/Users/juliusbusecke/.arraylake/\u001b[0m\u001b[2;95mtoken.json\u001b[0m\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">╭───────────────────────────────────────────────── <span style=\"font-weight: bold\">User Details</span> ──────────────────────────────────────────────────╮\n",
       "│ Name: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">Julius Busecke</span>                                                                                            │\n",
       "│ Email: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">julius@earthmover.io</span>                                                                                     │\n",
       "│ Id: <span style=\"color: #7f7f7f; text-decoration-color: #7f7f7f\">0d7588a7-a03e-4d1a-b862-fb8c52c95fab</span>                                                                        │\n",
       "╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\n",
       "</pre>\n"
      ],
      "text/plain": [
       "╭───────────────────────────────────────────────── \u001b[1mUser Details\u001b[0m ──────────────────────────────────────────────────╮\n",
       "│ Name: \u001b[2mJulius Busecke\u001b[0m                                                                                            │\n",
       "│ Email: \u001b[2mjulius@earthmover.io\u001b[0m                                                                                     │\n",
       "│ Id: \u001b[2m0d7588a7-a03e-4d1a-b862-fb8c52c95fab\u001b[0m                                                                        │\n",
       "╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Make a cool animation from GFDL\n",
    "import arraylake as al\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import matplotlib.pyplot as plt\n",
    "import cartopy.crs as ccrs\n",
    "import cartopy.feature as cfeature\n",
    "from matplotlib.animation import FuncAnimation\n",
    "import matplotlib.colors as mcolors\n",
    "\n",
    "client = al.Client()\n",
    "client.login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "8ef62db3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/juliusbusecke/Code/gfdl-spear/.pixi/envs/default/lib/python3.13/site-packages/zarr/codecs/numcodecs/_codecs.py:141: ZarrUserWarning: Numcodecs codecs are not in the Zarr version 3 specification and may not be supported by other zarr implementations.\n",
      "  super().__init__(**codec_config)\n"
     ]
    }
   ],
   "source": [
    "repo = client.get_repo(\"GFDL/noaa-gfdl-spear-large-ensembles-pds\")\n",
    "session = repo.readonly_session(branch=\"main\")\n",
    "ds_6hr = xr.open_zarr(session.store, zarr_format=3, group=\"historical/6hr\")\n",
    "ds_6hr\n",
    "\n",
    "# plot dataset (here precip)\n",
    "da = ds_6hr.pr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "ece4c72a-b02d-497c-ac42-45eaafcfeb2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Movie generation parameters\n",
    "N_FRAMES = 20 # for testing\n",
    "N_MEMBERS = 12\n",
    "FPS = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "84c71d0d-1c84-490e-88ef-35cd0eaae44d",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Define colormaps per member to distinguish between them when stacked (to use a single one, just replace \n",
    "\n",
    "def transparent_cmap(color, name=None, n=256):\n",
    "    \"\"\"Colormap where RGB stays constant and only alpha ramps from 0 to 1.\"\"\"\n",
    "    r, g, b = mcolors.to_rgb(color)\n",
    "    colors = np.zeros((n, 4))\n",
    "    colors[:, 0] = r\n",
    "    colors[:, 1] = g\n",
    "    colors[:, 2] = b\n",
    "    colors[:, 3] = np.linspace(0, 1, n)  # alpha ramp\n",
    "    return mcolors.ListedColormap(colors, name=name or f\"transparent_{color}\")\n",
    "\n",
    "\n",
    "# # Define individual colors for each member cmap\n",
    "# cmaps = [\n",
    "#     transparent_cmap(\"#1f77b4\", \"t_blue\"),      # blue\n",
    "#     transparent_cmap(\"#d62728\", \"t_red\"),       # red\n",
    "#     transparent_cmap(\"#2ca02c\", \"t_green\"),     # green\n",
    "#     transparent_cmap(\"#ff7f0e\", \"t_orange\"),    # orange\n",
    "#     transparent_cmap(\"#9467bd\", \"t_purple\"),    # purple\n",
    "#     transparent_cmap(\"#17becf\", \"t_cyan\"),      # cyan\n",
    "#     transparent_cmap(\"#e377c2\", \"t_pink\"),      # pink\n",
    "#     transparent_cmap(\"#bcbd22\", \"t_yellow\"),    # yellow-olive\n",
    "# ]\n",
    "\n",
    "# Sample 8 colors from magma and build transparent cmaps\n",
    "\n",
    "def sample_cmap_colors(cmap_name, n, lo=0, hi=1.0):\n",
    "    \"\"\"Pick n evenly-spaced colors from a colormap, avoiding the extremes.\"\"\"\n",
    "    cmap = plt.get_cmap(cmap_name)\n",
    "    return [cmap(x) for x in np.linspace(lo, hi, n)]\n",
    "\n",
    "base_colors = sample_cmap_colors(\"viridis\", n=N_MEMBERS)\n",
    "cmaps = [transparent_cmap(c, name=f\"transparent_cmap_{i}\") for i, c in enumerate(base_colors)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3094f9a5-5726-4286-869e-77443cb7b6c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "## frame animation\n",
    "\n",
    "sp_kw = dict(projection=ccrs.PlateCarree())\n",
    "plot_kwargs = dict(\n",
    "    x='lon', y='lat', vmax=0.00045,\n",
    "    transform=ccrs.PlateCarree(),\n",
    "    add_colorbar=False,\n",
    "    add_labels=False,\n",
    ")\n",
    "\n",
    "def update_plot_frame(frame, da, mesh):\n",
    "    da_frame = da.isel(time=frame).squeeze()\n",
    "    mesh.set_array(da_frame.values.ravel())\n",
    "    return [mesh]\n",
    "\n",
    "movies = []\n",
    "for member in range(N_MEMBERS):\n",
    "    cmap = cmaps[member]\n",
    "    filename = f\"spear_member_{member}.mp4\"\n",
    "    da_frame = da.isel(time=0, member_id=member)\n",
    "    fig, ax = plt.subplots(figsize=(19.2, 10.8), dpi=100, subplot_kw=sp_kw)\n",
    "    mesh = da_frame.plot(ax=ax, cmap=cmap, **plot_kwargs)\n",
    "    \n",
    "    ax.set_facecolor('0.4')\n",
    "    land_10m = cfeature.NaturalEarthFeature(\n",
    "        'physical', 'land', '10m',\n",
    "        facecolor='0.25', edgecolor='none'\n",
    "    )\n",
    "    ax.add_feature(land_10m, zorder=0)\n",
    "    \n",
    "    ax.set_aspect('auto')           # <-- add this\n",
    "    ax.set_position([0, 0, 1, 1])\n",
    "    ax.set_title('')\n",
    "    \n",
    "    fa = FuncAnimation(fig, update_plot_frame, N_FRAMES,\n",
    "                       fargs=[da.isel(member_id=member), mesh],\n",
    "                       interval=1000/FPS)\n",
    "    fa.save(filename, savefig_kwargs={'pad_inches': 0})\n",
    "    movies.append(filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20e23399-9501-4d4c-a178-d63e9a92b3b6",
   "metadata": {},
   "source": [
    "## Optional: Show multiple members in one GIF side by side"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "283db5d1-45d2-44e3-8fd6-b3ab2adc126e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import subprocess\n",
    "from pathlib import Path\n",
    "\n",
    "def stripe_gif(inputs, output=\"out.gif\", width=720, fps=15,\n",
    "               border_px=2, border_color=\"black\"):\n",
    "    n = len(inputs)\n",
    "    \n",
    "    crops = []\n",
    "    for i in range(n):\n",
    "        base = f\"[{i}:v]scale=1920:1080,setsar=1,crop=iw/{n}:ih:{i}*iw/{n}:0\"\n",
    "        # Add right-side border on all but the last stripe\n",
    "        if i < n - 1 and border_px > 0:\n",
    "            base += f\",pad=iw+{border_px}:ih:0:0:{border_color}\"\n",
    "        crops.append(f\"{base}[s{i}]\")\n",
    "    \n",
    "    stack_inputs = \"\".join(f\"[s{i}]\" for i in range(n))\n",
    "    filter_complex = (\n",
    "        \";\".join(crops)\n",
    "        + f\";{stack_inputs}hstack=inputs={n}\"\n",
    "        + f\",fps={fps},scale={width}:-1:flags=lanczos,split[a][b]\"\n",
    "        + \";[a]palettegen=stats_mode=diff[p]\"\n",
    "        + \";[b][p]paletteuse=dither=bayer:bayer_scale=5\"\n",
    "    )\n",
    "    \n",
    "    cmd = [\"ffmpeg\", \"-y\"]\n",
    "    for inp in inputs:\n",
    "        cmd += [\"-i\", str(inp)]\n",
    "    cmd += [\"-filter_complex\", filter_complex, \"-loop\", \"0\", output]\n",
    "    subprocess.run(cmd, check=True)\n",
    "\n",
    "stripe_gif(movies, \"result.gif\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73c417bb-fca5-4b87-8754-af141cd95793",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}