"""Fast multipole method (FMM) for the 2D log potential phi_i = sum_j q_j log|z_i - z_j|. Note: ../../Fast_multipole_method.md Source: Greengard & Rokhlin (1987), J. Comput. Phys. 73, 325 (Lemmas 2.1-2.5 give the five operators). The course (Scientific Computing 2, TUM, 2020, worksheet 5) built only the two ring approximations of Anderson's variant (see anderson_rings.py). This file is the complete algorithm, written for the vault, so that the O(N) structure can be run and timed. Points are complex numbers z = x + iy in the unit square. The five operators (p = number of expansion terms) P2M particles -> multipole of a leaf box phi(z) = a0 log(z-zc) + sum_k a_k (z-zc)^-k M2M child multipole -> parent multipole (upward pass, finest level to level 2) M2L multipole of a well-separated box -> local Taylor expansion of the target box L2L parent local -> child local (downward pass) L2P local expansion -> potential at particles; plus direct sums with the 9 neighbour boxes. Well separated = not touching. A box's interaction list is the children of its parent's neighbours that do not touch the box: at most 27 boxes in 2D. That fixed count is why the method is O(N). Run: python fmm_2d.py (checks against the direct sum, prints the error for several p) The loops over boxes are plain Python, so at N = 4000 the vectorised O(N^2) direct sum is still faster; the point of this file is the structure and the error, not the speed. """ from __future__ import annotations import time from math import comb import numpy as np def _binom_table(n): return np.array([[comb(i, j) for j in range(n)] for i in range(n)], float) class FMM2D: def __init__(self, z: np.ndarray, q: np.ndarray, levels: int = 4, p: int = 12): self.z, self.q, self.L, self.p = z, q, levels, p self.C = _binom_table(2 * p + 2) n = 2 ** levels ix = np.minimum((z.real * n).astype(int), n - 1) iy = np.minimum((z.imag * n).astype(int), n - 1) order = np.lexsort((ix, iy)) self.leaf_of = (ix, iy) # particle indices per leaf box self.members = {} for k in order: self.members.setdefault((ix[k], iy[k]), []).append(k) self.members = {b: np.array(v) for b, v in self.members.items()} @staticmethod def center(level, ix, iy): n = 2 ** level return complex((ix + 0.5) / n, (iy + 0.5) / n) # ---- the five operators ------------------------------------------------------------------- def p2m(self, zc, idx): d = self.z[idx] - zc k = np.arange(1, self.p + 1) a = np.empty(self.p + 1, complex) a[0] = self.q[idx].sum() a[1:] = -(self.q[idx, None] * d[:, None] ** k).sum(0) / k return a def m2m(self, a, z0): """Shift a multipole from a child centred at z0 (relative to the parent centre).""" p, C = self.p, self.C b = np.zeros(p + 1, complex) b[0] = a[0] for l in range(1, p + 1): k = np.arange(1, l + 1) b[l] = -a[0] * z0 ** l / l + (a[k] * z0 ** (l - k) * C[l - 1, k - 1]).sum() return b def m2l(self, a, z0): """Multipole centred at z0 (relative to the target centre) -> local expansion about the target.""" p, C = self.p, self.C k = np.arange(1, p + 1) s = a[k] / z0 ** k * (-1.0) ** k b = np.empty(p + 1, complex) b[0] = a[0] * np.log(-z0) + s.sum() for l in range(1, p + 1): b[l] = -a[0] / (l * z0 ** l) + (s * C[l + k - 1, k - 1]).sum() / z0 ** l return b def l2l(self, c, z0): """Local expansion about z0 (relative to the new centre) -> local expansion about the new centre.""" p, C = self.p, self.C d = np.zeros(p + 1, complex) for l in range(p + 1): k = np.arange(l, p + 1) d[l] = (c[k] * C[k, l] * (-z0) ** (k - l)).sum() return d def l2p(self, c, zc, idx): w = self.z[idx] - zc return np.polyval(c[::-1], w) # ---- the algorithm ------------------------------------------------------------------------ def potential(self) -> np.ndarray: L, p = self.L, self.p M = [dict() for _ in range(L + 1)] # upward pass for (ix, iy), idx in self.members.items(): M[L][(ix, iy)] = self.p2m(self.center(L, ix, iy), idx) for l in range(L, 2, -1): for (ix, iy), a in M[l].items(): par = (ix // 2, iy // 2) z0 = self.center(l, ix, iy) - self.center(l - 1, *par) M[l - 1][par] = M[l - 1].get(par, 0) + self.m2m(a, z0) # downward pass Loc = [dict() for _ in range(L + 1)] for l in range(2, L + 1): n = 2 ** l for ix in range(n): for iy in range(n): c = np.zeros(p + 1, complex) if l > 2 and (ix // 2, iy // 2) in Loc[l - 1]: z0 = self.center(l - 1, ix // 2, iy // 2) - self.center(l, ix, iy) c += self.l2l(Loc[l - 1][(ix // 2, iy // 2)], z0) px, py = ix // 2, iy // 2 for jx in range(2 * px - 2, 2 * px + 4): # children of parent's neighbours for jy in range(2 * py - 2, 2 * py + 4): if abs(jx - ix) <= 1 and abs(jy - iy) <= 1: continue # touching: handled later if (jx, jy) in M[l]: z0 = self.center(l, jx, jy) - self.center(l, ix, iy) c += self.m2l(M[l][(jx, jy)], z0) Loc[l][(ix, iy)] = c # evaluation: far field from the local expansion, near field directly phi = np.zeros(len(self.z)) for (ix, iy), idx in self.members.items(): phi[idx] += self.l2p(Loc[L][(ix, iy)], self.center(L, ix, iy), idx).real near = [self.members[(jx, jy)] for jx in range(ix - 1, ix + 2) for jy in range(iy - 1, iy + 2) if (jx, jy) in self.members] src = np.concatenate(near) r = np.abs(self.z[idx, None] - self.z[None, src]) with np.errstate(divide='ignore'): lg = np.where(r > 0, np.log(r), 0.0) phi[idx] += lg @ self.q[src] return phi def potential_direct(z, q): r = np.abs(z[:, None] - z[None, :]) np.fill_diagonal(r, 1.0) # log 1 = 0 removes self-interaction return np.log(r) @ q def main(): rng = np.random.default_rng(1) n = 4000 z = rng.random(n) + 1j * rng.random(n) q = rng.choice([-1.0, 1.0], n) t0 = time.perf_counter(); ref = potential_direct(z, q); t1 = time.perf_counter() print(f"direct O(N^2): {t1 - t0:.2f} s (N = {n})") for p in (4, 8, 12, 16, 20): t2 = time.perf_counter(); phi = FMM2D(z, q, levels=4, p=p).potential(); t3 = time.perf_counter() err = np.max(np.abs(phi - ref)) / np.max(np.abs(ref)) print(f"p = {p:2d}: max rel. error {err:.1e} time {t3 - t2:.2f} s") if __name__ == '__main__': main()