1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385 386 387 388 389
|
import numpy as np
import pandas as pd
import pytest
import xarray as xr
from xarray.indexes import PandasIndex, RangeIndex
from xarray.tests import assert_allclose, assert_equal, assert_identical
def create_dataset_arange(
start: float, stop: float, step: float, dim: str = "x"
) -> xr.Dataset:
index = RangeIndex.arange(start, stop, step, dim=dim)
return xr.Dataset(coords=xr.Coordinates.from_xindex(index))
@pytest.mark.parametrize(
"args,kwargs",
[
((10.0,), {}),
((), {"stop": 10.0}),
(
(
2.0,
10.0,
),
{},
),
((2.0,), {"stop": 10.0}),
((), {"start": 2.0, "stop": 10.0}),
((2.0, 10.0, 2.0), {}),
((), {"start": 2.0, "stop": 10.0, "step": 2.0}),
],
)
def test_range_index_arange(args, kwargs) -> None:
index = RangeIndex.arange(*args, **kwargs, dim="x")
actual = xr.Coordinates.from_xindex(index)
expected = xr.Coordinates({"x": np.arange(*args, **kwargs)})
assert_equal(actual, expected, check_default_indexes=False)
def test_range_index_arange_error() -> None:
with pytest.raises(TypeError, match=r".*requires stop to be specified"):
RangeIndex.arange(dim="x")
def test_range_index_arange_start_as_stop() -> None:
# Weird although probably very unlikely case where only `start` is given
# as keyword argument, which is interpreted as `stop`.
# This has been fixed in numpy (https://github.com/numpy/numpy/pull/17878)
# using Python C API. In pure Python it's more tricky as there's no easy way to know
# whether a value has been passed as positional or keyword argument.
# Note: `pandas.RangeIndex` constructor still has this weird behavior.
index = RangeIndex.arange(start=10.0, dim="x")
actual = xr.Coordinates.from_xindex(index)
expected = xr.Coordinates({"x": np.arange(10.0)})
assert_equal(actual, expected, check_default_indexes=False)
def test_range_index_arange_properties() -> None:
index = RangeIndex.arange(0.0, 1.0, 0.1, dim="x")
assert index.start == 0.0
assert index.stop == 1.0
assert index.step == 0.1
def test_range_index_linspace() -> None:
index = RangeIndex.linspace(0.0, 1.0, num=10, endpoint=False, dim="x")
actual = xr.Coordinates.from_xindex(index)
expected = xr.Coordinates({"x": np.linspace(0.0, 1.0, num=10, endpoint=False)})
assert_equal(actual, expected, check_default_indexes=False)
assert index.start == 0.0
assert index.stop == 1.0
assert index.step == 0.1
index = RangeIndex.linspace(0.0, 1.0, num=11, endpoint=True, dim="x")
actual = xr.Coordinates.from_xindex(index)
expected = xr.Coordinates({"x": np.linspace(0.0, 1.0, num=11, endpoint=True)})
assert_allclose(actual, expected, check_default_indexes=False)
assert index.start == 0.0
assert index.stop == 1.1
assert index.step == 0.1
def test_range_index_dtype() -> None:
index = RangeIndex.arange(0.0, 1.0, 0.1, dim="x", dtype=np.float32)
coords = xr.Coordinates.from_xindex(index)
assert coords["x"].dtype == np.dtype(np.float32)
def test_range_index_set_xindex() -> None:
coords = xr.Coordinates({"x": np.arange(0.0, 1.0, 0.1)}, indexes={})
ds = xr.Dataset(coords=coords)
with pytest.raises(
NotImplementedError, match=r"cannot create.*RangeIndex.*existing coordinate"
):
ds.set_xindex("x", RangeIndex)
def test_range_index_isel() -> None:
ds = create_dataset_arange(0.0, 1.0, 0.1)
# slicing
actual = ds.isel(x=slice(None))
assert_identical(actual, ds, check_default_indexes=False, check_indexes=True)
actual = ds.isel(x=slice(1, None))
expected = create_dataset_arange(0.1, 1.0, 0.1)
assert_identical(actual, expected, check_default_indexes=False, check_indexes=True)
actual = ds.isel(x=slice(None, 2))
expected = create_dataset_arange(0.0, 0.2, 0.1)
assert_identical(actual, expected, check_default_indexes=False, check_indexes=True)
actual = ds.isel(x=slice(1, 3))
expected = create_dataset_arange(0.1, 0.3, 0.1)
assert_identical(actual, expected, check_default_indexes=False, check_indexes=True)
actual = ds.isel(x=slice(None, None, 2))
expected = create_dataset_arange(0.0, 1.0, 0.2)
assert_identical(actual, expected, check_default_indexes=False, check_indexes=True)
actual = ds.isel(x=slice(None, None, -1))
expected = create_dataset_arange(0.9, -0.1, -0.1)
assert_identical(actual, expected, check_default_indexes=False, check_indexes=True)
actual = ds.isel(x=slice(None, 4, -1))
expected = create_dataset_arange(0.9, 0.4, -0.1)
assert_identical(actual, expected, check_default_indexes=False, check_indexes=True)
actual = ds.isel(x=slice(8, 4, -1))
expected = create_dataset_arange(0.8, 0.4, -0.1)
assert_identical(actual, expected, check_default_indexes=False, check_indexes=True)
actual = ds.isel(x=slice(8, None, -1))
expected = create_dataset_arange(0.8, -0.1, -0.1)
assert_identical(actual, expected, check_default_indexes=False, check_indexes=True)
# https://github.com/pydata/xarray/issues/10441
ds2 = create_dataset_arange(0.0, 3.0, 0.1)
actual = ds2.isel(x=slice(4, None, 3))
expected = create_dataset_arange(0.4, 3.0, 0.3)
assert_identical(actual, expected, check_default_indexes=False, check_indexes=True)
# scalar
actual = ds.isel(x=0)
expected = xr.Dataset(coords={"x": 0.0})
assert_identical(actual, expected)
# outer indexing with arbitrary array values
actual = ds.isel(x=[0, 2])
expected = xr.Dataset(coords={"x": [0.0, 0.2]})
assert_identical(actual, expected)
assert isinstance(actual.xindexes["x"], PandasIndex)
# fancy indexing with 1-d Variable
actual = ds.isel(x=xr.Variable("y", [0, 2]))
expected = xr.Dataset(coords={"x": ("y", [0.0, 0.2])}).set_xindex("x")
assert_identical(actual, expected, check_default_indexes=False, check_indexes=True)
assert isinstance(actual.xindexes["x"], PandasIndex)
# fancy indexing with n-d Variable
actual = ds.isel(x=xr.Variable(("u", "v"), [[0, 0], [2, 2]]))
expected = xr.Dataset(coords={"x": (("u", "v"), [[0.0, 0.0], [0.2, 0.2]])})
assert_identical(actual, expected)
def test_range_index_empty_slice() -> None:
"""Test that empty slices of RangeIndex are printable and preserve step.
Regression test for https://github.com/pydata/xarray/issues/10547
"""
# Test with linspace
n = 30
step = 1
da = xr.DataArray(np.zeros(n), dims=["x"])
da = da.assign_coords(
xr.Coordinates.from_xindex(RangeIndex.linspace(0, (n - 1) * step, n, dim="x"))
)
# This should not raise ZeroDivisionError
sub = da.isel(x=slice(0))
assert sub.sizes["x"] == 0
# Test that it's printable
repr_str = repr(sub)
assert "RangeIndex" in repr_str
assert "step=1" in repr_str
# Test with different step values
index = RangeIndex.arange(0, 10, 2.5, dim="y")
da2 = xr.DataArray(np.zeros(4), dims=["y"])
da2 = da2.assign_coords(xr.Coordinates.from_xindex(index))
empty = da2.isel(y=slice(0))
# Should preserve step
assert empty.sizes["y"] == 0
range_index_y = empty._indexes["y"]
assert isinstance(range_index_y, RangeIndex)
assert range_index_y.step == 2.5
# Test that it's printable
repr_str2 = repr(empty)
assert "RangeIndex" in repr_str2
assert "step=2.5" in repr_str2
# Test negative step
index3 = RangeIndex.arange(10, 0, -1, dim="z")
da3 = xr.DataArray(np.zeros(10), dims=["z"])
da3 = da3.assign_coords(xr.Coordinates.from_xindex(index3))
empty3 = da3.isel(z=slice(0))
assert empty3.sizes["z"] == 0
range_index_z = empty3._indexes["z"]
assert isinstance(range_index_z, RangeIndex)
assert range_index_z.step == -1.0
# Test that it's printable
repr_str3 = repr(empty3)
assert "RangeIndex" in repr_str3
assert "step=-1" in repr_str3
def test_range_index_sel() -> None:
ds = create_dataset_arange(0.0, 1.0, 0.1)
# start-stop slice
actual = ds.sel(x=slice(0.12, 0.28), method="nearest")
expected = create_dataset_arange(0.1, 0.3, 0.1)
assert_identical(actual, expected, check_default_indexes=False, check_indexes=True)
# start-stop-step slice
actual = ds.sel(x=slice(0.0, 1.0, 0.2), method="nearest")
expected = ds.isel(x=range(0, 10, 2))
assert_identical(actual, expected, check_default_indexes=False, check_indexes=True)
# basic indexing
actual = ds.sel(x=0.52, method="nearest")
expected = xr.Dataset(coords={"x": 0.5})
assert_allclose(actual, expected)
actual = ds.sel(x=0.58, method="nearest")
expected = xr.Dataset(coords={"x": 0.6})
assert_allclose(actual, expected)
# 1-d array indexing
actual = ds.sel(x=[0.52, 0.58], method="nearest")
expected = xr.Dataset(coords={"x": [0.5, 0.6]})
assert_allclose(actual, expected)
actual = ds.sel(x=xr.Variable("y", [0.52, 0.58]), method="nearest")
expected = xr.Dataset(coords={"x": ("y", [0.5, 0.6])}).set_xindex("x")
assert_allclose(actual, expected, check_default_indexes=False)
actual = ds.sel(x=xr.DataArray([0.52, 0.58], dims="y"), method="nearest")
expected = xr.Dataset(coords={"x": ("y", [0.5, 0.6])}).set_xindex("x")
assert_allclose(actual, expected, check_default_indexes=False)
with pytest.raises(ValueError, match=r"RangeIndex only supports.*method.*nearest"):
ds.sel(x=0.1)
with pytest.raises(ValueError, match=r"RangeIndex doesn't support.*tolerance"):
ds.sel(x=0.1, method="nearest", tolerance=1e-3)
def test_range_index_to_pandas_index() -> None:
ds = create_dataset_arange(0.0, 1.0, 0.1)
actual = ds.indexes["x"]
expected = pd.Index(np.arange(0.0, 1.0, 0.1))
assert actual.equals(expected)
def test_range_index_rename() -> None:
index = RangeIndex.arange(0.0, 1.0, 0.1, dim="x")
ds = xr.Dataset(coords=xr.Coordinates.from_xindex(index))
actual = ds.rename_vars(x="y")
idx = RangeIndex.arange(0.0, 1.0, 0.1, coord_name="y", dim="x")
expected = xr.Dataset(coords=xr.Coordinates.from_xindex(idx))
assert_identical(actual, expected, check_default_indexes=False, check_indexes=True)
actual = ds.rename_dims(x="y")
idx = RangeIndex.arange(0.0, 1.0, 0.1, coord_name="x", dim="y")
expected = xr.Dataset(coords=xr.Coordinates.from_xindex(idx))
assert_identical(actual, expected, check_default_indexes=False, check_indexes=True)
def test_range_index_repr() -> None:
index = RangeIndex.arange(0.0, 1.0, 0.1, dim="x")
actual = repr(index)
expected = (
"RangeIndex (start=0, stop=1, step=0.1, size=10, coord_name='x', dim='x')"
)
assert actual == expected
def test_range_index_repr_inline() -> None:
index = RangeIndex.arange(0.0, 1.0, 0.1, dim="x")
actual = index._repr_inline_(max_width=70)
expected = "RangeIndex (start=0, stop=1, step=0.1)"
assert actual == expected
def test_range_index_equals_floating_point_tolerance() -> None:
"""Test that equals() handles floating point precision errors correctly.
When slicing a RangeIndex, floating point errors can accumulate in the
internal state (e.g., stop=0.30000000000000004 vs stop=0.3), but the
indexes should still be considered equal if they represent the same values.
"""
# Create an index directly
index1 = RangeIndex.arange(0.0, 0.3, 0.1, dim="x")
# Create the same index by slicing a larger one
# This will accumulate floating point error: stop = 0.0 + 3 * 0.1 = 0.30000000000000004
index_large = RangeIndex.arange(0.0, 1.0, 0.1, dim="x")
ds_large = xr.Dataset(coords=xr.Coordinates.from_xindex(index_large))
ds_sliced = ds_large.isel(x=slice(3))
index2 = ds_sliced.xindexes["x"]
# They should be equal despite tiny floating point differences
assert index1.equals(index2)
assert index2.equals(index1)
# Verify they represent the same values
ds1 = xr.Dataset(coords=xr.Coordinates.from_xindex(index1))
ds2 = xr.Dataset(coords=xr.Coordinates.from_xindex(index2))
assert np.allclose(ds1["x"].values, ds2["x"].values)
def test_range_index_equals_different_sizes() -> None:
"""Test that equals() returns False for indexes with different sizes."""
index1 = RangeIndex.arange(0.0, 0.3, 0.1, dim="x")
index2 = RangeIndex.arange(0.0, 0.4, 0.1, dim="x")
assert not index1.equals(index2)
assert not index2.equals(index1)
def test_range_index_equals_different_start() -> None:
"""Test that equals() returns False for indexes with significantly different start values."""
index1 = RangeIndex.arange(0.0, 0.3, 0.1, dim="x")
index2 = RangeIndex.arange(0.1, 0.4, 0.1, dim="x")
assert not index1.equals(index2)
assert not index2.equals(index1)
def test_range_index_equals_different_stop() -> None:
"""Test that equals() returns False for indexes with significantly different stop values."""
index1 = RangeIndex.arange(0.0, 0.3, 0.1, dim="x")
index2 = RangeIndex.arange(0.0, 0.5, 0.1, dim="x")
assert not index1.equals(index2)
assert not index2.equals(index1)
def test_range_index_equals_different_type() -> None:
"""Test that equals() returns False when comparing with a different index type."""
index1 = RangeIndex.arange(0.0, 0.3, 0.1, dim="x")
index2 = PandasIndex(pd.Index([0.0, 0.1, 0.2]), dim="x")
assert not index1.equals(index2)
# Note: we don't test index2.equals(index1) because PandasIndex.equals()
# has its own logic
def test_range_index_equals_exact() -> None:
"""Test that equals(exact=True) requires exact floating point match."""
# Create an index directly
index1 = RangeIndex.arange(0.0, 0.3, 0.1, dim="x")
# Create the same index by slicing - this accumulates floating point error
index_large = RangeIndex.arange(0.0, 1.0, 0.1, dim="x")
ds_large = xr.Dataset(coords=xr.Coordinates.from_xindex(index_large))
ds_sliced = ds_large.isel(x=slice(3))
index2 = ds_sliced.xindexes["x"]
# Default (exact=False) should be equal due to np.isclose tolerance
assert index1.equals(index2)
# With exact=True, tiny floating point differences cause inequality
assert not index1.equals(index2, exact=True)
# But identical indexes should still be equal with exact=True
index3 = RangeIndex.arange(0.0, 0.3, 0.1, dim="x")
assert index1.equals(index3, exact=True)
|