Simple example of Wiener deconvolution in Python

wiener_deconvolution_example.py

#!/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()

This code is merely a simple, minimal example of Wiener filtering. You could certainly use it as a template from which to make other things, but that's not the goal here, and I'm afraid here is not the right place to ask.

Hi, I am certainly trying to hijack your code and you use it for a completely different purpose compared to what it was designed initially... but, I would like to perform a wiener deconvolution of a signal made of several Gaussians that are simply convolved with an ellipsoide profile. In this case though, the algorithm seems to completely fail the reconstruction, do you have any idea what could be wrong ? Is there any assumptions behind the shape of the signals (length, positivity,...) ?