def phase_correlation(signal, ref): fft_s = fft(signal) fft_r = fft(ref) prod_of_ffts = (fft_s * np.conjugate(fft_r)) pc = prod_of_ffts / np.abs(prod_of_ffts) return abs(ifft(pc))
def cryo_conv_vol(x, kernel_f): n = x.shape[0] n_ker = kernel_f.shape[0] if np.any(np.array(x.shape) != n): raise ValueError('Volume in `x` must be cubic') if np.any(np.array(kernel_f.shape) != n_ker): raise ValueError('Convolution kernel in `kernel_f` must be cubic') is_singleton = len(x.shape) == 3 shifted_kernel_f = np.fft.ifftshift(np.fft.ifftshift(np.fft.ifftshift(kernel_f, 0), 1), 2) if is_singleton: x = numpy_fft.fftn(x, [n_ker] * 3) else: x = numpy_fft.fft(x, n=n_ker, axis=0) x = numpy_fft.fft(x, n=n_ker, axis=1) x = numpy_fft.fft(x, n=n_ker, axis=2) x *= shifted_kernel_f if is_singleton: x = numpy_fft.ifftn(x) x = x[:n, :n, :n] else: x = numpy_fft.ifft(x, axis=0) x = numpy_fft.ifft(x, axis=1) x = numpy_fft.ifft(x, axis=2) x = x.real return x
def time_to_frequency(ad, ed, td, wd, del_t, q, compensate_window=True): """ Convert time domain data to discrete Fourier transformed data, with normalization by del_t x q x n / k1 Parameters ---------- ad ed td wd del_t q Returns ------- """ # ================================================================================================================== # Now we get extract the data and transform it to frequency domain # ================================================================================================================== if compensate_window: resc = ad.shape[0] / np.sum(wd) else: resc = 1.0 a_df = fft(wd * ad) * del_t * q * resc e_df = fft(wd * ed) * del_t * q * resc t_df = fft(wd * td) * del_t * q * resc return a_df, e_df, t_df
def ft1D(x,y,*,axis=0,zero_DC=False): """Takes in x and y = y(x), and returns k and the Fourier transform of y(x) -> f(k) along a single (1D) axis Handles all of the annoyances of fftshift and ifftshift, and gets the normalization right Args: x (np.ndarray) : independent variable, must be 1D y (np.ndarray) : dependent variable, can be nD Kwargs: axis (int) : which axis to perform FFT zero_DC (bool) : if true, sets f(0) = 0 """ dx = x[1]-x[0] k = fftshift(fftfreq(x.size,d=dx))*2*np.pi fft_norm = dx shifted_x = ifftshift(x) if np.isclose(shifted_x[0],0): f = fft(ifftshift(y,axes=(axis)),axis=axis)*fft_norm else: f = fft(y,axis=axis)*fft_norm if zero_DC: nd_slice = [slice(None) for i in range(len(f.shape))] nd_slice[axis] = slice(0,1,1) nd_slice = tuple(nd_slice) f[nd_slice] = 0 f = fftshift(f,axes=(axis)) return k, f
def toeplitz_multiplication(v, first_row, first_column): """ Performs the matrix-vector product T * v where T is a Toeplitz matrix, using FFT Parameters ---------- v : array_like input vector of size n_data first_row : array_like first row of the Toepltiz matrix (size n_data) first_column : array_like first column of the Toepltiz matrix (size n_data) Returns ------- y : numpy array vector such that y = T * v """ n = first_row.shape[0] a_2n_fft = fft(np.concatenate((first_row,[0],first_column[1:][::-1]))) return np.real(ifft(a_2n_fft*fft(v, 2*n))[0:n])
def toeplitz_matvec(col0, vec, n): """ col0: first column of the Toeplitz matrix with shape (n,) vec: vector with shape (n,) """ p = (ifft( fft(np.r_[col0, col0[-2:0:-1]]) * fft(np.r_[vec, np.zeros(n - 2)])).real)[:n] return p
def toeplitz_matvec_block(col0, vec, n): """ col0: first column of the Toeplitz matrix with shape (n,) vec: vector with shape (m, n) """ m, n = vec.shape padded_vec = np.zeros((m, n * 2 - 2)) padded_vec[:, :n] = vec p = ifft(fft(np.r_[col0, col0[-2:0:-1]]) * fft(padded_vec)).real[:, :n] return p
def csr_convolution(a, b): P = len(a) Q = len(b) L = P + Q - 1 K = 2**nextpow2(L) a_pad = np.pad(a, (0, K - P), 'constant', constant_values=(0)) b_pad = np.pad(b, (0, K - Q), 'constant', constant_values=(0)) c = ifft(fft(a_pad) * fft(b_pad)) c = c[0:L - 1].real return c
def multiple_toepltiz_inverse(c_mat, lambda_n, a): """ Efficiently compute several Toepltiz systems with the same Toepltiz matrix T xj = cj which is in matrix form T x_mat = c_mat Where T is a n_data x n_data Toeplitz matrix and c_mat is a n_data x n_knots matrix """ n = c_mat.shape[0] #zero_vect = np.zeros(n_data) # PRECOMPUTATIONS # Cf. Step 2 of Ref. [1] #ae_2n = np.concatenate(([1],a,zero_vect)) #ae_2n_fft = fft(ae_2n) ae_2n_fft = fft(np.concatenate(([1],a)),2*n) # using hermitian and real property of covariance matrices: # be_2n_fft = fft(be_2n) be_2n_fft = ae_2n_fft.conj() #np.real(ae_2n_fft) - 1j*np.imag(ae_2n_fft) signs = (-1)**(np.arange(2*n)+1) x_mat = np.empty(np.shape(c_mat)) print("shape of c_mat is " + str(c_mat.shape)) for j in range(c_mat.shape[1]): #ce_2n = np.zeros(2*n_data) #ce_2n[0:n_data] = c_mat[:,j] #ce_2n = np.concatenate((c_mat[:,j],zero_vect)) #ce_2n_fft = fft(ce_2n) ce_2n_fft = fft(c_mat[:,j],2*n) u_2n = ifft( ae_2n_fft*ce_2n_fft ) v_2n = ifft( be_2n_fft*ce_2n_fft ) #pe_2n_fft = fft( np.concatenate((v_2n[0:n_data],zero_vect)) ) #qe_2n_fft = fft( np.concatenate((u_2n[n_data:],zero_vect)) ) pe_2n_fft = fft( v_2n[0:n] , 2*n ) qe_2n_fft = fft( u_2n[n:] , 2*n ) we_2n = ifft( ae_2n_fft*pe_2n_fft + signs*be_2n_fft*qe_2n_fft ) x_mat[:,j] = np.real(we_2n[0:n]/lambda_n) return x_mat
def update_missing_data(self, y_gw_list): """ Update the value of missing data Parameters ---------- y_gw_list : list list of waveform in the time domain for each channel Returns ------- Nothing, just update the attribute self.data_dft """ # Update Gaussian process mean self.model.update_mean(y_gw_list) # Impute missing data y_imp = self.model.impute(self.data) # Transform back to Fourier domain, applying the windowing for complete # time series, with re-scaling resc = self.del_t * self.resc_full data_dft = [ fft(y_imp0 * self.wd_full)[self.inds] * resc for y_imp0 in y_imp ] # self.data_dft = data_dft[:] return y_imp, data_dft
def _encoder(self, x, decimate, axis=-1, planner_effort='FFTW_ESTIMATE'): """ Transform the time series data into a multiple sub-banded demodulated signal in frequency domain. """ Nch, Nsamp = x.shape Nsamp_dec = int(Nsamp / decimate) X = fft.fft(x, axis=axis, planner_effort=planner_effort) if Nsamp % 2 == 0: X[:, 1:Nsamp // 2] *= 2 X[:, Nsamp // 2:] = 0 else: X[:, 1:(Nsamp + 1) // 2] *= 2 X[:, (Nsamp + 1) // 2] = 0 X_ = fft.fftshift(X, axes=-1)[:, int((Nsamp - Nsamp_dec) // 2):int( (Nsamp + Nsamp_dec) // 2)] / decimate if self.n_processes > 1: func = self.pfunc.result else: func = self.multiply return func(X_[:, self.encoder_rule], np.atleast_2d(self.Hwin))
def gls_covariance_time(mat, psd): """ Compute the covariance of the generalized least-square estimator for a stationary noise. Parameters ---------- mat : ndarray Model design matrix of size n_data x n_params psd : ndarray Noise power spectrum = PSD * fs / 2 where fs is the sampling frequency and PSD is the one-sided noise power spectral density. Returns ------- type Description of returned object. """ # if type(wind) == str: # wd = np.hanning(mat.shape[0]) # elif type(wind) == np.ndarray: # wd = wind[:] # Apply time-windowing # mat_wind = mat * np.array([wd]).T # Fourier transform the model # mat_dft = fft(mat_wind, axis=0) return gsl_covariance_freq(fft(mat, axis=0), psd)
def check_efield_resolution(self, efield, *, plot_fields=False): efield_tail = np.max(np.abs([efield[0], efield[-1]])) if efield_tail > np.max(np.abs(efield)) / 100: warnings.warn( 'Consider using larger time interval, pulse does not decay to less than 1% of maximum value in time domain' ) efield_fft = fftshift(fft(ifftshift(efield))) * self.dt efield_fft_tail = np.max(np.abs([efield_fft[0], efield_fft[-1]])) if efield_fft_tail > np.max(np.abs(efield_fft)) / 100: warnings.warn( '''Consider using smaller value of dt, pulse does not decay to less than 1% of maximum value in frequency domain''' ) if plot_fields: fig, axes = plt.subplots(1, 2) l1, l2, = axes[0].plot(self.efield_t, np.real(efield), self.efield_t, np.imag(efield)) plt.legend([l1, l2], ['Real', 'Imag']) axes[1].plot(self.efield_w, np.real(efield_fft), self.efield_w, np.imag(efield_fft)) axes[0].set_ylabel('Electric field Amp') axes[0].set_xlabel('Time ($\omega_0^{-1})$') axes[1].set_xlabel('Frequency ($\omega_0$)') fig.suptitle( 'Check that efield is well-resolved in time and frequency') plt.show()
def apply_gap_convolution(self, y_gw_fft_pos): """ Transform the frequency-domain waveform to account for gaps. Parameters ---------- y_gw_fft_pos : list[ndarray] list of frequency-domain waveforms for all channels. Returns ------- y_gw_masked_fft : list[ndarray] list of distorted frequency-domain waveforms for all channels. """ # Convert waveform to time domain y_gw_list = [ self.frequency_to_time(y_gw_fft_pos[i]) for i in range(len(y_gw_fft_pos)) ] # Apply mask window and Fourier transform back y_gw_masked_fft = [ fft(self.wd * dat)[self.inds] * self.del_t * self.resc for dat in y_gw_list ] return y_gw_masked_fft
def track(self, im, pos, base_target_sz, current_scale_factor): """ track the scale using the scale filter """ # get scale filter features scales = current_scale_factor * self.scale_size_factors xs = self._extract_scale_sample(im, pos, base_target_sz, scales, self.scale_model_sz) # project xs = self.basis.dot(xs) * self.window # get scores xsf = fft(xs, axis=1) scale_responsef = np.sum(self.sf_num * xsf, 0) / (self.sf_den + config.lamBda) interp_scale_response = np.real(ifft(resize_dft(scale_responsef, config.number_of_interp_scales))) recovered_scale_index = np.argmax(interp_scale_response) if config.do_poly_interp: # fit a quadratic polynomial to get a refined scale estimate id1 = (recovered_scale_index - 1) % config.number_of_interp_scales id2 = (recovered_scale_index + 1) % config.number_of_interp_scales poly_x = np.array([self.interp_scale_factors[id1], self.interp_scale_factors[recovered_scale_index], self.interp_scale_factors[id2]]) poly_y = np.array([interp_scale_response[id1], interp_scale_response[recovered_scale_index], interp_scale_response[id2]]) poly_A = np.array([[poly_x[0]**2, poly_x[0], 1], [poly_x[1]**2, poly_x[1], 1], [poly_x[2]**2, poly_x[2], 1]], dtype=np.float32) poly = np.linalg.inv(poly_A).dot(poly_y.T) scale_change_factor = - poly[1] / (2 * poly[0]) else: scale_change_factor = self.interp_scale_factors[recovered_scale_index] return scale_change_factor
def toepltiz_mat_vect_prod(y, s_2n): """ Linear operator that calculate T y_in assuming that we can write: Com = F* Lambda F where Lambda is a P x P diagonal matrix and F is the P x n_data Discrete Fourier Transform matrix. Parameters ---------- y : numpy array input data vector of size n_data S_2N : numpy array (size P >= 2N) PSD vector Returns ------- y_out : numpy array y_out = T * y_in transformed output vector of size N_out """ return np.real(ifft(s_2n * fft(y, len(s_2n)))[0:len(y)])
def create_filter(order, cutoff, nyquist, N, ftype='fir', output='freq', shift=True): """ Create a prototype filter. """ h = firwin(order, cutoff, nyq=nyquist) if output == 'freq': w = fft.fftfreq(N) w *= (nyquist * 2) H = fft.fft(h, n=N, axis=-1, planner_effort='FFTW_ESTIMATE') if shift: return fft.fftshift(w), fft.fftshift(H) else: return w, H else: return h
def delay(x, time, fs, axis=-1, keeplength=False, pad=1): extra_pad = 200 # add 200 samples to prevent wrapping samps = int(np.floor(time * fs)) s = list(x.shape) sz_pre = np.copy(s) sz_post = np.copy(s) sz_fft = np.copy(s) sz_pre[axis] = samps sz_post[axis] = pad + extra_pad x = np.concatenate((np.zeros(sz_pre), x, np.zeros(sz_post)), axis) sz_fft[axis] = int(np.round(2 ** np.ceil(np.log2(x.shape[axis])) - x.shape[axis])) x = np.concatenate((x, np.zeros(sz_fft)), axis) new_len = sz_pre[axis] + s[axis] + sz_post[axis] + sz_fft[axis] # x[n-k] <--> X(jw)e^(-jwk) where w in [0, 2pi) if type(time) is not int: theta = (-np.arange(new_len).astype(float) * fs * 2 * np.pi / new_len * (time - np.float(samps) / fs)) theta[-(new_len // 2) + 1:] = -theta[(new_len // 2):1:-1] st = [1 for _ in range(x.ndim)] st[axis] = new_len x = np.real(fft.ifft(fft.fft(x, axis=axis) * np.exp(1j * theta.reshape(st)))) if keeplength: x = np.take(x, range(s[axis]), axis) else: x = np.take(x, range(s[axis] + samps + pad), axis) inds = tuple([slice(si) for si in sz_pre]) x[inds] = 0 return x
def detrend(y_data, detrend_method='poly', max_order=1, n_knots=10, psd=None): """ Remove linear or polynomial trend of order max_order Parameter ---------- y_data : 1d array input data of size n max_order : int maximum order of the polynomial to fit psd : 1d array noise PSD at Fourier frequencies (size n) Returns ------- y_detrend : 1d array output detrended data (size n) """ t_norm = np.arange(0, y_data.shape[0]) if detrend_method == 'poly': mat_linear = np.hstack( [np.array([t_norm**k]).T for k in range(0, max_order + 1)]) if psd is not None: amp = regression.generalized_least_squares(fft(mat_linear, axis=0), fft(y_data), psd) else: amp = regression.least_squares(mat_linear, y_data) trend = np.real(np.dot(mat_linear, amp)) y_detrend = y_data - trend elif detrend_method == 'spline': # Detrending using splines n_seg = y_data.shape[0] // n_knots t_knots = np.linspace(t_norm[n_seg], t_norm[-n_seg], n_knots) spl = interpolate.LSQUnivariateSpline(t_norm, y_data, t_knots, k=3, ext="const") trend = spl(t_norm) y_detrend = y_data - trend return y_detrend, trend
def plot2d_fft(self, *, delay_time_start=1, create_figure=True, color_range='auto', subtract_DC=True, draw_colorbar=True, frequency_range=[-1000, 1000], normalize=False, phase=False, save_fig=True, wT_frequency_range='auto'): w_ind = np.where((self.w > frequency_range[0]) & (self.w < frequency_range[1]))[0] w = self.w[w_ind] sig = self.signal_vs_delay_times[w_ind, :] delay_time_indices = np.where(self.delay_times > delay_time_start)[0] delay_times = self.delay_times[delay_time_indices] sig = sig[:, delay_time_indices] if normalize: sig /= np.dot(self.dipoles, self.dipoles)**2 wT = fftshift( fftfreq(delay_times.size, d=(delay_times[1] - delay_times[0]))) * 2 * np.pi sig_fft = fft(sig, axis=1) if subtract_DC: sig_fft[:, 0] = 0 sig_fft = fftshift(sig_fft, axes=(1)) ww, wTwT = np.meshgrid(wT, w) if create_figure: plt.figure() if phase: plt.title('Phase') plot_sig = np.arctan2(np.imag(sig_fft), np.real(sig_fft)) else: plt.title('Magnitude') plot_sig = np.abs(sig_fft) if color_range == 'auto': plt.pcolormesh(ww, wTwT, plot_sig) else: plt.pcolormesh(ww, wTwT, plot_sig, vmin=color_range[0], vmax=color_range[1]) if draw_colorbar: plt.colorbar() plt.xlabel('$\omega_T$ ($\omega_0$)', fontsize=16) plt.ylabel('Detection Frequency ($\omega_0$)', fontsize=16) if wT_frequency_range == 'auto': plt.xlim([0, np.max(wT)]) else: plt.xlim(wT_frequency_range) if save_fig: plt.savefig(self.base_path + 'TA_spectra_fft')
def plotTA_units(self, *, frequency_range=[-1000, 1000], subtract_DC=True, create_figure=True, color_range='auto', draw_colorbar=True, save_fig=True, omega_0=1): """Plots the transient absorption spectra with detection frequency on the y-axis and delay time on the x-axis. Args: frequency_range (list): sets the min (list[0]) and max (list[1]) detection frequency for y-axis subtract_DC (bool): if True subtract the DC component of the TA color_range (list): sets the min (list[0]) and max (list[1]) value for the colorbar draw_colorbar (bool): if True add a colorbar to the plot save_fig (bool): if True save the figure that is produced omega_0 (float): convert from unitless variables, omega_0 should be provided in wavenumbers """ self.load_eigen_params() f0_thz = omega_0 * 3E10 / 1.0E12 # omega_0 in wavenumbers T_ps = self.delay_times / f0_thz / (2 * np.pi) # Cut out unwanted detection frequency points self.w += self.ground_to_excited_transition + self.center self.w *= omega_0 w_ind = np.where((self.w > frequency_range[0]) & (self.w < frequency_range[1]))[0] w = self.w[w_ind] sig = self.signal[w_ind, :] if omega_0 == 1: xlab = r'Delay time ($\omega_0^{-1}$)' ylab = r'Detection Frequency ($\omega_0$)' else: xlab = 'Delay time (ps)' ylab = r'Detection Frequency (cm$^{-1}$)' if subtract_DC: sig_fft = fft(sig, axis=1) sig_fft[:, 0] = 0 sig = np.real(ifft(sig_fft)) ww, tt = np.meshgrid(T_ps, w) if create_figure: plt.figure() if color_range == 'auto': plt.pcolormesh(ww, tt, sig) else: plt.pcolormesh(ww, tt, sig, vmin=color_range[0], vmax=color_range[1]) if draw_colorbar: plt.colorbar() plt.xlabel(xlab, fontsize=16) plt.ylabel(ylab, fontsize=16) if save_fig: plt.savefig(self.base_path + 'TA_spectra_iso_ave')
def teopltiz_precompute(R,p=10,Nit = 1000,tol = 1e-4): """ Solve the system T y = e1 where T is symmetric Toepltiz. where e1 = [1 0 0 0 0 0].T to compute the vector a and lambda_n for further fast Toepltiz inversions. Compute the prefactor lambda_n and vector a_{n-1} such that the inverse of T writes T^{-1} = (1/lambda_n) * [ 1 a*_{n-1} a_{n-1} S_{n-1} ] Parameters ---------- R : array_like autocovariance function (first row of the Toepltiz matrix) Returns ------- lambda_n : scalar float prefactor of the inverse of T a : numpy array vector of size N-1 involved in the computation of the inverse of T p : scalar integer maximum number of lags for the preconditionning Nit : scalar integer maximum number of iterations for PCG tol : scalar float relative error convergence criterium for the PCG algorithm References ---------- [1] Jain, Fast Inversion of Banded Toeplitz Matrices by Circular Decompositions, 1978 """ N = len(R) # First basis vector (of orthonormal cartesian basis) e1 = np.concatenate(([1],np.zeros(N-1))) # Compute spectrum S_2N = fft(np.concatenate((R,[0],R[1:][::-1]))) # Linear operator correponding to the Toeplitz matrix T_op = toepltizLinearOp(N,S_2N) # Preconditionner to approximate T^{-1} Psolver = computeToepltizPrecond(R,p=p) # Build the associated linear operator P_op = matrixalgebra.precondLinearOp(Psolver,N,N) # Initial guess z,info = sparse.linalg.bicgstab(T_op, e1, x0=np.zeros(N),tol=tol, maxiter=Nit,M=P_op,callback=None) matrixalgebra.printPCGstatus(info) lambda_n = 1/z[0] a = lambda_n * z[1:] return lambda_n,a
def phaseran(signal): """ Performs phase randomization for coherence matrices. NOTE: Signal input has to be nTimepoints x nElectrodes (there is a check) """ # check that it is in right orientation if signal.shape[1] > signal.shape[0]: signal = signal.T # Get parameters nTimepoints = signal.shape[0] nElectrodes = signal.shape[1] # Check to make sure that it is an odd number of samples if nTimepoints % 2 == 0: signal = signal[:-1, :] nTimepoints = nTimepoints - 1 nTimepoints = signal.shape[0] len_ser = (nTimepoints - 1) / 2 interv1 = np.arange(1, len_ser + 1) interv2 = np.arange(len_ser + 1, nTimepoints) # fft_A = pfft.builders.fft(signal, axis=0) # FFT of original data try: fft_A = pfft.fft(signal, axis=0, threads=15) except: fft_A = pfft.fft(signal, axis=0) # Create the random phases for all the time series ph_rnd = np.random.rand(len_ser, nElectrodes) ph_interv1 = np.exp(2 * np.pi * 1j * ph_rnd) ph_interv2 = np.conj(np.flipud(ph_interv1)) # Randomize all time series simultaneously fft_recblk_surr = fft_A fft_recblk_surr[interv1, :] = fft_A[interv1, :] * ph_interv1 fft_recblk_surr[interv2, :] = fft_A[interv2, :] * ph_interv2 surrblk = np.float32(pfft.ifft(fft_recblk_surr, axis=0)).T return surrblk
def fourier_transform(array,axis,new_axis,inverse=False): import numpy as np from numpy.fft import fftshift,ifftshift from expresso.pycas import pi from ..coordinate_ndarray import CoordinateNDArray try: from pyfftw.interfaces.numpy_fft import fft,ifft except ImportError: from numpy.fft import fft,ifft axi = array.axis.index(axis) if len(array.axis) == 1: if not inverse: new_data = fftshift(fft(array.data,axis=axi),axes=[axi]) new_data *= 1/np.sqrt(2*np.pi) else: new_data = ifft(ifftshift(array.data,axes=[axi]),axis=axi) new_data *= np.sqrt(2*np.pi) else: from pypropagate.progressbar import ProgressBar axt = (axi + 1) % len(array.axis) axi = axi if axt > axi else axi - 1 transposed_data = np.rollaxis(array.data,axt,start=0) new_data = np.zeros(array.data.shape,dtype=complex) transposed_new_data = np.rollaxis(new_data,axt,start=0) for i in ProgressBar(range(transposed_data.shape[0])): if not inverse: transposed_new_data[i] = fftshift(fft(transposed_data[i],axis=axi),axes=[axi]) transposed_new_data[i] *= 1/np.sqrt(2*np.pi) else: transposed_new_data[i] = ifft(ifftshift(transposed_data[i],axes=[axi]),axis=axi) transposed_new_data[i] *= np.sqrt(2*np.pi) sw = array.bounds[axi][1] - array.bounds[axi][0] tmin,tmax = array.evaluate((-(pi*array.shape[axi])/sw, (pi*array.shape[axi])/sw)) new_bounds = [(b[0],b[1]) if i!=axi else (tmin,tmax) for i,b in enumerate(array.bounds)] new_axis = [a if i!=axi else new_axis for i,a in enumerate(array.axis)] return CoordinateNDArray(new_data,new_bounds,new_axis,array.evaluate)
def dft(self, wind='tukey', n_wind=500, normalized=True): self.w = self.compute_window(wind=wind, n_wind=n_wind) if normalized: norm = np.sum(self.w) / (self.del_t * 2) else: norm = 1.0 return fft(self * self.w) / norm
def integrated_ft(self,delay_time_start = 1,delay_time_stop = 300): delay_time_indices = np.where((self.delay_times > delay_time_start) & (self.delay_times < delay_time_stop))[0] delay_times = self.delay_times[delay_time_indices] sig = self.signal_vs_delay_times[:,delay_time_indices] integrated = np.trapz(sig,x=self.TA.w,axis=0) w_T = fftshift(fftfreq(delay_times.size,d=(delay_times[1] - delay_times[0])))*2*np.pi integrated_fft = fft(integrated) integrated_fft[0] = 0 integrated_fft = fftshift(integrated_fft) return w_T, integrated_ft
def update(self, im, pos, base_target_sz, current_scale_factor): """ update the scale filter """ # get scale filter features scales = current_scale_factor * self.scale_size_factors xs = self._extract_scale_sample(im, pos, base_target_sz, scales, self.scale_model_sz) first_frame = not hasattr(self, 's_num') if first_frame: self.s_num = xs else: self.s_num = (1 - config.scale_learning_rate ) * self.s_num + config.scale_learning_rate * xs # compute projection basis if self.max_scale_dim: self.basis, _ = scipy.linalg.qr(self.s_num, mode='economic') scale_basis_den, _ = scipy.linalg.qr(xs, mode='economic') else: U, _, _ = np.linalg.svd(self.s_num) self.basis = U[:, :self.s_num_compressed_dim] self.basis = self.basis.T # compute numerator feat_proj = self.basis.dot(self.s_num) * self.window sf_proj = fft(feat_proj, axis=1) self.sf_num = self.yf * np.conj(sf_proj) # update denominator xs = scale_basis_den.T.dot(xs) * self.window xsf = fft(xs, axis=1) new_sf_den = np.sum(np.real(xsf * np.conj(xsf)), 0) if first_frame: self.sf_den = new_sf_den else: self.sf_den = ( 1 - config.scale_learning_rate ) * self.sf_den + config.scale_learning_rate * new_sf_den
def _psd2invntt(psd, bandwidth, ncorr, fftw_flag='FFTW_MEASURE'): """ Compute the first row of the inverse of the noise time-time correlation matrix from two-sided PSD with frequency increment fsamp/psd.size. """ fftsize = psd.shape[-1] fsamp = bandwidth * fftsize psd = (1 / fsamp) / psd nthreads = multiprocessing.cpu_count() out = fft.fft(psd, axis=-1, planner_effort=fftw_flag, threads=nthreads) invntt = out[..., :ncorr+1].real / fftsize return invntt
def compute_periodogram(self, x): """ Compute the windowed periodogram from time series x, along with Fourier frequency grid """ # If size of analysed data is the same as the presets if (x.shape[0] == self.N): if self.wind == 'hanning': # Compute periodogram per = np.abs(fft(x * self.w))**2 / self.s2 elif self.wind == 'ones': per = np.abs(fft(x))**2 / self.s2 else: if self.wind == 'hanning': w = np.hanning(x.shape[0]) s2 = np.sum(w**2) per = np.abs(fft(x * w))**2 / s2 elif self.wind == 'ones': s2 = x.shape[0] per = np.abs(fft(x))**2 / s2 return per
def _psd2invntt(psd, bandwidth, ncorr, fftw_flag='FFTW_MEASURE'): """ Compute the first row of the inverse of the noise time-time correlation matrix from two-sided PSD with frequency increment fsamp/psd.size. """ fftsize = psd.shape[-1] fsamp = bandwidth * fftsize psd = (1 / fsamp) / psd nthreads = multiprocessing.cpu_count() out = fft.fft(psd, axis=-1, planner_effort=fftw_flag, threads=nthreads) invntt = out[..., :ncorr + 1].real / fftsize return invntt
def add_gaussian_linewidth(self, sigma): self.old_signal = self.signal.copy() sig_tau_t = fftshift(fft(ifftshift(self.old_signal, axes=(-1)), axis=-1), axes=(-1)) sig_tau_t = sig_tau_t * ( np.exp(-self.t**2 / (2 * sigma**2))[np.newaxis, np.newaxis, :] * np.exp(-self.t21_array**2 / (2 * sigma**2))[:, np.newaxis, np.newaxis]) sig_tau_w = fftshift(ifft(ifftshift(sig_tau_t, axes=(-1)), axis=-1), axes=(-1)) self.signal = sig_tau_w
def __init__( self, target_sz, ): init_target_sz = target_sz num_scales = config.number_of_scales_filter scale_step = config.scale_step_filter scale_sigma = config.number_of_interp_scales * config.scale_sigma_factor scale_exp = np.arange( -np.floor(num_scales - 1) / 2, np.ceil(num_scales - 1) / 2 + 1, dtype=np.float32) * config.number_of_interp_scales / num_scales scale_exp_shift = np.roll(scale_exp, (0, -int(np.floor((num_scales - 1) / 2)))) interp_scale_exp = np.arange( -np.floor((config.number_of_interp_scales - 1) / 2), np.ceil((config.number_of_interp_scales - 1) / 2) + 1, dtype=np.float32) interp_scale_exp_shift = np.roll( interp_scale_exp, [0, -int(np.floor(config.number_of_interp_scales - 1) / 2)]) self.scale_size_factors = scale_step**scale_exp self.interp_scale_factors = scale_step**interp_scale_exp_shift ys = np.exp(-0.5 * (scale_exp_shift**2) / (scale_sigma**2)) self.yf = np.real(fft(ys))[np.newaxis, :] self.window = signal.hann(ys.shape[0])[np.newaxis, :].astype( np.float32) # make sure the scale model is not to large, to save computation time if config.scale_model_factor**2 * np.prod( init_target_sz) > config.scale_model_max_area: scale_model_factor = np.sqrt(config.scale_model_max_area / np.prod(init_target_sz)) else: scale_model_factor = config.scale_model_factor # set the scale model size self.scale_model_sz = np.maximum( np.floor(init_target_sz * scale_model_factor), np.array([8, 8])) self.max_scale_dim = config.s_num_compressed_dim == 'MAX' if self.max_scale_dim: self.s_num_compressed_dim = len(self.scale_size_factors) self.num_scales = num_scales self.scale_step = scale_step self.scale_factors = np.array([1])
def _sampling2psd(sampling, sampling_frequency, fftw_flag='FFTW_MEASURE'): """ Return folded PSD without binning. """ sampling = np.asarray(sampling) n = sampling.size nthreads = multiprocessing.cpu_count() psd = np.abs(fft.fft(sampling, axis=-1, planner_effort=fftw_flag, threads=nthreads))**2 freq = np.fft.fftfreq(n, d=1/sampling_frequency)[:n//2+1] freq[-1] += sampling_frequency psd = np.concatenate([[0.], 2 * psd[1:n//2], [psd[n//2]]]) / \ (n * sampling_frequency) return freq, psd
def _gaussian_sample(nsamples, sampling_frequency, psd, twosided=False, out=None, fftw_flag='FFTW_MEASURE'): """ Generate a gaussian N-sample sampled at fs from a one- or two-sided Power Spectrum Density sampled at fs/N. Parameter --------- nsamples : int Number of time samples. sampling_frequency : float Sampling frequency [Hz]. psd : array-like One- or two-sided Power Spectrum Density [signal unit**2/Hz]. twosided : boolean, optional Whether or not the input psd is one-sided (only positive frequencies) or two-sided (positive and negative frequencies). out : ndarray Placeholder for the output buffer. """ psd = np.asarray(psd) if out is None: out = empty(psd.shape[:-1] + (nsamples,)) if not twosided: psd = _unfold_psd(psd) shape = psd.shape[:-1] + (nsamples,) gauss = np.random.randn(*shape) nthreads = multiprocessing.cpu_count() ftgauss = fft.fft(gauss, planner_effort=fftw_flag, threads=nthreads) ftgauss[..., 0] = 0 spec = ftgauss * np.sqrt(psd) out[...] = fft.ifft(spec, planner_effort=fftw_flag, threads=nthreads).real * np.sqrt(sampling_frequency) return out
# For pre-channelized data, data are always complex, # and should have shape (ntint, nchan, npol). # For baseband data, we wish to get to the same shape for # incoherent or by_channel, or just to fully channelized for coherent. ======= if fh.telescope in ('arochime','arochime-raw'): # take complex conjugate to ensure by-channel de-dispersion is applied correctly vals = np.conj(vals) >>>>>>> e6d25c1bfcc5a3426bc62becc3b3075fbea23ebc if fh.nchan == 1: # If we need coherent dedispersion, do FT of whole thing, # otherwise to output channels, mimicking pre-channelized data. if raw.dtype.kind == 'c': # complex data nsamp = len(vals) if dedisperse == 'coherent' else nchan vals = fft(vals.reshape(-1, nsamp, npol), axis=1, **_fftargs) else: # real data nsamp = len(vals) if dedisperse == 'coherent' else nchan * 2 vals = rfft(vals.reshape(-1, nsamp, npol), axis=1, **_rfftargs) # Sadly, the way data are stored depends on what FFT routine # one is using. We cannot deal with scipy's. if vals.dtype.kind == 'f': raise TypeError("Can no longer deal with scipy's format " "for storing FTs of real data.") if fedge_at_top: # take complex conjugate to ensure by-channel de-dispersion is # applied correctly. # This needs to be done for ARO data, since we are in 2nd Nyquist # zone; not clear it is needed for other telescopes.
def muenchetal(im, args): """Process a sinogram image with the Munch et al. de-striping algorithm. Parameters ---------- im : array_like Image data as numpy array. wlevel : int Levels of the wavelet decomposition. sigma : float Smoothing effect. (Parameters wlevel and sigma have to passed as a string separated by ;) Example (using tiffile.py) -------------------------- >>> im = imread('sino_orig.tif') >>> im = munchetal(im, '4;1.0') >>> imsave('sino_flt.tif', im) References ---------- B. Munch, P. Trtik, F. Marone, M. Stampanoni, Stripe and ring artifact removal with combined wavelet-Fourier filtering, Optics Express 17(10):8567-8591, 2009. """ # Disable a warning: simplefilter("ignore", ComplexWarning) # Get args: wlevel, sigma = args.split(";") wlevel = int(wlevel) sigma = float(sigma) # The wavelet transform to use : {'haar', 'db1'-'db20', 'sym2'-'sym20', 'coif1'-'coif5', 'dmey'} wname = "db5" # Wavelet decomposition: coeffs = wavedec2(im.astype(float32), wname, level=wlevel) coeffsFlt = [coeffs[0]] # FFT transform of horizontal frequency bands: for i in range(1, wlevel + 1): # Padding and windowing of input signal: n_byte_align(coeffs[i][1], simd_alignment) siz = coeffs[i][1].shape tmp = pad(coeffs[i][1], pad_width=((coeffs[i][1].shape[0] / 2, coeffs[i][1].shape[0] / 2), (0,0)), mode='constant') # or 'constant' for zero padding tmp = pad(tmp, pad_width=((0,0) ,(coeffs[i][1].shape[1] / 2, coeffs[i][1].shape[1] / 2)), mode='constant') # or 'constant' for zero padding tmp = _windowing_lr(tmp, siz[1]) tmp = _windowing_lr(tmp.T, siz[0]).T # FFT: fcV = fftshift(fft(tmp, axis=0, threads=2)) my, mx = fcV.shape # Damping of vertical stripes: damp = 1 - npexp(-(arange(-floor(my / 2.),-floor(my / 2.) + my) ** 2) / (2 * (sigma ** 2))) dampprime = kron(ones((1,mx)), damp.reshape((damp.shape[0],1))) fcV = fcV * dampprime # Inverse FFT: fcV = ifftshift(fcV) n_byte_align(fcV, simd_alignment) fcVflt = ifft(fcV, axis=0, threads=2) ## Crop image: tmp = fcVflt[fcVflt.shape[0] / 4:(fcVflt.shape[0] / 4 + siz[0]), fcVflt.shape[1] / 4:(fcVflt.shape[1] / 4 + siz[1])] # Dump back coefficients: cVHDtup = (coeffs[i][0], tmp, coeffs[i][2]) coeffsFlt.append(cVHDtup) # Get wavelet reconstruction: im_f = real(waverec2(coeffsFlt, wname)) # Return filtered image (an additional row and/or column might be present): return im_f[0:im.shape[0],0:im.shape[1]].astype(float32)
def fold(fh, comm, samplerate, fedge, fedge_at_top, nchan, nt, ntint, ngate, ntbin, ntw, dm, fref, phasepol, dedisperse='incoherent', do_waterfall=True, do_foldspec=True, verbose=True, progress_interval=100, rfi_filter_raw=None, rfi_filter_power=None, return_fits=False): """ FFT data, fold by phase/time and make a waterfall series Folding is done from the position the file is currently in Parameters ---------- fh : file handle handle to file holding voltage timeseries comm: MPI communicator or None will use size, rank attributes samplerate : Quantity rate at which samples were originally taken and thus double the band width (frequency units) fedge : float edge of the frequency band (frequency units) fedge_at_top: bool whether edge is at top (True) or bottom (False) nchan : int number of frequency channels for FFT nt, ntint : int total number nt of sets, each containing ntint samples in each file hence, total # of samples is nt*ntint, with each sample containing a single polarisation ngate, ntbin : int number of phase and time bins to use for folded spectrum ntbin should be an integer fraction of nt ntw : int number of time samples to combine for waterfall (does not have to be integer fraction of nt) dm : float dispersion measure of pulsar, used to correct for ism delay (column number density) fref: float reference frequency for dispersion measure phasepol : callable function that returns the pulsar phase for time in seconds relative to start of the file that is read. dedisperse : None or string (default: incoherent). None, 'incoherent', 'coherent', 'by-channel'. Note: None really does nothing do_waterfall, do_foldspec : bool whether to construct waterfall, folded spectrum (default: True) verbose : bool or int whether to give some progress information (default: True) progress_interval : int Ping every progress_interval sets return_fits : bool (default: False) return a subint fits table for rank == 0 (None otherwise) """ assert dedisperse in (None, 'incoherent', 'by-channel', 'coherent') need_fine_channels = dedisperse in ['by-channel', 'coherent'] assert nchan % fh.nchan == 0 if dedisperse in ['incoherent', 'by-channel'] and fh.nchan > 1: oversample = nchan // fh.nchan assert ntint % oversample == 0 else: oversample = 1 if dedisperse == 'coherent' and fh.nchan > 1: raise ValueError("Cannot coherently dedisperse channelized data.") if comm is None: mpi_rank = 0 mpi_size = 1 else: mpi_rank = comm.rank mpi_size = comm.size npol = getattr(fh, 'npol', 1) assert npol == 1 or npol == 2 if verbose > 1 and mpi_rank == 0: print("Number of polarisations={}".format(npol)) # initialize folded spectrum and waterfall # TODO: use estimated number of points to set dtype if do_foldspec: foldspec = np.zeros((ntbin, nchan, ngate, npol**2), dtype=np.float32) icount = np.zeros((ntbin, nchan, ngate), dtype=np.int32) else: foldspec = None icount = None if do_waterfall: nwsize = nt*ntint//ntw//oversample waterfall = np.zeros((nwsize, nchan, npol**2), dtype=np.float64) else: waterfall = None if verbose and mpi_rank == 0: print('Reading from {}'.format(fh)) nskip = fh.tell()/fh.blocksize if nskip > 0: if verbose and mpi_rank == 0: print('Starting {0} blocks = {1} bytes out from start.' .format(nskip, nskip*fh.blocksize)) dt1 = (1./samplerate).to(u.s) # need 2*nchan real-valued samples for each FFT if fh.telescope == 'lofar': dtsample = fh.dtsample else: dtsample = nchan // oversample * 2 * dt1 tstart = dtsample * ntint * nskip # pre-calculate time delay due to dispersion in coarse channels # for channelized data, frequencies are known tb = -1. if fedge_at_top else +1. if fh.nchan == 1: if getattr(fh, 'data_is_complex', False): # for complex data, really each complex sample consists of # 2 real ones, so multiply dt1 by 2. freq = fedge + tb * fftfreq(nchan, 2.*dt1) if dedisperse == 'coherent': fcoh = fedge + tb * fftfreq(nchan*ntint, 2.*dt1) fcoh.shape = (-1, 1) elif dedisperse == 'by-channel': fcoh = freq + tb * fftfreq(ntint, dtsample)[:, np.newaxis] else: # real data freq = fedge + tb * rfftfreq(nchan*2, dt1) if dedisperse == 'coherent': fcoh = fedge + tb * rfftfreq(ntint*nchan*2, dt1) fcoh.shape = (-1, 1) elif dedisperse == 'by-channel': fcoh = freq + tb * fftfreq(ntint, dtsample)[:, np.newaxis] freq_in = freq else: # Input frequencies may not be the ones going out. freq_in = fh.frequencies if oversample == 1: freq = freq_in else: freq = freq_in[:, np.newaxis] + tb * fftfreq(oversample, dtsample) fcoh = freq_in + tb * fftfreq(ntint, dtsample)[:, np.newaxis] # print('fedge_at_top={0}, tb={1}'.format(fedge_at_top, tb)) # By taking only up to nchan, we remove the top channel at the Nyquist # frequency for real, unchannelized data. ifreq = freq[:nchan].ravel().argsort() # pre-calculate time offsets in (input) channelized streams dt = dispersion_delay_constant * dm * (1./freq_in**2 - 1./fref**2) if need_fine_channels: # pre-calculate required turns due to dispersion. # # set frequency relative to which dispersion is coherently corrected if dedisperse == 'coherent': _fref = fref else: _fref = freq_in[np.newaxis, :] # (check via eq. 5.21 and following in # Lorimer & Kramer, Handbook of Pulsar Astronomy dang = (dispersion_delay_constant * dm * fcoh * (1./_fref-1./fcoh)**2) * u.cycle with u.set_enabled_equivalencies(u.dimensionless_angles()): dd_coh = np.exp(dang * 1j).conj().astype(np.complex64) # add dimension for polarisation dd_coh = dd_coh[..., np.newaxis] # Calculate the part of the whole file this node should handle. size_per_node = (nt-1)//mpi_size + 1 start_block = mpi_rank*size_per_node end_block = min((mpi_rank+1)*size_per_node, nt) for j in range(start_block, end_block): if verbose and j % progress_interval == 0: print('#{:4d}/{:4d} is doing {:6d}/{:6d} [={:6d}/{:6d}]; ' 'time={:18.12f}' .format(mpi_rank, mpi_size, j+1, nt, j-start_block+1, end_block-start_block, (tstart+dtsample*j*ntint).value)) # time since start # Just in case numbers were set wrong -- break if file ends; # better keep at least the work done. try: raw = fh.seek_record_read(int((nskip+j)*fh.blocksize), fh.blocksize) except(EOFError, IOError) as exc: print("Hit {0!r}; writing data collected.".format(exc)) break if verbose >= 2: print("#{:4d}/{:4d} read {} items" .format(mpi_rank, mpi_size, raw.size), end="") if npol == 2 and raw.dtype.fields is not None: raw = raw.view(raw.dtype.fields.values()[0][0]) if fh.nchan == 1: # raw.shape=(ntint*npol) raw = raw.reshape(-1, npol) else: # raw.shape=(ntint, nchan*npol) raw = raw.reshape(-1, fh.nchan, npol) if dedisperse == 'incoherent' and oversample > 1: raw = ifft(raw, axis=1, **_fftargs).reshape(-1, nchan, npol) raw = fft(raw, axis=1, **_fftargs) if rfi_filter_raw is not None: raw, ok = rfi_filter_raw(raw) if verbose >= 2: print("... raw RFI (zap {0}/{1})" .format(np.count_nonzero(~ok), ok.size), end="") if np.can_cast(raw.dtype, np.float32): vals = raw.astype(np.float32) else: assert raw.dtype.kind == 'c' vals = raw # For pre-channelized data, data are always complex, # and should have shape (ntint, nchan, npol). # For baseband data, we wish to get to the same shape for # incoherent or by_channel, or just to fully channelized for coherent. if fh.nchan == 1: # If we need coherent dedispersion, do FT of whole thing, # otherwise to output channels, mimicking pre-channelized data. if raw.dtype.kind == 'c': # complex data nsamp = len(vals) if dedisperse == 'coherent' else nchan vals = fft(vals.reshape(-1, nsamp, npol), axis=1, **_fftargs) else: # real data nsamp = len(vals) if dedisperse == 'coherent' else nchan * 2 vals = rfft(vals.reshape(-1, nsamp, npol), axis=1, **_rfftargs) # Sadly, the way data are stored depends on what FFT routine # one is using. We cannot deal with scipy's. if vals.dtype.kind == 'f': raise TypeError("Can no longer deal with scipy's format " "for storing FTs of real data.") if fedge_at_top: # take complex conjugate to ensure by-channel de-dispersion is # applied correctly. # This needs to be done for ARO data, since we are in 2nd Nyquist # zone; not clear it is needed for other telescopes. np.conj(vals, out=vals) # Now we coherently dedisperse, either all of it or by channel. if need_fine_channels: # for by_channel, we have vals.shape=(ntint, nchan, npol), # and want to FT over ntint to get fine channels; if vals.shape[0] > 1: fine = fft(vals, axis=0, **_fftargs) else: # for coherent, we just reshape: # (1, ntint*nchan, npol) -> (ntint*nchan, 1, npol) fine = vals.reshape(-1, 1, npol) # Dedisperse. fine *= dd_coh # Still have fine.shape=(ntint, nchan, npol), # w/ nchan=1 for coherent. if fine.shape[1] > 1 or raw.dtype.kind == 'c': vals = ifft(fine, axis=0, **_fftargs) else: vals = irfft(fine, axis=0, **_rfftargs) if fine.shape[1] == 1 and nchan > 1: # final FT to get requested channels if vals.dtype.kind == 'f': vals = vals.reshape(-1, nchan*2, npol) vals = rfft(vals, axis=1, **_rfftargs) else: vals = vals.reshape(-1, nchan, npol) vals = fft(vals, axis=1, **_fftargs) elif dedisperse == 'by-channel' and oversample > 1: vals = vals.reshape(-1, oversample, fh.nchan, npol) vals = fft(vals, axis=1, **_fftargs) vals = vals.transpose(0, 2, 1, 3).reshape(-1, nchan, npol) # vals[time, chan, pol] if verbose >= 2: print("... dedispersed", end="") if npol == 1: power = vals.real**2 + vals.imag**2 else: p0 = vals[..., 0] p1 = vals[..., 1] power = np.empty(vals.shape[:-1] + (4,), np.float32) power[..., 0] = p0.real**2 + p0.imag**2 power[..., 1] = p0.real*p1.real + p0.imag*p1.imag power[..., 2] = p0.imag*p1.real - p0.real*p1.imag power[..., 3] = p1.real**2 + p1.imag**2 if verbose >= 2: print("... power", end="") # current sample positions and corresponding time in stream isr = j*(ntint // oversample) + np.arange(ntint // oversample) tsr = (isr*dtsample*oversample)[:, np.newaxis] if rfi_filter_power is not None: power = rfi_filter_power(power, tsr.squeeze()) print("... power RFI", end="") # correct for delay if needed if dedisperse in ['incoherent', 'by-channel']: # tsample.shape=(ntint/oversample, nchan_in) tsr = tsr - dt if do_waterfall: # # loop over corresponding positions in waterfall # for iw in xrange(isr[0]//ntw, isr[-1]//ntw + 1): # if iw < nwsize: # add sum of corresponding samples # waterfall[iw, :] += np.sum(power[isr//ntw == iw], # axis=0)[ifreq] iw = np.round((tsr / dtsample / oversample).to(1) .value / ntw).astype(int) for k, kfreq in enumerate(ifreq): # sort in frequency while at it iwk = iw[:, (0 if iw.shape[1] == 1 else kfreq // oversample)] iwk = np.clip(iwk, 0, nwsize-1, out=iwk) iwkmin = iwk.min() iwkmax = iwk.max()+1 for ipow in range(npol**2): waterfall[iwkmin:iwkmax, k, ipow] += np.bincount( iwk-iwkmin, power[:, kfreq, ipow], iwkmax-iwkmin) if verbose >= 2: print("... waterfall", end="") if do_foldspec: ibin = (j*ntbin) // nt # bin in the time series: 0..ntbin-1 # times and cycles since start time of observation. tsample = tstart + tsr phase = (phasepol(tsample.to(u.s).value.ravel()) .reshape(tsample.shape)) # corresponding PSR phases iphase = np.remainder(phase*ngate, ngate).astype(np.int) for k, kfreq in enumerate(ifreq): # sort in frequency while at it iph = iphase[:, (0 if iphase.shape[1] == 1 else kfreq // oversample)] # sum and count samples by phase bin for ipow in range(npol**2): foldspec[ibin, k, :, ipow] += np.bincount( iph, power[:, kfreq, ipow], ngate) icount[ibin, k, :] += np.bincount( iph, power[:, kfreq, 0] != 0., ngate).astype(np.int32) if verbose >= 2: print("... folded", end="") if verbose >= 2: print("... done") #Commented out as workaround, this was causing "Referenced before assignment" errors with JB data #if verbose >= 2 or verbose and mpi_rank == 0: # print('#{:4d}/{:4d} read {:6d} out of {:6d}' # .format(mpi_rank, mpi_size, j+1, nt)) if npol == 1: if do_foldspec: foldspec = foldspec.reshape(foldspec.shape[:-1]) if do_waterfall: waterfall = waterfall.reshape(waterfall.shape[:-1]) return foldspec, icount, waterfall