updated pink noise function (not quite working if ndim >= 3)

pyret/noise.py

@@ -3,11 +3,12 @@
 
 """
 import numpy as np
+from scipy.stats import zscore
 
 __all__ = ['pink']
 
 
-def pink(t, nx, ny, dt=1.0):
+def pink(*args, power=1.):
     """Generates an array of pink noise
 
     Parameters
@@ -22,21 +23,19 @@ def pink(t, nx, ny, dt=1.0):
 
     """
 
-    # white noise
-    noise = np.random.randn(t, nx, ny)
+    # get dimensions from arguments
+    if len(args) == 1 and type(args[0]) is tuple:
+        shape = args[0]
 
-    tc = np.arange(t).reshape(-1, 1, 1).astype('float') * dt
-    xc = np.arange(nx).reshape(1, -1, 1).astype('float')
-    yc = np.arange(ny).reshape(1, 1, -1).astype('float')
+    elif len(args) >= 1:
+        shape = tuple(map(int, args))
 
-    tc -= np.mean(tc)
-    xc -= np.mean(xc)
-    yc -= np.mean(yc)
+    else:
+        return np.array([])
 
-    radii = np.sqrt(tc ** 2 + xc ** 2 + yc ** 2)
-
-    # normalize by 1/freq
-    noise /= np.fft.fftshift(radii)
-
-    # inverse Fourier and normalize
-    return np.real(np.fft.ifftn(noise))
+    axes = [np.linspace(-0.5 * n, 0.5 * n, n) for n in shape]
+    mesh = np.meshgrid(*axes, indexing='ij')
+    dist = np.sqrt(np.sum(list(map(np.square, mesh)), axis=0))
+    spectrum = np.ones(shape) / (dist ** power)
+    fourier_coefficients = spectrum * np.exp(1j * 2 * np.pi * np.random.rand(*shape))
+    return zscore(np.real(np.fft.ifftn(np.fft.fftshift(fourier_coefficients))))