rootfinding.f90

Go to the documentation of this file.

!##############################################################################
!     ________  _____  ______________ 
!    / ____/  |/  / / / /_  __/  _/ / 
!   / /_  / /|_/ / / / / / /  / // /  
!  / __/ / /  / / /_/ / / / _/ // /___
! /_/   /_/  /_/\____/ /_/ /___/_____/                                     
!
! Copyright 2020 Bharat Mahajan
!
! Licensed under the Apache License, Version 2.0 (the "License");
! you may not use this file except in compliance with the License.
! You may obtain a copy of the License at
!
!    http://www.apache.org/licenses/LICENSE-2.0
!
! Unless required by applicable law or agreed to in writing, software
! distributed under the License is distributed on an "AS IS" BASIS,
! WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
! See the License for the specific language governing permissions and
! limitations under the License.
!
!
!##############################################################################
 
    
module rootfinding
 
    use fmutilbase
    
    use lapack95, only: gebal, hseqr
 
    implicit none
 
    integer, parameter :: max_iterations = 50
    
    abstract interface
        pure function func(x) result(y)
            import :: wp
            implicit none
            real(wp), intent(in) :: x
            real(wp) :: y
        end function func
    end interface
    
    
    contains
    
    
    subroutine polyroots(c, r, error, BalanceOn)
    
        implicit none
        
        complex(WP), dimension(1:), intent(in)             :: c
        
        complex(WP), dimension(:), allocatable, intent(out)    :: r
        
        integer, intent(out)                            :: error
        
        logical, intent(in), optional           :: BalanceOn
        
        integer :: n, nnz, ctr
        integer, dimension(1) :: nzlocl, nzloct
        complex(WP), dimension(:), allocatable :: cnz
        complex(WP), dimension(:,:), allocatable :: H
        logical :: UseBalance
        integer :: ilo, ihi, info
        real(WP), dimension(:), allocatable, target :: scale
        
        error = 0 ! all ok so far!
        
        ! original polynomial degree
        n = size(c) - 1
        
        if (n <= 0) then
            error = -1
            return
        end if
        
        if (all(abs(c) == 0) .EQV. .true.) then
            error = -2
            return
        end if
    
        ! By default, balance the companion matrix
        usebalance = .true.
        if (present(balanceon)) usebalance = balanceon
        
        ! find the location of the first nonzero coefficient at the front and back
        block
            integer, dimension(size(c)) :: c1
            
            c1 = merge(1,0,real(c) /= 0 .OR. aimag(c) /= 0)
            nzlocl = findloc(c1, 1)
            nzloct = findloc(c1, 1, back=.true.)
        end block
    
        nnz = nzloct(1) - nzlocl(1)
        
        ! copy the nonzero coefficients
        cnz = c(nzlocl(1):nzloct(1))
        
        ! Relatively small leading coefficients are removed to avoid infinite
        ! values similar to MATLAB's "root" command, ctr will contain the
        ! location of the first valid significant coefficient
        ctr = 1
        do while (any(abs(cnz(ctr+1:)/cnz(ctr)) > inf))
            ctr = ctr + 1
            if (ctr == ubound(cnz,1)) exit
        end do
        nnz = nnz - (ctr - 1)
        
        ! Final polynomial degree is less than 1, so it has only zero roots
        ! corresponding to the trailing zero coefficients
        if (nnz <= 0) then
            r = [(0, ctr = 1,(n+1-nzloct(1)))]
            return
        elseif (nnz == 1) then ! Special case: 1st degree polynomial
            r = [complex(WP):: -cnz(1+ctr)/cnz(ctr), (0, ctr = 1,(n+1-nzloct(1)))]
            return
        end if
    
        ! Form the companion matrix that is an upper Hessenberg matrix
        allocate(h(nnz,nnz))
        h = (0,0)
        h(1,:) = -cnz(ctr+1:)/cnz(ctr)
        do ctr = 2,ubound(h,1)
            h(ctr, ctr-1) = (1,0)
        end do
    
        ! Balance the companion matrix for improved accuracy
        if (usebalance) then
            allocate(scale(nnz))
            call gebal(h, ilo=ilo, ihi=ihi, job='B',info=info)
            if (info /= 0) then
                ! balancing failed, return error
                error = -3
                return
            end if
        else
            ilo = 1
            ihi = nnz
        end if
        
        ! We have a balanced upper Hessenberg companion matrix, therefore
        ! we can use this LAPACK computational routine directly
        allocate(r(nnz))
        call hseqr(h, r, ilo=ilo, ihi=ihi, job='E', info=info)
        !call geevx(H, r, balanc='B', info=info)
    
        if (info /= 0) error = -4        
        
        ! Gather all the roots
        r = [complex(wp) :: r, (0, ctr = 1,(n+1-nzloct(1)))]
    
    end subroutine polyroots
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    pure subroutine root(a, b, ya, yb, tol, f, x, fx, niter, error)
    
        implicit none
        
        real(wp), intent(in)                 :: a
        real(wp), intent(in)                 :: b
        real(wp), intent(in)                 :: ya
        real(wp), intent(in)                 :: yb
        real(wp), intent(in)                 :: tol
        procedure(func)                      :: f
        real(wp), intent(out)                :: x
        real(wp), intent(out)                :: fx
        integer, intent(out)                 :: niter
        
        integer, intent(out)                 :: error 
 
        real(wp) :: aa, bb, cc, fa, fb, fc, xm, ss, dd, ee, pp, qq, rr
        
        real(wp) :: tolx
        
        real(wp), parameter :: nearzero = 1.0e-20_wp
        
        logical :: done
        integer :: itr
        
        aa = a
        bb = b
        fa = ya
        fb = yb
 
        error = 0
        niter = 0
 
        ! test
        fa = f(a)
        fb = f(b)
        
        ! Trivial cases
        if (abs(fa) <= nearzero) then
            x = aa
            fx = fa
            return
        else if (abs(fb) <= nearzero) then
            x = bb
            fx = fb
            return
        end if
        
        ! if root is not bracketed then return error
        if ((fa>0 .AND. fb>0) .OR. (fa<0 .AND. fb<0)) then
            error = 1
            niter = 0
            return
        end if
        
        ! algorithm start here
        fc = fb
        done = .false.
        itr = 0
        do while (done .EQV. .false. .AND. itr < max_iterations)
            
            ! if root is NOT bracketed between C and B, then C = A
            if ((fc>0 .AND. fb>0) .OR. (fc<0 .AND. fb<0)) then
                cc = aa
                fc = fa
                dd = bb - aa
                ee = dd
            end if
            ! Make sure |f(c)| is bigger than |f(b)|
            if (abs(fc) < abs(fb)) then
                aa = bb
                bb = cc
                cc = aa
                fa = fb
                fb = fc
                fc = fa
            end if
            ! tolerance on the independent variable
            tolx = 2.0*eps*abs(bb) + 0.5*tol
            ! middle point
            xm = 0.5*(cc - bb)
            
            if (abs(xm) <= tolx .OR. abs(fa) < nearzero) then
                ! root has been found
                x = bb
                done = .true.
                fx = f(x)
            else
                if (abs(ee) >= tolx .AND. abs(fa) > abs(fb)) then
                    ss = fb/fa ! make sure |ss| < 1
                    if (abs(aa - cc) < nearzero) then
                        ! linear interpolation
                        pp = 2.0*xm*ss
                        qq = 1.0 - ss
                    else
                        ! Inverse quadratic interpolation
                        qq = fa/fc
                        rr = fb/fc
                        pp = ss*(2.0*xm*qq*(qq-rr)-(bb-aa)*(rr-1.0))
                        qq = (qq - 1.0)*(rr - 1.0)*(ss - 1.0)
                    end if
                    
                    if (pp > nearzero) qq = -qq
                    pp = abs(pp)
                    if ((2.0*pp) < min(3.0*xm*qq-abs(tolx*qq), abs(ee*qq))) then
                        ! use linear or inverse quadratic interpolation
                        ee = dd
                        dd = pp/qq
                    else
                        ! use bisection
                        dd = xm
                        ee = dd
                    end if
                else
                    ! use bisection
                    dd = xm
                    ee = dd
                end if
 
                aa = bb
                fa = fb
                if (abs(dd) > tolx) then
                    bb = bb + dd
                else
                    ! Correction term too small, advance by at least tolx
                    if (xm > 0) then
                        bb = bb + abs(tolx)
                    else
                        bb = bb - abs(tolx)
                    end if
                end if
                fb = f(bb)
                itr = itr + 1
            end if
        end do
        
        if (itr >= max_iterations) error = 2
        niter = itr
 
    end subroutine root
 
    
    
end module rootfinding