demo of the current functionality

.gitignore

@@ -71,5 +71,9 @@ docs/_build/
 # Pyenv
 .python-version
 
+# jupyter
+.virtual_documents/
+.ipynb_checkpoints/
+
 # data files
 *.nc

.pre-commit-config.yaml

@@ -34,17 +34,25 @@ repos:
   - repo: https://github.com/psf/black-pre-commit-mirror
     rev: 24.10.0
     hooks:
-      - id: black
+      - id: black-jupyter
   - repo: https://github.com/keewis/blackdoc
     rev: v0.3.9
     hooks:
       - id: blackdoc
         additional_dependencies: ["black==24.10.0"]
       - id: blackdoc-autoupdate-black
+  - repo: https://github.com/kynan/nbstripout
+    rev: 0.8.1
+    hooks:
+      - id: nbstripout
+        args:
+          - "--extra-keys=metadata.kernelspec"
+          - "metadata.language_info.version"
   - repo: https://github.com/crate-ci/typos
     rev: v1.28.4
     hooks:
       - id: typos
+        exclude: ".*\\.ipynb$"
   - repo: local
     hooks:
       - id: cargo-fmt

docs/examples/infer-cell-geometries.ipynb

@@ -0,0 +1,133 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import cf_xarray  # noqa: F401\n",
+    "import lonboard\n",
+    "import xarray as xr\n",
+    "\n",
+    "from grid_indexing import infer_cell_geometries, infer_grid_type\n",
+    "\n",
+    "xr.set_options(keep_attrs=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def center_longitude(ds):\n",
+    "    lon_name = ds.cf.coordinates[\"longitude\"][0]\n",
+    "    longitude = (ds[lon_name] + 180) % 360 - 180\n",
+    "    return ds.assign_coords({lon_name: longitude})"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def visualize_grid(geoms, data, cmap=\"viridis\", alpha=0.8):\n",
+    "    from arro3.core import Array, ChunkedArray, Schema, Table\n",
+    "    from lonboard.colormap import apply_continuous_cmap\n",
+    "    from matplotlib import colormaps\n",
+    "    from matplotlib.colors import Normalize\n",
+    "\n",
+    "    array = Array.from_arrow(geoms)\n",
+    "    data_arrow = ChunkedArray([Array.from_numpy(data)])\n",
+    "    arrays = {\"geometry\": array, \"data\": data_arrow}\n",
+    "    fields = [array.field.with_name(name) for name, array in arrays.items()]\n",
+    "    schema = Schema(fields)\n",
+    "\n",
+    "    table = Table.from_arrays(list(arrays.values()), schema=schema)\n",
+    "\n",
+    "    normalizer = Normalize(vmin=data.min(skipna=True), vmax=data.max(skipna=True))\n",
+    "    normalized = normalizer(data.data)\n",
+    "    colormap = colormaps[cmap]\n",
+    "    colors = apply_continuous_cmap(normalized, colormap, alpha=alpha)\n",
+    "\n",
+    "    return lonboard.SolidPolygonLayer(table=table, filled=True, get_fill_color=colors)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "preprocessors = {\n",
+    "    \"air_temperature\": lambda ds: ds[\"air\"].isel(time=0).stack(cells=[\"lon\", \"lat\"]),\n",
+    "    \"rasm\": lambda ds: ds[\"Tair\"].isel(time=0).stack(cells=[\"y\", \"x\"]),\n",
+    "    \"ROMS_example\": lambda ds: ds[\"salt\"]\n",
+    "    .isel(ocean_time=0, s_rho=0)\n",
+    "    .stack(cells=[\"eta_rho\", \"xi_rho\"]),\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "datasets = preprocessors.keys()\n",
+    "cmaps = {\"ROMS_example\": \"viridis\", \"air_temperature\": \"plasma\", \"rasm\": \"cividis\"}\n",
+    "\n",
+    "dss = {\n",
+    "    name: xr.tutorial.open_dataset(name).pipe(center_longitude)\n",
+    "    for name in preprocessors\n",
+    "}\n",
+    "\n",
+    "print(\n",
+    "    \"grid types:\",\n",
+    "    *[f\"{name}: {infer_grid_type(ds)}\" for name, ds in dss.items()],\n",
+    "    sep=\"\\n\",\n",
+    ")\n",
+    "\n",
+    "layers = [\n",
+    "    visualize_grid(\n",
+    "        infer_cell_geometries(ds), ds.pipe(preprocessors[name]), cmap=cmaps[name]\n",
+    "    )\n",
+    "    for name, ds in dss.items()\n",
+    "]\n",
+    "\n",
+    "lonboard.Map(layers)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "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.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

docs/examples/overlap-query.ipynb

@@ -0,0 +1,111 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import xarray as xr\n",
+    "\n",
+    "import grid_indexing"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "source = xr.tutorial.open_dataset(\"air_temperature\")\n",
+    "source"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "min_lon = source[\"lon\"].min().item()\n",
+    "max_lon = source[\"lon\"].max().item()\n",
+    "min_lat = source[\"lat\"].min().item()\n",
+    "max_lat = source[\"lat\"].max().item()\n",
+    "\n",
+    "lon_attrs = {\"standard_name\": \"longitude\"}\n",
+    "lat_attrs = {\"standard_name\": \"latitude\"}\n",
+    "target = xr.Dataset(\n",
+    "    coords={\n",
+    "        \"lon\": (\"lon\", np.linspace(min_lon, max_lon, 1200), lon_attrs),\n",
+    "        \"lat\": (\"lat\", np.linspace(min_lat, max_lat, 1000), lat_attrs),\n",
+    "    }\n",
+    ")\n",
+    "target"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "source_cells = grid_indexing.infer_cell_geometries(source)\n",
+    "target_cells = grid_indexing.infer_cell_geometries(target)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "index = grid_indexing.Index(source_cells)\n",
+    "index"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%time\n",
+    "overlapping_cells = index.query_overlap(target_cells)\n",
+    "overlapping_cells"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "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.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

python/grid_indexing/__init__.py

@@ -1,7 +1,8 @@
 from grid_indexing import grid_indexing
 from grid_indexing.grid_indexing import Index  # noqa: F401
-from grid_indexing.grids import infer_grid_type
+from grid_indexing.grids import infer_cell_geometries, infer_grid_type
 
+__all__ = ["infer_grid_type", "infer_cell_geometries"]
 __doc__ = grid_indexing.__doc__
 if hasattr(grid_indexing, "__all__"):
-    __all__ = grid_indexing.__all__ + ["infer_grid_type"]
+    __all__.extend(grid_indexing.__all__)

python/grid_indexing/grids.py

@@ -22,7 +22,7 @@ def as_components(boundaries):
     geom_offsets = np.arange(np.prod(vertices.shape[:-2]) + 1, dtype="int32")
     ring_offsets = geom_offsets * coords_per_pixel
 
-    return coords, geom_offsets, ring_offsets
+    return coords.astype("float64"), geom_offsets, ring_offsets
 
 
 def infer_grid_type(ds: xr.Dataset):