rbf.h

#ifndef SOT_RBF_H
#define SOT_RBF_H

#include "common.h"
#include "utils.h"
#include "surrogate.h"

namespace sot {


    class Kernel {
        virtual inline int order() const = 0; 
        virtual inline int phiZero() const = 0; 

        virtual inline double eval(double dist) const = 0;

        virtual inline double deriv(double dist) const = 0;

        virtual inline mat eval(const mat &dists) const = 0;

        virtual inline mat deriv(const mat &dists) const = 0;
    };
    

    class CubicKernel : public Kernel {
    private:
        int mPhiZero = 0; 
        int mOrder = 2; 
    public:
        inline int order() const {
            return mOrder; }
        inline int phiZero() const { return mPhiZero;
        }
        inline double eval(double dist) const {
            return dist * dist * dist;
        }
        inline double deriv(double dist) const {
            return 3 * dist * dist;
        }
        inline mat eval(const mat &dists) const {
            return dists % dists % dists;
        }
        inline mat deriv(const mat &dists) const {
            return 3 * dists % dists;
        }
    };
    

    class ThinPlateKernel : public Kernel {
    private:
        int mPhiZero = 0; 
        int mOrder = 2; 
    public:
        inline int order() const {
            return mOrder;
        }
        inline int phiZero() const {
            return mPhiZero;
        }
        inline double eval(double dist) const {
            return dist * dist * std::log(dist + 1e-12);
        }
        inline double deriv(double dist) const {
            return dist * (1.0 + 2.0 * std::log(dist + 1e-12));
        }
        inline mat eval(const mat &dists) const {
            return dists % dists % arma::log(dists + 1e-12);
        }
        inline mat deriv(const mat &dists) const {
            return dists % (1 + 2.0 * arma::log(dists + 1e-12));
        }
    };
    

    class LinearKernel : public Kernel {
    private:
        int mPhiZero = 0; 
        int mOrder = 1; 
    public:
        inline int order() const {
            return mOrder;
        }
        inline int phiZero() const {
            return mPhiZero;
        }
        inline double eval(double dist) const {
            return dist;
        }
        inline double deriv(double dist) const {
            return 1.0;
        }
        inline mat eval(const mat &dists) const {
            return dists;
        }
        inline mat deriv(const mat &dists) const { 
            return arma::ones<mat>(dists.n_rows, dists.n_cols); 
        }
    };
    

    class Tail {
        virtual inline int degree() const = 0;  

        virtual inline int dimTail(int dim) const = 0;

        virtual inline vec eval(const vec &point) const = 0;

        virtual inline mat eval(const mat &points) const = 0;

        virtual inline mat deriv(const vec &point) const = 0;
    };
    

    class LinearTail : public Tail {
    private:
        int mDegree = 1; 
    public:
        inline int degree() const {
            return mDegree;
        }
        inline mat eval(const mat &X) const {
            return arma::join_vert(arma::ones<mat>(1, X.n_cols), X);
        }
        inline vec eval(const vec &x) const {
            vec tail = arma::ones<vec>(x.n_rows + 1);
            tail.tail(x.n_rows) = x;
            return tail;
        }
        inline mat deriv(const vec &x) const {
            return arma::join_vert(arma::zeros<mat>(1, x.n_rows),
                                   arma::eye<mat>(x.n_rows, x.n_rows));
        }
        inline int dimTail(int dim) const {
            return 1 + dim;
        }
    };
    

    class ConstantTail : public Tail {
    private:
        int mDegree = 0; 
    public:
        inline int degree() const {
            return mDegree;
        }
        inline mat eval(const mat &X) const {
            return arma::ones<mat>(1, X.n_cols);
        }
        inline vec eval(const vec &x) const {
            return arma::ones<mat>(1, 1);
        }
        inline mat deriv(const vec &x) const {
            return arma::zeros<mat>(1, x.n_rows);
        }
        inline int dimTail(int dim) const {
            return 1;
        }
    };
    

    template<class RBFKernel = CubicKernel, class PolyTail = LinearTail>
    class RBFInterpolant : public Surrogate {
        
    protected:
        RBFKernel mKernel; 
        PolyTail mTail; 
        mat mL; 
        mat mU; 
        uvec mp; 
        vec mCoeffs; 
        vec mfX; 
        mat mCenters; 
        int mMaxPoints; 
        int mDimTail; 
        int mNumPoints; 
        bool mDirty; 
        vec mxLow; 
        vec mxUp; 
        double mEta = 1e-8; 
        int mDim; 

