From: W. Trevor King Date: Mon, 17 May 2010 19:08:29 +0000 (-0400) Subject: Merged my unitary FFT wrappers (FFT_tools) as hooke.util.fft. X-Git-Url: http://git.tremily.us/?p=hooke.git;a=commitdiff_plain;h=7762de28cee60f98882d72db0c2ae2c6009ac465;hp=8273b2acd0162fd19e79cf52ab3822454d5b2c50 Merged my unitary FFT wrappers (FFT_tools) as hooke.util.fft. Origininally versioned in Git and licensed under BSD, but I wrote them, so I'm relicensing under the LGPL for Hooke. The original git repository was: http://www.physics.drexel.edu/~wking/code/git/FFT_tools.git/ --- diff --git a/hooke/util/fft.py b/hooke/util/fft.py new file mode 100644 index 0000000..9167078 --- /dev/null +++ b/hooke/util/fft.py @@ -0,0 +1,584 @@ +# Copyright + +"""Wrap :mod:`numpy.fft` to produce 1D unitary transforms and power spectra. + +Define some FFT wrappers to reduce clutter. +Provides a unitary discrete FFT and a windowed version. +Based on numpy.fft.rfft. + +Main entry functions: + unitary_rfft(data, freq=1.0) + power_spectrum(data, freq=1.0) + unitary_power_spectrum(data, freq=1.0) + avg_power_spectrum(data, freq=1.0, chunk_size=2048, overlap=True, window=window_hann) + unitary_avg_power_spectrum(data, freq=1.0, chunk_size=2048, overlap=True, window=window_hann) +""" + +from numpy import log2, floor, round, ceil, abs, pi, exp, cos, sin, sqrt, \ + sinc, arctan2, array, ones, arange, linspace, zeros, \ + uint16, float, concatenate, fromfile, argmax, complex +from numpy.fft import rfft + + +# print time- and freq- space plots of the test transforms if True +TEST_PLOTS = False +#TEST_PLOTS = True + +def floor_pow_of_two(num) : + "Round num down to the closest exact a power of two." + lnum = log2(num) + if int(lnum) != lnum : + num = 2**floor(lnum) + return num + +def round_pow_of_two(num) : + "Round num to the closest exact a power of two on a log scale." + lnum = log2(num) + if int(lnum) != lnum : + num = 2**round(lnum) + return num + +def ceil_pow_of_two(num) : + "Round num up to the closest exact a power of two." + lnum = log2(num) + if int(lnum) != lnum : + num = 2**ceil(lnum) + return num + +def _test_rfft(xs, Xs) : + # Numpy's FFT algoritm returns + # n-1 + # X[k] = SUM x[m] exp (-j 2pi km /n) + # m=0 + # (see http://www.tramy.us/numpybook.pdf) + j = complex(0,1) + n = len(xs) + Xa = [] + for k in range(n) : + Xa.append(sum([x*exp(-j*2*pi*k*m/n) for x,m in zip(xs,range(n))])) + if k < len(Xs): + assert (Xs[k]-Xa[k])/abs(Xa[k]) < 1e-6, \ + "rfft mismatch on element %d: %g != %g, relative error %g" \ + % (k, Xs[k], Xa[k], (Xs[k]-Xa[k])/abs(Xa[k])) + # Which should satisfy the discrete form of Parseval's theorem + # n-1 n-1 + # SUM |x_m|^2 = 1/n SUM |X_k|^2. + # m=0 k=0 + timeSum = sum([abs(x)**2 for x in xs]) + freqSum = sum([abs(X)**2 for X in Xa]) + assert abs(freqSum/float(n) - timeSum)/timeSum < 1e-6, \ + "Mismatch on Parseval's, %g != 1/%d * %g" % (timeSum, n, freqSum) + +def _test_rfft_suite() : + print "Test numpy rfft definition" + xs = [1,2,3,1,2,3,1,2,3,1,2,3,1,2,3,1] + _test_rfft(xs, rfft(xs)) + +def unitary_rfft(data, freq=1.0) : + """ + Compute the Fourier transform of real data. + Unitary (preserves power [Parseval's theorem]). + + If the units on your data are Volts, + and your sampling frequency is in Hz, + then freq_axis will be in Hz, + and trans will be in Volts. + """ + nsamps = floor_pow_of_two(len(data)) + # Which should satisfy the discrete form of Parseval's theorem + # n-1 n-1 + # SUM |x_m|^2 = 1/n SUM |X_k|^2. + # m=0 k=0 + # However, we want our FFT to satisfy the continuous Parseval eqn + # int_{-infty}^{infty} |x(t)|^2 dt = int_{-infty}^{infty} |X(f)|^2 df + # which has the discrete form + # n-1 n-1 + # SUM |x_m|^2 dt = SUM |X'_k|^2 df + # m=0 k=0 + # with X'_k = AX, this gives us + # n-1 n-1 + # SUM |x_m|^2 = A^2 df/dt SUM |X'_k|^2 + # m=0 k=0 + # so we see + # A^2 df/dt = 1/n + # A^2 = 1/n dt/df + # From Numerical Recipes (http://www.fizyka.umk.pl/nrbook/bookcpdf.html), + # Section 12.1, we see that for a sampling rate dt, the maximum frequency + # f_c in the transformed data is the Nyquist frequency (12.1.2) + # f_c = 1/2dt + # and the points are spaced out by (12.1.5) + # df = 1/ndt + # so + # dt = 1/ndf + # dt/df = 1/ndf^2 + # A^2 = 1/n^2df^2 + # A = 1/ndf = ndt/n = dt + # so we can convert the Numpy transformed data to match our unitary + # continuous transformed data with (also NR 12.1.8) + # X'_k = dtX = X / + trans = rfft(data[0:nsamps]) / float(freq) + freq_axis = linspace(0, freq/2, nsamps/2+1) + return (freq_axis, trans) + +def _test_unitary_rfft_parsevals(xs, freq, freqs, Xs): + # Which should satisfy the discretized integral form of Parseval's theorem + # n-1 n-1 + # SUM |x_m|^2 dt = SUM |X_k|^2 df + # m=0 k=0 + dt = 1.0/freq + df = freqs[1]-freqs[0] + assert (df - 1/(len(xs)*dt))/df < 1e-6, \ + "Mismatch in spacing, %g != 1/(%d*%g)" % (df, len(xs), dt) + Xa = list(Xs) + for k in range(len(Xs)-1,1,-1) : + Xa.append(Xa[k]) + assert len(xs) == len(Xa), "Length mismatch %d != %d" % (len(xs), len(Xa)) + lhs = sum([abs(x)**2 for x in xs]) * dt + rhs = sum([abs(X)**2 for X in Xa]) * df + assert abs(lhs - rhs)/lhs < 1e-4, "Mismatch on Parseval's, %g != %g" \ + % (lhs, rhs) + +def _test_unitary_rfft_parsevals_suite(): + print "Test unitary rfft on Parseval's theorem" + xs = [1,2,3,1,2,3,1,2,3,1,2,3,1,2,3,1] + dt = pi + freqs,Xs = unitary_rfft(xs, 1.0/dt) + _test_unitary_rfft_parsevals(xs, 1.0/dt, freqs, Xs) + +def _rect(t) : + if abs(t) < 0.5 : + return 1 + else : + return 0 + +def _test_unitary_rfft_rect(a=1.0, time_shift=5.0, samp_freq=25.6, samples=256) : + "Show fft(rect(at)) = 1/abs(a) * sinc(f/a)" + samp_freq = float(samp_freq) + a = float(a) + + x = zeros((samples,), dtype=float) + dt = 1.0/samp_freq + for i in range(samples) : + t = i*dt + x[i] = _rect(a*(t-time_shift)) + freq_axis, X = unitary_rfft(x, samp_freq) + #_test_unitary_rfft_parsevals(x, samp_freq, freq_axis, X) + + # remove the phase due to our time shift + j = complex(0.0,1.0) # sqrt(-1) + for i in range(len(freq_axis)) : + f = freq_axis[i] + inverse_phase_shift = exp(j*2.0*pi*time_shift*f) + X[i] *= inverse_phase_shift + + expected = zeros((len(freq_axis),), dtype=float) + # normalized sinc(x) = sin(pi x)/(pi x) + # so sinc(0.5) = sin(pi/2)/(pi/2) = 2/pi + assert sinc(0.5) == 2.0/pi, "abnormal sinc()" + for i in range(len(freq_axis)) : + f = freq_axis[i] + expected[i] = 1.0/abs(a) * sinc(f/a) + + if TEST_PLOTS : + pylab.figure() + pylab.subplot(211) + pylab.plot(arange(0, dt*samples, dt), x) + pylab.title('time series') + pylab.subplot(212) + pylab.plot(freq_axis, X.real, 'r.') + pylab.plot(freq_axis, X.imag, 'g.') + pylab.plot(freq_axis, expected, 'b-') + pylab.title('freq series') + +def _test_unitary_rfft_rect_suite() : + print "Test unitary FFTs on variously shaped rectangular functions" + _test_unitary_rfft_rect(a=0.5) + _test_unitary_rfft_rect(a=2.0) + _test_unitary_rfft_rect(a=0.7, samp_freq=50, samples=512) + _test_unitary_rfft_rect(a=3.0, samp_freq=60, samples=1024) + +def _gaussian(a, t) : + return exp(-a * t**2) + +def _test_unitary_rfft_gaussian(a=1.0, time_shift=5.0, samp_freq=25.6, samples=256) : + "Show fft(rect(at)) = 1/abs(a) * sinc(f/a)" + samp_freq = float(samp_freq) + a = float(a) + + x = zeros((samples,), dtype=float) + dt = 1.0/samp_freq + for i in range(samples) : + t = i*dt + x[i] = _gaussian(a, (t-time_shift)) + freq_axis, X = unitary_rfft(x, samp_freq) + #_test_unitary_rfft_parsevals(x, samp_freq, freq_axis, X) + + # remove the phase due to our time shift + j = complex(0.0,1.0) # sqrt(-1) + for i in range(len(freq_axis)) : + f = freq_axis[i] + inverse_phase_shift = exp(j*2.0*pi*time_shift*f) + X[i] *= inverse_phase_shift + + expected = zeros((len(freq_axis),), dtype=float) + for i in range(len(freq_axis)) : + f = freq_axis[i] + expected[i] = sqrt(pi/a) * _gaussian(1.0/a, pi*f) # see Wikipedia, or do the integral yourself. + + if TEST_PLOTS : + pylab.figure() + pylab.subplot(211) + pylab.plot(arange(0, dt*samples, dt), x) + pylab.title('time series') + pylab.subplot(212) + pylab.plot(freq_axis, X.real, 'r.') + pylab.plot(freq_axis, X.imag, 'g.') + pylab.plot(freq_axis, expected, 'b-') + pylab.title('freq series') + +def _test_unitary_rfft_gaussian_suite() : + print "Test unitary FFTs on variously shaped gaussian functions" + _test_unitary_rfft_gaussian(a=0.5) + _test_unitary_rfft_gaussian(a=2.0) + _test_unitary_rfft_gaussian(a=0.7, samp_freq=50, samples=512) + _test_unitary_rfft_gaussian(a=3.0, samp_freq=60, samples=1024) + + + +def power_spectrum(data, freq=1.0) : + """ + Compute the power spectrum of DATA taken with a sampling frequency FREQ. + DATA must be real (not complex). + Returns a tuple of two arrays, (freq_axis, power), suitable for plotting. + If the number of samples in data is not an integer power of two, + the FFT ignores some of the later points. + """ + nsamps = floor_pow_of_two(len(data)) + + freq_axis = linspace(0, freq/2, nsamps/2+1) + # nsamps/2+1 b/c zero-freq and nyqist-freq are both fully real. + # >>> help(numpy.fft.fftpack.rfft) for Numpy's explaination. + # See Numerical Recipies for a details. + trans = rfft(data[0:nsamps]) + power = trans * trans.conj() # We want the square of the amplitude. + return (freq_axis, power) + +def unitary_power_spectrum(data, freq=1.0) : + freq_axis,power = power_spectrum(data, freq) + # One sided power spectral density, so 2|H(f)|**2 (see NR 2nd edition 12.0.14, p498) + # + # numpy normalizes with 1/N on the inverse transform ifft, + # so we should normalize the freq-space representation with 1/sqrt(N). + # But we're using the rfft, where N points are like N/2 complex points, so 1/sqrt(N/2) + # So the power gets normalized by that twice and we have 2/N + # + # On top of this, the FFT assumes a sampling freq of 1 per second, + # and we want to preserve area under our curves. + # If our total time T = len(data)/freq is smaller than 1, + # our df_real = freq/len(data) is bigger that the FFT expects (dt_fft = 1/len(data)), + # and we need to scale the powers down to conserve area. + # df_fft * F_fft(f) = df_real *F_real(f) + # F_real = F_fft(f) * (1/len)/(freq/len) = F_fft(f)*freq + # So the power gets normalized by *that* twice and we have 2/N * freq**2 + + # power per unit time + # measure x(t) for time T + # X(f) = int_0^T x(t) exp(-2 pi ift) dt + # PSD(f) = 2 |X(f)|**2 / T + + # total_time = len(data)/float(freq) + # power *= 2.0 / float(freq)**2 / total_time + # power *= 2.0 / freq**2 * freq / len(data) + power *= 2.0 / (freq * float(len(data))) + + return (freq_axis, power) + +def _test_unitary_power_spectrum_sin(sin_freq=10, samp_freq=512, samples=1024) : + x = zeros((samples,), dtype=float) + samp_freq = float(samp_freq) + for i in range(samples) : + x[i] = sin(2.0 * pi * (i/samp_freq) * sin_freq) + freq_axis, power = unitary_power_spectrum(x, samp_freq) + imax = argmax(power) + + expected = zeros((len(freq_axis),), dtype=float) + df = samp_freq/float(samples) # df = 1/T, where T = total_time + i = int(sin_freq/df) + # average power per unit time is + # P = + # average value of sin(t)**2 = 0.5 (b/c sin**2+cos**2 == 1) + # so average value of (int sin(t)**2 dt) per unit time is 0.5 + # P = 0.5 + # we spread that power over a frequency bin of width df, sp + # P(f0) = 0.5/df + # where f0 is the sin's frequency + # + # or : + # FFT of sin(2*pi*t*f0) gives + # X(f) = 0.5 i (delta(f-f0) - delta(f-f0)), + # (area under x(t) = 0, area under X(f) = 0) + # so one sided power spectral density (PSD) per unit time is + # P(f) = 2 |X(f)|**2 / T + # = 2 * |0.5 delta(f-f0)|**2 / T + # = 0.5 * |delta(f-f0)|**2 / T + # but we're discrete and want the integral of the 'delta' to be 1, + # so 'delta'*df = 1 --> 'delta' = 1/df, and + # P(f) = 0.5 / (df**2 * T) + # = 0.5 / df (T = 1/df) + expected[i] = 0.5 / df + + print "The power should be a peak at %g Hz of %g (%g, %g)" % \ + (sin_freq, expected[i], freq_axis[imax], power[imax]) + Pexp = 0 + P = 0 + for i in range(len(freq_axis)) : + Pexp += expected[i] *df + P += power[i] * df + print " The total power should be %g (%g)" % (Pexp, P) + + if TEST_PLOTS : + pylab.figure() + pylab.subplot(211) + pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-') + pylab.title('time series') + pylab.subplot(212) + pylab.plot(freq_axis, power, 'r.') + pylab.plot(freq_axis, expected, 'b-') + pylab.title('%g samples of sin at %g Hz' % (samples, sin_freq)) + +def _test_unitary_power_spectrum_sin_suite() : + print "Test unitary power spectrums on variously shaped sin functions" + _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=1024) + _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=2048) + _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=4098) + _test_unitary_power_spectrum_sin(sin_freq=7, samp_freq=512, samples=1024) + _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=1024, samples=2048) + # finally, with some irrational numbers, to check that I'm not getting lucky + _test_unitary_power_spectrum_sin(sin_freq=pi, samp_freq=100*exp(1), samples=1024) + # test with non-integer number of periods + _test_unitary_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=256) + +def _test_unitary_power_spectrum_delta(amp=1, samp_freq=1, samples=256) : + x = zeros((samples,), dtype=float) + samp_freq = float(samp_freq) + x[0] = amp + freq_axis, power = unitary_power_spectrum(x, samp_freq) + + # power = = (amp)**2 * dt/T + # we spread that power over the entire freq_axis [0,fN], so + # P(f) = (amp)**2 dt / (T fN) + # where + # dt = 1/samp_freq (sample period) + # T = samples/samp_freq (total time of data aquisition) + # fN = 0.5 samp_freq (Nyquist frequency) + # so + # P(f) = amp**2 / (samp_freq * samples/samp_freq * 0.5 samp_freq) + # = 2 amp**2 / (samp_freq*samples) + expected_amp = 2.0 * amp**2 / (samp_freq * samples) + expected = ones((len(freq_axis),), dtype=float) * expected_amp + + print "The power should be flat at y = %g (%g)" % (expected_amp, power[0]) + + if TEST_PLOTS : + pylab.figure() + pylab.subplot(211) + pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-') + pylab.title('time series') + pylab.subplot(212) + pylab.plot(freq_axis, power, 'r.') + pylab.plot(freq_axis, expected, 'b-') + pylab.title('%g samples of delta amp %g' % (samples, amp)) + +def _test_unitary_power_spectrum_delta_suite() : + print "Test unitary power spectrums on various delta functions" + _test_unitary_power_spectrum_delta(amp=1, samp_freq=1.0, samples=1024) + _test_unitary_power_spectrum_delta(amp=1, samp_freq=1.0, samples=2048) + _test_unitary_power_spectrum_delta(amp=1, samp_freq=0.5, samples=2048)# expected = 2*computed + _test_unitary_power_spectrum_delta(amp=1, samp_freq=2.0, samples=2048)# expected = 0.5*computed + _test_unitary_power_spectrum_delta(amp=3, samp_freq=1.0, samples=1024) + _test_unitary_power_spectrum_delta(amp=pi, samp_freq=exp(1), samples=1024) + +def _gaussian2(area, mean, std, t) : + "Integral over all time = area (i.e. normalized for area=1)" + return area/(std*sqrt(2.0*pi)) * exp(-0.5*((t-mean)/std)**2) + +def _test_unitary_power_spectrum_gaussian(area=2.5, mean=5, std=1, samp_freq=10.24 ,samples=512) : #1024 + x = zeros((samples,), dtype=float) + mean = float(mean) + for i in range(samples) : + t = i/float(samp_freq) + x[i] = _gaussian2(area, mean, std, t) + freq_axis, power = unitary_power_spectrum(x, samp_freq) + + # generate the predicted curve + # by comparing our _gaussian2() form to _gaussian(), + # we see that the Fourier transform of x(t) has parameters: + # std' = 1/(2 pi std) (references declaring std' = 1/std are converting to angular frequency, not frequency like we are) + # area' = area/[std sqrt(2*pi)] (plugging into FT of _gaussian() above) + # mean' = 0 (changing the mean in the time-domain just changes the phase in the freq-domain) + # So our power spectral density per unit time is given by + # P(f) = 2 |X(f)|**2 / T + # Where + # T = samples/samp_freq (total time of data aquisition) + mean = 0.0 + area = area /(std*sqrt(2.0*pi)) + std = 1.0/(2.0*pi*std) + expected = zeros((len(freq_axis),), dtype=float) + df = float(samp_freq)/samples # 1/total_time ( = freq_axis[1]-freq_axis[0] = freq_axis[1]) + for i in range(len(freq_axis)) : + f = i*df + gaus = _gaussian2(area, mean, std, f) + expected[i] = 2.0 * gaus**2 * samp_freq/samples + print "The power should be a half-gaussian, ", + print "with a peak at 0 Hz with amplitude %g (%g)" % (expected[0], power[0]) + + if TEST_PLOTS : + pylab.figure() + pylab.subplot(211) + pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-') + pylab.title('time series') + pylab.subplot(212) + pylab.plot(freq_axis, power, 'r.') + pylab.plot(freq_axis, expected, 'b-') + pylab.title('freq series') + +def _test_unitary_power_spectrum_gaussian_suite() : + print "Test unitary power spectrums on various gaussian functions" + _test_unitary_power_spectrum_gaussian(area=1, std=1, samp_freq=10.0, samples=1024) + _test_unitary_power_spectrum_gaussian(area=1, std=2, samp_freq=10.0, samples=1024) + _test_unitary_power_spectrum_gaussian(area=1, std=1, samp_freq=10.0, samples=2048) + _test_unitary_power_spectrum_gaussian(area=1, std=1, samp_freq=20.0, samples=2048) + _test_unitary_power_spectrum_gaussian(area=3, std=1, samp_freq=10.0, samples=1024) + _test_unitary_power_spectrum_gaussian(area=pi, std=sqrt(2), samp_freq=exp(1), samples=1024) + +def window_hann(length) : + "Returns a Hann window array with length entries" + win = zeros((length,), dtype=float) + for i in range(length) : + win[i] = 0.5*(1.0-cos(2.0*pi*float(i)/(length))) + # avg value of cos over a period is 0 + # so average height of Hann window is 0.5 + return win + +def avg_power_spectrum(data, freq=1.0, chunk_size=2048, + overlap=True, window=window_hann) : + """ + Compute the avg power spectrum of DATA taken with a sampling frequency FREQ. + DATA must be real (not complex) by breaking DATA into chunks. + The chunks may or may not be overlapping (by setting OVERLAP). + The chunks are windowed by dotting with WINDOW(CHUNK_SIZE), FFTed, + and the resulting spectra are averaged together. + See NR 13.4 for rational. + + Returns a tuple of two arrays, (freq_axis, power), suitable for plotting. + CHUNK_SIZE should really be a power of 2. + If the number of samples in DATA is not an integer power of CHUNK_SIZE, + the FFT ignores some of the later points. + """ + assert chunk_size == floor_pow_of_two(chunk_size), \ + "chunk_size %d should be a power of 2" % chunk_size + + nchunks = len(data)/chunk_size # integer division = implicit floor + if overlap : + chunk_step = chunk_size/2 + else : + chunk_step = chunk_size + + win = window(chunk_size) # generate a window of the appropriate size + freq_axis = linspace(0, freq/2, chunk_size/2+1) + # nsamps/2+1 b/c zero-freq and nyqist-freq are both fully real. + # >>> help(numpy.fft.fftpack.rfft) for Numpy's explaination. + # See Numerical Recipies for a details. + power = zeros((chunk_size/2+1,), dtype=float) + for i in range(nchunks) : + starti = i*chunk_step + stopi = starti+chunk_size + fft_chunk = rfft(data[starti:stopi]*win) + p_chunk = fft_chunk * fft_chunk.conj() + power += p_chunk.astype(float) + power /= float(nchunks) + return (freq_axis, power) + +def unitary_avg_power_spectrum(data, freq=1.0, chunk_size=2048, + overlap=True, window=window_hann) : + """ + compute the average power spectrum, preserving normalization + """ + freq_axis,power = avg_power_spectrum(data, freq, chunk_size, + overlap, window) + # 2.0 / (freq * chunk_size) |rfft()|**2 --> unitary_power_spectrum + power *= 2.0 / (freq*float(chunk_size)) * 8/3 # see unitary_power_spectrum() + # * 8/3 to remove power from windowing + # <[x(t)*w(t)]**2> = ~= * + # where the ~= is because the frequency of x(t) >> the frequency of w(t). + # So our calulated power has and extra in it. + # For the Hann window, = <0.5(1 + 2cos + cos**2)> = 1/4 + 0 + 1/8 = 3/8 + # For low frequency components, where the frequency of x(t) is ~= the frequency of w(t), + # The normalization is not perfect. ?? + # The normalization approaches perfection as chunk_size -> infinity. + return (freq_axis, power) + +def _test_unitary_avg_power_spectrum_sin(sin_freq=10, samp_freq=512, samples=1024, + chunk_size=512, overlap=True, + window=window_hann) : + x = zeros((samples,), dtype=float) + samp_freq = float(samp_freq) + for i in range(samples) : + x[i] = sin(2.0 * pi * (i/samp_freq) * sin_freq) + freq_axis, power = unitary_avg_power_spectrum(x, samp_freq, chunk_size, + overlap, window) + imax = argmax(power) + + expected = zeros((len(freq_axis),), dtype=float) + df = samp_freq/float(chunk_size) # df = 1/T, where T = total_time + i = int(sin_freq/df) + expected[i] = 0.5 / df # see _test_unitary_power_spectrum_sin() + + print "The power should be a peak at %g Hz of %g (%g, %g)" % \ + (sin_freq, expected[i], freq_axis[imax], power[imax]) + Pexp = 0 + P = 0 + for i in range(len(freq_axis)) : + Pexp += expected[i] * df + P += power[i] * df + print " The total power should be %g (%g)" % (Pexp, P) + + if TEST_PLOTS : + pylab.figure() + pylab.subplot(211) + pylab.plot(arange(0, samples/samp_freq, 1.0/samp_freq), x, 'b-') + pylab.title('time series') + pylab.subplot(212) + pylab.plot(freq_axis, power, 'r.') + pylab.plot(freq_axis, expected, 'b-') + pylab.title('%g samples of sin at %g Hz' % (samples, sin_freq)) + +def _test_unitary_avg_power_spectrum_sin_suite() : + print "Test unitary avg power spectrums on variously shaped sin functions" + _test_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=1024) + _test_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=2048) + _test_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=512, samples=4098) + _test_unitary_avg_power_spectrum_sin(sin_freq=17, samp_freq=512, samples=1024) + _test_unitary_avg_power_spectrum_sin(sin_freq=5, samp_freq=1024, samples=2048) + # test long wavelenth sin, so be closer to window frequency + _test_unitary_avg_power_spectrum_sin(sin_freq=1, samp_freq=1024, samples=2048) + # finally, with some irrational numbers, to check that I'm not getting lucky + _test_unitary_avg_power_spectrum_sin(sin_freq=pi, samp_freq=100*exp(1), samples=1024) + + +def test() : + _test_rfft_suite() + _test_unitary_rfft_parsevals_suite() + _test_unitary_rfft_rect_suite() + _test_unitary_rfft_gaussian_suite() + _test_unitary_power_spectrum_sin_suite() + _test_unitary_power_spectrum_delta_suite() + _test_unitary_power_spectrum_gaussian_suite() + _test_unitary_avg_power_spectrum_sin_suite() + +if __name__ == "__main__" : + if TEST_PLOTS : + import pylab + test() + if TEST_PLOTS : + pylab.show()