Butterfly Factorization

Copyright (C) 2020 Andreas Kloeckner

MIT License

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

Following the paper "An algorithm for the rapid evaluation of special function transforms", by Michael O’Neil, Franco Woolfe, Vladimir Rokhlin.

import numpy as np
import scipy.linalg as la
import scipy.linalg.interpolative as sli
import matplotlib.pyplot as plt
import scipy.special as sps
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle

Set some parameters: (ignore nlevels for now)

#nlevels = 4
#nlevels = 7
nlevels = 9
n = 2**(nlevels + 2)

Part 1: Constant-Area Ranks

The Discrete Fourier Transform

def make_dft(n, power):
    omega = np.exp(2*np.pi*1j/n)

    ns = np.arange(n)
    exponents = ns.reshape(-1, 1) * ns
    
    return omega**(power*exponents)

dft = make_dft(n, power=1)
idft = make_dft(n, power=-1)

la.norm(np.abs(idft @ dft) - n*np.eye(n))

1.0665710446355353e-07

Verify the FFT property:

quotient = dft[::2, :n//2] / make_dft(n//2, power=1)
plt.imshow(quotient.real)
plt.colorbar()
print(np.max(np.abs(quotient - 1)))

7.629497034145061e-11

plt.imshow(dft.real)

<matplotlib.image.AxesImage at 0x7f661aba7170>

Consider the the claim that the numerical rank of the kernel $e^{ixt}$ for $x\in[0,X]$ and $t\in[0,T]$ depends only on the product $XT$:

T = 15
X = 15
resolution = 200

x, t = np.mgrid[0:X:resolution * 1j, 0:T:resolution * 1j]

mat = np.exp(1j*x*t)
plt.imshow(mat[:, ::-1].T.real, extent=(0, X, 0, T))

scale = 6
for exp in np.linspace(-1.25, 1.25, 5):
    subX = 2**exp * scale
    subT = 2**-exp * scale
    plt.gca().add_patch(Rectangle((0, 0), subX, subT, fill=False))

# Observe: These are all the same matrix!

Xfacs = np.linspace(1/2, 2, 30)
Tfacs = 1/Xfacs

# Change me
scale = np.pi  
resolution = 30

for Xfac, Tfac in zip(Xfacs, Tfacs):
    x, t = np.mgrid[0:Xfac*scale:resolution * 1j, 0:Tfac*scale:resolution * 1j]
    mat = np.exp(1j*x*t)
    
    _, sigma, _ = la.svd(mat)
    print(f"{Xfac:.2f} {Tfac:.2f}\t", np.sum(sigma > 1e-7))

0.50 2.00	 11
0.55 1.81	 11
0.60 1.66	 11
0.66 1.53	 11
0.71 1.41	 11
0.76 1.32	 11
0.81 1.23	 11
0.86 1.16	 11
0.91 1.09	 11
0.97 1.04	 11
1.02 0.98	 11
1.07 0.94	 11
1.12 0.89	 11
1.17 0.85	 11
1.22 0.82	 11
1.28 0.78	 11
1.33 0.75	 11
1.38 0.72	 11
1.43 0.70	 11
1.48 0.67	 11
1.53 0.65	 11
1.59 0.63	 11
1.64 0.61	 11
1.69 0.59	 11
1.74 0.57	 11
1.79 0.56	 11
1.84 0.54	 11
1.90 0.53	 11
1.95 0.51	 11
2.00 0.50	 11

The Legendre Vandermonde / Transform

lege_nodes = sps.legendre(n).weights[:, 0]
lege_vdm = np.array([sps.eval_legendre(i, lege_nodes) for i in range(n)]).T
plt.imshow(lege_vdm)

<matplotlib.image.AxesImage at 0x7f66135a4c50>

The Chebyshev Transform

k = n-1

i = np.arange(0, k+1)
x = np.linspace(-1, 1, 3000)

nodes = np.cos(i/k*np.pi)

i = np.arange(n, dtype=np.float64)
nodes = np.cos((2*(i+1)-1)/(2*n)*np.pi)
chebyshev_vdm = np.cos(i*np.arccos(nodes.reshape(-1, 1)))

plt.imshow(chebyshev_vdm)

<matplotlib.image.AxesImage at 0x7f6618179340>

A Random Matrix

randmat = np.random.randn(n, n)
plt.imshow(randmat)

<matplotlib.image.AxesImage at 0x7f66180aa000>

Part 2: The Butterfly Factorization

A Matrix Slicing Helper

class Level:
    def __init__(self, level, nlevels, n=None):
        self.level = level
        self.nlevels = nlevels
        
        if level > nlevels:
            raise ValueError("level too large")
            
        if n is None:
            n = 2**nlevels
        self.n = n
        
    @property
    def nblock_rows(self):
        return 2**self.level
        
    @property
    def block_nrows(self):
        return self.n//self.nblock_rows

    @property
    def nblock_cols(self):
        return 2**(self.nlevels-self.level)
        
    @property
    def block_ncols(self):
        return self.n//self.nblock_cols
    
    def matview(self, bi, bj, mat):
        br = self.block_nrows
        bc = self.block_ncols
        return mat[br*bi:br*(bi+1), bc*bj:bc*(bj+1)]
    
    def rowview(self, bi, vec):
        br = self.block_nrows
        return vec[br*bi:br*(bi+1)]
    
    def colview(self, bj, vec):
        bc = self.block_ncols
        return vec[bc*bj:bc*(bj+1)]

Level(0, nlevels, 256).matview(0, 0, dft).shape

(256, 0)

Rank-Revealing Factorization

epsilon = 1e-10

# ID

def id_decomp(A):
    k, idx, proj = sli.interp_decomp(A, epsilon)
    sort_idx = np.argsort(idx)
    B = A[:,idx[:k]]
    P = np.hstack([np.eye(k), proj])[:,np.argsort(idx)]
    return B, P

# Rank-Revealing Truncated QR

def qr_decomp(A):
    q, r, p = la.qr(A, pivoting=True, mode="economic")
    diag_r = np.diag(r)
    r = r[:, np.argsort(p)]
    
    flags = np.abs(diag_r) >= epsilon
    q = q[:, flags]
    r = r[flags]
    return q, r

#decomp = qr_decomp
decomp = id_decomp

def make_low_rank_matrix(n):
    A0 = np.random.randn(n, n)
    U0, sigma0, VT0 = la.svd(A0)
    sigma = np.exp(-np.arange(n))
    return (U0 * sigma).dot(VT0)

Atest = make_low_rank_matrix(100)
Btest, Ptest = decomp(Atest)

la.norm(Atest - Btest@Ptest)/la.norm(Atest)

1.0681674657170179e-10

plt.imshow(Ptest)

<matplotlib.image.AxesImage at 0x7f66182e2660>

Precomputation

A = dft

# keys: [level][i, j]
Ps = [{} for i in range(nlevels+1)]
Bs = [{} for i in range(nlevels+1)] 

Level 0

lev = Level(0, nlevels, n)

assert lev.nblock_rows == 1
for i in range(lev.nblock_rows):
    for j in range(lev.nblock_cols):
        b, p = Bs[0][i, j], Ps[0][i, j] = decomp(lev.matview(i, j, A))
        if (i, j) == (0, 0):
            print(f"B: {b.shape[0]}x{b.shape[1]} P: {p.shape[0]}x{p.shape[1]}")

B: 2048x4 P: 4x4

Levels 1, ..., L

for ilev in range(1, nlevels + 1):
    lev = Level(ilev, nlevels, n)

    for j in range(lev.nblock_rows):
        for k in range(lev.nblock_cols):
            # only process even j
            if j % 2 != 0:
                continue
                
            bblock = np.hstack((
                Bs[ilev-1][j//2, 2*k],
                Bs[ilev-1][j//2, 2*k+1],
                ))
            bblock_top = bblock[:lev.block_nrows]
            bblock_bottom = bblock[lev.block_nrows:]
            assert len(bblock_top)*2 == len(bblock)
            
            Bs[ilev][j, k], Ps[ilev][j, k] = decomp(bblock_top)
            Bs[ilev][j+1, k], Ps[ilev][j+1, k] = decomp(bblock_bottom)

            if (j, k) == (0, 0):
                tB = Bs[ilev][j, k].shape
                tP = Ps[ilev][j, k].shape
                print(
                    f"Level {ilev}: {lev.block_nrows}x{lev.block_ncols} "
                    f"Btop: {bblock_top.shape[0]}x{bblock_top.shape[1]} -> B: {tB[0]}x{tB[1]} P: {tP[0]}x{tP[1]}")

Level 1: 1024x8 Btop: 1024x8 -> B: 1024x8 P: 8x8
Level 2: 512x16 Btop: 512x16 -> B: 512x16 P: 16x16
Level 3: 256x32 Btop: 256x32 -> B: 256x17 P: 17x32
Level 4: 128x64 Btop: 128x34 -> B: 128x17 P: 17x34
Level 5: 64x128 Btop: 64x34 -> B: 64x17 P: 17x34
Level 6: 32x256 Btop: 32x34 -> B: 32x17 P: 17x34
Level 7: 16x512 Btop: 16x34 -> B: 16x16 P: 16x34
Level 8: 8x1024 Btop: 8x32 -> B: 8x8 P: 8x32
Level 9: 4x2048 Btop: 4x16 -> B: 4x4 P: 4x16

levels = []
ranks = []

for ilev in range(1, nlevels + 1):
    levels.append(ilev)
    ranks.append(Bs[ilev][0,0].shape[1])
    
plt.plot(levels, ranks, "o-")
plt.grid()
plt.xlabel("Level")
plt.ylabel("Rank")

Text(0, 0.5, 'Rank')

Only the last-level $B$ actually needs to be retained:

LLB = Bs[-1]
del Bs

Matvec

First, generate a random input:

x = np.random.randn(n)

Setup

# keys: [ilevel][i, j]
betas = [{} for i in range(nlevels+1)]

Level 0

lev = Level(0, nlevels, n)

assert lev.nblock_rows == 1
for i in range(lev.nblock_rows):
    for j in range(lev.nblock_cols):
        betas[0][i, j] = Ps[0][i, j] @ lev.colview(j, x)

Level 1, ..., L

for ilev in range(1, nlevels + 1):
    lev = Level(ilev, nlevels, n)

    for j in range(lev.nblock_rows):
        for k in range(lev.nblock_cols):
            beta_glued = np.hstack((
                betas[ilev-1][j//2, 2*k],
                betas[ilev-1][j//2, 2*k+1]
                ))
            
            betas[ilev][j, k] = Ps[ilev][j, k] @ beta_glued
            if (j, k) == (0, 0):
                p = Ps[ilev][j, k]
                print(f"P: {p.shape[0]}x{p.shape[1]} * ({lev.nblock_rows}x{lev.nblock_cols} = {lev.nblock_rows*lev.nblock_cols})")

P: 8x8 * (2x256 = 512)
P: 16x16 * (4x128 = 512)
P: 17x32 * (8x64 = 512)
P: 17x34 * (16x32 = 512)
P: 17x34 * (32x16 = 512)
P: 17x34 * (64x8 = 512)
P: 16x34 * (128x4 = 512)
P: 8x32 * (256x2 = 512)
P: 4x16 * (512x1 = 512)

Postprocess

Ax = np.zeros(n, dtype=np.complex128)

lev = Level(nlevels, nlevels, n)

assert lev.nblock_cols == 1
for j in range(lev.nblock_rows):
    for k in range(lev.nblock_cols):
        lev.rowview(j, Ax)[:] = LLB[j, k] @ betas[nlevels][j, k]

la.norm(Ax - A@x)/la.norm(A@x)

6.847661548741153e-11