Python dask.array.map_blocks() Examples

def atm_variables_finder(mus, muv, phi, height, tau, tO3, tH2O, taustep4sphalb, tO2=1.0):
    tau_step = da.linspace(taustep4sphalb, MAXNUMSPHALBVALUES * taustep4sphalb, MAXNUMSPHALBVALUES,
                           chunks=int(MAXNUMSPHALBVALUES / 2))
    sphalb0 = csalbr(tau_step)
    taur = tau * da.exp(-height / SCALEHEIGHT)
    rhoray, trdown, trup = chand(phi, muv, mus, taur)
    if isinstance(height, xr.DataArray):
        sphalb = da.map_blocks(_sphalb_index, (taur / taustep4sphalb + 0.5).astype(np.int32).data, sphalb0.compute(),
                               dtype=sphalb0.dtype)
    else:
        sphalb = sphalb0[(taur / taustep4sphalb + 0.5).astype(np.int32)]
    Ttotrayu = ((2 / 3. + muv) + (2 / 3. - muv) * trup) / (4 / 3. + taur)
    Ttotrayd = ((2 / 3. + mus) + (2 / 3. - mus) * trdown) / (4 / 3. + taur)
    TtotraytH2O = Ttotrayu * Ttotrayd * tH2O
    tOG = tO3 * tO2
    return sphalb, rhoray, TtotraytH2O, tOG 

def aggregate(d, y_size, x_size):
        """Average every 4 elements (2x2) in a 2D array."""
        if d.ndim != 2:
            # we can't guarantee what blocks we are getting and how
            # it should be reshaped to do the averaging.
            raise ValueError("Can't aggregrate (reduce) data arrays with "
                             "more than 2 dimensions.")
        if not (x_size.is_integer() and y_size.is_integer()):
            raise ValueError("Aggregation factors are not integers")
        for agg_size, chunks in zip([y_size, x_size], d.chunks):
            for chunk_size in chunks:
                if chunk_size % agg_size != 0:
                    raise ValueError("Aggregation requires arrays with "
                                     "shapes and chunks divisible by the "
                                     "factor")

        new_chunks = (tuple(int(x / y_size) for x in d.chunks[0]),
                      tuple(int(x / x_size) for x in d.chunks[1]))
        return da.core.map_blocks(_mean, d, y_size, x_size, dtype=d.dtype, chunks=new_chunks) 

def interpolate_xarray_linear(xpoints, ypoints, values, shape, chunks=CHUNK_SIZE):
    """Interpolate linearly, generating a dask array."""
    from scipy.interpolate.interpnd import (LinearNDInterpolator,
                                            _ndim_coords_from_arrays)

    if isinstance(chunks, (list, tuple)):
        vchunks, hchunks = chunks
    else:
        vchunks, hchunks = chunks, chunks

    points = _ndim_coords_from_arrays(np.vstack((np.asarray(ypoints),
                                                 np.asarray(xpoints))).T)

    interpolator = LinearNDInterpolator(points, values)

    grid_x, grid_y = da.meshgrid(da.arange(shape[1], chunks=hchunks),
                                 da.arange(shape[0], chunks=vchunks))

    # workaround for non-thread-safe first call of the interpolator:
    interpolator((0, 0))
    res = da.map_blocks(intp, grid_x, grid_y, interpolator=interpolator)

    return DataArray(res, dims=('y', 'x')) 

def interpolate_angles(self, angles, resolution):
        # FIXME: interpolate in cartesian coordinates if the lons or lats are
        # problematic
        from geotiepoints.multilinear import MultilinearInterpolator

        geocoding = self.root.find('.//Tile_Geocoding')
        rows = int(geocoding.find('Size[@resolution="' + str(resolution) + '"]/NROWS').text)
        cols = int(geocoding.find('Size[@resolution="' + str(resolution) + '"]/NCOLS').text)

        smin = [0, 0]
        smax = np.array(angles.shape) - 1
        orders = angles.shape
        minterp = MultilinearInterpolator(smin, smax, orders)
        minterp.set_values(da.atleast_2d(angles.ravel()))

        x = da.arange(rows, dtype=angles.dtype, chunks=CHUNK_SIZE) / (rows-1) * (angles.shape[0] - 1)
        y = da.arange(cols, dtype=angles.dtype, chunks=CHUNK_SIZE) / (cols-1) * (angles.shape[1] - 1)
        xcoord, ycoord = da.meshgrid(x, y)
        return da.map_blocks(self._do_interp, minterp, xcoord, ycoord, dtype=angles.dtype, chunks=xcoord.chunks) 

