gsBairstowSolver.h

#pragma once

#include <gismo.h>

#include "gsMonomialBasis.h"
#include "gsMonomialPoly.h"

namespace gismo
{

template<typename C> class gsBairstowSolver
{
    C eps, b0, c0;
    unsigned int max_iter;
    C* aux;    //this array is needed for division-with-remainder
    int auxLen;
    gsMatrix<C> Q[2];   //Q[0] and Q[1] store the coefficients of the polynomials generated by quadratic deflation
    std::vector< std::vector< std::complex<C> > > roots;   //we cannot use gsMatrix, since the individual lengths may differ
    int status; //0 = nothing happened yet, -1 = failure, 1 = success
    int init_mode;  //0 = (b0, c0), 1 = leading coefs
    bool polish_roots;
    int component;  //auxiliary variable, storing the component of the coefficient vector currently considered

public:

    // default constructor
    gsBairstowSolver() : eps(0.001), b0(0), c0(0), max_iter(50), aux(0), auxLen(0), status(-1), init_mode(0), polish_roots(true) { }

    // destructor
    ~gsBairstowSolver(){
        if (aux) delete[] aux;
    }

    /* \brief This is the main function of class 'gsBairstowSolver'.
       It tries to find the roots of all components of the n-dimensional polynomial 'poly' and stores them
       in vector 'roots' (which can then be accessed by 'getRoots').
       As another side effect, he function sets the value of 'status' to -1 if something went wrong,
       and to 1 if the method succeeded.
       \param poly is the n-dimensional polynomial whose roots shall be computed.
    */
    void solve(const gsMonomialPoly<C> & poly){
        Q[0] = poly.coefs();
        int n = Q[0].cols();
            
        if (n == 0){
            // nothing to do
            status = -1;
            return;
        }
        if (Q[0].rows() > auxLen){
            // resize the auxiliary array 'aux' which is needed for division-with-remainder.
            if (aux) delete[] aux;
            auxLen = Q[0].rows();
            aux = new C[auxLen];
        }

        // resize the vector containing the roots
        if (roots.size() != (unsigned int)n) roots.resize(n);

        /*
          The role of Q[0] and Q[1] is as follows:
          After the 'i'-th quadratic deflation of the 'j'-th component of 'poly',
          Q[i mod 2, j] contains the coefficients of the 'j'-th component of the
          current (i.e. the deflated) polynomial. Q[(i+1) mod 2] is then used in the next
          step to store the quotient, which will in turn become the next deflated polynomial.
          The memory needed by Q[0] and Q[1] is allocated only once, to speed up the computation.
        */
        Q[1] = gsMatrix<C>(Q[0].rows(), n);

        // we solve each component individually
        for (component = 0; component < n; ++component){
            if ((status = solveComponent()) < 0) return;
        }

        if (polish_roots){
            // if requested, the found roots are polished

            /* Since we have to use the already-computed roots as initial values,
               we temporarily set init_mode to 0 (i.e. use b0 and c0 as initial values).
               b0 and c0 are set to the already-computed roots in function 'polishRoots',
               and reset to their original values afterwards. */
            int init_mode_orig(init_mode);
            C b0_orig(b0), c0_orig(c0);
            init_mode = 0;
            Q[0] = poly.coefs();

            // we polish each component individually
            for (component = 0; component < n; ++component){
                if ((status = polishRoots()) < 0) break;
            }

            // init_mode, b0 and c0 are reset to their original values
            init_mode = init_mode_orig;
            b0 = b0_orig;
            c0 = c0_orig;
        }
    }

    /* \brief Returns the roots of the 'i'-th component that were computed by the last call to 'solve'.
       \param i refers to the component of the polynomial whose roots shall be returned.
    */
    std::vector< std::complex<C> > getRoots(unsigned int i){
        GISMO_ASSERT( i >= 0 && i < roots.size(),
                      "Root-vector which is out of range requested.");
        return roots[i];
    }

    /* \brief Returns the roots of the 'i'-th component that were computed by the last call to 'solve'.
       \param i refers to the component of the polynomial whose roots shall be returned
    */
    const std::vector< std::complex<C> > getRoots(unsigned int i) const{
        GISMO_ASSERT( i >= 0 && i < roots.size(),
                      "Root-vector which is out of range requested.");
        return roots[i];
    }

    // \brief Returns the roots of all components that were computed by the last call to 'solve'.
    std::vector< std::vector< std::complex<C> > > getRoots(){
        return this->roots;
    }

    // \brief Returns the roots of all components that were computed by the last call to 'solve'.
    const std::vector< std::vector< std::complex<C> > > getRoots() const{
        return this->roots;
    }

    // \brief Returns the maximum number of Newton-iterations in one step of quadratic deflation.
    unsigned int getMaxIterations() const { return max_iter; }

    /* \brief Sets the maximum number of Newton-iterations in one step of quadratic deflation.
       \param value is the new maximum number of iterations.
    */
    void setMaxIterations(unsigned int value) { max_iter = value; }

    // \brief Returns the tolerance for aborting Newton's method.
    C getEps() const { return eps; }

    /* \brief Sets the tolerance for aborting Newton's method.
       \param value is the new tolerance.
    */
    void setEps(const C & value) { eps = value; }

    // \brief Returns the 'i'-th initial coefficient for Newton's method, where 'i' has to be either 0 or 1.
    C getInitial(int i){
        GISMO_ASSERT( i >= 0 && i <= 1,
                      "Initial coefficient which is out of range requested.");
        switch (i){
        case 0:
            return c0;
        case 1:
            return b0;
        }
        return 0;
    }

    /* \brief Sets the two initial coefficients for Newton's method.
       Note that this only has an effect if the initialization mode is 0 (default).
       \param coef0 is the new value of the constant coefficient.
       \param coef1 is the new value of the coefficient of the linear term.
    */
    void setInitial(const C & coef0, const C & coef1){
        c0 = coef0;
        b0 = coef1;
    }

    // \brief Returns the mode for obtaining the initial values in Newton's method.
    int getInitMode(){ return init_mode; }

    /* \brief Sets the mode for obtaining the initial values in Newton's method.
       \param value is the new initialization mode.
    */
    void setInitMode(int value){ init_mode = value; }

    // \brief Returns the value which indicates whether the computed roots are polished or not.
    int getPolish(){ return polish_roots; }

    /* \brief Sets the value which indicates whether the computed roots are polished or not.
       \param value specifies whether the roots shall be polished or not.
    */
    void setPolish(bool value){ polish_roots = value; }

    /* \brief Returns the status of the last root computation.
       \param[out] The result is 0 if nothing happend yet, -1 if something went wrong, and 1 in case of success.
    */
    int stat() const { return status; }

private:

    /*
      This function computes the roots for a single component. It stores the results in
      member 'roots', and returns a value indicating the status of the computation (see member 'status').
    */
    int solveComponent(){
        /* Compute the degree of the 'component'-th component of Q[0].
           Note that 'component' is a private class member referring to the component currently under consideration. */
        int deg = this->degree(Q[0]);

        if (deg < 0) return -1; // zero polynomial

        roots[component] = std::vector< std::complex<C> >();

        if (deg > 0){
            // we reserve enough space to store all roots
            roots[component].reserve(deg);

            int i = 0, res;

            /* As long as the degree is > 2, we perform quadratic deflation.
               The current polynomial under consideration is Q[i], and the quotient of Q[i] and
               the quadratic factor found will be stored in Q[1 - i]. */
            for (; deg > 2; deg -= 2){
                res = findQuFactor(i, deg);
                if (res < 0) return res;
                i = 1 - i;
            }
                
            // depending on the degree of the remaining factor, we still have to solve a linear/quadratic equation
            switch (deg){
            case 1:
                roots[component].push_back(std::complex<C>(-Q[i](0, component) / Q[i](1, component)));
                break;
            case 2:
                quadraticSolve(Q[i](1, component) / Q[i](2, component), Q[i](0, component) / Q[i](2, component));
                break;
            }
        }

        return 1;
    }

    /*
      This version of 'findQuFactor' tries to find a quadratic factor of Q[i],
      computes its roots, adds them to 'roots', and sets Q[1-i] to the quotient.
      As 'solveComponent', it returns a value indicating the status of the computation.
    */
    int findQuFactor(int i, int d){
        std::complex<C> s1, s2;

        // we basically just call the other version of 'findQuFactor'
        int out = findQuFactor(i, d, s1, s2);
        if (out >= 0){
            roots[component].push_back(s1);
            roots[component].push_back(s2);
        }
        return out;
    }

    /*
      This version of 'findQuFactor' tries to find a quadratic factor of Q[i],
      computes its roots, stores them in 'sol1' and 'sol2', and sets Q[1-i] to the quotient.
      As 'solveComponent', it returns a value indicating the status of the computation.
    */
    int findQuFactor(int i, int d, std::complex<C> & sol1, std::complex<C> & sol2){
        C b, c, b2, c2, r, s, rb, sb, rc, sc, dv, db, dc, eps2(eps * eps);
        int out = -1;

        // depending on init_mode we initialize b and c ...
        switch (init_mode){
        case 1:
            // ... either by the leading coefficients of Q[i], or ...
            b = Q[i](d - 1, component) / Q[i](d, component);
            c = Q[i](d - 2, component) / Q[i](d, component);
            break;
        default:
            // ... by b0 and c0
            b = b0;
            c = c0;
            break;
        }
            
        /* implementation of Bairstow's method according to
           "Numerical Recipes in C: The Art of Scientific Computing", Chapter 9.6, page 378 */

        for (unsigned int k = 0; k < max_iter; ++k){
            // First division-with-remainder. The quotient is stored in Q[1-i], and the remainder is 'aux[1]*x + aux[0]'.
            div(Q[i], d, b, c, Q[1 - i]);
            r = aux[1];
            s = aux[0];

            if (r*r <= eps2 && s*s <= eps2){
                // we can stop immediately, since we've found a quadratic factor
                out = 1;
                break;
            }

            // Second division-with-remainder. The quotient is not needed, and the remainder is again 'aux[1]*x + aux[0]'.
            div(Q[1 - i], d - 2, b, c);

            // partial derivatives
            rc = -aux[1];
            sc = -aux[0];
            sb = -c * rc;
            rb = -b * rc + sc;

            // solve 2x2 system, finally yielding db, dc
            dv = sb * rc - sc * rb;
            if (dv == 0) return -1;

            db = (r * sc - s * rc) / dv;
            dc = (-r * sb + s * rb) / dv;

            // update b and c
            b += db;
            c += dc;

            if ((db*db <= eps2*(b2 = b*b) || b2 <= eps2) && (dc*dc <= eps2*(c2 = c*c) || c2 <= eps2)){
                out = 1;
                break;
            }
        }

        /* If we have found a quadratic factor, we add its roots to the root-vector.
           The quotient of Q[i] and the quadratic factor found is stored in Q[1-i]. */
        if (out > 0) quadraticSolve(b, c, sol1, sol2);
        return out;
    }

    /*
      This function polishes the roots of a single component.
      It does so by running again Bairstow's method for finding quadratic factors *of the original input polynomial*,
      using the already-computed roots as initial values.
    */
    int polishRoots(){
        if (roots[component].size() > 2){   //no need to polish the first two roots
            int deg = roots[component].size(), out;
            for (typename std::vector< std::complex<C> >::iterator it = roots[component].begin() + 2;
                 it != roots[component].end() && it + 1 != roots[component].end();
                 it += 2)
            {
                // set the initial values
                b0 = -(it->real() + (it + 1)->real());
                c0 = it->real() * (it + 1)->real() - it->imag() * (it + 1)->imag();

                /* We are not interested in the quadratic factor or the quotient,
                   but 'findQuFactor' also updates roots *it and *(it+1). */
                out = findQuFactor(0, deg, *it, *(it + 1));
                if (out < 0) return out;
            }
        }

        return 1;
    }

    /*
      This functions computes the degree of the 'component'-th component of the
      coefficient vector 'coefs'.
      Note that the degree might be strictly less than the length of the
      vector, since some leading coefficients might be 0.
    */
    short_t degree(const gsMatrix<C> & coefs){
        for (short_t d = coefs.rows() - 1; d >= 0; --d){
            if (coefs(d, component) != 0) return d;
        }
        return -1;
    }

    /*
      This version of 'quadraticSolve' solves the quadratic equation 'x^2 + b*x + c == 0'
      and stores the two solutions in vector 'roots' (at component 'component').
    */
    void quadraticSolve(const C & b, const C & c){
        std::complex<C> s1, s2;

        // we basically just call the other version of 'quadraticSolve'
        quadraticSolve(b, c, s1, s2);
        roots[component].push_back(s1);
        roots[component].push_back(s2);
    }

    /*
      This version of 'quadraticSolve' solves the quadratic equation 'x^2 + b*x + c == 0'
      and stores the two solutions in 's1' and 's2'.
    */
    void quadraticSolve(const C & b, const C & c, std::complex<C> & s1, std::complex<C> & s2){
        // straight-forward solving of 'x^2 + b*x + c == 0' by x_{1,2} = (-b +- sqrt(b^2 - 4*c)) / 2
        C discriminant(b * b - C(4) * c);
        C TWO = C(2);
        if (discriminant >= 0)
        {
            discriminant = sqrt(discriminant);
            s1 = std::complex<C>((-b - discriminant) / TWO);
            s2 = std::complex<C>((-b + discriminant) / TWO);
        }else{
            discriminant = sqrt(-discriminant) / TWO;
            s1 = std::complex<C>(-b / TWO, -discriminant);
            s2 = std::complex<C>(-b / TWO, discriminant);
        }
    }

    /*
      This version of 'div' divides an arbitrary polynomial 'p', given as a coefficient-vector,
      by a quadratic polynomial 'x^2+b*x+c'. The coefficient-vector of the quotient is stored in
      'quo', and the coefficients of the remainder are stored in the private class member 'aux';
      This array is also used for storing intermediate results.
    */
    void div(const gsMatrix<C> & p, int dp, const C & b, const C & c, gsMatrix<C> & quo){
        GISMO_ASSERT( p.rows() > dp && auxLen > dp && quo.rows() > dp - 2,
                      "Invalid division-with-remainder in gsBairstow.");
            
        /* division-with-remainder according to
           "Numerical Recipes in C: The Art of Scientific Computing", Chapter 5.3, page 175 */
        int k;
        for (k = 0; k <= dp; ++k) aux[k] = p(k, component);
        for (k = dp - 2; k >= 0; --k){
            quo(k, component) = aux[k + 2];
            aux[k + 1] -= quo(k, component) * b;
            aux[k] -= quo(k, component) * c;
        }
    }

    /*
      This version of 'div' divides an arbitrary polynomial 'p', given as a coefficient-vector,
      by a quadratic polynomial 'x^2+b*x+c'.
      Unlike with the other version above, no quotient, but only the remainder is computed.
    */
    void div(const gsMatrix<C> & p, int dp, const C & b, const C & c){
        GISMO_ASSERT( p.rows() > dp && auxLen > dp,
                      "Invalid division-with-remainder in gsBairstow.");
            
        /* division-with-remainder according to
           "Numerical Recipes in C: The Art of Scientific Computing", Chapter 5.3, page 175 */
        int k;
        for (k = 0; k <= dp; ++k) aux[k] = p(k, component);
        for (k = dp - 2; k >= 0; --k){
            aux[k + 1] -= aux[k + 2] * b;
            aux[k] -= aux[k + 2] * c;
        }
    }
};  //gsBairstowSolver

}   //namespave gismo