Skip to content

Commit 609af0f

Browse files
maxrjonesdcherian
andauthored
Add `rasterix.assign_index to public API (#23)
Co-authored-by: Deepak Cherian <deepak@cherian.net>
1 parent 0d6864c commit 609af0f

File tree

4 files changed

+257
-25
lines changed

4 files changed

+257
-25
lines changed

src/rasterix/__init__.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from .raster_index import RasterIndex
1+
from .raster_index import RasterIndex, assign_index
22

33

44
def _get_version():
@@ -12,4 +12,4 @@ def _get_version():
1212

1313
__version__ = _get_version()
1414

15-
__all__ = ["RasterIndex"]
15+
__all__ = ["RasterIndex", "assign_index"]

src/rasterix/raster_index.py

Lines changed: 28 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,36 @@
1+
from __future__ import annotations
2+
13
import textwrap
24
from collections.abc import Hashable, Mapping
3-
from typing import Any
5+
from typing import Any, TypeVar
46

57
import numpy as np
68
import pandas as pd
79
from affine import Affine
8-
from xarray import DataArray, Index, Variable
10+
from xarray import Coordinates, DataArray, Dataset, Index, Variable
911
from xarray.core.coordinate_transform import CoordinateTransform
1012

1113
# TODO: import from public API once it is available
1214
from xarray.core.indexes import CoordinateTransformIndex, PandasIndex
1315
from xarray.core.indexing import IndexSelResult, merge_sel_results
1416

17+
from rasterix.rioxarray_compat import guess_dims
18+
19+
T_Xarray = TypeVar("T_Xarray", "DataArray", "Dataset")
20+
21+
22+
def assign_index(obj: T_Xarray, *, x_dim: str | None = None, y_dim: str | None = None) -> T_Xarray:
23+
if x_dim is None or y_dim is None:
24+
guessed_x, guessed_y = guess_dims(obj)
25+
x_dim = x_dim or guessed_x
26+
y_dim = y_dim or guessed_y
27+
28+
index = RasterIndex.from_transform(
29+
obj.rio.transform(), obj.sizes[x_dim], obj.sizes[y_dim], x_dim=x_dim, y_dim=y_dim
30+
)
31+
coords = Coordinates.from_xindex(index)
32+
return obj.assign_coords(coords)
33+
1534

1635
class AffineTransform(CoordinateTransform):
1736
"""Affine 2D transform wrapper."""
@@ -126,7 +145,7 @@ def generate_coords(self, dims: tuple[str, ...] | None = None) -> dict[Hashable,
126145
assert dims is None or dims == self.dims
127146
return self.forward({self.dim: np.arange(self.size)})
128147

129-
def slice(self, slice: slice) -> "AxisAffineTransform":
148+
def slice(self, slice: slice) -> AxisAffineTransform:
130149
start = max(slice.start or 0, 0)
131150
stop = min(slice.stop or self.size, self.size)
132151
step = slice.step or 1
@@ -183,7 +202,7 @@ def __init__(self, transform: AxisAffineTransform):
183202

184203
def isel( # type: ignore[override]
185204
self, indexers: Mapping[Any, int | slice | np.ndarray | Variable]
186-
) -> "AxisAffineTransformIndex | PandasIndex | None":
205+
) -> AxisAffineTransformIndex | PandasIndex | None:
187206
idxer = indexers[self.dim]
188207

189208
# generate a new index with updated transform if a slice is given
@@ -243,7 +262,7 @@ def sel(self, labels, method=None, tolerance=None):
243262

244263
return result
245264

246-
def to_pandas_index(self) -> "pd.Index":
265+
def to_pandas_index(self) -> pd.Index:
247266
import pandas as pd
248267

249268
values = self.transform.generate_coords()
@@ -306,7 +325,7 @@ def __init__(self, indexes: Mapping[WrappedIndexCoords, WrappedIndex]):
306325
@classmethod
307326
def from_transform(
308327
cls, affine: Affine, width: int, height: int, x_dim: str = "x", y_dim: str = "y"
309-
) -> "RasterIndex":
328+
) -> RasterIndex:
310329
indexes: dict[WrappedIndexCoords, AxisAffineTransformIndex | CoordinateTransformIndex]
311330

312331
# pixel centered coordinates
@@ -331,7 +350,7 @@ def from_variables(
331350
variables: Mapping[Any, Variable],
332351
*,
333352
options: Mapping[str, Any],
334-
) -> "RasterIndex":
353+
) -> RasterIndex:
335354
# TODO: compute bounds, resolution and affine transform from explicit coordinates.
336355
raise NotImplementedError("Creating a RasterIndex from existing coordinates is not yet supported.")
337356

@@ -343,7 +362,7 @@ def create_variables(self, variables: Mapping[Any, Variable] | None = None) -> d
343362

344363
return new_variables
345364

346-
def isel(self, indexers: Mapping[Any, int | slice | np.ndarray | Variable]) -> "RasterIndex | None":
365+
def isel(self, indexers: Mapping[Any, int | slice | np.ndarray | Variable]) -> RasterIndex | None:
347366
new_indexes: dict[WrappedIndexCoords, WrappedIndex] = {}
348367

349368
for coord_names, index in self._wrapped_indexes.items():
@@ -388,7 +407,7 @@ def equals(self, other: Index) -> bool:
388407
for k, index in self._wrapped_indexes.items()
389408
)
390409

391-
def to_pandas_index(self) -> "pd.Index":
410+
def to_pandas_index(self) -> pd.Index:
392411
# conversion is possible only if this raster index encapsulates
393412
# exactly one AxisAffineTransformIndex or a PandasIndex associated
394413
# to either the x or y axis (1-dimensional) coordinate.

src/rasterix/rioxarray_compat.py

Lines changed: 224 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,224 @@
1+
# Code adapted from `rioxarray` under the terms of the Apache-2.0 license reproduced below
2+
# Copyright (c) 2019-2023, Corteva Agriscience™
3+
# All rights reserved.
4+
#
5+
# Licensed under the Apache License, Version 2.0 (the "License");
6+
# you may not use this file except in compliance with the License.
7+
# You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing, software
12+
# distributed under the License is distributed on an "AS IS" BASIS,
13+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
# See the License for the specific language governing permissions and
15+
# limitations under the License.
16+
#
17+
#
18+
# Apache License
19+
# Version 2.0, January 2004
20+
# http://www.apache.org/licenses/
21+
#
22+
# TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
23+
#
24+
# 1. Definitions.
25+
#
26+
# "License" shall mean the terms and conditions for use, reproduction,
27+
# and distribution as defined by Sections 1 through 9 of this document.
28+
#
29+
# "Licensor" shall mean the copyright owner or entity authorized by
30+
# the copyright owner that is granting the License.
31+
#
32+
# "Legal Entity" shall mean the union of the acting entity and all
33+
# other entities that control, are controlled by, or are under common
34+
# control with that entity. For the purposes of this definition,
35+
# "control" means (i) the power, direct or indirect, to cause the
36+
# direction or management of such entity, whether by contract or
37+
# otherwise, or (ii) ownership of fifty percent (50%) or more of the
38+
# outstanding shares, or (iii) beneficial ownership of such entity.
39+
#
40+
# "You" (or "Your") shall mean an individual or Legal Entity
41+
# exercising permissions granted by this License.
42+
#
43+
# "Source" form shall mean the preferred form for making modifications,
44+
# including but not limited to software source code, documentation
45+
# source, and configuration files.
46+
#
47+
# "Object" form shall mean any form resulting from mechanical
48+
# transformation or translation of a Source form, including but
49+
# not limited to compiled object code, generated documentation,
50+
# and conversions to other media types.
51+
#
52+
# "Work" shall mean the work of authorship, whether in Source or
53+
# Object form, made available under the License, as indicated by a
54+
# copyright notice that is included in or attached to the work
55+
# (an example is provided in the Appendix below).
56+
#
57+
# "Derivative Works" shall mean any work, whether in Source or Object
58+
# form, that is based on (or derived from) the Work and for which the
59+
# editorial revisions, annotations, elaborations, or other modifications
60+
# represent, as a whole, an original work of authorship. For the purposes
61+
# of this License, Derivative Works shall not include works that remain
62+
# separable from, or merely link (or bind by name) to the interfaces of,
63+
# the Work and Derivative Works thereof.
64+
#
65+
# "Contribution" shall mean any work of authorship, including
66+
# the original version of the Work and any modifications or additions
67+
# to that Work or Derivative Works thereof, that is intentionally
68+
# submitted to Licensor for inclusion in the Work by the copyright owner
69+
# or by an individual or Legal Entity authorized to submit on behalf of
70+
# the copyright owner. For the purposes of this definition, "submitted"
71+
# means any form of electronic, verbal, or written communication sent
72+
# to the Licensor or its representatives, including but not limited to
73+
# communication on electronic mailing lists, source code control systems,
74+
# and issue tracking systems that are managed by, or on behalf of, the
75+
# Licensor for the purpose of discussing and improving the Work, but
76+
# excluding communication that is conspicuously marked or otherwise
77+
# designated in writing by the copyright owner as "Not a Contribution."
78+
#
79+
# "Contributor" shall mean Licensor and any individual or Legal Entity
80+
# on behalf of whom a Contribution has been received by Licensor and
81+
# subsequently incorporated within the Work.
82+
#
83+
# 2. Grant of Copyright License. Subject to the terms and conditions of
84+
# this License, each Contributor hereby grants to You a perpetual,
85+
# worldwide, non-exclusive, no-charge, royalty-free, irrevocable
86+
# copyright license to reproduce, prepare Derivative Works of,
87+
# publicly display, publicly perform, sublicense, and distribute the
88+
# Work and such Derivative Works in Source or Object form.
89+
#
90+
# 3. Grant of Patent License. Subject to the terms and conditions of
91+
# this License, each Contributor hereby grants to You a perpetual,
92+
# worldwide, non-exclusive, no-charge, royalty-free, irrevocable
93+
# (except as stated in this section) patent license to make, have made,
94+
# use, offer to sell, sell, import, and otherwise transfer the Work,
95+
# where such license applies only to those patent claims licensable
96+
# by such Contributor that are necessarily infringed by their
97+
# Contribution(s) alone or by combination of their Contribution(s)
98+
# with the Work to which such Contribution(s) was submitted. If You
99+
# institute patent litigation against any entity (including a
100+
# cross-claim or counterclaim in a lawsuit) alleging that the Work
101+
# or a Contribution incorporated within the Work constitutes direct
102+
# or contributory patent infringement, then any patent licenses
103+
# granted to You under this License for that Work shall terminate
104+
# as of the date such litigation is filed.
105+
#
106+
# 4. Redistribution. You may reproduce and distribute copies of the
107+
# Work or Derivative Works thereof in any medium, with or without
108+
# modifications, and in Source or Object form, provided that You
109+
# meet the following conditions:
110+
#
111+
# (a) You must give any other recipients of the Work or
112+
# Derivative Works a copy of this License; and
113+
#
114+
# (b) You must cause any modified files to carry prominent notices
115+
# stating that You changed the files; and
116+
#
117+
# (c) You must retain, in the Source form of any Derivative Works
118+
# that You distribute, all copyright, patent, trademark, and
119+
# attribution notices from the Source form of the Work,
120+
# excluding those notices that do not pertain to any part of
121+
# the Derivative Works; and
122+
#
123+
# (d) If the Work includes a "NOTICE" text file as part of its
124+
# distribution, then any Derivative Works that You distribute must
125+
# include a readable copy of the attribution notices contained
126+
# within such NOTICE file, excluding those notices that do not
127+
# pertain to any part of the Derivative Works, in at least one
128+
# of the following places: within a NOTICE text file distributed
129+
# as part of the Derivative Works; within the Source form or
130+
# documentation, if provided along with the Derivative Works; or,
131+
# within a display generated by the Derivative Works, if and
132+
# wherever such third-party notices normally appear. The contents
133+
# of the NOTICE file are for informational purposes only and
134+
# do not modify the License. You may add Your own attribution
135+
# notices within Derivative Works that You distribute, alongside
136+
# or as an addendum to the NOTICE text from the Work, provided
137+
# that such additional attribution notices cannot be construed
138+
# as modifying the License.
139+
#
140+
# You may add Your own copyright statement to Your modifications and
141+
# may provide additional or different license terms and conditions
142+
# for use, reproduction, or distribution of Your modifications, or
143+
# for any such Derivative Works as a whole, provided Your use,
144+
# reproduction, and distribution of the Work otherwise complies with
145+
# the conditions stated in this License.
146+
#
147+
# 5. Submission of Contributions. Unless You explicitly state otherwise,
148+
# any Contribution intentionally submitted for inclusion in the Work
149+
# by You to the Licensor shall be under the terms and conditions of
150+
# this License, without any additional terms or conditions.
151+
# Notwithstanding the above, nothing herein shall supersede or modify
152+
# the terms of any separate license agreement you may have executed
153+
# with Licensor regarding such Contributions.
154+
#
155+
# 6. Trademarks. This License does not grant permission to use the trade
156+
# names, trademarks, service marks, or product names of the Licensor,
157+
# except as required for reasonable and customary use in describing the
158+
# origin of the Work and reproducing the content of the NOTICE file.
159+
#
160+
# 7. Disclaimer of Warranty. Unless required by applicable law or
161+
# agreed to in writing, Licensor provides the Work (and each
162+
# Contributor provides its Contributions) on an "AS IS" BASIS,
163+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
164+
# implied, including, without limitation, any warranties or conditions
165+
# of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
166+
# PARTICULAR PURPOSE. You are solely responsible for determining the
167+
# appropriateness of using or redistributing the Work and assume any
168+
# risks associated with Your exercise of permissions under this License.
169+
#
170+
# 8. Limitation of Liability. In no event and under no legal theory,
171+
# whether in tort (including negligence), contract, or otherwise,
172+
# unless required by applicable law (such as deliberate and grossly
173+
# negligent acts) or agreed to in writing, shall any Contributor be
174+
# liable to You for damages, including any direct, indirect, special,
175+
# incidental, or consequential damages of any character arising as a
176+
# result of this License or out of the use or inability to use the
177+
# Work (including but not limited to damages for loss of goodwill,
178+
# work stoppage, computer failure or malfunction, or any and all
179+
# other commercial damages or losses), even if such Contributor
180+
# has been advised of the possibility of such damages.
181+
#
182+
# 9. Accepting Warranty or Additional Liability. While redistributing
183+
# the Work or Derivative Works thereof, You may choose to offer,
184+
# and charge a fee for, acceptance of support, warranty, indemnity,
185+
# or other liability obligations and/or rights consistent with this
186+
# License. However, in accepting such obligations, You may act only
187+
# on Your own behalf and on Your sole responsibility, not on behalf
188+
# of any other Contributor, and only if You agree to indemnify,
189+
# defend, and hold each Contributor harmless for any liability
190+
# incurred by, or claims asserted against, such Contributor by reason
191+
# of your accepting any such warranty or additional liability.
192+
#
193+
# END OF TERMS AND CONDITIONS
194+
195+
from xarray.core.types import T_Xarray
196+
197+
198+
def guess_dims(obj: T_Xarray) -> tuple[str, str]:
199+
# Determine the spatial dimensions of the `xarray.DataArray`
200+
if "x" in obj.dims and "y" in obj.dims:
201+
x_dim = "x"
202+
y_dim = "y"
203+
elif "longitude" in obj.dims and "latitude" in obj.dims:
204+
x_dim = "longitude"
205+
y_dim = "latitude"
206+
else:
207+
# look for coordinates with CF attributes
208+
for coord in obj.coords:
209+
# make sure to only look in 1D coordinates
210+
# that has the same dimension name as the coordinate
211+
if obj.coords[coord].dims != (coord,):
212+
continue
213+
if (obj.coords[coord].attrs.get("axis", "").upper() == "X") or (
214+
obj.coords[coord].attrs.get("standard_name", "").lower()
215+
in ("longitude", "projection_x_coordinate")
216+
):
217+
x_dim = coord
218+
elif (obj.coords[coord].attrs.get("axis", "").upper() == "Y") or (
219+
obj.coords[coord].attrs.get("standard_name", "").lower()
220+
in ("latitude", "projection_y_coordinate")
221+
):
222+
y_dim = coord
223+
224+
return x_dim, y_dim

tests/test_raster_index.py

Lines changed: 3 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -3,24 +3,13 @@
33
import xarray as xr
44
from affine import Affine
55

6-
from rasterix import RasterIndex
7-
8-
9-
def set_raster_index(obj):
10-
x_dim = obj.rio.x_dim
11-
y_dim = obj.rio.y_dim
12-
13-
index = RasterIndex.from_transform(
14-
obj.rio.transform(), obj.sizes[x_dim], obj.sizes[y_dim], x_dim=x_dim, y_dim=y_dim
15-
)
16-
coords = xr.Coordinates.from_xindex(index)
17-
return obj.assign_coords(coords)
6+
from rasterix import RasterIndex, assign_index
187

198

209
def test_rectilinear():
2110
source = "/vsicurl/https://noaadata.apps.nsidc.org/NOAA/G02135/south/daily/geotiff/2024/01_Jan/S_20240101_concentration_v3.0.tif"
2211
da_no_raster_index = xr.open_dataarray(source, engine="rasterio")
23-
da_raster_index = set_raster_index(da_no_raster_index)
12+
da_raster_index = assign_index(da_no_raster_index)
2413
assert da_raster_index.equals(da_no_raster_index)
2514

2615

@@ -31,7 +20,7 @@ def test_sel_slice():
3120
ds = xr.Dataset({"foo": (("y", "x"), np.ones((10, 12)))})
3221
transform = Affine.identity()
3322
ds = ds.rio.write_transform(transform)
34-
ds = set_raster_index(ds)
23+
ds = assign_index(ds)
3524

3625
assert ds.xindexes["x"].transform() == transform
3726

0 commit comments

Comments
 (0)