HOME/Articles/

wiener_deconvolution_example

Article Outline

Example Python program wiener_deconvolution_example.py Python version 3.x or newer. To check the Python version use:

python --version

Modules

  • import numpy as np
  • from numpy.fft import fft, ifft, ifftshift
  • import matplotlib
  • import matplotlib.pyplot as plt
  • import matplotlib.cm as cm
  • from matplotlib.backends.backend_pdf import PdfPages

Methods

  • def gen_son(length):
  • def gen_ir(length):
  • def wiener_deconvolution(signal, kernel, lambd):

Code

Python example

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