pitch_tracking.py

import matplotlib.pyplot as plt
import matplotlib.collections as collections
plt.style.use('bmh')

import numpy as np
# from scipy.signal import fftconvolve

#
# pitch_tracking_sine_wave
#
def pitch_tracking_sine_wave(f0, fs, frame_size, phase=0):
	t = np.arange(frame_size) / fs
	return np.sin(2.0 * np.pi * f0 * t + phase)

#
# pitch_tracking_harmonic_wave
#
def pitch_tracking_harmonic_wave(f0, fs, frame_size, n_harmonics=10):
	""" creating a synthetic signal with a given fundamental """
	signal = np.zeros((frame_size,))
	for f in [f0 * i for i in range(1, n_harmonics + 1)]:
		signal += f0 / f * pitch_tracking_sine_wave(f, fs, frame_size, phase=f)
	return signal

#
# pitch_tracking_harmonic_wave_noisy
#
def pitch_tracking_harmonic_wave_noisy(noise, f0, fs, frame_size, n_harmonics=10):
	signal = pitch_tracking_harmonic_wave(f0, fs, frame_size, n_harmonics=n_harmonics)
	return signal + noise

#
# pitch_tracking_noise_generator_psd
#
def pitch_tracking_noise_generator_psd(frame_size, psd = lambda f: 1.0):
	white = np.fft.rfft(np.random.randn(frame_size));
	raw   = psd(np.fft.rfftfreq(frame_size))
	raw   = raw / np.sqrt(np.mean(raw**2))
	return np.fft.irfft(white * raw);

#
# pitch_tracking_noise_generator
#
def pitch_tracking_noise_generator(f):
	return lambda frame_size: pitch_tracking_noise_generator_psd(frame_size, f)

#
# pitch_tracking_brownian_noise
#
@pitch_tracking_noise_generator
def pitch_tracking_brownian_noise(f):
	return 1.0 / np.where(f == 0, float('inf'), f)

#
# pitch_tracking_pink_noise
#
@pitch_tracking_noise_generator
def pitch_tracking_pink_noise(f):
	return 1.0 / np.where(f == 0, float('inf'), np.sqrt(f))

#
# pitch_tracking_white_noise
#
@pitch_tracking_noise_generator
def pitch_tracking_white_noise(f):
	return 1.0

#
# pitch_tracking_blue_noise
#
@pitch_tracking_noise_generator
def pitch_tracking_blue_noise(f):
	return np.sqrt(f)

#
# pitch_tracking_violet_noise
#
@pitch_tracking_noise_generator
def pitch_tracking_violet_noise(f):
	return f

#
# pitch_tracking_show_noise
#
def pitch_tracking_show_noise():
	plt.figure(figsize=(16, 16))
	for G in [
			  pitch_tracking_brownian_noise
			, pitch_tracking_pink_noise
			, pitch_tracking_white_noise
			, pitch_tracking_blue_noise
			, pitch_tracking_violet_noise
	]:
		sp = G(2**8)
		fq = np.fft.rfftfreq(sp.size)
		plt.loglog(fq, np.abs(np.fft.rfft(sp)))
	plt.legend(['Brownian', 'Pink', 'White', 'Blue', 'Violet'])
	plt.ylim([1E-4, None])

#
# pitch_tracking_psd
#
def pitch_tracking_psd(signal, fs, window=True, threshold=0.0):
	""" computing a noisy raw psd """
	if window is True:
		wx  = np.hamming(signal.size) * signal
		psd = np.abs(np.fft.rfft(wx))**2 / wx.size;
	else:
		psd = np.abs(np.fft.rfft(signal))**2 / signal.size;
	if threshold != 0.0:
		""" de-noise/trim/zeroing everything under threshold """
		idx = psd > threshold
		psd = psd * idx
	rms = np.sqrt(np.mean(psd**2))
	return rms, psd

#
# pitch_tracking_autocorrelation_kernel
#
def pitch_tracking_autocorrelation_kernel(signal, threshold=0.0):
	""" autocorrelating input signal via FFT """
	var = numpy.var(signal)
	""" centering """
	y   = signal - numpy.mean(signal)
	y   = np.fft.rfft(y)
	psd = np.abs(y) ** 2
	if threshold != 0.0:
		""" de-noise/trim/zeroing everything under threshold """
		idx = psd > threshold
		psd = psd * idx
	return numpy.fft.ifft(psd).real / var / psd.size