def _get_solar_flux(self, band):
        """Get the solar flux for the band."""
        solar_flux = self.cal['solar_flux'].isel(bands=band).values
        d_index = self.cal['detector_index'].fillna(0).astype(int)

        return da.map_blocks(self._get_items, d_index.data,
                             solar_flux=solar_flux, dtype=solar_flux.dtype) 

def expansion_coefs(self):
        """Compute the expansion coefficients."""
        if self._expansion_coefs is not None:
            return self._expansion_coefs
        v_track = (np.arange(self.scans * self.scan_size) % self.scan_size + self.track_offset) / self.scan_size
        self.tpz_sizes = self.tpz_sizes.persist()
        self.nb_tpzs = self.nb_tpzs.persist()
        col_chunks = (self.tpz_sizes * self.nb_tpzs).compute()
        self._expansion_coefs = da.map_blocks(self.get_coefs, self.c_align, self.c_exp, self.tpz_sizes, self.nb_tpzs,
                                              dtype=np.float64, v_track=v_track,  new_axis=[0, 2],
                                              chunks=(self.scans * self.scan_size,
                                                      tuple(col_chunks), 4))

        return self._expansion_coefs 

def expand_angle_and_nav(self, arrays):
        """Expand angle and navigation datasets."""
        res = []
        for array in arrays:
            res.append(da.map_blocks(self.expand, array[:, :, np.newaxis], self.expansion_coefs,
                                     dtype=array.dtype, drop_axis=2, chunks=self.expansion_coefs.chunks[:-1]))
        return res 

def aggregate(self, **dims):
        """Aggregate the current swath definition by averaging.

        For example, averaging over 2x2 windows:
        `sd.aggregate(x=2, y=2)`
        """
        import pyproj
        import dask.array as da

        geocent = pyproj.Proj(proj='geocent')
        latlong = pyproj.Proj(proj='latlong')
        res = da.map_blocks(self._do_transform, latlong, geocent,
                            self.lons.data, self.lats.data,
                            da.zeros_like(self.lons.data), new_axis=[2],
                            chunks=(self.lons.chunks[0], self.lons.chunks[1], 3))
        res = DataArray(res, dims=['y', 'x', 'coord'], coords=self.lons.coords)
        res = res.coarsen(**dims).mean()
        lonlatalt = da.map_blocks(self._do_transform, geocent, latlong,
                                  res[:, :, 0].data, res[:, :, 1].data,
                                  res[:, :, 2].data, new_axis=[2],
                                  chunks=res.data.chunks)
        lons = DataArray(lonlatalt[:, :, 0], dims=self.lons.dims,
                         coords=res.coords, attrs=self.lons.attrs.copy())
        lats = DataArray(lonlatalt[:, :, 1], dims=self.lons.dims,
                         coords=res.coords, attrs=self.lons.attrs.copy())
        try:
            resolution = lons.attrs['resolution'] * ((dims.get('x', 1) + dims.get('y', 1)) / 2)
            lons.attrs['resolution'] = resolution
            lats.attrs['resolution'] = resolution
        except KeyError:
            pass
        return SwathDefinition(lons, lats) 

