program main
!=======================================================================
! Chapter 10, code 5: Schrodinger equation as a matrix eigenvalue problem.
! Dimensionless units: hbar = mass = 1.
!   -0.5*psi'' + V(x)*psi = E*psi, with psi(a) = psi(b) = 0.
! Method: centered finite differences and symmetric Jacobi rotations.
! Jacobi follows Chapter 8, using a stable small-angle rotation formula.
!
! Example: harmonic oscillator, V(x) = 0.5*x**2.
! Exact energies on the infinite line: E_j = j+0.5, j = 0,1,2,... .
! A finite interval approximates the infinite line for the lowest states.
! Check both grid spacing and interval size when testing accuracy.
!
! To solve another problem:
!   1. Edit the interval, grid size, and number of states below.
!   2. Edit potential(x), immediately below the driver.
!   3. Disable check_exact or replace the exact-energy comparison.
! The program computes all matrix eigenpairs and outputs the lowest states.
! Jacobi uses a dense matrix: keep grids modest for this educational code.
!=======================================================================
implicit none
integer, parameter :: n = 120, interior = n-1, states = 3, data_unit = 20
integer i, j, k, pivot, sweeps, rotations, info, max_sweeps
double precision a, b, h, abserr, offdiag_norm, potential, exact
double precision x(0:n), v(0:n), matrix(interior,interior)
double precision vectors(interior,interior), energies(interior)
double precision psi(0:n,states), norm, residual, max_residual
double precision overlap, defect, max_orthogonality
logical check_exact
external potential

!-------------------- Problem and grid: edit here -----------------------
a = -6.0d0
b = 6.0d0
abserr = 1.0d-18                ! Bound on SUM of squared off-diagonal entries
max_sweeps = 100
check_exact = .true.            ! Harmonic-oscillator energy comparison
! Change n and states in the declaration above.
! n intervals give n+1 grid points and n-1 interior unknowns.
!------------------------------------------------------------------------
if (n < 2 .or. states < 1 .or. states > interior .or. b <= a) &
    stop 'Use n >= 2, 1 <= states <= n-1, and b > a.'
h = (b-a)/dble(n)
do i = 0,n
    x(i) = a+dble(i)*h
    v(i) = potential(x(i))
end do
x(n) = b
v(n) = potential(b)

! Only interior values are unknown. Zero endpoint values contribute
! nothing to the first and last interior equations.
! H(i,i) = 1/h**2 + V(x_i); H(i,i+1) = H(i+1,i) = -1/(2*h**2).
matrix = 0.0d0
do i = 1,interior
    matrix(i,i) = 1.0d0/(h*h)+v(i)
end do
do i = 1,interior-1
    matrix(i,i+1) = -0.5d0/(h*h)
    matrix(i+1,i) = matrix(i,i+1)
end do

write(*,'(a)') 'Schrodinger equation: finite difference matrix and Jacobi eigenvalues'
write(*,'(a,i8)') 'Grid intervals: ', n
write(*,'(a,i8)') 'Matrix size:    ', interior
write(*,'(a,f12.6)') 'Grid spacing:   ', h
call jacobi(matrix,vectors,interior,abserr,max_sweeps,sweeps,rotations,offdiag_norm,info)
if (info /= 0) stop 'Jacobi failed: check controls or increase the sweep limit.'
do i = 1,interior
    energies(i) = matrix(i,i)
end do
call sort_eigenpairs(interior,energies,vectors)

! Matrix eigenvectors have sum(vectors**2) = 1.
! Wavefunctions need h*sum(psi**2) = 1. With zero endpoints, this is also
! the trapezoidal approximation to integral(psi**2 dx) over [a,b].
psi = 0.0d0
do j = 1,states
    norm = sqrt(h*sum(vectors(:,j)**2))
    psi(1:n-1,j) = vectors(:,j)/norm
    ! Fix an arbitrary overall sign for reproducible presentation.
    ! Make the component with largest magnitude positive.
    pivot = maxloc(abs(psi(:,j)),dim=1)-1
    if (psi(pivot,j) < 0.0d0) psi(:,j) = -psi(:,j)