        void setPoints(const mat &ppoints, const vec &funVals) {
            
            // Map point to be in the unit box
            mat points = toUnitBox(ppoints, mxLow, mxUp);
            
            mNumPoints = (int)points.n_cols;
            int n = mNumPoints + mDimTail;
            if(mNumPoints < mDimTail) { 
                throw std::logic_error("Not enough points"); 
            }
            mat px = mTail.eval(points);
            mat phi = mKernel.eval(arma::sqrt(squaredPairwiseDistance(points, points)));
                    
            mat A = arma::zeros<mat>(n, n);
            A(arma::span(mDimTail, n - 1), arma::span(0, mDimTail - 1)) = px.t();
            A(arma::span(0, mDimTail - 1), arma::span(mDimTail, n - 1)) = px;
            A(arma::span(mDimTail, n - 1), arma::span(mDimTail, n - 1)) = phi;
            mfX.rows(mDimTail, n - 1) = funVals;
            
            // REGULARIZATION
            A += mEta*arma::eye(n, n);
            
            // Compute the initial LU factorization of A
            mat L, U, P;
            arma::lu(L, U, P, A);
            
            mL(arma::span(0, n - 1), arma::span(0, n - 1)) = L;
            mU(arma::span(0, n - 1), arma::span(0, n - 1)) = U;
            
            // Convert P to a permutation vector
            for(int i = 0; i < n; i++) {
                uvec temp = find(P.row(i) > 0.5);
                mp(i) = temp(0);
            }
            
            mCenters.cols(0, mNumPoints - 1) = points;
            mDirty = true;
        }
        
    public:

        RBFInterpolant(int maxPoints, int dim) :
                RBFInterpolant(maxPoints, dim, arma::zeros(dim), arma::ones(dim)) {}

        RBFInterpolant(int maxPoints, int dim, double eta) :
                RBFInterpolant(maxPoints, dim, arma::zeros(dim), arma::ones(dim), eta) {}
        

        RBFInterpolant(int maxPoints, int dim, vec xLow, vec xUp, double eta) :
                RBFInterpolant(maxPoints, dim, xLow, xUp) {
            mEta = eta;
        }

        RBFInterpolant(int maxPoints, int dim, vec xLow, vec xUp) {
            mMaxPoints = maxPoints;
            mDim = dim;
            mDimTail = mTail.dimTail(dim);
            mNumPoints = 0;
            mCenters = arma::zeros<mat>(dim, maxPoints);
            mL = arma::zeros<mat>(maxPoints + mDimTail, maxPoints + mDimTail);
            mU = arma::zeros<mat>(maxPoints + mDimTail, maxPoints + mDimTail);
            mp = arma::zeros<uvec>(maxPoints + mDimTail);
            mfX = arma::zeros<vec>(maxPoints + mDimTail);
            mCoeffs = vec(maxPoints + mDimTail);
            mDirty = false;
            mxLow = xLow;
            mxUp = xUp;
            
            if ((mKernel.order() - 1) > mTail.degree()) {
                throw std::logic_error("Kernel and tail mismatch");
            }
        }
        
        int dim() const {
            return mDim;
        }
        
        // Reset RBF
        void reset() {
            mNumPoints = 0;
        }
        
        // Number of points
        int numPoints() const {
            return mNumPoints;
        }
        
        // Return points
        mat X() const {
            return fromUnitBox((mat)mCenters.cols(0, mNumPoints - 1), mxLow, mxUp);
        }
        
        // Return point
        vec X(int i) const {
            return fromUnitBox((mat)mCenters.col(i), mxLow, mxUp);
        }
        
        // Return function values
        vec fX() const {
            return mfX.rows(mDimTail, mDimTail + mNumPoints - 1);
        }
        
        // Return function value
        double fX(int i) const {
            return mfX(mDimTail + i);
        }
        

        vec coeffs() {
            if(mDirty) { throw std::logic_error("RBF not updated"); }
            return mCoeffs;
        }
       
        // Fit the RBF
        void fit() {
            if(mNumPoints < mDimTail) {
                throw std::logic_error("Not enough points");
            }
            if (mDirty) {
                int n = mNumPoints + mDimTail;
                mCoeffs = arma::solve(arma::trimatl(mL(arma::span(0, n - 1), arma::span(0, n - 1))), mfX(mp(arma::span(0, n - 1))));
                mCoeffs = arma::solve(arma::trimatu(mU(arma::span(0, n - 1), arma::span(0, n - 1))), mCoeffs);
                mDirty = false;
            }
        }
        

