#!/usr/bin/env python # Licensed under a 3-clause BSD style license # Created by Rene Breton, rene.p.breton@gmail.com, 2017-04-28 import numpy as np import scipy def Dedisperse(arr, freqs, DM, dt=1., taxis=0, faxis=1): """ Dedisperse an array to remove the effect of interstellar dispersion. The array is rotated using the convolution theorem: f(t-t0) = IFFT{exp(-2j*pi*w*t0)*FFT{f(t)}} where w are the frequencies of the Fourier transform attached to f(t) and t0 is the time shift Parameters ---------- arr : (...,M,...,N,...) ndarray Array to dedisperse. Must be two dimensional, or more, as long as the frequency axis matches the dimension of the frequency parameter. freqs : (N,) ndarray Array containing the frequencies to calculate the dedispersion at. Units MHz. DM : float Dispersion measure in pc cm^-3 dt : float Size of a time bin in seconds. taxis : int Time axis. faxis : int Frequency axis. """ ## Generating a shape structure shape_tmp = np.array([1]*arr.ndim) ## Calculate the time delay with respect to the top of the band delay = DM_delay(DM, freqs, freqs.max()) ## Convert the delay into a number of bins and shape to enable broadcasting delay /= dt shape = shape_tmp.copy() shape[faxis] = arr.shape[faxis] delay.shape = shape ## Calculate the 'time' frequencies and shape to enable broadcasting w = np.fft.fftfreq(arr.shape[taxis]) shape = shape_tmp.copy() shape[taxis] = arr.shape[taxis] w.shape = shape ## Calculate the phase shift phase_shift = np.exp(2j*np.pi*w*delay) ## Performing the dedispersion dedisp_arr = np.real(np.fft.ifft(phase_shift*np.fft.fft(arr,axis=taxis),axis=taxis)) return dedisp_arr def DM_delay(DM, freq, freq0=np.inf): """ Calculate the time delay in seconds due to dispersion for a given dispersion measure at a frequency with respect to a reference frequency. Parameters ---------- DM : float Dispersion measure in pc cm^-3 freq : float, ndarray Frequency to calculate the delay at. Can be an array. Units MHz. freq0 : float Reference frequency to calculate the time delay with respect to. Units MHz. Return ------ delay : float, ndarray Time delay in seconds. """ ## The DM constant in ms, for frequencies in MHz tau_dm = 4148.808 ## Calculating the delay delay = tau_dm * DM * (freq**(-2)-freq0**(-2)) return delay def Fold(f, t, period, n=128, npart=100, method='integrate'): """ f : ndarray Time series to fold. t : ndarray Times of the time series to fold. period : float Period to fold at. n : int Number of phase bins across the fold. npart : int Number of time bins across the fold. method : string Method used to do the resampling. 'integrate' Properly integrates over the phase bins. It is slower but accurate and works in any conditions. 'interpolate' Linearly interpolates the values at the desired phase bins. It is faster but only works reasonably well in cases where the phase bins are narrower than the sampling time. Example: >>> Make_pulse = lambda phs, phs0, sigma: np.exp(-0.5*((phs-phs0)/sigma)**2) >>> noise = np.random.normal(loc=0., scale=1., size=75000) >>> t = np.arange(noise.size, dtype=float)*0.008 >>> period = 1/11.1853482416 >>> phs = (t/period)%1 >>> signal = Make_pulse(phs, 0.5, 0.05) >>> flux = signal + noise >>> fold = Fold(flux, t, period, n=128, npart=100) """ ## Determining the spin phases phs = t / period ## Determining the number of spin cycles, rounding to the upper value maxphs = int(np.ceil(phs[-1])) ## Determining how many pts will be in the resampled data given the number ## of phase bins npts = n*maxphs+1 ## Generating the interpolated phase bins phs1 = np.linspace(0.,maxphs,npts) ## Performing the resampling according to the chosen method if method == 'integrate': f1 = Interp_linear_integrate(f, phs, phs1) else: #f1 = np.interp(phs1, phs, f) interp = scipy.interpolate.interp1d(phs, f, axis=0, bounds_error=False, fill_value=f[-1]) f1 = interp(phs1) ## The last bin is also extra, dropping it f1 = f1[:-1] ## Determining the new shape so we have a folding axis ## The time dimension is the first, other ones can be present ## such as frequency or polarisations shape = list(f1.shape) shape = [shape[0]/n,n] + shape[1:] f1.shape = shape ## Determining how many time intervals are averaged ncycles = f1.shape[0]/npart shape = [ncycles,npart] + shape[1:] ## Reshaping accordingly, dropping the extra, incomplete intervals f1 = np.resize(f1, shape) ## Averaging over time to downsample f1 = f1.mean(0) return f1 def Interp_linear_integrate(y, x, xnew): """ Resample a time series (x,y) at the values xnew by performing an integration within each new bin of the old time series using the Euler method. Here we assume that the new time series is undersampling the old one, otherwise it is just equivalent to linearly interpolating. Parameters ---------- y : (N,...) ndarray y values to interpolate from. The array can be multi-dimensional. The interpolation will be carried along the first axis. x : (N,) ndarray x values to interpolate from. y = f(x) xnew : (M,) ndarray x values to interpolate at. Return ------ ynew : (M,...) ndarray y values interpolated at. The first dimension is the same as xnew, while the other dimensions, if any, will match the other dimensions of x. >>> x = np.arange(100.) >>> y = np.sin(x/10) >>> xnew = np.arange(20.)*5+0.3 >>> ynew = Interp_linear_integrate(y, x, xnew) """ shape = list(y.shape) shape[0] = xnew.size ynew = np.zeros(shape, dtype=float) i = 0 ii = 0 while ii < xnew.size: weight = 0. val = 0. if ii == 0: xnewl = xnew[ii]-(xnew[ii+1]-xnew[ii])*0.5 else: xnewl = (xnew[ii]+xnew[ii-1])*0.5 if ii == xnew.size-1: xnewr = xnew[ii]+(xnew[ii]-xnew[ii-1])*0.5 else: xnewr = (xnew[ii+1]+xnew[ii])*0.5 while i < x.size: if i == 0: xl = x[i]-(x[i+1]-x[i])*0.5 else: xl = (x[i]+x[i-1])*0.5 if i == x.size-1: xr = x[i]+(x[i]-x[i-1])*0.5 else: xr = (x[i+1]+x[i])*0.5 ## Bin completely inside if xl >= xnewl and xr <= xnewr: weight += xr-xl val += y[i]*(xr-xl) ## Bin overlapping the right side elif xl < xnewr and xr > xnewr: weight += xnewr-xl val += y[i]*(xnewr-xl) ## Means we have to move to next xnew bin break ## Bin overlapping the left side elif xl < xnewl and xr > xnewl: weight += xr-xnewl val += y[i]*(xr-xnewl) ## Bin to the right elif xl >= xnewr: ## Means we are done break ## Bin to the left elif xr <= xnewl: pass ## This condition should not happen else: pass i += 1 ## Add the sum to ynew if weight != 0: ynew[ii] = val/weight ii += 1 return ynew if __name__ == "__main__": Make_pulse = lambda phs, phs0, sigma: np.exp(-0.5*((phs-phs0)/sigma)**2) noise = np.random.normal(loc=0., scale=1., size=75000) t = np.arange(noise.size, dtype=float)*0.008 period = 1/11.1853482416 phs = (t/period)%1 signal = Make_pulse(phs, 0.5, 0.05) flux = signal + noise fold = Fold(flux, t, period, n=128, npart=100)