diff --git a/pysteps/tests/test_utils_dimension.py b/pysteps/tests/test_utils_dimension.py index ab753ed7..a6788fbf 100644 --- a/pysteps/tests/test_utils_dimension.py +++ b/pysteps/tests/test_utils_dimension.py @@ -220,6 +220,64 @@ def test_aggregate_fields_space(R, metadata, space_window, ignore_nan, expected) (2, 4, 2, 4), np.ones((2, 2)), ), + # Clip a box that lies strictly inside the source domain. This used to raise + # a broadcasting ValueError because the source and destination windows had + # mismatched sizes (see GH issue on clip_domain). + ( + np.arange(16).reshape(4, 4).astype(float), + { + "x1": 0, + "x2": 4, + "y1": 0, + "y2": 4, + "xpixelsize": 1, + "ypixelsize": 1, + "zerovalue": 0, + "yorigin": "upper", + }, + (1, 3, 1, 3), + np.array([[5.0, 6.0], [9.0, 10.0]]), + ), + # Same interior clip with yorigin="lower": row 0 is the bottom of the domain. + ( + np.arange(16).reshape(4, 4).astype(float), + { + "x1": 0, + "x2": 4, + "y1": 0, + "y2": 4, + "xpixelsize": 1, + "ypixelsize": 1, + "zerovalue": 0, + "yorigin": "lower", + }, + (1, 3, 1, 3), + np.array([[5.0, 6.0], [9.0, 10.0]]), + ), + # An extent larger than the source domain is padded with the zero value + # instead of raising an error. + ( + np.ones((2, 2)), + { + "x1": 0, + "x2": 2, + "y1": 0, + "y2": 2, + "xpixelsize": 1, + "ypixelsize": 1, + "zerovalue": -1, + "yorigin": "upper", + }, + (-1, 3, -1, 3), + np.array( + [ + [-1.0, -1.0, -1.0, -1.0], + [-1.0, 1.0, 1.0, -1.0], + [-1.0, 1.0, 1.0, -1.0], + [-1.0, -1.0, -1.0, -1.0], + ] + ), + ), ] @@ -229,6 +287,30 @@ def test_clip_domain(R, metadata, extent, expected): assert_array_equal(dimension.clip_domain(R, metadata, extent)[0], expected) +def test_clip_domain_metadata(): + """The clipped metadata must describe the returned pixel grid.""" + R = np.arange(16).reshape(4, 4).astype(float) + metadata = { + "x1": 0, + "x2": 4, + "y1": 0, + "y2": 4, + "xpixelsize": 1, + "ypixelsize": 1, + "zerovalue": 0, + "yorigin": "upper", + } + R_clip, metadata_clip = dimension.clip_domain(R, metadata, (1, 3, 1, 3)) + assert R_clip.shape == (2, 2) + assert metadata_clip["x1"] == 1 + assert metadata_clip["x2"] == 3 + assert metadata_clip["y1"] == 1 + assert metadata_clip["y2"] == 3 + # the pixel size must be preserved + assert metadata_clip["xpixelsize"] == metadata["xpixelsize"] + assert metadata_clip["ypixelsize"] == metadata["ypixelsize"] + + # square_domain R = np.zeros((4, 2)) test_data = [ diff --git a/pysteps/utils/dimension.py b/pysteps/utils/dimension.py index 43b7e2ca..2d7b3fa6 100644 --- a/pysteps/utils/dimension.py +++ b/pysteps/utils/dimension.py @@ -375,75 +375,69 @@ def clip_domain(R, metadata, extent=None): if extent is None: return R, metadata - if len(R.shape) < 2: + if R.ndim < 2: raise ValueError("The number of dimension must be > 1") - if len(R.shape) == 2: + if R.ndim > 4: + raise ValueError("The number of dimension must be <= 4") + if R.ndim == 2: R = R[None, None, :, :] - if len(R.shape) == 3: + elif R.ndim == 3: R = R[None, :, :, :] - if len(R.shape) > 4: - raise ValueError("The number of dimension must be <= 4") - # extract original domain coordinates - left = metadata["x1"] - right = metadata["x2"] - bottom = metadata["y1"] - top = metadata["y2"] + # extract original domain geometry + x1 = metadata["x1"] + y1 = metadata["y1"] + y2 = metadata["y2"] + xpixelsize = metadata["xpixelsize"] + ypixelsize = metadata["ypixelsize"] # extract bounding box coordinates - left_ = extent[0] - right_ = extent[1] - bottom_ = extent[2] - top_ = extent[3] - - # compute its extent in pixels - dim_x_ = int((right_ - left_) / metadata["xpixelsize"]) - dim_y_ = int((top_ - bottom_) / metadata["ypixelsize"]) - R_ = np.ones((R.shape[0], R.shape[1], dim_y_, dim_x_)) * metadata["zerovalue"] - - # build set of coordinates for the original domain - y_coord = ( - np.linspace(bottom, top - metadata["ypixelsize"], R.shape[2]) - + metadata["ypixelsize"] / 2.0 - ) - x_coord = ( - np.linspace(left, right - metadata["xpixelsize"], R.shape[3]) - + metadata["xpixelsize"] / 2.0 - ) - - # build set of coordinates for the new domain - y_coord_ = ( - np.linspace(bottom_, top_ - metadata["ypixelsize"], R_.shape[2]) - + metadata["ypixelsize"] / 2.0 - ) - x_coord_ = ( - np.linspace(left_, right_ - metadata["xpixelsize"], R_.shape[3]) - + metadata["xpixelsize"] / 2.0 - ) - - # origin='upper' reverses the vertical axes direction + left_, right_, bottom_, top_ = extent + + # size of the clipped domain in pixels + dim_x = int(round((right_ - left_) / xpixelsize)) + dim_y = int(round((top_ - bottom_) / ypixelsize)) + if dim_x <= 0 or dim_y <= 0: + raise ValueError("The requested extent has a non-positive size") + + # Clipping preserves the pixel size, so it is an integer-pixel window + # operation. Using a single pixel offset between the clipped grid and the + # source array guarantees that the source and destination windows always + # have matching sizes, whether the extent lies inside the source domain or + # extends beyond it. + col_offset = int(round((left_ - x1) / xpixelsize)) + # The vertical origin determines which corner corresponds to row 0. if metadata["yorigin"] == "upper": - y_coord = y_coord[::-1] - y_coord_ = y_coord_[::-1] - - # extract original domain - idx_y = np.where(np.logical_and(y_coord < top_, y_coord > bottom_))[0] - idx_x = np.where(np.logical_and(x_coord < right_, x_coord > left_))[0] - - # extract new domain - idx_y_ = np.where(np.logical_and(y_coord_ < top, y_coord_ > bottom))[0] - idx_x_ = np.where(np.logical_and(x_coord_ < right, x_coord_ > left))[0] - - # compose the new array - R_[:, :, idx_y_[0] : (idx_y_[-1] + 1), idx_x_[0] : (idx_x_[-1] + 1)] = R[ - :, :, idx_y[0] : (idx_y[-1] + 1), idx_x[0] : (idx_x[-1] + 1) - ] - - # update coordinates - metadata["y1"] = bottom_ - metadata["y2"] = top_ - metadata["x1"] = left_ - metadata["x2"] = right_ + row_offset = int(round((y2 - top_) / ypixelsize)) + else: + row_offset = int(round((bottom_ - y1) / ypixelsize)) + + # Cells that fall outside the source domain keep the zero value. + R_ = np.ones((R.shape[0], R.shape[1], dim_y, dim_x)) * metadata["zerovalue"] + + # overlap window between the source array and the clipped domain + dst_x0 = max(0, -col_offset) + dst_x1 = min(dim_x, R.shape[3] - col_offset) + dst_y0 = max(0, -row_offset) + dst_y1 = min(dim_y, R.shape[2] - row_offset) + + if dst_x1 > dst_x0 and dst_y1 > dst_y0: + R_[:, :, dst_y0:dst_y1, dst_x0:dst_x1] = R[ + :, + :, + dst_y0 + row_offset : dst_y1 + row_offset, + dst_x0 + col_offset : dst_x1 + col_offset, + ] + + # update coordinates so that they stay consistent with the pixel grid + metadata["x1"] = x1 + col_offset * xpixelsize + metadata["x2"] = metadata["x1"] + dim_x * xpixelsize + if metadata["yorigin"] == "upper": + metadata["y2"] = y2 - row_offset * ypixelsize + metadata["y1"] = metadata["y2"] - dim_y * ypixelsize + else: + metadata["y1"] = y1 + row_offset * ypixelsize + metadata["y2"] = metadata["y1"] + dim_y * ypixelsize R_shape[-2] = R_.shape[-2] R_shape[-1] = R_.shape[-1]