program main
!=======================================================================
! Chapter 10, code 3: a nonlinear boundary value problem.
! Solve y'' = f(x,y,y'), with y(a) = alpha and y(b) = beta.
! Method: centered finite differences and Newton iteration.
! Each Newton correction is found with the Thomas algorithm.
!
! Example on [0,1]:
!   y'' = y**2 - sin(pi*x)**2 - pi**2*sin(pi*x)
!   y(0) = y(1) = 0; exact solution y = sin(pi*x).
! The forcing was chosen so that the exact solution is known.
!
! To solve another problem:
!   1. Edit the interval, endpoint values, grid, and initial guess.
!   2. Edit equation below, including its two partial derivatives.
!   3. Disable check_exact or replace exact_solution.
!=======================================================================
implicit none
integer, parameter :: n = 20, data_unit = 20
integer i, iterations, max_iterations
double precision a, b, alpha, beta, h, tolerance, residual_norm
double precision x(0:n), y(0:n), exact, error, max_error, exact_solution
logical check_exact
external equation, exact_solution

!-------------------- Problem and controls: edit here --------------------
a = 0.0d0
b = 1.0d0
alpha = 0.0d0                   ! y(a)
beta = 0.0d0                    ! y(b)
tolerance = 1.0d-10             ! Maximum absolute discrete-equation residual
max_iterations = 30
check_exact = .true.
! Change n in the declaration above to refine the grid.
!------------------------------------------------------------------------
if (n < 2 .or. b <= a) stop 'Use n >= 2 and b > a.'
if (tolerance <= 0.0d0 .or. max_iterations < 1) &
    stop 'Use positive tolerance and iteration limit.'
h = (b-a)/dble(n)
do i = 0,n
    x(i) = a+dble(i)*h
    ! Initial guess: a straight line between the two endpoint values.
    ! For another initial guess, change this expression.
    y(i) = alpha+(beta-alpha)*(x(i)-a)/(b-a)
end do
x(n) = b
y(0) = alpha
y(n) = beta

write(*,'(a)') 'Nonlinear boundary value problem: finite differences and Newton iteration'
write(*,'(a,i8)') 'Grid intervals:  ', n
write(*,'(a,f12.6)') 'Grid spacing:    ', h
write(*,'(a,es12.4)') 'Newton tolerance:', tolerance
call newton_bvp(equation,n,x,y,tolerance,max_iterations,iterations,residual_norm)

open(unit=data_unit,file='ch10_nonlinear_bvp.dat', &
     status='replace',action='write')
write(*,'(a)') ''
if (check_exact) then
    write(*,'(a10,2(1x,a16),1x,a12)') 'x','numerical y','exact y','abs error'
    write(data_unit,'(a)') '# x numerical_y exact_y abs_error'
else
    write(*,'(a10,1x,a16)') 'x','numerical y'
    write(data_unit,'(a)') '# x numerical_y'
end if
max_error = 0.0d0
do i = 0,n
    if (check_exact) then
        exact = exact_solution(x(i))
        error = abs(y(i)-exact)
        max_error = max(max_error,error)
        write(*,'(f10.6,2(1x,es16.8),1x,es12.4)') x(i),y(i),exact,error
        write(data_unit,'(4(1x,es24.16))') x(i),y(i),exact,error
    else
        write(*,'(f10.6,1x,es16.8)') x(i),y(i)
        write(data_unit,'(2(1x,es24.16))') x(i),y(i)
    end if
end do
close(data_unit)
write(*,'(a)') ''
write(*,'(a,i8)') 'Newton iterations: ', iterations
write(*,'(a,es12.4)') 'Final residual:    ', residual_norm
if (check_exact) write(*,'(a,es12.4)') 'Largest grid error:', max_error
write(*,'(a)') 'Data file: ch10_nonlinear_bvp.dat'
end program main


subroutine equation(x,y,v,f,dfdy,dfdv)
!=======================================================================
! Supply the original equation y'' = f(x,y,v), where v = y'.
! Also supply the partial derivatives df/dy and df/dv for Newton's method.
! These derivatives treat x, y, and v as independent variables.
! For an equation that does not depend on v, set dfdv = 0.
!=======================================================================
implicit none
double precision x, y, v, f, dfdy, dfdv, pi, s
pi = acos(-1.0d0)
s = sin(pi*x)
f = y*y-s*s-pi*pi*s
dfdy = 2.0d0*y
dfdv = 0.0d0
end subroutine equation


double precision function exact_solution(x)
! Exact solution for the supplied equation and default boundary values.
implicit none
double precision x, pi
pi = acos(-1.0d0)
exact_solution = sin(pi*x)
end function exact_solution


subroutine newton_bvp(fcn,n,x,y,tolerance,max_iterations,iterations,residual_norm)
!=======================================================================
! At each interior point define the discrete-equation residual
!   R_i = (y(i-1)-2*y(i)+y(i+1))/h**2 - f(x(i),y(i),v_i),
!   v_i = (y(i+1)-y(i-1))/(2*h).
!
! Newton's method solves J*correction = -R and updates y.
! J is tridiagonal because R_i depends only on three neighboring values.
! Multiply the Newton equations by h**2 to keep matrix entries near unity:
!   lower    =  1 + h*(df/dv)/2
!   diagonal = -2 - h**2*(df/dy)
!   upper    =  1 - h*(df/dv)/2
!   rhs      = -h**2*R_i.
!
! Endpoint corrections are zero: prescribed boundary values never change.
! Stop when max|R_i| <= tolerance. This measures convergence of the
! DISCRETE equations; grid refinement checks discretization error.
! Full Newton steps are used. An unsuitable initial guess may lead to
! failure or a different solution when multiple solutions exist.
!=======================================================================
implicit none
integer n, max_iterations, iterations, i
double precision x(0:n), y(0:n), tolerance, residual_norm
double precision lower(n-1), diagonal(n-1), upper(n-1), rhs(n-1)
double precision correction(n-1), h, v, f, dfdy, dfdv, residual, change
external fcn
h = (x(n)-x(0))/dble(n)
change = 0.0d0
write(*,'(a10,2(1x,a18))') 'iteration','max residual','max correction'

do iterations = 0,max_iterations
    residual_norm = 0.0d0
    do i = 1,n-1
        v = (y(i+1)-y(i-1))/(2.0d0*h)
        call fcn(x(i),y(i),v,f,dfdy,dfdv)
        if (.not. all(abs((/ f,dfdy,dfdv /)) <= huge(1.0d0))) &
            stop 'Nonfinite equation or partial derivative; check the initial guess.'
        residual = (y(i-1)-2.0d0*y(i)+y(i+1))/(h*h)-f
        residual_norm = max(residual_norm,abs(residual))
        lower(i) = 1.0d0+h*dfdv/2.0d0
        diagonal(i) = -2.0d0-h*h*dfdy
        upper(i) = 1.0d0-h*dfdv/2.0d0
        rhs(i) = -h*h*residual
    end do
    write(*,'(i10,2(1x,es18.10))') iterations,residual_norm,change
    if (residual_norm <= tolerance) return
    if (iterations == max_iterations) exit

    ! Boundary corrections are zero; only interior values are updated.
    lower(1) = 0.0d0
    upper(n-1) = 0.0d0
    call thomas(n-1,lower,diagonal,upper,rhs,correction)
    change = maxval(abs(correction))
    if (.not. all(abs(correction) <= huge(1.0d0))) &
        stop 'Nonfinite Newton correction; check the initial guess.'
    if (change <= epsilon(1.0d0)*max(1.0d0,maxval(abs(y)))) &
        stop 'Newton iteration cannot advance; check the tolerance and initial guess.'
    y(1:n-1) = y(1:n-1)+correction
end do
stop 'Newton iteration did not converge within the iteration limit.'
end subroutine newton_bvp


subroutine thomas(n,lower,diagonal,upper,rhs,y)
!=======================================================================
! Solve a tridiagonal linear system by forward elimination and
! backward substitution. diagonal and rhs are overwritten.
!
! No pivoting is used. A near-zero pivot stops the calculation.
! Such a failure can require a solver with pivoting or can indicate
! a singular boundary value problem.
!=======================================================================
implicit none
integer n, i
double precision lower(n), diagonal(n), upper(n), rhs(n), y(n)
double precision multiplier, scale

! Check each pivot before division.
do i = 1,n
    scale = max(abs(lower(i)),abs(diagonal(i)),abs(upper(i)))
    if (abs(diagonal(i)) <= epsilon(1.0d0)*scale) &
        stop 'Thomas algorithm: zero or near-zero pivot.'
    if (i == n) exit
    multiplier = lower(i+1)/diagonal(i)
    diagonal(i+1) = diagonal(i+1)-multiplier*upper(i)
    rhs(i+1) = rhs(i+1)-multiplier*rhs(i)
end do

y(n) = rhs(n)/diagonal(n)
do i = n-1,1,-1
    y(i) = (rhs(i)-upper(i)*y(i+1))/diagonal(i)
end do
end subroutine thomas