Python dask.array.Array() Examples

def interpolate_xarray(xpoints, ypoints, values, shape, kind='cubic',
                       blocksize=CHUNK_SIZE):
    """Interpolate, generating a dask array."""
    vchunks = range(0, shape[0], blocksize)
    hchunks = range(0, shape[1], blocksize)

    token = tokenize(blocksize, xpoints, ypoints, values, kind, shape)
    name = 'interpolate-' + token

    from scipy.interpolate import interp2d
    interpolator = interp2d(xpoints, ypoints, values, kind=kind)

    dskx = {(name, i, j): (interpolate_slice,
                           slice(vcs, min(vcs + blocksize, shape[0])),
                           slice(hcs, min(hcs + blocksize, shape[1])),
                           interpolator)
            for i, vcs in enumerate(vchunks)
            for j, hcs in enumerate(hchunks)
            }

    res = da.Array(dskx, name, shape=list(shape),
                   chunks=(blocksize, blocksize),
                   dtype=values.dtype)
    return DataArray(res, dims=('y', 'x')) 

def _dask_or_eager_func(name, eager_module=np, list_of_args=False, n_array_args=1):
    """Create a function that dispatches to dask for dask array inputs."""
    if has_dask:

        def f(*args, **kwargs):
            dispatch_args = args[0] if list_of_args else args
            if any(isinstance(a, dsa.Array) for a in dispatch_args[:n_array_args]):
                module = dsa
            else:
                module = eager_module
            return getattr(module, name)(*args, **kwargs)

    else:

        def f(data, *args, **kwargs):
            return getattr(eager_module, name)(data, *args, **kwargs)

    return f 

def StackColumns(*cols):
    """
    Stack the input dask arrays vertically, column by column.

    This uses :func:`dask.array.vstack`.

    Parameters
    ----------
    *cols : :class:`dask.array.Array`
        the dask arrays to stack vertically together

    Returns
    -------
    :class:`dask.array.Array` :
        the dask array where columns correspond to the input arrays

    Raises
    ------
    TypeError
        If the input columns are not dask arrays
    """
    cols = da.broadcast_arrays(*cols)
    return da.vstack(cols).T 

def _dask_or_eager_func(name, eager_module=np, list_of_args=False,
                        n_array_args=1):
    """Create a function that dispatches to dask for dask array inputs."""
    if has_dask:
        def f(*args, **kwargs):
            dispatch_args = args[0] if list_of_args else args
            if any(isinstance(a, dsa.Array)
                   for a in dispatch_args[:n_array_args]):
                module = dsa
            else:
                module = eager_module
            return getattr(module, name)(*args, **kwargs)
    else:
        def f(data, *args, **kwargs):
            return getattr(eager_module, name)(data, *args, **kwargs)
    return f 

def __new__(cls, dm, **kwargs):
        if isinstance(dm, da.Array):
            dm = DaskMeta.from_darray(dm)
        elif isinstance(dm, dict):
            dm = DaskMeta(**dm)
        elif isinstance(dm, DaskMeta):
            pass
        elif dm.__class__.__name__ in ("Op", "GraphMeta", "TmsMeta", "TemplateMeta"):
            itr = [dm.dask, dm.name, dm.chunks, dm.dtype, dm.shape]
            dm = DaskMeta._make(itr)
        else:
            raise ValueError("{} must be initialized with a DaskMeta, a dask array, or a dict with DaskMeta fields".format(cls.__name__))
        self = da.Array.__new__(cls, dm.dask, dm.name, dm.chunks, dtype=dm.dtype, shape=dm.shape)
        if "__geo_transform__" in kwargs:
            self.__geo_transform__ = kwargs["__geo_transform__"]
        if "__geo_interface__" in kwargs:
            self.__geo_interface__ = kwargs["__geo_interface__"]
        return self 

def _build_image_layer(self, image, image_bounds, cmap='viridis'):
        if image is not None:
            if isinstance(image, da.Array):
                if len(image.shape) == 2 or \
                    (image.shape[0] == 1 and len(image.shape) == 3):
                    arr = image.compute()
                else:
                    arr = image.rgb()
                coords = box(*image.bounds)
            else:
                assert image_bounds is not None, "Must pass image_bounds with ndarray images"
                arr = image
                coords = box(*image_bounds)
            b64 = self._encode_image(arr, cmap)
            return ImageLayer(b64, self._polygon_coords(coords))
        else:
            return 'false'; 