def _get_indices(self):
        """Calculate projection indices.

        Returns
        -------
        x_idxs : Dask array
            X indices of the target grid where the data are put
        y_idxs : Dask array
            Y indices of the target grid where the data are put
        """
        LOG.info("Determine bucket resampling indices")

        # Transform source lons/lats to target projection coordinates x/y
        lons = self.source_lons.ravel()
        lats = self.source_lats.ravel()
        result = da.map_blocks(self._get_proj_coordinates, lons, lats,
                               new_axis=0, chunks=(2,) + lons.chunks)
        proj_x = result[0, :]
        proj_y = result[1, :]

        # Calculate array indices. Orient so that 0-meridian is pointing down.
        adef = self.target_area
        x_res, y_res = adef.resolution
        x_idxs = da.floor((proj_x - adef.area_extent[0]) / x_res).astype(np.int)
        y_idxs = da.floor((adef.area_extent[3] - proj_y) / y_res).astype(np.int)

        # Get valid index locations
        mask = ((x_idxs >= 0) & (x_idxs < adef.width) &
                (y_idxs >= 0) & (y_idxs < adef.height))
        self.y_idxs = da.where(mask, y_idxs, -1)
        self.x_idxs = da.where(mask, x_idxs, -1)

        # Convert X- and Y-indices to raveled indexing
        target_shape = self.target_area.shape
        self.idxs = self.y_idxs * target_shape[1] + self.x_idxs 

def regrid_dask(self, indata):
        """See __call__()."""

        extra_chunk_shape = indata.chunksize[0:-2]

        output_chunk_shape = extra_chunk_shape + self.shape_out

        outdata = da.map_blocks(
            self.regrid_numpy,
            indata,
            dtype=float,
            chunks=output_chunk_shape
        )

        return outdata 

def detrend_wrap(detrend_func):
    """
    Wrapper function for `xrft.detrendn`.
    """
    def func(a, axes=None):
        if len(axes) > 3:
            raise ValueError("Detrending is only supported up to "
                            "3 dimensions.")
        if axes is None:
            axes = tuple(range(a.ndim))
        else:
            if len(set(axes)) < len(axes):
                raise ValueError("Duplicate axes are not allowed.")

        for each_axis in axes:
            if len(a.chunks[each_axis]) != 1:
                raise ValueError('The axis along the detrending is upon '
                                'cannot be chunked.')

        if len(axes) == 1:
            return dsar.map_blocks(sps.detrend, a, axis=axes[0],
                                   chunks=a.chunks, dtype=a.dtype
                                  )
        else:
            for each_axis in range(a.ndim):
                if each_axis not in axes:
                    if len(a.chunks[each_axis]) != a.shape[each_axis]:
                        raise ValueError("The axes other than ones to detrend "
                                        "over should have a chunk length of 1.")
            return dsar.map_blocks(detrend_func, a, axes,
                                   chunks=a.chunks, dtype=a.dtype
                                  )

    return func 

def _background_removal_median(dask_array, **kwargs):
    """Background removal using median filter.

    Parameters
    ----------
    dask_array : Dask array
        Must be at least 2 dimensions
    footprint : float

    Returns
    -------
    output_array : Dask array
        Same dimensions as dask_array

    Examples
    --------
    >>> import pyxem.utils.dask_tools as dt
    >>> import dask.array as da
    >>> s = pxm.dummy_data.dummy_data.get_cbed_signal()
    >>> dask_array = da.from_array(s.data+1e3, chunks=(5,5,25, 25))
    >>> s_rem = pxm.Diffraction2D(
    ...     dt._background_removal_median(dask_array), footprint=20)

    """
    dask_array_rechunked = _rechunk_signal2d_dim_one_chunk(dask_array)
    output_array = da.map_blocks(
        _background_removal_chunk_median,
        dask_array_rechunked,
        dtype=np.float32,
        **kwargs
    )
    return output_array 

def _background_removal_dog(dask_array, **kwargs):
    """Background removal using difference of Gaussians.

    Parameters
    ----------
    dask_array : Dask array
        At least 2 dimensions
    min_sigma : float
    max_sigma : float

    Returns
    -------
    output_array : Dask array
        Same dimensions as dask_array

    Examples
    --------
    >>> import pyxem.utils.dask_tools as dt
    >>> import dask.array as da
    >>> s = pxm.dummy_data.dummy_data.get_cbed_signal()
    >>> dask_array = da.from_array(s.data, chunks=(5,5,25, 25))
    >>> s_rem = pxm.Diffraction2D(
    ...     dt._background_removal_dog(dask_array))

    """
    dask_array_rechunked = _rechunk_signal2d_dim_one_chunk(dask_array)
    output_array = da.map_blocks(
        _background_removal_chunk_dog, dask_array_rechunked, dtype=np.float32, **kwargs
    )
    return output_array 

