diff --git a/cubed/overlap.py b/cubed/overlap.py index d93c2b89..8ae2785c 100644 --- a/cubed/overlap.py +++ b/cubed/overlap.py @@ -134,6 +134,6 @@ def _pad_boundaries(x, depth, boundary, numblocks, block_id): p = nxp.full_like(x, fill_value=boundary[i], shape=pad_shape) if block_id[i] == 0: # first block on axis i x = nxp.concat([p, x], axis=i) - elif block_id[i] == numblocks[i] - 1: # last block on axis i + if block_id[i] == numblocks[i] - 1: # last block on axis i x = nxp.concat([x, p], axis=i) return x diff --git a/cubed/tests/test_overlap.py b/cubed/tests/test_overlap.py index 0d8e3a18..5c29e3a7 100644 --- a/cubed/tests/test_overlap.py +++ b/cubed/tests/test_overlap.py @@ -22,6 +22,23 @@ def test_map_overlap_1d(): assert_array_equal(b.compute(), np.array([0, 0, 1, 2, 3, 2, 3, 4, 5, 0])) +def test_map_overlap_1d_single_chunk(): + x = np.arange(6) + a = xp.asarray(x, chunks=(6,)) + + b = cubed.map_overlap( + lambda x: x, + a, + dtype=a.dtype, + chunks=((8,),), + depth=1, + boundary=0, + trim=False, + ) + + assert_array_equal(b.compute(), np.array([0, 0, 1, 2, 3, 4, 5, 0])) + + def test_map_overlap_2d(): x = np.arange(36).reshape((6, 6)) a = xp.asarray(x, chunks=(3, 3))