def _load_GeoTransform(self):
        """Calculate latitude and longitude variable calculated from the
        gdal.Open.GetGeoTransform method"""
        def load_lon():
            return arange(ds.RasterXSize)*b[1]+b[0]

        def load_lat():
            return arange(ds.RasterYSize)*b[5]+b[3]
        ds = self.ds
        b = self.ds.GetGeoTransform()  # bbox, interval
        if with_dask:
            lat = Array(
                {('lat', 0): (load_lat,)}, 'lat', (self.ds.RasterYSize,),
                shape=(self.ds.RasterYSize,), dtype=float)
            lon = Array(
                {('lon', 0): (load_lon,)}, 'lon', (self.ds.RasterXSize,),
                shape=(self.ds.RasterXSize,), dtype=float)
        else:
            lat = load_lat()
            lon = load_lon()
        return Variable(('lat',), lat), Variable(('lon',), lon) 

def test_nearest_swath_1d_mask_to_grid_1n(self):
        """Test 1D swath definition to 2D grid definition; 1 neighbor."""
        from pyresample.kd_tree import XArrayResamplerNN
        import xarray as xr
        import dask.array as da
        resampler = XArrayResamplerNN(self.tswath_1d, self.tgrid,
                                      radius_of_influence=100000,
                                      neighbours=1)
        data = self.tdata_1d
        ninfo = resampler.get_neighbour_info(mask=data.isnull())
        for val in ninfo[:3]:
            # vii, ia, voi
            self.assertIsInstance(val, da.Array)
        res = resampler.get_sample_from_neighbour_info(data)
        self.assertIsInstance(res, xr.DataArray)
        self.assertIsInstance(res.data, da.Array)
        actual = res.values
        expected = np.array([
            [1., 2., 2.],
            [1., 2., 2.],
            [1., np.nan, 2.],
            [1., 2., 2.],
        ])
        np.testing.assert_allclose(actual, expected) 

def test_nearest_swath_2d_mask_to_area_1n(self):
        """Test 2D swath definition to 2D area definition; 1 neighbor."""
        from pyresample.kd_tree import XArrayResamplerNN
        import xarray as xr
        import dask.array as da
        swath_def = self.swath_def_2d
        data = self.data_2d
        resampler = XArrayResamplerNN(swath_def, self.area_def,
                                      radius_of_influence=50000,
                                      neighbours=1)
        ninfo = resampler.get_neighbour_info(mask=data.isnull())
        for val in ninfo[:3]:
            # vii, ia, voi
            self.assertIsInstance(val, da.Array)
        res = resampler.get_sample_from_neighbour_info(data)
        self.assertIsInstance(res, xr.DataArray)
        self.assertIsInstance(res.data, da.Array)
        res = res.values
        cross_sum = np.nansum(res)
        expected = 15874591.0
        self.assertEqual(cross_sum, expected) 

def test_nearest_swath_1d_mask_to_grid_8n(self):
        """Test 1D swath definition to 2D grid definition; 8 neighbors."""
        from pyresample.kd_tree import XArrayResamplerNN
        import xarray as xr
        import dask.array as da
        resampler = XArrayResamplerNN(self.tswath_1d, self.tgrid,
                                      radius_of_influence=100000,
                                      neighbours=8)
        data = self.tdata_1d
        ninfo = resampler.get_neighbour_info(mask=data.isnull())
        for val in ninfo[:3]:
            # vii, ia, voi
            self.assertIsInstance(val, da.Array)
        res = resampler.get_sample_from_neighbour_info(data)
        self.assertIsInstance(res, xr.DataArray)
        self.assertIsInstance(res.data, da.Array)
        # actual = res.values
        # expected = TODO
        # np.testing.assert_allclose(actual, expected) 

def _dask_or_eager_func(name, eager_module=np, list_of_args=False, n_array_args=1):
    """Create a function that dispatches to dask for dask array inputs."""
    if has_dask:

        def f(*args, **kwargs):
            dispatch_args = args[0] if list_of_args else args
            if any(isinstance(a, dsa.Array) for a in dispatch_args[:n_array_args]):
                module = dsa
            else:
                module = eager_module
            return getattr(module, name)(*args, **kwargs)

    else:

        def f(data, *args, **kwargs):
            return getattr(eager_module, name)(data, *args, **kwargs)

    return f 

