! ch13_gradient_descent.f95
! Companion code for Chapter 13: Optimization and Data Fitting.
!
! Local minimization by gradient descent with backtracking step control.
! Change the objective, gradient, starting point, and settings below.
! Variables are stored in a vector; the routine accepts any dimension.
! A small gradient indicates a stationary point, not necessarily a minimum.
! Scale variables sensibly: the gradient tolerance depends on their units.
!
! Compile (Fortran 95):
! gfortran -std=f95 -Wall -Wextra -O2 ch13_gradient_descent.f95 -o ch13_gradient_descent
! Run: ./ch13_gradient_descent (Windows: ch13_gradient_descent.exe)

module gradient_descent_methods
    implicit none
    private
    integer, parameter :: dp = selected_real_kind(15,307)
    public :: dp, descent_result, gradient_descent

    type :: descent_result
        real(dp) :: value, gradient_norm
        integer :: iterations, function_calls, gradient_calls
        logical :: converged
        character(len=90) :: status
    end type descent_result
contains
    subroutine gradient_descent(f,gradf,p,step_initial,tol_grad,max_iter,result)
        interface
            function f(p) result(value)
                integer, parameter :: dp = selected_real_kind(15,307)
                real(dp), intent(in) :: p(:)
                real(dp) :: value
            end function f
            subroutine gradf(p,g)
                integer, parameter :: dp = selected_real_kind(15,307)
                real(dp), intent(in) :: p(:)
                real(dp), intent(out) :: g(:)
            end subroutine gradf
        end interface
        real(dp), intent(inout) :: p(:)
        real(dp), intent(in) :: step_initial, tol_grad
        integer, intent(in) :: max_iter
        type(descent_result), intent(out) :: result
        real(dp) :: g(size(p)), trial_p(size(p)), fp, trial_f, alpha, gnorm
        real(dp), parameter :: armijo = 1.0e-4_dp
        integer, parameter :: max_backtrack = 50
        integer :: trial
        logical :: accepted

        if (size(p)==0 .or. .not. all(finite_value(p))) stop 'Invalid starting point.'
        if (.not. finite_value(step_initial) .or. step_initial<=0.0_dp) stop 'Invalid initial step.'
        if (.not. finite_value(tol_grad) .or. tol_grad<=0.0_dp) stop 'Invalid gradient tolerance.'
        if (max_iter<0) stop 'Iteration limit must be nonnegative.'
        fp=f(p)
        if (.not. finite_value(fp)) stop 'Objective must be finite at starting point.'
        call gradf(p,g)
        if (.not. all(finite_value(g))) stop 'Gradient must be finite.'
        gnorm=vector_norm(g)
        result%function_calls=1
        result%gradient_calls=1
        result%iterations=0
        result%status='Maximum iteration count reached.'

        do while (gnorm>tol_grad .and. result%iterations<max_iter)
            alpha=step_initial
            accepted=.false.
            do trial=1,max_backtrack
                trial_p=p-alpha*g
                if (all(trial_p==p)) exit
                if (all(finite_value(trial_p))) then
                    trial_f=f(trial_p)
                    result%function_calls=result%function_calls+1
                    if (finite_value(trial_f)) then
                        ! Armijo sufficient-decrease test for direction -g.
                        if (trial_f<=fp-armijo*alpha*sum(g*g)) then
                            accepted=.true.
                            exit
                        end if
                    end if
                end if
                alpha=alpha/2.0_dp
            end do
            if (.not. accepted) then
                result%status='Stopped: backtracking could not find a usable step.'
                exit
            end if
            p=trial_p
            fp=trial_f
            call gradf(p,g)
            result%gradient_calls=result%gradient_calls+1
            if (.not. all(finite_value(g))) stop 'Gradient must be finite.'
            gnorm=vector_norm(g)
            result%iterations=result%iterations+1
        end do
        result%value=fp
        result%gradient_norm=gnorm
        result%converged=gnorm<=tol_grad
        if (result%converged) result%status='Gradient-norm tolerance reached.'
    end subroutine gradient_descent

    function vector_norm(x) result(value)
        real(dp), intent(in) :: x(:)
        real(dp) :: value, scale
        ! Scale before squaring to avoid avoidable overflow or underflow.
        scale=maxval(abs(x))
        value=0.0_dp
        if (scale>0.0_dp) value=scale*sqrt(sum((x/scale)**2))
    end function vector_norm

    elemental logical function finite_value(x)
        real(dp), intent(in) :: x
        finite_value=abs(x)<=huge(x)
    end function finite_value
end module gradient_descent_methods

! Edit both the objective and its matching analytical gradient here.
module example_problem
    use gradient_descent_methods, only: dp
    implicit none
contains
    function objective(p) result(value)
        real(dp), intent(in) :: p(:)
        real(dp) :: value
        value=p(1)**2+4.0_dp*p(2)**2+2.0_dp*p(1)*p(2)
    end function objective

    subroutine gradient(p,g)
        real(dp), intent(in) :: p(:)
        real(dp), intent(out) :: g(:)
        g(1)=2.0_dp*p(1)+2.0_dp*p(2)
        g(2)=8.0_dp*p(2)+2.0_dp*p(1)
    end subroutine gradient
end module example_problem

program ch13_gradient_descent
    use gradient_descent_methods, only: dp, descent_result, gradient_descent
    use example_problem, only: objective, gradient
    implicit none

    ! Editable settings. Example minimum: p = (0,0), f = 0.
    integer, parameter :: n=2, max_iter=1000
    real(dp), parameter :: step_initial=1.0_dp, tol_grad=1.0e-6_dp
    real(dp) :: p(n)
    type(descent_result) :: answer
    integer :: j

    p=(/1.8_dp,-1.6_dp/)
    call gradient_descent(objective,gradient,p,step_initial,tol_grad,max_iter,answer)
    write(*,'(/,a)') 'Gradient descent with backtracking'
    do j=1,n
        write(*,'(a,i0,a,es23.15)') 'p(',j,') = ',p(j)
    end do
    write(*,'(a,es23.15)') 'Objective value: ',answer%value
    write(*,'(a,es14.6)') 'Gradient norm:   ',answer%gradient_norm
    write(*,'(a,i0)') 'Accepted steps:  ',answer%iterations
    write(*,'(a,i0)') 'Function calls:  ',answer%function_calls
    write(*,'(a,i0)') 'Gradient calls:  ',answer%gradient_calls
    write(*,'(a,a)') 'Status: ',trim(answer%status)
end program ch13_gradient_descent