def test_blockwise_stacked_new_axes_front(concatenate): def f(x): if isinstance(x, list): x = np.concatenate(x) return x[None, ...] * np.ones(7)[(slice(None), ) + (None, ) * x.ndim] x = da.ones(5, chunks=2) y = da.blockwise(f, "qa", x, "a", new_axes={"q": 7}, concatenate=concatenate, dtype=x.dtype) z = da.blockwise(f, "qab", y, "ab", new_axes={"q": 7}, concatenate=concatenate, dtype=x.dtype) assert z.chunks == ((7, ), (7, ), (2, 2, 1)) assert_eq(z, np.ones((7, 7, 5))) w = da.blockwise(lambda x: x[:, 0, 0], "a", z, "abc", dtype=x.dtype, concatenate=True) assert w.chunks == ((7, ), ) assert_eq(w, np.ones((7, )))
def window_stats(flag_window, ubls, chan_freqs, antenna_names, scan_no, field_name, ddid, nchanbins=10, prev_stats=None): """ Calculates stats for a chunk of a `flag_window`. Should be combined with the stats from other chunks. Parameters ---------- flag_window : :class:`dask.Array` Flag window of shape :code:`(bl, corr, time, chan)` ubls : :class:`dask.Array` Unique baselines of shape :code:`(bl, 3)` chan_freqs : :class:`dask.Array` Channel frequencies of shape :code:`(chan,)` antenna_names : list or :class:`numpy.ndarray` Antenna names of shape :code:`(ant,) scan_no : int Scan number field_name : str Field name ddid : int Data descriptor id nchanbins : int, optional Number of bins in a channel prev_stats : :class:`WindowStatistics`, optional Previous stats Returns ------- stats : :class:`dask.Array` Dask array containing a single :class:`WindowStatistics` object. `prev_stats` is merged into this result, if present. """ # Construct as array of per-baseline stats objects stats = da.blockwise(_window_stats, ("bl",), flag_window, _WINDOW_SCHEMA, ubls, ("bl", "bl-comp"), chan_freqs, ("chan",), antenna_names, None, scan_no, None, field_name, None, ddid, None, nchanbins, None, dtype=np.object) # Create an empty stats object if the user hasn't supplied one if prev_stats is None: def _window_stat_creator(): return WindowStatistics(nchanbins) prev_stats = da.blockwise(_window_stat_creator, (), dtype=np.object) # Combine per-baseline stats into a single stats object return da.blockwise(_combine_baseline_window_stats, (), stats, ("bl",), prev_stats, (), dtype=np.object)
def test_blockwise_new_axes(): def f(x): return x[:, None] * np.ones((1, 7)) x = da.ones(5, chunks=2) y = da.blockwise( f, "aq", x, "a", new_axes={"q": 7}, concatenate=True, dtype=x.dtype ) assert y.chunks == ((2, 2, 1), (7,)) assert_eq(y, np.ones((5, 7))) def f(x): return x[None, :] * np.ones((7, 1)) x = da.ones(5, chunks=2) y = da.blockwise( f, "qa", x, "a", new_axes={"q": 7}, concatenate=True, dtype=x.dtype ) assert y.chunks == ((7,), (2, 2, 1)) assert_eq(y, np.ones((7, 5))) def f(x): y = x.sum(axis=1) return y[:, None] * np.ones((1, 5)) x = da.ones((4, 6), chunks=(2, 2)) y = da.blockwise( f, "aq", x, "ab", new_axes={"q": 5}, concatenate=True, dtype=x.dtype ) assert y.chunks == ((2, 2), (5,)) assert_eq(y, np.ones((4, 5)) * 6)
def dde_factory(args, ms, ant, field, pol, lm, utime, frequency): if args.beam is None: return None # Beam is requested corr_type = tuple(pol.CORR_TYPE.data[0]) if not len(corr_type) == 4: raise ValueError("Need four correlations for DDEs") parangles = parallactic_angles(utime, ant.POSITION.data, field.PHASE_DIR.data[0][0]) corr_type_set = set(corr_type) if corr_type_set.issubset(set([9, 10, 11, 12])): pol_type = 'linear' elif corr_type_set.issubset(set([5, 6, 7, 8])): pol_type = 'circular' else: raise ValueError("Cannot determine polarisation type " "from correlations %s. Constructing " "a feed rotation matrix will not be " "possible." % (corr_type, )) # Construct feed rotation feed_rot = feed_rotation(parangles, pol_type) dtype = np.result_type(parangles, frequency) # Create zeroed pointing errors zpe = da.blockwise(_zero_pes, ("time", "ant", "chan", "comp"), parangles, ("time", "ant"), frequency, ("chan", ), dtype, None, new_axes={"comp": 2}, dtype=dtype) # Created zeroed antenna scaling factors zas = da.blockwise(_unity_ant_scales, ("ant", "chan", "comp"), parangles, ("time", "ant"), frequency, ("chan", ), dtype, None, new_axes={"comp": 2}, dtype=dtype) # Load the beam information beam, lm_ext, freq_map = load_beams(args.beam, corr_type, args.l_axis, args.m_axis) # Introduce the correlation axis beam = beam.reshape(beam.shape[:3] + (2, 2)) beam_dde = beam_cube_dde(beam, lm_ext, freq_map, lm, parangles, zpe, zas, frequency) # Multiply the beam by the feed rotation to form the DDE term return da.einsum("stafij,tajk->stafik", beam_dde, feed_rot)
def wsclean_predict(uvw, lm, source_type, flux, coeffs, log_poly, ref_freq, gauss_shape, frequency): spectrum_dtype = np.result_type(*(a.dtype for a in (flux, coeffs, log_poly, ref_freq, frequency))) spectrum = da.blockwise(wsclean_spectrum_wrapper, ("source", "chan"), flux, ("source",), coeffs, ("source", "comp"), log_poly, ("source",), ref_freq, ("source",), frequency, ("chan",), dtype=spectrum_dtype) out_dtype = np.result_type(uvw.dtype, lm.dtype, frequency.dtype, spectrum.dtype, np.complex64) vis = da.blockwise(wsclean_body_wrapper, ("source", "row", "chan", "corr"), uvw, ("row", "uvw"), lm, ("source", "lm"), source_type, ("source",), gauss_shape, ("source", "gauss"), frequency, ("chan",), spectrum, ("source", "chan"), out_dtype, None, adjust_chunks={"source": 1}, new_axes={"corr": 1}, dtype=out_dtype) return vis.sum(axis=0)
def _dask_merge_flags(flag_row, flag): """ Perform flag merging on dask arrays """ if flag_row is None and flag is not None: return da.blockwise(merge_flags, "r", flag_row, None, flag, "rfc", concatenate=True, dtype=flag.dtype) elif flag_row is not None and flag is None: return da.blockwise(merge_flags, "r", flag_row, "r", None, None, dtype=flag_row.dtype) elif flag_row is not None and flag is not None: return da.blockwise(merge_flags, "r", flag_row, "r", flag, "rfc", concatenate=True, dtype=flag_row.dtype) else: return None
def test_args_delayed(): x = da.arange(10, chunks=(5,)) y = dask.delayed(lambda: 100)() z = da.blockwise(add, 'i', x, 'i', y, None, dtype=x.dtype) assert_eq(z, np.arange(10) + 100) z = da.blockwise(lambda x, y: x + y, 'i', x, 'i', y=y, dtype=x.dtype) assert_eq(z, np.arange(10) + 100)
def test_args_delayed(): x = da.arange(10, chunks=(5, )) y = dask.delayed(lambda: 100)() z = da.blockwise(add, "i", x, "i", y, None, dtype=x.dtype) assert_eq(z, np.arange(10) + 100) z = da.blockwise(lambda x, y: x + y, "i", x, "i", y=y, dtype=x.dtype) assert_eq(z, np.arange(10) + 100)
def test_blockwise_stacked_new_axes(concatenate): def f(x): return x[..., None] * np.ones((1, 7)) x = da.ones(5, chunks=2) y = da.blockwise(f, 'aq', x, 'a', new_axes={'q': 7}, concatenate=concatenate, dtype=x.dtype) z = da.blockwise(f, 'abq', y, 'ab', new_axes={'q': 7}, concatenate=concatenate, dtype=x.dtype) assert z.chunks == ((2, 2, 1), (7,), (7,)) assert_eq(z, np.ones((5, 7, 7)))
def test_blockwise_stacked_new_axes_same_dim(concatenate): def f(x): return x[..., None] * np.ones((1, 7)) x = da.ones(5, chunks=2) y = da.zeros(5, chunks=2) a = da.blockwise(f, 'aq', x, 'a', new_axes={'q': 7}, concatenate=concatenate, dtype=x.dtype) b = da.blockwise(f, 'aq', y, 'a', new_axes={'q': 7}, concatenate=concatenate, dtype=x.dtype) c = a + b assert c.chunks == ((2, 2, 1), (7,)) assert_eq(c, np.ones((5, 7)))
def test_dont_merge_before_reductions(): x = da.ones(10, chunks=(5,)) y = da.blockwise(inc, "i", x, "i", dtype=x.dtype) z = da.blockwise(sum, "", y, "i", dtype=y.dtype) w = da.blockwise(sum, "", z, "", dtype=y.dtype) dsk = optimize_blockwise(w.dask) assert len([d for d in dsk.dicts.values() if isinstance(d, Blockwise)]) == 2 z.compute()
def test_dont_merge_before_reductions(): x = da.ones(10, chunks=(5,)) y = da.blockwise(inc, 'i', x, 'i', dtype=x.dtype) z = da.blockwise(sum, '', y, 'i', dtype=y.dtype) w = da.blockwise(sum, '', z, '', dtype=y.dtype) dsk = optimize_blockwise(w.dask) assert len([d for d in dsk.dicts.values() if isinstance(d, Blockwise)]) == 2 z.compute()
def test_blockwise_chunks(): x = da.ones((5, 5), chunks=((2, 1, 2), (3, 2))) def double(a, axis=0): return np.concatenate([a, a], axis=axis) y = da.blockwise( double, "ij", x, "ij", adjust_chunks={"i": lambda n: 2 * n}, axis=0, dtype=x.dtype, ) assert y.chunks == ((4, 2, 4), (3, 2)) assert_eq(y, np.ones((10, 5))) y = da.blockwise( double, "ij", x, "ij", adjust_chunks={"j": lambda n: 2 * n}, axis=1, dtype=x.dtype, ) assert y.chunks == ((2, 1, 2), (6, 4)) assert_eq(y, np.ones((5, 10))) x = da.ones((10, 10), chunks=(5, 5)) y = da.blockwise(double, "ij", x, "ij", axis=0, adjust_chunks={"i": 10}, dtype=x.dtype) assert y.chunks == ((10, 10), (5, 5)) assert_eq(y, np.ones((20, 10))) y = da.blockwise(double, "ij", x, "ij", axis=0, adjust_chunks={"i": (10, 10)}, dtype=x.dtype) assert y.chunks == ((10, 10), (5, 5)) assert_eq(y, np.ones((20, 10)))
def test_validate_top_inputs(): A = da.random.random((20, 20), chunks=(10, 10)) with pytest.raises(ValueError) as info: da.blockwise(inc, "jk", A, "ij", dtype=A.dtype) assert "unknown dimension" in str(info.value).lower() assert "k" in str(info.value) assert "j" not in str(info.value) with pytest.raises(ValueError) as info: da.blockwise(inc, "ii", A, "ij", dtype=A.dtype) assert "repeated" in str(info.value).lower() assert "i" in str(info.value)
def test_blockwise_stacked_new_axes_front(concatenate): def f(x): if isinstance(x, list): x = np.concatenate(x) return x[None, ...] * np.ones(7)[(slice(None),) + (None,) * x.ndim] x = da.ones(5, chunks=2) y = da.blockwise(f, 'qa', x, 'a', new_axes={'q': 7}, concatenate=concatenate, dtype=x.dtype) z = da.blockwise(f, 'qab', y, 'ab', new_axes={'q': 7}, concatenate=concatenate, dtype=x.dtype) assert z.chunks == ((7,), (7,), (2, 2, 1)) assert_eq(z, np.ones((7, 7, 5))) w = da.blockwise(lambda x: x[:, 0, 0], 'a', z, 'abc', dtype=x.dtype, concatenate=True) assert w.chunks == ((7,),) assert_eq(w, np.ones((7,)))
def test_validate_top_inputs(): A = da.random.random((20, 20), chunks=(10, 10)) with pytest.raises(ValueError) as info: da.blockwise(inc, 'jk', A, 'ij', dtype=A.dtype) assert 'unknown dimension' in str(info.value).lower() assert 'k' in str(info.value) assert 'j' not in str(info.value) with pytest.raises(ValueError) as info: da.blockwise(inc, 'ii', A, 'ij', dtype=A.dtype) assert 'repeated' in str(info.value).lower() assert 'i' in str(info.value)
def test_blockwise_no_array_args(): def f(dtype): return np.ones((2, 3), dtype) x = da.blockwise(f, 'ab', np.float32, None, new_axes={'a': 2, 'b': (3, 3)}, dtype=np.float32) assert x.chunks == ((2,), (3, 3)) assert_eq(x, np.ones((2, 6), np.float32))
def model_from_comps(comps, freq, mask, ref_freq): return da.blockwise(_model_from_comps_wrapper, ('chan', 'nx', 'ny'), comps, ('com', 'pix'), freq, ('chan',), mask, ('nx', 'ny'), ref_freq, None, dtype=comps.dtype)
def bda_mapper(time, interval, antenna1, antenna2, uvw, chan_width, chan_freq, max_uvw_dist, flag_row=None, max_fov=None, decorrelation=None, time_bin_secs=None, min_nchan=None): """ Createask row mapping structure for each row chunk """ return da.blockwise(_bda_mapper_wrapper, ("row", ), time, ("row", ), interval, ("row", ), antenna1, ("row", ), antenna2, ("row", ), uvw, ("row", "uvw"), chan_width, ("chan", ), chan_freq, ("chan", ), max_uvw_dist, None if max_uvw_dist is None else (), flag_row, None if flag_row is None else ("row", ), max_fov=max_fov, decorrelation=decorrelation, time_bin_secs=time_bin_secs, min_nchan=min_nchan, adjust_chunks={"row": lambda x: np.nan}, meta=np.empty((0, 0), dtype=np.object))
def restore_corrs(vis, ncorr): return da.blockwise(_restore_corrs, ('row', 'chan', 'corr'), vis, ('row', 'chan'), ncorr, None, new_axes={"corr": ncorr}, dtype=vis.dtype)
def test_atop_legacy(): x = da.ones(10, chunks=(5,)) with pytest.warns(None): y = da.atop(inc, 'i', x, 'i', dtype=x.dtype) z = da.blockwise(inc, 'i', x, 'i', dtype=x.dtype) assert_eq(y, z) assert y.name == z.name
def compute_counts(uvw, freqs, fbin_idx, fbin_counts, flag, nx, ny, cell_size_x, cell_size_y, dtype): counts = da.blockwise(compute_counts_wrapper, ('row', 'chan', 'nx', 'ny'), uvw, ('row', 'three'), freqs, ('chan', ), fbin_idx, ('chan', ), fbin_counts, ('chan', ), flag, ('row', 'chan'), nx, None, ny, None, cell_size_x, None, cell_size_y, None, dtype, None, new_axes={ "nx": nx, "ny": ny }, adjust_chunks={ 'chan': fbin_idx.chunks[0], 'row': (1, ) * len(uvw.chunks[0]) }, align_arrays=False, dtype=dtype) return counts.sum(axis=0)
def test_blockwise_kwargs(): def f(a, b=0): return a + b x = da.ones(5, chunks=(2,)) y = da.blockwise(f, 'i', x, 'i', b=10, dtype=x.dtype) assert_eq(y, np.ones(5) + 10)
def counts_to_weights(counts, uvw, freqs, fbin_idx, fbin_counts, nx, ny, cell_size_x, cell_size_y, dtype, robust): weights = da.blockwise(counts_to_weights_wrapper, ('row', 'chan'), counts, ('chan', 'nx', 'ny'), uvw, ('row', 'three'), freqs, ('chan', ), fbin_idx, ('chan', ), fbin_counts, ('chan', ), nx, None, ny, None, cell_size_x, None, cell_size_y, None, dtype, None, robust, None, adjust_chunks={'chan': freqs.chunks[0]}, align_arrays=False, dtype=dtype) return weights
def test_atop_legacy(): x = da.ones(10, chunks=(5, )) with pytest.warns(None): y = da.atop(inc, "i", x, "i", dtype=x.dtype) z = da.blockwise(inc, "i", x, "i", dtype=x.dtype) assert_eq(y, z) assert y.name == z.name
def _query_resample_kdtree(self, resample_kdtree, tlons, tlats, valid_input_index, valid_output_index, mask, neighbors, radius_of_influence, epsilon): """Query kd-tree on slice of target coordinates.""" if mask is None: args = tuple() else: ndims = self.source_geo_def.ndim dims = 'mn'[:ndims] args = (mask, dims, valid_input_index, dims) # res.shape = rows, cols, neighbors # j=rows, i=cols, k=neighbors, m=source rows, n=source cols res = da.blockwise(query_no_distance, 'jik', tlons, 'ji', tlats, 'ji', valid_output_index, 'ji', *args, kdtree=resample_kdtree, neighbours=neighbors, epsilon=epsilon, radius=radius_of_influence, dtype=np.int64, new_axes={'k': neighbors}, concatenate=True) return res
async def test_combo_of_layer_types(c, s, a, b): """Check pack/unpack of a HLG that has everything!""" def add(x, y, z, extra_arg): return x + y + z + extra_arg y = c.submit(lambda x: x, 2) z = c.submit(lambda x: x, 3) x = da.blockwise( add, "x", da.zeros((3, ), chunks=(1, )), "x", da.ones((3, ), chunks=(1, )), "x", y, None, concatenate=False, dtype=int, extra_arg=z, ) df = dd.from_pandas(pd.DataFrame({"a": np.arange(3)}), npartitions=3) df = df.shuffle("a", shuffle="tasks") df = df["a"].to_dask_array() res = x.sum() + df.sum() res = await c.compute(res, optimize_graph=False) assert res == 21
def test_blockwise_kwargs(): def f(a, b=0): return a + b x = da.ones(5, chunks=(2, )) y = da.blockwise(f, "i", x, "i", b=10, dtype=x.dtype) assert_eq(y, np.ones(5) + 10)
def _dask_row_chan_average(row_meta, chan_meta, flag_row=None, weight=None, vis=None, flag=None, weight_spectrum=None, sigma_spectrum=None, chan_bin_size=1): """ Average (row,chan,corr)-based dask arrays """ if chan_meta is None: return RowChanAverageOutput(None, None, None, None) # We don't know how many rows are in each row chunk, # but we can simply divide each channel chunk size by the bin size adjust_chunks = { "row": lambda r: np.nan, "chan": lambda c: (c + chan_bin_size - 1) // chan_bin_size } flag_row_dims = None if flag_row is None else ("row", ) weight_dims = None if weight is None else ("row", ) vis_dims = None if vis is None else _row_chan_avg_dims flag_dims = None if flag is None else _row_chan_avg_dims ws_dims = None if weight_spectrum is None else _row_chan_avg_dims ss_dims = None if sigma_spectrum is None else _row_chan_avg_dims # Lie just enough about the returned object (a namedtuple) # so that the blockwise passes kw = ({ "meta": np.empty((0, ) * len(_row_chan_avg_dims), dtype=np.object) } if PY3 else {}) avg = da.blockwise(row_chan_average, _row_chan_avg_dims, row_meta, ("row", ), chan_meta, ("chan", ), flag_row, flag_row_dims, weight, weight_dims, vis, vis_dims, flag, flag_dims, weight_spectrum, ws_dims, sigma_spectrum, ss_dims, align_arrays=False, adjust_chunks=adjust_chunks, dtype=np.object, **kw) tuple_gets = ( None if a is None else _getitem_row_chan(avg, i, a.dtype) for i, a in enumerate([vis, flag, weight_spectrum, sigma_spectrum])) return RowChanAverageOutput(*tuple_gets)
def dirty(grid, grid_config): """ Computes the dirty image from gridded visibilities and the gridding configuration. Parameters ---------- grid : :class:`dask.array.Array` Gridded visibilities of shape :code:`(nv, nu, ncorr)` grid_config : :class:`GridderConfigWrapper` Gridding configuration Returns ------- dirty : :class:`dask.array.Array` dirty image of shape :code:`(ny, nx, corr)` """ gc = grid_config.object nx = gc.Nxdirty() ny = gc.Nydirty() return da.blockwise(_nifty_dirty, ("nx", "ny", "corr"), grid, ("nx", "ny", "corr"), gc, None, adjust_chunks={ "nx": nx, "ny": ny }, dtype=grid.real.dtype)
def _dask_chan_average(chan_meta, chan_freq=None, chan_width=None, chan_bin_size=1): if chan_meta is None: return ChannelAverageOutput(None, None) adjust_chunks = { "chan": lambda c: (c + chan_bin_size - 1) // chan_bin_size } kw = {"meta": np.empty((0, ), dtype=np.object)} if PY3 else {} avg = da.blockwise(chan_average, ("chan", ), chan_meta, ("chan", ), chan_freq, None if chan_freq is None else ("chan", ), chan_width, None if chan_width is None else ("chan", ), adjust_chunks=adjust_chunks, dtype=np.object, **kw) tuple_gets = (None if a is None else _getitem_chan(avg, i, a.dtype) for i, a in enumerate([chan_freq, chan_width])) return ChannelAverageOutput(*tuple_gets)
def hdot(self, x): beta = np.pad(x[0].reshape(self.nband, self.nx * self.ny), ((0, 0), (0, self.nmax - self.nx * self.ny)), mode='constant') xdask = da.from_array(x[1], chunks=(self.nband, self.nx, self.ny)) alpha = da.blockwise(_hdot_internal_wrapper, ("basis", "band", "nmax"), xdask, ("band", "nx", "ny"), self.bases, ("basis", ), self.ntot, ("basis", ), self.nmax, None, self.nlevels, None, self.sqrtP, None, self.nx, None, self.ny, None, self.real_type, None, new_axes={"nmax": self.nmax}, dtype=self.real_type, align_arrays=False) alpha = alpha.compute(shedular='processes') return np.concatenate((beta[None], alpha), axis=0)
def dot(self, coeffs): x0 = coeffs[0, :, self.dpadding].reshape(self.nband, self.nx, self.ny) # Dirac components alpha_dask = da.from_array(coeffs[1::], chunks=(1, self.nband, self.nmax)) x1 = da.blockwise( _dot_internal_wrapper, ("basis", "band", "nx", "ny"), alpha_dask, ("basis", "band", "ntot"), self.bases, ("basis", ), self.padding, ("basis", ), self.iy, None, # ("basis",), self.sy, None, # ("basis",), self.sqrtP, None, self.nx, None, self.ny, None, self.real_type, None, new_axes={ "nx": self.nx, "ny": self.ny }, dtype=self.real_type, align_arrays=False) x1 = x1.sum(axis=0).compute(shedular='processes') return np.concatenate((x0[None], x1[None]), axis=0)
def model(image, grid_config): """ Computes model visibilities from an image and a gridding configuration. Parameters ---------- image : :class:`dask.array.Array` Image of shape :code:`(ny, nx, corr)`. grid_config : :class:`GridderConfigWrapper` nifty gridding configuration object Returns ------- model_vis : :class:`dask.array.Array` Model visibilities of shape :code:`(nu, nv, corr)`. """ gc = grid_config.object nu = gc.Nu() nv = gc.Nv() return da.blockwise(_nifty_model, ("nu", "nv", "corr"), image, ("nu", "nv", "corr"), gc, None, adjust_chunks={ "nu": nu, "nv": nv }, dtype=image.dtype)
def test_blockwise_no_args(): def f(): return np.ones((2, 3), np.float32) x = da.blockwise(f, "ab", new_axes={"a": 2, "b": (3, 3)}, dtype=np.float32) assert x.chunks == ((2, ), (3, 3)) assert_eq(x, np.ones((2, 6), np.float32))
def residual(uvw, freq, image, vis, freq_bin_idx, freq_bin_counts, cell, weights=None, flag=None, celly=None, epsilon=1e-5, nthreads=1, do_wstacking=True): if celly is None: celly = cell if not nthreads: import multiprocessing nthreads = multiprocessing.cpu_count() if weights is None: weight_out = None else: weight_out = ('row', 'chan') if flag is None: flag_out = None else: flag_out = ('row', 'chan') img = da.blockwise(_residual_wrapper, ('row', 'chan', 'nx', 'ny'), uvw, ('row', 'three'), freq, ('chan', ), image, ('chan', 'nx', 'ny'), vis, ('row', 'chan'), freq_bin_idx, ('chan', ), freq_bin_counts, ('chan', ), cell, None, weights, weight_out, flag, flag_out, celly, None, epsilon, None, nthreads, None, do_wstacking, None, adjust_chunks={ 'chan': freq_bin_idx.chunks[0], 'row': (1, ) * len(vis.chunks[0]) }, dtype=image.dtype, align_arrays=False) return img.sum(axis=0)
def test_blockwise_chunks(): x = da.ones((5, 5), chunks=((2, 1, 2), (3, 2))) def double(a, axis=0): return np.concatenate([a, a], axis=axis) y = da.blockwise( double, 'ij', x, 'ij', adjust_chunks={'i': lambda n: 2 * n}, axis=0, dtype=x.dtype ) assert y.chunks == ((4, 2, 4), (3, 2)) assert_eq(y, np.ones((10, 5))) y = da.blockwise( double, 'ij', x, 'ij', adjust_chunks={'j': lambda n: 2 * n}, axis=1, dtype=x.dtype ) assert y.chunks == ((2, 1, 2), (6, 4)) assert_eq(y, np.ones((5, 10))) x = da.ones((10, 10), chunks=(5, 5)) y = da.blockwise( double, 'ij', x, 'ij', axis=0, adjust_chunks={'i': 10}, dtype=x.dtype ) assert y.chunks == ((10, 10), (5, 5)) assert_eq(y, np.ones((20, 10))) y = da.blockwise( double, 'ij', x, 'ij', axis=0, adjust_chunks={'i': (10, 10)}, dtype=x.dtype ) assert y.chunks == ((10, 10), (5, 5)) assert_eq(y, np.ones((20, 10)))
def test_blockwise_new_axes_chunked(): def f(x): return x[None, :] * 2 x = da.arange(0, 6, 1, chunks=2, dtype=np.int32) y = da.blockwise(f, 'qa', x, 'a', new_axes={'q': (1, 1)}, dtype=x.dtype) assert y.chunks == ((1, 1), (2, 2, 2)) assert_eq(y, np.array([[0, 2, 4, 6, 8, 10], [0, 2, 4, 6, 8, 10]], np.int32))
def test_common_token_names_kwargs(name): x = np.array(['a', 'bb', 'ccc'], dtype=object) d = da.from_array(x, chunks=2) result = da.blockwise(lambda x, y: x + y, 'i', d, 'i', y=name, dtype=object) expected = x + name assert_eq(result, expected)
def test_namedtuple(tup): A = da.random.random((20, 20), chunks=(10, 10)) def f(data, x): return data B = da.blockwise( f, ("d1", "d2"), A, ("d1", "d2"), x=tup, dtype=A.dtype ) assert_eq(A, B)
def test_blockwise_new_axes(): def f(x): return x[:, None] * np.ones((1, 7)) x = da.ones(5, chunks=2) y = da.blockwise(f, 'aq', x, 'a', new_axes={'q': 7}, concatenate=True, dtype=x.dtype) assert y.chunks == ((2, 2, 1), (7,)) assert_eq(y, np.ones((5, 7))) def f(x): return x[None, :] * np.ones((7, 1)) x = da.ones(5, chunks=2) y = da.blockwise(f, 'qa', x, 'a', new_axes={'q': 7}, concatenate=True, dtype=x.dtype) assert y.chunks == ((7,), (2, 2, 1)) assert_eq(y, np.ones((7, 5))) def f(x): y = x.sum(axis=1) return y[:, None] * np.ones((1, 5)) x = da.ones((4, 6), chunks=(2, 2)) y = da.blockwise(f, 'aq', x, 'ab', new_axes={'q': 5}, concatenate=True, dtype=x.dtype) assert y.chunks == ((2, 2), (5,)) assert_eq(y, np.ones((4, 5)) * 6)
def test_blockwise_new_axes_2(): x = da.ones((2, 2), chunks=(1, 1)) def func(x): return np.stack([x, -x], axis=-1) y = da.blockwise( func, ('x', 'y', 'sign'), x, ('x', 'y'), dtype=x.dtype, concatenate=True, new_axes={'sign': 2} ) assert_eq(y, y)
def test_blockwise_names(): x = da.ones(5, chunks=(2,)) y = da.blockwise(add, 'i', x, 'i', dtype=x.dtype) assert y.name.startswith('add')