def test_structure():
s = ShareDict()
s.update(a)
s.update(b)
s.update(c)
assert all(any(d is x for d in s.dicts.values()) for x in [a, b, c])
def test_gh_4176():
with warnings.catch_warnings():
warnings.simplefilter('ignore')
from dask.sharedict import ShareDict
def foo(A):
return A[None, ...]
A = da.ones(shape=(10, 20, 4), chunks=(2, 5, 4))
name = 'D'
dsk = blockwise(foo,
name, ("nsrc", "ntime", "nbl", "npol"),
A.name, ("ntime", "nbl", "npol"),
new_axes={"nsrc": 1},
numblocks={a.name: a.numblocks
for a in (A, )})
array_dsk = ShareDict()
array_dsk.update(dsk)
array_dsk.update(A.__dask_graph__())
chunks = ((1, ), ) + A.chunks
D = da.Array(array_dsk, name, chunks, dtype=A.dtype)
D.sum(axis=0).compute()
def test_gh_4176():
from dask.sharedict import ShareDict
def foo(A):
return A[None, ...]
A = da.ones(shape=(10, 20, 4), chunks=(2, 5, 4))
name = 'D'
dsk = blockwise(
foo, name, ("nsrc", "ntime", "nbl", "npol"),
A.name, ("ntime", "nbl", "npol"),
new_axes={"nsrc": 1},
numblocks={a.name: a.numblocks for a in (A,)}
)
array_dsk = ShareDict()
array_dsk.update(dsk)
array_dsk.update(A.__dask_graph__())
chunks = ((1,),) + A.chunks
D = da.Array(array_dsk, name, chunks, dtype=A.dtype)
D.sum(axis=0).compute()
def test_structure():
s = ShareDict()
s.update(a)
s.update(b)
s.update(c)
assert all(any(d is x for d in s.dicts.values())
for x in [a, b, c])
def test_core():
s = ShareDict()
assert isinstance(s, Mapping)
s.update(a)
s.update(b)
assert s['x'] == 1
with pytest.raises(KeyError):
s['abc']
with pytest.raises((NotImplementedError, TypeError)):
s['abc'] = 123
def test_core():
s = ShareDict()
assert isinstance(s, Mapping)
s.update(a)
s.update(b)
assert s['x'] == 1
with pytest.raises(KeyError):
s['abc']
with pytest.raises((NotImplementedError, TypeError)):
s['abc'] = 123
def test_update_with_sharedict():
s = ShareDict()
s.update_with_key(a, key='a')
s.update_with_key(b, key='b')
s.update_with_key(c, key='c')
d = {'z': 5}
s2 = ShareDict()
s2.update_with_key(a, key='a')
s2.update_with_key(d, key='d')
s.update(s2)
assert s.dicts['a'] is s.dicts['a']
def test_ensure_dict():
d = {'x': 1}
assert ensure_dict(d) is d
sd = ShareDict()
sd.update(d)
assert type(ensure_dict(sd)) is dict
assert ensure_dict(sd) == d
class mydict(dict):
pass
md = mydict()
md['x'] = 1
assert type(ensure_dict(md)) is dict
assert ensure_dict(md) == d
def test_update_with_sharedict():
s = ShareDict()
s.update_with_key(a, key='a')
s.update_with_key(b, key='b')
s.update_with_key(c, key='c')
d = {'z': 5}
s2 = ShareDict()
s2.update_with_key(a, key='a')
s2.update_with_key(d, key='d')
s.update(s2)
assert s.dicts['a'] is s.dicts['a']
def test_ensure_dict():
d = {'x': 1}
assert ensure_dict(d) is d
sd = ShareDict()
sd.update(d)
assert type(ensure_dict(sd)) is dict
assert ensure_dict(sd) == d
class mydict(dict):
pass
md = mydict()
md['x'] = 1
assert type(ensure_dict(md)) is dict
assert ensure_dict(md) == d
def test_update_with_sharedict():
s = ShareDict()
s.update_with_key(a, key="a")
s.update_with_key(b, key="b")
s.update_with_key(c, key="c")
d = {"z": 5}
s2 = ShareDict()
s2.update_with_key(a, key="a")
s2.update_with_key(d, key="d")
s.update(s2)
assert s.dicts["a"] is s.dicts["a"]
def predict_vis(time_index, antenna1, antenna2,
dde1_jones=None, source_coh=None, dde2_jones=None,
die1_jones=None, base_vis=None, die2_jones=None):
have_a1 = dde1_jones is not None
have_a2 = dde2_jones is not None
have_bl = source_coh is not None
have_g1 = die1_jones is not None
have_coh = base_vis is not None
have_g2 = die2_jones is not None
if have_a1 ^ have_a2:
raise ValueError("Both dde1_jones and dde2_jones "
"must be present or absent")
have_ants = have_a1 and have_a2
if have_ants:
if dde1_jones.shape[2] != dde1_jones.chunks[2][0]:
raise ValueError("Subdivision of antenna dimension into "
"multiple chunks is not supported.")
if dde2_jones.shape[2] != dde2_jones.chunks[2][0]:
raise ValueError("Subdivision of antenna dimension into "
"multiple chunks is not supported.")
if dde1_jones.chunks != dde2_jones.chunks:
raise ValueError("dde1_jones.chunks != dde2_jones.chunks")
if len(dde1_jones.chunks[1]) != len(time_index.chunks[0]):
raise ValueError("Number of row chunks (%s) does not equal "
"number of time chunks (%s)." %
(time_index.chunks[0], dde1_jones.chunks[1]))
if have_g1 ^ have_g2:
raise ValueError("Both die1_jones and die2_jones "
"must be present or absent")
have_dies = have_g1 and have_g2
if have_dies:
if die1_jones.shape[1] != die1_jones.chunks[1][0]:
raise ValueError("Subdivision of antenna dimension into "
"multiple chunks is not supported.")
if die2_jones.shape[1] != die2_jones.chunks[1][0]:
raise ValueError("Subdivision of antenna dimension into "
"multiple chunks is not supported.")
if die1_jones.chunks != die2_jones.chunks:
raise ValueError("die1_jones.chunks != die2_jones.chunks")
if len(die1_jones.chunks[0]) != len(time_index.chunks[0]):
raise ValueError("Number of row chunks (%s) does not equal "
"number of time chunks (%s)." %
(time_index.chunks[0], die1_jones.chunks[1]))
# Generate strings for the correlation dimensions
if have_ants:
cdims = tuple("corr-%d" % i for i in range(len(dde1_jones.shape[4:])))
elif have_bl:
cdims = tuple("corr-%d" % i for i in range(len(source_coh.shape[3:])))
elif have_dies:
cdims = tuple("corr-%d" % i for i in range(len(die1_jones.shape[3:])))
else:
raise ValueError("Missing both antenna and baseline jones terms")
# Infer the output dtype
dtype_arrays = [dde1_jones, source_coh, dde2_jones, die1_jones, die2_jones]
out_dtype = np.result_type(*(np.dtype(a.dtype.name)
for a in dtype_arrays if a is not None))
# In the case of predict_vis, the "row" and "time" dimensions
# are intimately related -- a contiguous series of rows
# are related to a contiguous series of timesteps.
# This means that the number of chunks of these
# two dimensions must match even though the chunk sizes may not.
# da.core.atop insists on matching chunk sizes.
# For this reason, we use the lower level da.core.top and
# substitute "row" for "time" in arrays such as dde1_jones
# and die1_jones.
token = da.core.tokenize(time_index, antenna1, antenna2,
dde1_jones, source_coh, dde2_jones, base_vis)
ajones_dims = ("src", "row", "ant", "chan") + cdims
gjones_dims = ("row", "ant", "chan") + cdims
# Setup
# 1. Optional top arguments
# 2. Optional numblocks kwarg
# 3. dask graph inputs
array_dsk = ShareDict()
top_args = [time_index.name, ("row",),
antenna1.name, ("row",),
antenna2.name, ("row",)]
numblocks = {
time_index.name: time_index.numblocks,
antenna1.name: antenna1.numblocks,
antenna2.name: antenna2.numblocks
}
# Merge input graphs into the top graph
array_dsk.update(time_index.__dask_graph__())
array_dsk.update(antenna1.__dask_graph__())
array_dsk.update(antenna2.__dask_graph__())
# Handle presence/absence of dde1_jones
if have_ants:
top_args.extend([dde1_jones.name, ajones_dims])
numblocks[dde1_jones.name] = dde1_jones.numblocks
array_dsk.update(dde1_jones.__dask_graph__())
other_chunks = dde1_jones.chunks[3:]
src_chunks = dde1_jones.chunks[0]
else:
top_args.extend([None, None])
# Handle presence/absence of source_coh
if have_bl:
top_args.extend([source_coh.name, ("src", "row", "chan") + cdims])
numblocks[source_coh.name] = source_coh.numblocks
other_chunks = source_coh.chunks[2:]
src_chunks = source_coh.chunks[0]
array_dsk.update(source_coh.__dask_graph__())
else:
top_args.extend([None, None])
# Handle presence/absence of dde2_jones
if have_ants:
top_args.extend([dde2_jones.name, ajones_dims])
numblocks[dde2_jones.name] = dde2_jones.numblocks
other_chunks = dde1_jones.chunks[3:]
array_dsk.update(dde2_jones.__dask_graph__())
other_chunks = dde2_jones.chunks[3:]
src_chunks = dde1_jones.chunks[0]
else:
top_args.extend([None, None])
# die1_jones, base_vis and die2_jones absent for this part of the graph
top_args.extend([None, None, None, None, None, None])
assert len(top_args) // 2 == 9, len(top_args) // 2
name = "-".join(("predict_vis", token))
dsk = da.core.top(_predict_coh_wrapper,
name, ("src", "row", "chan") + cdims,
*top_args,
numblocks=numblocks)
array_dsk.update(dsk)
# We can infer output chunk sizes from source_coh
chunks = ((1,)*len(src_chunks), time_index.chunks[0],) + other_chunks
sum_coherencies = da.Array(array_dsk, name, chunks, dtype=out_dtype)
sum_coherencies = sum_coherencies.sum(axis=0)
if have_coh:
sum_coherencies += base_vis
if not have_dies:
return sum_coherencies
# Now apply any Direction Independent Effect Terms
# Setup
# 1. Optional top arguments
# 2. Optional numblocks kwarg
# 3. dask graph inputs
array_dsk = ShareDict()
top_args = [time_index.name, ("row",),
antenna1.name, ("row",),
antenna2.name, ("row",)]
numblocks = {
time_index.name: time_index.numblocks,
antenna1.name: antenna1.numblocks,
antenna2.name: antenna2.numblocks
}
array_dsk.update(time_index.__dask_graph__())
array_dsk.update(antenna1.__dask_graph__())
array_dsk.update(antenna2.__dask_graph__())
# dde1_jones, source_coh and dde2_jones not present
top_args.extend([None, None, None, None, None, None])
top_args.extend([die1_jones.name, gjones_dims])
top_args.extend([sum_coherencies.name, ("row", "chan") + cdims])
top_args.extend([die2_jones.name, gjones_dims])
numblocks[die1_jones.name] = die1_jones.numblocks
numblocks[sum_coherencies.name] = sum_coherencies.numblocks
numblocks[die2_jones.name] = die2_jones.numblocks
array_dsk.update(die1_jones.__dask_graph__())
array_dsk.update(sum_coherencies.__dask_graph__())
array_dsk.update(die2_jones.__dask_graph__())
assert len(top_args) // 2 == 9
token = da.core.tokenize(time_index, antenna1, antenna2,
die1_jones, sum_coherencies, die2_jones)
name = '-'.join(("predict_vis", token))
dsk = da.core.top(_predict_dies_wrapper,
name, ("row", "chan") + cdims,
*top_args, numblocks=numblocks)
array_dsk.update(dsk)
chunks = (time_index.chunks[0],) + other_chunks
return da.Array(array_dsk, name, chunks, dtype=out_dtype)