regrid_solvers.F90

!> Solvers of linear systems.
module regrid_solvers

! This file is part of MOM6. See LICENSE.md for the license.

use mom_error_handler, only : mom_error, fatal

implicit none ; private

public :: solve_linear_system, linear_solver, solve_tridiagonal_system, solve_diag_dominant_tridiag

contains

!> Solve the linear system AX = R by Gaussian elimination
!!
!! This routine uses Gauss's algorithm to transform the system's original
!! matrix into an upper triangular matrix. Back substitution yields the answer.
!! The matrix A must be square, with the first index varing down the column.
subroutine solve_linear_system( A, R, X, N, answers_2018 )
  integer,              intent(in)    :: n  !< The size of the system
  real, dimension(N,N), intent(inout) :: a  !< The matrix being inverted [nondim]
  real, dimension(N),   intent(inout) :: r  !< system right-hand side [A]
  real, dimension(N),   intent(inout) :: x  !< solution vector [A]
  logical,    optional, intent(in)    :: answers_2018 !< If true or absent use older, less efficient expressions.
  ! Local variables
  real, parameter       :: eps = 0.0        ! Minimum pivot magnitude allowed
  real    :: factor       ! The factor that eliminates the leading nonzero element in a row.
  real    :: pivot, i_pivot ! The pivot value and its reciprocal [nondim]
  real    :: swap_a, swap_b
  logical :: found_pivot  ! If true, a pivot has been found
  logical :: old_answers  ! If true, use expressions that give the original (2008 through 2018) MOM6 answers
  integer :: i, j, k

  old_answers = .true. ; if (present(answers_2018)) old_answers = answers_2018

  ! Loop on rows to transform the problem into multiplication by an upper-right matrix.
  do i = 1,n-1


    ! Start to look for a pivot in the current row, i.  If the pivot in row i is not valid,
    ! keep looking for a valid pivot by searching the entries of column i in rows below row i.
    ! Once a valid pivot is found (say in row k), rows i and k are swaped.
    found_pivot = .false.
    k = i
    do while ( ( .NOT. found_pivot ) .AND. ( k <= n ) )
      if ( abs( a(k,i) ) > eps ) then  ! A valid pivot has been found
        found_pivot = .true.
      else                             ! Seek a valid pivot in the next row
        k = k + 1
      endif
    enddo ! end loop to find pivot

    ! If no pivot could be found, the system is singular.
    if ( .NOT. found_pivot ) then
      write(0,*) ' A=',a
      call mom_error( fatal, 'The linear system is singular !' )
    endif

    ! If the pivot is in a row that is different than row i, that is if
    ! k is different than i, we need to swap those two rows
    if ( k /= i ) then
      do j = 1,n
        swap_a = a(i,j) ; a(i,j) = a(k,j) ; a(k,j) = swap_a
      enddo
      swap_b = r(i) ; r(i) = r(k) ; r(k) = swap_b
    endif

    ! Transform pivot to 1 by dividing the entire row (right-hand side included) by the pivot
    if (old_answers) then
      pivot = a(i,i)
      do j = i,n ; a(i,j) = a(i,j) / pivot ; enddo
      r(i) = r(i) / pivot
    else
      i_pivot = 1.0 / a(i,i)
      a(i,i) = 1.0
      do j = i+1,n ; a(i,j) = a(i,j) * i_pivot ; enddo
      r(i) = r(i) * i_pivot
    endif

    ! #INV: At this point, A(i,i) is a suitable pivot and it is equal to 1

    ! Put zeros in column for all rows below that contain the pivot (which is row i)
    do k = i+1,n    ! k is the row index
      factor = a(k,i)
      ! A(k,i) = 0.0  ! These elements are not used again, so this line can be skipped for speed.
      do j = i+1,n  ! j is the column index
        a(k,j) = a(k,j) - factor * a(i,j)
      enddo
      r(k) = r(k) - factor * r(i)
    enddo

  enddo ! end loop on i

  ! Solve system by back substituting in what is now an upper-right matrix.
  x(n) = r(n) / a(n,n)  ! The last row is now trivially solved.
  do i = n-1,1,-1 ! loop on rows, starting from second to last row
    x(i) = r(i)
    do j = i+1,n
      x(i) = x(i) - a(i,j) * x(j)
    enddo
    if (old_answers) x(i) = x(i) / a(i,i)
  enddo

