| 312 | |
| 313 | |
| 314 | def test_map_overlap(): |
| 315 | x = da.arange(10, chunks=5) |
| 316 | y = x.map_overlap(lambda x: x + len(x), depth=2, dtype=x.dtype, boundary="reflect") |
| 317 | assert_eq(y, np.arange(10) + 5 + 2 + 2) |
| 318 | |
| 319 | x = da.arange(10, chunks=5) |
| 320 | y = x.map_overlap( |
| 321 | lambda x: x + len(x), depth=np.int64(2), dtype=x.dtype, boundary="reflect" |
| 322 | ) |
| 323 | assert all([(type(s) is int) for s in y.shape]) |
| 324 | assert_eq(y, np.arange(10) + 5 + 2 + 2) |
| 325 | |
| 326 | x = da.ones((10, 10), chunks=(3, 4)) |
| 327 | z = x.map_overlap(lambda x: x, depth={1: 5}, boundary="reflect") |
| 328 | assert z.chunks[0] == x.chunks[0] # don't rechunk the first dimension |
| 329 | |
| 330 | x = np.arange(16).reshape((4, 4)) |
| 331 | d = da.from_array(x, chunks=(2, 2)) |
| 332 | exp1 = d.map_overlap( |
| 333 | lambda x: x + x.size, depth=1, dtype=d.dtype, boundary="reflect" |
| 334 | ) |
| 335 | exp2 = d.map_overlap( |
| 336 | lambda x: x + x.size, |
| 337 | depth={0: 1, 1: 1}, |
| 338 | boundary={0: "reflect", 1: "none"}, |
| 339 | dtype=d.dtype, |
| 340 | ) |
| 341 | exp3 = d.map_overlap( |
| 342 | lambda x: x + x.size, depth={1: 1}, boundary={1: "reflect"}, dtype=d.dtype |
| 343 | ) |
| 344 | exp4 = d.map_overlap( |
| 345 | lambda x: x + x.size, |
| 346 | depth={1: (1, 0)}, |
| 347 | boundary={0: "none", 1: "none"}, |
| 348 | dtype=d.dtype, |
| 349 | ) |
| 350 | assert_eq(exp1, x + 16) |
| 351 | assert_eq(exp2, x + 12) |
| 352 | assert_eq(exp3, x + 8) |
| 353 | assert_eq( |
| 354 | exp4, |
| 355 | np.block( |
| 356 | [[x[0:2, 0:2] + 4, x[0:2, 2:4] + 6], [x[2:4, 0:2] + 4, x[2:4, 2:4] + 6]] |
| 357 | ), |
| 358 | ) |
| 359 | |
| 360 | |
| 361 | def test_map_overlap_escapes_to_map_blocks_when_depth_is_zero(): |