FMS  2025.04
Flexible Modeling System
tridiagonal.inc
1 !***********************************************************************
2 !* Apache License 2.0
3 !*
4 !* This file is part of the GFDL Flexible Modeling System (FMS).
5 !*
6 !* Licensed under the Apache License, Version 2.0 (the "License");
7 !* you may not use this file except in compliance with the License.
8 !* You may obtain a copy of the License at
9 !*
10 !* http://www.apache.org/licenses/LICENSE-2.0
11 !*
12 !* FMS is distributed in the hope that it will be useful, but WITHOUT
13 !* WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied;
14 !* without even the implied warranty of MERCHANTABILITY or FITNESS FOR A
15 !* PARTICULAR PURPOSE. See the License for the specific language
16 !* governing permissions and limitations under the License.
17 !***********************************************************************
18 
19 !> @addtogroup tridiagonal_mod
20 !> @{
21 
22 !> @brief Sets up and solves the tridiagonal system of equations
23 !!
24 !> For simplicity, A and C are assumed to be dimensioned the same size
25 !! as B, D, and X, although any input values for A(N) and C(1) are ignored.
26 !! There are no checks to make sure the sizes agree.
27 !!
28 !! The value of A(N) is modified on output, and B and C are unchanged.
29 !!
30 !! For mixed precision, this routine uses the kind size macro(FMS_TRID_KIND_) to determine
31 !! which module variables are used/stored. This means a,b, and c values will only be stored for calls
32 !! of the same real kind value unless store_both_kinds is present and .true..
33 subroutine tri_invert_(x,d,a,b,c, store_both_kinds)
34 
35  real(FMS_TRID_KIND_), intent(out), dimension(:,:,:) :: x !< Solution to the tridiagonal system of equations
36  real(FMS_TRID_KIND_), intent(in), dimension(:,:,:) :: d !< The right-hand side term, see the schematic above.
37  real(FMS_TRID_KIND_), optional, dimension(:,:,:) :: a,b,c !< Left hand side terms(see schematic on module page).
38  !! If not provided, values from last call are used
39  logical, optional :: store_both_kinds !< Will save module state
40  !! variables for both kind types in order to be used in
41  !! subsequent calls with either kind.
42 
43  real(FMS_TRID_KIND_), dimension(size(x,1),size(x,2),size(x,3)) :: f
44  integer, parameter :: kindl = fms_trid_kind_
45 
46  integer :: k
47 
48  if(present(a)) then
49  !$OMP SINGLE
50  init_var = .true.
51  if(allocated(trid_real_type%e)) deallocate(trid_real_type%e)
52  if(allocated(trid_real_type%g)) deallocate(trid_real_type%g)
53  if(allocated(trid_real_type%bb)) deallocate(trid_real_type%bb)
54  if(allocated(trid_real_type%cc)) deallocate(trid_real_type%cc)
55  allocate(trid_real_type%e (size(x,1),size(x,2),size(x,3)))
56  allocate(trid_real_type%g (size(x,1),size(x,2),size(x,3)))
57  allocate(trid_real_type%bb(size(x,1),size(x,2)))
58  allocate(trid_real_type%cc(size(x,1),size(x,2),size(x,3)))
59  !$OMP END SINGLE
60 
61  trid_real_type%e(:,:,1) = - a(:,:,1) / b(:,:,1)
62  a(:,:,size(x,3)) = 0.0_kindl
63 
64  do k= 2,size(x,3)
65  trid_real_type%g(:,:,k) = 1.0_kindl/(b(:,:,k)+c(:,:,k)*trid_real_type%e(:,:,k-1))
66  trid_real_type%e(:,:,k) = - a(:,:,k)* trid_real_type%g(:,:,k)
67  end do
68  trid_real_type%cc = c
69  trid_real_type%bb = 1.0_kindl/b(:,:,1)
70 
71  end if
72 
73  if(.not.init_var) call mpp_error(fatal, 'tri_invert: a,b,and c args not provided or previously calculated.')
74 
75  f(:,:,1) = d(:,:,1)*trid_real_type%bb
76  do k= 2, size(x,3)
77  f(:,:,k) = (d(:,:,k) - trid_real_type%cc(:,:,k)*f(:,:,k-1))*trid_real_type%g(:,:,k)
78  end do
79 
80  x(:,:,size(x,3)) = f(:,:,size(x,3))
81  do k = size(x,3)-1,1,-1
82  x(:,:,k) = trid_real_type%e(:,:,k)*x(:,:,k+1)+f(:,:,k)
83  end do
84 
85  ! stores both kind values for subsequent calculations if running with option
86  if( present(store_both_kinds)) then
87  if( store_both_kinds ) then
88  if( fms_trid_kind_ .eq. r8_kind) then
89  tridiag_r4%e = real(trid_real_type%e, r4_kind)
90  tridiag_r4%g = real(trid_real_type%g, r4_kind)
91  tridiag_r4%cc = real(trid_real_type%cc, r4_kind)
92  tridiag_r4%bb = real(trid_real_type%bb, r4_kind)
93  init_tridiagonal_r4 = .true.
94  else
95  tridiag_r8%e = real(trid_real_type%e, r8_kind)
96  tridiag_r8%g = real(trid_real_type%g, r8_kind)
97  tridiag_r8%cc = real(trid_real_type%cc, r8_kind)
98  tridiag_r8%bb = real(trid_real_type%bb, r8_kind)
99  init_tridiagonal_r8 = .true.
100  endif
101  endif
102  endif
103 
104  return
105 end subroutine TRI_INVERT_
subroutine tri_invert_(x, d, a, b, c, store_both_kinds)
Sets up and solves the tridiagonal system of equations.
Definition: tridiagonal.inc:34