| 472 | |
| 473 | |
| 474 | def test_where(): |
| 475 | rng = np.random.default_rng() |
| 476 | # Copied and adapted from the da.where test. |
| 477 | x = rng.integers(10, size=(15, 14)) |
| 478 | mask = rng.choice(a=[False, True], size=(15, 14), p=[0.5, 0.5]) |
| 479 | x[5, 5] = x[4, 4] = 0 # Ensure some false elements |
| 480 | d = da.ma.masked_array(x, mask=mask, chunks=(4, 5)) |
| 481 | x = np.ma.array(x, mask=mask) |
| 482 | y = rng.integers(10, size=15).astype(np.uint8) |
| 483 | e = da.from_array(y, chunks=(4,)) |
| 484 | |
| 485 | # Nonzero test |
| 486 | sol = np.ma.where(x) |
| 487 | res = da.ma.where(d) |
| 488 | for i in range(len(sol)): |
| 489 | assert_eq(res[i], sol[i]) |
| 490 | |
| 491 | for c1, c2 in [ |
| 492 | (d > 5, x > 5), |
| 493 | (d, x), |
| 494 | (1, 1), |
| 495 | (5, 5), |
| 496 | (True, True), |
| 497 | (np.True_, np.True_), |
| 498 | (0, 0), |
| 499 | (False, False), |
| 500 | (np.False_, np.False_), |
| 501 | ]: |
| 502 | for b1, b2 in [(0, 0), (-e[:, None], -y[:, None]), (e[:14], y[:14])]: |
| 503 | w1 = da.ma.where(c1, d, b1) |
| 504 | w2 = np.ma.where(c2, x, b2) |
| 505 | assert_eq(w1, w2) |