{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Visualize xarray rasters with GeoLibre\n",
    "\n",
    "This notebook exercises `Map.add_raster` with both an `xarray.DataArray` and an `xarray.Dataset`. It creates synthetic geographic data, so no download is required."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run once if the raster dependencies are not installed.\n",
    "# %pip install \"geolibre[raster]\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import xarray as xr\n",
    "\n",
    "from geolibre import Map"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "## DataArray\n",
    "\n",
    "Coordinate dimensions named `lon` and `lat` are recognized automatically and default to EPSG:4326."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "lon = np.linspace(-125, -65, 240)\n",
    "lat = np.linspace(50, 24, 140)\n",
    "xx, yy = np.meshgrid(lon, lat)\n",
    "\n",
    "temperature = xr.DataArray(\n",
    "    28 - 0.45 * (yy - 24) + 5 * np.sin((xx + 100) / 8),\n",
    "    coords={\"lat\": lat, \"lon\": lon},\n",
    "    dims=(\"lat\", \"lon\"),\n",
    "    name=\"temperature\",\n",
    "    attrs={\"long_name\": \"Synthetic air temperature\", \"units\": \"°C\"},\n",
    ")\n",
    "temperature"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "m = Map(center=(-96, 37), zoom=3.5, height=\"700px\")\n",
    "m.add_raster(\n",
    "    temperature,\n",
    "    name=\"Temperature DataArray\",\n",
    "    colormap=\"turbo\",\n",
    "    rescale=[[-5, 35]],\n",
    "    array_args={\"nodata\": -9999, \"compress\": \"LZW\"},\n",
    ")\n",
    "m"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "## Dataset variable and dimension selection\n",
    "\n",
    "Use `array_args[\"variable\"]` to select a data variable and `array_args[\"isel\"]` to select an index from non-spatial dimensions such as time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "climate = xr.Dataset(\n",
    "    {\n",
    "        \"temperature\": xr.concat([temperature, temperature + 3], dim=\"time\"),\n",
    "        \"precipitation\": (\n",
    "            (\"time\", \"lat\", \"lon\"),\n",
    "            np.stack(\n",
    "                [\n",
    "                    80 + 60 * np.cos((xx + 95) / 10) ** 2,\n",
    "                    110 + 70 * np.cos((xx + 90) / 10) ** 2,\n",
    "                ]\n",
    "            ),\n",
    "        ),\n",
    "    },\n",
    "    coords={\"time\": [\"2026-01-01\", \"2026-07-01\"], \"lat\": lat, \"lon\": lon},\n",
    ")\n",
    "\n",
    "m.add_raster(\n",
    "    climate,\n",
    "    name=\"Precipitation Dataset\",\n",
    "    colormap=\"viridis\",\n",
    "    rescale=[[0, 180]],\n",
    "    array_args={\"variable\": \"precipitation\", \"isel\": {\"time\": 1}},\n",
    ")\n",
    "m"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "## Multivariable Dataset as RGB\n",
    "\n",
    "Compatible two-dimensional Dataset variables are written as separate bands. Here the three variables are displayed as red, green, and blue."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "red = np.clip(255 * (xx - lon.min()) / np.ptp(lon), 0, 255).astype(\"uint8\")\n",
    "green = np.clip(255 * (lat.max() - yy) / np.ptp(lat), 0, 255).astype(\"uint8\")\n",
    "blue = np.full_like(red, 140)\n",
    "rgb = xr.Dataset(\n",
    "    {\n",
    "        \"red\": ((\"lat\", \"lon\"), red),\n",
    "        \"green\": ((\"lat\", \"lon\"), green),\n",
    "        \"blue\": ((\"lat\", \"lon\"), blue),\n",
    "    },\n",
    "    coords={\"lat\": lat, \"lon\": lon},\n",
    ")\n",
    "\n",
    "m.add_raster(\n",
    "    rgb,\n",
    "    name=\"RGB Dataset\",\n",
    "    bands=[1, 2, 3],\n",
    "    rescale=[[0, 255], [0, 255], [0, 255]],\n",
    ")\n",
    "m"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "The xarray objects are converted to session-scoped temporary Cloud-Optimized GeoTIFFs. Close the widget when finished to remove them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "# m.close()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (geo)",
   "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.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}