diff --git a/sdp/challenger_sdp.py b/sdp/challenger_sdp.py index ce60e8d..7ee13af 100644 --- a/sdp/challenger_sdp.py +++ b/sdp/challenger_sdp.py @@ -1,14 +1,78 @@ +import math import numpy as np +from numba import njit -def setup(Q, T): - return +@njit(fastmath=True) +def _rfft_preprocess_TQ(T, Q): + n = len(T) + m = len(Q) + + Q0 = np.empty(n, dtype=np.float64) + Q0[:m] = Q[::-1] + Q0[m:] = 0.0 + + half_n = n // 2 + out = np.empty((2, half_n), dtype=np.complex128) + + m_stop = m // 2 + for i in range(m_stop): + out[0, i] = T[2 * i] + 1j * T[2 * i + 1] + out[1, i] = Q0[2 * i] + 1j * Q0[2 * i + 1] + + for i in range(m_stop, half_n): + out[0, i] = T[2 * i] + 1j * T[2 * i + 1] + out[1, m_stop:] = 0.0 + + return out + + +@njit(fastmath=True) +def _rfft_postprocess_TQ(TQ): + x_T = TQ[0] + x_Q = TQ[1] + + n_x = len(x_T) + half_n_x = n_x // 2 + + F = np.empty(n_x + 1, dtype=np.complex128) + + # 0th element, half_n_x, and n_x + F[0] = (x_T[0].real + x_T[0].imag) * (x_Q[0].real + x_Q[0].imag) + F[half_n_x] = (x_T[half_n_x] * x_Q[half_n_x]).conjugate() + F[n_x] = (x_T[0].real - x_T[0].imag) * (x_Q[0].real - x_Q[0].imag) + + theta0 = math.pi / n_x + factor = math.cos(theta0) - 1j * math.sin(theta0) + w = 0.5j + for k in range(1, half_n_x): + w = w * factor + + val_T = (x_T[k] - x_T[n_x - k].conjugate()) * (0.5 + w) + val_Q = (x_Q[k] - x_Q[n_x - k].conjugate()) * (0.5 + w) + + F[k] = (x_T[k] - val_T) * (x_Q[k] - val_Q) + F[n_x - k] = (x_T[n_x - k] + val_T.conjugate()) * (x_Q[n_x - k] + val_Q.conjugate()) + + return F + + +def _rfft_TQ(T, Q): + TQ = _rfft_preprocess_TQ(T, Q) + np.fft.fft(TQ, axis=1, out=TQ) + return _rfft_postprocess_TQ(TQ) def sliding_dot_product(Q, T): + n = len(T) m = len(Q) - l = T.shape[0] - m + 1 - out = np.empty(l) - for i in range(l): - out[i] = np.dot(Q, T[i : i + m]) - return out + + F = _rfft_TQ(T, Q) + + return np.fft.irfft(F)[m - 1 : n] + + +def setup(Q, T): + return sliding_dot_product(Q, T) + + diff --git a/test.sh b/test.sh index 7b853ba..c770552 100755 --- a/test.sh +++ b/test.sh @@ -1,5 +1,5 @@ #!/bin/bash - +python -m pytest test_challenger.py rm -rf sdp/__pycache__ ./test.py > timing.csv rm -rf sdp/__pycache__ diff --git a/test_challenger.py b/test_challenger.py new file mode 100644 index 0000000..28d11b2 --- /dev/null +++ b/test_challenger.py @@ -0,0 +1,27 @@ +import numpy as np +import pytest + +from sdp import challenger_sdp + +def naive_sliding_dot_product(Q, T): + m = len(Q) + l = T.shape[0] - m + 1 + out = np.empty(l) + for i in range(l): + out[i] = np.dot(Q, T[i : i + m]) + return out + + +def test_challenger(): + pmin = 2 + pmax = 10 + for q in range(pmin, pmax): + for p in range(q, pmax): + Q = np.random.rand(2 ** q) + T = np.random.rand(2 ** p) + ref = naive_sliding_dot_product(Q, T) + comp = challenger_sdp.sliding_dot_product(Q, T) + + np.testing.assert_allclose(comp, ref) + + return \ No newline at end of file