end subroutine solve_linear_system

!> Solve the linear system AX = R by Gaussian elimination
!!
!! This routine uses Gauss's algorithm to transform the system's original
!! matrix into an upper triangular matrix. Back substitution then yields the answer.
!! The matrix A must be square, with the first index varing along the row.
subroutine linear_solver( N, A, R, X )
  integer,              intent(in)    :: n  !< The size of the system
  real, dimension(N,N), intent(inout) :: a  !< The matrix being inverted [nondim]
  real, dimension(N),   intent(inout) :: r  !< system right-hand side [A]
  real, dimension(N),   intent(inout) :: x  !< solution vector [A]

  ! Local variables
  real, parameter :: eps = 0.0   ! Minimum pivot magnitude allowed
  real    :: factor       ! The factor that eliminates the leading nonzero element in a row.
  real    :: i_pivot      ! The reciprocal of the pivot value [inverse of the input units of a row of A]
  real    :: swap
  logical :: found_pivot  ! If true, a pivot has been found
  integer :: i, j, k

  ! Loop on rows to transform the problem into multiplication by an upper-right matrix.
  do i=1,n-1
    ! Seek a pivot for column i starting in row i, and continuing into the remaining rows.  If the
    ! pivot is in a row other than i, swap them.  If no valid pivot is found, i = N+1 after this loop.
    do k=i,n ; if ( abs( a(i,k) ) > eps ) exit ; enddo ! end loop to find pivot
    if ( k > n ) then  ! No pivot could be found and the system is singular.
      write(0,*) ' A=',a
      call mom_error( fatal, 'The linear system is singular !' )
    endif

    ! If the pivot is in a row that is different than row i, swap those two rows, noting that both
    ! rows start with i-1 zero values.
    if ( k /= i ) then
      do j=i,n ; swap = a(j,i) ; a(j,i) = a(j,k) ; a(j,k) = swap ; enddo
      swap = r(i) ; r(i) = r(k) ; r(k) = swap
    endif

    ! Transform the pivot to 1 by dividing the entire row (right-hand side included) by the pivot
    i_pivot = 1.0 / a(i,i)
    a(i,i) = 1.0
    do j=i+1,n ; a(j,i) = a(j,i) * i_pivot ; enddo
    r(i) = r(i) * i_pivot

    ! Put zeros in column for all rows below that contain the pivot (which is row i)
    do k=i+1,n    ! k is the row index
      factor = a(i,k)
      ! A(i,k) = 0.0  ! These elements are not used again, so this line can be skipped for speed.
      do j=i+1,n ; a(j,k) = a(j,k) - factor * a(j,i) ; enddo
      r(k) = r(k) - factor * r(i)
    enddo

  enddo ! end loop on i

  ! Solve the system by back substituting into what is now an upper-right matrix.
  x(n) = r(n) / a(n,n)  ! The last row is now trivially solved.
  do i=n-1,1,-1 ! loop on rows, starting from second to last row
    x(i) = r(i)
    do j=i+1,n ; x(i) = x(i) - a(j,i) * x(j) ; enddo
  enddo

end subroutine linear_solver