#
# pitch_tracking_cepstrum
#
def pitch_tracking_cepstrum(signal, fs, window=True, threshold=0.0):
	""" computing rms, psd and cepstrum of input signal """
	if window is True:
		wx = np.hamming(signal.size) * signal
		y  = np.fft.rfft(wx)
	else:
		y  = np.fft.rfft(signal)
	dt        = 1.0 / fs
	fq        = np.fft.rfftfreq(signal.size, d=dt)
	mag       = np.abs(y)
	psd       = mag**2 / signal.size;
	if threshold != 0.0:
		""" de-noise/trim/zeroing everything under threshold """
		idx = psd > threshold
		mag = mag * idx
		psd = psd * idx
	rms       = np.sqrt(np.mean(psd**2))
	mel       = np.log(mag)
	cepstrum  = np.fft.rfft(mel)
	df        = fq[1] - fq[0]
	quefrency = np.fft.rfftfreq(mel.size, df)
	return rms, psd, quefrency, cepstrum

#
# pitch_tracking_best_cepstrum
#
def pitch_tracking_best_cepstrum(x, y, z, fs, window=True, threshold=0.0):
	""" selecting spectrum having the most energy and returns its cepstrum """
	rms_x, _, quefrency_x, cepstrum_x = pitch_tracking_cepstrum(x, fs, window=window, threshold=threshold)
	rms_y, _, quefrency_y, cepstrum_y = pitch_tracking_cepstrum(y, fs, window=window, threshold=threshold)
	rms_z, _, quefrency_z, cepstrum_z = pitch_tracking_cepstrum(z, fs, window=window, threshold=threshold)
	if (rms_x >= rms_y) and (rms_x >= rms_z):
		return quefrency_x, cepstrum_x
	elif (rms_y >= rms_x) and (rms_y >= rms_z):
		return quefrency_y, cepstrum_y
	return quefrency_z, cepstrum_z

#
# pitch_tracking_show
#
def pitch_tracking_show(quefrency, cepstrum, f0, freq_lo=30.0, freq_hi=80.0):
	""" plotting cepstrum and best selection range for fundamental peak """
	fig, ax      = plt.subplots()
	ax.plot(quefrency, np.abs(cepstrum))
	ax.set_xlabel('quefrency (s)')
	ax.set_title('cepstrum')

	fig, ax      = plt.subplots()
	ax.vlines(1.0 / f0, 0, np.max(np.abs(cepstrum)), alpha=0.2, lw=3, label='expected peak')
	ax.plot(quefrency, np.abs(cepstrum))
	band_indices = (quefrency > 1.0 / freq_hi) & (quefrency <= 1.0 / freq_lo)
	collection   = collections.BrokenBarHCollection.span_where(
		  quefrency
		, ymin=0
		, ymax=np.abs(cepstrum).max()
		, where=band_indices
		, facecolor='green'
		, alpha=0.5
		, label='valid pitches'
	)
	ax.add_collection(collection)
	ax.set_xlabel('quefrency (s)')
	ax.set_title('cepstrum')
	ax.legend()

#
# pitch_tracking_f0
#
def pitch_tracking_f0(signal, fs, freq_lo=30.0, freq_hi=80.0, window=True, threshold=0.0):
	""" extrating fundamental from input signal """
	_, _, quefrency, cepstrum = pitch_tracking_cepstrum(signal, fs, window=window, threshold=threshold)
	band_indices              = (quefrency > 1.0 / freq_hi) & (quefrency <= 1.0 / freq_lo)
	quefrency_max             = np.argmax(np.abs(cepstrum)[band_indices])
	f0                        = 1.0 / quefrency[band_indices][quefrency_max]
	return quefrency, cepstrum, f0

#
# pitch_tracking_best_f0
#
def pitch_tracking_best_f0(x, y, z, fs, freq_lo=30.0, freq_hi=80.0, window=True, threshold=0.0):
	""" extrating fundamental from input axis having the most energy """
	quefrency, cepstrum = pitch_tracking_best_cepstrum(x, y, z, fs, window=window, threshold=threshold)
	band_indices        = (quefrency > 1.0 / freq_hi) & (quefrency <= 1.0 / freq_lo)
	quefrency_max       = np.argmax(np.abs(cepstrum)[band_indices])
	f0                  = 1.0 / quefrency[band_indices][quefrency_max]
	return quefrency, cepstrum, f0

# 	EOF