def model(uvw,
freq,
image,
freq_bin_idx,
freq_bin_counts,
cell,
weights=None,
flag=None,
celly=None,
epsilon=1e-5,
nthreads=1,
do_wstacking=True):
# determine output type
complex_type = da.result_type(image, np.complex64)
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')
vis = da.blockwise(_model_wrapper, ('row', 'chan'),
uvw, ('row', 'three'),
freq, ('chan', ),
image, ('chan', 'nx', 'ny'),
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.chunks[0]},
dtype=complex_type,
align_arrays=False)
return vis
def dask_array_select(condlist, choicelist, default=0):
# shim taken from dask.array 2021.6.1
# https://github.com/ibis-project/ibis/issues/2847
try:
from dask.array import select
return select(condlist, choicelist, default)
except ImportError:
def _select(*args, **kwargs):
split_at = len(args) // 2
condlist = args[:split_at]
choicelist = args[split_at:]
return np.select(condlist, choicelist, **kwargs)
if len(condlist) != len(choicelist):
raise ValueError(
"list of cases must be same length as list of conditions"
)
if len(condlist) == 0:
raise ValueError(
"select with an empty condition list is not possible"
)
choicelist = [da.asarray(choice) for choice in choicelist]
try:
intermediate_dtype = da.result_type(*choicelist)
except TypeError as e:
msg = "Choicelist elements do not have a common dtype."
raise TypeError(msg) from e
blockwise_shape = tuple(range(choicelist[0].ndim))
condargs = [
arg for elem in condlist for arg in (elem, blockwise_shape)
]
choiceargs = [
arg for elem in choicelist for arg in (elem, blockwise_shape)
]
return da.blockwise(
_select,
blockwise_shape,
*condargs,
*choiceargs,
dtype=intermediate_dtype,
name="select",
default=default,
)
def test_result_type():
a = da.from_array(np.ones(5, np.float32), chunks=(3, ))
b = da.from_array(np.ones(5, np.int16), chunks=(3, ))
c = da.from_array(np.ones(5, np.int64), chunks=(3, ))
x = np.ones(5, np.float32)
assert da.result_type(b, c) == np.int64
assert da.result_type(a, b, c) == np.float64
assert da.result_type(b, np.float32) == np.float32
assert da.result_type(b, np.dtype(np.float32)) == np.float32
assert da.result_type(b, x) == np.float32
# Effect of scalars depends on their value
assert da.result_type(1, b) == np.int16
assert da.result_type(1.0, a) == np.float32
assert da.result_type(np.int64(1), b) == np.int16
assert da.result_type(np.ones((), np.int64), b) == np.int16 # 0d array
assert da.result_type(1e200,
a) == np.float64 # 1e200 is too big for float32
# dask 0d-arrays are NOT treated like scalars
c = da.from_array(np.ones((), np.float64), chunks=())
assert da.result_type(a, c) == np.float64
def test_result_type():
a = da.from_array(np.ones(5, np.float32), chunks=(3,))
b = da.from_array(np.ones(5, np.int16), chunks=(3,))
c = da.from_array(np.ones(5, np.int64), chunks=(3,))
x = np.ones(5, np.float32)
assert da.result_type(b, c) == np.int64
assert da.result_type(a, b, c) == np.float64
assert da.result_type(b, np.float32) == np.float32
assert da.result_type(b, np.dtype(np.float32)) == np.float32
assert da.result_type(b, x) == np.float32
# Effect of scalars depends on their value
assert da.result_type(1, b) == np.int16
assert da.result_type(1.0, a) == np.float32
assert da.result_type(np.int64(1), b) == np.int16
assert da.result_type(np.ones((), np.int64), b) == np.int16 # 0d array
assert da.result_type(1e200, a) == np.float64 # 1e200 is too big for float32
# dask 0d-arrays are NOT treated like scalars
c = da.from_array(np.ones((), np.float64), chunks=())
assert da.result_type(a, c) == np.float64
