1c272f1080b734bb7ae891611660edf6c915d2c0
[pitch] / pitch.cpp
1 #include <stdio.h>
2 #include <stdlib.h>
3 #include <unistd.h>
4 #include <fcntl.h>
5 #include <complex>
6 #include <cassert>
7 #include <algorithm>
8 #include <fftw3.h>
9 #include <sys/ioctl.h>
10 #include <linux/soundcard.h>
11
12 #define SAMPLE_RATE     22050
13 #define FFT_LENGTH      4096     /* in samples */
14 #define PAD_FACTOR      2        /* 1/pf of the FFT samples are real samples, the rest are padding */
15 #define OVERLAP         4        /* 1/ol samples will be replaced in the buffer every frame. Should be
16                                   * a multiple of 2 for the Hamming window (see
17                                   * http://www-ccrma.stanford.edu/~jos/parshl/Choice_Hop_Size.html).
18                                   */
19
20 #define EQUAL_TEMPERAMENT     0
21 #define WELL_TEMPERED_GUITAR  1
22
23 #define TUNING WELL_TEMPERED_GUITAR
24
25 int get_dsp_fd();
26 void read_chunk(int fd, double *in, unsigned num_samples);
27 void apply_window(double *in, double *out, unsigned num_samples);
28 std::pair<double, double> find_peak(double *in, unsigned num_samples);
29 void find_peak_magnitudes(std::complex<double> *in, double *out, unsigned num_samples);
30 std::pair<double, double> adjust_for_overtones(std::pair<double, double> base, double *in, unsigned num_samples);
31 double bin_to_freq(double bin, unsigned num_samples);
32 double freq_to_bin(double freq, unsigned num_samples);
33 std::string freq_to_tonename(double freq);
34 std::pair<double, double> interpolate_peak(double ym1, double y0, double y1);
35 void print_spectrogram(double freq, double amp);
36 void write_sine(int dsp_fd, double freq, unsigned num_samples);
37
38 int main()
39 {
40         double *in, *in_windowed;
41         std::complex<double> *out;
42         double *bins;
43         fftw_plan p;
44
45         // allocate memory
46         in = reinterpret_cast<double *> (fftw_malloc(sizeof(double) * FFT_LENGTH / PAD_FACTOR));
47         in_windowed = reinterpret_cast<double *> (fftw_malloc(sizeof(double) * FFT_LENGTH));
48         out = reinterpret_cast<std::complex<double> *> (fftw_malloc(sizeof(std::complex<double>) * (FFT_LENGTH / 2 + 1)));
49         bins = reinterpret_cast<double *> (fftw_malloc(sizeof(double) * FFT_LENGTH / 2 + 1));
50
51         memset(in, 0, sizeof(double) * FFT_LENGTH / PAD_FACTOR);
52
53         // init FFTW
54         p = fftw_plan_dft_r2c_1d(FFT_LENGTH, in_windowed, reinterpret_cast<fftw_complex *> (out), FFTW_ESTIMATE);
55         
56         int fd = get_dsp_fd();
57         for ( ;; ) {
58                 read_chunk(fd, in, FFT_LENGTH);
59                 apply_window(in, in_windowed, FFT_LENGTH);
60                 fftw_execute(p);
61                 find_peak_magnitudes(out, bins, FFT_LENGTH);
62                 std::pair<double, double> peak = find_peak(bins, FFT_LENGTH);
63                 peak = adjust_for_overtones(peak, bins, FFT_LENGTH);
64
65                 if (peak.first < 50.0 || peak.second - log10(FFT_LENGTH) < 0.0) {
66 #if TUNING == WELL_TEMPERED_GUITAR
67                         printf("......\n");
68 #else           
69                         printf("............\n");
70 #endif
71                 } else {
72                         print_spectrogram(peak.first, peak.second - log10(FFT_LENGTH));
73                 }
74         }
75 }
76
77 int get_dsp_fd()
78 {
79         int fd = open("/dev/dsp", O_RDWR);
80         if (fd == -1) {
81                 perror("/dev/dsp");
82                 exit(1);
83         }
84         
85         ioctl(3, SNDCTL_DSP_RESET, 0);
86         
87         int fmt = AFMT_S16_LE;   // FIXME
88         ioctl(fd, SNDCTL_DSP_SETFMT, &fmt);
89
90         int chan = 1;
91         ioctl(fd, SOUND_PCM_WRITE_CHANNELS, &chan);
92         
93         int rate = SAMPLE_RATE;
94         ioctl(fd, SOUND_PCM_WRITE_RATE, &rate);
95
96         int max_fragments = 2;
97         int frag_shift = ffs(FFT_LENGTH / OVERLAP) - 1;
98         int fragments = (max_fragments << 16) | frag_shift;
99         ioctl(fd, SNDCTL_DSP_SETFRAGMENT, &fragments);
100         
101         ioctl(3, SNDCTL_DSP_SYNC, 0);
102         
103         return fd;
104 }
105
106 #if 1
107 void read_chunk(int fd, double *in, unsigned num_samples)
108 {
109         int ret;
110         unsigned to_read = num_samples / PAD_FACTOR / OVERLAP;
111         short buf[to_read];
112
113         memmove(in, in + to_read, (num_samples / PAD_FACTOR - to_read) * sizeof(double));
114         
115         ret = read(fd, buf, to_read * sizeof(short));
116         if (ret == 0) {
117                 printf("EOF\n");
118                 exit(0);
119         }
120
121         if (ret != int(to_read * sizeof(short))) {
122                 // blah
123                 perror("read");
124                 exit(1);
125         }
126         
127         for (unsigned i = 0; i < to_read; ++i)
128                 in[i + (num_samples / PAD_FACTOR - to_read)] = double(buf[i]);
129 }
130 #else
131 // make a pure 440hz sine for testing
132 void read_chunk(int fd, double *in, unsigned num_samples)
133 {
134         static double theta = 0.0;
135         for (unsigned i = 0; i < num_samples; ++i) {
136                 in[i] = cos(theta);
137                 theta += 2.0 * M_PI * 440.0 / double(SAMPLE_RATE);
138         }
139 }
140 #endif
141
142 void write_sine(int dsp_fd, double freq, unsigned num_samples) 
143 {
144         static double theta = 0.0;
145         short buf[num_samples];
146         
147         for (unsigned i = 0; i < num_samples; ++i) {
148                 buf[i] = short(cos(theta) * 16384.0);
149                 theta += 2.0 * M_PI * freq / double(SAMPLE_RATE);
150         }
151
152         write(dsp_fd, buf, num_samples * sizeof(short));
153 }
154
155 // Apply a standard Hamming window to our input data.
156 void apply_window(double *in, double *out, unsigned num_samples)
157 {
158         static double *win_data;
159         static unsigned win_len;
160         static bool win_inited = false;
161
162         // Initialize the window for the first time
163         if (!win_inited) {
164                 win_len = num_samples / PAD_FACTOR;
165                 win_data = new double[win_len];
166
167                 for (unsigned i = 0; i < win_len; ++i) {
168                         win_data[i] = 0.54 - 0.46 * cos(2.0 * M_PI * double(i) / double(win_len - 1));
169                 }
170
171                 win_inited = true;
172         }
173
174         assert(win_len == num_samples / PAD_FACTOR);
175         
176         for (unsigned i = 0; i < win_len; ++i) {
177                 out[i] = in[i] * win_data[i];
178         }
179         for (unsigned i = win_len; i < num_samples; ++i) {
180                 out[i] = 0.0;
181         }
182 }
183
184 void find_peak_magnitudes(std::complex<double> *in, double *out, unsigned num_samples)
185 {
186         for (unsigned i = 0; i < num_samples / 2 + 1; ++i)
187                 out[i] = abs(in[i]);
188 }
189
190 std::pair<double, double> find_peak(double *in, unsigned num_samples)
191 {
192         double best_peak = in[0];
193         unsigned best_bin = 0;
194
195         for (unsigned i = 1; i < num_samples / 2 + 1; ++i) {
196                 if (in[i] > best_peak) {
197                         best_peak = in[i];
198                         best_bin = i;
199                 }
200         }
201         
202         if (best_bin == 0 || best_bin == num_samples / 2) {
203                 return std::make_pair(-1.0, 0.0);
204         }
205
206 #if 0
207         printf("undertone strength: %+4.2f %+4.2f %+4.2f [%+4.2f] %+4.2f %+4.2f %+4.2f\n",
208                 20.0 * log10(in[best_bin*4] / FFT_LENGTH),
209                 20.0 * log10(in[best_bin*3] / FFT_LENGTH),
210                 20.0 * log10(in[best_bin*2] / FFT_LENGTH),
211                 20.0 * log10(in[best_bin] / FFT_LENGTH),
212                 20.0 * log10(in[best_bin/2] / FFT_LENGTH),
213                 20.0 * log10(in[best_bin/3] / FFT_LENGTH),
214                 20.0 * log10(in[best_bin/4] / FFT_LENGTH));
215 #endif
216
217         // see if we might have hit an overtone (set a limit of 5dB)
218         for (unsigned i = 4; i >= 1; --i) {
219                 if (best_bin != best_bin / i &&
220                     20.0 * log10(in[best_bin] / in[best_bin / i]) < 5.0f) {
221 #if 0
222                         printf("Overtone of degree %u!\n", i);
223 #endif
224                         best_bin /= i;
225
226                         // consider sliding one bin up or down
227                         if (best_bin > 0 && in[best_bin - 1] > in[best_bin] && in[best_bin - 1] > in[best_bin - 2]) {
228                                 --best_bin;
229                         } else if (best_bin < num_samples / 2 && in[best_bin + 1] > in[best_bin] && in[best_bin + 1] > in[best_bin + 2]) {
230                                 ++best_bin;
231                         }
232                            
233                         break;
234                 }
235         }
236
237         std::pair<double, double> peak = 
238                 interpolate_peak(in[best_bin - 1],
239                                  in[best_bin],
240                                  in[best_bin + 1]);
241                 
242         return std::make_pair(bin_to_freq(double(best_bin) + peak.first, num_samples), peak.second);
243 }
244
245 // it's perhaps not ideal to _first_ find the peak and _then_ the harmonics --
246 // ideally, one should find the set of all peaks and then determine the likely
247 // base from that... something for later. :-)
248 std::pair<double, double> adjust_for_overtones(std::pair<double, double> base, double *in, unsigned num_samples)
249 {
250         double mu = base.first, var = 1.0 / (base.second * base.second);
251                 
252         //printf("Base at %.2f (amp=%5.2fdB)\n", base.first, base.second);
253
254         for (unsigned i = 2; i < 10; ++i) {
255                 unsigned middle = unsigned(floor(freq_to_bin(base.first, num_samples) * i + 0.5));
256                 unsigned lower = middle - (i+1)/2, upper = middle + (i+1)/2;
257
258                 if (upper >= num_samples)
259                         upper = num_samples - 2;
260
261                 // printf("Searching in [%u,%u] = %f..%f\n", lower, upper, bin_to_freq(lower, num_samples), bin_to_freq(upper, num_samples));
262
263                 // search for a peak in this interval
264                 double best_harmonics_freq = -1.0;
265                 double best_harmonics_amp = -1.0;
266                 for (unsigned j = lower; j <= upper; ++j) {
267                         if (in[j] > in[j-1] && in[j] > in[j+1]) {
268                                 std::pair<double, double> peak =
269                                         interpolate_peak(in[j - 1],
270                                                 in[j],
271                                                 in[j + 1]);
272                                 
273                                 if (peak.second > best_harmonics_amp) {
274                                         best_harmonics_freq = bin_to_freq(j + peak.first, num_samples);
275                                         best_harmonics_amp = peak.second;
276                                 }
277                         }
278                 }
279
280                 if (best_harmonics_amp <= 0.0)
281                         continue;
282
283                 //printf("Found overtone %u at %.2f (amp=%5.2fdB)\n", i, best_harmonics_freq,
284                 //       best_harmonics_amp);
285
286                 double this_mu = best_harmonics_freq / double(i);
287                 double this_var = 1.0 / (best_harmonics_amp * best_harmonics_amp);
288
289                 double k = var / (var + this_var);
290                 mu = (1.0 - k) * mu + k * this_mu;
291                 var *= (1.0 - k);
292         }
293         return std::make_pair(mu, base.second);
294 }
295
296 double bin_to_freq(double bin, unsigned num_samples)
297 {
298         return bin * SAMPLE_RATE / double(num_samples);
299 }
300 double freq_to_bin(double freq, unsigned num_samples)
301 {
302         return freq * double(num_samples) / double(SAMPLE_RATE);
303 }
304
305 /*
306  * Given three bins, find the interpolated real peak based
307  * on their magnitudes. To do this, we execute the following
308  * plan:
309  * 
310  * Fit a polynomial of the form ax^2 + bx + c = 0 to the data
311  * we have. Maple readily yields our coefficients, assuming
312  * that we have the values at x=-1, x=0 and x=1:
313  *
314  * > f := x -> a*x^2 + b*x + c;                                
315  * 
316  *                               2
317  *           f := x -> a x  + b x + c
318  * 
319  * > cf := solve({ f(-1) = ym1, f(0) = y0, f(1) = y1 }, { a, b, c });
320  * 
321  *                            y1    ym1       y1    ym1
322  *        cf := {c = y0, b = ---- - ---, a = ---- + --- - y0}
323  *                            2      2        2      2
324  *
325  * Now let's find the maximum point for the polynomial (it has to be
326  * a maximum, since y0 is the greatest value):
327  *
328  * > xmax := solve(subs(cf, diff(f(x), x)) = 0, x);
329  * 
330  *                                -y1 + ym1
331  *                   xmax := -------------------
332  *                           2 (y1 + ym1 - 2 y0)
333  *
334  * We could go further and insert {fmax,a,b,c} into the original
335  * polynomial, but it just gets hairy. We calculate a, b and c separately
336  * instead.
337  *
338  * http://www-ccrma.stanford.edu/~jos/parshl/Peak_Detection_Steps_3.html
339  * claims that detection is almost twice as good when using dB scale instead
340  * of linear scale for the input values, so we use that. (As a tiny bonus,
341  * we get back dB scale from the function.)
342  */
343 std::pair<double, double> interpolate_peak(double ym1, double y0, double y1)
344 {
345         ym1 = log10(ym1);
346         y0 = log10(y0);
347         y1 = log10(y1);
348
349 #if 0
350         assert(y0 >= y1);
351         assert(y0 >= ym1);
352 #endif
353         
354         double a = 0.5 * y1 + 0.5 * ym1 - y0;
355         double b = 0.5 * y1 - 0.5 * ym1;
356         double c = y0;
357         
358         double xmax = (ym1 - y1) / (2.0 * (y1 + ym1 - 2.0 * y0));
359         double ymax = 20.0 * (a * xmax * xmax + b * xmax + c) - 90.0;
360
361         return std::make_pair(xmax, ymax);
362 }
363
364 std::string freq_to_tonename(double freq)
365 {
366         std::string notenames[] = { "C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B" };
367         double half_notes_away = 12.0 * log2(freq / 440.0) - 3.0;
368         int hnai = int(floor(half_notes_away + 0.5));
369         int octave = (hnai + 48) / 12;
370
371         char buf[256];
372         sprintf(buf, "%s%d + %.2f [%d]", notenames[((hnai % 12) + 12) % 12].c_str(), octave, half_notes_away - hnai, hnai);
373         return buf;
374 }
375
376 #if TUNING == EQUAL_TEMPERAMENT
377 void print_spectrogram(double freq, double amp)
378 {
379         std::string notenames[] = { "C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B" };
380         double half_notes_away = 12.0 * log2(freq / 440.0) - 3.0;
381         int hnai = int(floor(half_notes_away + 0.5));
382         int octave = (hnai + 48) / 12;
383
384         for (int i = 0; i < 12; ++i)
385                 if (i == ((hnai % 12) + 12) % 12)
386                         printf("#");
387                 else
388                         printf(".");
389
390         printf(" (%-2s%d %+.2f, %5.2fHz) [%5.2fdB]  [", notenames[((hnai % 12) + 12) % 12].c_str(), octave, half_notes_away - hnai,
391                 freq, amp);
392
393         double off = half_notes_away - hnai;
394         for (int i = -10; i <= 10; ++i) {
395                 if (off >= (i-0.5) * 0.05 && off < (i+0.5) * 0.05) {
396                         printf("#");
397                 } else {
398                         if (i == 0) {
399                                 printf("|");
400                         } else {
401                                 printf("-");
402                         }
403                 }
404         }
405         printf("]\n");
406
407 }
408 #else
409 struct note {
410         char notename[16];
411         double freq;
412 };
413 static note notes[] = {
414         { "E-3", 110.0 * (3.0/4.0) },
415         { "A-3", 110.0 },
416         { "D-4", 110.0 * (4.0/3.0) },
417         { "G-4", 110.0 * (4.0/3.0)*(4.0/3.0) },
418         { "B-4", 440.0 * (3.0/4.0)*(3.0/4.0) },
419         { "E-5", 440.0 * (3.0/4.0) }
420 };
421
422 void print_spectrogram(double freq, double amp)
423 {
424         double best_away = 999999999.9;
425         unsigned best_away_ind = 0;
426         
427         for (unsigned i = 0; i < sizeof(notes)/sizeof(note); ++i) {
428                 double half_notes_away = 12.0 * log2(freq / notes[i].freq);
429                 if (fabs(half_notes_away) < fabs(best_away)) {
430                         best_away = half_notes_away;
431                         best_away_ind = i;
432                 }
433         }
434         
435         for (unsigned i = 0; i < sizeof(notes)/sizeof(note); ++i)
436                 if (i == best_away_ind)
437                         printf("#");
438                 else
439                         printf(".");
440
441         printf(" (%s %+.2f, %5.2fHz) [%5.2fdB]  [", notes[best_away_ind].notename, best_away, freq, amp);
442
443         for (int i = -10; i <= 10; ++i) {
444                 if (best_away >= (i-0.5) * 0.05 && best_away < (i+0.5) * 0.05) {
445                         printf("#");
446                 } else {
447                         if (i == 0) {
448                                 printf("|");
449                         } else {
450                                 printf("-");
451                         }
452                 }
453         }
454         printf("]\n");
455 }
456 #endif