def _dask_array_vgrid(self, varname, klevels, k_chunksize):
        # return a dask array for a 1D vertical grid var

        # single chunk for 1D variables
        chunks = ((len(klevels),),)

        # manually build dask graph
        dsk = {}
        token = tokenize(varname, self.store)
        name = '-'.join([varname, token])

        nz = self.nz if _VAR_METADATA[varname]['dims'] != ['k_p1'] else self.nz+1
        task = (_get_1d_chunk, self.store, varname,
                list(klevels), nz, self.dtype)

        key = name, 0
        dsk[key] = task

        return dsa.Array(dsk, name, chunks, self.dtype) 

def compute_scaling(df, region1, region2=None, dmin=int(1e1), dmax=int(1e7), n_bins=50):

    import dask.array as da

    if region2 is None:
        region2 = region1

    distbins = numutils.logbins(dmin, dmax, N=n_bins)
    areas = contact_areas(distbins, region1, region2)

    df = df[
        (df["pos1"] >= region1[0])
        & (df["pos1"] < region1[1])
        & (df["pos2"] >= region2[0])
        & (df["pos2"] < region2[1])
    ]
    dists = (df["pos2"] - df["pos1"]).values

    if isinstance(dists, da.Array):
        obs, _ = da.histogram(dists[(dists >= dmin) & (dists < dmax)], bins=distbins)
    else:
        obs, _ = np.histogram(dists[(dists >= dmin) & (dists < dmax)], bins=distbins)

    return distbins, obs, areas 

def test_write_zarr(self, adata, adata_dist):
        import dask.array as da
        import zarr

        log1p(adata_dist)
        temp_store = zarr.TempStore()
        chunks = adata_dist.X.chunks
        if isinstance(chunks[0], tuple):
            chunks = (chunks[0][0],) + chunks[1]
        # write metadata using regular anndata
        adata.write_zarr(temp_store, chunks)
        if isinstance(adata_dist.X, da.Array):
            adata_dist.X.to_zarr(temp_store.dir_path("X"), overwrite=True)
        else:
            adata_dist.X.to_zarr(temp_store.dir_path("X"), chunks)
        # read back as zarr directly and check it is the same as adata.X
        adata_log1p = ad.read_zarr(temp_store)
        log1p(adata)
        npt.assert_allclose(adata_log1p.X, adata.X) 

def wrap_func(grid, data, dim, wrap, func="diff", idx=0):
    """interpolates data over discontuity (e.g. longitude values)

        TODO
        ----
        Write tests that runs np and dask arrays through this
    """
    if isinstance(data.data, da.Array):
        data.load()
        redask = 1
    else:
        redask = 0

    if func == "diff":
        out = grid.diff(data, dim)
    elif func == "interp":
        out = grid.interp(data, dim)
        # when interpolating the discontinuty gets halved
        wrap = -wrap / 2
    else:
        raise RuntimeError("`func` argument not recognized")

    target_dim = [
        a for a in out.dims if a in xgcm.comodo.get_axis_coords(grid._ds, dim)
    ]
    if len(target_dim) == 1:
        target_dim = target_dim[0]
    else:
        raise RuntimeError("more then one target dim found")

    # TODO the idx should be determined by a combo of the c grid shift and func
    out[{target_dim: idx}] = out[{target_dim: idx}] + wrap

    if redask:
        out.data = da.from_array(out.data, out.data.shape)
    return out 

def predict(client, model, data, proba=False, dtype=np.float32, **kwargs):
    if isinstance(data, dd._Frame):
        return data.map_partitions(_predict_part, model=model, proba=proba, **kwargs).values
    elif isinstance(data, da.Array):
        if proba:
            kwargs['chunks'] = (data.chunks[0], (model.n_classes_,))
        else:
            kwargs['drop_axis'] = 1
        return data.map_blocks(_predict_part, model=model, proba=proba, dtype=dtype, **kwargs)
    else:
        raise TypeError(f'Data must be either Dask array or dataframe. Got {type(data)}.') 

def is_dask_array(data):
    da = None
    if 'dask.array' in sys.modules:
        import dask.array as da
    return (da is not None and isinstance(data, da.Array)) 

def test_select_lazy(self):
        import dask.array as da
        arr = da.from_array(np.arange(1, 12), 3)
        ds = Dataset({'x': range(11), 'y': arr}, 'x', 'y')
        self.assertIsInstance(ds.select(x=(0, 5)).data['y'], da.Array) 

