Line data Source code
1 : module DGBSVX_wrapper
2 :
3 : use const_def, only: dp
4 : use pre_conditioners
5 :
6 : implicit none
7 :
8 : private
9 : public :: DGBSVX_tridiagonal_wrapper, DGBSVX_banded_wrapper
10 :
11 : contains
12 :
13 : !> Wraps the lapack DGBSVX routine for banded matrices with a preconditioner
14 : !! and a simpler input format.
15 : !!
16 : !! @param matrix_size The number of rows or columns in the square matrix.
17 : !! @params n_upper_bands The number of bands above the diagonal.
18 : !! @params n_lower_bands The number of bands below the diagonal.
19 : !! @params bands Array of shape (n_upper_bands + n_lower_bands + 1, matrix_size). Stores the bands going from upper-most to lower-most, aligned so the first entry in each band has index 1.
20 : !! @params x Array of shape (matrix_size). Stores the result of solving Matrix.x = b.
21 : !! @params b Array of shape (matrix_size). The right-hand side from the definition of x.
22 : !! @params pre_conditioner Pre-conditioning vector. Computed in this routine.
23 : !! @params ierr Integer-valued error code.
24 0 : subroutine DGBSVX_banded_wrapper(matrix_size, n_upper_bands, n_lower_bands, bands, x, b, pre_conditioner, ierr)
25 : ! Inputs
26 : real(dp), dimension(:,:), intent(in) :: bands
27 : real(dp), dimension(:), intent(in) :: b
28 : integer, intent(in) :: matrix_size, n_upper_bands, n_lower_bands
29 :
30 : ! Intermediates
31 : integer :: i, j
32 : character(len=1) :: FACT, TRANS, EQUED
33 : integer :: NRHS, LDAB, LDAFB, LDB, LDX
34 0 : integer :: IWORK(matrix_size), IPIV(matrix_size)
35 0 : real(dp) :: WORK(3 * matrix_size), BERR(1), FERR(1)
36 0 : real(dp) :: C(matrix_size), R(matrix_size)
37 0 : real(dp) :: b_conditioned(matrix_size, 1), x_tmp(matrix_size, 1)
38 0 : real(dp) :: AB(n_lower_bands+n_upper_bands+1, matrix_size)
39 0 : real(dp) :: AFB(2*n_lower_bands+n_upper_bands+1, matrix_size)
40 0 : real(dp) :: RCOND
41 :
42 : ! Outputs
43 : real(dp), dimension(:), intent(out) :: x
44 : real(dp), dimension(:), intent(out) :: pre_conditioner
45 : integer, intent(out) :: ierr
46 :
47 : ! Fixed arguments
48 :
49 0 : FACT = 'N' ! We've already equilibrated the matrix
50 0 : TRANS = 'N' ! This argument is irrelevant if FACT == 'N'
51 0 : EQUED = 'N' ! No equilibration
52 0 : NRHS = 1 ! Number of columns in equ for A.x=b.
53 0 : LDX = matrix_size ! Number of rows in flattened x
54 0 : LDB = matrix_size ! Number of rows in flattened b
55 0 : LDAB = n_lower_bands+n_upper_bands+1 ! Number of rows in AB format of banded matrix
56 0 : LDAFB = 2*n_lower_bands+n_upper_bands+1 ! Number of rows in AB format of LU-factored matrix
57 :
58 : ! Compute preconditioner
59 0 : call compute_band_preconditioner(matrix_size, n_upper_bands, n_lower_bands, bands, pre_conditioner)
60 :
61 : ! Pre-condition b
62 0 : do i=1,matrix_size
63 0 : b_conditioned(i,1) = b(i) / pre_conditioner(i)
64 : end do
65 :
66 : ! Put bands into AB form for LAPACK
67 : ! We store column j in the original matrix in column j of AB.
68 : ! In that column, row k in the original matrix is stored in row
69 : ! n_upper_bands+1+k-j of AB. So the mapping from the matrix to AB is
70 : ! (k,j) -> (n_upper_bands+1+k-j,j)
71 0 : AB = 0d0
72 :
73 : ! In the upper bands, upper band i at index j (bands(i,j)) corresponds to
74 : ! position (j, n_upper_bands - i + j + 1) in the matrix. This is then
75 : ! position (i, n_upper_bands - i + j + 1) in AB.
76 0 : do i=1,n_upper_bands
77 0 : do j=1,matrix_size + i - n_upper_bands - 1
78 0 : AB(i, n_upper_bands - i + j + 1) = bands(i,j) / pre_conditioner(j)
79 : end do
80 : end do
81 :
82 : ! In the lower bands, lower band i+1+n_upper_bands at index j corresponds to
83 : ! position (i+j, j) in the matrix. This is then
84 : ! position (n_upper_bands+1+i, j) in AB.
85 0 : do i=1,n_lower_bands
86 0 : do j=1,matrix_size + n_lower_bands - i - 1
87 0 : AB(n_upper_bands+1+i,j) = bands(i+n_upper_bands+1,j) / pre_conditioner(i+j)
88 : end do
89 : end do
90 :
91 : ! In the diagonal band, position k corresponds to index (k,k), which is position
92 : ! (n_upper_bands+1,k) in AB.
93 0 : do i=1,matrix_size
94 0 : AB(n_upper_bands + 1, i) = bands(n_upper_bands + 1, i) / pre_conditioner(i)
95 : end do
96 :
97 : ! Solve
98 : call DGBSVX(FACT, TRANS, matrix_size, n_lower_bands, n_upper_bands, NRHS, AB, LDAB, AFB, LDAFB, IPIV,&
99 : EQUED, R, C, b_conditioned, LDB, x_tmp, LDX, RCOND, FERR, BERR, WORK, IWORK,&
100 0 : ierr)
101 :
102 : ! Check for errors
103 0 : if (RCOND < 1d-12 .or. ierr == matrix_size+1) then
104 0 : write(*,*) 'Matrix is singular to machine precision.'
105 0 : write(*,*) 'RCOND = ', RCOND
106 0 : write(*,*) 'ierr = ', ierr
107 0 : write(*,*) 'N = ', matrix_size
108 :
109 0 : open(unit=10, file="bands.data")
110 0 : do j=1,LDAB
111 0 : do i=1,matrix_size
112 0 : write(10,*) bands(j,i)
113 : end do
114 : end do
115 0 : else if (ierr /= 0) then
116 0 : write(*,*) 'Matrix is exactly singular.'
117 0 : write(*,*) 'ierr = ', ierr
118 0 : write(*,*) 'N = ', matrix_size
119 : end if
120 :
121 : ! Store output
122 0 : do i=1,matrix_size
123 0 : x(i) = x_tmp(i,1)
124 : end do
125 :
126 0 : end subroutine DGBSVX_banded_wrapper
127 :
128 : !> Wraps the lapack DGBSVX routine for banded matrices to be used with block tridiagonal matrices.
129 : !! Also includes a pre-conditioning step and a simpler input format.
130 : !!
131 : !! @params ublk The upper blocks. Shape is (nvars, nvars, nblocks-1).
132 : !! @params lblk The lower blocks. Shape is (nvars, nvars, nblocks-1).
133 : !! @params dblk The diagonal blocks. Shape is (nvars, nvars, nblocks).
134 : !! @params nblocks The number of blocks on the diagonal.
135 : !! @params nvars The width and height of each block (e.g. each block is nvars by nvars).
136 : !! @params x Array of shape (nvar, nblocks). Stores the result of solving Matrix.x = b.
137 : !! @params b Array of shape (nvar, nblocks). The right-hand side from the definition of x.
138 : !! @params pre_conditioner Pre-conditioning vector. Computed in this routine.
139 : !! @params ierr Integer-valued error code.
140 0 : subroutine DGBSVX_tridiagonal_wrapper(ublk, lblk, dblk, pre_conditioner, x, b, nblocks, nvar, rcond, ierr)
141 : ! Inputs
142 : real(dp), dimension(:,:,:), intent(in) :: ublk, lblk, dblk
143 : real(dp), dimension(:,:), intent(in) :: b
144 : integer, intent(in) :: nblocks, nvar
145 :
146 : ! Intermediates
147 : integer :: counter, i, j, k, k1, k2
148 : character(len=1) :: FACT, TRANS, EQUED
149 : integer :: N, KL, KU, NRHS, LDAB, LDAFB, LDB, LDX
150 0 : integer :: IWORK(nvar * nblocks), IPIV(nvar * nblocks)
151 0 : real(dp) :: WORK(3 * nvar * nblocks), BERR(1), FERR(1)
152 0 : real(dp) :: C(nvar * nblocks), R(nvar * nblocks)
153 0 : real(dp) :: b_flat(nvar * nblocks, 1), x_flat(nvar * nblocks, 1)
154 0 : real(dp) :: AB(4 * nvar + 1, nvar * nblocks)
155 0 : real(dp) :: AFB(6 * nvar + 1, nblocks*nvar)
156 0 : real(dp) :: left_pre_conditioner(nvar, nblocks)
157 :
158 : ! Outputs
159 : real(dp), dimension(:,:), intent(out) :: pre_conditioner
160 : real(dp), dimension(:,:), intent(out) :: x
161 : real(dp), intent(out) :: rcond
162 : integer, intent(out) :: ierr
163 :
164 : ! Fixed arguments
165 :
166 0 : FACT = 'N' ! We've already equilibrated the matrix
167 0 : TRANS = 'N' ! This argument is irrelevant if FACT == 'N'
168 0 : EQUED = 'N' ! No equilibration
169 0 : N = nblocks * nvar ! Number of equations
170 0 : KL = 2 * nvar ! Number of lower bands. We round up to the max number
171 : ! which could be held by the block triadiagonal matrix.
172 : ! If performance is ever an issue this solve can be sped
173 : ! up by tightening KL.
174 0 : KU = 2 * nvar ! Number of upper bands. We round up to the max number
175 : ! which could be held by the block triadiagonal matrix.
176 : ! If performance is ever an issue this solve can be sped
177 : ! up by tightening KU.
178 0 : NRHS = 1 ! Number of columns in equ for A.x=b.
179 0 : LDX = N ! Number of rows in flattened x
180 0 : LDB = N ! Number of rows in flattened b
181 0 : LDAB = KL+KU+1 ! Number of rows in AB format of banded matrix
182 0 : LDAFB = 2*KL+KU+1 ! Number of rows in AB format of LU-factored matrix
183 :
184 : ! Compute preconditioners
185 0 : call compute_block_preconditioner(ublk, lblk, dblk, nblocks, nvar, pre_conditioner)
186 : call compute_block_left_preconditioner(ublk, lblk, dblk, nblocks, nvar, &
187 0 : pre_conditioner, left_pre_conditioner)
188 :
189 : ! We solve M_{ij} x_j = b_i
190 : ! We pre-condition by dividing b_i by p_i, so
191 : ! M_{ij} x_j / p_i = b_i / p_i
192 : ! We also pre-condition by dividing x_j by L_j, so
193 : ! (M_{ij} / (p_i L_j) ) (x_j L_j) = b_i / p_i
194 : ! We solve for the composite x_j' = x_j L_j, so in the end x_j = x_j' / L_j.
195 :
196 : ! Flatten b and precondition
197 0 : counter = 1
198 0 : do j=1,nblocks
199 0 : do k=1,nvar
200 0 : b_flat(counter,1) = b(k,j) / pre_conditioner(k,j)
201 0 : counter = counter + 1
202 : end do
203 : end do
204 :
205 : ! Set up banded matrix
206 0 : AB = 0d0
207 0 : do i=1,nblocks
208 0 : do j=1,nvar
209 0 : do k=1,nvar
210 : ! (j,k,i) in dblk corresponds to coordinate ((i-1)*nvar+j,(i-1)*nvar+k)
211 : ! in the original matrix.
212 : ! We store column m in the original matrix in column m of AB.
213 : ! In that column, row l in the original matrix is stored in row
214 : ! KU+1+l-m of AB. Hence
215 : ! (j,k,i) goes in spot (KU+1+(i-1)*nvar+j-(i-1)*nvar-k,(i-1)*nvar+k)
216 : ! = (KU+1+j-k,(i-1)*nvar+k) of AB.
217 0 : AB(KU+1+j-k, (i-1)*nvar+k) = dblk(j,k,i) / pre_conditioner(j,i) / left_pre_conditioner(k,i)
218 : end do
219 : end do
220 : end do
221 :
222 0 : do i=1,nblocks-1
223 0 : do j=1,nvar
224 0 : do k=1,nvar
225 : ! (j,k,i) in lblk corresponds to coordinate (i*nvar+j,(i-1)*nvar+k)
226 : ! in the original matrix.
227 : ! We store column m in the original matrix in column m of AB.
228 : ! In that column, row l in the original matrix is stored in row
229 : ! KU+1+l-m of AB. Hence
230 : ! (j,k,i) goes in spot (KU+1+i*nvar+j-(i-1)*nvar-k,(i-1)*nvar+k)
231 : ! = (KU+1+nvar+j-k,(i-1)*nvar+k) of AB.
232 0 : AB(KU+1+nvar+j-k, (i-1)*nvar+k) = lblk(j,k,i) / pre_conditioner(j,i+1) / left_pre_conditioner(k,i)
233 :
234 : ! (j,k,i) in ublk corresponds to coordinate ((i-1)*nvar+j,i*nvar+k)
235 : ! in the original matrix.
236 : ! We store column m in the original matrix in column m of AB.
237 : ! In that column, row l in the original matrix is stored in row
238 : ! KU+1+l-m of AB. Hence
239 : ! (j,k,i) goes in spot (KU+1+(i-1)*nvar+j-i*nvar-k,i*nvar+k)
240 : ! = (KU+1-nvar+j-k,(i-1)*nvar+k) of AB.
241 0 : AB(KU+1-nvar+j-k, i*nvar+k) = ublk(j,k,i) / pre_conditioner(j,i) / left_pre_conditioner(k,i+1)
242 : end do
243 : end do
244 : end do
245 :
246 : ! Solve
247 : call DGBSVX(FACT, TRANS, N, KL, KU, NRHS, AB, LDAB, AFB, LDAFB, IPIV,&
248 : EQUED, R, C, b_flat, LDB, x_flat, LDX, RCOND, FERR, BERR, WORK, IWORK,&
249 0 : ierr)
250 :
251 : ! Check for errors
252 0 : if (ierr /= 0 .or. rcond < 1d-10) then
253 0 : write(*,*) 'WARNING: Matrix is singular to machine precision.'
254 0 : write(*,*) 'RCOND = ', RCOND
255 0 : write(*,*) 'ierr = ', ierr
256 0 : write(*,*) 'N = ', N
257 0 : write(*,*) 'nvars, nblocks', nvar, nblocks
258 :
259 0 : open(unit=10, file="lower.data")
260 0 : open(unit=11, file="upper.data")
261 0 : open(unit=12, file="diagonal.data")
262 0 : do k=1,nblocks
263 0 : do k1=1,nvar
264 0 : do k2=1,nvar
265 0 : if (k < nblocks) then
266 0 : write(10,*) lblk(k1,k2,k)
267 0 : write(11,*) ublk(k1,k2,k)
268 : end if
269 0 : write(12,*) dblk(k1,k2,k)
270 : end do
271 : end do
272 : end do
273 0 : write(10,*) '---'
274 0 : write(11,*) '---'
275 0 : write(12,*) '---'
276 0 : ierr = N+1
277 : else if (ierr /= 0) then
278 : write(*,*) 'WARNING: Matrix is exactly singular.'
279 : write(*,*) 'ierr = ', ierr
280 : write(*,*) 'N = ', N
281 : end if
282 :
283 : ! Unflatten output
284 : counter = 1
285 0 : do j=1,nblocks
286 0 : do k=1,nvar
287 0 : x(k,j) = x_flat(counter,1) / left_pre_conditioner(k,j)
288 0 : counter = counter + 1
289 : end do
290 : end do
291 :
292 0 : end subroutine DGBSVX_tridiagonal_wrapper
293 :
294 : end module DGBSVX_wrapper
|
