mi@0: #!/usr/bin/env python
mi@0: # encoding: utf-8
mi@0: """
mi@0: acorr.py
mi@0: 
mi@0: This function provides similar fuctionality to Matlab's xcorr() but for computing autocorrelation only.
mi@0: 
mi@0: It gives equivalent results given the following Matlab test script:
mi@0: 
mi@0: %==== MATLAB ====
mi@0: a=[1,2,3,4,5,1,2,3,4,5,2,3,4,5,6,1,2,3,4,5];
mi@0: [xcr,lag] = xcorr(a,a,20,'unbiased');
mi@0: disp(xcr);
mi@0: plot(lag,xcr);
mi@0: %==== MATLAB ====
mi@0: 
mi@0: Created by George Fazekas on 2014-02-26.
mi@0: Copyright (c) 2014 . All rights reserved.
mi@0: """
mi@0: 
mi@0: import sys,os
mi@0: import numpy as np
mi@0: from scipy.fftpack import fft, ifft
mi@0: 
mi@0: class ACORR(object):
mi@0: 
mi@0: 	def nextpow2(self, n):
mi@0: 		'''Return the next power of 2 such as 2^p >= n'''
mi@0: 		if np.any(n < 0):
mi@0: 			raise ValueError("n should be > 0")
mi@0: 		if np.isscalar(n):
mi@0: 			f, p = np.frexp(n)
mi@0: 			if f == 0.5:
mi@0: 				return p-1
mi@0: 			elif np.isfinite(f):
mi@0: 				return p
mi@0: 			else:
mi@0: 				return f
mi@0: 		else:
mi@0: 			f, p = np.frexp(n)
mi@0: 			res = f
mi@0: 			bet = np.isfinite(f)
mi@0: 			exa = (f == 0.5)
mi@0: 			res[bet] = p[bet]
mi@0: 			res[exa] = p[exa] - 1
mi@0: 			return res
mi@0: 
mi@0: 	def acorr_ifft_fft(self, x, nfft, lag, onesided=False, scale='none'):
mi@0: 		'''Compute the actual autocorrelation via IFFT/FFT	'''
mi@0: 		ra = np.real(ifft(np.abs(fft(x, n=nfft) ** 2)))
mi@0: 		if onesided:
mi@0: 			b = ra[..., :lag]
mi@0: 		else:
mi@0: 			b = np.concatenate([ra[..., nfft-lag+1:nfft], ra[..., :lag]], axis=-1)
mi@0: 		#print b, ra[..., 0][..., np.newaxis], b / ra[..., 0][..., np.newaxis]
mi@0: 		if scale == 'coeff':
mi@0: 			return b / ra[..., 0][..., np.newaxis]
mi@0: 		elif scale == 'unbiased':
mi@0: 			'''scale = len(x)-abs(lags)'''
mi@0: 			lags = np.array(range(-lag+1,lag))
mi@0: 			# print lags,len(x)
mi@0: 			scale = len(x) - abs(lags)
mi@0: 			scale[scale<=0] = 1.0
mi@0: 			# print scale
mi@0: 			return b/scale
mi@0: 		else:
mi@0: 			return b
mi@0: 
mi@0: 	def acorr(self, x, axis=-1, onesided=False, lag=None, scale='none'):
mi@0: 		'''Compute autocorrelation of x along given axis.
mi@0: 			x : array-like
mi@0: 				signal to correlate.
mi@0: 			axis : int
mi@0: 				axis along which autocorrelation is computed.
mi@0: 			onesided: bool, optional
mi@0: 				if True, only returns the right side of the autocorrelation.
mi@0: 			scale: {'none', 'coeff'} scaling mode. 
mi@0: 			If 'coeff', the correlation is normalised such as the 0-lag is equal to 1.
mi@0: 			if 'unbiased' the correlation is normalised using len(x) -	abs(lags) '''
mi@0: 
mi@0: 		if not scale in ['none', 'coeff','unbiased']:
mi@0: 			raise ValueError("scale mode %s not understood" % scale)
mi@0: 		if not lag :
mi@0: 			lag = x.shape[axis]
mi@0: 		lag += 1
mi@0: 		nfft = 2 ** self.nextpow2(2 * lag - 1)
mi@0: 		lags = np.array(range(-lag+1, lag))
mi@0: 		if axis != -1:
mi@0: 			x = np.swapaxes(x, -1, axis)
mi@0: 		a = self.acorr_ifft_fft(x, nfft, lag, onesided, scale)
mi@0: 		if axis != -1:
mi@0: 			a = np.swapaxes(a, -1, axis)
mi@0: 		return a,lags
mi@0: