Chebyshev 1D/2D Approximation in C++

ChebyshevApproximation.h

/*
  ==============================================================================
    ChebyshevApproximation.h

    Creates Chebyshev-polynomial approximations for smooth 1d and 2d function
    lambdas and emits C++ source code to use in place of the lambda.

    More info and examples: https://apulsoft.ch/blog/chebyshev-approximation/

    (c) 2023 Adrian Pflugshaupt

    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.
  ==============================================================================
*/

#pragma once

#include <array>
#include <cassert>
#include <cmath>
#include <iomanip>
#include <iostream>
#include <vector>

namespace plap {

namespace cheby_priv {

// A helper class to hold Chebyshev approximation coefficients
// for multi-dimensional support this needs to be able to nest with itself
// therefore some operators for the necessary math are defined.
// Template params:
//  T the coefficient type (can be Cheby<>)
//  N the number of coefficients to use (polynomial order + 1)
//  Tquery (simple) type of the x,... query variable and of the cosines used to sum the values
template<typename T, int N, typename Tquery = T>
class Coeffs {
public:
    Coeffs() {}

    Coeffs(T x) {
        for (auto i = 0; i < N; ++i) c[i] = x;
    }

    Coeffs &operator+=(const Coeffs &b) {
        for (auto i = 0; i < N; ++i) c[i] += b.c[i];
        return *this;
    }
    Coeffs &operator-=(const Coeffs &b) {
        for (auto i = 0; i < N; ++i) c[i] -= b.c[i];
        return *this;
    }
    Coeffs &operator*=(const Coeffs &b) {
        for (auto i = 0; i < N; ++i) c[i] *= b.c[i];
        return *this;
    }
    Coeffs &operator*=(const T &b) {
        for (auto i = 0; i < N; ++i) c[i] *= b;
        return *this;
    }

    const Coeffs operator+(const Coeffs &b) const { return Coeffs(*this) += b; }
    const Coeffs operator-(const Coeffs &b) const { return Coeffs(*this) -= b; }
    const Coeffs operator*(const Coeffs &b) const { return Coeffs(*this) *= b; }
    const Coeffs operator*(const T &b) const { return Coeffs(*this) *= b; }

    /// build Chebyshev coefficients from values at Chebyshev nodes.
    /// if v.size() > N, the Chebyshev coeffs are trunctated to N
    void buildFromValues(const std::vector<T> &v) {
        auto M = v.size();
        auto invM = 1 / Tquery(M);
        const double theta = M_PI * invM;
        for (auto i = 0; i < N; ++i) {
            T sum = 0;
            auto angle = Tquery(0.5) * Tquery(i) * theta;
            const auto step = 2 * angle;
            for (auto j = 0; j < M; ++j) {
                sum += v[j] * std::cos(angle);
                angle += step;
            }
            c[i] = sum * (2 * invM);
        }
    }

    // estimate max error based on the truncated coeffs
    T estimatedErrorForValues(const std::vector<T> &v) {
        T res{0};
        auto M = v.size();
        auto invM = 1 / Tquery(M);
        const double theta = M_PI * invM;
        for (auto i = N; i < M; ++i) {
            T sum = 0;
            auto angle = Tquery(0.5) * Tquery(i) * theta;
            const auto step = 2 * angle;
            for (auto j = 0; j < M; ++j) {
                sum += v[j] * std::cos(angle);
                angle += step;
            }
            res += std::fabs(sum * (2 * invM));
        }
        return res;
    }

    // x needs to be inside -1..1
    // this uses a two-step Clenshaw algorithm to avoid temporary variables
    inline T eval(Tquery x_norm) const {
        static_assert(N >= 3);  // lower N not implemented (they make no sense)

        auto y = x_norm + x_norm;
        T b1, b2;
        if constexpr ((N & 1) == 0) {  // even number of coeffs
            b2 = c[N - 2] + c[N - 1] * y;
            b1 = c[N - 3] + b2 * y - c[N - 1];
            for (int i = N - 4; i > 0; i -= 2) {
                b2 = c[i] + b1 * y - b2;
                b1 = c[i - 1] + b2 * y - b1;
            }
        } else {  // odd number of coeffs
            b2 = c[N - 1];
            b1 = c[N - 2] + b2 * y;
            for (int i = N - 3; i > 0; i -= 2) {
                b2 = c[i] + b1 * y - b2;
                b1 = c[i - 1] + b2 * y - b1;
            }
        }
        return (b1 * y + c[0]) * 0.5 - b2;
    }

    std::array<T, N> polyCoeffs() const {
        std::array<T, N> pcf = {0};
        std::array<T, N> buf = {0};
        pcf[0] = c[N - 1];
        for (int i = N - 2; i > 0; --i) {
            for (int j = N - i; j > 0; --j) {
                std::swap(pcf[j], buf[j]);
                pcf[j] = 2 * pcf[j - 1] - pcf[j];
            }
            std::swap(pcf[0], buf[0]);
            pcf[0] = c[i] - pcf[0];
        }
        for (int i = N - 1; i > 0; --i) pcf[i] = pcf[i - 1] - buf[i];
        pcf[0] = 0.5 * c[0] - buf[0];

        return pcf;
    }

    friend std::ostream &operator<<(std::ostream &os, const Coeffs &a) {
        os << "{";
        for (auto i = 0; i < N - 1; ++i) os << a.c[i] << ", ";
        os << a.c[N - 1] << "}";
        return os;
    }

    std::array<T, N> c;
};

/// Helper class to hold the range of the approximation and normalizing values to -1..1 and back
template<typename T>
class Range {
public:
    Range() {}
    Range(T min, T max) {
        _mid = 0.5 * (min + max);
        _halfLen = 0.5 * (max - min);
        _invHalfLen = 1 / _halfLen;
    }

    // convert to -1 .. 1 norm
    inline T toNorm(T x) const { return (x - _mid) * _invHalfLen; }

    // convert from -1 .. 1 norm
    inline T fromNorm(T x) const { return x * _halfLen + _mid; }

    inline T min() const { return _mid - _halfLen; }
    inline T max() const { return _mid + _halfLen; }

private:
    T _mid, _halfLen, _invHalfLen;
};

/// Helper to emit the c++ type name.
template<typename T>
constexpr const char *getTypeName() {
    if constexpr (std::is_same<T, double>()) return "double";
    if constexpr (std::is_same<T, float>()) return "float";
    return "unknown";
}

}  // namespace cheby_priv

template<typename T, int N>
class alignas(64) Chebyshev_Approximation_1D {
public:
    Chebyshev_Approximation_1D() {}

    Chebyshev_Approximation_1D(const cheby_priv::Range<T> &range, const std::array<T, N> &coeffs) {
        _range = range;
        _cf.c = coeffs;
    }

    /// returns the approximate max error based on what gets truncated
    template<typename FN>
    T createForFunctionRange(const cheby_priv::Range<T> &range, FN func, int extraSamples = 4) {
        assert(extraSamples >= 0);

        _range = range;

        // evaluate at M points
        auto M = N + extraSamples;
        std::vector<T> values(M);
        auto invM = 1 / T(M);
        for (int i = 0; i < M; ++i) {
            auto y = std::cos(M_PI * (i + 0.5) * invM);
            values[i] = func(_range.fromNorm(y));
        }
        _cf.buildFromValues(values);
        return _cf.estimatedErrorForValues(values);
    }

    /// make beginning and end match a function perfectly
    template<typename FN>
    void normalizeBeginEndFunc(FN func) {
        normalizeBeginEnd(func(_range.min()), func(_range.max()));
    }

    /// make beginning and end match defined values by scaling the polynomial
    void normalizeBeginEnd(const T &value_begin, const T &value_end) {
        auto err_begin = value_begin - _cf.eval(-1);
        auto err_end = value_end - _cf.eval(1);
        auto err_center = 0.5 * (err_begin + err_end);
        auto err_linear = 0.5 * (err_end - err_begin);
        // apply to coeffs to make begin and end match.
        _cf.c[0] += 2 * err_center;
        _cf.c[1] += err_linear;
    }

    /// Emit coefficients of a polynomial equivalent to the approximation.
    /// This is less stable numerically, especially for N >= 12 and approximations
    /// getting close to the maximum type precision.
    /// the coeffients are in ascending order (the first one is for x^0).
    std::array<T, N> polyCoeffs() const {
        // get polygon from -1 .. 1
        auto pcf = _cf.polyCoeffs();
        /// transform to input range
        auto c = 2 / (_range.max() - _range.min());
        auto fact = c;
        for (int i = 1; i < N; ++i) {
            pcf[i] *= fact;
            fact *= c;
        }
        auto shift = 0.5 * (_range.max() + _range.min());
        for (int i = 0; i < N - 1; ++i) {
            for (int j = N - 2; j >= i; --j) pcf[j] -= shift * pcf[j + 1];
        }
        return pcf;
    }

    T eval(const T &x) const {
        auto xn = _range.toNorm(x);         // normalized to -1 .. 1
        assert(xn >= T(-1) && xn <= T(1));  // cannot evaluate outside approximation range
        return _cf.eval(xn);
    }

    friend std::ostream &operator<<(std::ostream &os, const Chebyshev_Approximation_1D &a) {
        if constexpr (std::is_same<T, float>())
            os << std::setprecision(9);
        else
            os << std::setprecision(18);
        os << "Chebyshev_Approximation_1D<" << cheby_priv::getTypeName<T>() << ", " << N << ">(";
        os << "{" << a._range.min() << ", " << a._range.max() << "}, ";
        os << "{" << a._cf << "})";
        return os;
    }

private:
    cheby_priv::Range<T> _range;
    cheby_priv::Coeffs<T, N> _cf;
};

template<typename T, int Nx, int Ny>
class alignas(64) Chebyshev_Approximation_2D {
public:
    Chebyshev_Approximation_2D() {}

    Chebyshev_Approximation_2D(const cheby_priv::Range<T> &rangeX,
        const cheby_priv::Range<T> &rangeY,
        const std::array<const std::array<T, Nx>, Ny> &coeffs) {
        _rangeX = rangeX;
        _rangeY = rangeY;

        for (auto i = 0; i < Ny; ++i) _cf.c[i].c = coeffs[i];
    }

    template<typename FN>
    void createForFunctionRange(const cheby_priv::Range<T> &rangeX,
        const cheby_priv::Range<T> &rangeY,
        FN f,
        int extraSamples = 4) {
        assert(extraSamples >= 0);

        _rangeX = rangeX;
        _rangeY = rangeY;

        // eval at Mx * My points
        auto Mx = Nx + extraSamples;
        auto My = Ny + extraSamples;
        auto invMx = 1 / T(Mx);
        auto invMy = 1 / T(My);

        std::vector<cheby_priv::Coeffs<T, Nx>> valY(My);
        std::vector<T> valX(Mx);

        // pre-calc x coords
        std::vector<T> xmap(Mx);
        for (auto i = 0; i < Mx; ++i) xmap[i] = _rangeX.fromNorm(std::cos(M_PI * (i + 0.5) * invMx));

        for (auto i = 0; i < My; ++i) {
            auto y = _rangeY.fromNorm(std::cos(M_PI * (i + 0.5) * invMy));
            for (auto j = 0; j < Mx; ++j) valX[j] = f(xmap[j], y);

            valY[i].buildFromValues(valX);
        }
        _cf.buildFromValues(valY);
    }

    T eval(T x, T y) const {
        auto xn = _rangeX.toNorm(x), yn = _rangeY.toNorm(y);
        assert(xn >= T(-1) && xn <= T(1) && yn >= T(-1) && yn <= T(1));  // cannot evaluate outside approximation range
        return _cf.eval(yn).eval(xn);
    }

    friend std::ostream &operator<<(std::ostream &os, const Chebyshev_Approximation_2D &a) {
        os << std::setprecision(15);
        os << "Chebyshev_Approximation_2D<" << cheby_priv::getTypeName<T>() << ", " << Nx << ", " << Ny << ">(";
        os << "{" << a._rangeX.min() << ", " << a._rangeX.max() << "}, ";
        os << "{" << a._rangeY.min() << ", " << a._rangeY.max() << "}, ";
        os << "{{";
        for (auto i = 0; i < Ny; ++i) {
            os << "{" << a._cf.c[i] << "}";
            if (i != Ny - 1) os << ", ";
        }
        os << "}})";
        return os;
    }

private:
    cheby_priv::Range<T> _rangeX, _rangeY;
    cheby_priv::Coeffs<cheby_priv::Coeffs<T, Nx>, Ny, T> _cf;
};

}  // namespace plap