selcal/scripts/selcald/correlate.py

85 lines
2.2 KiB
Python

# Run with "ipython -i --matplotlib=qt correlate.py"
#
from __future__ import print_function
import numpy as np
import pandas as pd
from scipy import signal
import matplotlib.pyplot as plt
from scipy.io.wavfile import write
RATE = 44100
Alpha = 312.6
# See http://www.dsprelated.com/showarticle/908.php
def voss(nrows, ncols=16):
"""Generates pink noise using the Voss-McCartney algorithm.
nrows: number of values to generate
rcols: number of random sources to add
returns: NumPy array
"""
array = np.empty((nrows, ncols))
array.fill(np.nan)
array[0, :] = np.random.random(ncols)
array[:, 0] = np.random.random(nrows)
# the total number of changes is nrows
n = nrows
cols = np.random.geometric(0.5, n)
cols[cols >= ncols] = 0
rows = np.random.randint(nrows, size=n)
array[rows, cols] = np.random.random(n)
df = pd.DataFrame(array)
df.fillna(method='ffill', axis=0, inplace=True)
total = df.sum(axis=1)
return (total.values - (ncols/2))
# tone synthesis
def note(freq, len, amp=1, rate=RATE):
t = np.linspace(0, len, int(len * rate))
if freq == 0:
data = np.zeros(int(len * rate))
else:
data = np.sin(2*np.pi*freq*t)*amp
return data
sig = note(Alpha, 0.10)
sig_tx = np.concatenate((note(0, 0.50),
note(Alpha, 0.20),
note(0, 0.10),
note(Alpha, 0.75),
note(0, 0.10),
note(Alpha, 0.10),
note(0, 0.10),))
# sig_noise = 0.33 * sig_tx + 0.66 * np.random.randn(len(sig_tx))
sig_noise = 0.50 * sig_tx + 0.50 * voss(len(sig_tx))
corr = np.abs(signal.correlate(sig_noise, sig, mode='same'))
fig, (ax_orig, ax_sig, ax_noise, ax_corr) = plt.subplots(4, 1, sharex=True)
ax_orig.plot(sig)
ax_orig.set_title('Reference')
ax_sig.plot(sig_tx)
ax_sig.set_title('Signal without noise')
ax_noise.plot(sig_noise)
ax_noise.set_title('Signal with noise')
ax_corr.plot(corr)
ax_corr.axhline(np.average(corr), ls=':')
ax_corr.set_title('Cross-correlated with Reference')
ax_orig.margins(0, 0.1)
fig.set_tight_layout(True)
plt.show()
write('test.wav', 44100, sig_noise)