diff --git a/.gitignore b/.gitignore index 5e928c6..d5c1321 100644 --- a/.gitignore +++ b/.gitignore @@ -71,5 +71,9 @@ docs/_build/ # Pyenv .python-version +# jupyter +.virtual_documents/ +.ipynb_checkpoints/ + # data files *.nc diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 02240bb..c670a80 100644 --- a/.pre-commit-config.yaml +++ b/.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 diff --git a/docs/examples/infer-cell-geometries.ipynb b/docs/examples/infer-cell-geometries.ipynb new file mode 100644 index 0000000..db34a5d --- /dev/null +++ b/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 +} diff --git a/docs/examples/overlap-query.ipynb b/docs/examples/overlap-query.ipynb new file mode 100644 index 0000000..89231f4 --- /dev/null +++ b/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 +} diff --git a/python/grid_indexing/__init__.py b/python/grid_indexing/__init__.py index 51e39d4..7428e4e 100644 --- a/python/grid_indexing/__init__.py +++ b/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__) diff --git a/python/grid_indexing/grids.py b/python/grid_indexing/grids.py index 178eccb..7e9e8f8 100644 --- a/python/grid_indexing/grids.py +++ b/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):