thermite-complex 0.2.0

#!/usr/bin/env python3
"""Generates `tables.rs` - Weideman's coefficients for the Faddeeva function w(z).

    pip install mpmath && python tables.py > tables.rs

A direct transcription of `reference/cef.m` (J. A. C. Weideman, 1995), the reference
implementation accompanying "Computation of the complex error function", SIAM J. Numer.
Anal. 31(5), 1497-1518, 1994:

    M = 2*N;  M2 = 2*M;  k = [-M+1:1:M-1]';
    L = sqrt(N/sqrt(2));                       % Optimal choice of L.
    theta = k*pi/M; t = L*tan(theta/2);        % Define variables theta and t.
    f = exp(-t.^2).*(L^2+t.^2); f = [0; f];    % Function to be transformed.
    a = real(fft(fftshift(f)))/M2;             % Coefficients of transform.
    a = flipud(a(2:N+1));                      % Reorder coefficients.

The transform is only 4N points at N <= 40, so a naive O(n^2) DFT is fine and no FFT
is needed - but it MUST be evaluated in extended precision, hence mpmath. Doing it in
f64 costs a factor of 10 in the final accuracy of w(z) at N = 40 (8.9e-15 rather than
8.7e-16, measured): the leading coefficients are down at 1e-15, so f64 rounding in the
samples destroys them completely, and no amount of careful summation (Kahan, fsum)
recovers it because the error is in the inputs rather than the accumulation. At 60
digits the result agrees with numpy's FFT to the last bit at every N.

The `flipud` is what leaves the coefficients in MATLAB `polyval` order, i.e.
LEADING-TERM-FIRST. That is the order `tables.rs` stores and the order `poly_rev`
expects. Do not "fix" it to constant-term-first: the reversed array is a different
polynomial that happens to agree at Z = 1, which is exactly how one of these got
shipped wrong in this workspace before.
"""

import struct

import mpmath as mp

mp.mp.dps = 60

NS = (8, 16, 24, 32, 40)


def coefficients(n):
    """Returns `(L, a)` with `a` leading-term-first, per cef.m."""
    m = 2 * n
    m2 = 2 * m
    ell = mp.sqrt(mp.mpf(n) / mp.sqrt(2))

    # f = [0; exp(-t^2) * (L^2 + t^2)] over the tangent grid, length M2 = 4N
    f = [mp.mpf(0)]
    for k in range(-m + 1, m):
        t = ell * mp.tan((mp.mpf(k) * mp.pi / m) / 2)
        f.append(mp.exp(-(t**2)) * (ell**2 + t**2))
    assert len(f) == m2

    # fftshift, then a forward DFT. Only the real part of entries 1..N is kept, and
    # the input is real, so the sine half of each twiddle never contributes.
    shifted = f[m2 // 2 :] + f[: m2 // 2]

    a = [
        mp.fsum(x * mp.cos(-2 * mp.pi * i * j / m2) for i, x in enumerate(shifted)) / m2
        for j in range(1, n + 1)
    ]

    a.reverse()  # flipud: ascending -> leading-term-first
    return ell, a


def as_f32(v):
    return struct.unpack("f", struct.pack("f", float(v)))[0]


def fmt(v, f32):
    """Shortest decimal that round-trips through the target format."""
    target = as_f32(v) if f32 else float(v)
    round_trip = as_f32 if f32 else float

    for p in range(1, 18):
        s = "%.*g" % (p, target)
        if round_trip(s) == target:
            break

    # Rust needs a decimal point or an exponent to infer a float literal
    if "." not in s and "e" not in s:
        s += ".0"

    return s


HEADER = '''//! Weideman coefficient tables for the Faddeeva function.
//!
//! **Generated by `tables.py`, in this directory. Do not edit by hand.**
//!
//! Leading-term-first (MATLAB `polyval` order), which is what
//! [`poly_rev`](thermite::math::specialized::SpecializedCoreMath::poly_rev) expects.
//! Reordering them silently yields a different polynomial that agrees at `Z = 1`.
//!
//! The f32 tables are the same values rounded, not a separate fit.
//! `Weideman<24|32|40> for f32` exist only to satisfy the
//! [`WeidemanTables`](super::WeidemanTables) supertrait bound - f32 clamps at
//! `MAX_N = 16` and never instantiates them.

use super::Weideman;
'''

FOOTER = """
impl super::WeidemanTables for f64 {
    const MAX_N: usize = 40;
    // sqrt(f64::MAX) is ~1.34e154; back off for the (L + y)^2 headroom.
    const HUGE: Self = 1.0e150;
    // Measured crossover: inside this box the correction is never worse than the direct
    // evaluation and up to 4 orders better; outside it, the direct one wins.
    const REAL_AXIS_Y: Self = 1.0e-5;
    const REAL_AXIS_X: Self = 1.0e3;
}

impl super::WeidemanTables for f32 {
    // f32 Horner roundoff floors at ~5e-7 (N = 16); more terms measure no better.
    const MAX_N: usize = 16;
    // sqrt(f32::MAX) is ~1.84e19.
    const HUGE: Self = 1.0e17;
    const REAL_AXIS_Y: Self = 1.0e-3;
    const REAL_AXIS_X: Self = 1.0e2;
}
"""


def main():
    out = [HEADER]

    tables = {n: coefficients(n) for n in NS}

    for ty in ("f64", "f32"):
        f32 = ty == "f32"
        for n in NS:
            ell, a = tables[n]
            out.append(f"impl Weideman<{n}> for {ty} {{")
            out.append(f"    const L: Self = {fmt(ell, f32)};")
            out.append(f"    const A: [Self; {n}] = [")
            out.extend(f"        {fmt(v, f32)}," for v in a)
            out.append("    ];")
            out.append("}")
            out.append("")

    out.append(FOOTER.strip())
    print("\n".join(out))


if __name__ == "__main__":
    main()