Support sparse

autocorr/multitau.py

@@ -25,14 +25,14 @@ def even(x):
         return x if x % 2 == 0 else x - 1
 
     # 1-D data hack
-    if len(signal.shape) == 1:
+    if signal.ndim == 1:
         N = even(signal.shape[0])
-        a = np.array(signal[np.newaxis, :N], copy=True)
-    elif len(signal.shape) == 2:
+        a = signal[np.newaxis, :N]
+    elif signal.ndim == 2:
         # copy data a local array
         N = even(signal.shape[1])
-        a = np.array(signal[:, :N], copy=True)
-    elif len(signal.shape) > 2:
+        a = signal[:, :N]
+    elif signal.ndim > 2:
         raise ValueError('Flatten the [2,3,..] dimensions before passing to autocorrelate.')
 
     if N < lags_per_level:

autocorr/multitau_sparse.py

@@ -0,0 +1,77 @@
+import numpy as np
+import sparse
+
+# if N is add subtract 1
+def even(x):
+    return x if x % 2 == 0 else x - 1
+
+
+def multitau(signal, lags_per_level=16):
+    """Autocorrelation of a signal using multi-tau method.
+
+    For details, please refer to D. Magatti and F. Ferri doi:10.1364/ao.40.004011
+
+    Parameters
+    ----------
+    signal: 2-D array
+        input signal, autocorrlation is calulated along `1-axis`
+    lags_per_level: integer, optional
+        number of lag-times per level, 16 is a as good number as any other
+
+    Returns:
+    --------
+    autocorrelations: numpy.ndarray
+        should be self-explanatory
+    lag-times: numpy.ndarray
+        lag times in log2, corresponding to autocorrelation values
+    """
+
+    # 1-D data hack
+    if signal.ndim == 1:
+        N = even(signal.shape[0])
+        a = signal[np.newaxis, :N]
+    elif signal.ndim == 2:
+        # copy data a local array
+        N = even(signal.shape[1])
+        a = signal[:, :N]
+    elif signal.ndim > 2:
+        raise ValueError('Flatten the [2,3,..] dimensions before passing to autocorrelate.')
+
+    if N < lags_per_level:
+        raise ValueError('Lag times per level must be greater than length of signal.')
+
+    #  shorthand for long names
+    m = lags_per_level
+
+    # calculate levels
+    levels = np.int_(np.log2(N / m)) + 1
+    dims = (a.shape[0], (levels + 1) * (m // 2))
+    g2 = [sparse.zeros(dims[0], dtype=np.float32) for _ in range(dims[1])]
+    tau = np.zeros(dims[1], dtype=np.float32)
+
+    # zero level
+    delta_t = 1
+    for i in range(m):
+        tau[i] = i * delta_t
+        t1 = np.mean(a[:, :N - i], axis=1)
+        t2 = np.mean(a[:, i:], axis=1)
+        g2[i] = np.mean(a[:, :N - i] * a[:, i:], axis=1) / t1 / t2
+    a = (a[:, :N:2] + a[:, 1:N:2]) / 2
+    N = even(N // 2)
+
+    for level in range(1, levels):
+        delta_t *= 2
+        for n in range(m // 2):
+            idx = m + (level - 1) * (m // 2) + n
+            shift = m // 2 + n
+            tau[idx] = tau[idx - 1] + delta_t
+            t1 = np.mean(a[:, :-shift], axis=1)
+            t2 = np.mean(a[:, shift:], axis=1)
+            g2[idx] = np.mean(a[:, :- shift] * a[:, shift:], axis=1) / t1 / t2
+        a = (a[:, :N:2] + a[:, 1:N:2]) / 2
+        N = even(N // 2)
+        if N < lags_per_level:
+            break
+
+    g2 = sparse.stack(g2)
+    return g2, tau

autocorr/tests/test_multitau.py

@@ -1,36 +1,46 @@
 import numpy as np
+from numpy.testing import assert_array_equal
 import autocorr
 import time
 
 N = 10240
 np.random.seed(0)
 t = np.arange(N)
 A = np.exp(-0.05 * t)[:, np.newaxis] + np.random.rand(N, 24) * 0.1
+original = A.copy()
 
 
 def test_multitau():
     t0 = time.time()
     g1, tau = autocorr.multitau(A.T, 16)
     t1 = time.time()
+    # Check that the input array has not been mutated.
+    assert_array_equal(A, original)
     print('pure python time = %f' % (t1 - t0))
 
 
 def test_multitau_mt():
     t0 = time.time()
     g1, tau = autocorr.multitau_mt(A.T, 16)
     t1 = time.time()
+    # Check that the input array has not been mutated.
+    assert_array_equal(A, original)
     print('accelerated version = %f' % (t1 - t0))
 
 
 def test_fftautocorr():
     t0 = time.time()
     g2, tau = autocorr.fftautocorr(A.T)
     t1 = time.time()
+    # Check that the input array has not been mutated.
+    assert_array_equal(A, original)
     print('fft version = %f' % (t1 - t0))
 
 
 def test_cfftautocorr():
     t0 = time.time()
     g2, tau = autocorr.fftautocorr_mt(A.T)
     t1 = time.time()
+    # Check that the input array has not been mutated.
+    assert_array_equal(A, original)
     print('fft version = %f' % (t1 - t0))

examples/sparse_memory_usage.py

@@ -0,0 +1,37 @@
+import numpy as np
+import sparse
+import autocorr
+import autocorr.multitau_sparse
+import os
+import psutil
+
+
+process = psutil.Process(os.getpid())
+
+def sizeof_fmt(num, suffix='B'):
+    # https://stackoverflow.com/a/1094933/1221924
+    for unit in ['','Ki','Mi','Gi','Ti','Pi','Ei','Zi']:
+        if abs(num) < 1024.0:
+            return "%3.1f%s%s" % (num, unit, suffix)
+        num /= 1024.0
+    return "%.1f%s%s" % (num, 'Yi', suffix)
+
+def usage():
+    return sizeof_fmt(process.memory_info().rss)
+
+print(f"Memory usage: {usage()} (baseline)")
+
+N = 2048 ** 2
+np.random.seed(0)
+t = np.arange(N)
+a = np.exp(-0.05 * t)[:, np.newaxis] + np.random.rand(N, 24) * 0.1
+DENSE_FRACTION = 0.01
+zero_mask = np.random.rand(*a.shape) > DENSE_FRACTION
+a[zero_mask] = 0
+result = autocorr.multitau(a)
+print(f"Memory usage: {usage()} (dense)")
+del result
+s = sparse.COO.from_numpy(a)
+del a
+autocorr.multitau_sparse.multitau(s)
+print(f"Memory usage: {usage()} (sparse)")