end do

write(*,'(a)') ''
if (check_exact) then
    write(*,'(a7,2(1x,a16),1x,a12)') 'state j','numerical E','exact E','abs error'
else
    write(*,'(a7,1x,a16)') 'state j','numerical E'
end if
do j = 1,states
    if (check_exact) then
        exact = dble(j-1)+0.5d0
        write(*,'(i7,2(1x,f16.8),1x,es12.4)') j-1,energies(j),exact,abs(energies(j)-exact)
    else
        write(*,'(i7,1x,f16.8)') j-1,energies(j)
    end if
end do

! Check the selected eigenpairs using the original tridiagonal operator,
! rather than the matrix overwritten by Jacobi.
max_residual = 0.0d0
do j = 1,states
    do i = 1,n-1
        residual = (1.0d0/(h*h)+v(i)-energies(j))*psi(i,j) &
                   -0.5d0*(psi(i-1,j)+psi(i+1,j))/(h*h)
        max_residual = max(max_residual,abs(residual))
    end do
end do
max_orthogonality = 0.0d0
do j = 1,states
    do k = 1,states
        overlap = h*sum(psi(:,j)*psi(:,k))
        defect = 0.0d0
        if (j == k) defect = 1.0d0
        max_orthogonality = max(max_orthogonality,abs(overlap-defect))
    end do
end do
write(*,'(a,i8)') 'Jacobi sweeps:            ', sweeps
write(*,'(a,es12.4)') 'Largest matrix residual: ', max_residual
write(*,'(a,es12.4)') 'Largest overlap defect:  ', max_orthogonality

! One data file: x, potential, psi_0, psi_1, ... .
open(unit=data_unit,file='ch10_schrodinger_matrix.dat',status='replace',action='write')
write(data_unit,'(a)') '# x V(x) psi_0 psi_1 ...; integral(psi_j**2 dx) = 1'
do j = 1,states
    write(data_unit,'(a,i6,a,es24.16)') '# state ',j-1,' E = ',energies(j)
end do
write(*,'(a)') ''
write(*,'(a)') 'Wavefunctions: columns are x, V(x), then the normalized states.'
do i = 0,n
    write(*,'(f10.6,1x,es16.8)',advance='no') x(i),v(i)
    do j = 1,states
        write(*,'(1x,es16.8)',advance='no') psi(i,j)
    end do
    write(*,'(a)') ''
    write(data_unit,'(1000(1x,es24.16))') x(i),v(i),psi(i,:)
end do
close(data_unit)
write(*,'(a)') 'Data file: ch10_schrodinger_matrix.dat'
end program main


double precision function potential(x)
!=======================================================================
! Supply the potential in -0.5*psi'' + V(x)*psi = E*psi.
! Default: dimensionless harmonic oscillator.
! For another potential, change this expression and disable check_exact.
! The zero endpoint conditions are imposed by the matrix construction.
!=======================================================================
implicit none
double precision x
potential = 0.5d0*x*x
end function potential


subroutine sort_eigenpairs(n,values,vectors)
! Sort energies in ascending order, moving their eigenvector columns too.
implicit none
integer n, i, j, smallest
double precision values(n), vectors(n,n), work(n), value
do i = 1,n-1
    smallest = i
    do j = i+1,n
        if (values(j) < values(smallest)) smallest = j
    end do
    if (smallest == i) cycle
    value = values(i)
    values(i) = values(smallest)
    values(smallest) = value
    work = vectors(:,i)
    vectors(:,i) = vectors(:,smallest)
    vectors(:,smallest) = work
end do
end subroutine sort_eigenpairs