def test_dataset_weighted_histogram_dask(self):
        import dask.array as da
        ds = Dataset((da.from_array(np.array(range(10), dtype='f'), chunks=3),
                      da.from_array([i/10. for i in range(10)], chunks=3)),
                     ['x', 'y'], datatype=['dask'])
        op_hist = histogram(ds, weight_dimension='y', num_bins=3)

        hist = Histogram(([0, 3, 6, 9], [0.022222, 0.088889, 0.222222]),
                         vdims='y')
        self.assertIsInstance(op_hist.data['y'], da.Array)
        self.assertEqual(op_hist, hist) 

def test_dataset_cumulative_histogram_dask(self):
        import dask.array as da
        ds = Dataset((da.from_array(np.array(range(10), dtype='f'), chunks=(3)),),
                     ['x'], datatype=['dask'])
        op_hist = histogram(ds, num_bins=3, cumulative=True)

        hist = Histogram(([0, 3, 6, 9], [0.3, 0.6, 1]),
                         vdims=('x_frequency', 'Frequency'))
        self.assertIsInstance(op_hist.data['x_frequency'], da.Array)
        self.assertEqual(op_hist, hist) 

def test_dataset_histogram_dask(self):
        import dask.array as da
        ds = Dataset((da.from_array(np.array(range(10), dtype='f'), chunks=(3)),),
                     ['x'], datatype=['dask'])
        op_hist = histogram(ds, num_bins=3)

        hist = Histogram(([0, 3, 6, 9], [0.1, 0.1, 0.133333]),
                         vdims=('x_frequency', 'Frequency'))
        self.assertIsInstance(op_hist.data['x_frequency'], da.Array)
        self.assertEqual(op_hist, hist) 

def test_rasterio_glob():
    import dask.array as da
    pytest.importorskip('rasterio')
    cat = intake.open_catalog(os.path.join(here, 'data', 'catalog.yaml'))
    s = cat.tiff_glob_source
    info = s.discover()
    assert info['shape'] == (1, 3, 718, 791)
    x = s.to_dask()
    assert isinstance(x.data, da.Array)
    x = s.read()
    assert x.data.shape == (1, 3, 718, 791) 

def test_rasterio():
    import dask.array as da
    pytest.importorskip('rasterio')
    cat = intake.open_catalog(os.path.join(here, 'data', 'catalog.yaml'))
    s = cat.tiff_source
    info = s.discover()
    assert info['shape'] == (3, 718, 791)
    x = s.to_dask()
    assert isinstance(x.data, da.Array)
    x = s.read()
    assert x.data.shape == (3, 718, 791) 

def test_grib_dask():
    pytest.importorskip('Nio')
    import dask.array as da
    cat = intake.open_catalog(os.path.join(here, 'data', 'catalog.yaml'))
    x = cat.grib.to_dask()
    assert len(x.fileno) == 2
    assert isinstance(x.APCP_P8_L1_GLL0_acc6h.data, da.Array)
    values = x.APCP_P8_L1_GLL0_acc6h.data.compute()
    x2 = cat.grib.read()
    assert (values == x2.APCP_P8_L1_GLL0_acc6h.values).all() 

def normalize(
    x: xr.DataArray,
    *,
    group: Union[str, Grouper] = "time",
    kind: str = ADDITIVE,
    norm: Optional[xr.DataArray] = None,
):
    """Normalize an array by removing its mean.

    Normalization if performed group-wise.

    Parameters
    ----------
    x : xr.DataArray
      Array to be normalized.
    group : Union[str, Grouper]
      Grouping information. See :py:class:`xclim.sdba.base.Grouper` for details.
    kind : {'+', '*'}
      How to apply the adjustment, either additively or multiplicatively.
    norm : xr.DataArray
      If the norm was already computed (for example with `group.apply("mean", x)`), skip the
        computation step. The array should have the same dimensions as `x` except for "time" that should
        be replaced by `group.prop`.

    Returns
    -------
    xr.DataArray or xr.Dataset
      Group-wise anomaly of x
    """

    def _normalize_group(grp, dim=["time"]):
        return apply_correction(grp, invert(grp.mean(dim=dim), kind), kind)

    if norm is None:
        return group.apply(_normalize_group, x)

    return apply_correction(
        x, broadcast(invert(norm, kind), x, group=group, interp="nearest"), kind,
    ) 

