fourier_extension.py

"""Generate C++ Fourier Extension coefficients for CKKS bootstrapping."""
from __future__ import annotations
import argparse
import math
import re
import sys
import time
import warnings
from typing import Callable, List
import numpy as np
from numpy.polynomial.polyutils import RankWarning
from scipy.integrate import IntegrationWarning, quad

def _function(function: str) -> Callable[[float], float]:
    functions = {
        "sin": np.sin, "cos": np.cos, "tan": np.tan,
        "sinh": np.sinh, "cosh": np.cosh, "tanh": np.tanh,
        "arcsin": np.arcsin, "arccos": np.arccos, "arctan": np.arctan,
        "asin": np.arcsin, "acos": np.arccos, "atan": np.arctan,
        "arcsinh": np.arcsinh, "arccosh": np.arccosh, "arctanh": np.arctanh,
        "asinh": np.arcsinh, "acosh": np.arccosh, "atanh": np.arctanh,
        "exp": np.exp, "expm1": np.expm1,
        "log": np.log, "ln": np.log, "log10": np.log10, "log2": np.log2,
        "log1p": np.log1p, "sqrt": np.sqrt, "cbrt": np.cbrt,
        "pow": np.power, "abs": np.abs, "absolute": np.absolute, "fabs": np.fabs,
        "floor": np.floor, "ceil": np.ceil, "trunc": np.trunc, "sign": np.sign,
        "minimum": np.minimum, "maximum": np.maximum, "heaviside": np.heaviside,
    }
    constants = {"pi": np.pi, "e": np.e, "tau": 2.0 * np.pi, "inf": np.inf}
    allowed = {**functions, **constants}
    try:
        code = compile(function, "<function>", "eval")
        numeric = lambda value: eval(code, {"__builtins__": {}}, {**allowed, "x": value})
    except (SyntaxError, TypeError) as exc:
        raise ValueError(f"invalid function expression: {function!r}") from exc
    def evaluate(value: float) -> float:
        result = numeric(value)
        if not np.isscalar(result) or not math.isfinite(float(result)):
            raise ValueError(f"function is not finite at x={value}")
        return float(result)
    return evaluate

def _integrate(function: Callable[[float], float], left: float, right: float,
               *, points: List[float] | None = None) -> float:
    tolerance = 1e-12
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", IntegrationWarning)
        value, _ = quad(function, left, right, epsabs=tolerance, epsrel=tolerance,
                        limit=300, points=points)
    if not math.isfinite(value):
        raise ValueError("quadrature did not converge")
    return float(value)

def _integrate_extension(target, approximation, factor, left, mid, right) -> float:
    """Integrate the original and Hermite portions separately."""
    original = _integrate(lambda value: target(value) * factor(value), left, mid)
    extension = _integrate(lambda value: approximation(value) * factor(value), mid, right)
    return original + extension

def format_cpp_coefficients(coefficients: list[complex], name: str) -> str:
    if not name.isidentifier(): raise ValueError("name must be a valid C++ identifier")
    lines = [f"static const inline std::vector<std::complex<double>> {name}{{"]
    for index, coefficient in enumerate(coefficients):
        suffix = "," if index < len(coefficients) - 1 else ""
        lines.append(f"        std::complex<double>({coefficient.real:.12e}, {coefficient.imag:.12e}){suffix}")
    lines.append("};")
    return "\n".join(lines)

def coefficient_name(function: str, lower_bound: float, upper_bound: float, degree: int) -> str:
    normalized = re.sub(r"\s+", "", function.lower())
    if "0.044715" in normalized and "tanh" in normalized:
        function_name = "gelu"
    elif "exp(-x)" in normalized and normalized.startswith("1/(1+"):
        function_name = "sigmoid"
    elif normalized == "x":
        function_name = "identity"
    else:
        match = re.search(r"(?:tanh|sinh|cosh|sin|cos|tan|exp|log|sqrt|abs)", normalized)
        function_name = match.group(0) if match else "function"
    interval = f"{abs(upper_bound):g}" if math.isclose(-lower_bound, upper_bound) else f"{lower_bound:g}_{upper_bound:g}"
    return f"coeff_{function_name}_{interval}_double_{degree}".replace("-", "m").replace(".", "p")

