view utils/acorr.py @ 19:890cfe424f4a tip

added annotations

#!/usr/bin/env python
# encoding: utf-8
"""
acorr.py

This function provides similar fuctionality to Matlab's xcorr() but for computing autocorrelation only.

It gives equivalent results given the following Matlab test script:

%==== MATLAB ====
a=[1,2,3,4,5,1,2,3,4,5,2,3,4,5,6,1,2,3,4,5];
[xcr,lag] = xcorr(a,a,20,'unbiased');
disp(xcr);
plot(lag,xcr);
%==== MATLAB ====

Created by George Fazekas on 2014-02-26.
Copyright (c) 2014 . All rights reserved.
"""

import sys,os
import numpy as np
from scipy.fftpack import fft, ifft

class ACORR(object):

	def nextpow2(self, n):
		'''Return the next power of 2 such as 2^p >= n'''
		if np.any(n < 0):
			raise ValueError("n should be > 0")
		if np.isscalar(n):
			f, p = np.frexp(n)
			if f == 0.5:
				return p-1
			elif np.isfinite(f):
				return p
			else:
				return f
		else:
			f, p = np.frexp(n)
			res = f
			bet = np.isfinite(f)
			exa = (f == 0.5)
			res[bet] = p[bet]
			res[exa] = p[exa] - 1
			return res

	def acorr_ifft_fft(self, x, nfft, lag, onesided=False, scale='none'):
		'''Compute the actual autocorrelation via IFFT/FFT	'''
		ra = np.real(ifft(np.abs(fft(x, n=nfft) ** 2)))
		if onesided:
			b = ra[..., :lag]
		else:
			b = np.concatenate([ra[..., nfft-lag+1:nfft], ra[..., :lag]], axis=-1)
		#print b, ra[..., 0][..., np.newaxis], b / ra[..., 0][..., np.newaxis]
		if scale == 'coeff':
			return b / ra[..., 0][..., np.newaxis]
		elif scale == 'unbiased':
			'''scale = len(x)-abs(lags)'''
			lags = np.array(range(-lag+1,lag))
			# print lags,len(x)
			scale = len(x) - abs(lags)
			scale[scale<=0] = 1.0
			# print scale
			return b/scale
		else:
			return b

	def acorr(self, x, axis=-1, onesided=False, lag=None, scale='none'):
		'''Compute autocorrelation of x along given axis.
			x : array-like
				signal to correlate.
			axis : int
				axis along which autocorrelation is computed.
			onesided: bool, optional
				if True, only returns the right side of the autocorrelation.
			scale: {'none', 'coeff'} scaling mode. 
			If 'coeff', the correlation is normalised such as the 0-lag is equal to 1.
			if 'unbiased' the correlation is normalised using len(x) -	abs(lags) '''

		if not scale in ['none', 'coeff','unbiased']:
			raise ValueError("scale mode %s not understood" % scale)
		if not lag :
			lag = x.shape[axis]
		lag += 1
		nfft = 2 ** self.nextpow2(2 * lag - 1)
		lags = np.array(range(-lag+1, lag))
		if axis != -1:
			x = np.swapaxes(x, -1, axis)
		a = self.acorr_ifft_fft(x, nfft, lag, onesided, scale)
		if axis != -1:
			a = np.swapaxes(a, -1, axis)
		return a,lags