def blstd(data, mask): """ Calculate std over baselines (ignoring zeros). Expects masked array as input """ nint, nbl, nchan, npol = data.shape # getting "data type not understood" if this typed as float32 blstd = np.zeros((nint, nchan, npol), dtype=complex64) for i in range(nint): for j in range(nchan): for k in range(npol): ss = complex64(0) weight = int64(0) for l in range(nbl): ss += data[i, l, j, k] if mask[i, l, j, k] is False: weight += 1 if weight > 0: mean = ss / weight ss = complex64(0) for l in range(nbl): if mask[i, l, j, k] is False: ss += np.abs((data[i, l, j, k] - mean)**2) blstd[i, j, k] = np.sqrt(ss / weight) else: blstd[i, j, k] = complex64(0) return blstd.real
def _meantsub_jit(data): """ Calculate mean in time (ignoring zeros) and subtract in place Could ultimately parallelize by computing only on subset of data. """ nint, nbl, nchan, npol = data.shape for i in range(nbl): for j in range(nchan): for k in range(npol): ss = complex64(0) weight = int64(0) for l in range(nint): ss += data[l, i, j, k] if data[l, i, j, k] != 0j: weight += 1 if weight > 0: mean = ss/weight else: mean = complex64(0) if mean: for l in range(nint): if data[l, i, j, k] != 0j: data[l, i, j, k] -= mean
def _meantsub_jit(data): """ Calculate mean in time (ignoring zeros) and subtract in place Could ultimately parallelize by computing only on subset of data. """ nint, nbl, nchan, npol = data.shape for i in range(nbl): for j in range(nchan): for k in range(npol): ss = complex64(0) weight = int64(0) for l in range(nint): ss += data[l, i, j, k] if data[l, i, j, k] != 0j: weight += 1 if weight > 0: mean = ss / weight else: mean = complex64(0) if mean: for l in range(nint): if data[l, i, j, k] != 0j: data[l, i, j, k] -= mean
def _sinewave(num, den): """Generate a complex sine wave of frequency sample_rate*num/den. Length is chosen such that a continuous sine wave can be made by repeating the returned signal.""" # The code below fails for num=0, so handle that as a special case if num == 0: return nb.complex64([1.0]) # "% den" is not absolutely necessary here, but wrapping the phase # using integers may avoid loss of floating point precision. phase = \ (np.arange(0, num*den, num, dtype = np.int64) % den) \ .astype(np.float32) * nb.float32(2.0 * np.pi / den) return np.cos(phase) + np.sin(phase) * nb.complex64(1j)
def _meantsub_gu(data): b""" Subtract time mean while ignoring zeros. Vectorizes over time axis. Assumes time axis is last so use np.swapaxis(0,3) when passing visibility array in **CURRENTLY NOT WORKING WITH FLAGS** """ ss = complex64(0) weight = int64(0) for i in range(data.shape[0]): ss += data[i] if data[i] != complex64(0): weight += 1 mean = ss / weight for i in range(data.shape[0]): data[i] -= mean
def _meantsub_gu(data): b""" Subtract time mean while ignoring zeros. Vectorizes over time axis. Assumes time axis is last so use np.swapaxis(0,3) when passing visibility array in **CURRENTLY NOT WORKING WITH FLAGS** """ ss = complex64(0) weight = int64(0) for i in range(data.shape[0]): ss += data[i] if data[i] != complex64(0): weight += 1 mean = ss/weight for i in range(data.shape[0]): data[i] -= mean
def _2ptsub_jit(data): """ Calculate 2-pt time trend and evaluate to subtract at each time. """ nint, nbl, nchan, npol = data.shape for i in range(nbl): for j in range(nchan): for k in range(npol): # first half mean ss1 = complex64(0) weight1 = int64(0) for l in range(0, nint // 2): ss1 += data[l, i, j, k] if data[l, i, j, k] != 0j: weight1 += 1 if weight1 > 0: mean1 = ss1 / weight1 else: mean1 = complex64(0) # second half mean ss2 = complex64(0) weight2 = int64(0) for l in range(nint // 2, nint): ss2 += data[l, i, j, k] if data[l, i, j, k] != 0j: weight2 += 1 if weight2 > 0: mean2 = ss2 / weight2 else: mean2 = complex64(0) # calc mean per int if mean1 and mean2: slope = (mean2 - mean1) / (nint // 2) mean0 = (mean2 + mean1) / 2 for l in range(nint): if data[l, i, j, k] != 0j: data[l, i, j, k] -= slope * (l - nint // 2) + mean0 else: # or just blank data for l in range(nint): data[l, i, j, k] = 0j
def execute(self, input): firlen = self.firlen interpolation = self.interpolation decimation = self.decimation lo_phase = self.lo_phase fir_phase = self.fir_phase fir_i = self.fir_i #output = np.zeros_like(input) output = np.zeros(1000, dtype=np.complex64) # TODO: size of the array outn = 0 # number of output samples produced for sample_in in input: # Mix input with local oscillator mixed = sample_in * self.lo_table[lo_phase] #lo_phase = (lo_phase + 1) % len(self.lo_table) lo_phase += 1 if lo_phase >= len(self.lo_table): lo_phase = 0 # Store to buffer. # Fake circular buffering by storing the samples twice. self.firbuf[fir_i ] = \ self.firbuf[fir_i + firlen] = mixed #fir_i = (fir_i + 1) % firlen fir_i += 1 if fir_i >= firlen: fir_i = 0 fir_phase += interpolation while fir_phase >= decimation: fir_phase -= decimation assert fir_phase >= 0 assert fir_phase < interpolation t = self.taps[fir_phase::interpolation] fb = self.firbuf[fir_i:fir_i + firlen] # This loop seems to be faster than calling np.dot o = nb.complex64(0.0) for j in range(firlen): o += t[j] * fb[j] output[outn] = o outn += 1 assert outn <= len(output) self.lo_phase = lo_phase self.fir_phase = fir_phase self.fir_i = fir_i return output[0:outn]
def _2ptinterp_jit(data): """ Calculate 2-pt time trend and evaluate to subtract at each time. """ nint, nbl, nchan, npol = data.shape for i in range(nbl): for j in range(nchan): for k in range(npol): # first half mean ss1 = complex64(0) weight1 = int64(0) for l in range(0, nint // 2): ss1 += data[l, i, j, k] if data[l, i, j, k] != 0j: weight1 += 1 if weight1 > 0: mean1 = ss1 / weight1 else: mean1 = complex64(0) # second half mean ss2 = complex64(0) weight2 = int64(0) for l in range(nint // 2, nint): ss2 += data[l, i, j, k] if data[l, i, j, k] != 0j: weight2 += 1 if weight2 > 0: mean2 = ss2 / weight2 else: mean2 = complex64(0) ff = interpolate.interp1d([nint // 4, 3 * nint // 4], [mean1, mean2], fill_value='extrapolate') for l in range(nint): if data[l, i, j, k] != 0j: data[l, i, j, k] -= slope * (l - nint // 2) + mean0
def _cssub1_jit(data, dataavginterp): """ Use interpolated data to subtract while ignoring zeros """ nint, nbl, nchan, npol = data.shape for i in range(nbl): for j in range(nchan): for k in range(npol): for l in range(nint): if data[l, i, j, k] == 0j: dataavginterp[l, i, j, k] = complex64(0) data[l, i, j, k] -= dataavginterp[l, i, j, k]
def meantsub_cuda(data, cache=True): """ Calculate mean in time (ignoring zeros) and subtract in place """ x, y, z = cuda.grid(3) nint, nbl, nchan, npol = data.shape if x < nbl and y < nchan and z < npol: sum = complex64(0) weight = 0 for i in range(nint): sum = sum + data[i, x, y, z] if data[i, x, y, z] == 0j: weight = weight + 1 mean = sum/weight for i in range(nint): data[i, x, y, z] = data[i, x, y, z] - mean
def meantsub_cuda(data, cache=True): """ Calculate mean in time (ignoring zeros) and subtract in place """ x, y, z = cuda.grid(3) nint, nbl, nchan, npol = data.shape if x < nbl and y < nchan and z < npol: sum = complex64(0) weight = 0 for i in range(nint): sum = sum + data[i, x, y, z] if data[i, x, y, z] == 0j: weight = weight + 1 mean = sum / weight for i in range(nint): data[i, x, y, z] = data[i, x, y, z] - mean
def _cssub0_jit(data, dataavg): """ Use scipy calculate 4-pt mean as input to spline estimate. zeroed data is treated as flagged """ nint, nbl, nchan, npol = data.shape piece = nint // 4 for i in range(nbl): for j in range(nchan): for k in range(npol): # mean in each piece for pp in range(4): ss = complex64(0) weight = int64(0) for l in range(pp * piece, (pp + 1) * piece): ss += data[l, i, j, k] if data[l, i, j, k] != 0j: weight += 1 if weight > 0: dataavg[pp, i, j, k] = ss / weight else: dataavg[pp, i, j, k] = complex64(0) # TODO: instead use nearest?
import numpy as np from numba import vectorize, complex64 from numpy import fft from scipy.fftpack import next_fast_len from scipy.ndimage import gaussian_filter1d try: from mkl_fft import fft2, ifft2 except ModuleNotFoundError: warnings.warn( "mkl_fft not installed. Install it with conda: conda install mkl_fft", ImportWarning) @vectorize([complex64(complex64, complex64)], nopython=True, target='parallel') def apply_dotnorm(Y, cfRefImg): return Y / (np.complex64(1e-5) + np.abs(Y)) * cfRefImg @vectorize([ 'complex64(int16, float32, float32)', 'complex64(float32, float32, float32)' ], nopython=True, target='parallel', cache=True) def addmultiply(x, mul, add): return np.complex64(np.float32(x) * mul + add)
import numba as nb import numpy as np from numba import jit, prange @nb.vectorize([nb.complex64(nb.float32), nb.complex128(nb.float64)]) def complex_exponential(x): return np.cos(x) + 1.j * np.sin(x) @nb.vectorize([nb.float32(nb.complex64), nb.float64(nb.complex128)]) def abs2(x): return x.real**2 + x.imag**2 @jit(nopython=True, nogil=True, parallel=True) def interpolate_radial_functions(array, disc_indices, positions, v, r, dvdr, sampling): n = r.shape[0] dt = np.log(r[-1] / r[0]) / (n - 1) for i in range(positions.shape[0]): for j in prange(disc_indices.shape[0]): k = int(round(positions[i, 0] / sampling[0]) + disc_indices[j, 0]) l = int(round(positions[i, 1] / sampling[1]) + disc_indices[j, 1]) if ((k < array.shape[0]) & (l < array.shape[1]) & (k >= 0) & (l >= 0)): r_interp = np.sqrt((k * sampling[0] - positions[i, 0])**2 + (l * sampling[1] - positions[i, 1])**2) idx = int(np.floor(np.log(r_interp / r[0] + 1e-7) / dt))
@classmethod def from_hdf5(cls, hdf5: h5py.File) -> 'PrimaryBeam': model_format = models.get_hdf5_attr(hdf5.attrs, 'model_format', str) or '' if model_format == 'aperture_plane': return PrimaryBeamAperturePlane.from_hdf5(hdf5) else: raise models.ModelFormatError( f'Unknown model_format {model_format!r} for {cls.model_type}') def to_hdf5(self, hdf5: h5py.File) -> None: raise NotImplementedError() # pragma: nocover @numba.vectorize([numba.complex64(numba.float32)]) def _expjm2pi(x): """Equivalent to ``exp(-2j * np.pi * x)`` where `x` is real. x is reduced to a small value before multiplication, which improves precision at a small cost in performance. """ y = np.float32(-2 * np.pi) * (x - np.rint(x)) return complex(np.cos(y), np.sin(y)) def _asarray(x: ArrayLike, dtype: Optional[DTypeLike] = None) -> np.ndarray: """Convert an array-like to an array. Unlike np.ndarray, this will reject astropy Quantities with dimensions and convert dimensionless quantities correctly even if they have scale.
from numba import vectorize, complex64, boolean, jit import time import numpy as np from stats import cpu_stats,gpu_stats from time import sleep @vectorize([boolean(complex64)]) def f(z): return (z.real*z.real + z.imag*z.imag) < 4.0 @vectorize([complex64(complex64, complex64)]) def g(z,c): return z*z + c @jit def mandelbrot_numpy(c, maxiter): output = np.zeros(c.shape, np.int) z = np.empty(c.shape, np.complex64) for it in range(maxiter): notdone = f(z) output[notdone] = it z[notdone] = g(z[notdone],c[notdone]) output[output == maxiter-1] = 0 return output def mandelbrot_set(width,height,xmin,xmax,ymin,ymax,maxiter): r1 = np.linspace(xmin, xmax, width, dtype=np.float32) r2 = np.linspace(ymin, ymax, height, dtype=np.float32) c = r1 + r2[:,None]*1j n3 = mandelbrot_numpy(c,maxiter) return (r1,r2,n3.T)
out[xs + d, ys + d] = err / cn return out.argmin() @_numba.vectorize( [_numba.float64(_numba.complex128), _numba.float32(_numba.complex64)], target="parallel", ) def abs2(x): return x.real * x.real + x.imag * x.imag @_numba.vectorize( [_numba.complex128(_numba.complex128), _numba.complex64(_numba.complex64)], target="parallel", ) def abs2c(x): return x.real * x.real + x.imag * x.imag + 0j def fill(data, invalid=None): """ fill invalid values by closest valid value. invalid: mask of invalid values, default: np.isnan(data) """ if invalid is None: invalid = _np.isnan(data) ind = _snd.distance_transform_edt(invalid, return_distances=False, return_indices=True)
def expj(x): return math.cos(x) + 1j * math.sin(x) def abs_sqd(x): return np.real(x * np.conj(x)) def sum_abs_sqd(x): return np.sum(np.real(x * np.conj(x))) def sum_abs_sqd_1d(x): return np.sum(np.real(x * np.conj(x)), axis=0) else: @numba.vectorize([numba.complex64(numba.float64)], nopython=True, cache=True) def expj(x): return math.cos(x) + 1j * math.sin(x) @numba.vectorize( cache=True ) # ([numba.float64(numba.complex128),numba.float32(numba.complex64)]) def abs_sqd(x): return x.real**2 + x.imag**2 @numba.jit(cache=True) def sum_abs_sqd_1d(x, cache=True): r = 0 for i in range(len(x)):
def smoothing_real_space(quantity,smoothing_scale): array_size = len(np.arange(0,cluster_grid,4)) smoothing_array = np.zeros([array_size,array_size,array_size],dtype=np.float32) for ix,i in enumerate(np.arange(0,cluster_grid,4)): for iy,j in enumerate(np.arange(0,cluster_grid,4)): for iz,k in enumerate(np.arange(0,cluster_grid,4)): i = int(i) j = int(j) k = int(k) smoothing_array[ix,iy,iz] = np.mean(quantity[i:i+smoothing_scale,j:j+smoothing_scale,k:k+smoothing_scale]) return smoothing_array @jit(complex64(complex64,float32,int32,int32)) def smoothing_k_space(fft_array,smoothing_scale,cluster_length,cluster_grid): smoothing_array = copy.deepcopy(fft_array) cluster_length = 40 # cluster_grid = 201 for ix in range(int(fft_array.shape[0]/2)): for iy in range(int(fft_array.shape[1]/2)): for iz in range(int(fft_array.shape[2]/2)): k = np.sqrt( (ix/cluster_length)**2 + (iy/cluster_length)**2 + (iz/cluster_length)**2 ) smoothing_array[ix,iy,iz] = fft_array[ix,iy,iz]*np.exp( (-(k**2)*smoothing_scale**2)/2) for iy in range(int(cluster_grid/2),cluster_grid-1): for iz in range(int(cluster_grid/2),cluster_grid-1):
__all__ = ['DadaReader', 'lookup_warn', '__version__', '__all__'] def lookup_warn(table, key, default=None): try: return table[key] except KeyError: if default is not None: print "#Warning: No key '%s'; using default value of %s" \ % (key, default) return default else: print "#Warning: No key '%s'" % key return None @jit(complex64(float32), target='cpu') def data2cplx(data): """ Convert data into complex form Helper function to allow for JIT compilation """ #shape: (200, 2, 600, 2176) data = data[:, 0, :, :] + np.complex(-1j) * data[:, 1, :, :] return data @jit(complex64(complex64), target='cpu') def reorderTranspose(mat, n_int, n_chans, n_station, n_pol): """ Do reorder and transpose operation """ # Reorder so that pol products change fastest mat = mat.reshape(n_int, n_chans, n_station, n_pol, n_station, n_pol) mat = mat.transpose([0, 2, 4, 1, 3, 5])
return k - (kx**2 + ky**2) / (2 * k) @numba.vectorize([numba.float64(numba.float64, numba.float64, numba.float64)]) def calc_kz_quadratic(k, kx, ky): return -(kx**2 + ky**2) / (2 * k) @numba.vectorize([numba.float64(numba.float64, numba.float64)]) def calc_kz_quadratic_1d(k, q): """kz not included.""" return -q**2 / (2 * k) @numba.vectorize([ numba.complex64(numba.float64, numba.float64, numba.float64, numba.float64) ]) def calc_propagator_exact(k, kx, ky, l): return mathx.expj(calc_kz_exact(k, kx, ky) * l) @numba.vectorize([ numba.complex64(numba.float64, numba.float64, numba.float64, numba.float64) ]) def calc_propagator_paraxial(k, kx, ky, l): return mathx.expj(calc_kz_paraxial(k, kx, ky) * l) @numba.vectorize([ numba.complex64(numba.float64, numba.float64, numba.float64, numba.float64) ])
# -*- coding: utf-8 -*- """ Created on Sat Jul 11 20:50:04 2015 @author: raghavsinghal """ import numpy as np import matplotlib.pyplot as plt from numba import jit, complex64, float64, int64 from math import pi @jit(complex64(float64, complex64[:, :], float64, float64, float64, float64)) def f(t, u, a, b, g, e): dx = np.zeros((18, 1), dtype='complex') ua = u[0, 0] ub = u[1, 0] ug = u[2, 0] va = u[3, 0] vb = u[4, 0] vg = u[5, 0] ha = u[6, 0] hb = u[7, 0] hg = u[8, 0] u_a = u[9, 0] u_b = u[10, 0] u_g = u[11, 0]
@numba.njit def norm_squared(x): return dot(x, x) @numba.njit def norm(x): return norm_squared(x)**0.5 @numba.njit def normalize(x): return np.asarray(x)/norm(x) @numba.vectorize([numba.complex64(numba.float64, numba.float64, numba.float64)], cache=True) def calc_ideal_lens_phase_paraxial(x, y, k_on_f): return cmath.exp(-0.5j*(x**2 + y**2)*k_on_f) @numba.vectorize([numba.complex64(numba.float64, numba.float64, numba.float64, numba.float64)], cache=True) def calc_ideal_lens_phase(x, y, k, f): return cmath.exp(-1j*k*((x**2 + y**2 + f**2)**0.5 - f)) @numba.vectorize([numba.complex64(numba.float64, numba.float64, numba.float64, numba.float64)], cache=True) def calc_ideal_square_lens_array_phase_paraxial(x, y, pitch, k_on_f): half_pitch = pitch/2 u = (x%pitch) - half_pitch v = (y%pitch) - half_pitch return cmath.exp(-0.5j*(u**2 + v**2)*k_on_f)
"""Variants of standard maths functions written for speed.""" import numpy as np import numba @numba.vectorize( [numba.complex64(numba.float32), numba.complex128(numba.float64)]) def expj2pi(x): """Equivalent to ``expj(2 * np.pi * x)`` where `x` is real. x is reduced to a small value before multiplication, which improves precision at a small cost in performance. """ y = 2 * np.pi * (x - np.rint(x)) return complex(np.cos(y), np.sin(y)) def nansum(x, *args, **kwargs): """Like np.nansum, provided `x` is a floating-point type.""" return np.sum(x, *args, where=~np.isnan(x), **kwargs)
num_args, mat_dim_list, num_reps, dtype=np.csingle) #%% ''' NUMBA TESTS ''' #%% Function declarations import cmath from numba import vectorize, complex128, complex64 @vectorize( [complex64(complex64, complex64), complex128(complex128, complex128)], target='parallel') def nb_add(x, y): return x + y @vectorize( [complex64(complex64, complex64), complex128(complex128, complex128)], target='parallel') def nb_subtract(x, y): return x + y @vectorize(