Forked from danstowell/wiener_deconvolution_example.py
Created
December 25, 2015 21:55
-
-
Save veco30/2b5667e62df8df0f251e to your computer and use it in GitHub Desktop.
Simple example of Wiener deconvolution in Python
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| #!/usr/bin/env python | |
| # Simple example of Wiener deconvolution in Python. | |
| # We use a fixed SNR across all frequencies in this example. | |
| # | |
| # Written 2015 by Dan Stowell. Public domain. | |
| import numpy as np | |
| from numpy.fft import fft, ifft, ifftshift | |
| import matplotlib | |
| #matplotlib.use('PDF') # http://www.astrobetter.com/plotting-to-a-file-in-python/ | |
| import matplotlib.pyplot as plt | |
| import matplotlib.cm as cm | |
| from matplotlib.backends.backend_pdf import PdfPages | |
| plt.rcParams.update({'font.size': 6}) | |
| ########################## | |
| # user config | |
| sonlen = 128 | |
| irlen = 64 | |
| lambd_est = 1e-3 # estimated noise lev | |
| ########################## | |
| def gen_son(length): | |
| "Generate a synthetic un-reverberated 'sound event' template" | |
| # (whitenoise -> integrate -> envelope -> normalise) | |
| son = np.cumsum(np.random.randn(length)) | |
| # apply envelope | |
| attacklen = length / 8 | |
| env = np.hstack((np.linspace(0.1, 1, attacklen), np.linspace(1, 0.1, length - attacklen))) | |
| son *= env | |
| son /= np.sqrt(np.sum(son * son)) | |
| return son | |
| def gen_ir(length): | |
| "Generate a synthetic impulse response" | |
| # First we generate a quietish tail | |
| son = np.random.randn(length) | |
| attacklen = length / 2 | |
| env = np.hstack((np.linspace(0.1, 1, attacklen), np.linspace(1, 0.1, length - attacklen))) | |
| son *= env | |
| son *= 0.05 | |
| # Here we add the "direct" signal | |
| son[0] = 1 | |
| # Now some early reflection spikes | |
| for _ in range(10): | |
| son[ int(length * (np.random.rand()**2)) ] += np.random.randn() * 0.5 | |
| # Normalise and return | |
| son /= np.sqrt(np.sum(son * son)) | |
| return son | |
| def wiener_deconvolution(signal, kernel, lambd): | |
| "lambd is the SNR" | |
| kernel = np.hstack((kernel, np.zeros(len(signal) - len(kernel)))) # zero pad the kernel to same length | |
| H = fft(kernel) | |
| deconvolved = np.real(ifft(fft(signal)*np.conj(H)/(H*np.conj(H) + lambd**2))) | |
| return deconvolved | |
| if __name__ == '__main__': | |
| "simple test: get one soundtype and one impulse response, convolve them, deconvolve them, and check the result (plot it!)" | |
| son = gen_son(sonlen) | |
| ir = gen_ir(irlen) | |
| obs = np.convolve(son, ir, mode='full') | |
| # let's add some noise to the obs | |
| obs += np.random.randn(*obs.shape) * lambd_est | |
| son_est = wiener_deconvolution(obs, ir, lambd=lambd_est)[:sonlen] | |
| ir_est = wiener_deconvolution(obs, son, lambd=lambd_est)[:irlen] | |
| # calc error | |
| son_err = np.sqrt(np.mean((son - son_est) ** 2)) | |
| ir_err = np.sqrt(np.mean((ir - ir_est) ** 2)) | |
| print("single_example_test(): RMS errors son %g, IR %g" % (son_err, ir_err)) | |
| # plot | |
| pdf = PdfPages('wiener_deconvolution_example.pdf') | |
| plt.figure(frameon=False) | |
| # | |
| plt.subplot(3,2,1) | |
| plt.plot(son) | |
| plt.title("son") | |
| plt.subplot(3,2,3) | |
| plt.plot(son_est) | |
| plt.title("son_est") | |
| plt.subplot(3,2,2) | |
| plt.plot(ir) | |
| plt.title("ir") | |
| plt.subplot(3,2,4) | |
| plt.plot(ir_est) | |
| plt.title("ir_est") | |
| plt.subplot(3,1,3) | |
| plt.plot(obs) | |
| plt.title("obs") | |
| # | |
| pdf.savefig() | |
| plt.close() | |
| pdf.close() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment