Added Python utilities (calculus, spectrogram)

o calculus.deriv2D   - 2D gradient using convolution

o calculus.simpson_integrate    - Integrate 2D data using convolution

o spectrogram.spectrogram  - Creates spectrograms using the Gabor transform

tools/pylib/boututils/calculus.py

@@ -6,13 +6,14 @@
 """
 
 try:
-    from numpy import zeros, arange, pi
+
+    from numpy import zeros, arange, pi, ones, array, transpose, sum, where, arange, multiply
     from numpy.fft import rfft, irfft
+    from scipy.signal import convolve
 except ImportError:
     print "ERROR: NumPy module not available"
     raise
 
-
 def deriv(*args, **kwargs):
     """Take derivative of 1D array
     
@@ -71,6 +72,65 @@ def deriv(*args, **kwargs):
             result[0] = 0.0
         return result
 
+def deriv2D(data,axis=-1,dx=1.0,noise_suppression=True):
+  """ Takes 1D or 2D Derivative of 2D array using convolution
+	
+	result = deriv2D(data)
+	result = deriv2D(data, dx)
+	
+	output is 2D (if only one axis specified)
+	output is 3D if no axis specified [nx,ny,2] with the third dimension being [dfdx, dfdy]
+	
+	keywords:
+	axis = 0/1  If no axis specified 2D derivative will be returned
+	dx = 1.0    axis spacing, must be 2D if 2D deriv is taken - default is [1.0,1.0]
+	noise_suppression = True   noise suppressing coefficients used to take derivative - default = True
+  """
+  s = data.shape
+  if axis > len(s)-1:
+    raise RuntimeError("ERROR: axis out of bounds for derivative")
+
+  if noise_suppression:  
+    if s[axis] < 11:
+      raise RuntimeError("Data too small to use 11th order method")
+    tmp = array([-1.0/512.0,-8.0/512.0,-27.0/512.0,-48.0/512.0,-42.0/512.0,0.0,42.0/512.0,48.0/512.0,27.0/512.0,8.0/512.0,1.0/512.0])
+  else:
+    if s[axis] < 9:
+      raise RuntimeError("Data too small to use 9th order method")
+    tmp = array([1.0/280.0,-4.0/105.0,1.0/5.0,-4.0/5.0,0.0,4.0/5.0,-1.0/5.0,4.0/105.0,-1.0/280.0])       
+
+  N = (tmp.size-1)/2
+  if axis==1:
+    W = transpose(tmp[:,None])
+    data_deriv = convolve(data,W,mode='same')/dx*-1.0
+    for i in range(s[0]):
+      data_deriv[i,0:N-1] = deriv(data[i,0:N-1])/dx
+      data_deriv[i,s[1]-N:] = deriv(data[i,s[1]-N:])/dx
+
+  elif axis==0:
+    W = tmp[:,None]
+    data_deriv = convolve(data,W,mode='same')/dx*-1.0
+    for i in range(s[1]):
+      data_deriv[0:N-1,i] = deriv(data[0:N-1,i])/dx
+      data_deriv[s[0]-N:,i] = deriv(data[s[0]-N:,i])/dx
+  else:
+    data_deriv = zeros((s[0],s[1],2))
+    if len(dx)==1:
+      dx = array([dx,dx])
+
+    W = tmp[:,None]#transpose(multiply(tmp,ones((s[1],tmp.size))))
+    data_deriv[:,:,0] = convolve(data,W,mode='same')/dx[0]*-1.0
+    for i in range(s[1]):
+      data_deriv[0:N-1,i,0]  =  deriv(data[0:N-1,i])/dx[0]
+      data_deriv[s[0]-N:s[0]+1,i,0] = deriv(data[s[0]-N:s[0]+1,i])/dx[0]
+
+    W = transpose(tmp[:,None])#multiply(tmp,ones((s[0],tmp.size)))
+    data_deriv[:,:,1] = convolve(data,W,mode='same')/dx[1]*-1.0
+    for i in range(s[0]):
+      data_deriv[i,0:N-1,1] = deriv(data[i,0:N-1])/dx[1]
+      data_deriv[i,s[1]-N:s[1]+1,1] = deriv(data[i,s[1]-N:s[1]+1])/dx[1]
+
+  return data_deriv
 
 def integrate(var, periodic=False):
     """Integrate a 1D array
@@ -132,3 +192,54 @@ def int_total(f):
             result[i] = result[-1] - int_total(var[i:])
         return result
 
+def simpson_integrate(data,dx,dy,kernel=0.0,weight=1.0):
+  """ Integrates 2D data to one value using the simpson method and matrix convolution
+
+      result = simpson_integrate(data,dx,dy)
+      
+      keywords:
+      
+      kernel - can be supplied if the simpson matrix is calculated ahead of time
+	    - if not supplied, is calculated within this function
+	    - if you need to integrate the same shape data over and over, calculated
+		it ahead of time using:
+		  kernel = simpson_matrix(Nx,Ny,dx,dy)
+
+      weight - can be used to scale data if single number
+	    - can be used to mask data if weight is array (same size as data)
+  """
+  s = data.shape
+  Nx = s[0]
+  Ny = s[1]
+
+  if len(kernel)==1:
+    kernel = simpson_matrix(Nx,Ny,dx,dy)
+
+  return sum(multiply(multiply(weight,kernel),data))/sum(multiply(weight,kernel))
+
+
+def simpson_matrix(Nx,Ny,dx,dy):
+  """
+      Creates a 2D matrix of coefficients for the simpson_integrate function
+      
+      Call ahead of time if you need to perform integration of the same size data with the
+	same dx and dy
+      
+      Otherwise, simpson_integrate will automatically call this
+
+  """
+  Wx = arange(Nx) + 2
+  Wx[where(arange(Nx) % 2 == 1)] = 4
+  Wx[0] = 1
+  Wx[Nx-1] = 1
+
+  Wy = arange(Ny) + 2
+  Wy[where(arange(Ny) % 2 == 1)] = 4
+  Wy[0] = 1
+  Wy[Ny-1] = 1
+
+  W = Wy[None,:] * Wx[:,None]
+
+  A = dx*dy/9.0
+
+  return W*A

tools/pylib/boututils/spectrogram.py

@@ -0,0 +1,132 @@
+"""
+    Creates spectrograms using the Gabor transform to maintain time and frequency resolution
+    
+    (spectrogram, frequency) = spectrogram(data, dt, sigma)
+    
+    three arguments are required:
+    
+	data  - the time series you want spectrogrammed
+	dt    - time resolution
+	sigma - used in the Gabor transform, will balance time and frequency resolution
+		  suggested value is 1.0, but may need to be adjusted manually
+		  until result is as desired
+		  
+		IF bands are too tall raise sigma
+		IF bands are too wide, lower sigma
+		  
+    optional keyword:
+    
+	clip = 1.0 - optional keyword, makes the spectrogram run faster, but decreases frequency resolution
+		      clip is by what factor the time spectrum should be clipped by --> N_new = N / clip
+	
+	
+    NOTE: Very early and very late times will have some issues due to the method - truncate them after
+	    taking the spectrogram if they are below your required standards
+	  
+    NOTE: If you are seeing issues at the top or bottom of the frequency range, you need a longer time series
+    
+    written by: Jarrod Leddy
+    updated:    4/10/2013
+
+"""
+
+try:
+
+    from numpy import arange,zeros,exp,power,transpose,sin,cos,linspace,min,max
+    from scipy import fftpack,pi
+    from scipy import pi
+
+except ImportError:
+    print "ERROR: NumPy or SciPy module not available"
+    raise
+
+def spectrogram(data, dx, sigma, clip=1.0):
+    n = data.size
+    xx = arange(n)*dx
+    sigma = sigma * 0.01 * dx
+
+    if clip: # clip to 6 std dev on either side for speed
+
+	n_clipped = int(n/clip)
+
+	spectra = zeros((n,n_clipped/2))
+
+	omega = fftpack.fftfreq(n_clipped, dx)
+	omega = omega[0:n_clipped/2]
+
+	for i in range(n):
+	    beg = i-n_clipped/2
+	    end = i+n_clipped/2-1
+	    if beg < 0:
+		end = end-beg
+		beg = 0
+	    elif end >= n:
+	        end = n-1
+	        beg = end - n_clipped + 1 
+	    gaussian = 1.0 / (sigma * 2.0 * pi) * exp(-0.5 * power(( xx[beg:end] - xx[i] ),2.0) / (2.0 * sigma) )
+	    fftt = abs(fftpack.fft(data[beg:end] * gaussian))
+	    fftt = fftt[:n_clipped/2]
+	    spectra[i,:] = fftt
+
+    else:
+
+	spectra = zeros((n,n/2))
+
+	omega = fftpack.fftfreq(n, dx)
+	omega = omega[0:n/2]
+
+	for i in range(n):
+	    gaussian = 1.0 / (sigma * 2.0 * pi) * exp(-0.5 * power(( xx - xx[i] ),2.0) / (2.0 * sigma) )
+	    fftt = abs(fftpack.fft(data * gaussian))
+	    fftt = fftt[:n/2]
+	    spectra[i,:] = fftt
+
+    return (transpose(spectra), omega)
+
+def test_spectrogram(n, d, s):
+
+  try:
+    import matplotlib.pyplot as plt
+  except ImportError:
+    print "ERROR: MatPlotLib module not available"
+    raise
+  """
+    Function used to test the performance of spectrogram with various values of sigma
+  """
+  xx = arange(n)/d
+  test_data = sin(2.0*pi*512.0*xx * ( 1.0 + 0.005*cos(xx*50.0))) + 0.5*exp(xx)*cos(2.0*pi*100.0*power(xx,2))
+  test_sigma = s
+  dx = 1.0/d
+
+  s1 = test_sigma*0.1
+  s2 = test_sigma
+  s3 = test_sigma*10.0
+
+  (spec2,omega2) = spectrogram(test_data, dx, s2, clip=1.0)
+  (spec3,omega3) = spectrogram(test_data, dx, s3)
+  (spec1,omega1) = spectrogram(test_data, dx, s1, clip=1.0)
+
+  levels = linspace(min(spec1),max(spec1),100)
+  plt.subplot(311)
+  plt.contourf(xx,omega1,spec1,levels=levels)
+  plt.ylabel("frequency")
+  plt.xlabel(r"$t$")
+  plt.title(r"Spectrogram of $sin(t + cos(t) )$ with $\sigma=$%3.1f"%s1)
+
+  levels = linspace(min(spec2),max(spec2),100)
+  plt.subplot(312)
+  plt.contourf(xx,omega2,spec2,levels=levels)
+  plt.ylabel("frequency")
+  plt.xlabel(r"$t$")
+  plt.title(r"Spectrogram of $sin(t + cos(t) )$ with $\sigma=$%3.1f"%s2)
+
+  levels = linspace(min(spec3),max(spec3),100)
+  plt.subplot(313)
+  plt.contourf(xx,omega3,spec3,levels=levels)
+  plt.ylabel("frequency")
+  plt.xlabel(r"$t$")
+  plt.title(r"Spectrogram of $sin(t + cos(t) )$ with $\sigma=$%3.1f"%s3)
+  plt.show()
+
+if __name__ == "__main__":
+  test_spectrogram(2048, 2048.0, 1.0) # array size, divisions per unit, sigma of gaussian