Skip to content

Commit e59053d

Browse files
keewiskeewis
authored
demo of the current functionality (#17)
* make sure we always have float64 coordinates * expose `infer_cell_geometries` * prepare for notebooks * demo notebook on `infer_*` The `visualize_grid` utility may be moved elsewhere at some point. * demo of the overlap query
1 parent 2433a09 commit e59053d

File tree

6 files changed

+261
-4
lines changed

6 files changed

+261
-4
lines changed

‎.gitignore‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,5 +71,9 @@ docs/_build/
7171
# Pyenv
7272
.python-version
7373

74+
# jupyter
75+
.virtual_documents/
76+
.ipynb_checkpoints/
77+
7478
# data files
7579
*.nc

‎.pre-commit-config.yaml‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,17 +34,25 @@ repos:
3434
- repo: https://github.com/psf/black-pre-commit-mirror
3535
rev: 24.10.0
3636
hooks:
37-
- id: black
37+
- id: black-jupyter
3838
- repo: https://github.com/keewis/blackdoc
3939
rev: v0.3.9
4040
hooks:
4141
- id: blackdoc
4242
additional_dependencies: ["black==24.10.0"]
4343
- id: blackdoc-autoupdate-black
44+
- repo: https://github.com/kynan/nbstripout
45+
rev: 0.8.1
46+
hooks:
47+
- id: nbstripout
48+
args:
49+
- "--extra-keys=metadata.kernelspec"
50+
- "metadata.language_info.version"
4451
- repo: https://github.com/crate-ci/typos
4552
rev: v1.28.4
4653
hooks:
4754
- id: typos
55+
exclude: ".*\\.ipynb$"
4856
- repo: local
4957
hooks:
5058
- id: cargo-fmt

‎docs/examples/infer-cell-geometries.ipynb‎

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "code",
5+
"execution_count": null,
6+
"id": "0",
7+
"metadata": {},
8+
"outputs": [],
9+
"source": [
10+
"import cf_xarray # noqa: F401\n",
11+
"import lonboard\n",
12+
"import xarray as xr\n",
13+
"\n",
14+
"from grid_indexing import infer_cell_geometries, infer_grid_type\n",
15+
"\n",
16+
"xr.set_options(keep_attrs=True)"
17+
]
18+
},
19+
{
20+
"cell_type": "code",
21+
"execution_count": null,
22+
"id": "1",
23+
"metadata": {},
24+
"outputs": [],
25+
"source": [
26+
"def center_longitude(ds):\n",
27+
" lon_name = ds.cf.coordinates[\"longitude\"][0]\n",
28+
" longitude = (ds[lon_name] + 180) % 360 - 180\n",
29+
" return ds.assign_coords({lon_name: longitude})"
30+
]
31+
},
32+
{
33+
"cell_type": "code",
34+
"execution_count": null,
35+
"id": "2",
36+
"metadata": {},
37+
"outputs": [],
38+
"source": [
39+
"def visualize_grid(geoms, data, cmap=\"viridis\", alpha=0.8):\n",
40+
" from arro3.core import Array, ChunkedArray, Schema, Table\n",
41+
" from lonboard.colormap import apply_continuous_cmap\n",
42+
" from matplotlib import colormaps\n",
43+
" from matplotlib.colors import Normalize\n",
44+
"\n",
45+
" array = Array.from_arrow(geoms)\n",
46+
" data_arrow = ChunkedArray([Array.from_numpy(data)])\n",
47+
" arrays = {\"geometry\": array, \"data\": data_arrow}\n",
48+
" fields = [array.field.with_name(name) for name, array in arrays.items()]\n",
49+
" schema = Schema(fields)\n",
50+
"\n",
51+
" table = Table.from_arrays(list(arrays.values()), schema=schema)\n",
52+
"\n",
53+
" normalizer = Normalize(vmin=data.min(skipna=True), vmax=data.max(skipna=True))\n",
54+
" normalized = normalizer(data.data)\n",
55+
" colormap = colormaps[cmap]\n",
56+
" colors = apply_continuous_cmap(normalized, colormap, alpha=alpha)\n",
57+
"\n",
58+
" return lonboard.SolidPolygonLayer(table=table, filled=True, get_fill_color=colors)"
59+
]
60+
},
61+
{
62+
"cell_type": "code",
63+
"execution_count": null,
64+
"id": "3",
65+
"metadata": {},
66+
"outputs": [],
67+
"source": [
68+
"preprocessors = {\n",
69+
" \"air_temperature\": lambda ds: ds[\"air\"].isel(time=0).stack(cells=[\"lon\", \"lat\"]),\n",
70+
" \"rasm\": lambda ds: ds[\"Tair\"].isel(time=0).stack(cells=[\"y\", \"x\"]),\n",
71+
" \"ROMS_example\": lambda ds: ds[\"salt\"]\n",
72+
" .isel(ocean_time=0, s_rho=0)\n",
73+
" .stack(cells=[\"eta_rho\", \"xi_rho\"]),\n",
74+
"}"
75+
]
76+
},
77+
{
78+
"cell_type": "code",
79+
"execution_count": null,
80+
"id": "4",
81+
"metadata": {},
82+
"outputs": [],
83+
"source": [
84+
"datasets = preprocessors.keys()\n",
85+
"cmaps = {\"ROMS_example\": \"viridis\", \"air_temperature\": \"plasma\", \"rasm\": \"cividis\"}\n",
86+
"\n",
87+
"dss = {\n",
88+
" name: xr.tutorial.open_dataset(name).pipe(center_longitude)\n",
89+
" for name in preprocessors\n",
90+
"}\n",
91+
"\n",
92+
"print(\n",
93+
" \"grid types:\",\n",
94+
" *[f\"{name}: {infer_grid_type(ds)}\" for name, ds in dss.items()],\n",
95+
" sep=\"\\n\",\n",
96+
")\n",
97+
"\n",
98+
"layers = [\n",
99+
" visualize_grid(\n",
100+
" infer_cell_geometries(ds), ds.pipe(preprocessors[name]), cmap=cmaps[name]\n",
101+
" )\n",
102+
" for name, ds in dss.items()\n",
103+
"]\n",
104+
"\n",
105+
"lonboard.Map(layers)"
106+
]
107+
},
108+
{
109+
"cell_type": "code",
110+
"execution_count": null,
111+
"id": "5",
112+
"metadata": {},
113+
"outputs": [],
114+
"source": []
115+
}
116+
],
117+
"metadata": {
118+
"language_info": {
119+
"codemirror_mode": {
120+
"name": "ipython",
121+
"version": 3
122+
},
123+
"file_extension": ".py",
124+
"mimetype": "text/x-python",
125+
"name": "python",
126+
"nbconvert_exporter": "python",
127+
"pygments_lexer": "ipython3",
128+
"version": "3.12.8"
129+
}
130+
},
131+
"nbformat": 4,
132+
"nbformat_minor": 5
133+
}

‎docs/examples/overlap-query.ipynb‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
{
2+
"cells": [
3+
{
4+
"cell_type": "code",
5+
"execution_count": null,
6+
"id": "0",
7+
"metadata": {},
8+
"outputs": [],
9+
"source": [
10+
"import numpy as np\n",
11+
"import xarray as xr\n",
12+
"\n",
13+
"import grid_indexing"
14+
]
15+
},
16+
{
17+
"cell_type": "code",
18+
"execution_count": null,
19+
"id": "1",
20+
"metadata": {},
21+
"outputs": [],
22+
"source": [
23+
"source = xr.tutorial.open_dataset(\"air_temperature\")\n",
24+
"source"
25+
]
26+
},
27+
{
28+
"cell_type": "code",
29+
"execution_count": null,
30+
"id": "2",
31+
"metadata": {},
32+
"outputs": [],
33+
"source": [
34+
"min_lon = source[\"lon\"].min().item()\n",
35+
"max_lon = source[\"lon\"].max().item()\n",
36+
"min_lat = source[\"lat\"].min().item()\n",
37+
"max_lat = source[\"lat\"].max().item()\n",
38+
"\n",
39+
"lon_attrs = {\"standard_name\": \"longitude\"}\n",
40+
"lat_attrs = {\"standard_name\": \"latitude\"}\n",
41+
"target = xr.Dataset(\n",
42+
" coords={\n",
43+
" \"lon\": (\"lon\", np.linspace(min_lon, max_lon, 1200), lon_attrs),\n",
44+
" \"lat\": (\"lat\", np.linspace(min_lat, max_lat, 1000), lat_attrs),\n",
45+
" }\n",
46+
")\n",
47+
"target"
48+
]
49+
},
50+
{
51+
"cell_type": "code",
52+
"execution_count": null,
53+
"id": "3",
54+
"metadata": {},
55+
"outputs": [],
56+
"source": [
57+
"%%time\n",
58+
"source_cells = grid_indexing.infer_cell_geometries(source)\n",
59+
"target_cells = grid_indexing.infer_cell_geometries(target)"
60+
]
61+
},
62+
{
63+
"cell_type": "code",
64+
"execution_count": null,
65+
"id": "4",
66+
"metadata": {},
67+
"outputs": [],
68+
"source": [
69+
"%%time\n",
70+
"index = grid_indexing.Index(source_cells)\n",
71+
"index"
72+
]
73+
},
74+
{
75+
"cell_type": "code",
76+
"execution_count": null,
77+
"id": "5",
78+
"metadata": {},
79+
"outputs": [],
80+
"source": [
81+
"%%time\n",
82+
"overlapping_cells = index.query_overlap(target_cells)\n",
83+
"overlapping_cells"
84+
]
85+
},
86+
{
87+
"cell_type": "code",
88+
"execution_count": null,
89+
"id": "6",
90+
"metadata": {},
91+
"outputs": [],
92+
"source": []
93+
}
94+
],
95+
"metadata": {
96+
"language_info": {
97+
"codemirror_mode": {
98+
"name": "ipython",
99+
"version": 3
100+
},
101+
"file_extension": ".py",
102+
"mimetype": "text/x-python",
103+
"name": "python",
104+
"nbconvert_exporter": "python",
105+
"pygments_lexer": "ipython3",
106+
"version": "3.12.8"
107+
}
108+
},
109+
"nbformat": 4,
110+
"nbformat_minor": 5
111+
}

‎python/grid_indexing/__init__.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
from grid_indexing import grid_indexing
22
from grid_indexing.grid_indexing import Index # noqa: F401
3-
from grid_indexing.grids import infer_grid_type
3+
from grid_indexing.grids import infer_cell_geometries, infer_grid_type
44

5+
__all__ = ["infer_grid_type", "infer_cell_geometries"]
56
__doc__ = grid_indexing.__doc__
67
if hasattr(grid_indexing, "__all__"):
7-
__all__ = grid_indexing.__all__ + ["infer_grid_type"]
8+
__all__.extend(grid_indexing.__all__)

‎python/grid_indexing/grids.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ def as_components(boundaries):
2222
geom_offsets = np.arange(np.prod(vertices.shape[:-2]) + 1, dtype="int32")
2323
ring_offsets = geom_offsets * coords_per_pixel
2424

25-
return coords, geom_offsets, ring_offsets
25+
return coords.astype("float64"), geom_offsets, ring_offsets
2626

2727

2828
def infer_grid_type(ds: xr.Dataset):

0 commit comments

Comments
 (0)