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