diff options
Diffstat (limited to 'python/receiver/chirpz.py')
| -rw-r--r-- | python/receiver/chirpz.py | 488 |
1 files changed, 488 insertions, 0 deletions
diff --git a/python/receiver/chirpz.py b/python/receiver/chirpz.py new file mode 100644 index 0000000..3043c44 --- /dev/null +++ b/python/receiver/chirpz.py @@ -0,0 +1,488 @@ +# This program is public domain +# Authors: Paul Kienzle, Nadav Horesh +""" +Chirp z-transform. + +CZT: callable (x,axis=-1)->array + define a chirp-z transform that can be applied to different signals +ZoomFFT: callable (x,axis=-1)->array + define a Fourier transform on a range of frequencies +ScaledFFT: callable (x,axis=-1)->array + define a limited frequency FFT + +czt: array + compute the chirp-z transform for a signal +zoomfft: array + compute the Fourier transform on a range of frequencies +scaledfft: array + compute a limited frequency FFT for a signal +""" +__all__ = ['czt', 'zoomfft', 'scaledfft'] + +import math, cmath + +import numpy as np +from numpy import pi, arange +from scipy.fftpack import fft, ifft, fftshift + +class CZT: + """ + Chirp-Z Transform. + + Transform to compute the frequency response around a spiral. + Objects of this class are callables which can compute the + chirp-z transform on their inputs. This object precalculates + constants for the given transform. + + If w does not lie on the unit circle, then the transform will be + around a spiral with exponentially increasing radius. Regardless, + angle will increase linearly. + + The chirp-z transform can be faster than an equivalent fft with + zero padding. Try it with your own array sizes to see. It is + theoretically faster for large prime fourier transforms, but not + in practice. + + The chirp-z transform is considerably less precise than the + equivalent zero-padded FFT, with differences on the order of 1e-11 + from the direct transform rather than the on the order of 1e-15 as + seen with zero-padding. + + See zoomfft for a friendlier interface to partial fft calculations. + """ + def __init__(self, n, m=None, w=1, a=1): + """ + Chirp-Z transform definition. + + Parameters: + ---------- + n: int + The size of the signal + m: int + The number of points desired. The default is the length of the input data. + a: complex + The starting point in the complex plane. The default is 1. + w: complex or float + If w is complex, it is the ratio between points in each step. + If w is float, it serves as a frequency scaling factor. for instance + when assigning w=0.5, the result FT will span half of frequncy range + (that fft would result) at half of the frequncy step size. + + Returns: + -------- + CZT: + callable object f(x,axis=-1) for computing the chirp-z transform on x + """ + if m is None: + m = n + if w is None: + w = cmath.exp(-1j*pi/m) + elif type(w) in (float, int): + w = cmath.exp(-1j*pi/m * w) + else: + w = cmath.sqrt(w) + self.w, self.a = w, a + self.m, self.n = m, n + + k = arange(max(m,n)) + wk2 = w**(k**2) + nfft = 2**nextpow2(n+m-1) + self._Awk2 = (a**-k * wk2)[:n] + self._nfft = nfft + self._Fwk2 = fft(1/np.hstack((wk2[n-1:0:-1], wk2[:m])), nfft) + self._wk2 = wk2[:m] + self._yidx = slice(n-1, n+m-1) + + def __call__(self, x, axis=-1): + """ + Parameters: + ---------- + x: array + The signal to transform. + axis: int + Array dimension to operate over. The default is the final + dimension. + + Returns: + ------- + An array of the same dimensions as x, but with the length of the + transformed axis set to m. Note that this is a view on a much + larger array. To save space, you may want to call it as + y = czt(x).copy() + """ + x = np.asarray(x) + if x.shape[axis] != self.n: + raise ValueError("CZT defined for length %d, not %d" % + (self.n, x.shape[axis])) + # Calculate transpose coordinates, to allow operation on any given axis + trnsp = np.arange(x.ndim) + trnsp[[axis, -1]] = [-1, axis] + x = x.transpose(*trnsp) + y = ifft(self._Fwk2 * fft(x*self._Awk2, self._nfft)) + y = y[..., self._yidx] * self._wk2 + return y.transpose(*trnsp) + + +def nextpow2(n): + """ + Return the smallest power of two greater than or equal to n. + """ + return int(math.ceil(math.log(n)/math.log(2))) + + +def ZoomFFT(n, f1, f2=None, m=None, Fs=2): + """ + Zoom FFT transform definition. + + Computes the Fourier transform for a set of equally spaced + frequencies. + + Parameters: + ---------- + n: int + size of the signal + m: int + size of the output + f1, f2: float + start and end frequencies; if f2 is not specified, use 0 to f1 + Fs: float + sampling frequency (default=2) + + Returns: + ------- + A CZT instance + A callable object f(x,axis=-1) for computing the zoom FFT on x. + + Sampling frequency is 1/dt, the time step between samples in the + signal x. The unit circle corresponds to frequencies from 0 up + to the sampling frequency. The default sampling frequency of 2 + means that f1,f2 values up to the Nyquist frequency are in the + range [0,1). For f1,f2 values expressed in radians, a sampling + frequency of 1/pi should be used. + + To graph the magnitude of the resulting transform, use:: + + plot(linspace(f1,f2,m), abs(zoomfft(x,f1,f2,m))). + + Use the zoomfft wrapper if you only need to compute one transform. + """ + if m is None: m = n + if f2 is None: f1, f2 = 0., f1 + w = cmath.exp(-2j * pi * (f2-f1) / ((m-1)*Fs)) + a = cmath.exp(2j * pi * f1/Fs) + return CZT(n, m=m, w=w, a=a) + +def ScaledFFT(n, m=None, scale=1.0): + """ + Scaled fft transform definition. + + Similar to fft, where the frequency range is scaled by a factor 'scale' and + divided into 'm-1' equal steps. Like the FFT, frequencies are arranged + from 0 to scale*Fs/2-delta followed by -scale*Fs/2 to -delta, where delta + is the step size scale*Fs/m for sampling frequence Fs. The intended use is in + a convolution of two signals, each has its own sampling step. + + This is equivalent to: + + fftshift(zoomfft(x, -scale, scale*(m-2.)/m, m=m)) + + For example: + + m,n = 10,len(x) + sf = ScaledFFT(n, m=m, scale=0.25) + X = fftshift(fft(x)) + W = linspace(-8, 8*(n-2.)/n, n) + SX = fftshift(sf(x)) + SW = linspace(-2, 2*(m-2.)/m, m) + plot(X,W,SX,SW) + + Parameters: + ---------- + n: int + Size of the signal + m: int + The size of the output. + Default: m=n + scale: float + Frequenct scaling factor. + Default: scale=1.0 + + Returns: + ------- + function + A callable f(x,axis=-1) for computing the scaled FFT on x. + """ + if m is None: + m = n + w = np.exp(-2j * pi / m * scale) + a = w**(m//2) + transform = CZT(n=n, m=m, a=a, w=w) + return lambda x, axis=-1: fftshift(transform(x, axis), axes=(axis,)) + +def scaledfft(x, m=None, scale=1.0, axis=-1): + """ + Partial with a frequency scaling. + See ScaledFFT doc for details + + Parameters: + ---------- + x: input array + m: int + The length of the output signal + scale: float + A frequency scaling factor + axis: int + The array dimension to operate over. The default is the + final dimension. + + Returns: + ------- + An array of the same rank of 'x', but with the size if + the 'axis' dimension set to 'm' + """ + return ScaledFFT(x.shape[axis], m, scale)(x,axis) + +def czt(x, m=None, w=1.0, a=1, axis=-1): + """ + Compute the frequency response around a spiral. + + Parameters: + ---------- + x: array + The set of data to transform. + m: int + The number of points desired. The default is the length of the input data. + a: complex + The starting point in the complex plane. The default is 1. + w: complex or float + If w is complex, it is the ratio between points in each step. + If w is float, it is the frequency step scale (relative to the + normal dft frquency step). + axis: int + Array dimension to operate over. The default is the final + dimension. + + Returns: + ------- + An array of the same dimensions as x, but with the length of the + transformed axis set to m. Note that this is a view on a much + larger array. To save space, you may want to call it as + y = ascontiguousarray(czt(x)) + + See zoomfft for a friendlier interface to partial fft calculations. + + If the transform needs to be repeated, use CZT to construct a + specialized transform function which can be reused without + recomputing constants. + """ + x = np.asarray(x) + transform = CZT(x.shape[axis], m=m, w=w, a=a) + return transform(x,axis=axis) + +def zoomfft(x, f1, f2=None, m=None, Fs=2, axis=-1): + """ + Compute the Fourier transform of x for frequencies in [f1, f2]. + + Parameters: + ---------- + m: int + The number of points to evaluate. The default is the length of x. + f1, f2: float + The frequency range. If f2 is not specified, the range 0-f1 is assumed. + Fs: float + The sampling frequency. With a sampling frequency of + 10kHz for example, the range f1 and f2 can be expressed in kHz. + The default sampling frequency is 2, so f1 and f2 should be + in the range 0,1 to keep the transform below the Nyquist + frequency. + x : array + The input signal. + axis: int + The array dimension the transform operates over. The default is the + final dimension. + + Returns: + ------- + array + The transformed signal. The fourier transform will be calculate + at the points f1, f1+df, f1+2df, ..., f2, where df=(f2-f1)/m. + + zoomfft(x,0,2-2./len(x)) is equivalent to fft(x). + + To graph the magnitude of the resulting transform, use:: + + plot(linspace(f1,f2,m), abs(zoomfit(x,f1,f2,m))). + + If the transform needs to be repeated, use ZoomFFT to construct a + specialized transform function which can be reused without + recomputing constants. + """ + x = np.asarray(x) + transform = ZoomFFT(x.shape[axis], f1, f2=f2, m=m, Fs=Fs) + return transform(x,axis=axis) + + +def _test1(x,show=False,plots=[1,2,3,4]): + norm = np.linalg.norm + + # Normal fft and zero-padded fft equivalent to 10x oversampling + over=10 + w = np.linspace(0,2-2./len(x),len(x)) + y = fft(x) + wover = np.linspace(0,2-2./(over*len(x)),over*len(x)) + yover = fft(x,over*len(x)) + + # Check that zoomfft is the equivalent of fft + y1 = zoomfft(x,0,2-2./len(y)) + + # Check that zoomfft with oversampling is equivalent to zero padding + y2 = zoomfft(x,0,2-2./len(yover), m=len(yover)) + + # Check that zoomfft works on a subrange + f1,f2 = w[3],w[6] + y3 = zoomfft(x,f1,f2,m=3*over+1) + w3 = np.linspace(f1,f2,len(y3)) + idx3 = slice(3*over,6*over+1) + + if not show: plots = [] + if plots != []: + import pylab + if 0 in plots: + pylab.figure(0) + pylab.plot(x) + pylab.ylabel('Intensity') + if 1 in plots: + pylab.figure(1) + pylab.subplot(311) + pylab.plot(w,abs(y),'o',w,abs(y1)) + pylab.legend(['fft','zoom']) + pylab.ylabel('Magnitude') + pylab.title('FFT equivalent') + pylab.subplot(312) + pylab.plot(w,np.angle(y),'o',w,np.angle(y1)) + pylab.legend(['fft','zoom']) + pylab.ylabel('Phase (radians)') + pylab.subplot(313) + pylab.plot(w,abs(y)-abs(y1)) #,w,np.angle(y)-np.angle(y1)) + #pylab.legend(['magnitude','phase']) + pylab.ylabel('Residuals') + if 2 in plots: + pylab.figure(2) + pylab.subplot(211) + pylab.plot(w,abs(y),'o',wover,abs(y2),wover,abs(yover)) + pylab.ylabel('Magnitude') + pylab.title('Oversampled FFT') + pylab.legend(['fft','zoom','pad']) + pylab.subplot(212) + pylab.plot(wover,abs(yover)-abs(y2), + w,abs(y)-abs(y2[0::over]),'o', + w,abs(y)-abs(yover[0::over]),'x') + pylab.legend(['pad-zoom','fft-zoom','fft-pad']) + pylab.ylabel('Residuals') + if 3 in plots: + pylab.figure(3) + ax1=pylab.subplot(211) + pylab.plot(w,abs(y),'o',w3,abs(y3),wover,abs(yover), + w[3:7],abs(y3[::over]),'x') + pylab.title('Zoomed FFT') + pylab.ylabel('Magnitude') + pylab.legend(['fft','zoom','pad']) + pylab.plot(w3,abs(y3),'x') + ax1.set_xlim(f1,f2) + ax2=pylab.subplot(212) + pylab.plot(wover[idx3],abs(yover[idx3])-abs(y3), + w[3:7],abs(y[3:7])-abs(y3[::over]),'o', + w[3:7],abs(y[3:7])-abs(yover[3*over:6*over+1:over]),'x') + pylab.legend(['pad-zoom','fft-zoom','fft-pad']) + ax2.set_xlim(f1,f2) + pylab.ylabel('Residuals') + if plots != []: + pylab.show() + + err = norm(y-y1)/norm(y) + #print "direct err %g"%err + assert err < 1e-10, "error for direct transform is %g"%(err,) + err = norm(yover-y2)/norm(yover) + #print "over err %g"%err + assert err < 1e-10, "error for oversampling is %g"%(err,) + err = norm(yover[idx3]-y3)/norm(yover[idx3]) + #print "range err %g"%err + assert err < 1e-10, "error for subrange is %g"%(err,) + +def _testscaled(x): + n = len(x) + norm = np.linalg.norm + assert norm(fft(x)-scaledfft(x)) < 1e-10 + assert norm(fftshift(fft(x))[n/4:3*n/4] - fftshift(scaledfft(x,scale=0.5,m=n/2))) < 1e-10 + +def test(demo=None,plots=[1,2,3]): + # 0: Gauss + t = np.linspace(-2,2,128) + x = np.exp(-t**2/0.01) + _test1(x, show=(demo==0), plots=plots) + + # 1: Linear + x=[1,2,3,4,5,6,7] + _test1(x, show=(demo==1), plots=plots) + + # Check near powers of two + _test1(range(126-31), show=False) + _test1(range(127-31), show=False) + _test1(range(128-31), show=False) + _test1(range(129-31), show=False) + _test1(range(130-31), show=False) + + # Check transform on n-D array input + x = np.reshape(np.arange(3*2*28),(3,2,28)) + y1 = zoomfft(x,0,2-2./28) + y2 = zoomfft(x[2,0,:],0,2-2./28) + err = np.linalg.norm(y2-y1[2,0]) + assert err < 1e-15, "error for n-D array is %g"%(err,) + + # 2: Random (not a test condition) + if demo==2: + x = np.random.rand(101) + _test1(x, show=True, plots=plots) + + # 3: Spikes + t=np.linspace(0,1,128) + x=np.sin(2*pi*t*5)+np.sin(2*pi*t*13) + _test1(x, show=(demo==3), plots=plots) + + # 4: Sines + x=np.zeros(100) + x[[1,5,21]]=1 + _test1(x, show=(demo==4), plots=plots) + + # 5: Sines plus complex component + x += 1j*np.linspace(0,0.5,x.shape[0]) + _test1(x, show=(demo==5), plots=plots) + + # 6: Scaled FFT on complex sines + x += 1j*np.linspace(0,0.5,x.shape[0]) + if demo == 6: + demo_scaledfft(x,0.25,200) + _testscaled(x) + + +def demo_scaledfft(v, scale, m): + import pylab + shift = pylab.fftshift + n = len(v) + x = pylab.linspace(-0.5, 0.5 - 1./n, n) + xz = pylab.linspace(-scale*0.5, scale*0.5*(m-2.)/m, m) + pylab.figure() + pylab.plot(x, shift(abs(fft(v))), label='fft') + pylab.plot(x, shift(abs(scaledfft(v))),'ro', label='x1 scaled fft') + pylab.plot(xz, abs(zoomfft(v, -scale, scale*(m-2.)/m, m=m)), + 'bo',label='zoomfft') + pylab.plot(xz, shift(abs(scaledfft(v, m=m, scale=scale))), + 'gx', label='x'+str(scale)+' scaled fft') + pylab.gca().set_yscale('log') + pylab.legend() + pylab.show() + +if __name__ == "__main__": + # Choose demo in [0,4] to show plot, or None for testing only + test(demo=None) + |