def rle_1d(
    arr: Union[int, float, bool, Sequence[Union[int, float, bool]]]
) -> Tuple[np.array, np.array, np.array]:
    """Return the length, starting position and value of consecutive identical values.

    Parameters
    ----------
    arr : Sequence[Union[int, float, bool]]
      Array of values to be parsed.

    Returns
    -------
    values : np.array
      The values taken by arr over each run
    run lengths : np.array
      The length of each run
    start position : np.array
      The starting index of each run

    Examples
    --------
    >>> a = [1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3]
    >>> rle_1d(a)
    (array([1, 2, 3]), array([2, 4, 6]), array([0, 2, 6]))
    """
    ia = np.asarray(arr)
    n = len(ia)

    if n == 0:
        e = "run length array empty"
        warn(e)
        # Returning None makes some other 1d func below fail.
        return np.array(np.nan), 0, np.array(np.nan)

    y = np.array(ia[1:] != ia[:-1])  # pairwise unequal (string safe)
    i = np.append(np.where(y), n - 1)  # must include last element position
    rl = np.diff(np.append(-1, i))  # run lengths
    pos = np.cumsum(np.append(0, rl))[:-1]  # positions
    return ia[i], rl, pos 

def _compute_ij_images_xarray_dask_block(dtype: np.dtype,
                                         block_id: int,
                                         block_shape: Tuple[int, int],
                                         block_slices: Tuple[Tuple[int, int], Tuple[int, int], Tuple[int, int]],
                                         src_x: xr.DataArray,
                                         src_y: xr.DataArray,
                                         src_ij_bboxes: np.ndarray,
                                         dst_x_min: float,
                                         dst_y_min: float,
                                         dst_xy_res: float,
                                         uv_delta: float) -> np.ndarray:
    """Compute dask.array.Array destination block with source pixel i,j coords from xarray.DataArray x,y sources """
    dst_src_ij_block = np.full(block_shape, np.nan, dtype=dtype)
    _, (dst_y_slice_start, _), (dst_x_slice_start, _) = block_slices
    src_ij_bbox = src_ij_bboxes[block_id]
    src_i_min, src_j_min, src_i_max, src_j_max = src_ij_bbox
    if src_i_min == -1:
        return dst_src_ij_block
    src_x_values = src_x[src_j_min:src_j_max + 1, src_i_min:src_i_max + 1].values
    src_y_values = src_y[src_j_min:src_j_max + 1, src_i_min:src_i_max + 1].values
    _compute_ij_images_numpy_sequential(src_x_values,
                                        src_y_values,
                                        src_i_min,
                                        src_j_min,
                                        dst_src_ij_block,
                                        dst_x_min + dst_x_slice_start * dst_xy_res,
                                        dst_y_min + dst_y_slice_start * dst_xy_res,
                                        dst_xy_res,
                                        uv_delta)
    return dst_src_ij_block 

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 get_variables(self):
        def load(band):
            band = ds.GetRasterBand(band)
            a = band.ReadAsArray()
            no_data = band.GetNoDataValue()
            if no_data is not None:
                try:
                    a[a == no_data] = a.dtype.type(nan)
                except ValueError:
                    pass
            return a
        ds = self.ds
        dims = ['lat', 'lon']
        chunks = ((ds.RasterYSize,), (ds.RasterXSize,))
        shape = (ds.RasterYSize, ds.RasterXSize)
        variables = OrderedDict()
        for iband in range(1, ds.RasterCount+1):
            band = ds.GetRasterBand(iband)
            dt = dtype(gdal_array.codes[band.DataType])
            if with_dask:
                dsk = {('x', 0, 0): (load, iband)}
                arr = Array(dsk, 'x', chunks, shape=shape, dtype=dt)
            else:
                arr = load(iband)
            attrs = band.GetMetadata_Dict()
            try:
                dt.type(nan)
                attrs['_FillValue'] = nan
            except ValueError:
                no_data = band.GetNoDataValue()
                attrs.update({'_FillValue': no_data} if no_data else {})
            variables['Band%i' % iband] = Variable(dims, arr, attrs)
        variables['lat'], variables['lon'] = self._load_GeoTransform()
        return FrozenOrderedDict(variables) 

def is_dask(array):
    da = dask_array_module()
    if da is None:
        return False
    return da and isinstance(array, da.Array)