def test_arithmetic(): x = np.arange(5).astype('f4') + 2 y = np.arange(5).astype('i8') + 2 z = np.arange(5).astype('i4') + 2 a = da.from_array(x, chunks=(2,)) b = da.from_array(y, chunks=(2,)) c = da.from_array(z, chunks=(2,)) assert eq(a + b, x + y) assert eq(a * b, x * y) assert eq(a - b, x - y) assert eq(a / b, x / y) assert eq(b & b, y & y) assert eq(b | b, y | y) assert eq(b ^ b, y ^ y) assert eq(a // b, x // y) assert eq(a ** b, x ** y) assert eq(a % b, x % y) assert eq(a > b, x > y) assert eq(a < b, x < y) assert eq(a >= b, x >= y) assert eq(a <= b, x <= y) assert eq(a == b, x == y) assert eq(a != b, x != y) assert eq(a + 2, x + 2) assert eq(a * 2, x * 2) assert eq(a - 2, x - 2) assert eq(a / 2, x / 2) assert eq(b & True, y & True) assert eq(b | True, y | True) assert eq(b ^ True, y ^ True) assert eq(a // 2, x // 2) assert eq(a ** 2, x ** 2) assert eq(a % 2, x % 2) assert eq(a > 2, x > 2) assert eq(a < 2, x < 2) assert eq(a >= 2, x >= 2) assert eq(a <= 2, x <= 2) assert eq(a == 2, x == 2) assert eq(a != 2, x != 2) assert eq(2 + b, 2 + y) assert eq(2 * b, 2 * y) assert eq(2 - b, 2 - y) assert eq(2 / b, 2 / y) assert eq(True & b, True & y) assert eq(True | b, True | y) assert eq(True ^ b, True ^ y) assert eq(2 // b, 2 // y) assert eq(2 ** b, 2 ** y) assert eq(2 % b, 2 % y) assert eq(2 > b, 2 > y) assert eq(2 < b, 2 < y) assert eq(2 >= b, 2 >= y) assert eq(2 <= b, 2 <= y) assert eq(2 == b, 2 == y) assert eq(2 != b, 2 != y) assert eq(-a, -x) assert eq(abs(a), abs(x)) assert eq(~(a == b), ~(x == y)) assert eq(~(a == b), ~(x == y)) assert eq(da.logaddexp(a, b), np.logaddexp(x, y)) assert eq(da.logaddexp2(a, b), np.logaddexp2(x, y)) assert eq(da.exp(b), np.exp(y)) assert eq(da.log(a), np.log(x)) assert eq(da.log10(a), np.log10(x)) assert eq(da.log1p(a), np.log1p(x)) assert eq(da.expm1(b), np.expm1(y)) assert eq(da.sqrt(a), np.sqrt(x)) assert eq(da.square(a), np.square(x)) assert eq(da.sin(a), np.sin(x)) assert eq(da.cos(b), np.cos(y)) assert eq(da.tan(a), np.tan(x)) assert eq(da.arcsin(b/10), np.arcsin(y/10)) assert eq(da.arccos(b/10), np.arccos(y/10)) assert eq(da.arctan(b/10), np.arctan(y/10)) assert eq(da.arctan2(b*10, a), np.arctan2(y*10, x)) assert eq(da.hypot(b, a), np.hypot(y, x)) assert eq(da.sinh(a), np.sinh(x)) assert eq(da.cosh(b), np.cosh(y)) assert eq(da.tanh(a), np.tanh(x)) assert eq(da.arcsinh(b*10), np.arcsinh(y*10)) assert eq(da.arccosh(b*10), np.arccosh(y*10)) assert eq(da.arctanh(b/10), np.arctanh(y/10)) assert eq(da.deg2rad(a), np.deg2rad(x)) assert eq(da.rad2deg(a), np.rad2deg(x)) assert eq(da.logical_and(a < 1, b < 4), np.logical_and(x < 1, y < 4)) assert eq(da.logical_or(a < 1, b < 4), np.logical_or(x < 1, y < 4)) assert eq(da.logical_xor(a < 1, b < 4), np.logical_xor(x < 1, y < 4)) assert eq(da.logical_not(a < 1), np.logical_not(x < 1)) assert eq(da.maximum(a, 5 - a), np.maximum(a, 5 - a)) assert eq(da.minimum(a, 5 - a), np.minimum(a, 5 - a)) assert eq(da.fmax(a, 5 - a), np.fmax(a, 5 - a)) assert eq(da.fmin(a, 5 - a), np.fmin(a, 5 - a)) assert eq(da.isreal(a + 1j * b), np.isreal(x + 1j * y)) assert eq(da.iscomplex(a + 1j * b), np.iscomplex(x + 1j * y)) assert eq(da.isfinite(a), np.isfinite(x)) assert eq(da.isinf(a), np.isinf(x)) assert eq(da.isnan(a), np.isnan(x)) assert eq(da.signbit(a - 3), np.signbit(x - 3)) assert eq(da.copysign(a - 3, b), np.copysign(x - 3, y)) assert eq(da.nextafter(a - 3, b), np.nextafter(x - 3, y)) assert eq(da.ldexp(c, c), np.ldexp(z, z)) assert eq(da.fmod(a * 12, b), np.fmod(x * 12, y)) assert eq(da.floor(a * 0.5), np.floor(x * 0.5)) assert eq(da.ceil(a), np.ceil(x)) assert eq(da.trunc(a / 2), np.trunc(x / 2)) assert eq(da.degrees(b), np.degrees(y)) assert eq(da.radians(a), np.radians(x)) assert eq(da.rint(a + 0.3), np.rint(x + 0.3)) assert eq(da.fix(a - 2.5), np.fix(x - 2.5)) assert eq(da.angle(a + 1j), np.angle(x + 1j)) assert eq(da.real(a + 1j), np.real(x + 1j)) assert eq((a + 1j).real, np.real(x + 1j)) assert eq(da.imag(a + 1j), np.imag(x + 1j)) assert eq((a + 1j).imag, np.imag(x + 1j)) assert eq(da.conj(a + 1j * b), np.conj(x + 1j * y)) assert eq((a + 1j * b).conj(), (x + 1j * y).conj()) assert eq(da.clip(b, 1, 4), np.clip(y, 1, 4)) assert eq(da.fabs(b), np.fabs(y)) assert eq(da.sign(b - 2), np.sign(y - 2)) l1, l2 = da.frexp(a) r1, r2 = np.frexp(x) assert eq(l1, r1) assert eq(l2, r2) l1, l2 = da.modf(a) r1, r2 = np.modf(x) assert eq(l1, r1) assert eq(l2, r2) assert eq(da.around(a, -1), np.around(x, -1))
def wiener(data, aux, fr, fr_npy, L): return uirdft2((da.conj(fr) * urdft2(data) + L * urdft2(aux)) / (da.absolute(fr_npy)**2 + L))
def wavg_ant(data, flags, weights, ant_array, bls_lookup, threshold=0.8): """Perform weighted average of data, flags and weights, over axis -1. This applies flags and is done per antenna. Parameters ---------- data : :class:`da.Array` complex (..., bls) flags : :class:`da.Array` int (..., bls) weights : :class:`da.Array` real (..., bls) ant_array : :class:`np.ndarray` array of strings representing antennas bls_lookup : :class:`np.ndarray` (bls x 2) array of antennas in each baseline threshold : float if fraction of flags in the input data array exceeds threshold then set output flag to True, else False Returns ------- av_data : :class:`da.Array` complex (..., n_ant), weighted average of data av_flags : :class:`da.Array` bool (n_ant), weighted average of flags av_weights : :class:`da.Array` real (..., n_ant), weighted average of weights """ av_data = [] av_flags = [] av_weights = [] weighted_data, flagged_weights = weight_data(data, flags, weights) for ant in range(len(ant_array)): # select all correlations with same antenna but ignore autocorrelations ant_idx = np.where((bls_lookup[:, 0] == ant) ^ (bls_lookup[:, 1] == ant))[0] # conjugate visibilities if antenna is 2nd on the baseline ant_data = weighted_data[..., ant_idx] ant1 = bls_lookup[ant_idx][:, 0] == ant ant1 = np.broadcast_to(ant1, ant_data.shape) ant_conj_data = where(ant1, ant_data, da.conj(ant_data)) ant_ave_data, ant_weights = _wavg_axis(ant_conj_data, flagged_weights[..., ant_idx], axis=-1) # update flags to include all invalid data updated_flags = flagged_weights[..., ant_idx] == 0 n_flags = da.sum(updated_flags, axis=-1) ant_flags = n_flags > ant_idx.shape[0] * threshold av_data.append(ant_ave_data) av_flags.append(ant_flags) av_weights.append(ant_weights) av_data = da.stack(av_data, axis=-1) av_flags = da.stack(av_flags, axis=-1) av_weights = da.stack(av_weights, axis=-1) return av_data, av_flags, av_weights
def test_arithmetic(): x = np.arange(5).astype('f4') + 2 y = np.arange(5).astype('i8') + 2 z = np.arange(5).astype('i4') + 2 a = da.from_array(x, chunks=(2, )) b = da.from_array(y, chunks=(2, )) c = da.from_array(z, chunks=(2, )) assert eq(a + b, x + y) assert eq(a * b, x * y) assert eq(a - b, x - y) assert eq(a / b, x / y) assert eq(b & b, y & y) assert eq(b | b, y | y) assert eq(b ^ b, y ^ y) assert eq(a // b, x // y) assert eq(a**b, x**y) assert eq(a % b, x % y) assert eq(a > b, x > y) assert eq(a < b, x < y) assert eq(a >= b, x >= y) assert eq(a <= b, x <= y) assert eq(a == b, x == y) assert eq(a != b, x != y) assert eq(a + 2, x + 2) assert eq(a * 2, x * 2) assert eq(a - 2, x - 2) assert eq(a / 2, x / 2) assert eq(b & True, y & True) assert eq(b | True, y | True) assert eq(b ^ True, y ^ True) assert eq(a // 2, x // 2) assert eq(a**2, x**2) assert eq(a % 2, x % 2) assert eq(a > 2, x > 2) assert eq(a < 2, x < 2) assert eq(a >= 2, x >= 2) assert eq(a <= 2, x <= 2) assert eq(a == 2, x == 2) assert eq(a != 2, x != 2) assert eq(2 + b, 2 + y) assert eq(2 * b, 2 * y) assert eq(2 - b, 2 - y) assert eq(2 / b, 2 / y) assert eq(True & b, True & y) assert eq(True | b, True | y) assert eq(True ^ b, True ^ y) assert eq(2 // b, 2 // y) assert eq(2**b, 2**y) assert eq(2 % b, 2 % y) assert eq(2 > b, 2 > y) assert eq(2 < b, 2 < y) assert eq(2 >= b, 2 >= y) assert eq(2 <= b, 2 <= y) assert eq(2 == b, 2 == y) assert eq(2 != b, 2 != y) assert eq(-a, -x) assert eq(abs(a), abs(x)) assert eq(~(a == b), ~(x == y)) assert eq(~(a == b), ~(x == y)) assert eq(da.logaddexp(a, b), np.logaddexp(x, y)) assert eq(da.logaddexp2(a, b), np.logaddexp2(x, y)) assert eq(da.exp(b), np.exp(y)) assert eq(da.log(a), np.log(x)) assert eq(da.log10(a), np.log10(x)) assert eq(da.log1p(a), np.log1p(x)) assert eq(da.expm1(b), np.expm1(y)) assert eq(da.sqrt(a), np.sqrt(x)) assert eq(da.square(a), np.square(x)) assert eq(da.sin(a), np.sin(x)) assert eq(da.cos(b), np.cos(y)) assert eq(da.tan(a), np.tan(x)) assert eq(da.arcsin(b / 10), np.arcsin(y / 10)) assert eq(da.arccos(b / 10), np.arccos(y / 10)) assert eq(da.arctan(b / 10), np.arctan(y / 10)) assert eq(da.arctan2(b * 10, a), np.arctan2(y * 10, x)) assert eq(da.hypot(b, a), np.hypot(y, x)) assert eq(da.sinh(a), np.sinh(x)) assert eq(da.cosh(b), np.cosh(y)) assert eq(da.tanh(a), np.tanh(x)) assert eq(da.arcsinh(b * 10), np.arcsinh(y * 10)) assert eq(da.arccosh(b * 10), np.arccosh(y * 10)) assert eq(da.arctanh(b / 10), np.arctanh(y / 10)) assert eq(da.deg2rad(a), np.deg2rad(x)) assert eq(da.rad2deg(a), np.rad2deg(x)) assert eq(da.logical_and(a < 1, b < 4), np.logical_and(x < 1, y < 4)) assert eq(da.logical_or(a < 1, b < 4), np.logical_or(x < 1, y < 4)) assert eq(da.logical_xor(a < 1, b < 4), np.logical_xor(x < 1, y < 4)) assert eq(da.logical_not(a < 1), np.logical_not(x < 1)) assert eq(da.maximum(a, 5 - a), np.maximum(a, 5 - a)) assert eq(da.minimum(a, 5 - a), np.minimum(a, 5 - a)) assert eq(da.fmax(a, 5 - a), np.fmax(a, 5 - a)) assert eq(da.fmin(a, 5 - a), np.fmin(a, 5 - a)) assert eq(da.isreal(a + 1j * b), np.isreal(x + 1j * y)) assert eq(da.iscomplex(a + 1j * b), np.iscomplex(x + 1j * y)) assert eq(da.isfinite(a), np.isfinite(x)) assert eq(da.isinf(a), np.isinf(x)) assert eq(da.isnan(a), np.isnan(x)) assert eq(da.signbit(a - 3), np.signbit(x - 3)) assert eq(da.copysign(a - 3, b), np.copysign(x - 3, y)) assert eq(da.nextafter(a - 3, b), np.nextafter(x - 3, y)) assert eq(da.ldexp(c, c), np.ldexp(z, z)) assert eq(da.fmod(a * 12, b), np.fmod(x * 12, y)) assert eq(da.floor(a * 0.5), np.floor(x * 0.5)) assert eq(da.ceil(a), np.ceil(x)) assert eq(da.trunc(a / 2), np.trunc(x / 2)) assert eq(da.degrees(b), np.degrees(y)) assert eq(da.radians(a), np.radians(x)) assert eq(da.rint(a + 0.3), np.rint(x + 0.3)) assert eq(da.fix(a - 2.5), np.fix(x - 2.5)) assert eq(da.angle(a + 1j), np.angle(x + 1j)) assert eq(da.real(a + 1j), np.real(x + 1j)) assert eq((a + 1j).real, np.real(x + 1j)) assert eq(da.imag(a + 1j), np.imag(x + 1j)) assert eq((a + 1j).imag, np.imag(x + 1j)) assert eq(da.conj(a + 1j * b), np.conj(x + 1j * y)) assert eq((a + 1j * b).conj(), (x + 1j * y).conj()) assert eq(da.clip(b, 1, 4), np.clip(y, 1, 4)) assert eq(da.fabs(b), np.fabs(y)) assert eq(da.sign(b - 2), np.sign(y - 2)) l1, l2 = da.frexp(a) r1, r2 = np.frexp(x) assert eq(l1, r1) assert eq(l2, r2) l1, l2 = da.modf(a) r1, r2 = np.modf(x) assert eq(l1, r1) assert eq(l2, r2) assert eq(da.around(a, -1), np.around(x, -1))
def _rmatvec(self, x): x1 = da.conj(x) y1 = da.conj(self.oOp._rmatvec(x1)) return y1
def cross_phase(da1, da2, spacing_tol=1e-3, dim=None, detrend=None, window=False, chunks_to_segments=False): """ Calculates the cross-phase between da1 and da2. Returned values are in [-pi, pi]. .. math:: da1' = da1 - \overline{da1};\ \ da2' = da2 - \overline{da2} .. math:: cp = \text{Arg} [\mathbb{F}(da1')^*, \mathbb{F}(da2')] Parameters ---------- da1 : `xarray.DataArray` The data to be transformed da2 : `xarray.DataArray` The data to be transformed spacing_tol: float, optional Spacing tolerance. Fourier transform should not be applied to uneven grid but this restriction can be relaxed with this setting. Use caution. dim : list, optional The dimension along which to take the real Fourier transformation. If `None`, all dimensions will be transformed. shift : bool, optional Whether to shift the fft output. detrend : str, optional If `constant`, the mean across the transform dimensions will be subtracted before calculating the Fourier transform (FT). If `linear`, the linear least-square fit along one axis will be subtracted before the FT. It will give an error if the length of `dim` is longer than one. window : bool, optional Whether to apply a Hann window to the data before the Fourier transform is taken Returns ------- cp : `xarray.DataArray` Cross-phase as a function of frequency. """ if dim is None: dim = da1.dims dim2 = da2.dims if dim != dim2: raise ValueError('The two datasets have different dimensions') elif not isinstance(dim, list): dim = [dim] if len(dim) > 1: raise ValueError('Cross phase calculation should only be done along ' 'a single dimension.') daft1 = dft(da1, spacing_tol, dim=dim, real=dim[0], shift=False, detrend=detrend, window=window, chunks_to_segments=chunks_to_segments) daft2 = dft(da2, spacing_tol, dim=dim, real=dim[0], shift=False, detrend=detrend, window=window, chunks_to_segments=chunks_to_segments) if daft1.chunks and daft2.chunks: _cross_phase = lambda a, b: dsar.angle(a * dsar.conj(b)) else: _cross_phase = lambda a, b: np.angle(a * np.conj(b)) cp = xr.apply_ufunc(_cross_phase, daft1, daft2, dask='allowed') if da1.name and da2.name: cp.name = "{}_{}_phase".format(da1.name, da2.name) return cp
def wiener(data, aux, fr, L): l_del = delayed(L) return uirdft2((da.conj(delayed(fr)) * urdft2(data) + l_del * urdft2(aux)) * ((da.absolute(fr_npy)**2 + l_del)**-1))
def cross_phase(da1, da2, spacing_tol=1e-3, dim=None, detrend=None, window=False, chunks_to_segments=False): """ Calculates the cross-phase between da1 and da2. Returned values are in [-pi, pi]. .. math:: da1' = da1 - \overline{da1};\ \ da2' = da2 - \overline{da2} .. math:: cp = \text{Arg} [\mathbb{F}(da1')^*, \mathbb{F}(da2')] Parameters ---------- da1 : `xarray.DataArray` The data to be transformed da2 : `xarray.DataArray` The data to be transformed spacing_tol: float, optional Spacing tolerance. Fourier transform should not be applied to uneven grid but this restriction can be relaxed with this setting. Use caution. dim : list, optional The dimension along which to take the real Fourier transformation. If `None`, all dimensions will be transformed. shift : bool, optional Whether to shift the fft output. detrend : str, optional If `constant`, the mean across the transform dimensions will be subtracted before calculating the Fourier transform (FT). If `linear`, the linear least-square fit along one axis will be subtracted before the FT. It will give an error if the length of `dim` is longer than one. window : bool, optional Whether to apply a Hann window to the data before the Fourier transform is taken Returns ------- cp : `xarray.DataArray` Cross-phase as a function of frequency. """ if dim is None: dim = da1.dims dim2 = da2.dims if dim != dim2: raise ValueError('The two datasets have different dimensions') elif not isinstance(dim, list): dim = [dim] if len(dim)>1: raise ValueError('Cross phase calculation should only be done along ' 'a single dimension.') daft1 = dft(da1, spacing_tol, dim=dim, real=dim[0], shift=False, detrend=detrend, window=window, chunks_to_segments=chunks_to_segments) daft2 = dft(da2, spacing_tol, dim=dim, real=dim[0], shift=False, detrend=detrend, window=window, chunks_to_segments=chunks_to_segments) if daft1.chunks and daft2.chunks: _cross_phase = lambda a, b: dsar.angle(a * dsar.conj(b)) else: _cross_phase = lambda a, b: np.angle(a * np.conj(b)) cp = xr.apply_ufunc(_cross_phase, daft1, daft2, dask='allowed') if da1.name and da2.name: cp.name = "{}_{}_phase".format(da1.name, da2.name) return cp