        void addPoint(const vec &ppoint, double funVal) {
            if(mNumPoints == 0) {
                vec fVal = {funVal};
                mat point = (mat)ppoint;
                return setPoints(point, fVal);
            }

            // Map point to be in the unit box
            vec point = toUnitBox(ppoint, mxLow, mxUp);
            
            int nAct = mDimTail + mNumPoints;
            if(mNumPoints + 1 > mMaxPoints) { 
                throw std::logic_error("Capacity exceeded"); 
            }
      
            vec vx = arma::join_vert(mTail.eval(point), 
                    mKernel.eval(arma::sqrt(squaredPointSetDistance(point, mCenters.cols(0, mNumPoints - 1)))));
            vec u12 = arma::solve(arma::trimatl(mL(arma::span(0, nAct - 1), arma::span(0, nAct - 1))), vx.rows(mp.head(nAct)));
            vec l21 = (arma::solve(arma::trimatl(mU(arma::span(0, nAct - 1), arma::span(0, nAct - 1)).t()), vx));
            double u22 = mKernel.phiZero() + mEta - arma::dot(u12, l21);

            mL(nAct, arma::span(0, nAct - 1)) = l21.t();
            mL(nAct, nAct) = 1;
            mU(arma::span(0, nAct - 1), nAct) = u12;
            mU(nAct, nAct) = u22;
            mp(nAct) = nAct;

            // Update F and add the centers
            mfX(nAct) = funVal;
            mCenters.col(mNumPoints) = point;
            mNumPoints++;
            
            mDirty = true;
        }
        

        void addPoints(const mat &ppoints, const vec &funVals) {
            if(mNumPoints == 0) {
                return setPoints(ppoints, funVals);
            }
            
            // Map point to be in the unit box
            mat points = toUnitBox(ppoints, mxLow, mxUp);
            
            int nAct = mDimTail + mNumPoints;
            int n = (int)funVals.n_rows;
            if(n < 2) {
                throw std::logic_error("Use add_point instead");
            }
            if(mNumPoints + n > mMaxPoints) {
                throw std::logic_error("Capacity exceeded");
            }
            
            auto px = mTail.eval(points);
            mat B = arma::zeros(nAct, n);
            B.rows(arma::span(0, mDim)) = px;
            B.rows(mDimTail, nAct - 1) = mKernel.eval(arma::sqrt(squaredPairwiseDistance<mat>(mCenters.cols(0, mNumPoints - 1), points)));
            mat K = mKernel.eval(arma::sqrt(squaredPairwiseDistance<mat>(points, points)));
            
            // REGULARIZATION
            K += mEta*arma::eye(n, n);

            // Update the LU factorization
            mat U12 = arma::solve(arma::trimatl(mL(arma::span(0, nAct - 1), arma::span(0, nAct - 1))), B.rows(mp.head(nAct)));
            mat L21 = (arma::solve(arma::trimatl(mU(arma::span(0, nAct - 1), arma::span(0, nAct - 1)).t()), B)).t();
            mat C;
            try {
                C = arma::chol(K - L21*U12, "lower");
            }
            catch (std::runtime_error) {
                std::cout << "Warning: Cholesky factorization failed, computing new LU from scratch..." << std::endl;
                // Add new points
                mfX.rows(nAct, nAct + n - 1) = funVals;
                mCenters.cols(mNumPoints, mNumPoints + n - 1) = points;
                mNumPoints += n;
                // Build LU from scratch
                setPoints(X(), mfX.rows(mDimTail, nAct + n - 1));
                return;
            }
            mL(arma::span(nAct, nAct + n - 1), arma::span(0, nAct - 1)) = L21;
            mL(arma::span(nAct, nAct + n - 1), arma::span(nAct, nAct + n - 1)) = C;
            mU(arma::span(0, nAct - 1), arma::span(nAct, nAct + n - 1)) = U12;
            mU(arma::span(nAct, nAct + n - 1), arma::span(nAct, nAct + n - 1)) = C.t();
            mp.rows(arma::span(nAct, nAct + n - 1)) = arma::linspace<uvec>(nAct, nAct + n - 1, n);
            
            // Update F and add the centers
            mfX.rows(nAct, nAct + n - 1) = funVals;
            mCenters.cols(mNumPoints, mNumPoints + n - 1) = points;
            mNumPoints += n;
            
            mDirty = true;
        }
        

        double eval(const vec &ppoint) const {
            if(mDirty) { throw std::logic_error("RBF not updated. You need to call fit() first"); }       

            // Map point to be in the unit box
            vec point = toUnitBox(ppoint, mxLow, mxUp);
            
            vec px = mTail.eval(point);
            vec phi = mKernel.eval(arma::sqrt(squaredPointSetDistance(point, (mat)mCenters.cols(0, mNumPoints - 1))));
            vec c = mCoeffs.head(mNumPoints + mDimTail);
            return arma::dot(c.head(mDimTail), px) + arma::dot(c.tail(mNumPoints), phi);
        }
        

