1"""Implicitly padded mixed-radix DFT and linear convolution."""
 2import numpy as np
 3
 4
 5def implicit_dft(x, M, m):
 6    """Positive-sign DFT of x zero-padded to M, without making padded x.
 7
 8    Returns the convention F[k]=sum_j exp(+2pi i*j*k/M)x[j].
 9    M must be divisible by m; x may have arbitrary length <= M.
10    """
11    x = np.asarray(x)
12    L = x.shape[-1]
13    if L > M or M % m:
14        raise ValueError("require len(x)<=M and M divisible by m")
15    p, q = (L + m - 1) // m, M // m
16    # Tiles are bounded and masking is done at load time. This is the only
17    # input-sized temporary; no length-M padded input is allocated.
18    tiles = np.zeros(x.shape[:-1] + (p, m), dtype=np.result_type(x, np.complex128))
19    for t in range(p):
20        lo, hi = t*m, min((t+1)*m, L)
21        tiles[..., t, :hi-lo] = x[..., lo:hi]
22    r = np.arange(q)[:, None]
23    t = np.arange(p)[None, :]
24    s = np.arange(m)[None, :]
25    zq = np.exp(2j*np.pi*r*t/q)
26    zqm = np.exp(2j*np.pi*r*s/(q*m))
27    # U[...,r,s] = sum_t zeta_q^(rt) x[t,s]
28    U = np.einsum('rt,...ts->...rs', zq, tiles) * zqm
29    # positive-sign m DFT in s; output is indexed q*ell+r
30    zm = np.exp(2j*np.pi*np.arange(m)[:, None]*np.arange(m)[None, :]/m)
31    F_rsell = np.einsum('ls,...rs->...rl', zm, U)
32    return np.transpose(F_rsell, tuple(range(F_rsell.ndim-2)) + (F_rsell.ndim-1, F_rsell.ndim-2)).reshape(x.shape[:-1] + (M,))
33
34
35def implicit_idft(F, L, m):
36    """Inverse of implicit_dft, returning the first L entries."""
37    F = np.asarray(F)
38    M = F.shape[-1]
39    if M % m:
40        raise ValueError("M must be divisible by m")
41    q, p = M // m, (L + m - 1) // m
42    # F[...,q*ell+r] -> F[...,r,ell]
43    A = F.reshape(F.shape[:-1] + (m, q)).swapaxes(-2, -1)
44    r = np.arange(q)[:, None]
45    t = np.arange(p)[None, :]
46    s = np.arange(m)[None, :]
47    zq = np.exp(-2j*np.pi*r*t/q)
48    zqm = np.exp(-2j*np.pi*r*s/(q*m))
49    zm = np.exp(-2j*np.pi*np.arange(m)[:, None]*np.arange(m)[None, :]/m)
50    # sum_l zeta_m^(-s*l) F[r,l], indexed r,s
51    H = np.einsum('sl,...rl->...rs', zm, A) * zqm
52    out = np.einsum('rt,...rs->...ts', zq, H) / M
53    return out.reshape(out.shape[:-2] + (p*m,))[..., :L]
54
55
56def implicit_convolve(x, g, m):
57    M = len(x) + len(g) - 1
58    # Select a convenient divisible transform length, while preserving exactness.
59    M2 = 1
60    while M2 < M or M2 % m:
61        M2 += 1
62    y = implicit_idft(implicit_dft(x, M2, m) * implicit_dft(g, M2, m), M2, m)
63    return y[:M].real