#A worksheet to analyze frequency content of a signal from sampled data
#using the discrete cosine transform (DCT).
#Parameters tlow and thigh control plot range for time domain plot.
#Parameters flow and fhigh control frequency range for DCT plot.

reset();

#Load in the appropriate data set, here in the form of (time,signal) pairs
#from the supplied files.

#load('gong.sage'); samprate = 16000.0; #sampled a 16 kHz
#load('sunspot_data.sage'); samprate = 1.0; #sampled at 1 per year
load('spring_mass_data_fourier.sage'); samprate = 50.0; #sampled at 50 Hz.

#The data is in the array "udat0".
#Number of data points "N" and time interval "T":
N = len(udat0);
T = N/samprate;

#Select a time interval for plot and plot.
tlow = 0.0; thigh = T;
jlow = floor(N*tlow/T); jhigh = min(ceil(N*thigh/T),N);
J2 = list(range(jlow,jhigh));
pdat = [udat0[j] for j in J2];
scatter_plot(pdat,facecolor='red',markersize=10,axes_labels=['Time','Signal Magnitude']).show();

#Compute the DCT of the signal, bootstrapping off of the fft command.

J = list(range(N)); #list 0 to N-1
A = [RR(udat0[i][1]) for i in J]; #Pick off signal components
A2 = [A[i] for i in J]; #Make a copy
A.reverse(); #Reverse A
A2.extend(A); #and append to A2
J2 = list(range(2*N));
udat = IndexedSequence(A2,J2); #Data structure for FFT command
C = udat.fft(); #Compute FFT
fC = C.list(); #Pick off components

#Modify fC to obtain DCT of data.
p2 = n(pi/2);
W = [exp(-p2*I*j) for j in J];
DC = [real(W[j]*fC[j])/sqrt(2.0*n(N)) for j in J];
DC[0] = DC[0]/sqrt(2.0); #Rescale DC[0] special.

#Plot DCT coefficient magnitudes.
J = list(range(0,N));
DCmag = [abs(DC[j]) for j in J];
freqs = [0.5*j*samprate/n(N) for j in J];

#Select frequency range for plot, minimum frequency 0, maximum samprate/2
flow = 1.0; fhigh = 25.0;
jlow = floor(2*flow*N/samprate); jhigh = min(ceil(2*fhigh*N/samprate),N);
J2 = list(range(jlow,jhigh));
pdat = [[freqs[j],DCmag[j]] for j in J2];
scatter_plot(pdat,facecolor='red',markersize=10,axes_labels=['Frequency (Hz)','DCT Magnitude']).show();