def precompute(self, cache_dir=None, swath_usage=0, **kwargs):
        """Generate row and column arrays and store it for later use."""
        if self.cache:
            # this resampler should be used for one SwathDefinition
            # no need to recompute ll2cr output again
            return None

        if kwargs.get('mask') is not None:
            LOG.warning("'mask' parameter has no affect during EWA "
                        "resampling")

        del kwargs
        source_geo_def = self.source_geo_def
        target_geo_def = self.target_geo_def

        if cache_dir:
            LOG.warning("'cache_dir' is not used by EWA resampling")

        # Satpy/PyResample don't support dynamic grids out of the box yet
        lons, lats = source_geo_def.get_lonlats()
        if isinstance(lons, xr.DataArray):
            # get dask arrays
            lons = lons.data
            lats = lats.data
        # we are remapping to a static unchanging grid/area with all of
        # its parameters specified
        chunks = (2,) + lons.chunks
        res = da.map_blocks(self._call_ll2cr, lons, lats,
                            target_geo_def, swath_usage,
                            dtype=lons.dtype, chunks=chunks, new_axis=[0])
        cols = res[0]
        rows = res[1]

        # save the dask arrays in the class instance cache
        # the on-disk cache will store the numpy arrays
        self.cache = {
            "rows": rows,
            "cols": cols,
        }

        return None 

def __ua_convert__(self, value, dispatch_type, coerce):
        if dispatch_type is not ufunc and value is None:
            return None

        if dispatch_type is ndarray:
            if not coerce and not isinstance(value, da.Array):
                return NotImplemented
            ret = da.asarray(value)
            with set_backend(self._inner):
                ret = ret.map_blocks(self._wrap_current_state(unumpy.asarray))

            return ret

        return value 

def wrap_map_blocks(self, func):
        @functools.wraps(func)
        def wrapped(*args, **kwargs):
            with set_backend(self._inner):
                return da.map_blocks(self._wrap_current_state(func), *args, **kwargs)

        return wrapped 

def HaloVelocityDispersion(mass, cosmo, redshift, mdef='vir'):
    """ Compute the velocity dispersion of halo from Mass.

        This is a simple model suggested by Martin White.

        See http://adsabs.harvard.edu/abs/2008ApJ...672..122E
    """

    mass, redshift = da.broadcast_arrays(mass, redshift)
    def compute_vdisp(mass, redshift):
        h = cosmo.efunc(redshift)
        return 1100. * (h * mass / 1e15) ** 0.33333

    return da.map_blocks(compute_vdisp, mass, redshift, dtype=mass.dtype) 

def _log_loss_inner(
    x: ArrayLike, y: ArrayLike, sample_weight: Optional[ArrayLike], **kwargs
):
    # da.map_blocks wasn't able to concatenate together the results
    # when we reduce down to a scalar per block. So we make an
    # array with 1 element.
    if sample_weight is not None:
        sample_weight = sample_weight.ravel()
    return np.array(
        [sklearn.metrics.log_loss(x, y, sample_weight=sample_weight, **kwargs)]
    ) 

def log_loss(
    y_true, y_pred, eps=1e-15, normalize=True, sample_weight=None, labels=None
):
    if not (dask.is_dask_collection(y_true) and dask.is_dask_collection(y_pred)):
        return sklearn.metrics.log_loss(
            y_true,
            y_pred,
            eps=eps,
            normalize=normalize,
            sample_weight=sample_weight,
            labels=labels,
        )

    if y_pred.ndim > 1 and y_true.ndim == 1:
        y_true = y_true.reshape(-1, 1)
        drop_axis: Optional[int] = 1
        if sample_weight is not None:
            sample_weight = sample_weight.reshape(-1, 1)
    else:
        drop_axis = None

    result = da.map_blocks(
        _log_loss_inner,
        y_true,
        y_pred,
        sample_weight,
        chunks=(1,),
        drop_axis=drop_axis,
        dtype="f8",
        eps=eps,
        normalize=normalize,
        labels=labels,
    )
    if normalize and sample_weight is not None:
        sample_weight = sample_weight.ravel()
        block_weights = sample_weight.map_blocks(np.sum, chunks=(1,), keepdims=True)
        return da.average(result, 0, weights=block_weights)
    elif normalize:
        return result.mean()
    else:
        return result.sum() 