subroutine jacobi(a,x,n,abserr,max_sweeps,sweeps,rotations,offdiag_norm,info)
!=======================================================================
! Diagonalize a real symmetric matrix by thresholded Jacobi rotations.
!
! Input/output:
!   a(n,n)       - symmetric matrix; overwritten by a nearly diagonal
!                  similar matrix whose diagonal holds the eigenvalues
! Output:
!   x(n,n)       - orthonormal eigenvectors in columns: A_original*x_j
!                  is approximately a(j,j)*x_j
!   sweeps       - completed passes over all off-diagonal pairs
!   rotations    - number of rotations actually performed
!   offdiag_norm - sqrt of the sum of all off-diagonal squares
!   info         - 0 converged; 1 invalid input; 2 sweep limit
! Controls:
!   abserr       - positive bound on SUM of squared off-diagonal entries
!   max_sweeps   - positive maximum number of complete passes
!
! Start X as the identity. For each off-diagonal pair above the
! threshold, rotate rows and columns of A by the same plane angle;
! rotate columns of X as well. Orthogonal similarity preserves the
! eigenvalues. Recompute the off-diagonal sum after each sweep rather
! than subtracting an idealized change that can accumulate roundoff.
!
! The stable tangent formula avoids cancellation when diagonal entries
! differ greatly. This matters more here than in the small Chapter 8
! example. The input matrix must be exactly symmetric.
!=======================================================================
implicit none
integer, intent(in) :: n, max_sweeps
double precision, intent(in) :: abserr
double precision, intent(inout) :: a(n,n)
double precision, intent(out) :: x(n,n), offdiag_norm
integer, intent(out) :: sweeps, rotations, info

double precision :: b2, bar, beta, tangent, c, s, cs, sc
integer :: i, j, k

x = 0.0d0
do i = 1,n
   x(i,i) = 1.0d0
end do
sweeps = 0
rotations = 0
offdiag_norm = huge(1.0d0)
info = 1
if (n < 1 .or. abserr <= 0.0d0 .or. max_sweeps < 1) return
do i = 1,n
   do j = i+1,n
      if (a(i,j) /= a(j,i)) return
   end do
end do

do
   ! Measure the actual off-diagonal entries at the start of each
   ! sweep. In exact arithmetic b2 decreases after every rotation.
   b2 = 0.0d0
   do i = 1,n
      do j = 1,n
         if (i /= j) b2 = b2 + a(i,j)**2
      end do
   end do
   offdiag_norm = sqrt(b2)
   if (b2 <= abserr) then
      info = 0
      return
   end if
   if (sweeps >= max_sweeps) then
      info = 2
      return
   end if

   ! Original code sets bar to half the mean off-diagonal square.
   ! As significant elements disappear, lower the threshold during
   ! this sweep using the expected reduction of b2.
   bar = 0.5d0*b2/dble(n*n)
   do i = 1,n-1
      do j = i+1,n
         if (a(j,i)**2 <= bar) cycle
         b2 = max(0.0d0,b2-2.0d0*a(j,i)**2)
         bar = 0.5d0*b2/dble(n*n)

         ! Choose a small rotation that eliminates a(j,i).
         ! The minus sign follows the row rotation [c s; -s c] below.
         beta = (a(j,j)-a(i,i))/(2.0d0*a(j,i))
         tangent = -sign(1.0d0,beta)/(abs(beta)+sqrt(1.0d0+beta**2))
         c = 1.0d0/sqrt(1.0d0+tangent**2)
         s = tangent*c

         ! First rotate rows i and j of A.
         do k = 1,n
            cs =  c*a(i,k)+s*a(j,k)
            sc = -s*a(i,k)+c*a(j,k)
            a(i,k) = cs
            a(j,k) = sc
         end do

         ! Then rotate columns i and j of A. Apply the same column
         ! rotation to X so its columns become the eigenvectors.
         do k = 1,n
            cs =  c*a(k,i)+s*a(k,j)
            sc = -s*a(k,i)+c*a(k,j)
            a(k,i) = cs
            a(k,j) = sc

            cs =  c*x(k,i)+s*x(k,j)
            sc = -s*x(k,i)+c*x(k,j)
            x(k,i) = cs
            x(k,j) = sc
         end do
         rotations = rotations + 1
      end do
   end do
   sweeps = sweeps + 1
end do
end subroutine jacobi