diff --git a/docs/overlap_add_math.ipynb b/docs/overlap_add_math.ipynb new file mode 100644 index 0000000..346a29e --- /dev/null +++ b/docs/overlap_add_math.ipynb @@ -0,0 +1,222 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "e984e902-bf76-43ae-a5d5-55e59357b6bf", + "metadata": {}, + "source": [ + "# From Linear Convolution to Overlap-Add Method!" + ] + }, + { + "cell_type": "markdown", + "id": "e9725572-a39b-4fb5-992c-5f3d22a80797", + "metadata": {}, + "source": [ + "This notebook explains the math behind the overlap-add method in convolution. The content is based on the book \"Discrete-Time Signal Processing\" by Alan V. Oppenheim and Ronald W. Schafer. The book is [publicly available in MIT OpenCourseWare](https://ocw.mit.edu/courses/res-6-dtsp-discrete-time-signal-processing/resources/mitres_6-dtsp_s26_thirdedition_pdf/)" + ] + }, + { + "cell_type": "markdown", + "id": "03f9d2f8-1560-4a5d-93a2-ec82fb422f6f", + "metadata": {}, + "source": [ + "Most textbooks use $x$ and $h$ notation in the convolution. However, in this notebook, we would like to use $Q$ and $T$ as those are the name of arrays in the realm of STUMPY. We would also like to introduce $Q'$, which is simply the reverse of $Q$." + ] + }, + { + "cell_type": "markdown", + "id": "3dda604a-3a3e-42f4-90c7-7b01ada16337", + "metadata": {}, + "source": [ + "# Linear Convolution of Two Finite-Length Sequences" + ] + }, + { + "cell_type": "markdown", + "id": "9ee474bd-1229-4749-80fe-f90c4a135931", + "metadata": {}, + "source": [ + "Consider two finte-length sequences $Q'$ (of length $m$) and $T$ (of length $n$). Their linear convolution, $C$, can be computed as follows:" + ] + }, + { + "cell_type": "markdown", + "id": "ad4a63b5-2d91-4875-b4c2-bc7ce11b6af9", + "metadata": {}, + "source": [ + "$$ C[idx] = \\sum_{i=-\\infty}^{\\infty}{T[i] \\times Q'[idx-i]} $$" + ] + }, + { + "cell_type": "markdown", + "id": "8028ee30-3244-4c79-86ca-22e2b8876694", + "metadata": {}, + "source": [ + "where $T[.]$ and $Q'[.]$ are zeros for any index that is outside of their range. So, $T[i]$ is zero if $i$ is outside of the range $0 \\le i \\le n-1$. Hence, the equation above is equivalent to:\n", + "\n", + "$$ C[idx] = \\sum_{i=0}^{n-1}{T[i] \\times Q'[idx-i]} $$" + ] + }, + { + "cell_type": "markdown", + "id": "593a8276-49cb-43d2-9300-8d29e67d3104", + "metadata": {}, + "source": [ + "Note that $Q'[idx-i]$ is not zero when $0 \\le idx-i \\le m-1$, or equivalently $i \\le idx \\le i+m-1$. \n", + "\n", + "Therefore:\n", + "\n", + "* $idx \\ge i, \\quad i \\ge 0 \\implies idx \\ge 0$\n", + "* $idx \\le i+m-1, \\quad i \\le n-1 \\implies idx \\le n+m-2$" + ] + }, + { + "cell_type": "markdown", + "id": "e638d3d7-cf2a-4355-aff3-a9192400defb", + "metadata": {}, + "source": [ + "This shows that the linear convolution $C$ has values at indices $\\set{0, 1, ..., n + m - 2}$, and it is 0 otherwise. Therefore, the length of output in linear convolution is $n+m-1$. \n", + "\n", + "Furthermore, let's check out the values for different indices:\n", + "* $idx=0 \\implies C[0]=T[0]Q'[0]$\n", + "* $idx=1 \\implies C[0]=T[0]Q'[1] + T[1]Q'[0]$\n", + "* ...\n", + "* $idx=m-1 \\implies C[m-1]=T[0]Q'[m-1] + T[1]Q'[m-2] + ... + T[m-1]Q'[0] = T_{0}.Q$\n", + "* $idx=m \\implies C[m]=T[1]Q'[m-1] + T[2]Q'[m-2] + ... + T[m]Q'[0] = T_{1}.Q$\n", + "* ...\n", + "* $idx=n-1 \\implies C[n-1]=T[n-m]Q'[(n-1)-(n-m)] + T[n-m+1]Q'[(n-1)-(n-m+1)] + ... + T[n-1]Q'[0] = T_{n-m+1}.Q$\n", + "* ...\n", + "* $idx=n+m-2 \\implies C[n+m-2]=T[n-1]Q'[(n+m-2)-(n-1)] = T[n-1]Q'[m-1]$\n", + "\n", + "As observed, when $m-1 \\le idx \\le n-1$, $C[idx]$ becomes the dot product between a subsequnce of $T$ and $Q$, which is the reverse of $Q'$. Therefore, the `range(m-1,n)` gives the sliding dot product between $Q$ and $T$." + ] + }, + { + "cell_type": "markdown", + "id": "48c7f9d8-1b7c-4abc-886b-27a015b78ed6", + "metadata": {}, + "source": [ + "**How can we leverage this to speed up the computation of sliding dot product?**" + ] + }, + { + "cell_type": "markdown", + "id": "52bcf566-3ac0-483d-81a3-57fa3b5766c3", + "metadata": {}, + "source": [ + "# Option I: Convert to Circular Convolution and use FFT-IFFT" + ] + }, + { + "cell_type": "markdown", + "id": "2c4176e2-60f3-410f-ad2d-6c1d29cca6cf", + "metadata": {}, + "source": [ + "Circular convolution is defined between two sequences that are both periodic and their period are the same, say $N$. Their circular convolution is also a periodic sequence, with period $N$, and it can be computed as follows:\n", + "\n", + "$$C_{N}[idx] = \\sum_{i=0}^{N-1} \\tilde{Q'}[i] \\times \\tilde{T}[(idx-i)_{N}], \\quad 0 \\le idx \\le N-1$$\n", + "\n", + "where, $\\tilde{Q'}$ and $\\tilde{T}$ are both periodic sequence with period $N$. $C_{N}$ represents one period of N-Circular convolution. This can also be computed via FFT, i.e. $$C_{N} = IFFT( FFT(\\tilde{Q'}_{N}) \\times FFT(\\tilde{T}_{N}) ) $$" + ] + }, + { + "cell_type": "markdown", + "id": "e10c5304-a1a1-4cbd-9b26-266e229bdc57", + "metadata": {}, + "source": [ + "If there is a way to compute the linear convolution via circular convolution, then we can take advantage of FFT by using eq (4). The good news is that there is a way! The linear convolution between $Q'$ and $T$ can be obtained by performing N-circular convolution between $\\tilde{Q'}$ and $\\tilde{T}$, where:\n", + "\n", + "* $N \\ge n + m - 2$\n", + "* $\\tilde{Q'}_{N}$ is $Q'$ but zero-padded with $N-m$ zeros\n", + "* $\\tilde{T}_{N}$ is $T$ but zero-padded with $N-n$ zeros" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "b6529137-fcdc-4d65-aa93-da3746739f94", + "metadata": {}, + "outputs": [], + "source": [ + "# WIP" + ] + }, + { + "cell_type": "markdown", + "id": "97fcc328-35f9-4b66-9598-1b5d8f1f97ac", + "metadata": {}, + "source": [ + "# Option II: Overlap-add method\n", + "This is a divide and conquer algorithm. It applies divide on \"linear convolution\" and conquer each via circular convolution (See: Option I)" + ] + }, + { + "cell_type": "markdown", + "id": "1e45ccf7-1564-4c41-b32e-4de9b1b201a3", + "metadata": {}, + "source": [ + "## Linearity in Linear Convolution" + ] + }, + { + "cell_type": "markdown", + "id": "a084578e-cb20-437c-a70c-3e33f8da8ac3", + "metadata": {}, + "source": [ + "Suppose the array $T$ can be written as $T_{1} + T_{2}$. In other words: $T[i] = T_{1}[i] + T_{2}[i]$, then:\n", + "\n", + "$$ C[idx] = \\sum_{i=-\\infty}^{\\infty}{T[i] \\times Q'[idx-i]} $$\n", + "\n", + "$$ C[idx] = \\sum_{i=-\\infty}^{\\infty}{(T_{1}[i] \\times Q'[idx-i] + T_{2}[i] \\times Q'[idx-i])} $$\n", + "\n", + "$$ C[idx] = \\sum_{i=-\\infty}^{\\infty}{T_{1}[i] \\times Q'[idx-i]} + \\sum_{i=-\\infty}^{\\infty}{T_{2}[i] \\times Q'[idx-i]}$$\n", + "\n", + "$$ C[idx] = C_{1}[idx] + C_{2}[idx] $$" + ] + }, + { + "cell_type": "markdown", + "id": "1f5be13b-82a7-46d9-9cb1-229c878e67f8", + "metadata": {}, + "source": [ + "This shows that a linear convolution has the property of \"linearity\". Now, we use this property to show that we can compute the linear convolution between $Q'$ and long $T$ by breaking $T$ into smaller parts." + ] + }, + { + "cell_type": "markdown", + "id": "fecee5fc-3c10-4c2e-8c66-bce17c3650bd", + "metadata": {}, + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "acf5516d-e52d-457c-a8c5-f4d1b7ae8f1c", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.14.4" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/sdp/challenger_sdp.py b/sdp/challenger_sdp.py index ce60e8d..938dc8e 100644 --- a/sdp/challenger_sdp.py +++ b/sdp/challenger_sdp.py @@ -1,14 +1,225 @@ +import math + import numpy as np +from scipy.fft import next_fast_len +from scipy.special import lambertw -def setup(Q, T): - return +from sdp import pocketfft_r2c_c2r_sdp + +# _duccfft replaced _pocketfft in scipy 1.18 +try: + from scipy.fft._duccfft.basic import c2r, r2c +except ModuleNotFoundError: # pragma: no cover + from scipy.fft._pocketfft.basic import c2r, r2c + + +def _compute_block_size(m, n, conv_block_size=None): + """ + Return a block size for the overlap-add method. + + Parameters + ---------- + m : int + Length of the query array Q. + + n : int + Length of the time series T. + + conv_block_size : int, default None + Block size for the convolution. When `conv_block_size` is None, + it will be automatically set to an optimal value, internally + computed based on the lengths of Q and T. + + Returns + ------- + conv_block_size : int + Block size for the convolution. Will be at least `m` and at most `n`. + """ + if conv_block_size is None: + # `conv_block_size < n` as, otherwise, there is no + # point in splitting the larger array of length `n` + # `conv_block_size >= 2 * (m-1)` so that + # the vectorized operation can be used later. + # Therefore: `m < n/2 + 1` + + # Note: + # A tighter upper bound can be computed + # by considering the range of values returned + # by the `lambertw(..., k=-1)` function for m>=3 + if m >= n / 2 + 1: + conv_block_size = n + else: + # To minimize Eq. 3 in + # https://en.wikipedia.org/wiki/Overlap–add_method + # ToDo: Revise `opt_size` by considering RFFT/IRFFT + # instead of FFT/IFFT in the computational cost + overlap = m - 1 + opt_size = -overlap * lambertw(-1 / (2 * math.e * overlap), k=-1).real + conv_block_size = next_fast_len(math.ceil(opt_size), real=True) + + # ToDo + # The computed (presumed) optimal `conv_block_size` is based on an approximate + # cost function (See Eq. 3 in https://en.wikipedia.org/wiki/Overlap–add_method). + # However, we should plug the obtained value into a "more accurate" cost function + # and compared it with the cost of regular circular convolution to see + # whether overlap-add should be considered or not. + + # Each chunk of `T` is padded with `m - 1` zeros to form a convolution block. + # Since a chunk (from `T`) must contain at least one element, + # the minimum block size is `m`. However, to take advantage of vectorized + # operation at a later step, the minimum block size is set to `2 * (m-1)` + conv_block_size = max(conv_block_size, 2 * (m - 1)) + + # `conv_block_size < n` as, otherwise, there is no + # point in splitting the larger array of length `n` + return min(conv_block_size, n) + + +def _pocketfft_circular_convolve_block(Q, T, conv_block_size): + m = Q.shape[0] + n = T.shape[0] + + # Each block in overlap-add method needs to be padded + # with `m-1` zeros. Therefore, the effective block size + # for T is `conv_block_size - (m-1)`. + T_block_size = conv_block_size - (m - 1) + n_blocks = math.ceil(n / T_block_size) + last_block_start = (n_blocks - 1) * T_block_size + + # To compute the circular convolution between the zero-padded Q + # and each zero-padded block of T, the data can be loaded into + # a 2D array with `n_blocks + 1` rows, where the first `n_blocks` + # rows correspond to the blocks of T, and the last row is the + # zero-padded Q. + tmp = np.empty((n_blocks + 1, conv_block_size), dtype=np.float64) + tmp[: n_blocks - 1, :T_block_size] = T[:last_block_start].reshape( + n_blocks - 1, T_block_size + ) + tmp[: n_blocks - 1, T_block_size:] = 0.0 + tmp[n_blocks - 1, : n - last_block_start] = T[last_block_start:] + tmp[n_blocks - 1, n - last_block_start :] = 0.0 + + tmp[n_blocks, :m] = Q + tmp[n_blocks, m:] = 0.0 + + fft_2d = r2c(True, tmp, axis=-1) + + return c2r(False, np.multiply(fft_2d[:-1], fft_2d[[-1]]), n=conv_block_size) + + +def _pocketfft_valid_oaconvolve(Q, T, conv_block_size): + """ + Compute the valid convolution between Q and T using the overlap-add method. + This method performs several circular convolutions between Q and blocks of T, + and then combines the results to obtain the valid convolution between Q and T + Parameters + ---------- + Q : numpy.ndarray + Query array or subsequence. -def sliding_dot_product(Q, T): + T : numpy.ndarray + Time series or sequence. + + conv_block_size : int + Block size for the overlap-add method. + The value cannot be less than len(Q). + + Returns + ------- + out : numpy.ndarray + The valid convolution between Q and T. + + Notes + ----- + Each block of the convolution contains part of `T`, padded with `len(Q)-1` + zeros. Therefore, `conv_block_size` must be at least `len(Q)` so that it + can cover at least one element of `T` in each block. However, The current + implementation requires the `conv_block_size` to be at least `2*(len(Q) - 1)` + """ + # performs several circular convolutions between + # zero-padded Q and zero-padded blocks of T + # and returns a 2D array of the results, + # where each row is associated with a block of T + QT_conv_blocks = _pocketfft_circular_convolve_block(Q, T, conv_block_size) + + # The subsequences at the boundaries of the blocks + # are shared between adjacent blocks. + # The following logic is needed to reconstruct + # the valid convolution between Q and T + overlap = len(Q) - 1 + out = QT_conv_blocks[:, :-overlap] + out[1:, :overlap] += QT_conv_blocks[:-1, -overlap:] + + return np.reshape(out, (-1,))[len(Q) - 1 : len(T)] + + +def _valid_convolve(Q, T, conv_block_size=None): + """ + Compute the valid convolution between Q and T + + Parameters + ---------- + Q : numpy.ndarray + Query array or subsequence. + + T : numpy.ndarray + Time series or sequence. + + conv_block_size : int, default None + Block size for the overlap-add method. When `conv_block_size` + is None, it will automatically be set to an optimal value, + internally computed based on the lengths of Q and T. + + Returns + ------- + out : numpy.ndarray + The valid convolution between Q and T. + + Notes + ----- + The valid convolution between ``Q`` and ``T`` is equivalent to + the sliding dot product between Q[::-1] and 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]) + n = len(T) + conv_block_size = _compute_block_size(m, n, conv_block_size=conv_block_size) + if conv_block_size >= n: + out = pocketfft_r2c_c2r_sdp._pocketfft_valid_convolve(Q, T) + else: + out = _pocketfft_valid_oaconvolve(Q, T, conv_block_size) + return out + + +def setup(Q, T): + return + + +def sliding_dot_product(Q, T, conv_block_size=None): + """ + Compute the sliding dot product between Q and T + + Parameters + ---------- + Q : numpy.ndarray + Query array or subsequence. + + T : numpy.ndarray + Time series or sequence. + + conv_block_size : int, default None + Block size for the overlap-add method. When `conv_block_size` + is None, it will automatically be set to an optimal value, + internally computed based on the lengths of Q and T. + + Returns + ------- + out : numpy.ndarray + The sliding dot product between Q and T. + """ + if len(Q) == len(T): + return np.dot(Q, T) + else: + return _valid_convolve(Q[::-1], T, conv_block_size=conv_block_size) diff --git a/sdp/pocketfft_r2c_c2r_sdp.py b/sdp/pocketfft_r2c_c2r_sdp.py index 1f14342..ac320ec 100644 --- a/sdp/pocketfft_r2c_c2r_sdp.py +++ b/sdp/pocketfft_r2c_c2r_sdp.py @@ -1,22 +1,37 @@ import numpy as np from scipy.fft import next_fast_len -from scipy.fft._pocketfft.basic import r2c, c2r - -def setup(Q, T): - return +# _duccfft replaced _pocketfft in scipy 1.18 +try: + from scipy.fft._duccfft.basic import c2r, r2c +except ModuleNotFoundError: # pragma: no cover + from scipy.fft._pocketfft.basic import c2r, r2c -def sliding_dot_product(Q, T): +def _pocketfft_valid_convolve(Q, T): + """ + Compute the valid convolution between ``Q`` and ``T`` + using circular convolution in the frequency domain + """ n = len(T) m = len(Q) next_fast_n = next_fast_len(n, real=True) tmp = np.empty((2, next_fast_n)) - tmp[0, :m] = Q[::-1] + tmp[0, :m] = Q tmp[0, m:] = 0.0 tmp[1, :n] = T tmp[1, n:] = 0.0 fft_2d = r2c(True, tmp, axis=-1) - return c2r(False, np.multiply(fft_2d[0], fft_2d[1]), n=next_fast_n)[m - 1 : n] + return c2r(False, np.multiply(fft_2d[0], fft_2d[1]), n=next_fast_n)[ + len(Q) - 1 : len(T) + ] + + +def setup(Q, T): + return + + +def sliding_dot_product(Q, T): + return _pocketfft_valid_convolve(Q[::-1], T) diff --git a/test.py b/test.py index d1b80e8..3244c3e 100644 --- a/test.py +++ b/test.py @@ -119,7 +119,7 @@ def test_sdp(n_T, remainder, comparator): 97, ] n_Q_power2 = [2, 4, 8, 16, 32, 64] - n_Q_values = n_Q_prime + n_Q_power2 + [n_T] + n_Q_values = n_Q_prime + n_Q_power2 + [n_T - 1, n_T] n_Q_values = sorted(n_Q for n_Q in set(n_Q_values) if n_Q <= n_T) modules = utils.import_sdp_mods() @@ -210,3 +210,18 @@ def test_pyfftw_sdp_max_n(): np.testing.assert_allclose(comp, ref) return + + +def test_oaconvolve_sdp_blocksize(): + from sdp.challenger_sdp import sliding_dot_product + + T = np.random.rand(2**10) + Q = np.random.rand(2**8) + conv_block_size = 2**9 + + comp = sliding_dot_product(Q, T, conv_block_size=conv_block_size) + ref = naive_sliding_dot_product(Q, T) + + np.testing.assert_allclose(comp, ref) + + return