# Source code for dipy.denoise.non_local_means

    from numbers import Number

    import numpy as np

    from dipy.denoise.nlmeans_block import nlmeans_block
    from dipy.testing.decorators import warning_for_keywords

    @warning_for_keywords()
    def non_local_means(
        arr, sigma, *, mask=None, patch_radius=1, block_radius=5, rician=True
    ):
        r"""Non-local means for denoising 3D and 4D images, using blockwise
        averaging approach.

        See :footcite:p:`Coupe2008` and :footcite:p:`Coupe2012` for further details
        about the method.

        Parameters
        ----------
        arr : 3D or 4D ndarray
            The array to be denoised
        mask : 3D ndarray
            Mask on data where the non-local means will be applied.
        sigma : float or ndarray
            standard deviation of the noise estimated from the data
        patch_radius : int
            patch size is ``2 x patch_radius + 1``. Default is 1.
        block_radius : int
            block size is ``2 x block_radius + 1``. Default is 5.
        rician : boolean
            If True the noise is estimated as Rician, otherwise Gaussian noise
            is assumed.

        Returns
        -------
        denoised_arr : ndarray
            the denoised ``arr`` which has the same shape as ``arr``.

        References
        ----------
        .. footbibliography::

        """
        if isinstance(sigma, np.ndarray) and sigma.size == 1:
            sigma = sigma.item()
        if isinstance(sigma, np.ndarray):
            if arr.ndim == 3:
                raise ValueError("sigma should be a scalar for 3D data", sigma)
            if not np.issubdtype(sigma.dtype, np.number):
                raise ValueError("sigma should be an array of floats", sigma)
            if arr.ndim == 4 and sigma.ndim != 1:
                raise ValueError("sigma should be a 1D array for 4D data", sigma)
            if arr.ndim == 4 and sigma.shape[0] != arr.shape[-1]:
                raise ValueError(
                    "sigma should have the same length as the last "
                    "dimension of arr for 4D data",
                    sigma,
                )
        else:
            if not isinstance(sigma, Number):
                raise ValueError("sigma should be a float", sigma)
            # if sigma is a scalar and arr is 4D, we assume the same sigma for all
            if arr.ndim == 4:
                sigma = np.array([sigma] * arr.shape[-1])

        if mask is None and arr.ndim > 2:
            mask = np.ones((arr.shape[0], arr.shape[1], arr.shape[2]), dtype="f8")
        else:
            mask = np.ascontiguousarray(mask, dtype="f8")

        if mask.ndim != 3:
            raise ValueError("mask needs to be a 3D ndarray", mask.shape)

        if arr.ndim == 3:
            return np.array(
                nlmeans_block(
                    np.double(arr), mask, patch_radius, block_radius, sigma, int(rician)
                )
            ).astype(arr.dtype)
        elif arr.ndim == 4:
            denoised_arr = np.zeros_like(arr)
            for i in range(arr.shape[-1]):
                denoised_arr[..., i] = np.array(
                    nlmeans_block(
                        np.double(arr[..., i]),
                        mask,
                        patch_radius,
                        block_radius,
                        sigma[i],
                        int(rician),
                    )
                ).astype(arr.dtype)

            return denoised_arr

        else:
            raise ValueError("Only 3D or 4D array are supported!", arr.shape)