def prepare_interpolator(nchunks=4):
    """
    Check that all the inputs have the correct shape, and that all values are
    finite. It also split the destination grid in  `nchunks` parts, and process each
    part independently.
    """

    def _preamble_interpolation(interpolator):
        @wraps(interpolator)
        def _interpolator_with_preamble(xy_coord, values, xgrid, ygrid, **kwargs):
            nonlocal nchunks  # https://stackoverflow.com/questions/5630409/

            values = values.copy()
            xy_coord = xy_coord.copy()

            input_ndims = values.ndim
            input_nvars = 1 if input_ndims == 1 else values.shape[1]
            input_nsamples = values.shape[0]

            coord_ndims = xy_coord.ndim
            coord_nsamples = xy_coord.shape[0]

            grid_shape = (ygrid.size, xgrid.size)

            if np.any(~np.isfinite(values)):
                raise ValueError("argument 'values' contains non-finite values")
            if np.any(~np.isfinite(xy_coord)):
                raise ValueError("argument 'xy_coord' contains non-finite values")

            if input_ndims > 2:
                raise ValueError(
                    "argument 'values' must have 1 (n) or 2 dimensions (n, m), "
                    f"but it has {input_ndims}"
                )
            if not coord_ndims == 2:
                raise ValueError(
                    "argument 'xy_coord' must have 2 dimensions (n, 2), "
                    f"but it has {coord_ndims}"
                )

            if not input_nsamples == coord_nsamples:
                raise ValueError(
                    "the number of samples in argument 'values' does not match the "
                    f"number of coordinates {input_nsamples}!={coord_nsamples}"
                )

            # only one sample, return uniform output
            if input_nsamples == 1:
                output_array = np.ones((input_nvars,) + grid_shape)
                for n, v in enumerate(values[0, ...]):
                    output_array[n, ...] *= v
                return output_array.squeeze()

            # all equal elements, return uniform output
            if values.max() == values.min():
                return np.ones((input_nvars,) + grid_shape) * values.ravel()[0]

            # split grid in n chunks
            nchunks = int(kwargs.get("nchunks", nchunks) ** 0.5)
            if nchunks > 1:
                subxgrids = np.array_split(xgrid, nchunks)
                subxgrids = [x for x in subxgrids if x.size > 0]
                subygrids = np.array_split(ygrid, nchunks)
                subygrids = [y for y in subygrids if y.size > 0]

                # generate a unique identifier to be used for caching
                # intermediate results
                kwargs["hkey"] = uuid.uuid1().int
            else:
                subxgrids = [xgrid]
                subygrids = [ygrid]

            interpolated = np.zeros((input_nvars,) + grid_shape)
            indx = 0
            for subxgrid in subxgrids:
                deltax = subxgrid.size
                indy = 0
                for subygrid in subygrids:
                    deltay = subygrid.size
                    interpolated[:, indy : (indy + deltay), indx : (indx + deltax)] = (
                        interpolator(xy_coord, values, subxgrid, subygrid, **kwargs)
                    )
                    indy += deltay
                indx += deltax

            return interpolated.squeeze()

        extra_kwargs_doc = """
            nchunks: int, optional
                Split and process the destination grid in nchunks.
                Useful for large grids to limit the memory footprint.
            """

        _add_extra_kwrds_to_docstrings(_interpolator_with_preamble, extra_kwargs_doc)

        return _interpolator_with_preamble

    return _preamble_interpolation
Check that all the inputs have the correct shape, and that all values are
finite. It also split the destination grid in  `nchunks` parts, and process each
part independently.
pysteps/decorators.py