def _compute_var_image_xarray_dask(src_var: xr.DataArray,
                                   dst_src_ij_images: np.ndarray,
                                   fill_value: Union[int, float, complex] = np.nan) -> da.Array:
    """Extract source pixels from xarray.DataArray source with dask.array.Array data"""
    return da.map_blocks(_compute_var_image_xarray_dask_block,
                         src_var.values,
                         dst_src_ij_images,
                         fill_value,
                         dtype=src_var.dtype,
                         drop_axis=0) 

def inverse_transform(self, y: Union[ArrayLike, SeriesType]):
        check_is_fitted(self, "classes_")
        y = self._check_array(y)

        if isinstance(y, da.Array):
            if getattr(self, "dtype_", None):
                # -> Series[category]
                if self.dtype_ is not None:
                    result = (
                        dd.from_dask_array(y)
                        .astype("category")
                        .cat.set_categories(np.arange(len(self.classes_)))
                        .cat.rename_categories(self.dtype_.categories)
                    )
                if self.dtype_.ordered:
                    result = result.cat.as_ordered()
                return result
            else:
                return da.map_blocks(
                    getitem,
                    self.classes_,
                    y,
                    dtype=self.classes_.dtype,
                    chunks=y.chunks,
                )
        else:
            y = np.asarray(y)
            if getattr(self, "dtype_", None):
                if self.dtype_ is not None:
                    return pd.Series(
                        pd.Categorical.from_codes(
                            y,
                            categories=self.dtype_.categories,
                            ordered=self.dtype_.ordered,
                        )
                    )
            else:
                return self.classes_[y] 

def _peak_refinement_centre_of_mass(dask_array, peak_array, square_size):
    """Refining the peak positions using the center of mass of the peaks

    Parameters
    ----------
    dask_array : Dask array
        The two last dimensions are the signal dimensions. Must have at least
        2 dimensions.
    peak_array : Dask array
    square_size : Even integer, this should be larger than the diffraction
        disc diameter. However not to large because then other discs will
        influence the center of mass.

    Returns
    -------
    peak_array : A dask array with the x and y positions of the refined peaks

    Examples
    --------
    >>> import pyxem.utils.dask_tools as dt
    >>> import dask.array as da
    >>> s = ps.dummy_data.get_cbed_signal()
    >>> dask_array = da.from_array(s.data, chunks=(5, 5, 25, 25))
    >>> peak_array = dt._peak_find_dog(dask_array)
    >>> r_p_array = dt._peak_refinement_centre_of_mass(
    ...     dask_array, peak_array, 20)
    >>> ref_peak_array = r_p_array.compute()
    """
    if dask_array.shape[:-2] != peak_array.shape:
        raise ValueError(
            "peak_array ({0}) must have the same shape as dask_array "
            "except the two last dimensions ({1})".format(
                peak_array.shape, dask_array.shape[:-2]
            )
        )
    if not hasattr(dask_array, "chunks"):
        raise ValueError("dask_array must be a Dask array")
    if not hasattr(peak_array, "chunks"):
        raise ValueError("peak_array must be a Dask array")

    dask_array_rechunked = _rechunk_signal2d_dim_one_chunk(dask_array)

    peak_array_rechunked = peak_array.rechunk(dask_array_rechunked.chunks[:-2])
    shape = list(peak_array_rechunked.shape)
    shape.extend([1, 1])
    peak_array_rechunked = peak_array_rechunked.reshape(shape)

    drop_axis = (dask_array_rechunked.ndim - 2, dask_array_rechunked.ndim - 1)
    kwargs_refinement = {"square_size": square_size}

    output_array = da.map_blocks(
        _peak_refinement_centre_of_mass_chunk,
        dask_array_rechunked,
        peak_array_rechunked,
        drop_axis=drop_axis,
        dtype=np.object,
        **kwargs_refinement
    )
    return output_array 