def _hermite(function, left, mid, right, speed):
    nodes, values = [], {}
    # Fit a local polynomial and read all derivatives at once. This avoids the
    # exponential recursive finite-difference tree and its severe cancellation.
    step = min((mid - left), (right - mid)) / max(4.0 * speed, 4.0)
    def derivatives(point):
        offsets = np.arange(-speed, speed + 1, dtype=float)
        samples = np.asarray([function(point + offset * step) for offset in offsets], dtype=float)
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", RankWarning)
            polynomial = np.polynomial.Polynomial.fit(offsets, samples, 2 * speed).convert()
        return [float(polynomial.deriv(order)(0.0) / step ** order)
                for order in range(speed)]
    for node, point in ((mid, mid), (right, left)):
        values[node] = derivatives(point)
        nodes.extend([node] * speed)
    table = [[0.0] * (len(nodes) - row) for row in range(len(nodes))]
    for row, node in enumerate(nodes): table[row][0] = values[node][0]
    for order in range(1, len(nodes)):
        for row in range(len(nodes) - order):
            if all(nodes[row] == nodes[row + offset] for offset in range(order + 1)):
                table[row][order] = values[nodes[row]][order] / math.factorial(order)
            else:
                table[row][order] = (table[row + 1][order - 1] - table[row][order - 1]) / (nodes[row + order] - nodes[row])
    def polynomial(value):
        result, product = 0.0, 1.0
        for order in range(len(nodes)):
            result += table[0][order] * product
            if order < len(nodes) - 1: product *= value - nodes[order]
        return result
    return polynomial

def _try_speed(function, lower_bound, upper_bound, degree, speed, debug=False):
    started = time.perf_counter()
    period = 1.0 / 16.0
    left, mid, right = -period / 4, period / 4, 3 * period / 4
    center = (upper_bound + lower_bound) / 2
    target_base = _function(function)
    target = lambda value: 0.5 * target_base((value - center) * 2 * (upper_bound - lower_bound) / period)
    approximation = _hermite(target, left, mid, right, speed)
    omega = 2 * math.pi / float(period)
    a0 = _integrate_extension(target, approximation, lambda _: 1.0, left, mid, right) / period
    coeffs = [complex(a0, 0)]
    for n in range(1, degree + 1):
        angle = omega * n
        real = _integrate_extension(target, approximation, lambda t: np.cos(angle * t), left, mid, right)
        imag = _integrate_extension(target, approximation, lambda t: np.sin(angle * t), left, mid, right)
        coeffs.append(complex(2 * real / float(period), -2 * imag / float(period)))
    # Match the reference implementation: precision is measured by Fourier
    # reconstruction error on the original-function half of the interval.
    grid = np.linspace(float(left), float(mid), 2000)
    target_values = np.asarray([target(float(value)) for value in grid], dtype=float)
    reconstructed = np.full(grid.shape, coeffs[0].real, dtype=float)
    for n, coefficient in enumerate(coeffs[1:], start=1):
        angle = omega * n
        reconstructed += coefficient.real * np.cos(angle * grid) - coefficient.imag * np.sin(angle * grid)
    error = float(np.mean(np.abs(reconstructed - target_values)))
    bits = float("inf") if error == 0 else -math.log2(error)
    if debug:
        print(f"[debug] ke={speed}: error={error:.3e}, bits={bits:.2f}, total={time.perf_counter() - started:.2f}s", file=sys.stderr, flush=True)
    return bits, coeffs

def calculate_fourier_coefficients(func_str: str, lower_bound: float, upper_bound: float,
                                   degree: int, debug: bool = False):
    """Return ``(ke, achieved_bits, coefficients)`` for a Fourier extension."""
    if degree < 0: raise ValueError("degree must be non-negative")
    if lower_bound >= upper_bound: raise ValueError("lower_bound must be smaller than upper_bound")
    function = func_str
    best = (-1.0, 0, [])
    stagnant = 0
    speed = 1
    while stagnant < 6:
        bits, coeffs = _try_speed(function, lower_bound, upper_bound, degree, speed, debug)
        if bits > best[0]:
            best = (bits, speed, coeffs)
            stagnant = 0
        else:
            stagnant += 1
        speed += 1
    return best[1], best[0], best[2]

if __name__ == "__main__":
    # Example target functions (the interval mapping is applied internally):
    #   tanh:
    #     python fourier_extension.py --func "tanh(x)" --left -1 --right 1 --degree 18
    #   sigmoid:
    #     python fourier_extension.py --func "1/(1+exp(-x))" --left -8 --right 8 --degree 32
    #   GELU using the tanh approximation:
    #     python fourier_extension.py \
    #       --func "0.5*x*(1+tanh(sqrt(2/pi)*(x+0.044715*x**3)))" \
    #       --left -4 --right 4 --degree 32
    parser = argparse.ArgumentParser(description="Fourier extension parameter search")
    parser.add_argument("--func", default="x")
    parser.add_argument("--left", type=float, default=-1.0)
    parser.add_argument("--right", type=float, default=1.0)
    parser.add_argument("--degree", type=int, default=20)
    parser.add_argument("--debug", action="store_true", help="print search progress to stderr")
    args = parser.parse_args()
    speed, achieved, coeffs = calculate_fourier_coefficients(
        args.func, args.left, args.right, args.degree, args.debug)
    print(f"// ke: {speed}, precision: {achieved:.4f} bits")
    print(format_cpp_coefficients(coeffs, coefficient_name(args.func, args.left, args.right, args.degree)))