tridiagonal.inc

!***********************************************************************
!*                   GNU Lesser General Public License
!*
!* This file is part of the GFDL Flexible Modeling System (FMS).
!*
!* FMS is free software: you can redistribute it and/or modify it under
!* the terms of the GNU Lesser General Public License as published by
!* the Free Software Foundation, either version 3 of the License, or (at
!* your option) any later version.
!*
!* FMS is distributed in the hope that it will be useful, but WITHOUT
!* ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
!* FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
!* for more details.
!*
!* You should have received a copy of the GNU Lesser General Public
!* License along with FMS.  If not, see <http://www.gnu.org/licenses/>.
!***********************************************************************
 
!> @addtogroup tridiagonal_mod
!> @{
 
!> @brief Sets up and solves the tridiagonal system of equations
!!
!> For simplicity, A and C are assumed to be dimensioned the same size
!! as B, D, and X, although any input values for A(N) and C(1) are ignored.
!! There are no checks to make sure the sizes agree.
!!
!! The value of A(N) is modified on output, and B and C are unchanged.
!!
!! For mixed precision, this routine uses the kind size macro(FMS_TRID_KIND_) to determine
!! which module variables are used/stored. This means a,b, and c values will only be stored for calls
!! of the same real kind value unless store_both_kinds is present and .true..
subroutine tri_invert_(x,d,a,b,c, store_both_kinds)
 
    real(FMS_TRID_KIND_), intent(out), dimension(:,:,:) :: x !< Solution to the tridiagonal system of equations
    real(FMS_TRID_KIND_), intent(in),  dimension(:,:,:) :: d !< The right-hand side term, see the schematic above.
    real(FMS_TRID_KIND_), optional,    dimension(:,:,:) :: a,b,c !< Left hand side terms(see schematic on module page).
                                                !! If not provided, values from last call are used
    logical, optional                                   :: store_both_kinds !< Will save module state
                                                         !! variables for both kind types in order to be used in
                                                         !! subsequent calls with either kind.
 
    real(FMS_TRID_KIND_), dimension(size(x,1),size(x,2),size(x,3)) :: f
    integer, parameter :: kindl = fms_trid_kind_
 
    integer :: k
 
    if(present(a)) then
        !$OMP SINGLE
        init_var = .true.
        if(allocated(trid_real_type%e))     deallocate(trid_real_type%e)
        if(allocated(trid_real_type%g))     deallocate(trid_real_type%g)
        if(allocated(trid_real_type%bb))    deallocate(trid_real_type%bb)
        if(allocated(trid_real_type%cc))    deallocate(trid_real_type%cc)
        allocate(trid_real_type%e (size(x,1),size(x,2),size(x,3)))
        allocate(trid_real_type%g (size(x,1),size(x,2),size(x,3)))
        allocate(trid_real_type%bb(size(x,1),size(x,2)))
        allocate(trid_real_type%cc(size(x,1),size(x,2),size(x,3)))
        !$OMP END SINGLE
 
        trid_real_type%e(:,:,1) = - a(:,:,1) / b(:,:,1)
        a(:,:,size(x,3)) = 0.0_kindl
 
        do  k= 2,size(x,3)
            trid_real_type%g(:,:,k) = 1.0_kindl/(b(:,:,k)+c(:,:,k)*trid_real_type%e(:,:,k-1))
            trid_real_type%e(:,:,k) = - a(:,:,k)* trid_real_type%g(:,:,k)
        end do
        trid_real_type%cc = c
        trid_real_type%bb = 1.0_kindl/b(:,:,1)
 
    end if
 
    if(.not.init_var) call mpp_error(fatal, 'tri_invert: a,b,and c args not provided or previously calculated.')
 
    f(:,:,1) =  d(:,:,1)*trid_real_type%bb
    do k= 2, size(x,3)
        f(:,:,k) = (d(:,:,k) - trid_real_type%cc(:,:,k)*f(:,:,k-1))*trid_real_type%g(:,:,k)
    end do
 
    x(:,:,size(x,3)) = f(:,:,size(x,3))
    do k = size(x,3)-1,1,-1
        x(:,:,k) = trid_real_type%e(:,:,k)*x(:,:,k+1)+f(:,:,k)
    end do
 
    ! stores both kind values for subsequent calculations if running with option
    if( present(store_both_kinds)) then
      if( store_both_kinds ) then
        if( fms_trid_kind_ .eq. r8_kind) then
          tridiag_r4%e = real(trid_real_type%e, r4_kind)
          tridiag_r4%g = real(trid_real_type%g, r4_kind)
          tridiag_r4%cc = real(trid_real_type%cc, r4_kind)
          tridiag_r4%bb = real(trid_real_type%bb, r4_kind)
          init_tridiagonal_r4 = .true.
        else
          tridiag_r8%e = real(trid_real_type%e, r8_kind)
          tridiag_r8%g = real(trid_real_type%g, r8_kind)
          tridiag_r8%cc = real(trid_real_type%cc, r8_kind)
          tridiag_r8%bb = real(trid_real_type%bb, r8_kind)
          init_tridiagonal_r8 = .true.
        endif
      endif
    endif
 
    return
end subroutine TRI_INVERT_