def _peak_find_log(dask_array, **kwargs):
    """Find peaks in a dask array using skimage's blob_log function.

    Parameters
    ----------
    dask_array : Dask array
        Must be at least 2 dimensions.
    min_sigma : float, optional
    max_sigma : float, optional
    num_sigma : float, optional
    threshold : float, optional
    overlap : float, optional
    normalize_value : float, optional
        All the values in dask_array will be divided by this value.
        If no value is specified, the max value in each individual image will
        be used.

    Returns
    -------
    peak_array : dask object array
        Same size as the two last dimensions in data.
        The peak positions themselves are stored in 2D NumPy arrays
        inside each position in peak_array. This is done instead of
        making a 4D NumPy array, since the number of found peaks can
        vary in each position.

    Example
    -------
    >>> s = pxm.dummy_data.dummy_data.get_cbed_signal()
    >>> import dask.array as da
    >>> dask_array = da.from_array(s.data, chunks=(5, 5, 25, 25))
    >>> import pyxem.utils.dask_tools as dt
    >>> peak_array = dt._peak_find_log(dask_array)
    >>> peak_array_computed = peak_array.compute()

    """
    dask_array_rechunked = _rechunk_signal2d_dim_one_chunk(dask_array)
    drop_axis = (dask_array_rechunked.ndim - 2, dask_array_rechunked.ndim - 1)
    output_array = da.map_blocks(
        _peak_find_log_chunk,
        dask_array_rechunked,
        drop_axis=drop_axis,
        dtype=np.object,
        **kwargs
    )
    return output_array 

def _peak_find_dog(dask_array, **kwargs):
    """Find peaks in a dask array using skimage's blob_dog function.

    Parameters
    ----------
    dask_array : Dask array
        Must be at least 2 dimensions.
    min_sigma : float, optional
    max_sigma : float, optional
    sigma_ratio : float, optional
    threshold : float, optional
    overlap : float, optional
    normalize_value : float, optional
        All the values in dask_array will be divided by this value.
        If no value is specified, the max value in each individual image will
        be used.

    Returns
    -------
    peak_array : dask object array
        Same size as the two last dimensions in data.
        The peak positions themselves are stored in 2D NumPy arrays
        inside each position in peak_array. This is done instead of
        making a 4D NumPy array, since the number of found peaks can
        vary in each position.

    Example
    -------
    >>> s = pxm.dummy_data.dummy_data.get_cbed_signal()
    >>> import dask.array as da
    >>> dask_array = da.from_array(s.data, chunks=(5, 5, 25, 25))
    >>> import pyxem.utils.dask_tools as dt
    >>> peak_array = _peak_find_dog(dask_array)
    >>> peak_array_computed = peak_array.compute()

    """
    dask_array_rechunked = _rechunk_signal2d_dim_one_chunk(dask_array)
    drop_axis = (dask_array_rechunked.ndim - 2, dask_array_rechunked.ndim - 1)
    output_array = da.map_blocks(
        _peak_find_dog_chunk,
        dask_array_rechunked,
        drop_axis=drop_axis,
        dtype=np.object,
        **kwargs
    )
    return output_array 

def _template_match_with_binary_image(dask_array, binary_image):
    """Template match a dask array with a binary image (template).

    Parameters
    ----------
    dask_array : Dask array
        The two last dimensions are the signal dimensions. Must have at least
        2 dimensions.
    binary_image : 2D NumPy array
        Must be smaller than the two last dimensions in dask_array

    Returns
    -------
    template_match : Dask array
        Same size as input data

    Examples
    --------
    >>> data = np.random.randint(1000, size=(20, 20, 256, 256))
    >>> import dask.array as da
    >>> dask_array = da.from_array(data, chunks=(5, 5, 128, 128))
    >>> from skimage import morphology
    >>> binary_image = morphology.disk(4, np.uint16)
    >>> import pyxem.utils.dask_tools as dt
    >>> template_match = dt._template_match_with_binary_image(
    ...     dask_array, binary_image)

    """
    if len(binary_image.shape) != 2:
        raise ValueError(
            "binary_image must have two dimensions, not {0}".format(
                len(binary_image.shape)
            )
        )
    dask_array_rechunked = _rechunk_signal2d_dim_one_chunk(dask_array)
    output_array = da.map_blocks(
        _template_match_binary_image_chunk,
        dask_array_rechunked,
        binary_image,
        dtype=np.float32,
    )
    return output_array 

