Line data    Source code

       1             : module trsolv_mod
       2             :   use shr_kind_mod, only: r8 => shr_kind_r8
       3             :   
       4             : contains
       5             : !-----------------------------------------------------------------------
       6      109440 :   subroutine trsolv(a,b,c,f,x,lev0,lev1,k1,k2,lon0,lon1)
       7             : !
       8             : ! Tri-diagonal solver.
       9             : !   a(k,i)*x(k-1,i) + b(k,i)*x(k,i) + c(k,i)*x(k+1,i) = f(k,i)
      10             : !
      11             :     implicit none
      12             : !
      13             : ! Args:
      14             :     integer,intent(in) :: lev0,lev1,k1,k2,lon0,lon1
      15             :     real(r8),dimension(lev0:lev1,lon0:lon1),intent(in) :: &
      16             :       a, & ! input coefficients
      17             :       b, & ! input coefficients
      18             :       c, & ! input coefficients
      19             :       f    ! input RHS
      20             :     real(r8),dimension(lev0:lev1,lon0:lon1),intent(out) :: &
      21             :       x  ! output
      22             : !
      23             : ! Local:
      24             :     integer :: k,kk,i
      25      218880 :     real(r8),dimension(lev0:lev1,lon0:lon1) :: w1,w2,w3  ! work arrays
      26             : 
      27             : !
      28             : ! Lower boundary (W(K1)=B(K1):
      29     1422720 :     do i=lon0,lon1
      30     1422720 :       w1(lev0,i) = b(lev0,i) 
      31             :     enddo
      32             : !
      33             : ! Set up work arrays:
      34     1422720 :     do i=lon0,lon1
      35    60520320 :       do k=k1+1,k2
      36             : !
      37             : ! W(KF+K-1)=C(K-1)/W(K-1):
      38    59097600 :         w2(k-1,i) = c(k-1,i) / w1(k-1,i)
      39             : !
      40             : ! W(K)=A(K)*W(KF+K-1)
      41    59097600 :         w1(k,i) = a(k,i) * w2(k-1,i)
      42             : !
      43             : ! W(K)=B(K)-W(K)
      44    60410880 :         w1(k,i) = b(k,i) - w1(k,i)
      45             :       enddo ! k=k1+1,k2
      46             :     enddo ! i=lon0,lon1
      47             : !
      48             : ! Lower boundary (W(2*KF+K1)=F(K1)/W(K1)):
      49     1422720 :     do i=lon0,lon1
      50     1422720 :       w3(k1,i) = f(k1,i) / w1(k1,i)
      51             :     enddo
      52             : !
      53     1422720 :     do i=lon0,lon1
      54    60520320 :       do k=k1+1,k2
      55             : !
      56             : ! W(2*KF+K)=A(K)*W(2*KF+K-1)
      57    59097600 :         w3(k,i) = a(k,i) * w3(k-1,i)
      58             : !
      59             : ! W(2*KF+K)=F(K)-W(2*KF+K)
      60    59097600 :         w3(k,i) = f(k,i) - w3(k,i)         
      61             : !
      62             : ! W(2*KF+K)=W(2*KF+K)/W(K)
      63    60410880 :         w3(k,i) = w3(k,i) / w1(k,i)
      64             :       enddo ! k=k1+1,k2
      65             :     enddo ! i=lon0,lon1
      66             : !
      67             : ! Upper boundary (X(K2)=W(2*KF+K2)):
      68     1422720 :     do i=lon0,lon1
      69     1422720 :       x(k2,i) = w3(k2,i)       
      70             :     enddo
      71             : !
      72             : 
      73             : ! Back substitution:
      74     1422720 :     do i=lon0,lon1
      75    60520320 :       do kk=k1+1,k2          
      76    59097600 :         k = k1+k2-kk ! k2-1,k1,-1
      77             : !
      78             : ! X(K)=W(KF+K)*X(K+1)
      79    59097600 :         x(k,i) = w2(k,i) * x(k+1,i)
      80             : !
      81             : ! X(K)=W(2*KF+K)-X(K)
      82    60410880 :         x(k,i) = w3(k,i) - x(k,i)
      83             :       enddo ! k=k1+1,k2
      84             :     enddo
      85      109440 :   end subroutine trsolv
      86             : !-----------------------------------------------------------------------
      87             : end module trsolv_mod