1import pickle 2import sys 3import warnings 4from copy import deepcopy 5from textwrap import dedent 6 7import numpy as np 8import pandas as pd 9import pytest 10from pandas.core.computation.ops import UndefinedVariableError 11 12import xarray as xr 13from xarray import ( 14 DataArray, 15 Dataset, 16 IndexVariable, 17 Variable, 18 align, 19 broadcast, 20 set_options, 21) 22from xarray.coding.times import CFDatetimeCoder 23from xarray.convert import from_cdms2 24from xarray.core import dtypes 25from xarray.core.common import full_like 26from xarray.core.indexes import Index, PandasIndex, propagate_indexes 27from xarray.core.utils import is_scalar 28from xarray.tests import ( 29 LooseVersion, 30 ReturnItem, 31 assert_allclose, 32 assert_array_equal, 33 assert_equal, 34 assert_identical, 35 has_dask, 36 raise_if_dask_computes, 37 requires_bottleneck, 38 requires_cupy, 39 requires_dask, 40 requires_iris, 41 requires_numbagg, 42 requires_numexpr, 43 requires_pint, 44 requires_scipy, 45 requires_sparse, 46 source_ndarray, 47) 48 49pytestmark = [ 50 pytest.mark.filterwarnings("error:Mean of empty slice"), 51 pytest.mark.filterwarnings("error:All-NaN (slice|axis) encountered"), 52] 53 54 55class TestDataArray: 56 @pytest.fixture(autouse=True) 57 def setup(self): 58 self.attrs = {"attr1": "value1", "attr2": 2929} 59 self.x = np.random.random((10, 20)) 60 self.v = Variable(["x", "y"], self.x) 61 self.va = Variable(["x", "y"], self.x, self.attrs) 62 self.ds = Dataset({"foo": self.v}) 63 self.dv = self.ds["foo"] 64 65 self.mindex = pd.MultiIndex.from_product( 66 [["a", "b"], [1, 2]], names=("level_1", "level_2") 67 ) 68 self.mda = DataArray([0, 1, 2, 3], coords={"x": self.mindex}, dims="x") 69 70 def test_repr(self): 71 v = Variable(["time", "x"], [[1, 2, 3], [4, 5, 6]], {"foo": "bar"}) 72 coords = {"x": np.arange(3, dtype=np.int64), "other": np.int64(0)} 73 data_array = DataArray(v, coords, name="my_variable") 74 expected = dedent( 75 """\ 76 <xarray.DataArray 'my_variable' (time: 2, x: 3)> 77 array([[1, 2, 3], 78 [4, 5, 6]]) 79 Coordinates: 80 * x (x) int64 0 1 2 81 other int64 0 82 Dimensions without coordinates: time 83 Attributes: 84 foo: bar""" 85 ) 86 assert expected == repr(data_array) 87 88 def test_repr_multiindex(self): 89 expected = dedent( 90 """\ 91 <xarray.DataArray (x: 4)> 92 array([0, 1, 2, 3]) 93 Coordinates: 94 * x (x) MultiIndex 95 - level_1 (x) object 'a' 'a' 'b' 'b' 96 - level_2 (x) int64 1 2 1 2""" 97 ) 98 assert expected == repr(self.mda) 99 100 def test_repr_multiindex_long(self): 101 mindex_long = pd.MultiIndex.from_product( 102 [["a", "b", "c", "d"], [1, 2, 3, 4, 5, 6, 7, 8]], 103 names=("level_1", "level_2"), 104 ) 105 mda_long = DataArray(list(range(32)), coords={"x": mindex_long}, dims="x") 106 expected = dedent( 107 """\ 108 <xarray.DataArray (x: 32)> 109 array([ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 110 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31]) 111 Coordinates: 112 * x (x) MultiIndex 113 - level_1 (x) object 'a' 'a' 'a' 'a' 'a' 'a' 'a' ... 'd' 'd' 'd' 'd' 'd' 'd' 114 - level_2 (x) int64 1 2 3 4 5 6 7 8 1 2 3 4 5 6 ... 4 5 6 7 8 1 2 3 4 5 6 7 8""" 115 ) 116 assert expected == repr(mda_long) 117 118 def test_properties(self): 119 assert_equal(self.dv.variable, self.v) 120 assert_array_equal(self.dv.values, self.v.values) 121 for attr in ["dims", "dtype", "shape", "size", "nbytes", "ndim", "attrs"]: 122 assert getattr(self.dv, attr) == getattr(self.v, attr) 123 assert len(self.dv) == len(self.v) 124 assert_equal(self.dv.variable, self.v) 125 assert set(self.dv.coords) == set(self.ds.coords) 126 for k, v in self.dv.coords.items(): 127 assert_array_equal(v, self.ds.coords[k]) 128 with pytest.raises(AttributeError): 129 self.dv.dataset 130 assert isinstance(self.ds["x"].to_index(), pd.Index) 131 with pytest.raises(ValueError, match=r"must be 1-dimensional"): 132 self.ds["foo"].to_index() 133 with pytest.raises(AttributeError): 134 self.dv.variable = self.v 135 136 def test_data_property(self): 137 array = DataArray(np.zeros((3, 4))) 138 actual = array.copy() 139 actual.values = np.ones((3, 4)) 140 assert_array_equal(np.ones((3, 4)), actual.values) 141 actual.data = 2 * np.ones((3, 4)) 142 assert_array_equal(2 * np.ones((3, 4)), actual.data) 143 assert_array_equal(actual.data, actual.values) 144 145 def test_indexes(self): 146 array = DataArray(np.zeros((2, 3)), [("x", [0, 1]), ("y", ["a", "b", "c"])]) 147 expected_indexes = {"x": pd.Index([0, 1]), "y": pd.Index(["a", "b", "c"])} 148 expected_xindexes = { 149 k: PandasIndex(idx, k) for k, idx in expected_indexes.items() 150 } 151 assert array.xindexes.keys() == expected_xindexes.keys() 152 assert array.indexes.keys() == expected_indexes.keys() 153 assert all([isinstance(idx, pd.Index) for idx in array.indexes.values()]) 154 assert all([isinstance(idx, Index) for idx in array.xindexes.values()]) 155 for k in expected_indexes: 156 assert array.xindexes[k].equals(expected_xindexes[k]) 157 assert array.indexes[k].equals(expected_indexes[k]) 158 159 def test_get_index(self): 160 array = DataArray(np.zeros((2, 3)), coords={"x": ["a", "b"]}, dims=["x", "y"]) 161 assert array.get_index("x").equals(pd.Index(["a", "b"])) 162 assert array.get_index("y").equals(pd.Index([0, 1, 2])) 163 with pytest.raises(KeyError): 164 array.get_index("z") 165 166 def test_get_index_size_zero(self): 167 array = DataArray(np.zeros((0,)), dims=["x"]) 168 actual = array.get_index("x") 169 expected = pd.Index([], dtype=np.int64) 170 assert actual.equals(expected) 171 assert actual.dtype == expected.dtype 172 173 def test_struct_array_dims(self): 174 """ 175 This test checks subraction of two DataArrays for the case 176 when dimension is a structured array. 177 """ 178 # GH837, GH861 179 # checking array subtraction when dims are the same 180 p_data = np.array( 181 [("Abe", 180), ("Stacy", 150), ("Dick", 200)], 182 dtype=[("name", "|S256"), ("height", object)], 183 ) 184 weights_0 = DataArray( 185 [80, 56, 120], dims=["participant"], coords={"participant": p_data} 186 ) 187 weights_1 = DataArray( 188 [81, 52, 115], dims=["participant"], coords={"participant": p_data} 189 ) 190 actual = weights_1 - weights_0 191 192 expected = DataArray( 193 [1, -4, -5], dims=["participant"], coords={"participant": p_data} 194 ) 195 196 assert_identical(actual, expected) 197 198 # checking array subraction when dims are not the same 199 p_data_alt = np.array( 200 [("Abe", 180), ("Stacy", 151), ("Dick", 200)], 201 dtype=[("name", "|S256"), ("height", object)], 202 ) 203 weights_1 = DataArray( 204 [81, 52, 115], dims=["participant"], coords={"participant": p_data_alt} 205 ) 206 actual = weights_1 - weights_0 207 208 expected = DataArray( 209 [1, -5], dims=["participant"], coords={"participant": p_data[[0, 2]]} 210 ) 211 212 assert_identical(actual, expected) 213 214 # checking array subraction when dims are not the same and one 215 # is np.nan 216 p_data_nan = np.array( 217 [("Abe", 180), ("Stacy", np.nan), ("Dick", 200)], 218 dtype=[("name", "|S256"), ("height", object)], 219 ) 220 weights_1 = DataArray( 221 [81, 52, 115], dims=["participant"], coords={"participant": p_data_nan} 222 ) 223 actual = weights_1 - weights_0 224 225 expected = DataArray( 226 [1, -5], dims=["participant"], coords={"participant": p_data[[0, 2]]} 227 ) 228 229 assert_identical(actual, expected) 230 231 def test_name(self): 232 arr = self.dv 233 assert arr.name == "foo" 234 235 copied = arr.copy() 236 arr.name = "bar" 237 assert arr.name == "bar" 238 assert_equal(copied, arr) 239 240 actual = DataArray(IndexVariable("x", [3])) 241 actual.name = "y" 242 expected = DataArray([3], [("x", [3])], name="y") 243 assert_identical(actual, expected) 244 245 def test_dims(self): 246 arr = self.dv 247 assert arr.dims == ("x", "y") 248 249 with pytest.raises(AttributeError, match=r"you cannot assign"): 250 arr.dims = ("w", "z") 251 252 def test_sizes(self): 253 array = DataArray(np.zeros((3, 4)), dims=["x", "y"]) 254 assert array.sizes == {"x": 3, "y": 4} 255 assert tuple(array.sizes) == array.dims 256 with pytest.raises(TypeError): 257 array.sizes["foo"] = 5 258 259 def test_encoding(self): 260 expected = {"foo": "bar"} 261 self.dv.encoding["foo"] = "bar" 262 assert expected == self.dv.encoding 263 264 expected = {"baz": 0} 265 self.dv.encoding = expected 266 267 assert expected is not self.dv.encoding 268 269 def test_constructor(self): 270 data = np.random.random((2, 3)) 271 272 actual = DataArray(data) 273 expected = Dataset({None: (["dim_0", "dim_1"], data)})[None] 274 assert_identical(expected, actual) 275 276 actual = DataArray(data, [["a", "b"], [-1, -2, -3]]) 277 expected = Dataset( 278 { 279 None: (["dim_0", "dim_1"], data), 280 "dim_0": ("dim_0", ["a", "b"]), 281 "dim_1": ("dim_1", [-1, -2, -3]), 282 } 283 )[None] 284 assert_identical(expected, actual) 285 286 actual = DataArray( 287 data, [pd.Index(["a", "b"], name="x"), pd.Index([-1, -2, -3], name="y")] 288 ) 289 expected = Dataset( 290 {None: (["x", "y"], data), "x": ("x", ["a", "b"]), "y": ("y", [-1, -2, -3])} 291 )[None] 292 assert_identical(expected, actual) 293 294 coords = [["a", "b"], [-1, -2, -3]] 295 actual = DataArray(data, coords, ["x", "y"]) 296 assert_identical(expected, actual) 297 298 coords = [pd.Index(["a", "b"], name="A"), pd.Index([-1, -2, -3], name="B")] 299 actual = DataArray(data, coords, ["x", "y"]) 300 assert_identical(expected, actual) 301 302 coords = {"x": ["a", "b"], "y": [-1, -2, -3]} 303 actual = DataArray(data, coords, ["x", "y"]) 304 assert_identical(expected, actual) 305 306 actual = DataArray(data, coords) 307 assert_identical(expected, actual) 308 309 coords = [("x", ["a", "b"]), ("y", [-1, -2, -3])] 310 actual = DataArray(data, coords) 311 assert_identical(expected, actual) 312 313 expected = Dataset({None: (["x", "y"], data), "x": ("x", ["a", "b"])})[None] 314 actual = DataArray(data, {"x": ["a", "b"]}, ["x", "y"]) 315 assert_identical(expected, actual) 316 317 actual = DataArray(data, dims=["x", "y"]) 318 expected = Dataset({None: (["x", "y"], data)})[None] 319 assert_identical(expected, actual) 320 321 actual = DataArray(data, dims=["x", "y"], name="foo") 322 expected = Dataset({"foo": (["x", "y"], data)})["foo"] 323 assert_identical(expected, actual) 324 325 actual = DataArray(data, name="foo") 326 expected = Dataset({"foo": (["dim_0", "dim_1"], data)})["foo"] 327 assert_identical(expected, actual) 328 329 actual = DataArray(data, dims=["x", "y"], attrs={"bar": 2}) 330 expected = Dataset({None: (["x", "y"], data, {"bar": 2})})[None] 331 assert_identical(expected, actual) 332 333 actual = DataArray(data, dims=["x", "y"]) 334 expected = Dataset({None: (["x", "y"], data, {}, {"bar": 2})})[None] 335 assert_identical(expected, actual) 336 337 actual = DataArray([1, 2, 3], coords={"x": [0, 1, 2]}) 338 expected = DataArray([1, 2, 3], coords=[("x", [0, 1, 2])]) 339 assert_identical(expected, actual) 340 341 def test_constructor_invalid(self): 342 data = np.random.randn(3, 2) 343 344 with pytest.raises(ValueError, match=r"coords is not dict-like"): 345 DataArray(data, [[0, 1, 2]], ["x", "y"]) 346 347 with pytest.raises(ValueError, match=r"not a subset of the .* dim"): 348 DataArray(data, {"x": [0, 1, 2]}, ["a", "b"]) 349 with pytest.raises(ValueError, match=r"not a subset of the .* dim"): 350 DataArray(data, {"x": [0, 1, 2]}) 351 352 with pytest.raises(TypeError, match=r"is not a string"): 353 DataArray(data, dims=["x", None]) 354 355 with pytest.raises(ValueError, match=r"conflicting sizes for dim"): 356 DataArray([1, 2, 3], coords=[("x", [0, 1])]) 357 with pytest.raises(ValueError, match=r"conflicting sizes for dim"): 358 DataArray([1, 2], coords={"x": [0, 1], "y": ("x", [1])}, dims="x") 359 360 with pytest.raises(ValueError, match=r"conflicting MultiIndex"): 361 DataArray(np.random.rand(4, 4), [("x", self.mindex), ("y", self.mindex)]) 362 with pytest.raises(ValueError, match=r"conflicting MultiIndex"): 363 DataArray(np.random.rand(4, 4), [("x", self.mindex), ("level_1", range(4))]) 364 365 with pytest.raises(ValueError, match=r"matching the dimension size"): 366 DataArray(data, coords={"x": 0}, dims=["x", "y"]) 367 368 def test_constructor_from_self_described(self): 369 data = [[-0.1, 21], [0, 2]] 370 expected = DataArray( 371 data, 372 coords={"x": ["a", "b"], "y": [-1, -2]}, 373 dims=["x", "y"], 374 name="foobar", 375 attrs={"bar": 2}, 376 ) 377 actual = DataArray(expected) 378 assert_identical(expected, actual) 379 380 actual = DataArray(expected.values, actual.coords) 381 assert_equal(expected, actual) 382 383 frame = pd.DataFrame( 384 data, 385 index=pd.Index(["a", "b"], name="x"), 386 columns=pd.Index([-1, -2], name="y"), 387 ) 388 actual = DataArray(frame) 389 assert_equal(expected, actual) 390 391 series = pd.Series(data[0], index=pd.Index([-1, -2], name="y")) 392 actual = DataArray(series) 393 assert_equal(expected[0].reset_coords("x", drop=True), actual) 394 395 expected = DataArray( 396 data, 397 coords={"x": ["a", "b"], "y": [-1, -2], "a": 0, "z": ("x", [-0.5, 0.5])}, 398 dims=["x", "y"], 399 ) 400 actual = DataArray(expected) 401 assert_identical(expected, actual) 402 403 actual = DataArray(expected.values, expected.coords) 404 assert_identical(expected, actual) 405 406 expected = Dataset({"foo": ("foo", ["a", "b"])})["foo"] 407 actual = DataArray(pd.Index(["a", "b"], name="foo")) 408 assert_identical(expected, actual) 409 410 actual = DataArray(IndexVariable("foo", ["a", "b"])) 411 assert_identical(expected, actual) 412 413 def test_constructor_from_0d(self): 414 expected = Dataset({None: ([], 0)})[None] 415 actual = DataArray(0) 416 assert_identical(expected, actual) 417 418 @requires_dask 419 def test_constructor_dask_coords(self): 420 # regression test for GH1684 421 import dask.array as da 422 423 coord = da.arange(8, chunks=(4,)) 424 data = da.random.random((8, 8), chunks=(4, 4)) + 1 425 actual = DataArray(data, coords={"x": coord, "y": coord}, dims=["x", "y"]) 426 427 ecoord = np.arange(8) 428 expected = DataArray(data, coords={"x": ecoord, "y": ecoord}, dims=["x", "y"]) 429 assert_equal(actual, expected) 430 431 def test_equals_and_identical(self): 432 orig = DataArray(np.arange(5.0), {"a": 42}, dims="x") 433 434 expected = orig 435 actual = orig.copy() 436 assert expected.equals(actual) 437 assert expected.identical(actual) 438 439 actual = expected.rename("baz") 440 assert expected.equals(actual) 441 assert not expected.identical(actual) 442 443 actual = expected.rename({"x": "xxx"}) 444 assert not expected.equals(actual) 445 assert not expected.identical(actual) 446 447 actual = expected.copy() 448 actual.attrs["foo"] = "bar" 449 assert expected.equals(actual) 450 assert not expected.identical(actual) 451 452 actual = expected.copy() 453 actual["x"] = ("x", -np.arange(5)) 454 assert not expected.equals(actual) 455 assert not expected.identical(actual) 456 457 actual = expected.reset_coords(drop=True) 458 assert not expected.equals(actual) 459 assert not expected.identical(actual) 460 461 actual = orig.copy() 462 actual[0] = np.nan 463 expected = actual.copy() 464 assert expected.equals(actual) 465 assert expected.identical(actual) 466 467 actual[:] = np.nan 468 assert not expected.equals(actual) 469 assert not expected.identical(actual) 470 471 actual = expected.copy() 472 actual["a"] = 100000 473 assert not expected.equals(actual) 474 assert not expected.identical(actual) 475 476 def test_equals_failures(self): 477 orig = DataArray(np.arange(5.0), {"a": 42}, dims="x") 478 assert not orig.equals(np.arange(5)) 479 assert not orig.identical(123) 480 assert not orig.broadcast_equals({1: 2}) 481 482 def test_broadcast_equals(self): 483 a = DataArray([0, 0], {"y": 0}, dims="x") 484 b = DataArray([0, 0], {"y": ("x", [0, 0])}, dims="x") 485 assert a.broadcast_equals(b) 486 assert b.broadcast_equals(a) 487 assert not a.equals(b) 488 assert not a.identical(b) 489 490 c = DataArray([0], coords={"x": 0}, dims="y") 491 assert not a.broadcast_equals(c) 492 assert not c.broadcast_equals(a) 493 494 def test_getitem(self): 495 # strings pull out dataarrays 496 assert_identical(self.dv, self.ds["foo"]) 497 x = self.dv["x"] 498 y = self.dv["y"] 499 assert_identical(self.ds["x"], x) 500 assert_identical(self.ds["y"], y) 501 502 arr = ReturnItem() 503 for i in [ 504 arr[:], 505 arr[...], 506 arr[x.values], 507 arr[x.variable], 508 arr[x], 509 arr[x, y], 510 arr[x.values > -1], 511 arr[x.variable > -1], 512 arr[x > -1], 513 arr[x > -1, y > -1], 514 ]: 515 assert_equal(self.dv, self.dv[i]) 516 for i in [ 517 arr[0], 518 arr[:, 0], 519 arr[:3, :2], 520 arr[x.values[:3]], 521 arr[x.variable[:3]], 522 arr[x[:3]], 523 arr[x[:3], y[:4]], 524 arr[x.values > 3], 525 arr[x.variable > 3], 526 arr[x > 3], 527 arr[x > 3, y > 3], 528 ]: 529 assert_array_equal(self.v[i], self.dv[i]) 530 531 def test_getitem_dict(self): 532 actual = self.dv[{"x": slice(3), "y": 0}] 533 expected = self.dv.isel(x=slice(3), y=0) 534 assert_identical(expected, actual) 535 536 def test_getitem_coords(self): 537 orig = DataArray( 538 [[10], [20]], 539 { 540 "x": [1, 2], 541 "y": [3], 542 "z": 4, 543 "x2": ("x", ["a", "b"]), 544 "y2": ("y", ["c"]), 545 "xy": (["y", "x"], [["d", "e"]]), 546 }, 547 dims=["x", "y"], 548 ) 549 550 assert_identical(orig, orig[:]) 551 assert_identical(orig, orig[:, :]) 552 assert_identical(orig, orig[...]) 553 assert_identical(orig, orig[:2, :1]) 554 assert_identical(orig, orig[[0, 1], [0]]) 555 556 actual = orig[0, 0] 557 expected = DataArray( 558 10, {"x": 1, "y": 3, "z": 4, "x2": "a", "y2": "c", "xy": "d"} 559 ) 560 assert_identical(expected, actual) 561 562 actual = orig[0, :] 563 expected = DataArray( 564 [10], 565 { 566 "x": 1, 567 "y": [3], 568 "z": 4, 569 "x2": "a", 570 "y2": ("y", ["c"]), 571 "xy": ("y", ["d"]), 572 }, 573 dims="y", 574 ) 575 assert_identical(expected, actual) 576 577 actual = orig[:, 0] 578 expected = DataArray( 579 [10, 20], 580 { 581 "x": [1, 2], 582 "y": 3, 583 "z": 4, 584 "x2": ("x", ["a", "b"]), 585 "y2": "c", 586 "xy": ("x", ["d", "e"]), 587 }, 588 dims="x", 589 ) 590 assert_identical(expected, actual) 591 592 def test_getitem_dataarray(self): 593 # It should not conflict 594 da = DataArray(np.arange(12).reshape((3, 4)), dims=["x", "y"]) 595 ind = DataArray([[0, 1], [0, 1]], dims=["x", "z"]) 596 actual = da[ind] 597 assert_array_equal(actual, da.values[[[0, 1], [0, 1]], :]) 598 599 da = DataArray( 600 np.arange(12).reshape((3, 4)), 601 dims=["x", "y"], 602 coords={"x": [0, 1, 2], "y": ["a", "b", "c", "d"]}, 603 ) 604 ind = xr.DataArray([[0, 1], [0, 1]], dims=["X", "Y"]) 605 actual = da[ind] 606 expected = da.values[[[0, 1], [0, 1]], :] 607 assert_array_equal(actual, expected) 608 assert actual.dims == ("X", "Y", "y") 609 610 # boolean indexing 611 ind = xr.DataArray([True, True, False], dims=["x"]) 612 assert_equal(da[ind], da[[0, 1], :]) 613 assert_equal(da[ind], da[[0, 1]]) 614 assert_equal(da[ind], da[ind.values]) 615 616 def test_getitem_empty_index(self): 617 da = DataArray(np.arange(12).reshape((3, 4)), dims=["x", "y"]) 618 assert_identical(da[{"x": []}], DataArray(np.zeros((0, 4)), dims=["x", "y"])) 619 assert_identical( 620 da.loc[{"y": []}], DataArray(np.zeros((3, 0)), dims=["x", "y"]) 621 ) 622 assert_identical(da[[]], DataArray(np.zeros((0, 4)), dims=["x", "y"])) 623 624 def test_setitem(self): 625 # basic indexing should work as numpy's indexing 626 tuples = [ 627 (0, 0), 628 (0, slice(None, None)), 629 (slice(None, None), slice(None, None)), 630 (slice(None, None), 0), 631 ([1, 0], slice(None, None)), 632 (slice(None, None), [1, 0]), 633 ] 634 for t in tuples: 635 expected = np.arange(6).reshape(3, 2) 636 orig = DataArray( 637 np.arange(6).reshape(3, 2), 638 { 639 "x": [1, 2, 3], 640 "y": ["a", "b"], 641 "z": 4, 642 "x2": ("x", ["a", "b", "c"]), 643 "y2": ("y", ["d", "e"]), 644 }, 645 dims=["x", "y"], 646 ) 647 orig[t] = 1 648 expected[t] = 1 649 assert_array_equal(orig.values, expected) 650 651 def test_setitem_fancy(self): 652 # vectorized indexing 653 da = DataArray(np.ones((3, 2)), dims=["x", "y"]) 654 ind = Variable(["a"], [0, 1]) 655 da[dict(x=ind, y=ind)] = 0 656 expected = DataArray([[0, 1], [1, 0], [1, 1]], dims=["x", "y"]) 657 assert_identical(expected, da) 658 # assign another 0d-variable 659 da[dict(x=ind, y=ind)] = Variable((), 0) 660 expected = DataArray([[0, 1], [1, 0], [1, 1]], dims=["x", "y"]) 661 assert_identical(expected, da) 662 # assign another 1d-variable 663 da[dict(x=ind, y=ind)] = Variable(["a"], [2, 3]) 664 expected = DataArray([[2, 1], [1, 3], [1, 1]], dims=["x", "y"]) 665 assert_identical(expected, da) 666 667 # 2d-vectorized indexing 668 da = DataArray(np.ones((3, 2)), dims=["x", "y"]) 669 ind_x = DataArray([[0, 1]], dims=["a", "b"]) 670 ind_y = DataArray([[1, 0]], dims=["a", "b"]) 671 da[dict(x=ind_x, y=ind_y)] = 0 672 expected = DataArray([[1, 0], [0, 1], [1, 1]], dims=["x", "y"]) 673 assert_identical(expected, da) 674 675 da = DataArray(np.ones((3, 2)), dims=["x", "y"]) 676 ind = Variable(["a"], [0, 1]) 677 da[ind] = 0 678 expected = DataArray([[0, 0], [0, 0], [1, 1]], dims=["x", "y"]) 679 assert_identical(expected, da) 680 681 def test_setitem_dataarray(self): 682 def get_data(): 683 return DataArray( 684 np.ones((4, 3, 2)), 685 dims=["x", "y", "z"], 686 coords={ 687 "x": np.arange(4), 688 "y": ["a", "b", "c"], 689 "non-dim": ("x", [1, 3, 4, 2]), 690 }, 691 ) 692 693 da = get_data() 694 # indexer with inconsistent coordinates. 695 ind = DataArray(np.arange(1, 4), dims=["x"], coords={"x": np.random.randn(3)}) 696 with pytest.raises(IndexError, match=r"dimension coordinate 'x'"): 697 da[dict(x=ind)] = 0 698 699 # indexer with consistent coordinates. 700 ind = DataArray(np.arange(1, 4), dims=["x"], coords={"x": np.arange(1, 4)}) 701 da[dict(x=ind)] = 0 # should not raise 702 assert np.allclose(da[dict(x=ind)].values, 0) 703 assert_identical(da["x"], get_data()["x"]) 704 assert_identical(da["non-dim"], get_data()["non-dim"]) 705 706 da = get_data() 707 # conflict in the assigning values 708 value = xr.DataArray( 709 np.zeros((3, 3, 2)), 710 dims=["x", "y", "z"], 711 coords={"x": [0, 1, 2], "non-dim": ("x", [0, 2, 4])}, 712 ) 713 with pytest.raises(IndexError, match=r"dimension coordinate 'x'"): 714 da[dict(x=ind)] = value 715 716 # consistent coordinate in the assigning values 717 value = xr.DataArray( 718 np.zeros((3, 3, 2)), 719 dims=["x", "y", "z"], 720 coords={"x": [1, 2, 3], "non-dim": ("x", [0, 2, 4])}, 721 ) 722 da[dict(x=ind)] = value 723 assert np.allclose(da[dict(x=ind)].values, 0) 724 assert_identical(da["x"], get_data()["x"]) 725 assert_identical(da["non-dim"], get_data()["non-dim"]) 726 727 # Conflict in the non-dimension coordinate 728 value = xr.DataArray( 729 np.zeros((3, 3, 2)), 730 dims=["x", "y", "z"], 731 coords={"x": [1, 2, 3], "non-dim": ("x", [0, 2, 4])}, 732 ) 733 da[dict(x=ind)] = value # should not raise 734 735 # conflict in the assigning values 736 value = xr.DataArray( 737 np.zeros((3, 3, 2)), 738 dims=["x", "y", "z"], 739 coords={"x": [0, 1, 2], "non-dim": ("x", [0, 2, 4])}, 740 ) 741 with pytest.raises(IndexError, match=r"dimension coordinate 'x'"): 742 da[dict(x=ind)] = value 743 744 # consistent coordinate in the assigning values 745 value = xr.DataArray( 746 np.zeros((3, 3, 2)), 747 dims=["x", "y", "z"], 748 coords={"x": [1, 2, 3], "non-dim": ("x", [0, 2, 4])}, 749 ) 750 da[dict(x=ind)] = value # should not raise 751 752 def test_contains(self): 753 data_array = DataArray([1, 2]) 754 assert 1 in data_array 755 assert 3 not in data_array 756 757 def test_attr_sources_multiindex(self): 758 # make sure attr-style access for multi-index levels 759 # returns DataArray objects 760 assert isinstance(self.mda.level_1, DataArray) 761 762 def test_pickle(self): 763 data = DataArray(np.random.random((3, 3)), dims=("id", "time")) 764 roundtripped = pickle.loads(pickle.dumps(data)) 765 assert_identical(data, roundtripped) 766 767 @requires_dask 768 def test_chunk(self): 769 unblocked = DataArray(np.ones((3, 4))) 770 assert unblocked.chunks is None 771 772 blocked = unblocked.chunk() 773 assert blocked.chunks == ((3,), (4,)) 774 first_dask_name = blocked.data.name 775 776 blocked = unblocked.chunk(chunks=((2, 1), (2, 2))) 777 assert blocked.chunks == ((2, 1), (2, 2)) 778 assert blocked.data.name != first_dask_name 779 780 blocked = unblocked.chunk(chunks=(3, 3)) 781 assert blocked.chunks == ((3,), (3, 1)) 782 assert blocked.data.name != first_dask_name 783 784 # name doesn't change when rechunking by same amount 785 # this fails if ReprObject doesn't have __dask_tokenize__ defined 786 assert unblocked.chunk(2).data.name == unblocked.chunk(2).data.name 787 788 assert blocked.load().chunks is None 789 790 # Check that kwargs are passed 791 import dask.array as da 792 793 blocked = unblocked.chunk(name_prefix="testname_") 794 assert isinstance(blocked.data, da.Array) 795 assert "testname_" in blocked.data.name 796 797 def test_isel(self): 798 assert_identical(self.dv[0], self.dv.isel(x=0)) 799 assert_identical(self.dv, self.dv.isel(x=slice(None))) 800 assert_identical(self.dv[:3], self.dv.isel(x=slice(3))) 801 assert_identical(self.dv[:3, :5], self.dv.isel(x=slice(3), y=slice(5))) 802 with pytest.raises( 803 ValueError, 804 match=r"Dimensions {'not_a_dim'} do not exist. Expected " 805 r"one or more of \('x', 'y'\)", 806 ): 807 self.dv.isel(not_a_dim=0) 808 with pytest.warns( 809 UserWarning, 810 match=r"Dimensions {'not_a_dim'} do not exist. " 811 r"Expected one or more of \('x', 'y'\)", 812 ): 813 self.dv.isel(not_a_dim=0, missing_dims="warn") 814 assert_identical(self.dv, self.dv.isel(not_a_dim=0, missing_dims="ignore")) 815 816 def test_isel_types(self): 817 # regression test for #1405 818 da = DataArray([1, 2, 3], dims="x") 819 # uint64 820 assert_identical( 821 da.isel(x=np.array([0], dtype="uint64")), da.isel(x=np.array([0])) 822 ) 823 # uint32 824 assert_identical( 825 da.isel(x=np.array([0], dtype="uint32")), da.isel(x=np.array([0])) 826 ) 827 # int64 828 assert_identical( 829 da.isel(x=np.array([0], dtype="int64")), da.isel(x=np.array([0])) 830 ) 831 832 @pytest.mark.filterwarnings("ignore::DeprecationWarning") 833 def test_isel_fancy(self): 834 shape = (10, 7, 6) 835 np_array = np.random.random(shape) 836 da = DataArray( 837 np_array, dims=["time", "y", "x"], coords={"time": np.arange(0, 100, 10)} 838 ) 839 y = [1, 3] 840 x = [3, 0] 841 842 expected = da.values[:, y, x] 843 844 actual = da.isel(y=(("test_coord",), y), x=(("test_coord",), x)) 845 assert actual.coords["test_coord"].shape == (len(y),) 846 assert list(actual.coords) == ["time"] 847 assert actual.dims == ("time", "test_coord") 848 849 np.testing.assert_equal(actual, expected) 850 851 # a few corner cases 852 da.isel( 853 time=(("points",), [1, 2]), x=(("points",), [2, 2]), y=(("points",), [3, 4]) 854 ) 855 np.testing.assert_allclose( 856 da.isel( 857 time=(("p",), [1]), x=(("p",), [2]), y=(("p",), [4]) 858 ).values.squeeze(), 859 np_array[1, 4, 2].squeeze(), 860 ) 861 da.isel(time=(("points",), [1, 2])) 862 y = [-1, 0] 863 x = [-2, 2] 864 expected = da.values[:, y, x] 865 actual = da.isel(x=(("points",), x), y=(("points",), y)).values 866 np.testing.assert_equal(actual, expected) 867 868 # test that the order of the indexers doesn't matter 869 assert_identical( 870 da.isel(y=(("points",), y), x=(("points",), x)), 871 da.isel(x=(("points",), x), y=(("points",), y)), 872 ) 873 874 # make sure we're raising errors in the right places 875 with pytest.raises(IndexError, match=r"Dimensions of indexers mismatch"): 876 da.isel(y=(("points",), [1, 2]), x=(("points",), [1, 2, 3])) 877 878 # tests using index or DataArray as indexers 879 stations = Dataset() 880 stations["station"] = (("station",), ["A", "B", "C"]) 881 stations["dim1s"] = (("station",), [1, 2, 3]) 882 stations["dim2s"] = (("station",), [4, 5, 1]) 883 884 actual = da.isel(x=stations["dim1s"], y=stations["dim2s"]) 885 assert "station" in actual.coords 886 assert "station" in actual.dims 887 assert_identical(actual["station"], stations["station"]) 888 889 with pytest.raises(ValueError, match=r"conflicting values for "): 890 da.isel( 891 x=DataArray([0, 1, 2], dims="station", coords={"station": [0, 1, 2]}), 892 y=DataArray([0, 1, 2], dims="station", coords={"station": [0, 1, 3]}), 893 ) 894 895 # multi-dimensional selection 896 stations = Dataset() 897 stations["a"] = (("a",), ["A", "B", "C"]) 898 stations["b"] = (("b",), [0, 1]) 899 stations["dim1s"] = (("a", "b"), [[1, 2], [2, 3], [3, 4]]) 900 stations["dim2s"] = (("a",), [4, 5, 1]) 901 902 actual = da.isel(x=stations["dim1s"], y=stations["dim2s"]) 903 assert "a" in actual.coords 904 assert "a" in actual.dims 905 assert "b" in actual.coords 906 assert "b" in actual.dims 907 assert_identical(actual["a"], stations["a"]) 908 assert_identical(actual["b"], stations["b"]) 909 expected = da.variable[ 910 :, stations["dim2s"].variable, stations["dim1s"].variable 911 ] 912 assert_array_equal(actual, expected) 913 914 def test_sel(self): 915 self.ds["x"] = ("x", np.array(list("abcdefghij"))) 916 da = self.ds["foo"] 917 assert_identical(da, da.sel(x=slice(None))) 918 assert_identical(da[1], da.sel(x="b")) 919 assert_identical(da[:3], da.sel(x=slice("c"))) 920 assert_identical(da[:3], da.sel(x=["a", "b", "c"])) 921 assert_identical(da[:, :4], da.sel(y=(self.ds["y"] < 4))) 922 # verify that indexing with a dataarray works 923 b = DataArray("b") 924 assert_identical(da[1], da.sel(x=b)) 925 assert_identical(da[[1]], da.sel(x=slice(b, b))) 926 927 def test_sel_dataarray(self): 928 # indexing with DataArray 929 self.ds["x"] = ("x", np.array(list("abcdefghij"))) 930 da = self.ds["foo"] 931 932 ind = DataArray(["a", "b", "c"], dims=["x"]) 933 actual = da.sel(x=ind) 934 assert_identical(actual, da.isel(x=[0, 1, 2])) 935 936 # along new dimension 937 ind = DataArray(["a", "b", "c"], dims=["new_dim"]) 938 actual = da.sel(x=ind) 939 assert_array_equal(actual, da.isel(x=[0, 1, 2])) 940 assert "new_dim" in actual.dims 941 942 # with coordinate 943 ind = DataArray( 944 ["a", "b", "c"], dims=["new_dim"], coords={"new_dim": [0, 1, 2]} 945 ) 946 actual = da.sel(x=ind) 947 assert_array_equal(actual, da.isel(x=[0, 1, 2])) 948 assert "new_dim" in actual.dims 949 assert "new_dim" in actual.coords 950 assert_equal(actual["new_dim"].drop_vars("x"), ind["new_dim"]) 951 952 def test_sel_invalid_slice(self): 953 array = DataArray(np.arange(10), [("x", np.arange(10))]) 954 with pytest.raises(ValueError, match=r"cannot use non-scalar arrays"): 955 array.sel(x=slice(array.x)) 956 957 def test_sel_dataarray_datetime_slice(self): 958 # regression test for GH1240 959 times = pd.date_range("2000-01-01", freq="D", periods=365) 960 array = DataArray(np.arange(365), [("time", times)]) 961 result = array.sel(time=slice(array.time[0], array.time[-1])) 962 assert_equal(result, array) 963 964 array = DataArray(np.arange(365), [("delta", times - times[0])]) 965 result = array.sel(delta=slice(array.delta[0], array.delta[-1])) 966 assert_equal(result, array) 967 968 def test_sel_float(self): 969 data_values = np.arange(4) 970 971 # case coords are float32 and label is list of floats 972 float_values = [0.0, 0.111, 0.222, 0.333] 973 coord_values = np.asarray(float_values, dtype="float32") 974 array = DataArray(data_values, [("float32_coord", coord_values)]) 975 expected = DataArray(data_values[1:3], [("float32_coord", coord_values[1:3])]) 976 actual = array.sel(float32_coord=float_values[1:3]) 977 # case coords are float16 and label is list of floats 978 coord_values_16 = np.asarray(float_values, dtype="float16") 979 expected_16 = DataArray( 980 data_values[1:3], [("float16_coord", coord_values_16[1:3])] 981 ) 982 array_16 = DataArray(data_values, [("float16_coord", coord_values_16)]) 983 actual_16 = array_16.sel(float16_coord=float_values[1:3]) 984 985 # case coord, label are scalars 986 expected_scalar = DataArray( 987 data_values[2], coords={"float32_coord": coord_values[2]} 988 ) 989 actual_scalar = array.sel(float32_coord=float_values[2]) 990 991 assert_equal(expected, actual) 992 assert_equal(expected_scalar, actual_scalar) 993 assert_equal(expected_16, actual_16) 994 995 def test_sel_no_index(self): 996 array = DataArray(np.arange(10), dims="x") 997 assert_identical(array[0], array.sel(x=0)) 998 assert_identical(array[:5], array.sel(x=slice(5))) 999 assert_identical(array[[0, -1]], array.sel(x=[0, -1])) 1000 assert_identical(array[array < 5], array.sel(x=(array < 5))) 1001 1002 def test_sel_method(self): 1003 data = DataArray(np.random.randn(3, 4), [("x", [0, 1, 2]), ("y", list("abcd"))]) 1004 1005 expected = data.sel(y=["a", "b"]) 1006 actual = data.sel(y=["ab", "ba"], method="pad") 1007 assert_identical(expected, actual) 1008 1009 expected = data.sel(x=[1, 2]) 1010 actual = data.sel(x=[0.9, 1.9], method="backfill", tolerance=1) 1011 assert_identical(expected, actual) 1012 1013 def test_sel_drop(self): 1014 data = DataArray([1, 2, 3], [("x", [0, 1, 2])]) 1015 expected = DataArray(1) 1016 selected = data.sel(x=0, drop=True) 1017 assert_identical(expected, selected) 1018 1019 expected = DataArray(1, {"x": 0}) 1020 selected = data.sel(x=0, drop=False) 1021 assert_identical(expected, selected) 1022 1023 data = DataArray([1, 2, 3], dims=["x"]) 1024 expected = DataArray(1) 1025 selected = data.sel(x=0, drop=True) 1026 assert_identical(expected, selected) 1027 1028 def test_isel_drop(self): 1029 data = DataArray([1, 2, 3], [("x", [0, 1, 2])]) 1030 expected = DataArray(1) 1031 selected = data.isel(x=0, drop=True) 1032 assert_identical(expected, selected) 1033 1034 expected = DataArray(1, {"x": 0}) 1035 selected = data.isel(x=0, drop=False) 1036 assert_identical(expected, selected) 1037 1038 def test_head(self): 1039 assert_equal(self.dv.isel(x=slice(5)), self.dv.head(x=5)) 1040 assert_equal(self.dv.isel(x=slice(0)), self.dv.head(x=0)) 1041 assert_equal( 1042 self.dv.isel({dim: slice(6) for dim in self.dv.dims}), self.dv.head(6) 1043 ) 1044 assert_equal( 1045 self.dv.isel({dim: slice(5) for dim in self.dv.dims}), self.dv.head() 1046 ) 1047 with pytest.raises(TypeError, match=r"either dict-like or a single int"): 1048 self.dv.head([3]) 1049 with pytest.raises(TypeError, match=r"expected integer type"): 1050 self.dv.head(x=3.1) 1051 with pytest.raises(ValueError, match=r"expected positive int"): 1052 self.dv.head(-3) 1053 1054 def test_tail(self): 1055 assert_equal(self.dv.isel(x=slice(-5, None)), self.dv.tail(x=5)) 1056 assert_equal(self.dv.isel(x=slice(0)), self.dv.tail(x=0)) 1057 assert_equal( 1058 self.dv.isel({dim: slice(-6, None) for dim in self.dv.dims}), 1059 self.dv.tail(6), 1060 ) 1061 assert_equal( 1062 self.dv.isel({dim: slice(-5, None) for dim in self.dv.dims}), self.dv.tail() 1063 ) 1064 with pytest.raises(TypeError, match=r"either dict-like or a single int"): 1065 self.dv.tail([3]) 1066 with pytest.raises(TypeError, match=r"expected integer type"): 1067 self.dv.tail(x=3.1) 1068 with pytest.raises(ValueError, match=r"expected positive int"): 1069 self.dv.tail(-3) 1070 1071 def test_thin(self): 1072 assert_equal(self.dv.isel(x=slice(None, None, 5)), self.dv.thin(x=5)) 1073 assert_equal( 1074 self.dv.isel({dim: slice(None, None, 6) for dim in self.dv.dims}), 1075 self.dv.thin(6), 1076 ) 1077 with pytest.raises(TypeError, match=r"either dict-like or a single int"): 1078 self.dv.thin([3]) 1079 with pytest.raises(TypeError, match=r"expected integer type"): 1080 self.dv.thin(x=3.1) 1081 with pytest.raises(ValueError, match=r"expected positive int"): 1082 self.dv.thin(-3) 1083 with pytest.raises(ValueError, match=r"cannot be zero"): 1084 self.dv.thin(time=0) 1085 1086 def test_loc(self): 1087 self.ds["x"] = ("x", np.array(list("abcdefghij"))) 1088 da = self.ds["foo"] 1089 assert_identical(da[:3], da.loc[:"c"]) 1090 assert_identical(da[1], da.loc["b"]) 1091 assert_identical(da[1], da.loc[{"x": "b"}]) 1092 assert_identical(da[1], da.loc["b", ...]) 1093 assert_identical(da[:3], da.loc[["a", "b", "c"]]) 1094 assert_identical(da[:3, :4], da.loc[["a", "b", "c"], np.arange(4)]) 1095 assert_identical(da[:, :4], da.loc[:, self.ds["y"] < 4]) 1096 1097 def test_loc_datetime64_value(self): 1098 # regression test for https://github.com/pydata/xarray/issues/4283 1099 t = np.array(["2017-09-05T12", "2017-09-05T15"], dtype="datetime64[ns]") 1100 array = DataArray(np.ones(t.shape), dims=("time",), coords=(t,)) 1101 assert_identical(array.loc[{"time": t[0]}], array[0]) 1102 1103 def test_loc_assign(self): 1104 self.ds["x"] = ("x", np.array(list("abcdefghij"))) 1105 da = self.ds["foo"] 1106 # assignment 1107 da.loc["a":"j"] = 0 1108 assert np.all(da.values == 0) 1109 da.loc[{"x": slice("a", "j")}] = 2 1110 assert np.all(da.values == 2) 1111 1112 da.loc[{"x": slice("a", "j")}] = 2 1113 assert np.all(da.values == 2) 1114 1115 # Multi dimensional case 1116 da = DataArray(np.arange(12).reshape(3, 4), dims=["x", "y"]) 1117 da.loc[0, 0] = 0 1118 assert da.values[0, 0] == 0 1119 assert da.values[0, 1] != 0 1120 1121 da = DataArray(np.arange(12).reshape(3, 4), dims=["x", "y"]) 1122 da.loc[0] = 0 1123 assert np.all(da.values[0] == np.zeros(4)) 1124 assert da.values[1, 0] != 0 1125 1126 def test_loc_assign_dataarray(self): 1127 def get_data(): 1128 return DataArray( 1129 np.ones((4, 3, 2)), 1130 dims=["x", "y", "z"], 1131 coords={ 1132 "x": np.arange(4), 1133 "y": ["a", "b", "c"], 1134 "non-dim": ("x", [1, 3, 4, 2]), 1135 }, 1136 ) 1137 1138 da = get_data() 1139 # indexer with inconsistent coordinates. 1140 ind = DataArray(np.arange(1, 4), dims=["y"], coords={"y": np.random.randn(3)}) 1141 with pytest.raises(IndexError, match=r"dimension coordinate 'y'"): 1142 da.loc[dict(x=ind)] = 0 1143 1144 # indexer with consistent coordinates. 1145 ind = DataArray(np.arange(1, 4), dims=["x"], coords={"x": np.arange(1, 4)}) 1146 da.loc[dict(x=ind)] = 0 # should not raise 1147 assert np.allclose(da[dict(x=ind)].values, 0) 1148 assert_identical(da["x"], get_data()["x"]) 1149 assert_identical(da["non-dim"], get_data()["non-dim"]) 1150 1151 da = get_data() 1152 # conflict in the assigning values 1153 value = xr.DataArray( 1154 np.zeros((3, 3, 2)), 1155 dims=["x", "y", "z"], 1156 coords={"x": [0, 1, 2], "non-dim": ("x", [0, 2, 4])}, 1157 ) 1158 with pytest.raises(IndexError, match=r"dimension coordinate 'x'"): 1159 da.loc[dict(x=ind)] = value 1160 1161 # consistent coordinate in the assigning values 1162 value = xr.DataArray( 1163 np.zeros((3, 3, 2)), 1164 dims=["x", "y", "z"], 1165 coords={"x": [1, 2, 3], "non-dim": ("x", [0, 2, 4])}, 1166 ) 1167 da.loc[dict(x=ind)] = value 1168 assert np.allclose(da[dict(x=ind)].values, 0) 1169 assert_identical(da["x"], get_data()["x"]) 1170 assert_identical(da["non-dim"], get_data()["non-dim"]) 1171 1172 def test_loc_single_boolean(self): 1173 data = DataArray([0, 1], coords=[[True, False]]) 1174 assert data.loc[True] == 0 1175 assert data.loc[False] == 1 1176 1177 def test_loc_dim_name_collision_with_sel_params(self): 1178 da = xr.DataArray( 1179 [[0, 0], [1, 1]], 1180 dims=["dim1", "method"], 1181 coords={"dim1": ["x", "y"], "method": ["a", "b"]}, 1182 ) 1183 np.testing.assert_array_equal( 1184 da.loc[dict(dim1=["x", "y"], method=["a"])], [[0], [1]] 1185 ) 1186 1187 def test_selection_multiindex(self): 1188 mindex = pd.MultiIndex.from_product( 1189 [["a", "b"], [1, 2], [-1, -2]], names=("one", "two", "three") 1190 ) 1191 mdata = DataArray(range(8), [("x", mindex)]) 1192 1193 def test_sel(lab_indexer, pos_indexer, replaced_idx=False, renamed_dim=None): 1194 da = mdata.sel(x=lab_indexer) 1195 expected_da = mdata.isel(x=pos_indexer) 1196 if not replaced_idx: 1197 assert_identical(da, expected_da) 1198 else: 1199 if renamed_dim: 1200 assert da.dims[0] == renamed_dim 1201 da = da.rename({renamed_dim: "x"}) 1202 assert_identical(da.variable, expected_da.variable) 1203 assert not da["x"].equals(expected_da["x"]) 1204 1205 test_sel(("a", 1, -1), 0) 1206 test_sel(("b", 2, -2), -1) 1207 test_sel(("a", 1), [0, 1], replaced_idx=True, renamed_dim="three") 1208 test_sel(("a",), range(4), replaced_idx=True) 1209 test_sel("a", range(4), replaced_idx=True) 1210 test_sel([("a", 1, -1), ("b", 2, -2)], [0, 7]) 1211 test_sel(slice("a", "b"), range(8)) 1212 test_sel(slice(("a", 1), ("b", 1)), range(6)) 1213 test_sel({"one": "a", "two": 1, "three": -1}, 0) 1214 test_sel({"one": "a", "two": 1}, [0, 1], replaced_idx=True, renamed_dim="three") 1215 test_sel({"one": "a"}, range(4), replaced_idx=True) 1216 1217 assert_identical(mdata.loc["a"], mdata.sel(x="a")) 1218 assert_identical(mdata.loc[("a", 1), ...], mdata.sel(x=("a", 1))) 1219 assert_identical(mdata.loc[{"one": "a"}, ...], mdata.sel(x={"one": "a"})) 1220 with pytest.raises(IndexError): 1221 mdata.loc[("a", 1)] 1222 1223 assert_identical(mdata.sel(x={"one": "a", "two": 1}), mdata.sel(one="a", two=1)) 1224 1225 def test_selection_multiindex_remove_unused(self): 1226 # GH2619. For MultiIndex, we need to call remove_unused. 1227 ds = xr.DataArray( 1228 np.arange(40).reshape(8, 5), 1229 dims=["x", "y"], 1230 coords={"x": np.arange(8), "y": np.arange(5)}, 1231 ) 1232 ds = ds.stack(xy=["x", "y"]) 1233 ds_isel = ds.isel(xy=ds["x"] < 4) 1234 with pytest.raises(KeyError): 1235 ds_isel.sel(x=5) 1236 1237 actual = ds_isel.unstack() 1238 expected = ds.reset_index("xy").isel(xy=ds["x"] < 4) 1239 expected = expected.set_index(xy=["x", "y"]).unstack() 1240 assert_identical(expected, actual) 1241 1242 def test_selection_multiindex_from_level(self): 1243 # GH: 3512 1244 da = DataArray([0, 1], dims=["x"], coords={"x": [0, 1], "y": "a"}) 1245 db = DataArray([2, 3], dims=["x"], coords={"x": [0, 1], "y": "b"}) 1246 data = xr.concat([da, db], dim="x").set_index(xy=["x", "y"]) 1247 assert data.dims == ("xy",) 1248 actual = data.sel(y="a") 1249 expected = data.isel(xy=[0, 1]).unstack("xy").squeeze("y").drop_vars("y") 1250 assert_equal(actual, expected) 1251 1252 def test_virtual_default_coords(self): 1253 array = DataArray(np.zeros((5,)), dims="x") 1254 expected = DataArray(range(5), dims="x", name="x") 1255 assert_identical(expected, array["x"]) 1256 assert_identical(expected, array.coords["x"]) 1257 1258 def test_virtual_time_components(self): 1259 dates = pd.date_range("2000-01-01", periods=10) 1260 da = DataArray(np.arange(1, 11), [("time", dates)]) 1261 1262 assert_array_equal(da["time.dayofyear"], da.values) 1263 assert_array_equal(da.coords["time.dayofyear"], da.values) 1264 1265 def test_coords(self): 1266 # use int64 to ensure repr() consistency on windows 1267 coords = [ 1268 IndexVariable("x", np.array([-1, -2], "int64")), 1269 IndexVariable("y", np.array([0, 1, 2], "int64")), 1270 ] 1271 da = DataArray(np.random.randn(2, 3), coords, name="foo") 1272 1273 assert 2 == len(da.coords) 1274 1275 assert ["x", "y"] == list(da.coords) 1276 1277 assert coords[0].identical(da.coords["x"]) 1278 assert coords[1].identical(da.coords["y"]) 1279 1280 assert "x" in da.coords 1281 assert 0 not in da.coords 1282 assert "foo" not in da.coords 1283 1284 with pytest.raises(KeyError): 1285 da.coords[0] 1286 with pytest.raises(KeyError): 1287 da.coords["foo"] 1288 1289 expected = dedent( 1290 """\ 1291 Coordinates: 1292 * x (x) int64 -1 -2 1293 * y (y) int64 0 1 2""" 1294 ) 1295 actual = repr(da.coords) 1296 assert expected == actual 1297 1298 del da.coords["x"] 1299 da._indexes = propagate_indexes(da._indexes, exclude="x") 1300 expected = DataArray(da.values, {"y": [0, 1, 2]}, dims=["x", "y"], name="foo") 1301 assert_identical(da, expected) 1302 1303 with pytest.raises(ValueError, match=r"conflicting MultiIndex"): 1304 self.mda["level_1"] = np.arange(4) 1305 self.mda.coords["level_1"] = np.arange(4) 1306 1307 def test_coords_to_index(self): 1308 da = DataArray(np.zeros((2, 3)), [("x", [1, 2]), ("y", list("abc"))]) 1309 1310 with pytest.raises(ValueError, match=r"no valid index"): 1311 da[0, 0].coords.to_index() 1312 1313 expected = pd.Index(["a", "b", "c"], name="y") 1314 actual = da[0].coords.to_index() 1315 assert expected.equals(actual) 1316 1317 expected = pd.MultiIndex.from_product( 1318 [[1, 2], ["a", "b", "c"]], names=["x", "y"] 1319 ) 1320 actual = da.coords.to_index() 1321 assert expected.equals(actual) 1322 1323 expected = pd.MultiIndex.from_product( 1324 [["a", "b", "c"], [1, 2]], names=["y", "x"] 1325 ) 1326 actual = da.coords.to_index(["y", "x"]) 1327 assert expected.equals(actual) 1328 1329 with pytest.raises(ValueError, match=r"ordered_dims must match"): 1330 da.coords.to_index(["x"]) 1331 1332 def test_coord_coords(self): 1333 orig = DataArray( 1334 [10, 20], {"x": [1, 2], "x2": ("x", ["a", "b"]), "z": 4}, dims="x" 1335 ) 1336 1337 actual = orig.coords["x"] 1338 expected = DataArray( 1339 [1, 2], {"z": 4, "x2": ("x", ["a", "b"]), "x": [1, 2]}, dims="x", name="x" 1340 ) 1341 assert_identical(expected, actual) 1342 1343 del actual.coords["x2"] 1344 assert_identical(expected.reset_coords("x2", drop=True), actual) 1345 1346 actual.coords["x3"] = ("x", ["a", "b"]) 1347 expected = DataArray( 1348 [1, 2], {"z": 4, "x3": ("x", ["a", "b"]), "x": [1, 2]}, dims="x", name="x" 1349 ) 1350 assert_identical(expected, actual) 1351 1352 def test_reset_coords(self): 1353 data = DataArray( 1354 np.zeros((3, 4)), 1355 {"bar": ("x", ["a", "b", "c"]), "baz": ("y", range(4)), "y": range(4)}, 1356 dims=["x", "y"], 1357 name="foo", 1358 ) 1359 1360 actual = data.reset_coords() 1361 expected = Dataset( 1362 { 1363 "foo": (["x", "y"], np.zeros((3, 4))), 1364 "bar": ("x", ["a", "b", "c"]), 1365 "baz": ("y", range(4)), 1366 "y": range(4), 1367 } 1368 ) 1369 assert_identical(actual, expected) 1370 1371 actual = data.reset_coords(["bar", "baz"]) 1372 assert_identical(actual, expected) 1373 1374 actual = data.reset_coords("bar") 1375 expected = Dataset( 1376 {"foo": (["x", "y"], np.zeros((3, 4))), "bar": ("x", ["a", "b", "c"])}, 1377 {"baz": ("y", range(4)), "y": range(4)}, 1378 ) 1379 assert_identical(actual, expected) 1380 1381 actual = data.reset_coords(["bar"]) 1382 assert_identical(actual, expected) 1383 1384 actual = data.reset_coords(drop=True) 1385 expected = DataArray( 1386 np.zeros((3, 4)), coords={"y": range(4)}, dims=["x", "y"], name="foo" 1387 ) 1388 assert_identical(actual, expected) 1389 1390 actual = data.copy() 1391 actual = actual.reset_coords(drop=True) 1392 assert_identical(actual, expected) 1393 1394 actual = data.reset_coords("bar", drop=True) 1395 expected = DataArray( 1396 np.zeros((3, 4)), 1397 {"baz": ("y", range(4)), "y": range(4)}, 1398 dims=["x", "y"], 1399 name="foo", 1400 ) 1401 assert_identical(actual, expected) 1402 1403 with pytest.raises(ValueError, match=r"cannot be found"): 1404 data.reset_coords("foo", drop=True) 1405 with pytest.raises(ValueError, match=r"cannot be found"): 1406 data.reset_coords("not_found") 1407 with pytest.raises(ValueError, match=r"cannot remove index"): 1408 data.reset_coords("y") 1409 1410 def test_assign_coords(self): 1411 array = DataArray(10) 1412 actual = array.assign_coords(c=42) 1413 expected = DataArray(10, {"c": 42}) 1414 assert_identical(actual, expected) 1415 1416 with pytest.raises(ValueError, match=r"conflicting MultiIndex"): 1417 self.mda.assign_coords(level_1=range(4)) 1418 1419 # GH: 2112 1420 da = xr.DataArray([0, 1, 2], dims="x") 1421 with pytest.raises(ValueError): 1422 da["x"] = [0, 1, 2, 3] # size conflict 1423 with pytest.raises(ValueError): 1424 da.coords["x"] = [0, 1, 2, 3] # size conflict 1425 1426 def test_coords_alignment(self): 1427 lhs = DataArray([1, 2, 3], [("x", [0, 1, 2])]) 1428 rhs = DataArray([2, 3, 4], [("x", [1, 2, 3])]) 1429 lhs.coords["rhs"] = rhs 1430 1431 expected = DataArray( 1432 [1, 2, 3], coords={"rhs": ("x", [np.nan, 2, 3]), "x": [0, 1, 2]}, dims="x" 1433 ) 1434 assert_identical(lhs, expected) 1435 1436 def test_set_coords_update_index(self): 1437 actual = DataArray([1, 2, 3], [("x", [1, 2, 3])]) 1438 actual.coords["x"] = ["a", "b", "c"] 1439 assert actual.xindexes["x"].to_pandas_index().equals(pd.Index(["a", "b", "c"])) 1440 1441 def test_coords_replacement_alignment(self): 1442 # regression test for GH725 1443 arr = DataArray([0, 1, 2], dims=["abc"]) 1444 new_coord = DataArray([1, 2, 3], dims=["abc"], coords=[[1, 2, 3]]) 1445 arr["abc"] = new_coord 1446 expected = DataArray([0, 1, 2], coords=[("abc", [1, 2, 3])]) 1447 assert_identical(arr, expected) 1448 1449 def test_coords_non_string(self): 1450 arr = DataArray(0, coords={1: 2}) 1451 actual = arr.coords[1] 1452 expected = DataArray(2, coords={1: 2}, name=1) 1453 assert_identical(actual, expected) 1454 1455 def test_coords_delitem_delete_indexes(self): 1456 # regression test for GH3746 1457 arr = DataArray(np.ones((2,)), dims="x", coords={"x": [0, 1]}) 1458 del arr.coords["x"] 1459 assert "x" not in arr.xindexes 1460 1461 def test_broadcast_like(self): 1462 arr1 = DataArray( 1463 np.ones((2, 3)), 1464 dims=["x", "y"], 1465 coords={"x": ["a", "b"], "y": ["a", "b", "c"]}, 1466 ) 1467 arr2 = DataArray( 1468 np.ones((3, 2)), 1469 dims=["x", "y"], 1470 coords={"x": ["a", "b", "c"], "y": ["a", "b"]}, 1471 ) 1472 orig1, orig2 = broadcast(arr1, arr2) 1473 new1 = arr1.broadcast_like(arr2) 1474 new2 = arr2.broadcast_like(arr1) 1475 1476 assert_identical(orig1, new1) 1477 assert_identical(orig2, new2) 1478 1479 orig3 = DataArray(np.random.randn(5), [("x", range(5))]) 1480 orig4 = DataArray(np.random.randn(6), [("y", range(6))]) 1481 new3, new4 = broadcast(orig3, orig4) 1482 1483 assert_identical(orig3.broadcast_like(orig4), new3.transpose("y", "x")) 1484 assert_identical(orig4.broadcast_like(orig3), new4) 1485 1486 def test_reindex_like(self): 1487 foo = DataArray(np.random.randn(5, 6), [("x", range(5)), ("y", range(6))]) 1488 bar = foo[:2, :2] 1489 assert_identical(foo.reindex_like(bar), bar) 1490 1491 expected = foo.copy() 1492 expected[:] = np.nan 1493 expected[:2, :2] = bar 1494 assert_identical(bar.reindex_like(foo), expected) 1495 1496 def test_reindex_like_no_index(self): 1497 foo = DataArray(np.random.randn(5, 6), dims=["x", "y"]) 1498 assert_identical(foo, foo.reindex_like(foo)) 1499 1500 bar = foo[:4] 1501 with pytest.raises(ValueError, match=r"different size for unlabeled"): 1502 foo.reindex_like(bar) 1503 1504 def test_reindex_regressions(self): 1505 da = DataArray(np.random.randn(5), coords=[("time", range(5))]) 1506 time2 = DataArray(np.arange(5), dims="time2") 1507 with pytest.raises(ValueError): 1508 da.reindex(time=time2) 1509 1510 # regression test for #736, reindex can not change complex nums dtype 1511 x = np.array([1, 2, 3], dtype=complex) 1512 x = DataArray(x, coords=[[0.1, 0.2, 0.3]]) 1513 y = DataArray([2, 5, 6, 7, 8], coords=[[-1.1, 0.21, 0.31, 0.41, 0.51]]) 1514 re_dtype = x.reindex_like(y, method="pad").dtype 1515 assert x.dtype == re_dtype 1516 1517 def test_reindex_method(self): 1518 x = DataArray([10, 20], dims="y", coords={"y": [0, 1]}) 1519 y = [-0.1, 0.5, 1.1] 1520 actual = x.reindex(y=y, method="backfill", tolerance=0.2) 1521 expected = DataArray([10, np.nan, np.nan], coords=[("y", y)]) 1522 assert_identical(expected, actual) 1523 1524 alt = Dataset({"y": y}) 1525 actual = x.reindex_like(alt, method="backfill") 1526 expected = DataArray([10, 20, np.nan], coords=[("y", y)]) 1527 assert_identical(expected, actual) 1528 1529 @pytest.mark.parametrize("fill_value", [dtypes.NA, 2, 2.0, {None: 2, "u": 1}]) 1530 def test_reindex_fill_value(self, fill_value): 1531 x = DataArray([10, 20], dims="y", coords={"y": [0, 1], "u": ("y", [1, 2])}) 1532 y = [0, 1, 2] 1533 if fill_value == dtypes.NA: 1534 # if we supply the default, we expect the missing value for a 1535 # float array 1536 fill_value_var = fill_value_u = np.nan 1537 elif isinstance(fill_value, dict): 1538 fill_value_var = fill_value[None] 1539 fill_value_u = fill_value["u"] 1540 else: 1541 fill_value_var = fill_value_u = fill_value 1542 actual = x.reindex(y=y, fill_value=fill_value) 1543 expected = DataArray( 1544 [10, 20, fill_value_var], 1545 dims="y", 1546 coords={"y": y, "u": ("y", [1, 2, fill_value_u])}, 1547 ) 1548 assert_identical(expected, actual) 1549 1550 @pytest.mark.parametrize("dtype", [str, bytes]) 1551 def test_reindex_str_dtype(self, dtype): 1552 1553 data = DataArray( 1554 [1, 2], dims="x", coords={"x": np.array(["a", "b"], dtype=dtype)} 1555 ) 1556 1557 actual = data.reindex(x=data.x) 1558 expected = data 1559 1560 assert_identical(expected, actual) 1561 assert actual.dtype == expected.dtype 1562 1563 def test_rename(self): 1564 renamed = self.dv.rename("bar") 1565 assert_identical(renamed.to_dataset(), self.ds.rename({"foo": "bar"})) 1566 assert renamed.name == "bar" 1567 1568 renamed = self.dv.x.rename({"x": "z"}).rename("z") 1569 assert_identical(renamed, self.ds.rename({"x": "z"}).z) 1570 assert renamed.name == "z" 1571 assert renamed.dims == ("z",) 1572 1573 renamed_kwargs = self.dv.x.rename(x="z").rename("z") 1574 assert_identical(renamed, renamed_kwargs) 1575 1576 def test_init_value(self): 1577 expected = DataArray( 1578 np.full((3, 4), 3), dims=["x", "y"], coords=[range(3), range(4)] 1579 ) 1580 actual = DataArray(3, dims=["x", "y"], coords=[range(3), range(4)]) 1581 assert_identical(expected, actual) 1582 1583 expected = DataArray( 1584 np.full((1, 10, 2), 0), 1585 dims=["w", "x", "y"], 1586 coords={"x": np.arange(10), "y": ["north", "south"]}, 1587 ) 1588 actual = DataArray(0, dims=expected.dims, coords=expected.coords) 1589 assert_identical(expected, actual) 1590 1591 expected = DataArray( 1592 np.full((10, 2), np.nan), coords=[("x", np.arange(10)), ("y", ["a", "b"])] 1593 ) 1594 actual = DataArray(coords=[("x", np.arange(10)), ("y", ["a", "b"])]) 1595 assert_identical(expected, actual) 1596 1597 with pytest.raises(ValueError, match=r"different number of dim"): 1598 DataArray(np.array(1), coords={"x": np.arange(10)}, dims=["x"]) 1599 with pytest.raises(ValueError, match=r"does not match the 0 dim"): 1600 DataArray(np.array(1), coords=[("x", np.arange(10))]) 1601 1602 def test_swap_dims(self): 1603 array = DataArray(np.random.randn(3), {"x": list("abc")}, "x") 1604 expected = DataArray(array.values, {"x": ("y", list("abc"))}, dims="y") 1605 actual = array.swap_dims({"x": "y"}) 1606 assert_identical(expected, actual) 1607 for dim_name in set().union(expected.xindexes.keys(), actual.xindexes.keys()): 1608 pd.testing.assert_index_equal( 1609 expected.xindexes[dim_name].to_pandas_index(), 1610 actual.xindexes[dim_name].to_pandas_index(), 1611 ) 1612 1613 # as kwargs 1614 array = DataArray(np.random.randn(3), {"x": list("abc")}, "x") 1615 expected = DataArray(array.values, {"x": ("y", list("abc"))}, dims="y") 1616 actual = array.swap_dims(x="y") 1617 assert_identical(expected, actual) 1618 for dim_name in set().union(expected.xindexes.keys(), actual.xindexes.keys()): 1619 pd.testing.assert_index_equal( 1620 expected.xindexes[dim_name].to_pandas_index(), 1621 actual.xindexes[dim_name].to_pandas_index(), 1622 ) 1623 1624 # multiindex case 1625 idx = pd.MultiIndex.from_arrays([list("aab"), list("yzz")], names=["y1", "y2"]) 1626 array = DataArray(np.random.randn(3), {"y": ("x", idx)}, "x") 1627 expected = DataArray(array.values, {"y": idx}, "y") 1628 actual = array.swap_dims({"x": "y"}) 1629 assert_identical(expected, actual) 1630 for dim_name in set().union(expected.xindexes.keys(), actual.xindexes.keys()): 1631 pd.testing.assert_index_equal( 1632 expected.xindexes[dim_name].to_pandas_index(), 1633 actual.xindexes[dim_name].to_pandas_index(), 1634 ) 1635 1636 def test_expand_dims_error(self): 1637 array = DataArray( 1638 np.random.randn(3, 4), 1639 dims=["x", "dim_0"], 1640 coords={"x": np.linspace(0.0, 1.0, 3)}, 1641 attrs={"key": "entry"}, 1642 ) 1643 1644 with pytest.raises(TypeError, match=r"dim should be hashable or"): 1645 array.expand_dims(0) 1646 with pytest.raises(ValueError, match=r"lengths of dim and axis"): 1647 # dims and axis argument should be the same length 1648 array.expand_dims(dim=["a", "b"], axis=[1, 2, 3]) 1649 with pytest.raises(ValueError, match=r"Dimension x already"): 1650 # Should not pass the already existing dimension. 1651 array.expand_dims(dim=["x"]) 1652 # raise if duplicate 1653 with pytest.raises(ValueError, match=r"duplicate values"): 1654 array.expand_dims(dim=["y", "y"]) 1655 with pytest.raises(ValueError, match=r"duplicate values"): 1656 array.expand_dims(dim=["y", "z"], axis=[1, 1]) 1657 with pytest.raises(ValueError, match=r"duplicate values"): 1658 array.expand_dims(dim=["y", "z"], axis=[2, -2]) 1659 1660 # out of bounds error, axis must be in [-4, 3] 1661 with pytest.raises(IndexError): 1662 array.expand_dims(dim=["y", "z"], axis=[2, 4]) 1663 with pytest.raises(IndexError): 1664 array.expand_dims(dim=["y", "z"], axis=[2, -5]) 1665 # Does not raise an IndexError 1666 array.expand_dims(dim=["y", "z"], axis=[2, -4]) 1667 array.expand_dims(dim=["y", "z"], axis=[2, 3]) 1668 1669 array = DataArray( 1670 np.random.randn(3, 4), 1671 dims=["x", "dim_0"], 1672 coords={"x": np.linspace(0.0, 1.0, 3)}, 1673 attrs={"key": "entry"}, 1674 ) 1675 with pytest.raises(TypeError): 1676 array.expand_dims({"new_dim": 3.2}) 1677 1678 # Attempt to use both dim and kwargs 1679 with pytest.raises(ValueError): 1680 array.expand_dims({"d": 4}, e=4) 1681 1682 def test_expand_dims(self): 1683 array = DataArray( 1684 np.random.randn(3, 4), 1685 dims=["x", "dim_0"], 1686 coords={"x": np.linspace(0.0, 1.0, 3)}, 1687 attrs={"key": "entry"}, 1688 ) 1689 # pass only dim label 1690 actual = array.expand_dims(dim="y") 1691 expected = DataArray( 1692 np.expand_dims(array.values, 0), 1693 dims=["y", "x", "dim_0"], 1694 coords={"x": np.linspace(0.0, 1.0, 3)}, 1695 attrs={"key": "entry"}, 1696 ) 1697 assert_identical(expected, actual) 1698 roundtripped = actual.squeeze("y", drop=True) 1699 assert_identical(array, roundtripped) 1700 1701 # pass multiple dims 1702 actual = array.expand_dims(dim=["y", "z"]) 1703 expected = DataArray( 1704 np.expand_dims(np.expand_dims(array.values, 0), 0), 1705 dims=["y", "z", "x", "dim_0"], 1706 coords={"x": np.linspace(0.0, 1.0, 3)}, 1707 attrs={"key": "entry"}, 1708 ) 1709 assert_identical(expected, actual) 1710 roundtripped = actual.squeeze(["y", "z"], drop=True) 1711 assert_identical(array, roundtripped) 1712 1713 # pass multiple dims and axis. Axis is out of order 1714 actual = array.expand_dims(dim=["z", "y"], axis=[2, 1]) 1715 expected = DataArray( 1716 np.expand_dims(np.expand_dims(array.values, 1), 2), 1717 dims=["x", "y", "z", "dim_0"], 1718 coords={"x": np.linspace(0.0, 1.0, 3)}, 1719 attrs={"key": "entry"}, 1720 ) 1721 assert_identical(expected, actual) 1722 # make sure the attrs are tracked 1723 assert actual.attrs["key"] == "entry" 1724 roundtripped = actual.squeeze(["z", "y"], drop=True) 1725 assert_identical(array, roundtripped) 1726 1727 # Negative axis and they are out of order 1728 actual = array.expand_dims(dim=["y", "z"], axis=[-1, -2]) 1729 expected = DataArray( 1730 np.expand_dims(np.expand_dims(array.values, -1), -1), 1731 dims=["x", "dim_0", "z", "y"], 1732 coords={"x": np.linspace(0.0, 1.0, 3)}, 1733 attrs={"key": "entry"}, 1734 ) 1735 assert_identical(expected, actual) 1736 assert actual.attrs["key"] == "entry" 1737 roundtripped = actual.squeeze(["y", "z"], drop=True) 1738 assert_identical(array, roundtripped) 1739 1740 def test_expand_dims_with_scalar_coordinate(self): 1741 array = DataArray( 1742 np.random.randn(3, 4), 1743 dims=["x", "dim_0"], 1744 coords={"x": np.linspace(0.0, 1.0, 3), "z": 1.0}, 1745 attrs={"key": "entry"}, 1746 ) 1747 actual = array.expand_dims(dim="z") 1748 expected = DataArray( 1749 np.expand_dims(array.values, 0), 1750 dims=["z", "x", "dim_0"], 1751 coords={"x": np.linspace(0.0, 1.0, 3), "z": np.ones(1)}, 1752 attrs={"key": "entry"}, 1753 ) 1754 assert_identical(expected, actual) 1755 roundtripped = actual.squeeze(["z"], drop=False) 1756 assert_identical(array, roundtripped) 1757 1758 def test_expand_dims_with_greater_dim_size(self): 1759 array = DataArray( 1760 np.random.randn(3, 4), 1761 dims=["x", "dim_0"], 1762 coords={"x": np.linspace(0.0, 1.0, 3), "z": 1.0}, 1763 attrs={"key": "entry"}, 1764 ) 1765 actual = array.expand_dims({"y": 2, "z": 1, "dim_1": ["a", "b", "c"]}) 1766 1767 expected_coords = { 1768 "y": [0, 1], 1769 "z": [1.0], 1770 "dim_1": ["a", "b", "c"], 1771 "x": np.linspace(0, 1, 3), 1772 "dim_0": range(4), 1773 } 1774 expected = DataArray( 1775 array.values * np.ones([2, 1, 3, 3, 4]), 1776 coords=expected_coords, 1777 dims=list(expected_coords.keys()), 1778 attrs={"key": "entry"}, 1779 ).drop_vars(["y", "dim_0"]) 1780 assert_identical(expected, actual) 1781 1782 # Test with kwargs instead of passing dict to dim arg. 1783 1784 other_way = array.expand_dims(dim_1=["a", "b", "c"]) 1785 1786 other_way_expected = DataArray( 1787 array.values * np.ones([3, 3, 4]), 1788 coords={ 1789 "dim_1": ["a", "b", "c"], 1790 "x": np.linspace(0, 1, 3), 1791 "dim_0": range(4), 1792 "z": 1.0, 1793 }, 1794 dims=["dim_1", "x", "dim_0"], 1795 attrs={"key": "entry"}, 1796 ).drop_vars("dim_0") 1797 assert_identical(other_way_expected, other_way) 1798 1799 def test_set_index(self): 1800 indexes = [self.mindex.get_level_values(n) for n in self.mindex.names] 1801 coords = {idx.name: ("x", idx) for idx in indexes} 1802 array = DataArray(self.mda.values, coords=coords, dims="x") 1803 expected = self.mda.copy() 1804 level_3 = ("x", [1, 2, 3, 4]) 1805 array["level_3"] = level_3 1806 expected["level_3"] = level_3 1807 1808 obj = array.set_index(x=self.mindex.names) 1809 assert_identical(obj, expected) 1810 1811 obj = obj.set_index(x="level_3", append=True) 1812 expected = array.set_index(x=["level_1", "level_2", "level_3"]) 1813 assert_identical(obj, expected) 1814 1815 array = array.set_index(x=["level_1", "level_2", "level_3"]) 1816 assert_identical(array, expected) 1817 1818 array2d = DataArray( 1819 np.random.rand(2, 2), 1820 coords={"x": ("x", [0, 1]), "level": ("y", [1, 2])}, 1821 dims=("x", "y"), 1822 ) 1823 with pytest.raises(ValueError, match=r"dimension mismatch"): 1824 array2d.set_index(x="level") 1825 1826 # Issue 3176: Ensure clear error message on key error. 1827 with pytest.raises(ValueError) as excinfo: 1828 obj.set_index(x="level_4") 1829 assert str(excinfo.value) == "level_4 is not the name of an existing variable." 1830 1831 def test_reset_index(self): 1832 indexes = [self.mindex.get_level_values(n) for n in self.mindex.names] 1833 coords = {idx.name: ("x", idx) for idx in indexes} 1834 expected = DataArray(self.mda.values, coords=coords, dims="x") 1835 1836 obj = self.mda.reset_index("x") 1837 assert_identical(obj, expected) 1838 obj = self.mda.reset_index(self.mindex.names) 1839 assert_identical(obj, expected) 1840 obj = self.mda.reset_index(["x", "level_1"]) 1841 assert_identical(obj, expected) 1842 1843 coords = { 1844 "x": ("x", self.mindex.droplevel("level_1")), 1845 "level_1": ("x", self.mindex.get_level_values("level_1")), 1846 } 1847 expected = DataArray(self.mda.values, coords=coords, dims="x") 1848 obj = self.mda.reset_index(["level_1"]) 1849 assert_identical(obj, expected) 1850 1851 expected = DataArray(self.mda.values, dims="x") 1852 obj = self.mda.reset_index("x", drop=True) 1853 assert_identical(obj, expected) 1854 1855 array = self.mda.copy() 1856 array = array.reset_index(["x"], drop=True) 1857 assert_identical(array, expected) 1858 1859 # single index 1860 array = DataArray([1, 2], coords={"x": ["a", "b"]}, dims="x") 1861 expected = DataArray([1, 2], coords={"x_": ("x", ["a", "b"])}, dims="x") 1862 assert_identical(array.reset_index("x"), expected) 1863 1864 def test_reset_index_keep_attrs(self): 1865 coord_1 = DataArray([1, 2], dims=["coord_1"], attrs={"attrs": True}) 1866 da = DataArray([1, 0], [coord_1]) 1867 expected = DataArray([1, 0], {"coord_1_": coord_1}, dims=["coord_1"]) 1868 obj = da.reset_index("coord_1") 1869 assert_identical(expected, obj) 1870 1871 def test_reorder_levels(self): 1872 midx = self.mindex.reorder_levels(["level_2", "level_1"]) 1873 expected = DataArray(self.mda.values, coords={"x": midx}, dims="x") 1874 1875 obj = self.mda.reorder_levels(x=["level_2", "level_1"]) 1876 assert_identical(obj, expected) 1877 1878 array = DataArray([1, 2], dims="x") 1879 with pytest.raises(KeyError): 1880 array.reorder_levels(x=["level_1", "level_2"]) 1881 1882 array["x"] = [0, 1] 1883 with pytest.raises(ValueError, match=r"has no MultiIndex"): 1884 array.reorder_levels(x=["level_1", "level_2"]) 1885 1886 def test_dataset_getitem(self): 1887 dv = self.ds["foo"] 1888 assert_identical(dv, self.dv) 1889 1890 def test_array_interface(self): 1891 assert_array_equal(np.asarray(self.dv), self.x) 1892 # test patched in methods 1893 assert_array_equal(self.dv.astype(float), self.v.astype(float)) 1894 assert_array_equal(self.dv.argsort(), self.v.argsort()) 1895 assert_array_equal(self.dv.clip(2, 3), self.v.clip(2, 3)) 1896 # test ufuncs 1897 expected = deepcopy(self.ds) 1898 expected["foo"][:] = np.sin(self.x) 1899 assert_equal(expected["foo"], np.sin(self.dv)) 1900 assert_array_equal(self.dv, np.maximum(self.v, self.dv)) 1901 bar = Variable(["x", "y"], np.zeros((10, 20))) 1902 assert_equal(self.dv, np.maximum(self.dv, bar)) 1903 1904 def test_astype_attrs(self): 1905 for v in [self.va.copy(), self.mda.copy(), self.ds.copy()]: 1906 v.attrs["foo"] = "bar" 1907 assert v.attrs == v.astype(float).attrs 1908 assert not v.astype(float, keep_attrs=False).attrs 1909 1910 def test_astype_dtype(self): 1911 original = DataArray([-1, 1, 2, 3, 1000]) 1912 converted = original.astype(float) 1913 assert_array_equal(original, converted) 1914 assert np.issubdtype(original.dtype, np.integer) 1915 assert np.issubdtype(converted.dtype, np.floating) 1916 1917 def test_astype_order(self): 1918 original = DataArray([[1, 2], [3, 4]]) 1919 converted = original.astype("d", order="F") 1920 assert_equal(original, converted) 1921 assert original.values.flags["C_CONTIGUOUS"] 1922 assert converted.values.flags["F_CONTIGUOUS"] 1923 1924 def test_astype_subok(self): 1925 class NdArraySubclass(np.ndarray): 1926 pass 1927 1928 original = DataArray(NdArraySubclass(np.arange(3))) 1929 converted_not_subok = original.astype("d", subok=False) 1930 converted_subok = original.astype("d", subok=True) 1931 if not isinstance(original.data, NdArraySubclass): 1932 pytest.xfail("DataArray cannot be backed yet by a subclasses of np.ndarray") 1933 assert isinstance(converted_not_subok.data, np.ndarray) 1934 assert not isinstance(converted_not_subok.data, NdArraySubclass) 1935 assert isinstance(converted_subok.data, NdArraySubclass) 1936 1937 def test_is_null(self): 1938 x = np.random.RandomState(42).randn(5, 6) 1939 x[x < 0] = np.nan 1940 original = DataArray(x, [-np.arange(5), np.arange(6)], ["x", "y"]) 1941 expected = DataArray(pd.isnull(x), [-np.arange(5), np.arange(6)], ["x", "y"]) 1942 assert_identical(expected, original.isnull()) 1943 assert_identical(~expected, original.notnull()) 1944 1945 def test_math(self): 1946 x = self.x 1947 v = self.v 1948 a = self.dv 1949 # variable math was already tested extensively, so let's just make sure 1950 # that all types are properly converted here 1951 assert_equal(a, +a) 1952 assert_equal(a, a + 0) 1953 assert_equal(a, 0 + a) 1954 assert_equal(a, a + 0 * v) 1955 assert_equal(a, 0 * v + a) 1956 assert_equal(a, a + 0 * x) 1957 assert_equal(a, 0 * x + a) 1958 assert_equal(a, a + 0 * a) 1959 assert_equal(a, 0 * a + a) 1960 1961 def test_math_automatic_alignment(self): 1962 a = DataArray(range(5), [("x", range(5))]) 1963 b = DataArray(range(5), [("x", range(1, 6))]) 1964 expected = DataArray(np.ones(4), [("x", [1, 2, 3, 4])]) 1965 assert_identical(a - b, expected) 1966 1967 def test_non_overlapping_dataarrays_return_empty_result(self): 1968 1969 a = DataArray(range(5), [("x", range(5))]) 1970 result = a.isel(x=slice(2)) + a.isel(x=slice(2, None)) 1971 assert len(result["x"]) == 0 1972 1973 def test_empty_dataarrays_return_empty_result(self): 1974 1975 a = DataArray(data=[]) 1976 result = a * a 1977 assert len(result["dim_0"]) == 0 1978 1979 def test_inplace_math_basics(self): 1980 x = self.x 1981 a = self.dv 1982 v = a.variable 1983 b = a 1984 b += 1 1985 assert b is a 1986 assert b.variable is v 1987 assert_array_equal(b.values, x) 1988 assert source_ndarray(b.values) is x 1989 1990 def test_inplace_math_automatic_alignment(self): 1991 a = DataArray(range(5), [("x", range(5))]) 1992 b = DataArray(range(1, 6), [("x", range(1, 6))]) 1993 with pytest.raises(xr.MergeError, match="Automatic alignment is not supported"): 1994 a += b 1995 with pytest.raises(xr.MergeError, match="Automatic alignment is not supported"): 1996 b += a 1997 1998 def test_math_name(self): 1999 # Verify that name is preserved only when it can be done unambiguously. 2000 # The rule (copied from pandas.Series) is keep the current name only if 2001 # the other object has the same name or no name attribute and this 2002 # object isn't a coordinate; otherwise reset to None. 2003 a = self.dv 2004 assert (+a).name == "foo" 2005 assert (a + 0).name == "foo" 2006 assert (a + a.rename(None)).name is None 2007 assert (a + a.rename("bar")).name is None 2008 assert (a + a).name == "foo" 2009 assert (+a["x"]).name == "x" 2010 assert (a["x"] + 0).name == "x" 2011 assert (a + a["x"]).name is None 2012 2013 def test_math_with_coords(self): 2014 coords = { 2015 "x": [-1, -2], 2016 "y": ["ab", "cd", "ef"], 2017 "lat": (["x", "y"], [[1, 2, 3], [-1, -2, -3]]), 2018 "c": -999, 2019 } 2020 orig = DataArray(np.random.randn(2, 3), coords, dims=["x", "y"]) 2021 2022 actual = orig + 1 2023 expected = DataArray(orig.values + 1, orig.coords) 2024 assert_identical(expected, actual) 2025 2026 actual = 1 + orig 2027 assert_identical(expected, actual) 2028 2029 actual = orig + orig[0, 0] 2030 exp_coords = {k: v for k, v in coords.items() if k != "lat"} 2031 expected = DataArray( 2032 orig.values + orig.values[0, 0], exp_coords, dims=["x", "y"] 2033 ) 2034 assert_identical(expected, actual) 2035 2036 actual = orig[0, 0] + orig 2037 assert_identical(expected, actual) 2038 2039 actual = orig[0, 0] + orig[-1, -1] 2040 expected = DataArray(orig.values[0, 0] + orig.values[-1, -1], {"c": -999}) 2041 assert_identical(expected, actual) 2042 2043 actual = orig[:, 0] + orig[0, :] 2044 exp_values = orig[:, 0].values[:, None] + orig[0, :].values[None, :] 2045 expected = DataArray(exp_values, exp_coords, dims=["x", "y"]) 2046 assert_identical(expected, actual) 2047 2048 actual = orig[0, :] + orig[:, 0] 2049 assert_identical(expected.transpose(transpose_coords=True), actual) 2050 2051 actual = orig - orig.transpose(transpose_coords=True) 2052 expected = DataArray(np.zeros((2, 3)), orig.coords) 2053 assert_identical(expected, actual) 2054 2055 actual = orig.transpose(transpose_coords=True) - orig 2056 assert_identical(expected.transpose(transpose_coords=True), actual) 2057 2058 alt = DataArray([1, 1], {"x": [-1, -2], "c": "foo", "d": 555}, "x") 2059 actual = orig + alt 2060 expected = orig + 1 2061 expected.coords["d"] = 555 2062 del expected.coords["c"] 2063 assert_identical(expected, actual) 2064 2065 actual = alt + orig 2066 assert_identical(expected, actual) 2067 2068 def test_index_math(self): 2069 orig = DataArray(range(3), dims="x", name="x") 2070 actual = orig + 1 2071 expected = DataArray(1 + np.arange(3), dims="x", name="x") 2072 assert_identical(expected, actual) 2073 2074 # regression tests for #254 2075 actual = orig[0] < orig 2076 expected = DataArray([False, True, True], dims="x", name="x") 2077 assert_identical(expected, actual) 2078 2079 actual = orig > orig[0] 2080 assert_identical(expected, actual) 2081 2082 def test_dataset_math(self): 2083 # more comprehensive tests with multiple dataset variables 2084 obs = Dataset( 2085 {"tmin": ("x", np.arange(5)), "tmax": ("x", 10 + np.arange(5))}, 2086 {"x": ("x", 0.5 * np.arange(5)), "loc": ("x", range(-2, 3))}, 2087 ) 2088 2089 actual = 2 * obs["tmax"] 2090 expected = DataArray(2 * (10 + np.arange(5)), obs.coords, name="tmax") 2091 assert_identical(actual, expected) 2092 2093 actual = obs["tmax"] - obs["tmin"] 2094 expected = DataArray(10 * np.ones(5), obs.coords) 2095 assert_identical(actual, expected) 2096 2097 sim = Dataset( 2098 { 2099 "tmin": ("x", 1 + np.arange(5)), 2100 "tmax": ("x", 11 + np.arange(5)), 2101 # does *not* include 'loc' as a coordinate 2102 "x": ("x", 0.5 * np.arange(5)), 2103 } 2104 ) 2105 2106 actual = sim["tmin"] - obs["tmin"] 2107 expected = DataArray(np.ones(5), obs.coords, name="tmin") 2108 assert_identical(actual, expected) 2109 2110 actual = -obs["tmin"] + sim["tmin"] 2111 assert_identical(actual, expected) 2112 2113 actual = sim["tmin"].copy() 2114 actual -= obs["tmin"] 2115 assert_identical(actual, expected) 2116 2117 actual = sim.copy() 2118 actual["tmin"] = sim["tmin"] - obs["tmin"] 2119 expected = Dataset( 2120 {"tmin": ("x", np.ones(5)), "tmax": ("x", sim["tmax"].values)}, obs.coords 2121 ) 2122 assert_identical(actual, expected) 2123 2124 actual = sim.copy() 2125 actual["tmin"] -= obs["tmin"] 2126 assert_identical(actual, expected) 2127 2128 def test_stack_unstack(self): 2129 orig = DataArray([[0, 1], [2, 3]], dims=["x", "y"], attrs={"foo": 2}) 2130 assert_identical(orig, orig.unstack()) 2131 2132 # test GH3000 2133 a = orig[:0, :1].stack(dim=("x", "y")).dim.to_index() 2134 if pd.__version__ < "0.24.0": 2135 b = pd.MultiIndex( 2136 levels=[pd.Int64Index([]), pd.Int64Index([0])], 2137 labels=[[], []], 2138 names=["x", "y"], 2139 ) 2140 else: 2141 b = pd.MultiIndex( 2142 levels=[pd.Int64Index([]), pd.Int64Index([0])], 2143 codes=[[], []], 2144 names=["x", "y"], 2145 ) 2146 pd.testing.assert_index_equal(a, b) 2147 2148 actual = orig.stack(z=["x", "y"]).unstack("z").drop_vars(["x", "y"]) 2149 assert_identical(orig, actual) 2150 2151 actual = orig.stack(z=[...]).unstack("z").drop_vars(["x", "y"]) 2152 assert_identical(orig, actual) 2153 2154 dims = ["a", "b", "c", "d", "e"] 2155 orig = xr.DataArray(np.random.rand(1, 2, 3, 2, 1), dims=dims) 2156 stacked = orig.stack(ab=["a", "b"], cd=["c", "d"]) 2157 2158 unstacked = stacked.unstack(["ab", "cd"]) 2159 roundtripped = unstacked.drop_vars(["a", "b", "c", "d"]).transpose(*dims) 2160 assert_identical(orig, roundtripped) 2161 2162 unstacked = stacked.unstack() 2163 roundtripped = unstacked.drop_vars(["a", "b", "c", "d"]).transpose(*dims) 2164 assert_identical(orig, roundtripped) 2165 2166 def test_stack_unstack_decreasing_coordinate(self): 2167 # regression test for GH980 2168 orig = DataArray( 2169 np.random.rand(3, 4), 2170 dims=("y", "x"), 2171 coords={"x": np.arange(4), "y": np.arange(3, 0, -1)}, 2172 ) 2173 stacked = orig.stack(allpoints=["y", "x"]) 2174 actual = stacked.unstack("allpoints") 2175 assert_identical(orig, actual) 2176 2177 def test_unstack_pandas_consistency(self): 2178 df = pd.DataFrame({"foo": range(3), "x": ["a", "b", "b"], "y": [0, 0, 1]}) 2179 s = df.set_index(["x", "y"])["foo"] 2180 expected = DataArray(s.unstack(), name="foo") 2181 actual = DataArray(s, dims="z").unstack("z") 2182 assert_identical(expected, actual) 2183 2184 def test_stack_nonunique_consistency(self, da): 2185 da = da.isel(time=0, drop=True) # 2D 2186 actual = da.stack(z=["a", "x"]) 2187 expected = DataArray(da.to_pandas().stack(), dims="z") 2188 assert_identical(expected, actual) 2189 2190 def test_to_unstacked_dataset_raises_value_error(self): 2191 data = DataArray([0, 1], dims="x", coords={"x": [0, 1]}) 2192 with pytest.raises(ValueError, match="'x' is not a stacked coordinate"): 2193 data.to_unstacked_dataset("x", 0) 2194 2195 def test_transpose(self): 2196 da = DataArray( 2197 np.random.randn(3, 4, 5), 2198 dims=("x", "y", "z"), 2199 coords={ 2200 "x": range(3), 2201 "y": range(4), 2202 "z": range(5), 2203 "xy": (("x", "y"), np.random.randn(3, 4)), 2204 }, 2205 ) 2206 2207 actual = da.transpose(transpose_coords=False) 2208 expected = DataArray(da.values.T, dims=("z", "y", "x"), coords=da.coords) 2209 assert_equal(expected, actual) 2210 2211 actual = da.transpose("z", "y", "x", transpose_coords=True) 2212 expected = DataArray( 2213 da.values.T, 2214 dims=("z", "y", "x"), 2215 coords={ 2216 "x": da.x.values, 2217 "y": da.y.values, 2218 "z": da.z.values, 2219 "xy": (("y", "x"), da.xy.values.T), 2220 }, 2221 ) 2222 assert_equal(expected, actual) 2223 2224 # same as previous but with ellipsis 2225 actual = da.transpose("z", ..., "x", transpose_coords=True) 2226 assert_equal(expected, actual) 2227 2228 # same as previous but with a missing dimension 2229 actual = da.transpose( 2230 "z", "y", "x", "not_a_dim", transpose_coords=True, missing_dims="ignore" 2231 ) 2232 assert_equal(expected, actual) 2233 2234 with pytest.raises(ValueError): 2235 da.transpose("x", "y") 2236 2237 with pytest.raises(ValueError): 2238 da.transpose("not_a_dim", "z", "x", ...) 2239 2240 with pytest.warns(UserWarning): 2241 da.transpose("not_a_dim", "y", "x", ..., missing_dims="warn") 2242 2243 def test_squeeze(self): 2244 assert_equal(self.dv.variable.squeeze(), self.dv.squeeze().variable) 2245 2246 def test_squeeze_drop(self): 2247 array = DataArray([1], [("x", [0])]) 2248 expected = DataArray(1) 2249 actual = array.squeeze(drop=True) 2250 assert_identical(expected, actual) 2251 2252 expected = DataArray(1, {"x": 0}) 2253 actual = array.squeeze(drop=False) 2254 assert_identical(expected, actual) 2255 2256 array = DataArray([[[0.0, 1.0]]], dims=["dim_0", "dim_1", "dim_2"]) 2257 expected = DataArray([[0.0, 1.0]], dims=["dim_1", "dim_2"]) 2258 actual = array.squeeze(axis=0) 2259 assert_identical(expected, actual) 2260 2261 array = DataArray([[[[0.0, 1.0]]]], dims=["dim_0", "dim_1", "dim_2", "dim_3"]) 2262 expected = DataArray([[0.0, 1.0]], dims=["dim_1", "dim_3"]) 2263 actual = array.squeeze(axis=(0, 2)) 2264 assert_identical(expected, actual) 2265 2266 array = DataArray([[[0.0, 1.0]]], dims=["dim_0", "dim_1", "dim_2"]) 2267 with pytest.raises(ValueError): 2268 array.squeeze(axis=0, dim="dim_1") 2269 2270 def test_drop_coordinates(self): 2271 expected = DataArray(np.random.randn(2, 3), dims=["x", "y"]) 2272 arr = expected.copy() 2273 arr.coords["z"] = 2 2274 actual = arr.drop_vars("z") 2275 assert_identical(expected, actual) 2276 2277 with pytest.raises(ValueError): 2278 arr.drop_vars("not found") 2279 2280 actual = expected.drop_vars("not found", errors="ignore") 2281 assert_identical(actual, expected) 2282 2283 with pytest.raises(ValueError, match=r"cannot be found"): 2284 arr.drop_vars("w") 2285 2286 actual = expected.drop_vars("w", errors="ignore") 2287 assert_identical(actual, expected) 2288 2289 renamed = arr.rename("foo") 2290 with pytest.raises(ValueError, match=r"cannot be found"): 2291 renamed.drop_vars("foo") 2292 2293 actual = renamed.drop_vars("foo", errors="ignore") 2294 assert_identical(actual, renamed) 2295 2296 def test_drop_index_labels(self): 2297 arr = DataArray(np.random.randn(2, 3), coords={"y": [0, 1, 2]}, dims=["x", "y"]) 2298 actual = arr.drop_sel(y=[0, 1]) 2299 expected = arr[:, 2:] 2300 assert_identical(actual, expected) 2301 2302 with pytest.raises((KeyError, ValueError), match=r"not .* in axis"): 2303 actual = arr.drop_sel(y=[0, 1, 3]) 2304 2305 actual = arr.drop_sel(y=[0, 1, 3], errors="ignore") 2306 assert_identical(actual, expected) 2307 2308 with pytest.warns(DeprecationWarning): 2309 arr.drop([0, 1, 3], dim="y", errors="ignore") 2310 2311 def test_drop_index_positions(self): 2312 arr = DataArray(np.random.randn(2, 3), dims=["x", "y"]) 2313 actual = arr.drop_isel(y=[0, 1]) 2314 expected = arr[:, 2:] 2315 assert_identical(actual, expected) 2316 2317 def test_dropna(self): 2318 x = np.random.randn(4, 4) 2319 x[::2, 0] = np.nan 2320 arr = DataArray(x, dims=["a", "b"]) 2321 2322 actual = arr.dropna("a") 2323 expected = arr[1::2] 2324 assert_identical(actual, expected) 2325 2326 actual = arr.dropna("b", how="all") 2327 assert_identical(actual, arr) 2328 2329 actual = arr.dropna("a", thresh=1) 2330 assert_identical(actual, arr) 2331 2332 actual = arr.dropna("b", thresh=3) 2333 expected = arr[:, 1:] 2334 assert_identical(actual, expected) 2335 2336 def test_where(self): 2337 arr = DataArray(np.arange(4), dims="x") 2338 expected = arr.sel(x=slice(2)) 2339 actual = arr.where(arr.x < 2, drop=True) 2340 assert_identical(actual, expected) 2341 2342 def test_where_lambda(self): 2343 arr = DataArray(np.arange(4), dims="y") 2344 expected = arr.sel(y=slice(2)) 2345 actual = arr.where(lambda x: x.y < 2, drop=True) 2346 assert_identical(actual, expected) 2347 2348 def test_where_string(self): 2349 array = DataArray(["a", "b"]) 2350 expected = DataArray(np.array(["a", np.nan], dtype=object)) 2351 actual = array.where([True, False]) 2352 assert_identical(actual, expected) 2353 2354 def test_cumops(self): 2355 coords = { 2356 "x": [-1, -2], 2357 "y": ["ab", "cd", "ef"], 2358 "lat": (["x", "y"], [[1, 2, 3], [-1, -2, -3]]), 2359 "c": -999, 2360 } 2361 orig = DataArray([[-1, 0, 1], [-3, 0, 3]], coords, dims=["x", "y"]) 2362 2363 actual = orig.cumsum() 2364 expected = DataArray([[-1, -1, 0], [-4, -4, 0]], coords, dims=["x", "y"]) 2365 assert_identical(expected, actual) 2366 2367 actual = orig.cumsum("x") 2368 expected = DataArray([[-1, 0, 1], [-4, 0, 4]], coords, dims=["x", "y"]) 2369 assert_identical(expected, actual) 2370 2371 actual = orig.cumsum("y") 2372 expected = DataArray([[-1, -1, 0], [-3, -3, 0]], coords, dims=["x", "y"]) 2373 assert_identical(expected, actual) 2374 2375 actual = orig.cumprod("x") 2376 expected = DataArray([[-1, 0, 1], [3, 0, 3]], coords, dims=["x", "y"]) 2377 assert_identical(expected, actual) 2378 2379 actual = orig.cumprod("y") 2380 expected = DataArray([[-1, 0, 0], [-3, 0, 0]], coords, dims=["x", "y"]) 2381 assert_identical(expected, actual) 2382 2383 def test_reduce(self): 2384 coords = { 2385 "x": [-1, -2], 2386 "y": ["ab", "cd", "ef"], 2387 "lat": (["x", "y"], [[1, 2, 3], [-1, -2, -3]]), 2388 "c": -999, 2389 } 2390 orig = DataArray([[-1, 0, 1], [-3, 0, 3]], coords, dims=["x", "y"]) 2391 2392 actual = orig.mean() 2393 expected = DataArray(0, {"c": -999}) 2394 assert_identical(expected, actual) 2395 2396 actual = orig.mean(["x", "y"]) 2397 assert_identical(expected, actual) 2398 2399 actual = orig.mean("x") 2400 expected = DataArray([-2, 0, 2], {"y": coords["y"], "c": -999}, "y") 2401 assert_identical(expected, actual) 2402 2403 actual = orig.mean(["x"]) 2404 assert_identical(expected, actual) 2405 2406 actual = orig.mean("y") 2407 expected = DataArray([0, 0], {"x": coords["x"], "c": -999}, "x") 2408 assert_identical(expected, actual) 2409 2410 assert_equal(self.dv.reduce(np.mean, "x").variable, self.v.reduce(np.mean, "x")) 2411 2412 orig = DataArray([[1, 0, np.nan], [3, 0, 3]], coords, dims=["x", "y"]) 2413 actual = orig.count() 2414 expected = DataArray(5, {"c": -999}) 2415 assert_identical(expected, actual) 2416 2417 # uint support 2418 orig = DataArray(np.arange(6).reshape(3, 2).astype("uint"), dims=["x", "y"]) 2419 assert orig.dtype.kind == "u" 2420 actual = orig.mean(dim="x", skipna=True) 2421 expected = DataArray(orig.values.astype(int), dims=["x", "y"]).mean("x") 2422 assert_equal(actual, expected) 2423 2424 def test_reduce_keepdims(self): 2425 coords = { 2426 "x": [-1, -2], 2427 "y": ["ab", "cd", "ef"], 2428 "lat": (["x", "y"], [[1, 2, 3], [-1, -2, -3]]), 2429 "c": -999, 2430 } 2431 orig = DataArray([[-1, 0, 1], [-3, 0, 3]], coords, dims=["x", "y"]) 2432 2433 # Mean on all axes loses non-constant coordinates 2434 actual = orig.mean(keepdims=True) 2435 expected = DataArray( 2436 orig.data.mean(keepdims=True), 2437 dims=orig.dims, 2438 coords={k: v for k, v in coords.items() if k in ["c"]}, 2439 ) 2440 assert_equal(actual, expected) 2441 2442 assert actual.sizes["x"] == 1 2443 assert actual.sizes["y"] == 1 2444 2445 # Mean on specific axes loses coordinates not involving that axis 2446 actual = orig.mean("y", keepdims=True) 2447 expected = DataArray( 2448 orig.data.mean(axis=1, keepdims=True), 2449 dims=orig.dims, 2450 coords={k: v for k, v in coords.items() if k not in ["y", "lat"]}, 2451 ) 2452 assert_equal(actual, expected) 2453 2454 @requires_bottleneck 2455 def test_reduce_keepdims_bottleneck(self): 2456 import bottleneck 2457 2458 coords = { 2459 "x": [-1, -2], 2460 "y": ["ab", "cd", "ef"], 2461 "lat": (["x", "y"], [[1, 2, 3], [-1, -2, -3]]), 2462 "c": -999, 2463 } 2464 orig = DataArray([[-1, 0, 1], [-3, 0, 3]], coords, dims=["x", "y"]) 2465 2466 # Bottleneck does not have its own keepdims implementation 2467 actual = orig.reduce(bottleneck.nanmean, keepdims=True) 2468 expected = orig.mean(keepdims=True) 2469 assert_equal(actual, expected) 2470 2471 def test_reduce_dtype(self): 2472 coords = { 2473 "x": [-1, -2], 2474 "y": ["ab", "cd", "ef"], 2475 "lat": (["x", "y"], [[1, 2, 3], [-1, -2, -3]]), 2476 "c": -999, 2477 } 2478 orig = DataArray([[-1, 0, 1], [-3, 0, 3]], coords, dims=["x", "y"]) 2479 2480 for dtype in [np.float16, np.float32, np.float64]: 2481 assert orig.astype(float).mean(dtype=dtype).dtype == dtype 2482 2483 def test_reduce_out(self): 2484 coords = { 2485 "x": [-1, -2], 2486 "y": ["ab", "cd", "ef"], 2487 "lat": (["x", "y"], [[1, 2, 3], [-1, -2, -3]]), 2488 "c": -999, 2489 } 2490 orig = DataArray([[-1, 0, 1], [-3, 0, 3]], coords, dims=["x", "y"]) 2491 2492 with pytest.raises(TypeError): 2493 orig.mean(out=np.ones(orig.shape)) 2494 2495 @pytest.mark.parametrize("skipna", [True, False]) 2496 @pytest.mark.parametrize("q", [0.25, [0.50], [0.25, 0.75]]) 2497 @pytest.mark.parametrize( 2498 "axis, dim", zip([None, 0, [0], [0, 1]], [None, "x", ["x"], ["x", "y"]]) 2499 ) 2500 def test_quantile(self, q, axis, dim, skipna): 2501 actual = DataArray(self.va).quantile(q, dim=dim, keep_attrs=True, skipna=skipna) 2502 _percentile_func = np.nanpercentile if skipna else np.percentile 2503 expected = _percentile_func(self.dv.values, np.array(q) * 100, axis=axis) 2504 np.testing.assert_allclose(actual.values, expected) 2505 if is_scalar(q): 2506 assert "quantile" not in actual.dims 2507 else: 2508 assert "quantile" in actual.dims 2509 2510 assert actual.attrs == self.attrs 2511 2512 def test_reduce_keep_attrs(self): 2513 # Test dropped attrs 2514 vm = self.va.mean() 2515 assert len(vm.attrs) == 0 2516 assert vm.attrs == {} 2517 2518 # Test kept attrs 2519 vm = self.va.mean(keep_attrs=True) 2520 assert len(vm.attrs) == len(self.attrs) 2521 assert vm.attrs == self.attrs 2522 2523 def test_assign_attrs(self): 2524 expected = DataArray([], attrs=dict(a=1, b=2)) 2525 expected.attrs["a"] = 1 2526 expected.attrs["b"] = 2 2527 new = DataArray([]) 2528 actual = DataArray([]).assign_attrs(a=1, b=2) 2529 assert_identical(actual, expected) 2530 assert new.attrs == {} 2531 2532 expected.attrs["c"] = 3 2533 new_actual = actual.assign_attrs({"c": 3}) 2534 assert_identical(new_actual, expected) 2535 assert actual.attrs == {"a": 1, "b": 2} 2536 2537 @pytest.mark.parametrize( 2538 "func", [lambda x: x.clip(0, 1), lambda x: np.float64(1.0) * x, np.abs, abs] 2539 ) 2540 def test_propagate_attrs(self, func): 2541 da = DataArray(self.va) 2542 2543 # test defaults 2544 assert func(da).attrs == da.attrs 2545 2546 with set_options(keep_attrs=False): 2547 assert func(da).attrs == {} 2548 2549 with set_options(keep_attrs=True): 2550 assert func(da).attrs == da.attrs 2551 2552 def test_fillna(self): 2553 a = DataArray([np.nan, 1, np.nan, 3], coords={"x": range(4)}, dims="x") 2554 actual = a.fillna(-1) 2555 expected = DataArray([-1, 1, -1, 3], coords={"x": range(4)}, dims="x") 2556 assert_identical(expected, actual) 2557 2558 b = DataArray(range(4), coords={"x": range(4)}, dims="x") 2559 actual = a.fillna(b) 2560 expected = b.copy() 2561 assert_identical(expected, actual) 2562 2563 actual = a.fillna(range(4)) 2564 assert_identical(expected, actual) 2565 2566 actual = a.fillna(b[:3]) 2567 assert_identical(expected, actual) 2568 2569 actual = a.fillna(b[:0]) 2570 assert_identical(a, actual) 2571 2572 with pytest.raises(TypeError, match=r"fillna on a DataArray"): 2573 a.fillna({0: 0}) 2574 2575 with pytest.raises(ValueError, match=r"broadcast"): 2576 a.fillna([1, 2]) 2577 2578 def test_align(self): 2579 array = DataArray( 2580 np.random.random((6, 8)), coords={"x": list("abcdef")}, dims=["x", "y"] 2581 ) 2582 array1, array2 = align(array, array[:5], join="inner") 2583 assert_identical(array1, array[:5]) 2584 assert_identical(array2, array[:5]) 2585 2586 def test_align_dtype(self): 2587 # regression test for #264 2588 x1 = np.arange(30) 2589 x2 = np.arange(5, 35) 2590 a = DataArray(np.random.random((30,)).astype(np.float32), [("x", x1)]) 2591 b = DataArray(np.random.random((30,)).astype(np.float32), [("x", x2)]) 2592 c, d = align(a, b, join="outer") 2593 assert c.dtype == np.float32 2594 2595 def test_align_copy(self): 2596 x = DataArray([1, 2, 3], coords=[("a", [1, 2, 3])]) 2597 y = DataArray([1, 2], coords=[("a", [3, 1])]) 2598 2599 expected_x2 = x 2600 expected_y2 = DataArray([2, np.nan, 1], coords=[("a", [1, 2, 3])]) 2601 2602 x2, y2 = align(x, y, join="outer", copy=False) 2603 assert_identical(expected_x2, x2) 2604 assert_identical(expected_y2, y2) 2605 assert source_ndarray(x2.data) is source_ndarray(x.data) 2606 2607 x2, y2 = align(x, y, join="outer", copy=True) 2608 assert_identical(expected_x2, x2) 2609 assert_identical(expected_y2, y2) 2610 assert source_ndarray(x2.data) is not source_ndarray(x.data) 2611 2612 # Trivial align - 1 element 2613 x = DataArray([1, 2, 3], coords=[("a", [1, 2, 3])]) 2614 (x2,) = align(x, copy=False) 2615 assert_identical(x, x2) 2616 assert source_ndarray(x2.data) is source_ndarray(x.data) 2617 2618 (x2,) = align(x, copy=True) 2619 assert_identical(x, x2) 2620 assert source_ndarray(x2.data) is not source_ndarray(x.data) 2621 2622 def test_align_override(self): 2623 left = DataArray([1, 2, 3], dims="x", coords={"x": [0, 1, 2]}) 2624 right = DataArray( 2625 np.arange(9).reshape((3, 3)), 2626 dims=["x", "y"], 2627 coords={"x": [0.1, 1.1, 2.1], "y": [1, 2, 3]}, 2628 ) 2629 2630 expected_right = DataArray( 2631 np.arange(9).reshape(3, 3), 2632 dims=["x", "y"], 2633 coords={"x": [0, 1, 2], "y": [1, 2, 3]}, 2634 ) 2635 2636 new_left, new_right = align(left, right, join="override") 2637 assert_identical(left, new_left) 2638 assert_identical(new_right, expected_right) 2639 2640 new_left, new_right = align(left, right, exclude="x", join="override") 2641 assert_identical(left, new_left) 2642 assert_identical(right, new_right) 2643 2644 new_left, new_right = xr.align( 2645 left.isel(x=0, drop=True), right, exclude="x", join="override" 2646 ) 2647 assert_identical(left.isel(x=0, drop=True), new_left) 2648 assert_identical(right, new_right) 2649 2650 with pytest.raises(ValueError, match=r"Indexes along dimension 'x' don't have"): 2651 align(left.isel(x=0).expand_dims("x"), right, join="override") 2652 2653 @pytest.mark.parametrize( 2654 "darrays", 2655 [ 2656 [ 2657 DataArray(0), 2658 DataArray([1], [("x", [1])]), 2659 DataArray([2, 3], [("x", [2, 3])]), 2660 ], 2661 [ 2662 DataArray([2, 3], [("x", [2, 3])]), 2663 DataArray([1], [("x", [1])]), 2664 DataArray(0), 2665 ], 2666 ], 2667 ) 2668 def test_align_override_error(self, darrays): 2669 with pytest.raises(ValueError, match=r"Indexes along dimension 'x' don't have"): 2670 xr.align(*darrays, join="override") 2671 2672 def test_align_exclude(self): 2673 x = DataArray([[1, 2], [3, 4]], coords=[("a", [-1, -2]), ("b", [3, 4])]) 2674 y = DataArray([[1, 2], [3, 4]], coords=[("a", [-1, 20]), ("b", [5, 6])]) 2675 z = DataArray([1], dims=["a"], coords={"a": [20], "b": 7}) 2676 2677 x2, y2, z2 = align(x, y, z, join="outer", exclude=["b"]) 2678 expected_x2 = DataArray( 2679 [[3, 4], [1, 2], [np.nan, np.nan]], 2680 coords=[("a", [-2, -1, 20]), ("b", [3, 4])], 2681 ) 2682 expected_y2 = DataArray( 2683 [[np.nan, np.nan], [1, 2], [3, 4]], 2684 coords=[("a", [-2, -1, 20]), ("b", [5, 6])], 2685 ) 2686 expected_z2 = DataArray( 2687 [np.nan, np.nan, 1], dims=["a"], coords={"a": [-2, -1, 20], "b": 7} 2688 ) 2689 assert_identical(expected_x2, x2) 2690 assert_identical(expected_y2, y2) 2691 assert_identical(expected_z2, z2) 2692 2693 def test_align_indexes(self): 2694 x = DataArray([1, 2, 3], coords=[("a", [-1, 10, -2])]) 2695 y = DataArray([1, 2], coords=[("a", [-2, -1])]) 2696 2697 x2, y2 = align(x, y, join="outer", indexes={"a": [10, -1, -2]}) 2698 expected_x2 = DataArray([2, 1, 3], coords=[("a", [10, -1, -2])]) 2699 expected_y2 = DataArray([np.nan, 2, 1], coords=[("a", [10, -1, -2])]) 2700 assert_identical(expected_x2, x2) 2701 assert_identical(expected_y2, y2) 2702 2703 (x2,) = align(x, join="outer", indexes={"a": [-2, 7, 10, -1]}) 2704 expected_x2 = DataArray([3, np.nan, 2, 1], coords=[("a", [-2, 7, 10, -1])]) 2705 assert_identical(expected_x2, x2) 2706 2707 def test_align_without_indexes_exclude(self): 2708 arrays = [DataArray([1, 2, 3], dims=["x"]), DataArray([1, 2], dims=["x"])] 2709 result0, result1 = align(*arrays, exclude=["x"]) 2710 assert_identical(result0, arrays[0]) 2711 assert_identical(result1, arrays[1]) 2712 2713 def test_align_mixed_indexes(self): 2714 array_no_coord = DataArray([1, 2], dims=["x"]) 2715 array_with_coord = DataArray([1, 2], coords=[("x", ["a", "b"])]) 2716 result0, result1 = align(array_no_coord, array_with_coord) 2717 assert_identical(result0, array_with_coord) 2718 assert_identical(result1, array_with_coord) 2719 2720 result0, result1 = align(array_no_coord, array_with_coord, exclude=["x"]) 2721 assert_identical(result0, array_no_coord) 2722 assert_identical(result1, array_with_coord) 2723 2724 def test_align_without_indexes_errors(self): 2725 with pytest.raises(ValueError, match=r"cannot be aligned"): 2726 align(DataArray([1, 2, 3], dims=["x"]), DataArray([1, 2], dims=["x"])) 2727 2728 with pytest.raises(ValueError, match=r"cannot be aligned"): 2729 align( 2730 DataArray([1, 2, 3], dims=["x"]), 2731 DataArray([1, 2], coords=[("x", [0, 1])]), 2732 ) 2733 2734 def test_align_str_dtype(self): 2735 2736 a = DataArray([0, 1], dims=["x"], coords={"x": ["a", "b"]}) 2737 b = DataArray([1, 2], dims=["x"], coords={"x": ["b", "c"]}) 2738 2739 expected_a = DataArray( 2740 [0, 1, np.NaN], dims=["x"], coords={"x": ["a", "b", "c"]} 2741 ) 2742 expected_b = DataArray( 2743 [np.NaN, 1, 2], dims=["x"], coords={"x": ["a", "b", "c"]} 2744 ) 2745 2746 actual_a, actual_b = xr.align(a, b, join="outer") 2747 2748 assert_identical(expected_a, actual_a) 2749 assert expected_a.x.dtype == actual_a.x.dtype 2750 2751 assert_identical(expected_b, actual_b) 2752 assert expected_b.x.dtype == actual_b.x.dtype 2753 2754 def test_broadcast_arrays(self): 2755 x = DataArray([1, 2], coords=[("a", [-1, -2])], name="x") 2756 y = DataArray([1, 2], coords=[("b", [3, 4])], name="y") 2757 x2, y2 = broadcast(x, y) 2758 expected_coords = [("a", [-1, -2]), ("b", [3, 4])] 2759 expected_x2 = DataArray([[1, 1], [2, 2]], expected_coords, name="x") 2760 expected_y2 = DataArray([[1, 2], [1, 2]], expected_coords, name="y") 2761 assert_identical(expected_x2, x2) 2762 assert_identical(expected_y2, y2) 2763 2764 x = DataArray(np.random.randn(2, 3), dims=["a", "b"]) 2765 y = DataArray(np.random.randn(3, 2), dims=["b", "a"]) 2766 x2, y2 = broadcast(x, y) 2767 expected_x2 = x 2768 expected_y2 = y.T 2769 assert_identical(expected_x2, x2) 2770 assert_identical(expected_y2, y2) 2771 2772 def test_broadcast_arrays_misaligned(self): 2773 # broadcast on misaligned coords must auto-align 2774 x = DataArray([[1, 2], [3, 4]], coords=[("a", [-1, -2]), ("b", [3, 4])]) 2775 y = DataArray([1, 2], coords=[("a", [-1, 20])]) 2776 expected_x2 = DataArray( 2777 [[3, 4], [1, 2], [np.nan, np.nan]], 2778 coords=[("a", [-2, -1, 20]), ("b", [3, 4])], 2779 ) 2780 expected_y2 = DataArray( 2781 [[np.nan, np.nan], [1, 1], [2, 2]], 2782 coords=[("a", [-2, -1, 20]), ("b", [3, 4])], 2783 ) 2784 x2, y2 = broadcast(x, y) 2785 assert_identical(expected_x2, x2) 2786 assert_identical(expected_y2, y2) 2787 2788 def test_broadcast_arrays_nocopy(self): 2789 # Test that input data is not copied over in case 2790 # no alteration is needed 2791 x = DataArray([1, 2], coords=[("a", [-1, -2])], name="x") 2792 y = DataArray(3, name="y") 2793 expected_x2 = DataArray([1, 2], coords=[("a", [-1, -2])], name="x") 2794 expected_y2 = DataArray([3, 3], coords=[("a", [-1, -2])], name="y") 2795 2796 x2, y2 = broadcast(x, y) 2797 assert_identical(expected_x2, x2) 2798 assert_identical(expected_y2, y2) 2799 assert source_ndarray(x2.data) is source_ndarray(x.data) 2800 2801 # single-element broadcast (trivial case) 2802 (x2,) = broadcast(x) 2803 assert_identical(x, x2) 2804 assert source_ndarray(x2.data) is source_ndarray(x.data) 2805 2806 def test_broadcast_arrays_exclude(self): 2807 x = DataArray([[1, 2], [3, 4]], coords=[("a", [-1, -2]), ("b", [3, 4])]) 2808 y = DataArray([1, 2], coords=[("a", [-1, 20])]) 2809 z = DataArray(5, coords={"b": 5}) 2810 2811 x2, y2, z2 = broadcast(x, y, z, exclude=["b"]) 2812 expected_x2 = DataArray( 2813 [[3, 4], [1, 2], [np.nan, np.nan]], 2814 coords=[("a", [-2, -1, 20]), ("b", [3, 4])], 2815 ) 2816 expected_y2 = DataArray([np.nan, 1, 2], coords=[("a", [-2, -1, 20])]) 2817 expected_z2 = DataArray( 2818 [5, 5, 5], dims=["a"], coords={"a": [-2, -1, 20], "b": 5} 2819 ) 2820 assert_identical(expected_x2, x2) 2821 assert_identical(expected_y2, y2) 2822 assert_identical(expected_z2, z2) 2823 2824 def test_broadcast_coordinates(self): 2825 # regression test for GH649 2826 ds = Dataset({"a": (["x", "y"], np.ones((5, 6)))}) 2827 x_bc, y_bc, a_bc = broadcast(ds.x, ds.y, ds.a) 2828 assert_identical(ds.a, a_bc) 2829 2830 X, Y = np.meshgrid(np.arange(5), np.arange(6), indexing="ij") 2831 exp_x = DataArray(X, dims=["x", "y"], name="x") 2832 exp_y = DataArray(Y, dims=["x", "y"], name="y") 2833 assert_identical(exp_x, x_bc) 2834 assert_identical(exp_y, y_bc) 2835 2836 def test_to_pandas(self): 2837 # 0d 2838 actual = DataArray(42).to_pandas() 2839 expected = np.array(42) 2840 assert_array_equal(actual, expected) 2841 2842 # 1d 2843 values = np.random.randn(3) 2844 index = pd.Index(["a", "b", "c"], name="x") 2845 da = DataArray(values, coords=[index]) 2846 actual = da.to_pandas() 2847 assert_array_equal(actual.values, values) 2848 assert_array_equal(actual.index, index) 2849 assert_array_equal(actual.index.name, "x") 2850 2851 # 2d 2852 values = np.random.randn(3, 2) 2853 da = DataArray( 2854 values, coords=[("x", ["a", "b", "c"]), ("y", [0, 1])], name="foo" 2855 ) 2856 actual = da.to_pandas() 2857 assert_array_equal(actual.values, values) 2858 assert_array_equal(actual.index, ["a", "b", "c"]) 2859 assert_array_equal(actual.columns, [0, 1]) 2860 2861 # roundtrips 2862 for shape in [(3,), (3, 4)]: 2863 dims = list("abc")[: len(shape)] 2864 da = DataArray(np.random.randn(*shape), dims=dims) 2865 roundtripped = DataArray(da.to_pandas()).drop_vars(dims) 2866 assert_identical(da, roundtripped) 2867 2868 with pytest.raises(ValueError, match=r"cannot convert"): 2869 DataArray(np.random.randn(1, 2, 3, 4, 5)).to_pandas() 2870 2871 def test_to_dataframe(self): 2872 # regression test for #260 2873 arr_np = np.random.randn(3, 4) 2874 2875 arr = DataArray(arr_np, [("B", [1, 2, 3]), ("A", list("cdef"))], name="foo") 2876 expected = arr.to_series() 2877 actual = arr.to_dataframe()["foo"] 2878 assert_array_equal(expected.values, actual.values) 2879 assert_array_equal(expected.name, actual.name) 2880 assert_array_equal(expected.index.values, actual.index.values) 2881 2882 actual = arr.to_dataframe(dim_order=["A", "B"])["foo"] 2883 assert_array_equal(arr_np.transpose().reshape(-1), actual.values) 2884 2885 # regression test for coords with different dimensions 2886 arr.coords["C"] = ("B", [-1, -2, -3]) 2887 expected = arr.to_series().to_frame() 2888 expected["C"] = [-1] * 4 + [-2] * 4 + [-3] * 4 2889 expected = expected[["C", "foo"]] 2890 actual = arr.to_dataframe() 2891 assert_array_equal(expected.values, actual.values) 2892 assert_array_equal(expected.columns.values, actual.columns.values) 2893 assert_array_equal(expected.index.values, actual.index.values) 2894 2895 with pytest.raises(ValueError, match="does not match the set of dimensions"): 2896 arr.to_dataframe(dim_order=["B", "A", "C"]) 2897 2898 with pytest.raises(ValueError, match=r"cannot convert a scalar"): 2899 arr.sel(A="c", B=2).to_dataframe() 2900 2901 arr.name = None # unnamed 2902 with pytest.raises(ValueError, match=r"unnamed"): 2903 arr.to_dataframe() 2904 2905 def test_to_dataframe_multiindex(self): 2906 # regression test for #3008 2907 arr_np = np.random.randn(4, 3) 2908 2909 mindex = pd.MultiIndex.from_product([[1, 2], list("ab")], names=["A", "B"]) 2910 2911 arr = DataArray(arr_np, [("MI", mindex), ("C", [5, 6, 7])], name="foo") 2912 2913 actual = arr.to_dataframe() 2914 assert_array_equal(actual["foo"].values, arr_np.flatten()) 2915 assert_array_equal(actual.index.names, list("ABC")) 2916 assert_array_equal(actual.index.levels[0], [1, 2]) 2917 assert_array_equal(actual.index.levels[1], ["a", "b"]) 2918 assert_array_equal(actual.index.levels[2], [5, 6, 7]) 2919 2920 def test_to_dataframe_0length(self): 2921 # regression test for #3008 2922 arr_np = np.random.randn(4, 0) 2923 2924 mindex = pd.MultiIndex.from_product([[1, 2], list("ab")], names=["A", "B"]) 2925 2926 arr = DataArray(arr_np, [("MI", mindex), ("C", [])], name="foo") 2927 2928 actual = arr.to_dataframe() 2929 assert len(actual) == 0 2930 assert_array_equal(actual.index.names, list("ABC")) 2931 2932 def test_to_pandas_name_matches_coordinate(self): 2933 # coordinate with same name as array 2934 arr = DataArray([1, 2, 3], dims="x", name="x") 2935 series = arr.to_series() 2936 assert_array_equal([1, 2, 3], series.values) 2937 assert_array_equal([0, 1, 2], series.index.values) 2938 assert "x" == series.name 2939 assert "x" == series.index.name 2940 2941 frame = arr.to_dataframe() 2942 expected = series.to_frame() 2943 assert expected.equals(frame) 2944 2945 def test_to_and_from_series(self): 2946 expected = self.dv.to_dataframe()["foo"] 2947 actual = self.dv.to_series() 2948 assert_array_equal(expected.values, actual.values) 2949 assert_array_equal(expected.index.values, actual.index.values) 2950 assert "foo" == actual.name 2951 # test roundtrip 2952 assert_identical(self.dv, DataArray.from_series(actual).drop_vars(["x", "y"])) 2953 # test name is None 2954 actual.name = None 2955 expected_da = self.dv.rename(None) 2956 assert_identical( 2957 expected_da, DataArray.from_series(actual).drop_vars(["x", "y"]) 2958 ) 2959 2960 def test_from_series_multiindex(self): 2961 # GH:3951 2962 df = pd.DataFrame({"B": [1, 2, 3], "A": [4, 5, 6]}) 2963 df = df.rename_axis("num").rename_axis("alpha", axis=1) 2964 actual = df.stack("alpha").to_xarray() 2965 assert (actual.sel(alpha="B") == [1, 2, 3]).all() 2966 assert (actual.sel(alpha="A") == [4, 5, 6]).all() 2967 2968 @requires_sparse 2969 def test_from_series_sparse(self): 2970 import sparse 2971 2972 series = pd.Series([1, 2], index=[("a", 1), ("b", 2)]) 2973 2974 actual_sparse = DataArray.from_series(series, sparse=True) 2975 actual_dense = DataArray.from_series(series, sparse=False) 2976 2977 assert isinstance(actual_sparse.data, sparse.COO) 2978 actual_sparse.data = actual_sparse.data.todense() 2979 assert_identical(actual_sparse, actual_dense) 2980 2981 @requires_sparse 2982 def test_from_multiindex_series_sparse(self): 2983 # regression test for GH4019 2984 import sparse 2985 2986 idx = pd.MultiIndex.from_product([np.arange(3), np.arange(5)], names=["a", "b"]) 2987 series = pd.Series(np.random.RandomState(0).random(len(idx)), index=idx).sample( 2988 n=5, random_state=3 2989 ) 2990 2991 dense = DataArray.from_series(series, sparse=False) 2992 expected_coords = sparse.COO.from_numpy(dense.data, np.nan).coords 2993 2994 actual_sparse = xr.DataArray.from_series(series, sparse=True) 2995 actual_coords = actual_sparse.data.coords 2996 2997 np.testing.assert_equal(actual_coords, expected_coords) 2998 2999 def test_to_and_from_empty_series(self): 3000 # GH697 3001 expected = pd.Series([], dtype=np.float64) 3002 da = DataArray.from_series(expected) 3003 assert len(da) == 0 3004 actual = da.to_series() 3005 assert len(actual) == 0 3006 assert expected.equals(actual) 3007 3008 def test_series_categorical_index(self): 3009 # regression test for GH700 3010 if not hasattr(pd, "CategoricalIndex"): 3011 pytest.skip("requires pandas with CategoricalIndex") 3012 3013 s = pd.Series(np.arange(5), index=pd.CategoricalIndex(list("aabbc"))) 3014 arr = DataArray(s) 3015 assert "'a'" in repr(arr) # should not error 3016 3017 def test_to_and_from_dict(self): 3018 array = DataArray( 3019 np.random.randn(2, 3), {"x": ["a", "b"]}, ["x", "y"], name="foo" 3020 ) 3021 expected = { 3022 "name": "foo", 3023 "dims": ("x", "y"), 3024 "data": array.values.tolist(), 3025 "attrs": {}, 3026 "coords": {"x": {"dims": ("x",), "data": ["a", "b"], "attrs": {}}}, 3027 } 3028 actual = array.to_dict() 3029 3030 # check that they are identical 3031 assert expected == actual 3032 3033 # check roundtrip 3034 assert_identical(array, DataArray.from_dict(actual)) 3035 3036 # a more bare bones representation still roundtrips 3037 d = { 3038 "name": "foo", 3039 "dims": ("x", "y"), 3040 "data": array.values.tolist(), 3041 "coords": {"x": {"dims": "x", "data": ["a", "b"]}}, 3042 } 3043 assert_identical(array, DataArray.from_dict(d)) 3044 3045 # and the most bare bones representation still roundtrips 3046 d = {"name": "foo", "dims": ("x", "y"), "data": array.values} 3047 assert_identical(array.drop_vars("x"), DataArray.from_dict(d)) 3048 3049 # missing a dims in the coords 3050 d = { 3051 "dims": ("x", "y"), 3052 "data": array.values, 3053 "coords": {"x": {"data": ["a", "b"]}}, 3054 } 3055 with pytest.raises( 3056 ValueError, 3057 match=r"cannot convert dict when coords are missing the key 'dims'", 3058 ): 3059 DataArray.from_dict(d) 3060 3061 # this one is missing some necessary information 3062 d = {"dims": ("t")} 3063 with pytest.raises( 3064 ValueError, match=r"cannot convert dict without the key 'data'" 3065 ): 3066 DataArray.from_dict(d) 3067 3068 # check the data=False option 3069 expected_no_data = expected.copy() 3070 del expected_no_data["data"] 3071 del expected_no_data["coords"]["x"]["data"] 3072 endiantype = "<U1" if sys.byteorder == "little" else ">U1" 3073 expected_no_data["coords"]["x"].update({"dtype": endiantype, "shape": (2,)}) 3074 expected_no_data.update({"dtype": "float64", "shape": (2, 3)}) 3075 actual_no_data = array.to_dict(data=False) 3076 assert expected_no_data == actual_no_data 3077 3078 def test_to_and_from_dict_with_time_dim(self): 3079 x = np.random.randn(10, 3) 3080 t = pd.date_range("20130101", periods=10) 3081 lat = [77.7, 83.2, 76] 3082 da = DataArray(x, {"t": t, "lat": lat}, dims=["t", "lat"]) 3083 roundtripped = DataArray.from_dict(da.to_dict()) 3084 assert_identical(da, roundtripped) 3085 3086 def test_to_and_from_dict_with_nan_nat(self): 3087 y = np.random.randn(10, 3) 3088 y[2] = np.nan 3089 t = pd.Series(pd.date_range("20130101", periods=10)) 3090 t[2] = np.nan 3091 lat = [77.7, 83.2, 76] 3092 da = DataArray(y, {"t": t, "lat": lat}, dims=["t", "lat"]) 3093 roundtripped = DataArray.from_dict(da.to_dict()) 3094 assert_identical(da, roundtripped) 3095 3096 def test_to_dict_with_numpy_attrs(self): 3097 # this doesn't need to roundtrip 3098 x = np.random.randn(10, 3) 3099 t = list("abcdefghij") 3100 lat = [77.7, 83.2, 76] 3101 attrs = { 3102 "created": np.float64(1998), 3103 "coords": np.array([37, -110.1, 100]), 3104 "maintainer": "bar", 3105 } 3106 da = DataArray(x, {"t": t, "lat": lat}, dims=["t", "lat"], attrs=attrs) 3107 expected_attrs = { 3108 "created": attrs["created"].item(), 3109 "coords": attrs["coords"].tolist(), 3110 "maintainer": "bar", 3111 } 3112 actual = da.to_dict() 3113 3114 # check that they are identical 3115 assert expected_attrs == actual["attrs"] 3116 3117 def test_to_masked_array(self): 3118 rs = np.random.RandomState(44) 3119 x = rs.random_sample(size=(10, 20)) 3120 x_masked = np.ma.masked_where(x < 0.5, x) 3121 da = DataArray(x_masked) 3122 3123 # Test round trip 3124 x_masked_2 = da.to_masked_array() 3125 da_2 = DataArray(x_masked_2) 3126 assert_array_equal(x_masked, x_masked_2) 3127 assert_equal(da, da_2) 3128 3129 da_masked_array = da.to_masked_array(copy=True) 3130 assert isinstance(da_masked_array, np.ma.MaskedArray) 3131 # Test masks 3132 assert_array_equal(da_masked_array.mask, x_masked.mask) 3133 # Test that mask is unpacked correctly 3134 assert_array_equal(da.values, x_masked.filled(np.nan)) 3135 # Test that the underlying data (including nans) hasn't changed 3136 assert_array_equal(da_masked_array, x_masked.filled(np.nan)) 3137 3138 # Test that copy=False gives access to values 3139 masked_array = da.to_masked_array(copy=False) 3140 masked_array[0, 0] = 10.0 3141 assert masked_array[0, 0] == 10.0 3142 assert da[0, 0].values == 10.0 3143 assert masked_array.base is da.values 3144 assert isinstance(masked_array, np.ma.MaskedArray) 3145 3146 # Test with some odd arrays 3147 for v in [4, np.nan, True, "4", "four"]: 3148 da = DataArray(v) 3149 ma = da.to_masked_array() 3150 assert isinstance(ma, np.ma.MaskedArray) 3151 3152 # Fix GH issue 684 - masked arrays mask should be an array not a scalar 3153 N = 4 3154 v = range(N) 3155 da = DataArray(v) 3156 ma = da.to_masked_array() 3157 assert len(ma.mask) == N 3158 3159 def test_to_and_from_cdms2_classic(self): 3160 """Classic with 1D axes""" 3161 pytest.importorskip("cdms2") 3162 3163 original = DataArray( 3164 np.arange(6).reshape(2, 3), 3165 [ 3166 ("distance", [-2, 2], {"units": "meters"}), 3167 ("time", pd.date_range("2000-01-01", periods=3)), 3168 ], 3169 name="foo", 3170 attrs={"baz": 123}, 3171 ) 3172 expected_coords = [ 3173 IndexVariable("distance", [-2, 2]), 3174 IndexVariable("time", [0, 1, 2]), 3175 ] 3176 actual = original.to_cdms2() 3177 assert_array_equal(actual.asma(), original) 3178 assert actual.id == original.name 3179 assert tuple(actual.getAxisIds()) == original.dims 3180 for axis, coord in zip(actual.getAxisList(), expected_coords): 3181 assert axis.id == coord.name 3182 assert_array_equal(axis, coord.values) 3183 assert actual.baz == original.attrs["baz"] 3184 3185 component_times = actual.getAxis(1).asComponentTime() 3186 assert len(component_times) == 3 3187 assert str(component_times[0]) == "2000-1-1 0:0:0.0" 3188 3189 roundtripped = DataArray.from_cdms2(actual) 3190 assert_identical(original, roundtripped) 3191 3192 back = from_cdms2(actual) 3193 assert original.dims == back.dims 3194 assert original.coords.keys() == back.coords.keys() 3195 for coord_name in original.coords.keys(): 3196 assert_array_equal(original.coords[coord_name], back.coords[coord_name]) 3197 3198 def test_to_and_from_cdms2_sgrid(self): 3199 """Curvilinear (structured) grid 3200 3201 The rectangular grid case is covered by the classic case 3202 """ 3203 pytest.importorskip("cdms2") 3204 3205 lonlat = np.mgrid[:3, :4] 3206 lon = DataArray(lonlat[1], dims=["y", "x"], name="lon") 3207 lat = DataArray(lonlat[0], dims=["y", "x"], name="lat") 3208 x = DataArray(np.arange(lon.shape[1]), dims=["x"], name="x") 3209 y = DataArray(np.arange(lon.shape[0]), dims=["y"], name="y") 3210 original = DataArray( 3211 lonlat.sum(axis=0), 3212 dims=["y", "x"], 3213 coords=dict(x=x, y=y, lon=lon, lat=lat), 3214 name="sst", 3215 ) 3216 actual = original.to_cdms2() 3217 assert tuple(actual.getAxisIds()) == original.dims 3218 assert_array_equal(original.coords["lon"], actual.getLongitude().asma()) 3219 assert_array_equal(original.coords["lat"], actual.getLatitude().asma()) 3220 3221 back = from_cdms2(actual) 3222 assert original.dims == back.dims 3223 assert set(original.coords.keys()) == set(back.coords.keys()) 3224 assert_array_equal(original.coords["lat"], back.coords["lat"]) 3225 assert_array_equal(original.coords["lon"], back.coords["lon"]) 3226 3227 def test_to_and_from_cdms2_ugrid(self): 3228 """Unstructured grid""" 3229 pytest.importorskip("cdms2") 3230 3231 lon = DataArray(np.random.uniform(size=5), dims=["cell"], name="lon") 3232 lat = DataArray(np.random.uniform(size=5), dims=["cell"], name="lat") 3233 cell = DataArray(np.arange(5), dims=["cell"], name="cell") 3234 original = DataArray( 3235 np.arange(5), dims=["cell"], coords={"lon": lon, "lat": lat, "cell": cell} 3236 ) 3237 actual = original.to_cdms2() 3238 assert tuple(actual.getAxisIds()) == original.dims 3239 assert_array_equal(original.coords["lon"], actual.getLongitude().getValue()) 3240 assert_array_equal(original.coords["lat"], actual.getLatitude().getValue()) 3241 3242 back = from_cdms2(actual) 3243 assert set(original.dims) == set(back.dims) 3244 assert set(original.coords.keys()) == set(back.coords.keys()) 3245 assert_array_equal(original.coords["lat"], back.coords["lat"]) 3246 assert_array_equal(original.coords["lon"], back.coords["lon"]) 3247 3248 def test_to_dataset_whole(self): 3249 unnamed = DataArray([1, 2], dims="x") 3250 with pytest.raises(ValueError, match=r"unable to convert unnamed"): 3251 unnamed.to_dataset() 3252 3253 actual = unnamed.to_dataset(name="foo") 3254 expected = Dataset({"foo": ("x", [1, 2])}) 3255 assert_identical(expected, actual) 3256 3257 named = DataArray([1, 2], dims="x", name="foo", attrs={"y": "testattr"}) 3258 actual = named.to_dataset() 3259 expected = Dataset({"foo": ("x", [1, 2], {"y": "testattr"})}) 3260 assert_identical(expected, actual) 3261 3262 # Test promoting attrs 3263 actual = named.to_dataset(promote_attrs=True) 3264 expected = Dataset( 3265 {"foo": ("x", [1, 2], {"y": "testattr"})}, attrs={"y": "testattr"} 3266 ) 3267 assert_identical(expected, actual) 3268 3269 with pytest.raises(TypeError): 3270 actual = named.to_dataset("bar") 3271 3272 def test_to_dataset_split(self): 3273 array = DataArray([1, 2, 3], coords=[("x", list("abc"))], attrs={"a": 1}) 3274 expected = Dataset({"a": 1, "b": 2, "c": 3}, attrs={"a": 1}) 3275 actual = array.to_dataset("x") 3276 assert_identical(expected, actual) 3277 3278 with pytest.raises(TypeError): 3279 array.to_dataset("x", name="foo") 3280 3281 roundtripped = actual.to_array(dim="x") 3282 assert_identical(array, roundtripped) 3283 3284 array = DataArray([1, 2, 3], dims="x") 3285 expected = Dataset({0: 1, 1: 2, 2: 3}) 3286 actual = array.to_dataset("x") 3287 assert_identical(expected, actual) 3288 3289 def test_to_dataset_retains_keys(self): 3290 3291 # use dates as convenient non-str objects. Not a specific date test 3292 import datetime 3293 3294 dates = [datetime.date(2000, 1, d) for d in range(1, 4)] 3295 3296 array = DataArray([1, 2, 3], coords=[("x", dates)], attrs={"a": 1}) 3297 3298 # convert to dateset and back again 3299 result = array.to_dataset("x").to_array(dim="x") 3300 3301 assert_equal(array, result) 3302 3303 def test__title_for_slice(self): 3304 array = DataArray( 3305 np.ones((4, 3, 2)), 3306 dims=["a", "b", "c"], 3307 coords={"a": range(4), "b": range(3), "c": range(2)}, 3308 ) 3309 assert "" == array._title_for_slice() 3310 assert "c = 0" == array.isel(c=0)._title_for_slice() 3311 title = array.isel(b=1, c=0)._title_for_slice() 3312 assert "b = 1, c = 0" == title or "c = 0, b = 1" == title 3313 3314 a2 = DataArray(np.ones((4, 1)), dims=["a", "b"]) 3315 assert "" == a2._title_for_slice() 3316 3317 def test__title_for_slice_truncate(self): 3318 array = DataArray(np.ones(4)) 3319 array.coords["a"] = "a" * 100 3320 array.coords["b"] = "b" * 100 3321 3322 nchar = 80 3323 title = array._title_for_slice(truncate=nchar) 3324 3325 assert nchar == len(title) 3326 assert title.endswith("...") 3327 3328 def test_dataarray_diff_n1(self): 3329 da = DataArray(np.random.randn(3, 4), dims=["x", "y"]) 3330 actual = da.diff("y") 3331 expected = DataArray(np.diff(da.values, axis=1), dims=["x", "y"]) 3332 assert_equal(expected, actual) 3333 3334 def test_coordinate_diff(self): 3335 # regression test for GH634 3336 arr = DataArray(range(0, 20, 2), dims=["lon"], coords=[range(10)]) 3337 lon = arr.coords["lon"] 3338 expected = DataArray([1] * 9, dims=["lon"], coords=[range(1, 10)], name="lon") 3339 actual = lon.diff("lon") 3340 assert_equal(expected, actual) 3341 3342 @pytest.mark.parametrize("offset", [-5, 0, 1, 2]) 3343 @pytest.mark.parametrize("fill_value, dtype", [(2, int), (dtypes.NA, float)]) 3344 def test_shift(self, offset, fill_value, dtype): 3345 arr = DataArray([1, 2, 3], dims="x") 3346 actual = arr.shift(x=1, fill_value=fill_value) 3347 if fill_value == dtypes.NA: 3348 # if we supply the default, we expect the missing value for a 3349 # float array 3350 fill_value = np.nan 3351 expected = DataArray([fill_value, 1, 2], dims="x") 3352 assert_identical(expected, actual) 3353 assert actual.dtype == dtype 3354 3355 arr = DataArray([1, 2, 3], [("x", ["a", "b", "c"])]) 3356 expected = DataArray(arr.to_pandas().shift(offset)) 3357 actual = arr.shift(x=offset) 3358 assert_identical(expected, actual) 3359 3360 def test_roll_coords(self): 3361 arr = DataArray([1, 2, 3], coords={"x": range(3)}, dims="x") 3362 actual = arr.roll(x=1, roll_coords=True) 3363 expected = DataArray([3, 1, 2], coords=[("x", [2, 0, 1])]) 3364 assert_identical(expected, actual) 3365 3366 def test_roll_no_coords(self): 3367 arr = DataArray([1, 2, 3], coords={"x": range(3)}, dims="x") 3368 actual = arr.roll(x=1) 3369 expected = DataArray([3, 1, 2], coords=[("x", [0, 1, 2])]) 3370 assert_identical(expected, actual) 3371 3372 def test_copy_with_data(self): 3373 orig = DataArray( 3374 np.random.random(size=(2, 2)), 3375 dims=("x", "y"), 3376 attrs={"attr1": "value1"}, 3377 coords={"x": [4, 3]}, 3378 name="helloworld", 3379 ) 3380 new_data = np.arange(4).reshape(2, 2) 3381 actual = orig.copy(data=new_data) 3382 expected = orig.copy() 3383 expected.data = new_data 3384 assert_identical(expected, actual) 3385 3386 @pytest.mark.xfail(raises=AssertionError) 3387 @pytest.mark.parametrize( 3388 "deep, expected_orig", 3389 [ 3390 [ 3391 True, 3392 xr.DataArray( 3393 xr.IndexVariable("a", np.array([1, 2])), 3394 coords={"a": [1, 2]}, 3395 dims=["a"], 3396 ), 3397 ], 3398 [ 3399 False, 3400 xr.DataArray( 3401 xr.IndexVariable("a", np.array([999, 2])), 3402 coords={"a": [999, 2]}, 3403 dims=["a"], 3404 ), 3405 ], 3406 ], 3407 ) 3408 def test_copy_coords(self, deep, expected_orig): 3409 """The test fails for the shallow copy, and apparently only on Windows 3410 for some reason. In windows coords seem to be immutable unless it's one 3411 dataarray deep copied from another.""" 3412 da = xr.DataArray( 3413 np.ones([2, 2, 2]), 3414 coords={"a": [1, 2], "b": ["x", "y"], "c": [0, 1]}, 3415 dims=["a", "b", "c"], 3416 ) 3417 da_cp = da.copy(deep) 3418 da_cp["a"].data[0] = 999 3419 3420 expected_cp = xr.DataArray( 3421 xr.IndexVariable("a", np.array([999, 2])), 3422 coords={"a": [999, 2]}, 3423 dims=["a"], 3424 ) 3425 assert_identical(da_cp["a"], expected_cp) 3426 3427 assert_identical(da["a"], expected_orig) 3428 3429 def test_real_and_imag(self): 3430 array = DataArray(1 + 2j) 3431 assert_identical(array.real, DataArray(1)) 3432 assert_identical(array.imag, DataArray(2)) 3433 3434 def test_setattr_raises(self): 3435 array = DataArray(0, coords={"scalar": 1}, attrs={"foo": "bar"}) 3436 with pytest.raises(AttributeError, match=r"cannot set attr"): 3437 array.scalar = 2 3438 with pytest.raises(AttributeError, match=r"cannot set attr"): 3439 array.foo = 2 3440 with pytest.raises(AttributeError, match=r"cannot set attr"): 3441 array.other = 2 3442 3443 def test_full_like(self): 3444 # For more thorough tests, see test_variable.py 3445 da = DataArray( 3446 np.random.random(size=(2, 2)), 3447 dims=("x", "y"), 3448 attrs={"attr1": "value1"}, 3449 coords={"x": [4, 3]}, 3450 name="helloworld", 3451 ) 3452 3453 actual = full_like(da, 2) 3454 expect = da.copy(deep=True) 3455 expect.values = [[2.0, 2.0], [2.0, 2.0]] 3456 assert_identical(expect, actual) 3457 3458 # override dtype 3459 actual = full_like(da, fill_value=True, dtype=bool) 3460 expect.values = [[True, True], [True, True]] 3461 assert expect.dtype == bool 3462 assert_identical(expect, actual) 3463 3464 with pytest.raises(ValueError, match="'dtype' cannot be dict-like"): 3465 full_like(da, fill_value=True, dtype={"x": bool}) 3466 3467 def test_dot(self): 3468 x = np.linspace(-3, 3, 6) 3469 y = np.linspace(-3, 3, 5) 3470 z = range(4) 3471 da_vals = np.arange(6 * 5 * 4).reshape((6, 5, 4)) 3472 da = DataArray(da_vals, coords=[x, y, z], dims=["x", "y", "z"]) 3473 3474 dm_vals = range(4) 3475 dm = DataArray(dm_vals, coords=[z], dims=["z"]) 3476 3477 # nd dot 1d 3478 actual = da.dot(dm) 3479 expected_vals = np.tensordot(da_vals, dm_vals, [2, 0]) 3480 expected = DataArray(expected_vals, coords=[x, y], dims=["x", "y"]) 3481 assert_equal(expected, actual) 3482 3483 # all shared dims 3484 actual = da.dot(da) 3485 expected_vals = np.tensordot(da_vals, da_vals, axes=([0, 1, 2], [0, 1, 2])) 3486 expected = DataArray(expected_vals) 3487 assert_equal(expected, actual) 3488 3489 # multiple shared dims 3490 dm_vals = np.arange(20 * 5 * 4).reshape((20, 5, 4)) 3491 j = np.linspace(-3, 3, 20) 3492 dm = DataArray(dm_vals, coords=[j, y, z], dims=["j", "y", "z"]) 3493 actual = da.dot(dm) 3494 expected_vals = np.tensordot(da_vals, dm_vals, axes=([1, 2], [1, 2])) 3495 expected = DataArray(expected_vals, coords=[x, j], dims=["x", "j"]) 3496 assert_equal(expected, actual) 3497 3498 # Ellipsis: all dims are shared 3499 actual = da.dot(da, dims=...) 3500 expected = da.dot(da) 3501 assert_equal(expected, actual) 3502 3503 # Ellipsis: not all dims are shared 3504 actual = da.dot(dm, dims=...) 3505 expected = da.dot(dm, dims=("j", "x", "y", "z")) 3506 assert_equal(expected, actual) 3507 3508 with pytest.raises(NotImplementedError): 3509 da.dot(dm.to_dataset(name="dm")) 3510 with pytest.raises(TypeError): 3511 da.dot(dm.values) 3512 3513 def test_dot_align_coords(self): 3514 # GH 3694 3515 3516 x = np.linspace(-3, 3, 6) 3517 y = np.linspace(-3, 3, 5) 3518 z_a = range(4) 3519 da_vals = np.arange(6 * 5 * 4).reshape((6, 5, 4)) 3520 da = DataArray(da_vals, coords=[x, y, z_a], dims=["x", "y", "z"]) 3521 3522 z_m = range(2, 6) 3523 dm_vals = range(4) 3524 dm = DataArray(dm_vals, coords=[z_m], dims=["z"]) 3525 3526 with xr.set_options(arithmetic_join="exact"): 3527 with pytest.raises(ValueError, match=r"indexes along dimension"): 3528 da.dot(dm) 3529 3530 da_aligned, dm_aligned = xr.align(da, dm, join="inner") 3531 3532 # nd dot 1d 3533 actual = da.dot(dm) 3534 expected_vals = np.tensordot(da_aligned.values, dm_aligned.values, [2, 0]) 3535 expected = DataArray(expected_vals, coords=[x, da_aligned.y], dims=["x", "y"]) 3536 assert_equal(expected, actual) 3537 3538 # multiple shared dims 3539 dm_vals = np.arange(20 * 5 * 4).reshape((20, 5, 4)) 3540 j = np.linspace(-3, 3, 20) 3541 dm = DataArray(dm_vals, coords=[j, y, z_m], dims=["j", "y", "z"]) 3542 da_aligned, dm_aligned = xr.align(da, dm, join="inner") 3543 actual = da.dot(dm) 3544 expected_vals = np.tensordot( 3545 da_aligned.values, dm_aligned.values, axes=([1, 2], [1, 2]) 3546 ) 3547 expected = DataArray(expected_vals, coords=[x, j], dims=["x", "j"]) 3548 assert_equal(expected, actual) 3549 3550 def test_matmul(self): 3551 3552 # copied from above (could make a fixture) 3553 x = np.linspace(-3, 3, 6) 3554 y = np.linspace(-3, 3, 5) 3555 z = range(4) 3556 da_vals = np.arange(6 * 5 * 4).reshape((6, 5, 4)) 3557 da = DataArray(da_vals, coords=[x, y, z], dims=["x", "y", "z"]) 3558 3559 result = da @ da 3560 expected = da.dot(da) 3561 assert_identical(result, expected) 3562 3563 def test_matmul_align_coords(self): 3564 # GH 3694 3565 3566 x_a = np.arange(6) 3567 x_b = np.arange(2, 8) 3568 da_vals = np.arange(6) 3569 da_a = DataArray(da_vals, coords=[x_a], dims=["x"]) 3570 da_b = DataArray(da_vals, coords=[x_b], dims=["x"]) 3571 3572 # only test arithmetic_join="inner" (=default) 3573 result = da_a @ da_b 3574 expected = da_a.dot(da_b) 3575 assert_identical(result, expected) 3576 3577 with xr.set_options(arithmetic_join="exact"): 3578 with pytest.raises(ValueError, match=r"indexes along dimension"): 3579 da_a @ da_b 3580 3581 def test_binary_op_propagate_indexes(self): 3582 # regression test for GH2227 3583 self.dv["x"] = np.arange(self.dv.sizes["x"]) 3584 expected = self.dv.xindexes["x"] 3585 3586 actual = (self.dv * 10).xindexes["x"] 3587 assert expected is actual 3588 3589 actual = (self.dv > 10).xindexes["x"] 3590 assert expected is actual 3591 3592 def test_binary_op_join_setting(self): 3593 dim = "x" 3594 align_type = "outer" 3595 coords_l, coords_r = [0, 1, 2], [1, 2, 3] 3596 missing_3 = xr.DataArray(coords_l, [(dim, coords_l)]) 3597 missing_0 = xr.DataArray(coords_r, [(dim, coords_r)]) 3598 with xr.set_options(arithmetic_join=align_type): 3599 actual = missing_0 + missing_3 3600 missing_0_aligned, missing_3_aligned = xr.align( 3601 missing_0, missing_3, join=align_type 3602 ) 3603 expected = xr.DataArray([np.nan, 2, 4, np.nan], [(dim, [0, 1, 2, 3])]) 3604 assert_equal(actual, expected) 3605 3606 def test_combine_first(self): 3607 ar0 = DataArray([[0, 0], [0, 0]], [("x", ["a", "b"]), ("y", [-1, 0])]) 3608 ar1 = DataArray([[1, 1], [1, 1]], [("x", ["b", "c"]), ("y", [0, 1])]) 3609 ar2 = DataArray([2], [("x", ["d"])]) 3610 3611 actual = ar0.combine_first(ar1) 3612 expected = DataArray( 3613 [[0, 0, np.nan], [0, 0, 1], [np.nan, 1, 1]], 3614 [("x", ["a", "b", "c"]), ("y", [-1, 0, 1])], 3615 ) 3616 assert_equal(actual, expected) 3617 3618 actual = ar1.combine_first(ar0) 3619 expected = DataArray( 3620 [[0, 0, np.nan], [0, 1, 1], [np.nan, 1, 1]], 3621 [("x", ["a", "b", "c"]), ("y", [-1, 0, 1])], 3622 ) 3623 assert_equal(actual, expected) 3624 3625 actual = ar0.combine_first(ar2) 3626 expected = DataArray( 3627 [[0, 0], [0, 0], [2, 2]], [("x", ["a", "b", "d"]), ("y", [-1, 0])] 3628 ) 3629 assert_equal(actual, expected) 3630 3631 def test_sortby(self): 3632 da = DataArray( 3633 [[1, 2], [3, 4], [5, 6]], [("x", ["c", "b", "a"]), ("y", [1, 0])] 3634 ) 3635 3636 sorted1d = DataArray( 3637 [[5, 6], [3, 4], [1, 2]], [("x", ["a", "b", "c"]), ("y", [1, 0])] 3638 ) 3639 3640 sorted2d = DataArray( 3641 [[6, 5], [4, 3], [2, 1]], [("x", ["a", "b", "c"]), ("y", [0, 1])] 3642 ) 3643 3644 expected = sorted1d 3645 dax = DataArray([100, 99, 98], [("x", ["c", "b", "a"])]) 3646 actual = da.sortby(dax) 3647 assert_equal(actual, expected) 3648 3649 # test descending order sort 3650 actual = da.sortby(dax, ascending=False) 3651 assert_equal(actual, da) 3652 3653 # test alignment (fills in nan for 'c') 3654 dax_short = DataArray([98, 97], [("x", ["b", "a"])]) 3655 actual = da.sortby(dax_short) 3656 assert_equal(actual, expected) 3657 3658 # test multi-dim sort by 1D dataarray values 3659 expected = sorted2d 3660 dax = DataArray([100, 99, 98], [("x", ["c", "b", "a"])]) 3661 day = DataArray([90, 80], [("y", [1, 0])]) 3662 actual = da.sortby([day, dax]) 3663 assert_equal(actual, expected) 3664 3665 expected = sorted1d 3666 actual = da.sortby("x") 3667 assert_equal(actual, expected) 3668 3669 expected = sorted2d 3670 actual = da.sortby(["x", "y"]) 3671 assert_equal(actual, expected) 3672 3673 @requires_bottleneck 3674 def test_rank(self): 3675 # floats 3676 ar = DataArray([[3, 4, np.nan, 1]]) 3677 expect_0 = DataArray([[1, 1, np.nan, 1]]) 3678 expect_1 = DataArray([[2, 3, np.nan, 1]]) 3679 assert_equal(ar.rank("dim_0"), expect_0) 3680 assert_equal(ar.rank("dim_1"), expect_1) 3681 # int 3682 x = DataArray([3, 2, 1]) 3683 assert_equal(x.rank("dim_0"), x) 3684 # str 3685 y = DataArray(["c", "b", "a"]) 3686 assert_equal(y.rank("dim_0"), x) 3687 3688 x = DataArray([3.0, 1.0, np.nan, 2.0, 4.0], dims=("z",)) 3689 y = DataArray([0.75, 0.25, np.nan, 0.5, 1.0], dims=("z",)) 3690 assert_equal(y.rank("z", pct=True), y) 3691 3692 @pytest.mark.parametrize("use_dask", [True, False]) 3693 @pytest.mark.parametrize("use_datetime", [True, False]) 3694 @pytest.mark.filterwarnings("ignore:overflow encountered in multiply") 3695 def test_polyfit(self, use_dask, use_datetime): 3696 if use_dask and not has_dask: 3697 pytest.skip("requires dask") 3698 xcoord = xr.DataArray( 3699 pd.date_range("1970-01-01", freq="D", periods=10), dims=("x",), name="x" 3700 ) 3701 x = xr.core.missing.get_clean_interp_index(xcoord, "x") 3702 if not use_datetime: 3703 xcoord = x 3704 3705 da_raw = DataArray( 3706 np.stack( 3707 (10 + 1e-15 * x + 2e-28 * x ** 2, 30 + 2e-14 * x + 1e-29 * x ** 2) 3708 ), 3709 dims=("d", "x"), 3710 coords={"x": xcoord, "d": [0, 1]}, 3711 ) 3712 3713 if use_dask: 3714 da = da_raw.chunk({"d": 1}) 3715 else: 3716 da = da_raw 3717 3718 out = da.polyfit("x", 2) 3719 expected = DataArray( 3720 [[2e-28, 1e-15, 10], [1e-29, 2e-14, 30]], 3721 dims=("d", "degree"), 3722 coords={"degree": [2, 1, 0], "d": [0, 1]}, 3723 ).T 3724 assert_allclose(out.polyfit_coefficients, expected, rtol=1e-3) 3725 3726 # Full output and deficient rank 3727 with warnings.catch_warnings(): 3728 warnings.simplefilter("ignore", np.RankWarning) 3729 out = da.polyfit("x", 12, full=True) 3730 assert out.polyfit_residuals.isnull().all() 3731 3732 # With NaN 3733 da_raw[0, 1:3] = np.nan 3734 if use_dask: 3735 da = da_raw.chunk({"d": 1}) 3736 else: 3737 da = da_raw 3738 out = da.polyfit("x", 2, skipna=True, cov=True) 3739 assert_allclose(out.polyfit_coefficients, expected, rtol=1e-3) 3740 assert "polyfit_covariance" in out 3741 3742 # Skipna + Full output 3743 out = da.polyfit("x", 2, skipna=True, full=True) 3744 assert_allclose(out.polyfit_coefficients, expected, rtol=1e-3) 3745 assert out.x_matrix_rank == 3 3746 np.testing.assert_almost_equal(out.polyfit_residuals, [0, 0]) 3747 3748 with warnings.catch_warnings(): 3749 warnings.simplefilter("ignore", np.RankWarning) 3750 out = da.polyfit("x", 8, full=True) 3751 np.testing.assert_array_equal(out.polyfit_residuals.isnull(), [True, False]) 3752 3753 def test_pad_constant(self): 3754 ar = DataArray(np.arange(3 * 4 * 5).reshape(3, 4, 5)) 3755 actual = ar.pad(dim_0=(1, 3)) 3756 expected = DataArray( 3757 np.pad( 3758 np.arange(3 * 4 * 5).reshape(3, 4, 5).astype(np.float32), 3759 mode="constant", 3760 pad_width=((1, 3), (0, 0), (0, 0)), 3761 constant_values=np.nan, 3762 ) 3763 ) 3764 assert actual.shape == (7, 4, 5) 3765 assert_identical(actual, expected) 3766 3767 ar = xr.DataArray([9], dims="x") 3768 3769 actual = ar.pad(x=1) 3770 expected = xr.DataArray([np.NaN, 9, np.NaN], dims="x") 3771 assert_identical(actual, expected) 3772 3773 actual = ar.pad(x=1, constant_values=1.23456) 3774 expected = xr.DataArray([1, 9, 1], dims="x") 3775 assert_identical(actual, expected) 3776 3777 if LooseVersion(np.__version__) >= "1.20": 3778 with pytest.raises(ValueError, match="cannot convert float NaN to integer"): 3779 ar.pad(x=1, constant_values=np.NaN) 3780 else: 3781 actual = ar.pad(x=1, constant_values=np.NaN) 3782 expected = xr.DataArray( 3783 [-9223372036854775808, 9, -9223372036854775808], dims="x" 3784 ) 3785 assert_identical(actual, expected) 3786 3787 def test_pad_coords(self): 3788 ar = DataArray( 3789 np.arange(3 * 4 * 5).reshape(3, 4, 5), 3790 [("x", np.arange(3)), ("y", np.arange(4)), ("z", np.arange(5))], 3791 ) 3792 actual = ar.pad(x=(1, 3), constant_values=1) 3793 expected = DataArray( 3794 np.pad( 3795 np.arange(3 * 4 * 5).reshape(3, 4, 5), 3796 mode="constant", 3797 pad_width=((1, 3), (0, 0), (0, 0)), 3798 constant_values=1, 3799 ), 3800 [ 3801 ( 3802 "x", 3803 np.pad( 3804 np.arange(3).astype(np.float32), 3805 mode="constant", 3806 pad_width=(1, 3), 3807 constant_values=np.nan, 3808 ), 3809 ), 3810 ("y", np.arange(4)), 3811 ("z", np.arange(5)), 3812 ], 3813 ) 3814 assert_identical(actual, expected) 3815 3816 @pytest.mark.parametrize("mode", ("minimum", "maximum", "mean", "median")) 3817 @pytest.mark.parametrize( 3818 "stat_length", (None, 3, (1, 3), {"dim_0": (2, 1), "dim_2": (4, 2)}) 3819 ) 3820 def test_pad_stat_length(self, mode, stat_length): 3821 ar = DataArray(np.arange(3 * 4 * 5).reshape(3, 4, 5)) 3822 actual = ar.pad(dim_0=(1, 3), dim_2=(2, 2), mode=mode, stat_length=stat_length) 3823 if isinstance(stat_length, dict): 3824 stat_length = (stat_length["dim_0"], (4, 4), stat_length["dim_2"]) 3825 expected = DataArray( 3826 np.pad( 3827 np.arange(3 * 4 * 5).reshape(3, 4, 5), 3828 pad_width=((1, 3), (0, 0), (2, 2)), 3829 mode=mode, 3830 stat_length=stat_length, 3831 ) 3832 ) 3833 assert actual.shape == (7, 4, 9) 3834 assert_identical(actual, expected) 3835 3836 @pytest.mark.parametrize( 3837 "end_values", (None, 3, (3, 5), {"dim_0": (2, 1), "dim_2": (4, 2)}) 3838 ) 3839 def test_pad_linear_ramp(self, end_values): 3840 ar = DataArray(np.arange(3 * 4 * 5).reshape(3, 4, 5)) 3841 actual = ar.pad( 3842 dim_0=(1, 3), dim_2=(2, 2), mode="linear_ramp", end_values=end_values 3843 ) 3844 if end_values is None: 3845 end_values = 0 3846 elif isinstance(end_values, dict): 3847 end_values = (end_values["dim_0"], (4, 4), end_values["dim_2"]) 3848 expected = DataArray( 3849 np.pad( 3850 np.arange(3 * 4 * 5).reshape(3, 4, 5), 3851 pad_width=((1, 3), (0, 0), (2, 2)), 3852 mode="linear_ramp", 3853 end_values=end_values, 3854 ) 3855 ) 3856 assert actual.shape == (7, 4, 9) 3857 assert_identical(actual, expected) 3858 3859 @pytest.mark.parametrize("mode", ("reflect", "symmetric")) 3860 @pytest.mark.parametrize("reflect_type", (None, "even", "odd")) 3861 def test_pad_reflect(self, mode, reflect_type): 3862 3863 ar = DataArray(np.arange(3 * 4 * 5).reshape(3, 4, 5)) 3864 actual = ar.pad( 3865 dim_0=(1, 3), dim_2=(2, 2), mode=mode, reflect_type=reflect_type 3866 ) 3867 np_kwargs = { 3868 "array": np.arange(3 * 4 * 5).reshape(3, 4, 5), 3869 "pad_width": ((1, 3), (0, 0), (2, 2)), 3870 "mode": mode, 3871 } 3872 # numpy does not support reflect_type=None 3873 if reflect_type is not None: 3874 np_kwargs["reflect_type"] = reflect_type 3875 expected = DataArray(np.pad(**np_kwargs)) 3876 3877 assert actual.shape == (7, 4, 9) 3878 assert_identical(actual, expected) 3879 3880 @pytest.mark.parametrize("parser", ["pandas", "python"]) 3881 @pytest.mark.parametrize( 3882 "engine", ["python", None, pytest.param("numexpr", marks=[requires_numexpr])] 3883 ) 3884 @pytest.mark.parametrize( 3885 "backend", ["numpy", pytest.param("dask", marks=[requires_dask])] 3886 ) 3887 def test_query(self, backend, engine, parser): 3888 """Test querying a dataset.""" 3889 3890 # setup test data 3891 np.random.seed(42) 3892 a = np.arange(0, 10, 1) 3893 b = np.random.randint(0, 100, size=10) 3894 c = np.linspace(0, 1, 20) 3895 d = np.random.choice(["foo", "bar", "baz"], size=30, replace=True).astype( 3896 object 3897 ) 3898 if backend == "numpy": 3899 aa = DataArray(data=a, dims=["x"], name="a") 3900 bb = DataArray(data=b, dims=["x"], name="b") 3901 cc = DataArray(data=c, dims=["y"], name="c") 3902 dd = DataArray(data=d, dims=["z"], name="d") 3903 3904 elif backend == "dask": 3905 import dask.array as da 3906 3907 aa = DataArray(data=da.from_array(a, chunks=3), dims=["x"], name="a") 3908 bb = DataArray(data=da.from_array(b, chunks=3), dims=["x"], name="b") 3909 cc = DataArray(data=da.from_array(c, chunks=7), dims=["y"], name="c") 3910 dd = DataArray(data=da.from_array(d, chunks=12), dims=["z"], name="d") 3911 3912 # query single dim, single variable 3913 actual = aa.query(x="a > 5", engine=engine, parser=parser) 3914 expect = aa.isel(x=(a > 5)) 3915 assert_identical(expect, actual) 3916 3917 # query single dim, single variable, via dict 3918 actual = aa.query(dict(x="a > 5"), engine=engine, parser=parser) 3919 expect = aa.isel(dict(x=(a > 5))) 3920 assert_identical(expect, actual) 3921 3922 # query single dim, single variable 3923 actual = bb.query(x="b > 50", engine=engine, parser=parser) 3924 expect = bb.isel(x=(b > 50)) 3925 assert_identical(expect, actual) 3926 3927 # query single dim, single variable 3928 actual = cc.query(y="c < .5", engine=engine, parser=parser) 3929 expect = cc.isel(y=(c < 0.5)) 3930 assert_identical(expect, actual) 3931 3932 # query single dim, single string variable 3933 if parser == "pandas": 3934 # N.B., this query currently only works with the pandas parser 3935 # xref https://github.com/pandas-dev/pandas/issues/40436 3936 actual = dd.query(z='d == "bar"', engine=engine, parser=parser) 3937 expect = dd.isel(z=(d == "bar")) 3938 assert_identical(expect, actual) 3939 3940 # test error handling 3941 with pytest.raises(ValueError): 3942 aa.query("a > 5") # must be dict or kwargs 3943 with pytest.raises(ValueError): 3944 aa.query(x=(a > 5)) # must be query string 3945 with pytest.raises(UndefinedVariableError): 3946 aa.query(x="spam > 50") # name not present 3947 3948 @requires_scipy 3949 @pytest.mark.parametrize("use_dask", [True, False]) 3950 def test_curvefit(self, use_dask): 3951 if use_dask and not has_dask: 3952 pytest.skip("requires dask") 3953 3954 def exp_decay(t, n0, tau=1): 3955 return n0 * np.exp(-t / tau) 3956 3957 t = np.arange(0, 5, 0.5) 3958 da = DataArray( 3959 np.stack([exp_decay(t, 3, 3), exp_decay(t, 5, 4), np.nan * t], axis=-1), 3960 dims=("t", "x"), 3961 coords={"t": t, "x": [0, 1, 2]}, 3962 ) 3963 da[0, 0] = np.nan 3964 3965 expected = DataArray( 3966 [[3, 3], [5, 4], [np.nan, np.nan]], 3967 dims=("x", "param"), 3968 coords={"x": [0, 1, 2], "param": ["n0", "tau"]}, 3969 ) 3970 3971 if use_dask: 3972 da = da.chunk({"x": 1}) 3973 3974 fit = da.curvefit( 3975 coords=[da.t], func=exp_decay, p0={"n0": 4}, bounds={"tau": [2, 6]} 3976 ) 3977 assert_allclose(fit.curvefit_coefficients, expected, rtol=1e-3) 3978 3979 da = da.compute() 3980 fit = da.curvefit(coords="t", func=np.power, reduce_dims="x", param_names=["a"]) 3981 assert "a" in fit.param 3982 assert "x" not in fit.dims 3983 3984 def test_curvefit_helpers(self): 3985 def exp_decay(t, n0, tau=1): 3986 return n0 * np.exp(-t / tau) 3987 3988 params, func_args = xr.core.dataset._get_func_args(exp_decay, []) 3989 assert params == ["n0", "tau"] 3990 param_defaults, bounds_defaults = xr.core.dataset._initialize_curvefit_params( 3991 params, {"n0": 4}, {"tau": [5, np.inf]}, func_args 3992 ) 3993 assert param_defaults == {"n0": 4, "tau": 6} 3994 assert bounds_defaults == {"n0": (-np.inf, np.inf), "tau": (5, np.inf)} 3995 3996 param_names = ["a"] 3997 params, func_args = xr.core.dataset._get_func_args(np.power, param_names) 3998 assert params == param_names 3999 with pytest.raises(ValueError): 4000 xr.core.dataset._get_func_args(np.power, []) 4001 4002 4003class TestReduce: 4004 @pytest.fixture(autouse=True) 4005 def setup(self): 4006 self.attrs = {"attr1": "value1", "attr2": 2929} 4007 4008 4009@pytest.mark.parametrize( 4010 "x, minindex, maxindex, nanindex", 4011 [ 4012 (np.array([0, 1, 2, 0, -2, -4, 2]), 5, 2, None), 4013 (np.array([0.0, 1.0, 2.0, 0.0, -2.0, -4.0, 2.0]), 5, 2, None), 4014 (np.array([1.0, np.NaN, 2.0, np.NaN, -2.0, -4.0, 2.0]), 5, 2, 1), 4015 ( 4016 np.array([1.0, np.NaN, 2.0, np.NaN, -2.0, -4.0, 2.0]).astype("object"), 4017 5, 4018 2, 4019 1, 4020 ), 4021 (np.array([np.NaN, np.NaN]), np.NaN, np.NaN, 0), 4022 ( 4023 np.array( 4024 ["2015-12-31", "2020-01-02", "2020-01-01", "2016-01-01"], 4025 dtype="datetime64[ns]", 4026 ), 4027 0, 4028 1, 4029 None, 4030 ), 4031 ], 4032) 4033class TestReduce1D(TestReduce): 4034 def test_min(self, x, minindex, maxindex, nanindex): 4035 ar = xr.DataArray( 4036 x, dims=["x"], coords={"x": np.arange(x.size) * 4}, attrs=self.attrs 4037 ) 4038 4039 if np.isnan(minindex): 4040 minindex = 0 4041 4042 expected0 = ar.isel(x=minindex, drop=True) 4043 result0 = ar.min(keep_attrs=True) 4044 assert_identical(result0, expected0) 4045 4046 result1 = ar.min() 4047 expected1 = expected0.copy() 4048 expected1.attrs = {} 4049 assert_identical(result1, expected1) 4050 4051 result2 = ar.min(skipna=False) 4052 if nanindex is not None and ar.dtype.kind != "O": 4053 expected2 = ar.isel(x=nanindex, drop=True) 4054 expected2.attrs = {} 4055 else: 4056 expected2 = expected1 4057 4058 assert_identical(result2, expected2) 4059 4060 def test_max(self, x, minindex, maxindex, nanindex): 4061 ar = xr.DataArray( 4062 x, dims=["x"], coords={"x": np.arange(x.size) * 4}, attrs=self.attrs 4063 ) 4064 4065 if np.isnan(minindex): 4066 maxindex = 0 4067 4068 expected0 = ar.isel(x=maxindex, drop=True) 4069 result0 = ar.max(keep_attrs=True) 4070 assert_identical(result0, expected0) 4071 4072 result1 = ar.max() 4073 expected1 = expected0.copy() 4074 expected1.attrs = {} 4075 assert_identical(result1, expected1) 4076 4077 result2 = ar.max(skipna=False) 4078 if nanindex is not None and ar.dtype.kind != "O": 4079 expected2 = ar.isel(x=nanindex, drop=True) 4080 expected2.attrs = {} 4081 else: 4082 expected2 = expected1 4083 4084 assert_identical(result2, expected2) 4085 4086 @pytest.mark.filterwarnings( 4087 "ignore:Behaviour of argmin/argmax with neither dim nor :DeprecationWarning" 4088 ) 4089 def test_argmin(self, x, minindex, maxindex, nanindex): 4090 ar = xr.DataArray( 4091 x, dims=["x"], coords={"x": np.arange(x.size) * 4}, attrs=self.attrs 4092 ) 4093 indarr = xr.DataArray(np.arange(x.size, dtype=np.intp), dims=["x"]) 4094 4095 if np.isnan(minindex): 4096 with pytest.raises(ValueError): 4097 ar.argmin() 4098 return 4099 4100 expected0 = indarr[minindex] 4101 result0 = ar.argmin() 4102 assert_identical(result0, expected0) 4103 4104 result1 = ar.argmin(keep_attrs=True) 4105 expected1 = expected0.copy() 4106 expected1.attrs = self.attrs 4107 assert_identical(result1, expected1) 4108 4109 result2 = ar.argmin(skipna=False) 4110 if nanindex is not None and ar.dtype.kind != "O": 4111 expected2 = indarr.isel(x=nanindex, drop=True) 4112 expected2.attrs = {} 4113 else: 4114 expected2 = expected0 4115 4116 assert_identical(result2, expected2) 4117 4118 @pytest.mark.filterwarnings( 4119 "ignore:Behaviour of argmin/argmax with neither dim nor :DeprecationWarning" 4120 ) 4121 def test_argmax(self, x, minindex, maxindex, nanindex): 4122 ar = xr.DataArray( 4123 x, dims=["x"], coords={"x": np.arange(x.size) * 4}, attrs=self.attrs 4124 ) 4125 indarr = xr.DataArray(np.arange(x.size, dtype=np.intp), dims=["x"]) 4126 4127 if np.isnan(maxindex): 4128 with pytest.raises(ValueError): 4129 ar.argmax() 4130 return 4131 4132 expected0 = indarr[maxindex] 4133 result0 = ar.argmax() 4134 assert_identical(result0, expected0) 4135 4136 result1 = ar.argmax(keep_attrs=True) 4137 expected1 = expected0.copy() 4138 expected1.attrs = self.attrs 4139 assert_identical(result1, expected1) 4140 4141 result2 = ar.argmax(skipna=False) 4142 if nanindex is not None and ar.dtype.kind != "O": 4143 expected2 = indarr.isel(x=nanindex, drop=True) 4144 expected2.attrs = {} 4145 else: 4146 expected2 = expected0 4147 4148 assert_identical(result2, expected2) 4149 4150 @pytest.mark.parametrize("use_dask", [True, False]) 4151 def test_idxmin(self, x, minindex, maxindex, nanindex, use_dask): 4152 if use_dask and not has_dask: 4153 pytest.skip("requires dask") 4154 if use_dask and x.dtype.kind == "M": 4155 pytest.xfail("dask operation 'argmin' breaks when dtype is datetime64 (M)") 4156 ar0_raw = xr.DataArray( 4157 x, dims=["x"], coords={"x": np.arange(x.size) * 4}, attrs=self.attrs 4158 ) 4159 4160 if use_dask: 4161 ar0 = ar0_raw.chunk({}) 4162 else: 4163 ar0 = ar0_raw 4164 4165 # dim doesn't exist 4166 with pytest.raises(KeyError): 4167 ar0.idxmin(dim="spam") 4168 4169 # Scalar Dataarray 4170 with pytest.raises(ValueError): 4171 xr.DataArray(5).idxmin() 4172 4173 coordarr0 = xr.DataArray(ar0.coords["x"], dims=["x"]) 4174 coordarr1 = coordarr0.copy() 4175 4176 hasna = np.isnan(minindex) 4177 if np.isnan(minindex): 4178 minindex = 0 4179 4180 if hasna: 4181 coordarr1[...] = 1 4182 fill_value_0 = np.NaN 4183 else: 4184 fill_value_0 = 1 4185 4186 expected0 = ( 4187 (coordarr1 * fill_value_0).isel(x=minindex, drop=True).astype("float") 4188 ) 4189 expected0.name = "x" 4190 4191 # Default fill value (NaN) 4192 result0 = ar0.idxmin() 4193 assert_identical(result0, expected0) 4194 4195 # Manually specify NaN fill_value 4196 result1 = ar0.idxmin(fill_value=np.NaN) 4197 assert_identical(result1, expected0) 4198 4199 # keep_attrs 4200 result2 = ar0.idxmin(keep_attrs=True) 4201 expected2 = expected0.copy() 4202 expected2.attrs = self.attrs 4203 assert_identical(result2, expected2) 4204 4205 # skipna=False 4206 if nanindex is not None and ar0.dtype.kind != "O": 4207 expected3 = coordarr0.isel(x=nanindex, drop=True).astype("float") 4208 expected3.name = "x" 4209 expected3.attrs = {} 4210 else: 4211 expected3 = expected0.copy() 4212 4213 result3 = ar0.idxmin(skipna=False) 4214 assert_identical(result3, expected3) 4215 4216 # fill_value should be ignored with skipna=False 4217 result4 = ar0.idxmin(skipna=False, fill_value=-100j) 4218 assert_identical(result4, expected3) 4219 4220 # Float fill_value 4221 if hasna: 4222 fill_value_5 = -1.1 4223 else: 4224 fill_value_5 = 1 4225 4226 expected5 = (coordarr1 * fill_value_5).isel(x=minindex, drop=True) 4227 expected5.name = "x" 4228 4229 result5 = ar0.idxmin(fill_value=-1.1) 4230 assert_identical(result5, expected5) 4231 4232 # Integer fill_value 4233 if hasna: 4234 fill_value_6 = -1 4235 else: 4236 fill_value_6 = 1 4237 4238 expected6 = (coordarr1 * fill_value_6).isel(x=minindex, drop=True) 4239 expected6.name = "x" 4240 4241 result6 = ar0.idxmin(fill_value=-1) 4242 assert_identical(result6, expected6) 4243 4244 # Complex fill_value 4245 if hasna: 4246 fill_value_7 = -1j 4247 else: 4248 fill_value_7 = 1 4249 4250 expected7 = (coordarr1 * fill_value_7).isel(x=minindex, drop=True) 4251 expected7.name = "x" 4252 4253 result7 = ar0.idxmin(fill_value=-1j) 4254 assert_identical(result7, expected7) 4255 4256 @pytest.mark.parametrize("use_dask", [True, False]) 4257 def test_idxmax(self, x, minindex, maxindex, nanindex, use_dask): 4258 if use_dask and not has_dask: 4259 pytest.skip("requires dask") 4260 if use_dask and x.dtype.kind == "M": 4261 pytest.xfail("dask operation 'argmax' breaks when dtype is datetime64 (M)") 4262 ar0_raw = xr.DataArray( 4263 x, dims=["x"], coords={"x": np.arange(x.size) * 4}, attrs=self.attrs 4264 ) 4265 4266 if use_dask: 4267 ar0 = ar0_raw.chunk({}) 4268 else: 4269 ar0 = ar0_raw 4270 4271 # dim doesn't exist 4272 with pytest.raises(KeyError): 4273 ar0.idxmax(dim="spam") 4274 4275 # Scalar Dataarray 4276 with pytest.raises(ValueError): 4277 xr.DataArray(5).idxmax() 4278 4279 coordarr0 = xr.DataArray(ar0.coords["x"], dims=["x"]) 4280 coordarr1 = coordarr0.copy() 4281 4282 hasna = np.isnan(maxindex) 4283 if np.isnan(maxindex): 4284 maxindex = 0 4285 4286 if hasna: 4287 coordarr1[...] = 1 4288 fill_value_0 = np.NaN 4289 else: 4290 fill_value_0 = 1 4291 4292 expected0 = ( 4293 (coordarr1 * fill_value_0).isel(x=maxindex, drop=True).astype("float") 4294 ) 4295 expected0.name = "x" 4296 4297 # Default fill value (NaN) 4298 result0 = ar0.idxmax() 4299 assert_identical(result0, expected0) 4300 4301 # Manually specify NaN fill_value 4302 result1 = ar0.idxmax(fill_value=np.NaN) 4303 assert_identical(result1, expected0) 4304 4305 # keep_attrs 4306 result2 = ar0.idxmax(keep_attrs=True) 4307 expected2 = expected0.copy() 4308 expected2.attrs = self.attrs 4309 assert_identical(result2, expected2) 4310 4311 # skipna=False 4312 if nanindex is not None and ar0.dtype.kind != "O": 4313 expected3 = coordarr0.isel(x=nanindex, drop=True).astype("float") 4314 expected3.name = "x" 4315 expected3.attrs = {} 4316 else: 4317 expected3 = expected0.copy() 4318 4319 result3 = ar0.idxmax(skipna=False) 4320 assert_identical(result3, expected3) 4321 4322 # fill_value should be ignored with skipna=False 4323 result4 = ar0.idxmax(skipna=False, fill_value=-100j) 4324 assert_identical(result4, expected3) 4325 4326 # Float fill_value 4327 if hasna: 4328 fill_value_5 = -1.1 4329 else: 4330 fill_value_5 = 1 4331 4332 expected5 = (coordarr1 * fill_value_5).isel(x=maxindex, drop=True) 4333 expected5.name = "x" 4334 4335 result5 = ar0.idxmax(fill_value=-1.1) 4336 assert_identical(result5, expected5) 4337 4338 # Integer fill_value 4339 if hasna: 4340 fill_value_6 = -1 4341 else: 4342 fill_value_6 = 1 4343 4344 expected6 = (coordarr1 * fill_value_6).isel(x=maxindex, drop=True) 4345 expected6.name = "x" 4346 4347 result6 = ar0.idxmax(fill_value=-1) 4348 assert_identical(result6, expected6) 4349 4350 # Complex fill_value 4351 if hasna: 4352 fill_value_7 = -1j 4353 else: 4354 fill_value_7 = 1 4355 4356 expected7 = (coordarr1 * fill_value_7).isel(x=maxindex, drop=True) 4357 expected7.name = "x" 4358 4359 result7 = ar0.idxmax(fill_value=-1j) 4360 assert_identical(result7, expected7) 4361 4362 @pytest.mark.filterwarnings( 4363 "ignore:Behaviour of argmin/argmax with neither dim nor :DeprecationWarning" 4364 ) 4365 def test_argmin_dim(self, x, minindex, maxindex, nanindex): 4366 ar = xr.DataArray( 4367 x, dims=["x"], coords={"x": np.arange(x.size) * 4}, attrs=self.attrs 4368 ) 4369 indarr = xr.DataArray(np.arange(x.size, dtype=np.intp), dims=["x"]) 4370 4371 if np.isnan(minindex): 4372 with pytest.raises(ValueError): 4373 ar.argmin() 4374 return 4375 4376 expected0 = {"x": indarr[minindex]} 4377 result0 = ar.argmin(...) 4378 for key in expected0: 4379 assert_identical(result0[key], expected0[key]) 4380 4381 result1 = ar.argmin(..., keep_attrs=True) 4382 expected1 = deepcopy(expected0) 4383 for da in expected1.values(): 4384 da.attrs = self.attrs 4385 for key in expected1: 4386 assert_identical(result1[key], expected1[key]) 4387 4388 result2 = ar.argmin(..., skipna=False) 4389 if nanindex is not None and ar.dtype.kind != "O": 4390 expected2 = {"x": indarr.isel(x=nanindex, drop=True)} 4391 expected2["x"].attrs = {} 4392 else: 4393 expected2 = expected0 4394 4395 for key in expected2: 4396 assert_identical(result2[key], expected2[key]) 4397 4398 @pytest.mark.filterwarnings( 4399 "ignore:Behaviour of argmin/argmax with neither dim nor :DeprecationWarning" 4400 ) 4401 def test_argmax_dim(self, x, minindex, maxindex, nanindex): 4402 ar = xr.DataArray( 4403 x, dims=["x"], coords={"x": np.arange(x.size) * 4}, attrs=self.attrs 4404 ) 4405 indarr = xr.DataArray(np.arange(x.size, dtype=np.intp), dims=["x"]) 4406 4407 if np.isnan(maxindex): 4408 with pytest.raises(ValueError): 4409 ar.argmax() 4410 return 4411 4412 expected0 = {"x": indarr[maxindex]} 4413 result0 = ar.argmax(...) 4414 for key in expected0: 4415 assert_identical(result0[key], expected0[key]) 4416 4417 result1 = ar.argmax(..., keep_attrs=True) 4418 expected1 = deepcopy(expected0) 4419 for da in expected1.values(): 4420 da.attrs = self.attrs 4421 for key in expected1: 4422 assert_identical(result1[key], expected1[key]) 4423 4424 result2 = ar.argmax(..., skipna=False) 4425 if nanindex is not None and ar.dtype.kind != "O": 4426 expected2 = {"x": indarr.isel(x=nanindex, drop=True)} 4427 expected2["x"].attrs = {} 4428 else: 4429 expected2 = expected0 4430 4431 for key in expected2: 4432 assert_identical(result2[key], expected2[key]) 4433 4434 4435@pytest.mark.parametrize( 4436 "x, minindex, maxindex, nanindex", 4437 [ 4438 ( 4439 np.array( 4440 [ 4441 [0, 1, 2, 0, -2, -4, 2], 4442 [1, 1, 1, 1, 1, 1, 1], 4443 [0, 0, -10, 5, 20, 0, 0], 4444 ] 4445 ), 4446 [5, 0, 2], 4447 [2, 0, 4], 4448 [None, None, None], 4449 ), 4450 ( 4451 np.array( 4452 [ 4453 [2.0, 1.0, 2.0, 0.0, -2.0, -4.0, 2.0], 4454 [-4.0, np.NaN, 2.0, np.NaN, -2.0, -4.0, 2.0], 4455 [np.NaN] * 7, 4456 ] 4457 ), 4458 [5, 0, np.NaN], 4459 [0, 2, np.NaN], 4460 [None, 1, 0], 4461 ), 4462 ( 4463 np.array( 4464 [ 4465 [2.0, 1.0, 2.0, 0.0, -2.0, -4.0, 2.0], 4466 [-4.0, np.NaN, 2.0, np.NaN, -2.0, -4.0, 2.0], 4467 [np.NaN] * 7, 4468 ] 4469 ).astype("object"), 4470 [5, 0, np.NaN], 4471 [0, 2, np.NaN], 4472 [None, 1, 0], 4473 ), 4474 ( 4475 np.array( 4476 [ 4477 ["2015-12-31", "2020-01-02", "2020-01-01", "2016-01-01"], 4478 ["2020-01-02", "2020-01-02", "2020-01-02", "2020-01-02"], 4479 ["1900-01-01", "1-02-03", "1900-01-02", "1-02-03"], 4480 ], 4481 dtype="datetime64[ns]", 4482 ), 4483 [0, 0, 1], 4484 [1, 0, 2], 4485 [None, None, None], 4486 ), 4487 ], 4488) 4489class TestReduce2D(TestReduce): 4490 def test_min(self, x, minindex, maxindex, nanindex): 4491 ar = xr.DataArray( 4492 x, 4493 dims=["y", "x"], 4494 coords={"x": np.arange(x.shape[1]) * 4, "y": 1 - np.arange(x.shape[0])}, 4495 attrs=self.attrs, 4496 ) 4497 4498 minindex = [x if not np.isnan(x) else 0 for x in minindex] 4499 expected0 = [ 4500 ar.isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(minindex) 4501 ] 4502 expected0 = xr.concat(expected0, dim="y") 4503 4504 result0 = ar.min(dim="x", keep_attrs=True) 4505 assert_identical(result0, expected0) 4506 4507 result1 = ar.min(dim="x") 4508 expected1 = expected0 4509 expected1.attrs = {} 4510 assert_identical(result1, expected1) 4511 4512 result2 = ar.min(axis=1) 4513 assert_identical(result2, expected1) 4514 4515 minindex = [ 4516 x if y is None or ar.dtype.kind == "O" else y 4517 for x, y in zip(minindex, nanindex) 4518 ] 4519 expected2 = [ 4520 ar.isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(minindex) 4521 ] 4522 expected2 = xr.concat(expected2, dim="y") 4523 expected2.attrs = {} 4524 4525 result3 = ar.min(dim="x", skipna=False) 4526 4527 assert_identical(result3, expected2) 4528 4529 def test_max(self, x, minindex, maxindex, nanindex): 4530 ar = xr.DataArray( 4531 x, 4532 dims=["y", "x"], 4533 coords={"x": np.arange(x.shape[1]) * 4, "y": 1 - np.arange(x.shape[0])}, 4534 attrs=self.attrs, 4535 ) 4536 4537 maxindex = [x if not np.isnan(x) else 0 for x in maxindex] 4538 expected0 = [ 4539 ar.isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(maxindex) 4540 ] 4541 expected0 = xr.concat(expected0, dim="y") 4542 4543 result0 = ar.max(dim="x", keep_attrs=True) 4544 assert_identical(result0, expected0) 4545 4546 result1 = ar.max(dim="x") 4547 expected1 = expected0.copy() 4548 expected1.attrs = {} 4549 assert_identical(result1, expected1) 4550 4551 result2 = ar.max(axis=1) 4552 assert_identical(result2, expected1) 4553 4554 maxindex = [ 4555 x if y is None or ar.dtype.kind == "O" else y 4556 for x, y in zip(maxindex, nanindex) 4557 ] 4558 expected2 = [ 4559 ar.isel(y=yi).isel(x=indi, drop=True) for yi, indi in enumerate(maxindex) 4560 ] 4561 expected2 = xr.concat(expected2, dim="y") 4562 expected2.attrs = {} 4563 4564 result3 = ar.max(dim="x", skipna=False) 4565 4566 assert_identical(result3, expected2) 4567 4568 def test_argmin(self, x, minindex, maxindex, nanindex): 4569 ar = xr.DataArray( 4570 x, 4571 dims=["y", "x"], 4572 coords={"x": np.arange(x.shape[1]) * 4, "y": 1 - np.arange(x.shape[0])}, 4573 attrs=self.attrs, 4574 ) 4575 indarr = np.tile(np.arange(x.shape[1], dtype=np.intp), [x.shape[0], 1]) 4576 indarr = xr.DataArray(indarr, dims=ar.dims, coords=ar.coords) 4577 4578 if np.isnan(minindex).any(): 4579 with pytest.raises(ValueError): 4580 ar.argmin(dim="x") 4581 return 4582 4583 expected0 = [ 4584 indarr.isel(y=yi).isel(x=indi, drop=True) 4585 for yi, indi in enumerate(minindex) 4586 ] 4587 expected0 = xr.concat(expected0, dim="y") 4588 4589 result0 = ar.argmin(dim="x") 4590 assert_identical(result0, expected0) 4591 4592 result1 = ar.argmin(axis=1) 4593 assert_identical(result1, expected0) 4594 4595 result2 = ar.argmin(dim="x", keep_attrs=True) 4596 expected1 = expected0.copy() 4597 expected1.attrs = self.attrs 4598 assert_identical(result2, expected1) 4599 4600 minindex = [ 4601 x if y is None or ar.dtype.kind == "O" else y 4602 for x, y in zip(minindex, nanindex) 4603 ] 4604 expected2 = [ 4605 indarr.isel(y=yi).isel(x=indi, drop=True) 4606 for yi, indi in enumerate(minindex) 4607 ] 4608 expected2 = xr.concat(expected2, dim="y") 4609 expected2.attrs = {} 4610 4611 result3 = ar.argmin(dim="x", skipna=False) 4612 4613 assert_identical(result3, expected2) 4614 4615 def test_argmax(self, x, minindex, maxindex, nanindex): 4616 ar = xr.DataArray( 4617 x, 4618 dims=["y", "x"], 4619 coords={"x": np.arange(x.shape[1]) * 4, "y": 1 - np.arange(x.shape[0])}, 4620 attrs=self.attrs, 4621 ) 4622 indarr = np.tile(np.arange(x.shape[1], dtype=np.intp), [x.shape[0], 1]) 4623 indarr = xr.DataArray(indarr, dims=ar.dims, coords=ar.coords) 4624 4625 if np.isnan(maxindex).any(): 4626 with pytest.raises(ValueError): 4627 ar.argmax(dim="x") 4628 return 4629 4630 expected0 = [ 4631 indarr.isel(y=yi).isel(x=indi, drop=True) 4632 for yi, indi in enumerate(maxindex) 4633 ] 4634 expected0 = xr.concat(expected0, dim="y") 4635 4636 result0 = ar.argmax(dim="x") 4637 assert_identical(result0, expected0) 4638 4639 result1 = ar.argmax(axis=1) 4640 assert_identical(result1, expected0) 4641 4642 result2 = ar.argmax(dim="x", keep_attrs=True) 4643 expected1 = expected0.copy() 4644 expected1.attrs = self.attrs 4645 assert_identical(result2, expected1) 4646 4647 maxindex = [ 4648 x if y is None or ar.dtype.kind == "O" else y 4649 for x, y in zip(maxindex, nanindex) 4650 ] 4651 expected2 = [ 4652 indarr.isel(y=yi).isel(x=indi, drop=True) 4653 for yi, indi in enumerate(maxindex) 4654 ] 4655 expected2 = xr.concat(expected2, dim="y") 4656 expected2.attrs = {} 4657 4658 result3 = ar.argmax(dim="x", skipna=False) 4659 4660 assert_identical(result3, expected2) 4661 4662 @pytest.mark.parametrize("use_dask", [True, False]) 4663 def test_idxmin(self, x, minindex, maxindex, nanindex, use_dask): 4664 if use_dask and not has_dask: 4665 pytest.skip("requires dask") 4666 if use_dask and x.dtype.kind == "M": 4667 pytest.xfail("dask operation 'argmin' breaks when dtype is datetime64 (M)") 4668 4669 if x.dtype.kind == "O": 4670 # TODO: nanops._nan_argminmax_object computes once to check for all-NaN slices. 4671 max_computes = 1 4672 else: 4673 max_computes = 0 4674 4675 ar0_raw = xr.DataArray( 4676 x, 4677 dims=["y", "x"], 4678 coords={"x": np.arange(x.shape[1]) * 4, "y": 1 - np.arange(x.shape[0])}, 4679 attrs=self.attrs, 4680 ) 4681 4682 if use_dask: 4683 ar0 = ar0_raw.chunk({}) 4684 else: 4685 ar0 = ar0_raw 4686 4687 assert_identical(ar0, ar0) 4688 4689 # No dimension specified 4690 with pytest.raises(ValueError): 4691 ar0.idxmin() 4692 4693 # dim doesn't exist 4694 with pytest.raises(KeyError): 4695 ar0.idxmin(dim="Y") 4696 4697 assert_identical(ar0, ar0) 4698 4699 coordarr0 = xr.DataArray( 4700 np.tile(ar0.coords["x"], [x.shape[0], 1]), dims=ar0.dims, coords=ar0.coords 4701 ) 4702 4703 hasna = [np.isnan(x) for x in minindex] 4704 coordarr1 = coordarr0.copy() 4705 coordarr1[hasna, :] = 1 4706 minindex0 = [x if not np.isnan(x) else 0 for x in minindex] 4707 4708 nan_mult_0 = np.array([np.NaN if x else 1 for x in hasna])[:, None] 4709 expected0 = [ 4710 (coordarr1 * nan_mult_0).isel(y=yi).isel(x=indi, drop=True) 4711 for yi, indi in enumerate(minindex0) 4712 ] 4713 expected0 = xr.concat(expected0, dim="y") 4714 expected0.name = "x" 4715 4716 # Default fill value (NaN) 4717 with raise_if_dask_computes(max_computes=max_computes): 4718 result0 = ar0.idxmin(dim="x") 4719 assert_identical(result0, expected0) 4720 4721 # Manually specify NaN fill_value 4722 with raise_if_dask_computes(max_computes=max_computes): 4723 result1 = ar0.idxmin(dim="x", fill_value=np.NaN) 4724 assert_identical(result1, expected0) 4725 4726 # keep_attrs 4727 with raise_if_dask_computes(max_computes=max_computes): 4728 result2 = ar0.idxmin(dim="x", keep_attrs=True) 4729 expected2 = expected0.copy() 4730 expected2.attrs = self.attrs 4731 assert_identical(result2, expected2) 4732 4733 # skipna=False 4734 minindex3 = [ 4735 x if y is None or ar0.dtype.kind == "O" else y 4736 for x, y in zip(minindex0, nanindex) 4737 ] 4738 expected3 = [ 4739 coordarr0.isel(y=yi).isel(x=indi, drop=True) 4740 for yi, indi in enumerate(minindex3) 4741 ] 4742 expected3 = xr.concat(expected3, dim="y") 4743 expected3.name = "x" 4744 expected3.attrs = {} 4745 4746 with raise_if_dask_computes(max_computes=max_computes): 4747 result3 = ar0.idxmin(dim="x", skipna=False) 4748 assert_identical(result3, expected3) 4749 4750 # fill_value should be ignored with skipna=False 4751 with raise_if_dask_computes(max_computes=max_computes): 4752 result4 = ar0.idxmin(dim="x", skipna=False, fill_value=-100j) 4753 assert_identical(result4, expected3) 4754 4755 # Float fill_value 4756 nan_mult_5 = np.array([-1.1 if x else 1 for x in hasna])[:, None] 4757 expected5 = [ 4758 (coordarr1 * nan_mult_5).isel(y=yi).isel(x=indi, drop=True) 4759 for yi, indi in enumerate(minindex0) 4760 ] 4761 expected5 = xr.concat(expected5, dim="y") 4762 expected5.name = "x" 4763 4764 with raise_if_dask_computes(max_computes=max_computes): 4765 result5 = ar0.idxmin(dim="x", fill_value=-1.1) 4766 assert_identical(result5, expected5) 4767 4768 # Integer fill_value 4769 nan_mult_6 = np.array([-1 if x else 1 for x in hasna])[:, None] 4770 expected6 = [ 4771 (coordarr1 * nan_mult_6).isel(y=yi).isel(x=indi, drop=True) 4772 for yi, indi in enumerate(minindex0) 4773 ] 4774 expected6 = xr.concat(expected6, dim="y") 4775 expected6.name = "x" 4776 4777 with raise_if_dask_computes(max_computes=max_computes): 4778 result6 = ar0.idxmin(dim="x", fill_value=-1) 4779 assert_identical(result6, expected6) 4780 4781 # Complex fill_value 4782 nan_mult_7 = np.array([-5j if x else 1 for x in hasna])[:, None] 4783 expected7 = [ 4784 (coordarr1 * nan_mult_7).isel(y=yi).isel(x=indi, drop=True) 4785 for yi, indi in enumerate(minindex0) 4786 ] 4787 expected7 = xr.concat(expected7, dim="y") 4788 expected7.name = "x" 4789 4790 with raise_if_dask_computes(max_computes=max_computes): 4791 result7 = ar0.idxmin(dim="x", fill_value=-5j) 4792 assert_identical(result7, expected7) 4793 4794 @pytest.mark.parametrize("use_dask", [True, False]) 4795 def test_idxmax(self, x, minindex, maxindex, nanindex, use_dask): 4796 if use_dask and not has_dask: 4797 pytest.skip("requires dask") 4798 if use_dask and x.dtype.kind == "M": 4799 pytest.xfail("dask operation 'argmax' breaks when dtype is datetime64 (M)") 4800 4801 if x.dtype.kind == "O": 4802 # TODO: nanops._nan_argminmax_object computes once to check for all-NaN slices. 4803 max_computes = 1 4804 else: 4805 max_computes = 0 4806 4807 ar0_raw = xr.DataArray( 4808 x, 4809 dims=["y", "x"], 4810 coords={"x": np.arange(x.shape[1]) * 4, "y": 1 - np.arange(x.shape[0])}, 4811 attrs=self.attrs, 4812 ) 4813 4814 if use_dask: 4815 ar0 = ar0_raw.chunk({}) 4816 else: 4817 ar0 = ar0_raw 4818 4819 # No dimension specified 4820 with pytest.raises(ValueError): 4821 ar0.idxmax() 4822 4823 # dim doesn't exist 4824 with pytest.raises(KeyError): 4825 ar0.idxmax(dim="Y") 4826 4827 ar1 = ar0.copy() 4828 del ar1.coords["y"] 4829 with pytest.raises(KeyError): 4830 ar1.idxmax(dim="y") 4831 4832 coordarr0 = xr.DataArray( 4833 np.tile(ar0.coords["x"], [x.shape[0], 1]), dims=ar0.dims, coords=ar0.coords 4834 ) 4835 4836 hasna = [np.isnan(x) for x in maxindex] 4837 coordarr1 = coordarr0.copy() 4838 coordarr1[hasna, :] = 1 4839 maxindex0 = [x if not np.isnan(x) else 0 for x in maxindex] 4840 4841 nan_mult_0 = np.array([np.NaN if x else 1 for x in hasna])[:, None] 4842 expected0 = [ 4843 (coordarr1 * nan_mult_0).isel(y=yi).isel(x=indi, drop=True) 4844 for yi, indi in enumerate(maxindex0) 4845 ] 4846 expected0 = xr.concat(expected0, dim="y") 4847 expected0.name = "x" 4848 4849 # Default fill value (NaN) 4850 with raise_if_dask_computes(max_computes=max_computes): 4851 result0 = ar0.idxmax(dim="x") 4852 assert_identical(result0, expected0) 4853 4854 # Manually specify NaN fill_value 4855 with raise_if_dask_computes(max_computes=max_computes): 4856 result1 = ar0.idxmax(dim="x", fill_value=np.NaN) 4857 assert_identical(result1, expected0) 4858 4859 # keep_attrs 4860 with raise_if_dask_computes(max_computes=max_computes): 4861 result2 = ar0.idxmax(dim="x", keep_attrs=True) 4862 expected2 = expected0.copy() 4863 expected2.attrs = self.attrs 4864 assert_identical(result2, expected2) 4865 4866 # skipna=False 4867 maxindex3 = [ 4868 x if y is None or ar0.dtype.kind == "O" else y 4869 for x, y in zip(maxindex0, nanindex) 4870 ] 4871 expected3 = [ 4872 coordarr0.isel(y=yi).isel(x=indi, drop=True) 4873 for yi, indi in enumerate(maxindex3) 4874 ] 4875 expected3 = xr.concat(expected3, dim="y") 4876 expected3.name = "x" 4877 expected3.attrs = {} 4878 4879 with raise_if_dask_computes(max_computes=max_computes): 4880 result3 = ar0.idxmax(dim="x", skipna=False) 4881 assert_identical(result3, expected3) 4882 4883 # fill_value should be ignored with skipna=False 4884 with raise_if_dask_computes(max_computes=max_computes): 4885 result4 = ar0.idxmax(dim="x", skipna=False, fill_value=-100j) 4886 assert_identical(result4, expected3) 4887 4888 # Float fill_value 4889 nan_mult_5 = np.array([-1.1 if x else 1 for x in hasna])[:, None] 4890 expected5 = [ 4891 (coordarr1 * nan_mult_5).isel(y=yi).isel(x=indi, drop=True) 4892 for yi, indi in enumerate(maxindex0) 4893 ] 4894 expected5 = xr.concat(expected5, dim="y") 4895 expected5.name = "x" 4896 4897 with raise_if_dask_computes(max_computes=max_computes): 4898 result5 = ar0.idxmax(dim="x", fill_value=-1.1) 4899 assert_identical(result5, expected5) 4900 4901 # Integer fill_value 4902 nan_mult_6 = np.array([-1 if x else 1 for x in hasna])[:, None] 4903 expected6 = [ 4904 (coordarr1 * nan_mult_6).isel(y=yi).isel(x=indi, drop=True) 4905 for yi, indi in enumerate(maxindex0) 4906 ] 4907 expected6 = xr.concat(expected6, dim="y") 4908 expected6.name = "x" 4909 4910 with raise_if_dask_computes(max_computes=max_computes): 4911 result6 = ar0.idxmax(dim="x", fill_value=-1) 4912 assert_identical(result6, expected6) 4913 4914 # Complex fill_value 4915 nan_mult_7 = np.array([-5j if x else 1 for x in hasna])[:, None] 4916 expected7 = [ 4917 (coordarr1 * nan_mult_7).isel(y=yi).isel(x=indi, drop=True) 4918 for yi, indi in enumerate(maxindex0) 4919 ] 4920 expected7 = xr.concat(expected7, dim="y") 4921 expected7.name = "x" 4922 4923 with raise_if_dask_computes(max_computes=max_computes): 4924 result7 = ar0.idxmax(dim="x", fill_value=-5j) 4925 assert_identical(result7, expected7) 4926 4927 @pytest.mark.filterwarnings( 4928 "ignore:Behaviour of argmin/argmax with neither dim nor :DeprecationWarning" 4929 ) 4930 def test_argmin_dim(self, x, minindex, maxindex, nanindex): 4931 ar = xr.DataArray( 4932 x, 4933 dims=["y", "x"], 4934 coords={"x": np.arange(x.shape[1]) * 4, "y": 1 - np.arange(x.shape[0])}, 4935 attrs=self.attrs, 4936 ) 4937 indarr = np.tile(np.arange(x.shape[1], dtype=np.intp), [x.shape[0], 1]) 4938 indarr = xr.DataArray(indarr, dims=ar.dims, coords=ar.coords) 4939 4940 if np.isnan(minindex).any(): 4941 with pytest.raises(ValueError): 4942 ar.argmin(dim="x") 4943 return 4944 4945 expected0 = [ 4946 indarr.isel(y=yi).isel(x=indi, drop=True) 4947 for yi, indi in enumerate(minindex) 4948 ] 4949 expected0 = {"x": xr.concat(expected0, dim="y")} 4950 4951 result0 = ar.argmin(dim=["x"]) 4952 for key in expected0: 4953 assert_identical(result0[key], expected0[key]) 4954 4955 result1 = ar.argmin(dim=["x"], keep_attrs=True) 4956 expected1 = deepcopy(expected0) 4957 expected1["x"].attrs = self.attrs 4958 for key in expected1: 4959 assert_identical(result1[key], expected1[key]) 4960 4961 minindex = [ 4962 x if y is None or ar.dtype.kind == "O" else y 4963 for x, y in zip(minindex, nanindex) 4964 ] 4965 expected2 = [ 4966 indarr.isel(y=yi).isel(x=indi, drop=True) 4967 for yi, indi in enumerate(minindex) 4968 ] 4969 expected2 = {"x": xr.concat(expected2, dim="y")} 4970 expected2["x"].attrs = {} 4971 4972 result2 = ar.argmin(dim=["x"], skipna=False) 4973 4974 for key in expected2: 4975 assert_identical(result2[key], expected2[key]) 4976 4977 result3 = ar.argmin(...) 4978 min_xind = ar.isel(expected0).argmin() 4979 expected3 = { 4980 "y": DataArray(min_xind), 4981 "x": DataArray(minindex[min_xind.item()]), 4982 } 4983 4984 for key in expected3: 4985 assert_identical(result3[key], expected3[key]) 4986 4987 @pytest.mark.filterwarnings( 4988 "ignore:Behaviour of argmin/argmax with neither dim nor :DeprecationWarning" 4989 ) 4990 def test_argmax_dim(self, x, minindex, maxindex, nanindex): 4991 ar = xr.DataArray( 4992 x, 4993 dims=["y", "x"], 4994 coords={"x": np.arange(x.shape[1]) * 4, "y": 1 - np.arange(x.shape[0])}, 4995 attrs=self.attrs, 4996 ) 4997 indarr = np.tile(np.arange(x.shape[1], dtype=np.intp), [x.shape[0], 1]) 4998 indarr = xr.DataArray(indarr, dims=ar.dims, coords=ar.coords) 4999 5000 if np.isnan(maxindex).any(): 5001 with pytest.raises(ValueError): 5002 ar.argmax(dim="x") 5003 return 5004 5005 expected0 = [ 5006 indarr.isel(y=yi).isel(x=indi, drop=True) 5007 for yi, indi in enumerate(maxindex) 5008 ] 5009 expected0 = {"x": xr.concat(expected0, dim="y")} 5010 5011 result0 = ar.argmax(dim=["x"]) 5012 for key in expected0: 5013 assert_identical(result0[key], expected0[key]) 5014 5015 result1 = ar.argmax(dim=["x"], keep_attrs=True) 5016 expected1 = deepcopy(expected0) 5017 expected1["x"].attrs = self.attrs 5018 for key in expected1: 5019 assert_identical(result1[key], expected1[key]) 5020 5021 maxindex = [ 5022 x if y is None or ar.dtype.kind == "O" else y 5023 for x, y in zip(maxindex, nanindex) 5024 ] 5025 expected2 = [ 5026 indarr.isel(y=yi).isel(x=indi, drop=True) 5027 for yi, indi in enumerate(maxindex) 5028 ] 5029 expected2 = {"x": xr.concat(expected2, dim="y")} 5030 expected2["x"].attrs = {} 5031 5032 result2 = ar.argmax(dim=["x"], skipna=False) 5033 5034 for key in expected2: 5035 assert_identical(result2[key], expected2[key]) 5036 5037 result3 = ar.argmax(...) 5038 max_xind = ar.isel(expected0).argmax() 5039 expected3 = { 5040 "y": DataArray(max_xind), 5041 "x": DataArray(maxindex[max_xind.item()]), 5042 } 5043 5044 for key in expected3: 5045 assert_identical(result3[key], expected3[key]) 5046 5047 5048@pytest.mark.parametrize( 5049 "x, minindices_x, minindices_y, minindices_z, minindices_xy, " 5050 "minindices_xz, minindices_yz, minindices_xyz, maxindices_x, " 5051 "maxindices_y, maxindices_z, maxindices_xy, maxindices_xz, maxindices_yz, " 5052 "maxindices_xyz, nanindices_x, nanindices_y, nanindices_z, nanindices_xy, " 5053 "nanindices_xz, nanindices_yz, nanindices_xyz", 5054 [ 5055 ( 5056 np.array( 5057 [ 5058 [[0, 1, 2, 0], [-2, -4, 2, 0]], 5059 [[1, 1, 1, 1], [1, 1, 1, 1]], 5060 [[0, 0, -10, 5], [20, 0, 0, 0]], 5061 ] 5062 ), 5063 {"x": np.array([[0, 2, 2, 0], [0, 0, 2, 0]])}, 5064 {"y": np.array([[1, 1, 0, 0], [0, 0, 0, 0], [0, 0, 0, 1]])}, 5065 {"z": np.array([[0, 1], [0, 0], [2, 1]])}, 5066 {"x": np.array([0, 0, 2, 0]), "y": np.array([1, 1, 0, 0])}, 5067 {"x": np.array([2, 0]), "z": np.array([2, 1])}, 5068 {"y": np.array([1, 0, 0]), "z": np.array([1, 0, 2])}, 5069 {"x": np.array(2), "y": np.array(0), "z": np.array(2)}, 5070 {"x": np.array([[1, 0, 0, 2], [2, 1, 0, 1]])}, 5071 {"y": np.array([[0, 0, 0, 0], [0, 0, 0, 0], [1, 0, 1, 0]])}, 5072 {"z": np.array([[2, 2], [0, 0], [3, 0]])}, 5073 {"x": np.array([2, 0, 0, 2]), "y": np.array([1, 0, 0, 0])}, 5074 {"x": np.array([2, 2]), "z": np.array([3, 0])}, 5075 {"y": np.array([0, 0, 1]), "z": np.array([2, 0, 0])}, 5076 {"x": np.array(2), "y": np.array(1), "z": np.array(0)}, 5077 {"x": np.array([[None, None, None, None], [None, None, None, None]])}, 5078 { 5079 "y": np.array( 5080 [ 5081 [None, None, None, None], 5082 [None, None, None, None], 5083 [None, None, None, None], 5084 ] 5085 ) 5086 }, 5087 {"z": np.array([[None, None], [None, None], [None, None]])}, 5088 { 5089 "x": np.array([None, None, None, None]), 5090 "y": np.array([None, None, None, None]), 5091 }, 5092 {"x": np.array([None, None]), "z": np.array([None, None])}, 5093 {"y": np.array([None, None, None]), "z": np.array([None, None, None])}, 5094 {"x": np.array(None), "y": np.array(None), "z": np.array(None)}, 5095 ), 5096 ( 5097 np.array( 5098 [ 5099 [[2.0, 1.0, 2.0, 0.0], [-2.0, -4.0, 2.0, 0.0]], 5100 [[-4.0, np.NaN, 2.0, np.NaN], [-2.0, -4.0, 2.0, 0.0]], 5101 [[np.NaN] * 4, [np.NaN] * 4], 5102 ] 5103 ), 5104 {"x": np.array([[1, 0, 0, 0], [0, 0, 0, 0]])}, 5105 { 5106 "y": np.array( 5107 [[1, 1, 0, 0], [0, 1, 0, 1], [np.NaN, np.NaN, np.NaN, np.NaN]] 5108 ) 5109 }, 5110 {"z": np.array([[3, 1], [0, 1], [np.NaN, np.NaN]])}, 5111 {"x": np.array([1, 0, 0, 0]), "y": np.array([0, 1, 0, 0])}, 5112 {"x": np.array([1, 0]), "z": np.array([0, 1])}, 5113 {"y": np.array([1, 0, np.NaN]), "z": np.array([1, 0, np.NaN])}, 5114 {"x": np.array(0), "y": np.array(1), "z": np.array(1)}, 5115 {"x": np.array([[0, 0, 0, 0], [0, 0, 0, 0]])}, 5116 { 5117 "y": np.array( 5118 [[0, 0, 0, 0], [1, 1, 0, 1], [np.NaN, np.NaN, np.NaN, np.NaN]] 5119 ) 5120 }, 5121 {"z": np.array([[0, 2], [2, 2], [np.NaN, np.NaN]])}, 5122 {"x": np.array([0, 0, 0, 0]), "y": np.array([0, 0, 0, 0])}, 5123 {"x": np.array([0, 0]), "z": np.array([2, 2])}, 5124 {"y": np.array([0, 0, np.NaN]), "z": np.array([0, 2, np.NaN])}, 5125 {"x": np.array(0), "y": np.array(0), "z": np.array(0)}, 5126 {"x": np.array([[2, 1, 2, 1], [2, 2, 2, 2]])}, 5127 { 5128 "y": np.array( 5129 [[None, None, None, None], [None, 0, None, 0], [0, 0, 0, 0]] 5130 ) 5131 }, 5132 {"z": np.array([[None, None], [1, None], [0, 0]])}, 5133 {"x": np.array([2, 1, 2, 1]), "y": np.array([0, 0, 0, 0])}, 5134 {"x": np.array([1, 2]), "z": np.array([1, 0])}, 5135 {"y": np.array([None, 0, 0]), "z": np.array([None, 1, 0])}, 5136 {"x": np.array(1), "y": np.array(0), "z": np.array(1)}, 5137 ), 5138 ( 5139 np.array( 5140 [ 5141 [[2.0, 1.0, 2.0, 0.0], [-2.0, -4.0, 2.0, 0.0]], 5142 [[-4.0, np.NaN, 2.0, np.NaN], [-2.0, -4.0, 2.0, 0.0]], 5143 [[np.NaN] * 4, [np.NaN] * 4], 5144 ] 5145 ).astype("object"), 5146 {"x": np.array([[1, 0, 0, 0], [0, 0, 0, 0]])}, 5147 { 5148 "y": np.array( 5149 [[1, 1, 0, 0], [0, 1, 0, 1], [np.NaN, np.NaN, np.NaN, np.NaN]] 5150 ) 5151 }, 5152 {"z": np.array([[3, 1], [0, 1], [np.NaN, np.NaN]])}, 5153 {"x": np.array([1, 0, 0, 0]), "y": np.array([0, 1, 0, 0])}, 5154 {"x": np.array([1, 0]), "z": np.array([0, 1])}, 5155 {"y": np.array([1, 0, np.NaN]), "z": np.array([1, 0, np.NaN])}, 5156 {"x": np.array(0), "y": np.array(1), "z": np.array(1)}, 5157 {"x": np.array([[0, 0, 0, 0], [0, 0, 0, 0]])}, 5158 { 5159 "y": np.array( 5160 [[0, 0, 0, 0], [1, 1, 0, 1], [np.NaN, np.NaN, np.NaN, np.NaN]] 5161 ) 5162 }, 5163 {"z": np.array([[0, 2], [2, 2], [np.NaN, np.NaN]])}, 5164 {"x": np.array([0, 0, 0, 0]), "y": np.array([0, 0, 0, 0])}, 5165 {"x": np.array([0, 0]), "z": np.array([2, 2])}, 5166 {"y": np.array([0, 0, np.NaN]), "z": np.array([0, 2, np.NaN])}, 5167 {"x": np.array(0), "y": np.array(0), "z": np.array(0)}, 5168 {"x": np.array([[2, 1, 2, 1], [2, 2, 2, 2]])}, 5169 { 5170 "y": np.array( 5171 [[None, None, None, None], [None, 0, None, 0], [0, 0, 0, 0]] 5172 ) 5173 }, 5174 {"z": np.array([[None, None], [1, None], [0, 0]])}, 5175 {"x": np.array([2, 1, 2, 1]), "y": np.array([0, 0, 0, 0])}, 5176 {"x": np.array([1, 2]), "z": np.array([1, 0])}, 5177 {"y": np.array([None, 0, 0]), "z": np.array([None, 1, 0])}, 5178 {"x": np.array(1), "y": np.array(0), "z": np.array(1)}, 5179 ), 5180 ( 5181 np.array( 5182 [ 5183 [["2015-12-31", "2020-01-02"], ["2020-01-01", "2016-01-01"]], 5184 [["2020-01-02", "2020-01-02"], ["2020-01-02", "2020-01-02"]], 5185 [["1900-01-01", "1-02-03"], ["1900-01-02", "1-02-03"]], 5186 ], 5187 dtype="datetime64[ns]", 5188 ), 5189 {"x": np.array([[2, 2], [2, 2]])}, 5190 {"y": np.array([[0, 1], [0, 0], [0, 0]])}, 5191 {"z": np.array([[0, 1], [0, 0], [1, 1]])}, 5192 {"x": np.array([2, 2]), "y": np.array([0, 0])}, 5193 {"x": np.array([2, 2]), "z": np.array([1, 1])}, 5194 {"y": np.array([0, 0, 0]), "z": np.array([0, 0, 1])}, 5195 {"x": np.array(2), "y": np.array(0), "z": np.array(1)}, 5196 {"x": np.array([[1, 0], [1, 1]])}, 5197 {"y": np.array([[1, 0], [0, 0], [1, 0]])}, 5198 {"z": np.array([[1, 0], [0, 0], [0, 0]])}, 5199 {"x": np.array([1, 0]), "y": np.array([0, 0])}, 5200 {"x": np.array([0, 1]), "z": np.array([1, 0])}, 5201 {"y": np.array([0, 0, 1]), "z": np.array([1, 0, 0])}, 5202 {"x": np.array(0), "y": np.array(0), "z": np.array(1)}, 5203 {"x": np.array([[None, None], [None, None]])}, 5204 {"y": np.array([[None, None], [None, None], [None, None]])}, 5205 {"z": np.array([[None, None], [None, None], [None, None]])}, 5206 {"x": np.array([None, None]), "y": np.array([None, None])}, 5207 {"x": np.array([None, None]), "z": np.array([None, None])}, 5208 {"y": np.array([None, None, None]), "z": np.array([None, None, None])}, 5209 {"x": np.array(None), "y": np.array(None), "z": np.array(None)}, 5210 ), 5211 ], 5212) 5213class TestReduce3D(TestReduce): 5214 def test_argmin_dim( 5215 self, 5216 x, 5217 minindices_x, 5218 minindices_y, 5219 minindices_z, 5220 minindices_xy, 5221 minindices_xz, 5222 minindices_yz, 5223 minindices_xyz, 5224 maxindices_x, 5225 maxindices_y, 5226 maxindices_z, 5227 maxindices_xy, 5228 maxindices_xz, 5229 maxindices_yz, 5230 maxindices_xyz, 5231 nanindices_x, 5232 nanindices_y, 5233 nanindices_z, 5234 nanindices_xy, 5235 nanindices_xz, 5236 nanindices_yz, 5237 nanindices_xyz, 5238 ): 5239 5240 ar = xr.DataArray( 5241 x, 5242 dims=["x", "y", "z"], 5243 coords={ 5244 "x": np.arange(x.shape[0]) * 4, 5245 "y": 1 - np.arange(x.shape[1]), 5246 "z": 2 + 3 * np.arange(x.shape[2]), 5247 }, 5248 attrs=self.attrs, 5249 ) 5250 xindarr = np.tile( 5251 np.arange(x.shape[0], dtype=np.intp)[:, np.newaxis, np.newaxis], 5252 [1, x.shape[1], x.shape[2]], 5253 ) 5254 xindarr = xr.DataArray(xindarr, dims=ar.dims, coords=ar.coords) 5255 yindarr = np.tile( 5256 np.arange(x.shape[1], dtype=np.intp)[np.newaxis, :, np.newaxis], 5257 [x.shape[0], 1, x.shape[2]], 5258 ) 5259 yindarr = xr.DataArray(yindarr, dims=ar.dims, coords=ar.coords) 5260 zindarr = np.tile( 5261 np.arange(x.shape[2], dtype=np.intp)[np.newaxis, np.newaxis, :], 5262 [x.shape[0], x.shape[1], 1], 5263 ) 5264 zindarr = xr.DataArray(zindarr, dims=ar.dims, coords=ar.coords) 5265 5266 for inds in [ 5267 minindices_x, 5268 minindices_y, 5269 minindices_z, 5270 minindices_xy, 5271 minindices_xz, 5272 minindices_yz, 5273 minindices_xyz, 5274 ]: 5275 if np.array([np.isnan(i) for i in inds.values()]).any(): 5276 with pytest.raises(ValueError): 5277 ar.argmin(dim=[d for d in inds]) 5278 return 5279 5280 result0 = ar.argmin(dim=["x"]) 5281 expected0 = { 5282 key: xr.DataArray(value, dims=("y", "z")) 5283 for key, value in minindices_x.items() 5284 } 5285 for key in expected0: 5286 assert_identical(result0[key].drop_vars(["y", "z"]), expected0[key]) 5287 5288 result1 = ar.argmin(dim=["y"]) 5289 expected1 = { 5290 key: xr.DataArray(value, dims=("x", "z")) 5291 for key, value in minindices_y.items() 5292 } 5293 for key in expected1: 5294 assert_identical(result1[key].drop_vars(["x", "z"]), expected1[key]) 5295 5296 result2 = ar.argmin(dim=["z"]) 5297 expected2 = { 5298 key: xr.DataArray(value, dims=("x", "y")) 5299 for key, value in minindices_z.items() 5300 } 5301 for key in expected2: 5302 assert_identical(result2[key].drop_vars(["x", "y"]), expected2[key]) 5303 5304 result3 = ar.argmin(dim=("x", "y")) 5305 expected3 = { 5306 key: xr.DataArray(value, dims=("z")) for key, value in minindices_xy.items() 5307 } 5308 for key in expected3: 5309 assert_identical(result3[key].drop_vars("z"), expected3[key]) 5310 5311 result4 = ar.argmin(dim=("x", "z")) 5312 expected4 = { 5313 key: xr.DataArray(value, dims=("y")) for key, value in minindices_xz.items() 5314 } 5315 for key in expected4: 5316 assert_identical(result4[key].drop_vars("y"), expected4[key]) 5317 5318 result5 = ar.argmin(dim=("y", "z")) 5319 expected5 = { 5320 key: xr.DataArray(value, dims=("x")) for key, value in minindices_yz.items() 5321 } 5322 for key in expected5: 5323 assert_identical(result5[key].drop_vars("x"), expected5[key]) 5324 5325 result6 = ar.argmin(...) 5326 expected6 = {key: xr.DataArray(value) for key, value in minindices_xyz.items()} 5327 for key in expected6: 5328 assert_identical(result6[key], expected6[key]) 5329 5330 minindices_x = { 5331 key: xr.where( 5332 nanindices_x[key] == None, # noqa: E711 5333 minindices_x[key], 5334 nanindices_x[key], 5335 ) 5336 for key in minindices_x 5337 } 5338 expected7 = { 5339 key: xr.DataArray(value, dims=("y", "z")) 5340 for key, value in minindices_x.items() 5341 } 5342 5343 result7 = ar.argmin(dim=["x"], skipna=False) 5344 for key in expected7: 5345 assert_identical(result7[key].drop_vars(["y", "z"]), expected7[key]) 5346 5347 minindices_y = { 5348 key: xr.where( 5349 nanindices_y[key] == None, # noqa: E711 5350 minindices_y[key], 5351 nanindices_y[key], 5352 ) 5353 for key in minindices_y 5354 } 5355 expected8 = { 5356 key: xr.DataArray(value, dims=("x", "z")) 5357 for key, value in minindices_y.items() 5358 } 5359 5360 result8 = ar.argmin(dim=["y"], skipna=False) 5361 for key in expected8: 5362 assert_identical(result8[key].drop_vars(["x", "z"]), expected8[key]) 5363 5364 minindices_z = { 5365 key: xr.where( 5366 nanindices_z[key] == None, # noqa: E711 5367 minindices_z[key], 5368 nanindices_z[key], 5369 ) 5370 for key in minindices_z 5371 } 5372 expected9 = { 5373 key: xr.DataArray(value, dims=("x", "y")) 5374 for key, value in minindices_z.items() 5375 } 5376 5377 result9 = ar.argmin(dim=["z"], skipna=False) 5378 for key in expected9: 5379 assert_identical(result9[key].drop_vars(["x", "y"]), expected9[key]) 5380 5381 minindices_xy = { 5382 key: xr.where( 5383 nanindices_xy[key] == None, # noqa: E711 5384 minindices_xy[key], 5385 nanindices_xy[key], 5386 ) 5387 for key in minindices_xy 5388 } 5389 expected10 = { 5390 key: xr.DataArray(value, dims="z") for key, value in minindices_xy.items() 5391 } 5392 5393 result10 = ar.argmin(dim=("x", "y"), skipna=False) 5394 for key in expected10: 5395 assert_identical(result10[key].drop_vars("z"), expected10[key]) 5396 5397 minindices_xz = { 5398 key: xr.where( 5399 nanindices_xz[key] == None, # noqa: E711 5400 minindices_xz[key], 5401 nanindices_xz[key], 5402 ) 5403 for key in minindices_xz 5404 } 5405 expected11 = { 5406 key: xr.DataArray(value, dims="y") for key, value in minindices_xz.items() 5407 } 5408 5409 result11 = ar.argmin(dim=("x", "z"), skipna=False) 5410 for key in expected11: 5411 assert_identical(result11[key].drop_vars("y"), expected11[key]) 5412 5413 minindices_yz = { 5414 key: xr.where( 5415 nanindices_yz[key] == None, # noqa: E711 5416 minindices_yz[key], 5417 nanindices_yz[key], 5418 ) 5419 for key in minindices_yz 5420 } 5421 expected12 = { 5422 key: xr.DataArray(value, dims="x") for key, value in minindices_yz.items() 5423 } 5424 5425 result12 = ar.argmin(dim=("y", "z"), skipna=False) 5426 for key in expected12: 5427 assert_identical(result12[key].drop_vars("x"), expected12[key]) 5428 5429 minindices_xyz = { 5430 key: xr.where( 5431 nanindices_xyz[key] == None, # noqa: E711 5432 minindices_xyz[key], 5433 nanindices_xyz[key], 5434 ) 5435 for key in minindices_xyz 5436 } 5437 expected13 = {key: xr.DataArray(value) for key, value in minindices_xyz.items()} 5438 5439 result13 = ar.argmin(..., skipna=False) 5440 for key in expected13: 5441 assert_identical(result13[key], expected13[key]) 5442 5443 def test_argmax_dim( 5444 self, 5445 x, 5446 minindices_x, 5447 minindices_y, 5448 minindices_z, 5449 minindices_xy, 5450 minindices_xz, 5451 minindices_yz, 5452 minindices_xyz, 5453 maxindices_x, 5454 maxindices_y, 5455 maxindices_z, 5456 maxindices_xy, 5457 maxindices_xz, 5458 maxindices_yz, 5459 maxindices_xyz, 5460 nanindices_x, 5461 nanindices_y, 5462 nanindices_z, 5463 nanindices_xy, 5464 nanindices_xz, 5465 nanindices_yz, 5466 nanindices_xyz, 5467 ): 5468 5469 ar = xr.DataArray( 5470 x, 5471 dims=["x", "y", "z"], 5472 coords={ 5473 "x": np.arange(x.shape[0]) * 4, 5474 "y": 1 - np.arange(x.shape[1]), 5475 "z": 2 + 3 * np.arange(x.shape[2]), 5476 }, 5477 attrs=self.attrs, 5478 ) 5479 xindarr = np.tile( 5480 np.arange(x.shape[0], dtype=np.intp)[:, np.newaxis, np.newaxis], 5481 [1, x.shape[1], x.shape[2]], 5482 ) 5483 xindarr = xr.DataArray(xindarr, dims=ar.dims, coords=ar.coords) 5484 yindarr = np.tile( 5485 np.arange(x.shape[1], dtype=np.intp)[np.newaxis, :, np.newaxis], 5486 [x.shape[0], 1, x.shape[2]], 5487 ) 5488 yindarr = xr.DataArray(yindarr, dims=ar.dims, coords=ar.coords) 5489 zindarr = np.tile( 5490 np.arange(x.shape[2], dtype=np.intp)[np.newaxis, np.newaxis, :], 5491 [x.shape[0], x.shape[1], 1], 5492 ) 5493 zindarr = xr.DataArray(zindarr, dims=ar.dims, coords=ar.coords) 5494 5495 for inds in [ 5496 maxindices_x, 5497 maxindices_y, 5498 maxindices_z, 5499 maxindices_xy, 5500 maxindices_xz, 5501 maxindices_yz, 5502 maxindices_xyz, 5503 ]: 5504 if np.array([np.isnan(i) for i in inds.values()]).any(): 5505 with pytest.raises(ValueError): 5506 ar.argmax(dim=[d for d in inds]) 5507 return 5508 5509 result0 = ar.argmax(dim=["x"]) 5510 expected0 = { 5511 key: xr.DataArray(value, dims=("y", "z")) 5512 for key, value in maxindices_x.items() 5513 } 5514 for key in expected0: 5515 assert_identical(result0[key].drop_vars(["y", "z"]), expected0[key]) 5516 5517 result1 = ar.argmax(dim=["y"]) 5518 expected1 = { 5519 key: xr.DataArray(value, dims=("x", "z")) 5520 for key, value in maxindices_y.items() 5521 } 5522 for key in expected1: 5523 assert_identical(result1[key].drop_vars(["x", "z"]), expected1[key]) 5524 5525 result2 = ar.argmax(dim=["z"]) 5526 expected2 = { 5527 key: xr.DataArray(value, dims=("x", "y")) 5528 for key, value in maxindices_z.items() 5529 } 5530 for key in expected2: 5531 assert_identical(result2[key].drop_vars(["x", "y"]), expected2[key]) 5532 5533 result3 = ar.argmax(dim=("x", "y")) 5534 expected3 = { 5535 key: xr.DataArray(value, dims=("z")) for key, value in maxindices_xy.items() 5536 } 5537 for key in expected3: 5538 assert_identical(result3[key].drop_vars("z"), expected3[key]) 5539 5540 result4 = ar.argmax(dim=("x", "z")) 5541 expected4 = { 5542 key: xr.DataArray(value, dims=("y")) for key, value in maxindices_xz.items() 5543 } 5544 for key in expected4: 5545 assert_identical(result4[key].drop_vars("y"), expected4[key]) 5546 5547 result5 = ar.argmax(dim=("y", "z")) 5548 expected5 = { 5549 key: xr.DataArray(value, dims=("x")) for key, value in maxindices_yz.items() 5550 } 5551 for key in expected5: 5552 assert_identical(result5[key].drop_vars("x"), expected5[key]) 5553 5554 result6 = ar.argmax(...) 5555 expected6 = {key: xr.DataArray(value) for key, value in maxindices_xyz.items()} 5556 for key in expected6: 5557 assert_identical(result6[key], expected6[key]) 5558 5559 maxindices_x = { 5560 key: xr.where( 5561 nanindices_x[key] == None, # noqa: E711 5562 maxindices_x[key], 5563 nanindices_x[key], 5564 ) 5565 for key in maxindices_x 5566 } 5567 expected7 = { 5568 key: xr.DataArray(value, dims=("y", "z")) 5569 for key, value in maxindices_x.items() 5570 } 5571 5572 result7 = ar.argmax(dim=["x"], skipna=False) 5573 for key in expected7: 5574 assert_identical(result7[key].drop_vars(["y", "z"]), expected7[key]) 5575 5576 maxindices_y = { 5577 key: xr.where( 5578 nanindices_y[key] == None, # noqa: E711 5579 maxindices_y[key], 5580 nanindices_y[key], 5581 ) 5582 for key in maxindices_y 5583 } 5584 expected8 = { 5585 key: xr.DataArray(value, dims=("x", "z")) 5586 for key, value in maxindices_y.items() 5587 } 5588 5589 result8 = ar.argmax(dim=["y"], skipna=False) 5590 for key in expected8: 5591 assert_identical(result8[key].drop_vars(["x", "z"]), expected8[key]) 5592 5593 maxindices_z = { 5594 key: xr.where( 5595 nanindices_z[key] == None, # noqa: E711 5596 maxindices_z[key], 5597 nanindices_z[key], 5598 ) 5599 for key in maxindices_z 5600 } 5601 expected9 = { 5602 key: xr.DataArray(value, dims=("x", "y")) 5603 for key, value in maxindices_z.items() 5604 } 5605 5606 result9 = ar.argmax(dim=["z"], skipna=False) 5607 for key in expected9: 5608 assert_identical(result9[key].drop_vars(["x", "y"]), expected9[key]) 5609 5610 maxindices_xy = { 5611 key: xr.where( 5612 nanindices_xy[key] == None, # noqa: E711 5613 maxindices_xy[key], 5614 nanindices_xy[key], 5615 ) 5616 for key in maxindices_xy 5617 } 5618 expected10 = { 5619 key: xr.DataArray(value, dims="z") for key, value in maxindices_xy.items() 5620 } 5621 5622 result10 = ar.argmax(dim=("x", "y"), skipna=False) 5623 for key in expected10: 5624 assert_identical(result10[key].drop_vars("z"), expected10[key]) 5625 5626 maxindices_xz = { 5627 key: xr.where( 5628 nanindices_xz[key] == None, # noqa: E711 5629 maxindices_xz[key], 5630 nanindices_xz[key], 5631 ) 5632 for key in maxindices_xz 5633 } 5634 expected11 = { 5635 key: xr.DataArray(value, dims="y") for key, value in maxindices_xz.items() 5636 } 5637 5638 result11 = ar.argmax(dim=("x", "z"), skipna=False) 5639 for key in expected11: 5640 assert_identical(result11[key].drop_vars("y"), expected11[key]) 5641 5642 maxindices_yz = { 5643 key: xr.where( 5644 nanindices_yz[key] == None, # noqa: E711 5645 maxindices_yz[key], 5646 nanindices_yz[key], 5647 ) 5648 for key in maxindices_yz 5649 } 5650 expected12 = { 5651 key: xr.DataArray(value, dims="x") for key, value in maxindices_yz.items() 5652 } 5653 5654 result12 = ar.argmax(dim=("y", "z"), skipna=False) 5655 for key in expected12: 5656 assert_identical(result12[key].drop_vars("x"), expected12[key]) 5657 5658 maxindices_xyz = { 5659 key: xr.where( 5660 nanindices_xyz[key] == None, # noqa: E711 5661 maxindices_xyz[key], 5662 nanindices_xyz[key], 5663 ) 5664 for key in maxindices_xyz 5665 } 5666 expected13 = {key: xr.DataArray(value) for key, value in maxindices_xyz.items()} 5667 5668 result13 = ar.argmax(..., skipna=False) 5669 for key in expected13: 5670 assert_identical(result13[key], expected13[key]) 5671 5672 5673class TestReduceND(TestReduce): 5674 @pytest.mark.parametrize("op", ["idxmin", "idxmax"]) 5675 @pytest.mark.parametrize("ndim", [3, 5]) 5676 def test_idxminmax_dask(self, op, ndim): 5677 if not has_dask: 5678 pytest.skip("requires dask") 5679 5680 ar0_raw = xr.DataArray( 5681 np.random.random_sample(size=[10] * ndim), 5682 dims=[i for i in "abcdefghij"[: ndim - 1]] + ["x"], 5683 coords={"x": np.arange(10)}, 5684 attrs=self.attrs, 5685 ) 5686 5687 ar0_dsk = ar0_raw.chunk({}) 5688 # Assert idx is the same with dask and without 5689 assert_equal(getattr(ar0_dsk, op)(dim="x"), getattr(ar0_raw, op)(dim="x")) 5690 5691 5692@pytest.fixture(params=[1]) 5693def da(request, backend): 5694 if request.param == 1: 5695 times = pd.date_range("2000-01-01", freq="1D", periods=21) 5696 da = DataArray( 5697 np.random.random((3, 21, 4)), 5698 dims=("a", "time", "x"), 5699 coords=dict(time=times), 5700 ) 5701 5702 if request.param == 2: 5703 da = DataArray([0, np.nan, 1, 2, np.nan, 3, 4, 5, np.nan, 6, 7], dims="time") 5704 5705 if request.param == "repeating_ints": 5706 da = DataArray( 5707 np.tile(np.arange(12), 5).reshape(5, 4, 3), 5708 coords={"x": list("abc"), "y": list("defg")}, 5709 dims=list("zyx"), 5710 ) 5711 5712 if backend == "dask": 5713 return da.chunk() 5714 elif backend == "numpy": 5715 return da 5716 else: 5717 raise ValueError 5718 5719 5720@pytest.mark.parametrize("da", ("repeating_ints",), indirect=True) 5721def test_isin(da): 5722 expected = DataArray( 5723 np.asarray([[0, 0, 0], [1, 0, 0]]), 5724 dims=list("yx"), 5725 coords={"x": list("abc"), "y": list("de")}, 5726 ).astype("bool") 5727 5728 result = da.isin([3]).sel(y=list("de"), z=0) 5729 assert_equal(result, expected) 5730 5731 expected = DataArray( 5732 np.asarray([[0, 0, 1], [1, 0, 0]]), 5733 dims=list("yx"), 5734 coords={"x": list("abc"), "y": list("de")}, 5735 ).astype("bool") 5736 result = da.isin([2, 3]).sel(y=list("de"), z=0) 5737 assert_equal(result, expected) 5738 5739 5740@pytest.mark.parametrize("da", (1, 2), indirect=True) 5741def test_rolling_iter(da): 5742 rolling_obj = da.rolling(time=7) 5743 rolling_obj_mean = rolling_obj.mean() 5744 5745 assert len(rolling_obj.window_labels) == len(da["time"]) 5746 assert_identical(rolling_obj.window_labels, da["time"]) 5747 5748 for i, (label, window_da) in enumerate(rolling_obj): 5749 assert label == da["time"].isel(time=i) 5750 5751 actual = rolling_obj_mean.isel(time=i) 5752 expected = window_da.mean("time") 5753 5754 # TODO add assert_allclose_with_nan, which compares nan position 5755 # as well as the closeness of the values. 5756 assert_array_equal(actual.isnull(), expected.isnull()) 5757 if (~actual.isnull()).sum() > 0: 5758 np.allclose( 5759 actual.values[actual.values.nonzero()], 5760 expected.values[expected.values.nonzero()], 5761 ) 5762 5763 5764@pytest.mark.parametrize("da", (1,), indirect=True) 5765def test_rolling_repr(da): 5766 rolling_obj = da.rolling(time=7) 5767 assert repr(rolling_obj) == "DataArrayRolling [time->7]" 5768 rolling_obj = da.rolling(time=7, center=True) 5769 assert repr(rolling_obj) == "DataArrayRolling [time->7(center)]" 5770 rolling_obj = da.rolling(time=7, x=3, center=True) 5771 assert repr(rolling_obj) == "DataArrayRolling [time->7(center),x->3(center)]" 5772 5773 5774@requires_dask 5775def test_repeated_rolling_rechunks(): 5776 5777 # regression test for GH3277, GH2514 5778 dat = DataArray(np.random.rand(7653, 300), dims=("day", "item")) 5779 dat_chunk = dat.chunk({"item": 20}) 5780 dat_chunk.rolling(day=10).mean().rolling(day=250).std() 5781 5782 5783def test_rolling_doc(da): 5784 rolling_obj = da.rolling(time=7) 5785 5786 # argument substitution worked 5787 assert "`mean`" in rolling_obj.mean.__doc__ 5788 5789 5790def test_rolling_properties(da): 5791 rolling_obj = da.rolling(time=4) 5792 5793 assert rolling_obj.obj.get_axis_num("time") == 1 5794 5795 # catching invalid args 5796 with pytest.raises(ValueError, match="window must be > 0"): 5797 da.rolling(time=-2) 5798 5799 with pytest.raises(ValueError, match="min_periods must be greater than zero"): 5800 da.rolling(time=2, min_periods=0) 5801 5802 5803@pytest.mark.parametrize("name", ("sum", "mean", "std", "min", "max", "median")) 5804@pytest.mark.parametrize("center", (True, False, None)) 5805@pytest.mark.parametrize("min_periods", (1, None)) 5806@pytest.mark.parametrize("backend", ["numpy"], indirect=True) 5807def test_rolling_wrapped_bottleneck(da, name, center, min_periods): 5808 bn = pytest.importorskip("bottleneck", minversion="1.1") 5809 5810 # Test all bottleneck functions 5811 rolling_obj = da.rolling(time=7, min_periods=min_periods) 5812 5813 func_name = f"move_{name}" 5814 actual = getattr(rolling_obj, name)() 5815 expected = getattr(bn, func_name)( 5816 da.values, window=7, axis=1, min_count=min_periods 5817 ) 5818 assert_array_equal(actual.values, expected) 5819 5820 with pytest.warns(DeprecationWarning, match="Reductions are applied"): 5821 getattr(rolling_obj, name)(dim="time") 5822 5823 # Test center 5824 rolling_obj = da.rolling(time=7, center=center) 5825 actual = getattr(rolling_obj, name)()["time"] 5826 assert_equal(actual, da["time"]) 5827 5828 5829@requires_dask 5830@pytest.mark.parametrize("name", ("mean", "count")) 5831@pytest.mark.parametrize("center", (True, False, None)) 5832@pytest.mark.parametrize("min_periods", (1, None)) 5833@pytest.mark.parametrize("window", (7, 8)) 5834@pytest.mark.parametrize("backend", ["dask"], indirect=True) 5835def test_rolling_wrapped_dask(da, name, center, min_periods, window): 5836 # dask version 5837 rolling_obj = da.rolling(time=window, min_periods=min_periods, center=center) 5838 actual = getattr(rolling_obj, name)().load() 5839 if name != "count": 5840 with pytest.warns(DeprecationWarning, match="Reductions are applied"): 5841 getattr(rolling_obj, name)(dim="time") 5842 # numpy version 5843 rolling_obj = da.load().rolling(time=window, min_periods=min_periods, center=center) 5844 expected = getattr(rolling_obj, name)() 5845 5846 # using all-close because rolling over ghost cells introduces some 5847 # precision errors 5848 assert_allclose(actual, expected) 5849 5850 # with zero chunked array GH:2113 5851 rolling_obj = da.chunk().rolling( 5852 time=window, min_periods=min_periods, center=center 5853 ) 5854 actual = getattr(rolling_obj, name)().load() 5855 assert_allclose(actual, expected) 5856 5857 5858@pytest.mark.parametrize("center", (True, None)) 5859def test_rolling_wrapped_dask_nochunk(center): 5860 # GH:2113 5861 pytest.importorskip("dask.array") 5862 5863 da_day_clim = xr.DataArray( 5864 np.arange(1, 367), coords=[np.arange(1, 367)], dims="dayofyear" 5865 ) 5866 expected = da_day_clim.rolling(dayofyear=31, center=center).mean() 5867 actual = da_day_clim.chunk().rolling(dayofyear=31, center=center).mean() 5868 assert_allclose(actual, expected) 5869 5870 5871@pytest.mark.parametrize("center", (True, False)) 5872@pytest.mark.parametrize("min_periods", (None, 1, 2, 3)) 5873@pytest.mark.parametrize("window", (1, 2, 3, 4)) 5874def test_rolling_pandas_compat(center, window, min_periods): 5875 s = pd.Series(np.arange(10)) 5876 da = DataArray.from_series(s) 5877 5878 if min_periods is not None and window < min_periods: 5879 min_periods = window 5880 5881 s_rolling = s.rolling(window, center=center, min_periods=min_periods).mean() 5882 da_rolling = da.rolling(index=window, center=center, min_periods=min_periods).mean() 5883 da_rolling_np = da.rolling( 5884 index=window, center=center, min_periods=min_periods 5885 ).reduce(np.nanmean) 5886 5887 np.testing.assert_allclose(s_rolling.values, da_rolling.values) 5888 np.testing.assert_allclose(s_rolling.index, da_rolling["index"]) 5889 np.testing.assert_allclose(s_rolling.values, da_rolling_np.values) 5890 np.testing.assert_allclose(s_rolling.index, da_rolling_np["index"]) 5891 5892 5893@pytest.mark.parametrize("center", (True, False)) 5894@pytest.mark.parametrize("window", (1, 2, 3, 4)) 5895def test_rolling_construct(center, window): 5896 s = pd.Series(np.arange(10)) 5897 da = DataArray.from_series(s) 5898 5899 s_rolling = s.rolling(window, center=center, min_periods=1).mean() 5900 da_rolling = da.rolling(index=window, center=center, min_periods=1) 5901 5902 da_rolling_mean = da_rolling.construct("window").mean("window") 5903 np.testing.assert_allclose(s_rolling.values, da_rolling_mean.values) 5904 np.testing.assert_allclose(s_rolling.index, da_rolling_mean["index"]) 5905 5906 # with stride 5907 da_rolling_mean = da_rolling.construct("window", stride=2).mean("window") 5908 np.testing.assert_allclose(s_rolling.values[::2], da_rolling_mean.values) 5909 np.testing.assert_allclose(s_rolling.index[::2], da_rolling_mean["index"]) 5910 5911 # with fill_value 5912 da_rolling_mean = da_rolling.construct("window", stride=2, fill_value=0.0).mean( 5913 "window" 5914 ) 5915 assert da_rolling_mean.isnull().sum() == 0 5916 assert (da_rolling_mean == 0.0).sum() >= 0 5917 5918 5919@pytest.mark.parametrize("da", (1, 2), indirect=True) 5920@pytest.mark.parametrize("center", (True, False)) 5921@pytest.mark.parametrize("min_periods", (None, 1, 2, 3)) 5922@pytest.mark.parametrize("window", (1, 2, 3, 4)) 5923@pytest.mark.parametrize("name", ("sum", "mean", "std", "max")) 5924def test_rolling_reduce(da, center, min_periods, window, name): 5925 if min_periods is not None and window < min_periods: 5926 min_periods = window 5927 5928 if da.isnull().sum() > 1 and window == 1: 5929 # this causes all nan slices 5930 window = 2 5931 5932 rolling_obj = da.rolling(time=window, center=center, min_periods=min_periods) 5933 5934 # add nan prefix to numpy methods to get similar # behavior as bottleneck 5935 actual = rolling_obj.reduce(getattr(np, "nan%s" % name)) 5936 expected = getattr(rolling_obj, name)() 5937 assert_allclose(actual, expected) 5938 assert actual.dims == expected.dims 5939 5940 5941@pytest.mark.parametrize("center", (True, False)) 5942@pytest.mark.parametrize("min_periods", (None, 1, 2, 3)) 5943@pytest.mark.parametrize("window", (1, 2, 3, 4)) 5944@pytest.mark.parametrize("name", ("sum", "max")) 5945def test_rolling_reduce_nonnumeric(center, min_periods, window, name): 5946 da = DataArray( 5947 [0, np.nan, 1, 2, np.nan, 3, 4, 5, np.nan, 6, 7], dims="time" 5948 ).isnull() 5949 5950 if min_periods is not None and window < min_periods: 5951 min_periods = window 5952 5953 rolling_obj = da.rolling(time=window, center=center, min_periods=min_periods) 5954 5955 # add nan prefix to numpy methods to get similar behavior as bottleneck 5956 actual = rolling_obj.reduce(getattr(np, "nan%s" % name)) 5957 expected = getattr(rolling_obj, name)() 5958 assert_allclose(actual, expected) 5959 assert actual.dims == expected.dims 5960 5961 5962def test_rolling_count_correct(): 5963 da = DataArray([0, np.nan, 1, 2, np.nan, 3, 4, 5, np.nan, 6, 7], dims="time") 5964 5965 kwargs = [ 5966 {"time": 11, "min_periods": 1}, 5967 {"time": 11, "min_periods": None}, 5968 {"time": 7, "min_periods": 2}, 5969 ] 5970 expecteds = [ 5971 DataArray([1, 1, 2, 3, 3, 4, 5, 6, 6, 7, 8], dims="time"), 5972 DataArray( 5973 [ 5974 np.nan, 5975 np.nan, 5976 np.nan, 5977 np.nan, 5978 np.nan, 5979 np.nan, 5980 np.nan, 5981 np.nan, 5982 np.nan, 5983 np.nan, 5984 np.nan, 5985 ], 5986 dims="time", 5987 ), 5988 DataArray([np.nan, np.nan, 2, 3, 3, 4, 5, 5, 5, 5, 5], dims="time"), 5989 ] 5990 5991 for kwarg, expected in zip(kwargs, expecteds): 5992 result = da.rolling(**kwarg).count() 5993 assert_equal(result, expected) 5994 5995 result = da.to_dataset(name="var1").rolling(**kwarg).count()["var1"] 5996 assert_equal(result, expected) 5997 5998 5999@pytest.mark.parametrize("da", (1,), indirect=True) 6000@pytest.mark.parametrize("center", (True, False)) 6001@pytest.mark.parametrize("min_periods", (None, 1)) 6002@pytest.mark.parametrize("name", ("sum", "mean", "max")) 6003def test_ndrolling_reduce(da, center, min_periods, name): 6004 rolling_obj = da.rolling(time=3, x=2, center=center, min_periods=min_periods) 6005 6006 actual = getattr(rolling_obj, name)() 6007 expected = getattr( 6008 getattr( 6009 da.rolling(time=3, center=center, min_periods=min_periods), name 6010 )().rolling(x=2, center=center, min_periods=min_periods), 6011 name, 6012 )() 6013 6014 assert_allclose(actual, expected) 6015 assert actual.dims == expected.dims 6016 6017 if name in ["mean"]: 6018 # test our reimplementation of nanmean using np.nanmean 6019 expected = getattr(rolling_obj.construct({"time": "tw", "x": "xw"}), name)( 6020 ["tw", "xw"] 6021 ) 6022 count = rolling_obj.count() 6023 if min_periods is None: 6024 min_periods = 1 6025 assert_allclose(actual, expected.where(count >= min_periods)) 6026 6027 6028@pytest.mark.parametrize("center", (True, False, (True, False))) 6029@pytest.mark.parametrize("fill_value", (np.nan, 0.0)) 6030def test_ndrolling_construct(center, fill_value): 6031 da = DataArray( 6032 np.arange(5 * 6 * 7).reshape(5, 6, 7).astype(float), 6033 dims=["x", "y", "z"], 6034 coords={"x": ["a", "b", "c", "d", "e"], "y": np.arange(6)}, 6035 ) 6036 actual = da.rolling(x=3, z=2, center=center).construct( 6037 x="x1", z="z1", fill_value=fill_value 6038 ) 6039 if not isinstance(center, tuple): 6040 center = (center, center) 6041 expected = ( 6042 da.rolling(x=3, center=center[0]) 6043 .construct(x="x1", fill_value=fill_value) 6044 .rolling(z=2, center=center[1]) 6045 .construct(z="z1", fill_value=fill_value) 6046 ) 6047 assert_allclose(actual, expected) 6048 6049 6050@pytest.mark.parametrize( 6051 "funcname, argument", 6052 [ 6053 ("reduce", (np.mean,)), 6054 ("mean", ()), 6055 ("construct", ("window_dim",)), 6056 ("count", ()), 6057 ], 6058) 6059def test_rolling_keep_attrs(funcname, argument): 6060 attrs_da = {"da_attr": "test"} 6061 6062 data = np.linspace(10, 15, 100) 6063 coords = np.linspace(1, 10, 100) 6064 6065 da = DataArray( 6066 data, dims=("coord"), coords={"coord": coords}, attrs=attrs_da, name="name" 6067 ) 6068 6069 # attrs are now kept per default 6070 func = getattr(da.rolling(dim={"coord": 5}), funcname) 6071 result = func(*argument) 6072 assert result.attrs == attrs_da 6073 assert result.name == "name" 6074 6075 # discard attrs 6076 func = getattr(da.rolling(dim={"coord": 5}), funcname) 6077 result = func(*argument, keep_attrs=False) 6078 assert result.attrs == {} 6079 assert result.name == "name" 6080 6081 # test discard attrs using global option 6082 func = getattr(da.rolling(dim={"coord": 5}), funcname) 6083 with set_options(keep_attrs=False): 6084 result = func(*argument) 6085 assert result.attrs == {} 6086 assert result.name == "name" 6087 6088 # keyword takes precedence over global option 6089 func = getattr(da.rolling(dim={"coord": 5}), funcname) 6090 with set_options(keep_attrs=False): 6091 result = func(*argument, keep_attrs=True) 6092 assert result.attrs == attrs_da 6093 assert result.name == "name" 6094 6095 func = getattr(da.rolling(dim={"coord": 5}), funcname) 6096 with set_options(keep_attrs=True): 6097 result = func(*argument, keep_attrs=False) 6098 assert result.attrs == {} 6099 assert result.name == "name" 6100 6101 6102def test_raise_no_warning_for_nan_in_binary_ops(): 6103 with pytest.warns(None) as record: 6104 xr.DataArray([1, 2, np.NaN]) > 0 6105 assert len(record) == 0 6106 6107 6108@pytest.mark.filterwarnings("error") 6109def test_no_warning_for_all_nan(): 6110 _ = xr.DataArray([np.NaN, np.NaN]).mean() 6111 6112 6113def test_name_in_masking(): 6114 name = "RingoStarr" 6115 da = xr.DataArray(range(10), coords=[("x", range(10))], name=name) 6116 assert da.where(da > 5).name == name 6117 assert da.where((da > 5).rename("YokoOno")).name == name 6118 assert da.where(da > 5, drop=True).name == name 6119 assert da.where((da > 5).rename("YokoOno"), drop=True).name == name 6120 6121 6122class TestIrisConversion: 6123 @requires_iris 6124 def test_to_and_from_iris(self): 6125 import cf_units # iris requirement 6126 import iris 6127 6128 # to iris 6129 coord_dict = {} 6130 coord_dict["distance"] = ("distance", [-2, 2], {"units": "meters"}) 6131 coord_dict["time"] = ("time", pd.date_range("2000-01-01", periods=3)) 6132 coord_dict["height"] = 10 6133 coord_dict["distance2"] = ("distance", [0, 1], {"foo": "bar"}) 6134 coord_dict["time2"] = (("distance", "time"), [[0, 1, 2], [2, 3, 4]]) 6135 6136 original = DataArray( 6137 np.arange(6, dtype="float").reshape(2, 3), 6138 coord_dict, 6139 name="Temperature", 6140 attrs={ 6141 "baz": 123, 6142 "units": "Kelvin", 6143 "standard_name": "fire_temperature", 6144 "long_name": "Fire Temperature", 6145 }, 6146 dims=("distance", "time"), 6147 ) 6148 6149 # Set a bad value to test the masking logic 6150 original.data[0, 2] = np.NaN 6151 6152 original.attrs["cell_methods"] = "height: mean (comment: A cell method)" 6153 actual = original.to_iris() 6154 assert_array_equal(actual.data, original.data) 6155 assert actual.var_name == original.name 6156 assert tuple(d.var_name for d in actual.dim_coords) == original.dims 6157 assert actual.cell_methods == ( 6158 iris.coords.CellMethod( 6159 method="mean", 6160 coords=("height",), 6161 intervals=(), 6162 comments=("A cell method",), 6163 ), 6164 ) 6165 6166 for coord, orginal_key in zip((actual.coords()), original.coords): 6167 original_coord = original.coords[orginal_key] 6168 assert coord.var_name == original_coord.name 6169 assert_array_equal( 6170 coord.points, CFDatetimeCoder().encode(original_coord).values 6171 ) 6172 assert actual.coord_dims(coord) == original.get_axis_num( 6173 original.coords[coord.var_name].dims 6174 ) 6175 6176 assert ( 6177 actual.coord("distance2").attributes["foo"] 6178 == original.coords["distance2"].attrs["foo"] 6179 ) 6180 assert actual.coord("distance").units == cf_units.Unit( 6181 original.coords["distance"].units 6182 ) 6183 assert actual.attributes["baz"] == original.attrs["baz"] 6184 assert actual.standard_name == original.attrs["standard_name"] 6185 6186 roundtripped = DataArray.from_iris(actual) 6187 assert_identical(original, roundtripped) 6188 6189 actual.remove_coord("time") 6190 auto_time_dimension = DataArray.from_iris(actual) 6191 assert auto_time_dimension.dims == ("distance", "dim_1") 6192 6193 @requires_iris 6194 @requires_dask 6195 def test_to_and_from_iris_dask(self): 6196 import cf_units # iris requirement 6197 import dask.array as da 6198 import iris 6199 6200 coord_dict = {} 6201 coord_dict["distance"] = ("distance", [-2, 2], {"units": "meters"}) 6202 coord_dict["time"] = ("time", pd.date_range("2000-01-01", periods=3)) 6203 coord_dict["height"] = 10 6204 coord_dict["distance2"] = ("distance", [0, 1], {"foo": "bar"}) 6205 coord_dict["time2"] = (("distance", "time"), [[0, 1, 2], [2, 3, 4]]) 6206 6207 original = DataArray( 6208 da.from_array(np.arange(-1, 5, dtype="float").reshape(2, 3), 3), 6209 coord_dict, 6210 name="Temperature", 6211 attrs=dict( 6212 baz=123, 6213 units="Kelvin", 6214 standard_name="fire_temperature", 6215 long_name="Fire Temperature", 6216 ), 6217 dims=("distance", "time"), 6218 ) 6219 6220 # Set a bad value to test the masking logic 6221 original.data = da.ma.masked_less(original.data, 0) 6222 6223 original.attrs["cell_methods"] = "height: mean (comment: A cell method)" 6224 actual = original.to_iris() 6225 6226 # Be careful not to trigger the loading of the iris data 6227 actual_data = ( 6228 actual.core_data() if hasattr(actual, "core_data") else actual.data 6229 ) 6230 assert_array_equal(actual_data, original.data) 6231 assert actual.var_name == original.name 6232 assert tuple(d.var_name for d in actual.dim_coords) == original.dims 6233 assert actual.cell_methods == ( 6234 iris.coords.CellMethod( 6235 method="mean", 6236 coords=("height",), 6237 intervals=(), 6238 comments=("A cell method",), 6239 ), 6240 ) 6241 6242 for coord, orginal_key in zip((actual.coords()), original.coords): 6243 original_coord = original.coords[orginal_key] 6244 assert coord.var_name == original_coord.name 6245 assert_array_equal( 6246 coord.points, CFDatetimeCoder().encode(original_coord).values 6247 ) 6248 assert actual.coord_dims(coord) == original.get_axis_num( 6249 original.coords[coord.var_name].dims 6250 ) 6251 6252 assert ( 6253 actual.coord("distance2").attributes["foo"] 6254 == original.coords["distance2"].attrs["foo"] 6255 ) 6256 assert actual.coord("distance").units == cf_units.Unit( 6257 original.coords["distance"].units 6258 ) 6259 assert actual.attributes["baz"] == original.attrs["baz"] 6260 assert actual.standard_name == original.attrs["standard_name"] 6261 6262 roundtripped = DataArray.from_iris(actual) 6263 assert_identical(original, roundtripped) 6264 6265 # If the Iris version supports it then we should have a dask array 6266 # at each stage of the conversion 6267 if hasattr(actual, "core_data"): 6268 assert isinstance(original.data, type(actual.core_data())) 6269 assert isinstance(original.data, type(roundtripped.data)) 6270 6271 actual.remove_coord("time") 6272 auto_time_dimension = DataArray.from_iris(actual) 6273 assert auto_time_dimension.dims == ("distance", "dim_1") 6274 6275 @requires_iris 6276 @pytest.mark.parametrize( 6277 "var_name, std_name, long_name, name, attrs", 6278 [ 6279 ( 6280 "var_name", 6281 "height", 6282 "Height", 6283 "var_name", 6284 {"standard_name": "height", "long_name": "Height"}, 6285 ), 6286 ( 6287 None, 6288 "height", 6289 "Height", 6290 "height", 6291 {"standard_name": "height", "long_name": "Height"}, 6292 ), 6293 (None, None, "Height", "Height", {"long_name": "Height"}), 6294 (None, None, None, None, {}), 6295 ], 6296 ) 6297 def test_da_name_from_cube(self, std_name, long_name, var_name, name, attrs): 6298 from iris.cube import Cube 6299 6300 data = [] 6301 cube = Cube( 6302 data, var_name=var_name, standard_name=std_name, long_name=long_name 6303 ) 6304 result = xr.DataArray.from_iris(cube) 6305 expected = xr.DataArray(data, name=name, attrs=attrs) 6306 xr.testing.assert_identical(result, expected) 6307 6308 @requires_iris 6309 @pytest.mark.parametrize( 6310 "var_name, std_name, long_name, name, attrs", 6311 [ 6312 ( 6313 "var_name", 6314 "height", 6315 "Height", 6316 "var_name", 6317 {"standard_name": "height", "long_name": "Height"}, 6318 ), 6319 ( 6320 None, 6321 "height", 6322 "Height", 6323 "height", 6324 {"standard_name": "height", "long_name": "Height"}, 6325 ), 6326 (None, None, "Height", "Height", {"long_name": "Height"}), 6327 (None, None, None, "unknown", {}), 6328 ], 6329 ) 6330 def test_da_coord_name_from_cube(self, std_name, long_name, var_name, name, attrs): 6331 from iris.coords import DimCoord 6332 from iris.cube import Cube 6333 6334 latitude = DimCoord( 6335 [-90, 0, 90], standard_name=std_name, var_name=var_name, long_name=long_name 6336 ) 6337 data = [0, 0, 0] 6338 cube = Cube(data, dim_coords_and_dims=[(latitude, 0)]) 6339 result = xr.DataArray.from_iris(cube) 6340 expected = xr.DataArray(data, coords=[(name, [-90, 0, 90], attrs)]) 6341 xr.testing.assert_identical(result, expected) 6342 6343 @requires_iris 6344 def test_prevent_duplicate_coord_names(self): 6345 from iris.coords import DimCoord 6346 from iris.cube import Cube 6347 6348 # Iris enforces unique coordinate names. Because we use a different 6349 # name resolution order a valid iris Cube with coords that have the 6350 # same var_name would lead to duplicate dimension names in the 6351 # DataArray 6352 longitude = DimCoord([0, 360], standard_name="longitude", var_name="duplicate") 6353 latitude = DimCoord( 6354 [-90, 0, 90], standard_name="latitude", var_name="duplicate" 6355 ) 6356 data = [[0, 0, 0], [0, 0, 0]] 6357 cube = Cube(data, dim_coords_and_dims=[(longitude, 0), (latitude, 1)]) 6358 with pytest.raises(ValueError): 6359 xr.DataArray.from_iris(cube) 6360 6361 @requires_iris 6362 @pytest.mark.parametrize( 6363 "coord_values", 6364 [["IA", "IL", "IN"], [0, 2, 1]], # non-numeric values # non-monotonic values 6365 ) 6366 def test_fallback_to_iris_AuxCoord(self, coord_values): 6367 from iris.coords import AuxCoord 6368 from iris.cube import Cube 6369 6370 data = [0, 0, 0] 6371 da = xr.DataArray(data, coords=[coord_values], dims=["space"]) 6372 result = xr.DataArray.to_iris(da) 6373 expected = Cube( 6374 data, aux_coords_and_dims=[(AuxCoord(coord_values, var_name="space"), 0)] 6375 ) 6376 assert result == expected 6377 6378 6379@requires_numbagg 6380@pytest.mark.parametrize("dim", ["time", "x"]) 6381@pytest.mark.parametrize( 6382 "window_type, window", [["span", 5], ["alpha", 0.5], ["com", 0.5], ["halflife", 5]] 6383) 6384@pytest.mark.parametrize("backend", ["numpy"], indirect=True) 6385@pytest.mark.parametrize("func", ["mean", "sum"]) 6386def test_rolling_exp_runs(da, dim, window_type, window, func): 6387 import numbagg 6388 6389 if ( 6390 LooseVersion(getattr(numbagg, "__version__", "0.1.0")) < "0.2.1" 6391 and func == "sum" 6392 ): 6393 pytest.skip("rolling_exp.sum requires numbagg 0.2.1") 6394 6395 da = da.where(da > 0.2) 6396 6397 rolling_exp = da.rolling_exp(window_type=window_type, **{dim: window}) 6398 result = getattr(rolling_exp, func)() 6399 assert isinstance(result, DataArray) 6400 6401 6402@requires_numbagg 6403@pytest.mark.parametrize("dim", ["time", "x"]) 6404@pytest.mark.parametrize( 6405 "window_type, window", [["span", 5], ["alpha", 0.5], ["com", 0.5], ["halflife", 5]] 6406) 6407@pytest.mark.parametrize("backend", ["numpy"], indirect=True) 6408def test_rolling_exp_mean_pandas(da, dim, window_type, window): 6409 da = da.isel(a=0).where(lambda x: x > 0.2) 6410 6411 result = da.rolling_exp(window_type=window_type, **{dim: window}).mean() 6412 assert isinstance(result, DataArray) 6413 6414 pandas_array = da.to_pandas() 6415 assert pandas_array.index.name == "time" 6416 if dim == "x": 6417 pandas_array = pandas_array.T 6418 expected = xr.DataArray(pandas_array.ewm(**{window_type: window}).mean()).transpose( 6419 *da.dims 6420 ) 6421 6422 assert_allclose(expected.variable, result.variable) 6423 6424 6425@requires_numbagg 6426@pytest.mark.parametrize("backend", ["numpy"], indirect=True) 6427@pytest.mark.parametrize("func", ["mean", "sum"]) 6428def test_rolling_exp_keep_attrs(da, func): 6429 import numbagg 6430 6431 if ( 6432 LooseVersion(getattr(numbagg, "__version__", "0.1.0")) < "0.2.1" 6433 and func == "sum" 6434 ): 6435 pytest.skip("rolling_exp.sum requires numbagg 0.2.1") 6436 6437 attrs = {"attrs": "da"} 6438 da.attrs = attrs 6439 6440 # Equivalent of `da.rolling_exp(time=10).mean` 6441 rolling_exp_func = getattr(da.rolling_exp(time=10), func) 6442 6443 # attrs are kept per default 6444 result = rolling_exp_func() 6445 assert result.attrs == attrs 6446 6447 # discard attrs 6448 result = rolling_exp_func(keep_attrs=False) 6449 assert result.attrs == {} 6450 6451 # test discard attrs using global option 6452 with set_options(keep_attrs=False): 6453 result = rolling_exp_func() 6454 assert result.attrs == {} 6455 6456 # keyword takes precedence over global option 6457 with set_options(keep_attrs=False): 6458 result = rolling_exp_func(keep_attrs=True) 6459 assert result.attrs == attrs 6460 6461 with set_options(keep_attrs=True): 6462 result = rolling_exp_func(keep_attrs=False) 6463 assert result.attrs == {} 6464 6465 with pytest.warns( 6466 UserWarning, match="Passing ``keep_attrs`` to ``rolling_exp`` has no effect." 6467 ): 6468 da.rolling_exp(time=10, keep_attrs=True) 6469 6470 6471def test_no_dict(): 6472 d = DataArray() 6473 with pytest.raises(AttributeError): 6474 d.__dict__ 6475 6476 6477def test_subclass_slots(): 6478 """Test that DataArray subclasses must explicitly define ``__slots__``. 6479 6480 .. note:: 6481 As of 0.13.0, this is actually mitigated into a FutureWarning for any class 6482 defined outside of the xarray package. 6483 """ 6484 with pytest.raises(AttributeError) as e: 6485 6486 class MyArray(DataArray): 6487 pass 6488 6489 assert str(e.value) == "MyArray must explicitly define __slots__" 6490 6491 6492def test_weakref(): 6493 """Classes with __slots__ are incompatible with the weakref module unless they 6494 explicitly state __weakref__ among their slots 6495 """ 6496 from weakref import ref 6497 6498 a = DataArray(1) 6499 r = ref(a) 6500 assert r() is a 6501 6502 6503def test_delete_coords(): 6504 """Make sure that deleting a coordinate doesn't corrupt the DataArray. 6505 See issue #3899. 6506 6507 Also test that deleting succeeds and produces the expected output. 6508 """ 6509 a0 = DataArray( 6510 np.array([[1, 2, 3], [4, 5, 6]]), 6511 dims=["y", "x"], 6512 coords={"x": ["a", "b", "c"], "y": [-1, 1]}, 6513 ) 6514 assert_identical(a0, a0) 6515 6516 a1 = a0.copy() 6517 del a1.coords["y"] 6518 6519 # This test will detect certain sorts of corruption in the DataArray 6520 assert_identical(a0, a0) 6521 6522 assert a0.dims == ("y", "x") 6523 assert a1.dims == ("y", "x") 6524 assert set(a0.coords.keys()) == {"x", "y"} 6525 assert set(a1.coords.keys()) == {"x"} 6526 6527 6528def test_deepcopy_obj_array(): 6529 x0 = DataArray(np.array([object()])) 6530 x1 = deepcopy(x0) 6531 assert x0.values[0] is not x1.values[0] 6532 6533 6534def test_clip(da): 6535 with raise_if_dask_computes(): 6536 result = da.clip(min=0.5) 6537 assert result.min(...) >= 0.5 6538 6539 result = da.clip(max=0.5) 6540 assert result.max(...) <= 0.5 6541 6542 result = da.clip(min=0.25, max=0.75) 6543 assert result.min(...) >= 0.25 6544 assert result.max(...) <= 0.75 6545 6546 with raise_if_dask_computes(): 6547 result = da.clip(min=da.mean("x"), max=da.mean("a")) 6548 assert result.dims == da.dims 6549 assert_array_equal( 6550 result.data, 6551 np.clip(da.data, da.mean("x").data[:, :, np.newaxis], da.mean("a").data), 6552 ) 6553 6554 with_nans = da.isel(time=[0, 1]).reindex_like(da) 6555 with raise_if_dask_computes(): 6556 result = da.clip(min=da.mean("x"), max=da.mean("a")) 6557 result = da.clip(with_nans) 6558 # The values should be the same where there were NaNs. 6559 assert_array_equal(result.isel(time=[0, 1]), with_nans.isel(time=[0, 1])) 6560 6561 # Unclear whether we want this work, OK to adjust the test when we have decided. 6562 with pytest.raises(ValueError, match="arguments without labels along dimension"): 6563 result = da.clip(min=da.mean("x"), max=da.mean("a").isel(x=[0, 1])) 6564 6565 6566@pytest.mark.parametrize("keep", ["first", "last", False]) 6567def test_drop_duplicates(keep): 6568 ds = xr.DataArray( 6569 [0, 5, 6, 7], dims="time", coords={"time": [0, 0, 1, 2]}, name="test" 6570 ) 6571 6572 if keep == "first": 6573 data = [0, 6, 7] 6574 time = [0, 1, 2] 6575 elif keep == "last": 6576 data = [5, 6, 7] 6577 time = [0, 1, 2] 6578 else: 6579 data = [6, 7] 6580 time = [1, 2] 6581 6582 expected = xr.DataArray(data, dims="time", coords={"time": time}, name="test") 6583 result = ds.drop_duplicates("time", keep=keep) 6584 assert_equal(expected, result) 6585 6586 6587class TestNumpyCoercion: 6588 # TODO once flexible indexes refactor complete also test coercion of dimension coords 6589 def test_from_numpy(self): 6590 da = xr.DataArray([1, 2, 3], dims="x", coords={"lat": ("x", [4, 5, 6])}) 6591 6592 assert_identical(da.as_numpy(), da) 6593 np.testing.assert_equal(da.to_numpy(), np.array([1, 2, 3])) 6594 np.testing.assert_equal(da["lat"].to_numpy(), np.array([4, 5, 6])) 6595 6596 @requires_dask 6597 def test_from_dask(self): 6598 da = xr.DataArray([1, 2, 3], dims="x", coords={"lat": ("x", [4, 5, 6])}) 6599 da_chunked = da.chunk(1) 6600 6601 assert_identical(da_chunked.as_numpy(), da.compute()) 6602 np.testing.assert_equal(da.to_numpy(), np.array([1, 2, 3])) 6603 np.testing.assert_equal(da["lat"].to_numpy(), np.array([4, 5, 6])) 6604 6605 @requires_pint 6606 def test_from_pint(self): 6607 from pint import Quantity 6608 6609 arr = np.array([1, 2, 3]) 6610 da = xr.DataArray( 6611 Quantity(arr, units="Pa"), 6612 dims="x", 6613 coords={"lat": ("x", Quantity(arr + 3, units="m"))}, 6614 ) 6615 6616 expected = xr.DataArray(arr, dims="x", coords={"lat": ("x", arr + 3)}) 6617 assert_identical(da.as_numpy(), expected) 6618 np.testing.assert_equal(da.to_numpy(), arr) 6619 np.testing.assert_equal(da["lat"].to_numpy(), arr + 3) 6620 6621 @requires_sparse 6622 def test_from_sparse(self): 6623 import sparse 6624 6625 arr = np.diagflat([1, 2, 3]) 6626 sparr = sparse.COO.from_numpy(arr) 6627 da = xr.DataArray( 6628 sparr, dims=["x", "y"], coords={"elev": (("x", "y"), sparr + 3)} 6629 ) 6630 6631 expected = xr.DataArray( 6632 arr, dims=["x", "y"], coords={"elev": (("x", "y"), arr + 3)} 6633 ) 6634 assert_identical(da.as_numpy(), expected) 6635 np.testing.assert_equal(da.to_numpy(), arr) 6636 6637 @requires_cupy 6638 def test_from_cupy(self): 6639 import cupy as cp 6640 6641 arr = np.array([1, 2, 3]) 6642 da = xr.DataArray( 6643 cp.array(arr), dims="x", coords={"lat": ("x", cp.array(arr + 3))} 6644 ) 6645 6646 expected = xr.DataArray(arr, dims="x", coords={"lat": ("x", arr + 3)}) 6647 assert_identical(da.as_numpy(), expected) 6648 np.testing.assert_equal(da.to_numpy(), arr) 6649 6650 @requires_dask 6651 @requires_pint 6652 def test_from_pint_wrapping_dask(self): 6653 import dask 6654 from pint import Quantity 6655 6656 arr = np.array([1, 2, 3]) 6657 d = dask.array.from_array(arr) 6658 da = xr.DataArray( 6659 Quantity(d, units="Pa"), 6660 dims="x", 6661 coords={"lat": ("x", Quantity(d, units="m") * 2)}, 6662 ) 6663 6664 result = da.as_numpy() 6665 result.name = None # remove dask-assigned name 6666 expected = xr.DataArray(arr, dims="x", coords={"lat": ("x", arr * 2)}) 6667 assert_identical(result, expected) 6668 np.testing.assert_equal(da.to_numpy(), arr) 6669