1# Owner(s): ["module: dynamo"] 2 3"""Test functions for 1D array set operations. 4 5""" 6from unittest import expectedFailure as xfail, skipIf 7 8import numpy 9from pytest import raises as assert_raises 10 11from torch.testing._internal.common_utils import ( 12 instantiate_parametrized_tests, 13 parametrize, 14 run_tests, 15 subtest, 16 TEST_WITH_TORCHDYNAMO, 17 TestCase, 18 xpassIfTorchDynamo, 19) 20 21 22# If we are going to trace through these, we should use NumPy 23# If testing on eager mode, we use torch._numpy 24if TEST_WITH_TORCHDYNAMO: 25 import numpy as np 26 from numpy import ediff1d, in1d, intersect1d, setdiff1d, setxor1d, union1d, unique 27 from numpy.testing import assert_array_equal, assert_equal, assert_raises_regex 28 29else: 30 import torch._numpy as np 31 from torch._numpy import unique 32 from torch._numpy.testing import assert_array_equal, assert_equal 33 34 35@skipIf(numpy.__version__ < "1.24", reason="NP_VER: fails on NumPy 1.23.x") 36@skipIf(True, reason="TODO implement these ops") 37@instantiate_parametrized_tests 38class TestSetOps(TestCase): 39 def test_intersect1d(self): 40 # unique inputs 41 a = np.array([5, 7, 1, 2]) 42 b = np.array([2, 4, 3, 1, 5]) 43 44 ec = np.array([1, 2, 5]) 45 c = intersect1d(a, b, assume_unique=True) 46 assert_array_equal(c, ec) 47 48 # non-unique inputs 49 a = np.array([5, 5, 7, 1, 2]) 50 b = np.array([2, 1, 4, 3, 3, 1, 5]) 51 52 ed = np.array([1, 2, 5]) 53 c = intersect1d(a, b) 54 assert_array_equal(c, ed) 55 assert_array_equal([], intersect1d([], [])) 56 57 def test_intersect1d_array_like(self): 58 # See gh-11772 59 class Test: 60 def __array__(self): 61 return np.arange(3) 62 63 a = Test() 64 res = intersect1d(a, a) 65 assert_array_equal(res, a) 66 res = intersect1d([1, 2, 3], [1, 2, 3]) 67 assert_array_equal(res, [1, 2, 3]) 68 69 def test_intersect1d_indices(self): 70 # unique inputs 71 a = np.array([1, 2, 3, 4]) 72 b = np.array([2, 1, 4, 6]) 73 c, i1, i2 = intersect1d(a, b, assume_unique=True, return_indices=True) 74 ee = np.array([1, 2, 4]) 75 assert_array_equal(c, ee) 76 assert_array_equal(a[i1], ee) 77 assert_array_equal(b[i2], ee) 78 79 # non-unique inputs 80 a = np.array([1, 2, 2, 3, 4, 3, 2]) 81 b = np.array([1, 8, 4, 2, 2, 3, 2, 3]) 82 c, i1, i2 = intersect1d(a, b, return_indices=True) 83 ef = np.array([1, 2, 3, 4]) 84 assert_array_equal(c, ef) 85 assert_array_equal(a[i1], ef) 86 assert_array_equal(b[i2], ef) 87 88 # non1d, unique inputs 89 a = np.array([[2, 4, 5, 6], [7, 8, 1, 15]]) 90 b = np.array([[3, 2, 7, 6], [10, 12, 8, 9]]) 91 c, i1, i2 = intersect1d(a, b, assume_unique=True, return_indices=True) 92 ui1 = np.unravel_index(i1, a.shape) 93 ui2 = np.unravel_index(i2, b.shape) 94 ea = np.array([2, 6, 7, 8]) 95 assert_array_equal(ea, a[ui1]) 96 assert_array_equal(ea, b[ui2]) 97 98 # non1d, not assumed to be uniqueinputs 99 a = np.array([[2, 4, 5, 6, 6], [4, 7, 8, 7, 2]]) 100 b = np.array([[3, 2, 7, 7], [10, 12, 8, 7]]) 101 c, i1, i2 = intersect1d(a, b, return_indices=True) 102 ui1 = np.unravel_index(i1, a.shape) 103 ui2 = np.unravel_index(i2, b.shape) 104 ea = np.array([2, 7, 8]) 105 assert_array_equal(ea, a[ui1]) 106 assert_array_equal(ea, b[ui2]) 107 108 def test_setxor1d(self): 109 a = np.array([5, 7, 1, 2]) 110 b = np.array([2, 4, 3, 1, 5]) 111 112 ec = np.array([3, 4, 7]) 113 c = setxor1d(a, b) 114 assert_array_equal(c, ec) 115 116 a = np.array([1, 2, 3]) 117 b = np.array([6, 5, 4]) 118 119 ec = np.array([1, 2, 3, 4, 5, 6]) 120 c = setxor1d(a, b) 121 assert_array_equal(c, ec) 122 123 a = np.array([1, 8, 2, 3]) 124 b = np.array([6, 5, 4, 8]) 125 126 ec = np.array([1, 2, 3, 4, 5, 6]) 127 c = setxor1d(a, b) 128 assert_array_equal(c, ec) 129 130 assert_array_equal([], setxor1d([], [])) 131 132 def test_ediff1d(self): 133 zero_elem = np.array([]) 134 one_elem = np.array([1]) 135 two_elem = np.array([1, 2]) 136 137 assert_array_equal([], ediff1d(zero_elem)) 138 assert_array_equal([0], ediff1d(zero_elem, to_begin=0)) 139 assert_array_equal([0], ediff1d(zero_elem, to_end=0)) 140 assert_array_equal([-1, 0], ediff1d(zero_elem, to_begin=-1, to_end=0)) 141 assert_array_equal([], ediff1d(one_elem)) 142 assert_array_equal([1], ediff1d(two_elem)) 143 assert_array_equal([7, 1, 9], ediff1d(two_elem, to_begin=7, to_end=9)) 144 assert_array_equal( 145 [5, 6, 1, 7, 8], ediff1d(two_elem, to_begin=[5, 6], to_end=[7, 8]) 146 ) 147 assert_array_equal([1, 9], ediff1d(two_elem, to_end=9)) 148 assert_array_equal([1, 7, 8], ediff1d(two_elem, to_end=[7, 8])) 149 assert_array_equal([7, 1], ediff1d(two_elem, to_begin=7)) 150 assert_array_equal([5, 6, 1], ediff1d(two_elem, to_begin=[5, 6])) 151 152 @parametrize( 153 "ary, prepend, append, expected", 154 [ 155 # should fail because trying to cast 156 # np.nan standard floating point value 157 # into an integer array: 158 (np.array([1, 2, 3], dtype=np.int64), None, np.nan, "to_end"), 159 # should fail because attempting 160 # to downcast to int type: 161 subtest( 162 ( 163 np.array([1, 2, 3], dtype=np.int64), 164 np.array([5, 7, 2], dtype=np.float32), 165 None, 166 "to_begin", 167 ), 168 ), 169 # should fail because attempting to cast 170 # two special floating point values 171 # to integers (on both sides of ary), 172 # `to_begin` is in the error message as the impl checks this first: 173 (np.array([1.0, 3.0, 9.0], dtype=np.int8), np.nan, np.nan, "to_begin"), 174 ], 175 ) 176 def test_ediff1d_forbidden_type_casts(self, ary, prepend, append, expected): 177 # verify resolution of gh-11490 178 179 # specifically, raise an appropriate 180 # Exception when attempting to append or 181 # prepend with an incompatible type 182 msg = f"dtype of `{expected}` must be compatible" 183 with assert_raises_regex(TypeError, msg): 184 ediff1d(ary=ary, to_end=append, to_begin=prepend) 185 186 @parametrize( 187 "ary,prepend,append,expected", 188 [ 189 ( 190 np.array([1, 2, 3], dtype=np.int16), 191 2**16, # will be cast to int16 under same kind rule. 192 2**16 + 4, 193 np.array([0, 1, 1, 4], dtype=np.int16), 194 ), 195 ( 196 np.array([1, 2, 3], dtype=np.float32), 197 np.array([5], dtype=np.float64), 198 None, 199 np.array([5, 1, 1], dtype=np.float32), 200 ), 201 ( 202 np.array([1, 2, 3], dtype=np.int32), 203 0, 204 0, 205 np.array([0, 1, 1, 0], dtype=np.int32), 206 ), 207 ( 208 np.array([1, 2, 3], dtype=np.int64), 209 3, 210 -9, 211 np.array([3, 1, 1, -9], dtype=np.int64), 212 ), 213 ], 214 ) 215 def test_ediff1d_scalar_handling(self, ary, prepend, append, expected): 216 # maintain backwards-compatibility 217 # of scalar prepend / append behavior 218 # in ediff1d following fix for gh-11490 219 actual = np.ediff1d(ary=ary, to_end=append, to_begin=prepend) 220 assert_equal(actual, expected) 221 assert actual.dtype == expected.dtype 222 223 @skipIf(True, reason="NP_VER: fails with NumPy 1.22.x") 224 @parametrize("kind", [None, "sort", "table"]) 225 def test_isin(self, kind): 226 # the tests for in1d cover most of isin's behavior 227 # if in1d is removed, would need to change those tests to test 228 # isin instead. 229 def _isin_slow(a, b): 230 b = np.asarray(b).flatten().tolist() 231 return a in b 232 233 isin_slow = np.vectorize(_isin_slow, otypes=[bool], excluded={1}) 234 235 def assert_isin_equal(a, b): 236 x = np.isin(a, b, kind=kind) 237 y = isin_slow(a, b) 238 assert_array_equal(x, y) 239 240 # multidimensional arrays in both arguments 241 a = np.arange(24).reshape([2, 3, 4]) 242 b = np.array([[10, 20, 30], [0, 1, 3], [11, 22, 33]]) 243 assert_isin_equal(a, b) 244 245 # array-likes as both arguments 246 c = [(9, 8), (7, 6)] 247 d = (9, 7) 248 assert_isin_equal(c, d) 249 250 # zero-d array: 251 f = np.array(3) 252 assert_isin_equal(f, b) 253 assert_isin_equal(a, f) 254 assert_isin_equal(f, f) 255 256 # scalar: 257 assert_isin_equal(5, b) 258 assert_isin_equal(a, 6) 259 assert_isin_equal(5, 6) 260 261 # empty array-like: 262 if kind != "table": 263 # An empty list will become float64, 264 # which is invalid for kind="table" 265 x = [] 266 assert_isin_equal(x, b) 267 assert_isin_equal(a, x) 268 assert_isin_equal(x, x) 269 270 # empty array with various types: 271 for dtype in [bool, np.int64, np.float64]: 272 if kind == "table" and dtype == np.float64: 273 continue 274 275 if dtype in {np.int64, np.float64}: 276 ar = np.array([10, 20, 30], dtype=dtype) 277 elif dtype in {bool}: 278 ar = np.array([True, False, False]) 279 280 empty_array = np.array([], dtype=dtype) 281 282 assert_isin_equal(empty_array, ar) 283 assert_isin_equal(ar, empty_array) 284 assert_isin_equal(empty_array, empty_array) 285 286 @parametrize("kind", [None, "sort", "table"]) 287 def test_in1d(self, kind): 288 # we use two different sizes for the b array here to test the 289 # two different paths in in1d(). 290 for mult in (1, 10): 291 # One check without np.array to make sure lists are handled correct 292 a = [5, 7, 1, 2] 293 b = [2, 4, 3, 1, 5] * mult 294 ec = np.array([True, False, True, True]) 295 c = in1d(a, b, assume_unique=True, kind=kind) 296 assert_array_equal(c, ec) 297 298 a[0] = 8 299 ec = np.array([False, False, True, True]) 300 c = in1d(a, b, assume_unique=True, kind=kind) 301 assert_array_equal(c, ec) 302 303 a[0], a[3] = 4, 8 304 ec = np.array([True, False, True, False]) 305 c = in1d(a, b, assume_unique=True, kind=kind) 306 assert_array_equal(c, ec) 307 308 a = np.array([5, 4, 5, 3, 4, 4, 3, 4, 3, 5, 2, 1, 5, 5]) 309 b = [2, 3, 4] * mult 310 ec = [ 311 False, 312 True, 313 False, 314 True, 315 True, 316 True, 317 True, 318 True, 319 True, 320 False, 321 True, 322 False, 323 False, 324 False, 325 ] 326 c = in1d(a, b, kind=kind) 327 assert_array_equal(c, ec) 328 329 b = b + [5, 5, 4] * mult 330 ec = [ 331 True, 332 True, 333 True, 334 True, 335 True, 336 True, 337 True, 338 True, 339 True, 340 True, 341 True, 342 False, 343 True, 344 True, 345 ] 346 c = in1d(a, b, kind=kind) 347 assert_array_equal(c, ec) 348 349 a = np.array([5, 7, 1, 2]) 350 b = np.array([2, 4, 3, 1, 5] * mult) 351 ec = np.array([True, False, True, True]) 352 c = in1d(a, b, kind=kind) 353 assert_array_equal(c, ec) 354 355 a = np.array([5, 7, 1, 1, 2]) 356 b = np.array([2, 4, 3, 3, 1, 5] * mult) 357 ec = np.array([True, False, True, True, True]) 358 c = in1d(a, b, kind=kind) 359 assert_array_equal(c, ec) 360 361 a = np.array([5, 5]) 362 b = np.array([2, 2] * mult) 363 ec = np.array([False, False]) 364 c = in1d(a, b, kind=kind) 365 assert_array_equal(c, ec) 366 367 a = np.array([5]) 368 b = np.array([2]) 369 ec = np.array([False]) 370 c = in1d(a, b, kind=kind) 371 assert_array_equal(c, ec) 372 373 if kind in {None, "sort"}: 374 assert_array_equal(in1d([], [], kind=kind), []) 375 376 def test_in1d_char_array(self): 377 a = np.array(["a", "b", "c", "d", "e", "c", "e", "b"]) 378 b = np.array(["a", "c"]) 379 380 ec = np.array([True, False, True, False, False, True, False, False]) 381 c = in1d(a, b) 382 383 assert_array_equal(c, ec) 384 385 @parametrize("kind", [None, "sort", "table"]) 386 def test_in1d_invert(self, kind): 387 "Test in1d's invert parameter" 388 # We use two different sizes for the b array here to test the 389 # two different paths in in1d(). 390 for mult in (1, 10): 391 a = np.array([5, 4, 5, 3, 4, 4, 3, 4, 3, 5, 2, 1, 5, 5]) 392 b = [2, 3, 4] * mult 393 assert_array_equal( 394 np.invert(in1d(a, b, kind=kind)), in1d(a, b, invert=True, kind=kind) 395 ) 396 397 # float: 398 if kind in {None, "sort"}: 399 for mult in (1, 10): 400 a = np.array( 401 [5, 4, 5, 3, 4, 4, 3, 4, 3, 5, 2, 1, 5, 5], dtype=np.float32 402 ) 403 b = [2, 3, 4] * mult 404 b = np.array(b, dtype=np.float32) 405 assert_array_equal( 406 np.invert(in1d(a, b, kind=kind)), in1d(a, b, invert=True, kind=kind) 407 ) 408 409 @parametrize("kind", [None, "sort", "table"]) 410 def test_in1d_ravel(self, kind): 411 # Test that in1d ravels its input arrays. This is not documented 412 # behavior however. The test is to ensure consistentency. 413 a = np.arange(6).reshape(2, 3) 414 b = np.arange(3, 9).reshape(3, 2) 415 long_b = np.arange(3, 63).reshape(30, 2) 416 ec = np.array([False, False, False, True, True, True]) 417 418 assert_array_equal(in1d(a, b, assume_unique=True, kind=kind), ec) 419 assert_array_equal(in1d(a, b, assume_unique=False, kind=kind), ec) 420 assert_array_equal(in1d(a, long_b, assume_unique=True, kind=kind), ec) 421 assert_array_equal(in1d(a, long_b, assume_unique=False, kind=kind), ec) 422 423 def test_in1d_hit_alternate_algorithm(self): 424 """Hit the standard isin code with integers""" 425 # Need extreme range to hit standard code 426 # This hits it without the use of kind='table' 427 a = np.array([5, 4, 5, 3, 4, 4, 1e9], dtype=np.int64) 428 b = np.array([2, 3, 4, 1e9], dtype=np.int64) 429 expected = np.array([0, 1, 0, 1, 1, 1, 1], dtype=bool) 430 assert_array_equal(expected, in1d(a, b)) 431 assert_array_equal(np.invert(expected), in1d(a, b, invert=True)) 432 433 a = np.array([5, 7, 1, 2], dtype=np.int64) 434 b = np.array([2, 4, 3, 1, 5, 1e9], dtype=np.int64) 435 ec = np.array([True, False, True, True]) 436 c = in1d(a, b, assume_unique=True) 437 assert_array_equal(c, ec) 438 439 @parametrize("kind", [None, "sort", "table"]) 440 def test_in1d_boolean(self, kind): 441 """Test that in1d works for boolean input""" 442 a = np.array([True, False]) 443 b = np.array([False, False, False]) 444 expected = np.array([False, True]) 445 assert_array_equal(expected, in1d(a, b, kind=kind)) 446 assert_array_equal(np.invert(expected), in1d(a, b, invert=True, kind=kind)) 447 448 @parametrize("kind", [None, "sort"]) 449 def test_in1d_timedelta(self, kind): 450 """Test that in1d works for timedelta input""" 451 rstate = np.random.RandomState(0) 452 a = rstate.randint(0, 100, size=10) 453 b = rstate.randint(0, 100, size=10) 454 truth = in1d(a, b) 455 a_timedelta = a.astype("timedelta64[s]") 456 b_timedelta = b.astype("timedelta64[s]") 457 assert_array_equal(truth, in1d(a_timedelta, b_timedelta, kind=kind)) 458 459 def test_in1d_table_timedelta_fails(self): 460 a = np.array([0, 1, 2], dtype="timedelta64[s]") 461 b = a 462 # Make sure it raises a value error: 463 with assert_raises(ValueError): 464 in1d(a, b, kind="table") 465 466 @parametrize( 467 "dtype1,dtype2", 468 [ 469 (np.int8, np.int16), 470 (np.int16, np.int8), 471 ], 472 ) 473 @parametrize("kind", [None, "sort", "table"]) 474 def test_in1d_mixed_dtype(self, dtype1, dtype2, kind): 475 """Test that in1d works as expected for mixed dtype input.""" 476 is_dtype2_signed = np.issubdtype(dtype2, np.signedinteger) 477 ar1 = np.array([0, 0, 1, 1], dtype=dtype1) 478 479 if is_dtype2_signed: 480 ar2 = np.array([-128, 0, 127], dtype=dtype2) 481 else: 482 ar2 = np.array([127, 0, 255], dtype=dtype2) 483 484 expected = np.array([True, True, False, False]) 485 486 expect_failure = kind == "table" and any( 487 ( 488 dtype1 == np.int8 and dtype2 == np.int16, 489 dtype1 == np.int16 and dtype2 == np.int8, 490 ) 491 ) 492 493 if expect_failure: 494 with assert_raises(RuntimeError, match="exceed the maximum"): 495 in1d(ar1, ar2, kind=kind) 496 else: 497 assert_array_equal(in1d(ar1, ar2, kind=kind), expected) 498 499 @parametrize("kind", [None, "sort", "table"]) 500 def test_in1d_mixed_boolean(self, kind): 501 """Test that in1d works as expected for bool/int input.""" 502 for dtype in np.typecodes["AllInteger"]: 503 a = np.array([True, False, False], dtype=bool) 504 b = np.array([0, 0, 0, 0], dtype=dtype) 505 expected = np.array([False, True, True], dtype=bool) 506 assert_array_equal(in1d(a, b, kind=kind), expected) 507 508 a, b = b, a 509 expected = np.array([True, True, True, True], dtype=bool) 510 assert_array_equal(in1d(a, b, kind=kind), expected) 511 512 def test_in1d_first_array_is_object(self): 513 ar1 = [None] 514 ar2 = np.array([1] * 10) 515 expected = np.array([False]) 516 result = np.in1d(ar1, ar2) 517 assert_array_equal(result, expected) 518 519 def test_in1d_second_array_is_object(self): 520 ar1 = 1 521 ar2 = np.array([None] * 10) 522 expected = np.array([False]) 523 result = np.in1d(ar1, ar2) 524 assert_array_equal(result, expected) 525 526 def test_in1d_both_arrays_are_object(self): 527 ar1 = [None] 528 ar2 = np.array([None] * 10) 529 expected = np.array([True]) 530 result = np.in1d(ar1, ar2) 531 assert_array_equal(result, expected) 532 533 @xfail 534 def test_in1d_both_arrays_have_structured_dtype(self): 535 # Test arrays of a structured data type containing an integer field 536 # and a field of dtype `object` allowing for arbitrary Python objects 537 dt = np.dtype([("field1", int), ("field2", object)]) 538 ar1 = np.array([(1, None)], dtype=dt) 539 ar2 = np.array([(1, None)] * 10, dtype=dt) 540 expected = np.array([True]) 541 result = np.in1d(ar1, ar2) 542 assert_array_equal(result, expected) 543 544 def test_in1d_with_arrays_containing_tuples(self): 545 ar1 = np.array([(1,), 2], dtype=object) 546 ar2 = np.array([(1,), 2], dtype=object) 547 expected = np.array([True, True]) 548 result = np.in1d(ar1, ar2) 549 assert_array_equal(result, expected) 550 result = np.in1d(ar1, ar2, invert=True) 551 assert_array_equal(result, np.invert(expected)) 552 553 # An integer is added at the end of the array to make sure 554 # that the array builder will create the array with tuples 555 # and after it's created the integer is removed. 556 # There's a bug in the array constructor that doesn't handle 557 # tuples properly and adding the integer fixes that. 558 ar1 = np.array([(1,), (2, 1), 1], dtype=object) 559 ar1 = ar1[:-1] 560 ar2 = np.array([(1,), (2, 1), 1], dtype=object) 561 ar2 = ar2[:-1] 562 expected = np.array([True, True]) 563 result = np.in1d(ar1, ar2) 564 assert_array_equal(result, expected) 565 result = np.in1d(ar1, ar2, invert=True) 566 assert_array_equal(result, np.invert(expected)) 567 568 ar1 = np.array([(1,), (2, 3), 1], dtype=object) 569 ar1 = ar1[:-1] 570 ar2 = np.array([(1,), 2], dtype=object) 571 expected = np.array([True, False]) 572 result = np.in1d(ar1, ar2) 573 assert_array_equal(result, expected) 574 result = np.in1d(ar1, ar2, invert=True) 575 assert_array_equal(result, np.invert(expected)) 576 577 def test_in1d_errors(self): 578 """Test that in1d raises expected errors.""" 579 580 # Error 1: `kind` is not one of 'sort' 'table' or None. 581 ar1 = np.array([1, 2, 3, 4, 5]) 582 ar2 = np.array([2, 4, 6, 8, 10]) 583 assert_raises(ValueError, in1d, ar1, ar2, kind="quicksort") 584 585 # Error 2: `kind="table"` does not work for non-integral arrays. 586 obj_ar1 = np.array([1, "a", 3, "b", 5], dtype=object) 587 obj_ar2 = np.array([1, "a", 3, "b", 5], dtype=object) 588 assert_raises(ValueError, in1d, obj_ar1, obj_ar2, kind="table") 589 590 for dtype in [np.int32, np.int64]: 591 ar1 = np.array([-1, 2, 3, 4, 5], dtype=dtype) 592 # The range of this array will overflow: 593 overflow_ar2 = np.array([-1, np.iinfo(dtype).max], dtype=dtype) 594 595 # Error 3: `kind="table"` will trigger a runtime error 596 # if there is an integer overflow expected when computing the 597 # range of ar2 598 assert_raises(RuntimeError, in1d, ar1, overflow_ar2, kind="table") 599 600 # Non-error: `kind=None` will *not* trigger a runtime error 601 # if there is an integer overflow, it will switch to 602 # the `sort` algorithm. 603 result = np.in1d(ar1, overflow_ar2, kind=None) 604 assert_array_equal(result, [True] + [False] * 4) 605 result = np.in1d(ar1, overflow_ar2, kind="sort") 606 assert_array_equal(result, [True] + [False] * 4) 607 608 def test_union1d(self): 609 a = np.array([5, 4, 7, 1, 2]) 610 b = np.array([2, 4, 3, 3, 2, 1, 5]) 611 612 ec = np.array([1, 2, 3, 4, 5, 7]) 613 c = union1d(a, b) 614 assert_array_equal(c, ec) 615 616 # Tests gh-10340, arguments to union1d should be 617 # flattened if they are not already 1D 618 x = np.array([[0, 1, 2], [3, 4, 5]]) 619 y = np.array([0, 1, 2, 3, 4]) 620 ez = np.array([0, 1, 2, 3, 4, 5]) 621 z = union1d(x, y) 622 assert_array_equal(z, ez) 623 624 assert_array_equal([], union1d([], [])) 625 626 def test_setdiff1d(self): 627 a = np.array([6, 5, 4, 7, 1, 2, 7, 4]) 628 b = np.array([2, 4, 3, 3, 2, 1, 5]) 629 630 ec = np.array([6, 7]) 631 c = setdiff1d(a, b) 632 assert_array_equal(c, ec) 633 634 a = np.arange(21) 635 b = np.arange(19) 636 ec = np.array([19, 20]) 637 c = setdiff1d(a, b) 638 assert_array_equal(c, ec) 639 640 assert_array_equal([], setdiff1d([], [])) 641 a = np.array((), np.uint32) 642 assert_equal(setdiff1d(a, []).dtype, np.uint32) 643 644 def test_setdiff1d_unique(self): 645 a = np.array([3, 2, 1]) 646 b = np.array([7, 5, 2]) 647 expected = np.array([3, 1]) 648 actual = setdiff1d(a, b, assume_unique=True) 649 assert_equal(actual, expected) 650 651 def test_setdiff1d_char_array(self): 652 a = np.array(["a", "b", "c"]) 653 b = np.array(["a", "b", "s"]) 654 assert_array_equal(setdiff1d(a, b), np.array(["c"])) 655 656 def test_manyways(self): 657 a = np.array([5, 7, 1, 2, 8]) 658 b = np.array([9, 8, 2, 4, 3, 1, 5]) 659 660 c1 = setxor1d(a, b) 661 aux1 = intersect1d(a, b) 662 aux2 = union1d(a, b) 663 c2 = setdiff1d(aux2, aux1) 664 assert_array_equal(c1, c2) 665 666 667@instantiate_parametrized_tests 668class TestUnique(TestCase): 669 def test_unique_1d(self): 670 def check_all(a, b, i1, i2, c, dt): 671 base_msg = "check {0} failed for type {1}" 672 673 msg = base_msg.format("values", dt) 674 v = unique(a) 675 assert_array_equal(v, b, msg) 676 677 # msg = base_msg.format('return_index', dt) 678 # v, j = unique(a, True, False, False) 679 # assert_array_equal(v, b, msg) 680 # assert_array_equal(j, i1, msg) 681 682 msg = base_msg.format("return_inverse", dt) 683 v, j = unique(a, False, True, False) 684 assert_array_equal(v, b, msg) 685 assert_array_equal(j, i2, msg) 686 687 msg = base_msg.format("return_counts", dt) 688 v, j = unique(a, False, False, True) 689 assert_array_equal(v, b, msg) 690 assert_array_equal(j, c, msg) 691 692 # msg = base_msg.format('return_index and return_inverse', dt) 693 # v, j1, j2 = unique(a, True, True, False) 694 # assert_array_equal(v, b, msg) 695 # assert_array_equal(j1, i1, msg) 696 # assert_array_equal(j2, i2, msg) 697 698 # msg = base_msg.format('return_index and return_counts', dt) 699 # v, j1, j2 = unique(a, True, False, True) 700 # assert_array_equal(v, b, msg) 701 # assert_array_equal(j1, i1, msg) 702 # assert_array_equal(j2, c, msg) 703 704 msg = base_msg.format("return_inverse and return_counts", dt) 705 v, j1, j2 = unique(a, False, True, True) 706 assert_array_equal(v, b, msg) 707 assert_array_equal(j1, i2, msg) 708 assert_array_equal(j2, c, msg) 709 710 # msg = base_msg.format(('return_index, return_inverse ' 711 # 'and return_counts'), dt) 712 # v, j1, j2, j3 = unique(a, True, True, True) 713 # assert_array_equal(v, b, msg) 714 # assert_array_equal(j1, i1, msg) 715 # assert_array_equal(j2, i2, msg) 716 # assert_array_equal(j3, c, msg) 717 718 a = [5, 7, 1, 2, 1, 5, 7] * 10 719 b = [1, 2, 5, 7] 720 i1 = [2, 3, 0, 1] 721 i2 = [2, 3, 0, 1, 0, 2, 3] * 10 722 c = np.multiply([2, 1, 2, 2], 10) 723 724 # test for numeric arrays 725 types = [] 726 types.extend(np.typecodes["AllInteger"]) 727 types.extend(np.typecodes["AllFloat"]) 728 for dt in types: 729 if dt in "FD": 730 # RuntimeError: "unique" not implemented for 'ComplexFloat' 731 continue 732 733 aa = np.array(a, dt) 734 bb = np.array(b, dt) 735 check_all(aa, bb, i1, i2, c, dt) 736 737 # test for ticket #2799 738 # RuntimeError: "unique" not implemented for 'ComplexFloat' 739 # aa = [1. + 0.j, 1 - 1.j, 1] 740 # assert_array_equal(np.unique(aa), [1. - 1.j, 1. + 0.j]) 741 742 # test for ticket #4785 743 a = [(1, 2), (1, 2), (2, 3)] 744 unq = [1, 2, 3] 745 inv = [0, 1, 0, 1, 1, 2] 746 a1 = unique(a) 747 assert_array_equal(a1, unq) 748 a2, a2_inv = unique(a, return_inverse=True) 749 assert_array_equal(a2, unq) 750 assert_array_equal(a2_inv, inv) 751 752 # test for ticket #9137 753 a = [] 754 # a1_idx = np.unique(a, return_index=True)[1] 755 a2_inv = np.unique(a, return_inverse=True)[1] 756 # a3_idx, a3_inv = np.unique(a, return_index=True, 757 # return_inverse=True)[1:] 758 # assert_equal(a1_idx.dtype, np.intp) 759 assert_equal(a2_inv.dtype, np.intp) 760 761 # assert_equal(a3_idx.dtype, np.intp) 762 # assert_equal(a3_inv.dtype, np.intp) 763 764 @xpassIfTorchDynamo # (reason="unique with nans") 765 def test_unique_1d_2(self): 766 # test for ticket 2111 - float 767 a = [2.0, np.nan, 1.0, np.nan] 768 ua = [1.0, 2.0, np.nan] 769 ua_idx = [2, 0, 1] 770 ua_inv = [1, 2, 0, 2] 771 ua_cnt = [1, 1, 2] 772 assert_equal(np.unique(a), ua) 773 assert_equal(np.unique(a, return_index=True), (ua, ua_idx)) 774 assert_equal(np.unique(a, return_inverse=True), (ua, ua_inv)) 775 assert_equal(np.unique(a, return_counts=True), (ua, ua_cnt)) 776 777 # test for ticket 2111 - complex 778 a = [2.0 - 1j, np.nan, 1.0 + 1j, complex(0.0, np.nan), complex(1.0, np.nan)] 779 ua = [1.0 + 1j, 2.0 - 1j, complex(0.0, np.nan)] 780 ua_idx = [2, 0, 3] 781 ua_inv = [1, 2, 0, 2, 2] 782 ua_cnt = [1, 1, 3] 783 assert_equal(np.unique(a), ua) 784 assert_equal(np.unique(a, return_index=True), (ua, ua_idx)) 785 assert_equal(np.unique(a, return_inverse=True), (ua, ua_inv)) 786 assert_equal(np.unique(a, return_counts=True), (ua, ua_cnt)) 787 788 # test for gh-19300 789 all_nans = [np.nan] * 4 790 ua = [np.nan] 791 ua_idx = [0] 792 ua_inv = [0, 0, 0, 0] 793 ua_cnt = [4] 794 assert_equal(np.unique(all_nans), ua) 795 assert_equal(np.unique(all_nans, return_index=True), (ua, ua_idx)) 796 assert_equal(np.unique(all_nans, return_inverse=True), (ua, ua_inv)) 797 assert_equal(np.unique(all_nans, return_counts=True), (ua, ua_cnt)) 798 799 def test_unique_axis_errors(self): 800 assert_raises(np.AxisError, unique, np.arange(10), axis=2) 801 assert_raises(np.AxisError, unique, np.arange(10), axis=-2) 802 803 def test_unique_axis_list(self): 804 msg = "Unique failed on list of lists" 805 inp = [[0, 1, 0], [0, 1, 0]] 806 inp_arr = np.asarray(inp) 807 assert_array_equal(unique(inp, axis=0), unique(inp_arr, axis=0), msg) 808 assert_array_equal(unique(inp, axis=1), unique(inp_arr, axis=1), msg) 809 810 @xpassIfTorchDynamo # _run_axis_tests xfails with the message 811 # torch has different unique ordering behaviour" 812 def test_unique_axis(self): 813 types = [] 814 types.extend(np.typecodes["AllInteger"]) 815 types.extend(np.typecodes["AllFloat"]) 816 817 for dtype in types: 818 self._run_axis_tests(dtype) 819 820 msg = "Non-bitwise-equal booleans test failed" 821 data = np.arange(10, dtype=np.uint8).reshape(-1, 2).view(bool) 822 result = np.array([[False, True], [True, True]], dtype=bool) 823 assert_array_equal(unique(data, axis=0), result, msg) 824 825 msg = "Negative zero equality test failed" 826 data = np.array([[-0.0, 0.0], [0.0, -0.0], [-0.0, 0.0], [0.0, -0.0]]) 827 result = np.array([[-0.0, 0.0]]) 828 assert_array_equal(unique(data, axis=0), result, msg) 829 830 @parametrize("axis", [0, -1]) 831 def test_unique_1d_with_axis(self, axis): 832 x = np.array([4, 3, 2, 3, 2, 1, 2, 2]) 833 uniq = unique(x, axis=axis) 834 assert_array_equal(uniq, [1, 2, 3, 4]) 835 836 @xpassIfTorchDynamo # (reason="unique / return_index") 837 def test_unique_axis_zeros(self): 838 # issue 15559 839 single_zero = np.empty(shape=(2, 0), dtype=np.int8) 840 uniq, idx, inv, cnt = unique( 841 single_zero, 842 axis=0, 843 return_index=True, 844 return_inverse=True, 845 return_counts=True, 846 ) 847 848 # there's 1 element of shape (0,) along axis 0 849 assert_equal(uniq.dtype, single_zero.dtype) 850 assert_array_equal(uniq, np.empty(shape=(1, 0))) 851 assert_array_equal(idx, np.array([0])) 852 assert_array_equal(inv, np.array([0, 0])) 853 assert_array_equal(cnt, np.array([2])) 854 855 # there's 0 elements of shape (2,) along axis 1 856 uniq, idx, inv, cnt = unique( 857 single_zero, 858 axis=1, 859 return_index=True, 860 return_inverse=True, 861 return_counts=True, 862 ) 863 864 assert_equal(uniq.dtype, single_zero.dtype) 865 assert_array_equal(uniq, np.empty(shape=(2, 0))) 866 assert_array_equal(idx, np.array([])) 867 assert_array_equal(inv, np.array([])) 868 assert_array_equal(cnt, np.array([])) 869 870 # test a "complicated" shape 871 shape = (0, 2, 0, 3, 0, 4, 0) 872 multiple_zeros = np.empty(shape=shape) 873 for axis in range(len(shape)): 874 expected_shape = list(shape) 875 if shape[axis] == 0: 876 expected_shape[axis] = 0 877 else: 878 expected_shape[axis] = 1 879 880 assert_array_equal( 881 unique(multiple_zeros, axis=axis), np.empty(shape=expected_shape) 882 ) 883 884 def test_unique_sort_order_with_axis(self): 885 # These tests fail if sorting along axis is done by treating subarrays 886 # as unsigned byte strings. See gh-10495. 887 fmt = "sort order incorrect for integer type '%s'" 888 for dt in "bhil": 889 a = np.array([[-1], [0]], dt) 890 b = np.unique(a, axis=0) 891 assert_array_equal(a, b, fmt % dt) 892 893 def _run_axis_tests(self, dtype): 894 data = np.array( 895 [[0, 1, 0, 0], [1, 0, 0, 0], [0, 1, 0, 0], [1, 0, 0, 0]] 896 ).astype(dtype) 897 898 msg = "Unique with 1d array and axis=0 failed" 899 result = np.array([0, 1]) 900 assert_array_equal(unique(data), result.astype(dtype), msg) 901 902 msg = "Unique with 2d array and axis=0 failed" 903 result = np.array([[0, 1, 0, 0], [1, 0, 0, 0]]) 904 assert_array_equal(unique(data, axis=0), result.astype(dtype), msg) 905 906 msg = "Unique with 2d array and axis=1 failed" 907 result = np.array([[0, 0, 1], [0, 1, 0], [0, 0, 1], [0, 1, 0]]) 908 assert_array_equal(unique(data, axis=1), result.astype(dtype), msg) 909 910 # e.g. 911 # 912 # >>> x = np.array([[[1, 1], [0, 1]], [[1, 0], [0, 0]]]) 913 # >>> np.unique(x, axis=2) 914 # [[1, 1], [0, 1]], [[1, 0], [0, 0]] 915 # >>> torch.unique(torch.as_tensor(x), dim=2) 916 # [[1, 1], [1, 0]], [[0, 1], [0, 0]] 917 # 918 msg = "Unique with 3d array and axis=2 failed" 919 data3d = np.array([[[1, 1], [1, 0]], [[0, 1], [0, 0]]]).astype(dtype) 920 result = np.take(data3d, [1, 0], axis=2) 921 assert_array_equal(unique(data3d, axis=2), result, msg) 922 923 uniq, idx, inv, cnt = unique( 924 data, axis=0, return_index=True, return_inverse=True, return_counts=True 925 ) 926 msg = "Unique's return_index=True failed with axis=0" 927 assert_array_equal(data[idx], uniq, msg) 928 msg = "Unique's return_inverse=True failed with axis=0" 929 assert_array_equal(uniq[inv], data) 930 msg = "Unique's return_counts=True failed with axis=0" 931 assert_array_equal(cnt, np.array([2, 2]), msg) 932 933 uniq, idx, inv, cnt = unique( 934 data, axis=1, return_index=True, return_inverse=True, return_counts=True 935 ) 936 msg = "Unique's return_index=True failed with axis=1" 937 assert_array_equal(data[:, idx], uniq) 938 msg = "Unique's return_inverse=True failed with axis=1" 939 assert_array_equal(uniq[:, inv], data) 940 msg = "Unique's return_counts=True failed with axis=1" 941 assert_array_equal(cnt, np.array([2, 1, 1]), msg) 942 943 @skipIf(True, reason="NP_VER: fails on CI with older NumPy") 944 @xpassIfTorchDynamo # (reason="unique / return_index / nans") 945 def test_unique_nanequals(self): 946 # issue 20326 947 a = np.array([1, 1, np.nan, np.nan, np.nan]) 948 unq = np.unique(a) 949 not_unq = np.unique(a, equal_nan=False) 950 assert_array_equal(unq, np.array([1, np.nan])) 951 assert_array_equal(not_unq, np.array([1, np.nan, np.nan, np.nan])) 952 953 954if __name__ == "__main__": 955 run_tests() 956