def HaloRadius(mass, cosmo, redshift, mdef='vir'):
    r"""
    Return proper halo radius from halo mass, based on the specified mass definition.
    This is independent of halo profile, and simply returns

    .. math::

        R = \left [ 3 M /(4\pi\Delta) \right]^{1/3}

    where :math:`\Delta` is the density threshold, which depends on cosmology,
    redshift, and mass definition

    .. note::
        The units of the input mass are assumed to be :math:`M_{\odot}/h`

    Parameters
    ----------
    mass : array_like
        either a numpy or dask array specifying the halo mass; units
        assumed to be :math:`M_{\odot}/h`
    cosmo : :class:`~nbodykit.cosmology.cosmology.Cosmology`
        the cosmology instance
    redshift : float
        compute the density threshold which determines the R(M) relation
        at this redshift
    mdef : str, optional
        string specifying the halo mass definition to use; should be
        'vir' or 'XXXc' or 'XXXm' where 'XXX' is an int specifying the
        overdensity

    Returns
    -------
    radius : :class:`dask.array.Array`
        a dask array holding the halo radius in 'physical Mpc/h [sic]'.
        This is proper Mpc/h, to convert to comoving, divide this by scaling factor.

    """
    from halotools.empirical_models import halo_mass_to_halo_radius

    mass, redshift = da.broadcast_arrays(mass, redshift)

    kws = {'cosmology':cosmo.to_astropy(), 'mdef':mdef}

    def mass_to_radius(mass, redshift):
        return halo_mass_to_halo_radius(mass=mass, redshift=redshift, **kws)

    return da.map_blocks(mass_to_radius, mass, redshift, dtype=mass.dtype)

# deprecated functions 

def HaloConcentration(mass, cosmo, redshift, mdef='vir'):
    """
    Return halo concentration from halo mass, based on the analytic fitting
    formulas presented in
    `Dutton and Maccio 2014 <https://arxiv.org/abs/1402.7073>`_.

    .. note::
        The units of the input mass are assumed to be :math:`M_{\odot}/h`

    Parameters
    ----------
    mass : array_like
        either a numpy or dask array specifying the halo mass; units
        assumed to be :math:`M_{\odot}/h`
    cosmo : :class:`~nbodykit.cosmology.cosmology.Cosmology`
        the cosmology instance used in the analytic formula
    redshift : float
        compute the c(M) relation at this redshift
    mdef : str, optional
        string specifying the halo mass definition to use; should be
        'vir' or 'XXXc' or 'XXXm' where 'XXX' is an int specifying the
        overdensity

    Returns
    -------
    concen : :class:`dask.array.Array`
        a dask array holding the analytic concentration values

    References
    ----------
    Dutton and Maccio, "Cold dark matter haloes in the Planck era: evolution
    of structural parameters for Einasto and NFW profiles", 2014, arxiv:1402.7073

    """
    from halotools.empirical_models import NFWProfile

    mass, redshift = da.broadcast_arrays(mass, redshift)

    kws = {'cosmology':cosmo.to_astropy(), 'conc_mass_model':'dutton_maccio14', 'mdef':mdef}

    def get_nfw_conc(mass, redshift):
        kw1 = {}
        kw1.update(kws)
        kw1['redshift'] = redshift
        model = NFWProfile(**kw1)
        return model.conc_NFWmodel(prim_haloprop=mass)

    return da.map_blocks(get_nfw_conc, mass, redshift, dtype=mass.dtype) 

