Fast discrete cosine transform in Python

Computes the DCT in O(n log n) by reducing it to an FFT of the same length.

File: fast_dct.py, 82 lines
Level: I0 not checked yet
Toolchain: Python 3.14, pytest 9
Test framework: pytest
Run the tests: python3 -m pytest

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/fast-dct/. The authors are listed in prov/EWEX4W.jsonl.

import cmath
import math


def _fft(x: list[complex], inverse: bool = False) -> None:
    """In-place iterative radix-2 FFT with the root exp(-2j*pi/n), or its inverse.

    len(x) must be a power of two. The inverse includes the factor 1/n.
    """
    n = len(x)
    t = n.bit_length() - 1
    sign = 1.0 if inverse else -1.0
    table = [cmath.exp(sign * 2j * math.pi * k / n) for k in range(max(n // 2, 1))]
    rev = [0] * n
    for j in range(1, n):
        rev[j] = (rev[j >> 1] >> 1) | ((j & 1) << (t - 1))
    for j in range(n):
        if j < rev[j]:
            x[j], x[rev[j]] = x[rev[j]], x[j]
    h = 1
    while h < n:
        step = n // (2 * h)
        for s in range(0, n, 2 * h):
            for k in range(h):
                u = x[s + k]
                v = table[k * step] * x[s + k + h]
                x[s + k] = u + v
                x[s + k + h] = u - v
        h *= 2
    if inverse:
        for j in range(n):
            x[j] /= n


def fast_dct(x: list[float]) -> list[float]:
    """Return the DCT-II, X[k] = sum of x[j] * cos(pi * (2j + 1) * k / (2n)).

    Precondition: n = len(x) is a power of two and the entries are real.
    Method of Makhoul (1980): reorder, one complex FFT of length n, and a
    rotation of each output. Cost: one complex FFT of length n, about
    5 n lg n real operations, plus 3n real operations for the rotations.
    """
    n = len(x)  # @line L1
    if n < 1 or n & (n - 1):
        raise ValueError("the length must be a power of two")
    if n == 1:  # @line L2
        return [float(x[0])]
    v = [0j] * n
    for j in range(n // 2):  # @line L3
        v[j] = complex(x[2 * j])  # @line L4
        v[n - 1 - j] = complex(x[2 * j + 1])  # @line L5
    _fft(v)  # @line L6
    out = [0.0] * n
    for k in range(n):  # @line L7
        out[k] = (cmath.exp(-1j * math.pi * k / (2 * n)) * v[k]).real  # @line L8
    return out  # @line L9


def inverse_fast_dct(big_x: list[float]) -> list[float]:
    """Return x with fast_dct(x) == big_x (the scaled DCT-III).

    Precondition: as for fast_dct. Cost: one inverse complex FFT of length
    n and O(n) further real operations.
    """
    n = len(big_x)
    if n < 1 or n & (n - 1):
        raise ValueError("the length must be a power of two")
    if n == 1:
        return [float(big_x[0])]
    w = [complex(big_x[0])] + [
        cmath.exp(1j * math.pi * k / (2 * n)) * complex(big_x[k], -big_x[n - k])
        for k in range(1, n)
    ]  # @line L10
    _fft(w, inverse=True)  # @line L11
    x = [0.0] * n
    for j in range(n // 2):  # @line L12
        x[2 * j] = w[j].real
        x[2 * j + 1] = w[n - 1 - j].real
    return x

test_fast_dct.py 93 lines

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/fast-dct/. The authors are listed in prov/EWEX4W.jsonl.

import cmath
import math
import random

import pytest

from fast_dct import _fft, fast_dct, inverse_fast_dct


def naive_dct(x: list[float]) -> list[float]:
    n = len(x)
    return [
        sum(x[j] * math.cos(math.pi * (2 * j + 1) * k / (2 * n)) for j in range(n))
        for k in range(n)
    ]


def naive_dft(x: list[complex]) -> list[complex]:
    n = len(x)
    return [sum(x[j] * cmath.exp(-2j * math.pi * j * k / n) for j in range(n)) for k in range(n)]


def close(a: list, b: list, scale: float) -> bool:
    n = max(len(a), 1)
    tol = 1e-12 * n * max(n.bit_length(), 1) * max(scale, 1.0)
    return len(a) == len(b) and all(abs(u - v) <= tol for u, v in zip(a, b))


def test_examples() -> None:
    c, s = math.cos(math.pi / 8), math.sin(math.pi / 8)
    assert close(fast_dct([1, 2, 3, 4]), [10, -3 * c - s, 0, -3 * s + c], 4)
    assert all(abs(u - v) < 1e-4 for u, v in zip(fast_dct([1, 2, 3, 4]), [10, -3.1543, 0, -0.2242]))
    out = fast_dct([1, 2, 3, 4, 5, 6, 7, 8])
    assert abs(out[0] - 36) < 1e-12
    assert all(abs(out[k]) < 1e-12 for k in (2, 4, 6))
    assert close(out, naive_dct([1, 2, 3, 4, 5, 6, 7, 8]), 8)


def test_internal_fft_against_naive_dft() -> None:
    rng = random.Random(4)
    for t in range(0, 7):
        n = 1 << t
        x = [complex(rng.uniform(-1, 1), rng.uniform(-1, 1)) for _ in range(n)]
        y = list(x)
        _fft(y)
        assert close(y, naive_dft(x), n)
        _fft(y, inverse=True)
        assert close(y, x, 1)


def test_edges() -> None:
    assert fast_dct([2.5]) == [2.5]
    assert inverse_fast_dct([2.5]) == [2.5]
    assert close(fast_dct([1, 1]), [2, 0], 2)
    assert close(fast_dct([1, -1]), [0, math.sqrt(2)], 2)
    for t in range(0, 7):
        n = 1 << t
        assert close(fast_dct([0.0] * n), [0.0] * n, 1)
        assert close(fast_dct([3.0] * n), [3.0 * n] + [0.0] * (n - 1), 3 * n)
        for pos in (0, n - 1):
            imp = [0.0] * n
            imp[pos] = 1.0
            assert close(fast_dct(imp), naive_dct(imp), 1)


def test_not_power_of_two() -> None:
    for n in (0, 3, 6, 10):
        with pytest.raises(ValueError):
            fast_dct([1.0] * n)
        with pytest.raises(ValueError):
            inverse_fast_dct([1.0] * n)


def test_random_against_naive() -> None:
    rng = random.Random(1974)
    for _ in range(600):
        n = 1 << rng.randint(0, 6)
        scale = rng.choice([1.0, 1e3])
        x = [rng.uniform(-scale, scale) for _ in range(n)]
        out = fast_dct(x)
        assert close(out, naive_dct(x), scale * n)
        assert close(inverse_fast_dct(out), x, scale)


def test_inverse_of_naive() -> None:
    rng = random.Random(1980)
    for _ in range(100):
        n = 1 << rng.randint(1, 6)
        x = [rng.uniform(-1, 1) for _ in range(n)]
        assert close(inverse_fast_dct(naive_dct(x)), x, n)

What the next level needs

the tests did not run in this build

Read the full page about Fast discrete cosine transform