        double eval(const vec &ppoint, const vec &dists) const {
            if(mDirty) { throw std::logic_error("RBF not updated. You need to call fit() first"); }       
            if(not (arma::all(mxLow == 0) and arma::all(mxUp == 1))) { 
                std::cout << "RBF uses internal scaling so distances have to be recomputed" << std::endl;
                return eval(ppoint);
            }     
            
            // Map point to be in the unit box
            vec point = toUnitBox(ppoint, mxLow, mxUp);
            
            vec px = mTail.eval(point);
            vec phi = mKernel.eval(dists);
            vec c = mCoeffs.head(mNumPoints + mDimTail);
            return arma::dot(c.head(mDimTail), px) + arma::dot(c.tail(mNumPoints), phi);
        }
        

        vec evals(const mat &ppoints) const {
            if(mDirty) { throw std::logic_error("RBF not updated. You need to call fit() first"); }       

            // Map point to be in the unit box
            mat points = toUnitBox(ppoints, mxLow, mxUp);
            
            mat px = mTail.eval(points);
            mat phi = mKernel.eval(arma::sqrt(squaredPairwiseDistance(
                    (mat)mCenters.cols(0, mNumPoints - 1), points)));
            vec c = mCoeffs.head(mNumPoints + mDimTail);
            return px.t() * c.head(mDimTail) + phi.t() * c.tail(mNumPoints);
        }
        

        vec evals(const mat &ppoints, const mat &dists) const {
            if(mDirty) { throw std::logic_error("RBF not updated. You need to call fit() first"); }       
            if(not (arma::all(mxLow == 0) and arma::all(mxUp == 1))) { 
                std::cout << "RBF uses internal scaling so distances have to be recomputed" << std::endl;
                return evals(ppoints);
            }       

            // Map point to be in the unit box
            mat points = toUnitBox(ppoints, mxLow, mxUp);
            
            mat px = mTail.eval(points);
            mat phi = mKernel.eval(dists);
            vec c = mCoeffs.head(mNumPoints + mDimTail);
            return px.t() * c.head(mDimTail) + phi.t() * c.tail(mNumPoints);
        }
        

        vec deriv(const vec &ppoint) const {
            if(mDirty) { throw std::logic_error("RBF not updated. You need to call fit() first"); }       

            // Map point to be in the unit box
            vec point = toUnitBox(ppoint, mxLow, mxUp);
            
            mat dpx = mTail.deriv(point);
            vec c = mCoeffs.head(mNumPoints + mDimTail);
            vec dists = arma::sqrt(squaredPairwiseDistance<mat>(
                    mCenters.cols(0, mNumPoints - 1), point));
            dists.elem(arma::find(dists < 1e-10)).fill(1e-10); // Better safe than sorry
            mat dsx = - mCenters.cols(0, mNumPoints - 1);
            dsx.each_col() += point;
            dsx.each_row() %= (mKernel.deriv(dists) % c.tail(mNumPoints)/dists).t();
            return arma::sum(dsx, 1) + dpx.t() * c.head(mDimTail);
        }
    };


    template<class RBFKernel, class PolyTail>
    class RBFInterpolantCap : public RBFInterpolant<RBFKernel,PolyTail> {
    public:

        RBFInterpolantCap(int maxPoints, int dim, vec xLow, vec xUp, double eta) :
            RBFInterpolant<RBFKernel,PolyTail>(maxPoints, dim, xLow, xUp, eta) {}

        RBFInterpolantCap(int maxPoints, int dim, vec xLow, vec xUp) :
            RBFInterpolant<Kernel,Tail>(maxPoints, dim, xLow, xUp) {}

        void fit() {
            if(this->mNumPoints < this->mDimTail) {
                throw std::logic_error("Not enough points");
            }
            if (this->mDirty) {
                int n = this->mNumPoints + this->mDimTail;
                vec ff = this->mfX;
                double medf = arma::median(ff.rows(this->mDimTail, n-1)); // Computes the median
                for(int i=this->mDimTail; i < n; i++) { // Apply the capping
                    if (ff(i) > medf) {
                        ff(i) = medf;
                    }
                }
                
                // Solve for the coefficients
                this->mCoeffs = arma::solve(arma::trimatl(
                        this->mL(arma::span(0, n - 1), arma::span(0, n - 1))), ff(this->mp(arma::span(0, n - 1))));
                this->mCoeffs = arma::solve(arma::trimatu(
                        this->mU(arma::span(0, n - 1), arma::span(0, n - 1))), this->mCoeffs );
                this->mDirty = false;
            }
        }
    };      
}

#endif