!> Solve the tridiagonal system AX = R
!!
!! This routine uses Thomas's algorithm to solve the tridiagonal system AX = R.
!! (A is made up of lower, middle and upper diagonals)
subroutine solve_tridiagonal_system( Al, Ad, Au, R, X, N, answers_2018 )
  integer,            intent(in)  :: n   !< The size of the system
  real, dimension(N), intent(in)  :: ad  !< Matrix center diagonal
  real, dimension(N), intent(in)  :: al  !< Matrix lower diagonal
  real, dimension(N), intent(in)  :: au  !< Matrix upper diagonal
  real, dimension(N), intent(in)  :: r   !< system right-hand side
  real, dimension(N), intent(out) :: x   !< solution vector
  logical,  optional, intent(in)  :: answers_2018 !< If true use older, less acccurate expressions.
  ! Local variables
  real, dimension(N) :: pivot, al_piv
  real, dimension(N) :: c1       ! Au / pivot for the backward sweep
  real    :: i_pivot  ! The inverse of the most recent pivot
  integer :: k        ! Loop index
  logical :: old_answers  ! If true, use expressions that give the original (2008 through 2018) MOM6 answers

  old_answers = .true. ; if (present(answers_2018)) old_answers = answers_2018

  if (old_answers) then
    ! This version gives the same answers as the original (2008 through 2018) MOM6 code
    ! Factorization and forward sweep
    pivot(1) = ad(1)
    x(1) = r(1)
    do k = 2,n
      al_piv(k) = al(k) / pivot(k-1)
      pivot(k) = ad(k) - al_piv(k) * au(k-1)
      x(k) = r(k) - al_piv(k) * x(k-1)
    enddo

    ! Backward sweep
    x(n) = r(n) / pivot(n)  ! This should be X(N) / pivot(N), but is OK if Al(N) = 0.
    do k = n-1,1,-1
      x(k) = ( x(k) - au(k)*x(k+1) ) / pivot(k)
    enddo
  else
    ! This is a more typical implementation of a tridiagonal solver than the one above.
    ! It is mathematically equivalent but differs at roundoff, which can cascade up to larger values.

    ! Factorization and forward sweep
    i_pivot = 1.0 / ad(1)
    x(1) = r(1) * i_pivot
    do k = 2,n
      c1(k-1) = au(k-1) * i_pivot
      i_pivot = 1.0 / (ad(k) - al(k) * c1(k-1))
      x(k) = (r(k) - al(k) * x(k-1)) * i_pivot
    enddo
    ! Backward sweep
    do k = n-1,1,-1
      x(k) = x(k) - c1(k) * x(k+1)
    enddo

  endif

end subroutine solve_tridiagonal_system


!> Solve the tridiagonal system AX = R
!!
!! This routine uses a variant of Thomas's algorithm to solve the tridiagonal system AX = R, in
!! a form that is guaranteed to avoid dividing by a zero pivot.  The matrix A is made up of
!! lower (Al) and upper diagonals (Au) and a central diagonal Ad = Ac+Al+Au, where
!! Al, Au, and Ac are all positive (or negative) definite.  However when Ac is smaller than
!! roundoff compared with (Al+Au), the answers are prone to inaccuracy.
subroutine solve_diag_dominant_tridiag( Al, Ac, Au, R, X, N )
  integer,            intent(in)  :: n   !< The size of the system
  real, dimension(N), intent(in)  :: ac  !< Matrix center diagonal offset from Al + Au
  real, dimension(N), intent(in)  :: al  !< Matrix lower diagonal
  real, dimension(N), intent(in)  :: au  !< Matrix upper diagonal
  real, dimension(N), intent(in)  :: r   !< system right-hand side
  real, dimension(N), intent(out) :: x   !< solution vector
  ! Local variables
  real, dimension(N) :: c1       ! Au / pivot for the backward sweep
  real               :: d1       ! The next value of 1.0 - c1
  real               :: i_pivot  ! The inverse of the most recent pivot
  real               :: denom_t1 ! The first term in the denominator of the inverse of the pivot.
  integer            :: k        ! Loop index

  ! Factorization and forward sweep, in a form that will never give a division by a
  ! zero pivot for positive definite Ac, Al, and Au.
  i_pivot = 1.0 / (ac(1) + au(1))
  d1 = ac(1) * i_pivot
  c1(1) = au(1) * i_pivot
  x(1) = r(1) * i_pivot
  do k=2,n-1
    denom_t1 = ac(k) + d1 * al(k)
    i_pivot = 1.0 / (denom_t1 + au(k))
    d1 = denom_t1 * i_pivot
    c1(k) = au(k) * i_pivot
    x(k) = (r(k) - al(k) * x(k-1)) * i_pivot
  enddo
  i_pivot = 1.0 / (ac(n) + d1 * al(n))
  x(n) = (r(n) - al(n) * x(n-1)) * i_pivot
  ! Backward sweep
  do k=n-1,1,-1
    x(k) = x(k) - c1(k) * x(k+1)
  enddo

end subroutine solve_diag_dominant_tridiag


!> \namespace regrid_solvers
!!
!! Date of creation: 2008.06.12
!! L. White
!!
!! This module contains solvers of linear systems.
!! These routines have now been updated for greater efficiency, especially in special cases.

end module regrid_solvers