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(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_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_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_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')