def SkyToCartesian(ra, dec, redshift, cosmo, observer=[0, 0, 0], degrees=True, frame='icrs'):
    """
    Convert sky coordinates (``ra``, ``dec``, ``redshift``) to a
    Cartesian ``Position`` column.

    .. warning::

        The returned Cartesian position is in units of Mpc/h.

    Parameters
    -----------
    ra : :class:`dask.array.Array`; shape: (N,)
        the right ascension angular coordinate
    dec : :class:`dask.array.Array`; shape: (N,)
        the declination angular coordinate
    redshift : :class:`dask.array.Array`; shape: (N,)
        the redshift coordinate
    cosmo : :class:`~nbodykit.cosmology.cosmology.Cosmology`
        the cosmology used to meausre the comoving distance from ``redshift``
    degrees : bool, optional
        specifies whether ``ra`` and ``dec`` are in degrees
    frame : string ('icrs' or 'galactic')
        speciefies which frame the Cartesian coordinates is. 

    Returns
    -------
    pos : :class:`dask.array.Array`; shape: (N,3)
        the cartesian position coordinates, where columns represent
        ``x``, ``y``, and ``z`` in units of Mpc/h

    Raises
    ------
    TypeError
        If the input columns are not dask arrays
    """
    ra, dec, redshift = da.broadcast_arrays(ra, dec, redshift)

    # pos on the unit sphere
    pos = SkyToUnitSphere(ra, dec, degrees=degrees, frame=frame)

    # multiply by the comoving distance in Mpc/h
    r = redshift.map_blocks(lambda z: cosmo.comoving_distance(z), dtype=redshift.dtype)

    return r[:,None] * pos + observer 

def _fit_xrf_block(data, data_sel_indices,
                   matv, snip_param, use_snip):
    """
    Spectrum fitting for a block of XRF dataset. The function is intended to be
    called using `map_blocks` function for parallel processing using Dask distributed
    package.

    Parameters
    ----------
    data : ndarray
        block of an XRF dataset. Shape=(ny, nx, ne).
    data_sel_indices: tuple
        tuple `(n_start, n_end)` which defines the indices along axis 2 of `data` array
        that are used for fitting. Note that `ne` (in `data`) and `ne_model` (in `matv`)
        are not equal. But `n_end - n_start` MUST be equal to `ne_model`! Indexes
        `n_start .. n_end - 1` will be selected from each pixel.
    matv: ndarray
        Matrix of spectra of the selected elements (emission lines). Shape=(ne_model, n_lines)
    snip_param: dict
        Dictionary of parameters forwarded to 'snip' method for background removal.
        Keys: `e_offset`, `e_linear`, `e_quadratic` (parameters of the energy axis approximation),
        `b_width` (width of the window that defines resolution of the snip algorithm).
    use_snip: bool, optional
        enable/disable background removal using snip algorithm

    Returns
    -------
    data_out: ndarray
        array with fitting results. Shape: `(ny, nx, ne_model + 4)`. For each pixel
        the output data contains: `ne_model` values that represent area under the emission
        line spectra; background area (only in the selected energy range), error (R-factor),
        total count in the selected energy range, total count of the full experimental spectrum.
    """
    spec = data
    spec_sel = spec[:, :, data_sel_indices[0]: data_sel_indices[1]]

    if use_snip:
        bg_sel = np.apply_along_axis(snip_method_numba, 2, spec_sel,
                                     snip_param['e_offset'],
                                     snip_param['e_linear'],
                                     snip_param['e_quadratic'],
                                     width=snip_param['b_width'])

        y = spec_sel - bg_sel
        bg_sum = np.sum(bg_sel, axis=2)

    else:
        y = spec_sel
        bg_sum = np.zeros(shape=data.shape[0:2])

    weights, rfactor, _ = fit_spectrum(y, matv, axis=2, method="nnls")

    total_cnt = np.sum(spec, axis=2)
    sel_cnt = np.sum(spec_sel, axis=2)

    # Stack depth-wise (along axis 2)
    data_out = np.dstack((weights, bg_sum, rfactor, sel_cnt, total_cnt))

    return data_out