| 1 | ! radiation_matrix.F90 - SPARTACUS matrix operations |
|---|
| 2 | ! |
|---|
| 3 | ! (C) Copyright 2014- ECMWF. |
|---|
| 4 | ! |
|---|
| 5 | ! This software is licensed under the terms of the Apache Licence Version 2.0 |
|---|
| 6 | ! which can be obtained at http://www.apache.org/licenses/LICENSE-2.0. |
|---|
| 7 | ! |
|---|
| 8 | ! In applying this licence, ECMWF does not waive the privileges and immunities |
|---|
| 9 | ! granted to it by virtue of its status as an intergovernmental organisation |
|---|
| 10 | ! nor does it submit to any jurisdiction. |
|---|
| 11 | ! |
|---|
| 12 | ! Author: Robin Hogan |
|---|
| 13 | ! Email: r.j.hogan@ecmwf.int |
|---|
| 14 | ! |
|---|
| 15 | ! Modifications |
|---|
| 16 | ! 2018-10-15 R. Hogan Added fast_expm_exchange_[23] |
|---|
| 17 | ! |
|---|
| 18 | ! This module provides the neccessary mathematical functions for the |
|---|
| 19 | ! SPARTACUS radiation scheme: matrix multiplication, matrix solvers |
|---|
| 20 | ! and matrix exponentiation, but (a) multiple matrices are operated on |
|---|
| 21 | ! at once with array access indended to facilitate vectorization, and |
|---|
| 22 | ! (b) optimization for 2x2 and 3x3 matrices. There is probably |
|---|
| 23 | ! considerable scope for further optimization. Note that this module |
|---|
| 24 | ! is not used by the McICA solver. |
|---|
| 25 | |
|---|
| 26 | module radiation_matrix |
|---|
| 27 | |
|---|
| 28 | use parkind1, only : jprb |
|---|
| 29 | |
|---|
| 30 | implicit none |
|---|
| 31 | public |
|---|
| 32 | |
|---|
| 33 | ! Codes to describe sparseness pattern, where the SHORTWAVE |
|---|
| 34 | ! pattern is of the form: |
|---|
| 35 | ! (x x x) |
|---|
| 36 | ! (x x x) |
|---|
| 37 | ! (0 0 x) |
|---|
| 38 | ! where each element may itself be a square matrix. |
|---|
| 39 | integer, parameter :: IMatrixPatternDense = 0 |
|---|
| 40 | integer, parameter :: IMatrixPatternShortwave = 1 |
|---|
| 41 | |
|---|
| 42 | public :: mat_x_vec, singlemat_x_vec, mat_x_mat, & |
|---|
| 43 | & singlemat_x_mat, mat_x_singlemat, & |
|---|
| 44 | & identity_minus_mat_x_mat, solve_vec, solve_mat, expm, & |
|---|
| 45 | & fast_expm_exchange_2, fast_expm_exchange_3, & |
|---|
| 46 | & sparse_x_dense |
|---|
| 47 | |
|---|
| 48 | private :: solve_vec_2, solve_vec_3, solve_mat_2, & |
|---|
| 49 | & solve_mat_3, lu_factorization, lu_substitution, solve_mat_n, & |
|---|
| 50 | & diag_mat_right_divide_3 |
|---|
| 51 | |
|---|
| 52 | interface fast_expm_exchange |
|---|
| 53 | module procedure fast_expm_exchange_2, fast_expm_exchange_3 |
|---|
| 54 | end interface fast_expm_exchange |
|---|
| 55 | |
|---|
| 56 | contains |
|---|
| 57 | |
|---|
| 58 | ! --- MATRIX-VECTOR MULTIPLICATION --- |
|---|
| 59 | |
|---|
| 60 | !--------------------------------------------------------------------- |
|---|
| 61 | ! Treat A as n m-by-m square matrices (with the n dimension varying |
|---|
| 62 | ! fastest) and b as n m-element vectors, and perform matrix-vector |
|---|
| 63 | ! multiplications on first iend pairs |
|---|
| 64 | function mat_x_vec(n,iend,m,A,b,do_top_left_only_in) |
|---|
| 65 | |
|---|
| 66 | use yomhook, only : lhook, dr_hook, jphook |
|---|
| 67 | |
|---|
| 68 | integer, intent(in) :: n, m, iend |
|---|
| 69 | real(jprb), intent(in), dimension(:,:,:) :: A |
|---|
| 70 | real(jprb), intent(in), dimension(:,:) :: b |
|---|
| 71 | logical, intent(in), optional :: do_top_left_only_in |
|---|
| 72 | real(jprb), dimension(iend,m):: mat_x_vec |
|---|
| 73 | |
|---|
| 74 | integer :: j1, j2 |
|---|
| 75 | logical :: do_top_left_only |
|---|
| 76 | |
|---|
| 77 | real(jphook) :: hook_handle |
|---|
| 78 | |
|---|
| 79 | if (lhook) call dr_hook('radiation_matrix:mat_x_vec',0,hook_handle) |
|---|
| 80 | |
|---|
| 81 | if (present(do_top_left_only_in)) then |
|---|
| 82 | do_top_left_only = do_top_left_only_in |
|---|
| 83 | else |
|---|
| 84 | do_top_left_only = .false. |
|---|
| 85 | end if |
|---|
| 86 | |
|---|
| 87 | ! Array-wise assignment |
|---|
| 88 | mat_x_vec = 0.0_jprb |
|---|
| 89 | |
|---|
| 90 | if (do_top_left_only) then |
|---|
| 91 | mat_x_vec(1:iend,1) = A(1:iend,1,1)*b(1:iend,1) |
|---|
| 92 | else |
|---|
| 93 | do j1 = 1,m |
|---|
| 94 | do j2 = 1,m |
|---|
| 95 | mat_x_vec(1:iend,j1) = mat_x_vec(1:iend,j1) & |
|---|
| 96 | & + A(1:iend,j1,j2)*b(1:iend,j2) |
|---|
| 97 | end do |
|---|
| 98 | end do |
|---|
| 99 | end if |
|---|
| 100 | |
|---|
| 101 | if (lhook) call dr_hook('radiation_matrix:mat_x_vec',1,hook_handle) |
|---|
| 102 | |
|---|
| 103 | end function mat_x_vec |
|---|
| 104 | |
|---|
| 105 | |
|---|
| 106 | !--------------------------------------------------------------------- |
|---|
| 107 | ! Treat A as an m-by-m square matrix and b as n m-element vectors |
|---|
| 108 | ! (with the n dimension varying fastest), and perform matrix-vector |
|---|
| 109 | ! multiplications on first iend pairs |
|---|
| 110 | function singlemat_x_vec(n,iend,m,A,b) |
|---|
| 111 | |
|---|
| 112 | ! use yomhook, only : lhook, dr_hook, jphook |
|---|
| 113 | |
|---|
| 114 | integer, intent(in) :: n, m, iend |
|---|
| 115 | real(jprb), intent(in), dimension(m,m) :: A |
|---|
| 116 | real(jprb), intent(in), dimension(:,:) :: b |
|---|
| 117 | real(jprb), dimension(iend,m) :: singlemat_x_vec |
|---|
| 118 | |
|---|
| 119 | integer :: j1, j2 |
|---|
| 120 | ! real(jphook) :: hook_handle |
|---|
| 121 | |
|---|
| 122 | ! if (lhook) call dr_hook('radiation_matrix:single_mat_x_vec',0,hook_handle) |
|---|
| 123 | |
|---|
| 124 | ! Array-wise assignment |
|---|
| 125 | singlemat_x_vec = 0.0_jprb |
|---|
| 126 | |
|---|
| 127 | do j1 = 1,m |
|---|
| 128 | do j2 = 1,m |
|---|
| 129 | singlemat_x_vec(1:iend,j1) = singlemat_x_vec(1:iend,j1) & |
|---|
| 130 | & + A(j1,j2)*b(1:iend,j2) |
|---|
| 131 | end do |
|---|
| 132 | end do |
|---|
| 133 | |
|---|
| 134 | ! if (lhook) call dr_hook('radiation_matrix:single_mat_x_vec',1,hook_handle) |
|---|
| 135 | |
|---|
| 136 | end function singlemat_x_vec |
|---|
| 137 | |
|---|
| 138 | |
|---|
| 139 | ! --- SQUARE MATRIX-MATRIX MULTIPLICATION --- |
|---|
| 140 | |
|---|
| 141 | !--------------------------------------------------------------------- |
|---|
| 142 | ! Treat A and B each as n m-by-m square matrices (with the n |
|---|
| 143 | ! dimension varying fastest) and perform matrix multiplications on |
|---|
| 144 | ! all n matrix pairs |
|---|
| 145 | function mat_x_mat(n,iend,m,A,B,i_matrix_pattern) |
|---|
| 146 | |
|---|
| 147 | use yomhook, only : lhook, dr_hook, jphook |
|---|
| 148 | |
|---|
| 149 | integer, intent(in) :: n, m, iend |
|---|
| 150 | integer, intent(in), optional :: i_matrix_pattern |
|---|
| 151 | real(jprb), intent(in), dimension(:,:,:) :: A, B |
|---|
| 152 | |
|---|
| 153 | real(jprb), dimension(iend,m,m) :: mat_x_mat |
|---|
| 154 | integer :: j1, j2, j3 |
|---|
| 155 | integer :: mblock, m2block |
|---|
| 156 | integer :: i_actual_matrix_pattern |
|---|
| 157 | real(jphook) :: hook_handle |
|---|
| 158 | |
|---|
| 159 | if (lhook) call dr_hook('radiation_matrix:mat_x_mat',0,hook_handle) |
|---|
| 160 | |
|---|
| 161 | if (present(i_matrix_pattern)) then |
|---|
| 162 | i_actual_matrix_pattern = i_matrix_pattern |
|---|
| 163 | else |
|---|
| 164 | i_actual_matrix_pattern = IMatrixPatternDense |
|---|
| 165 | end if |
|---|
| 166 | |
|---|
| 167 | ! Array-wise assignment |
|---|
| 168 | mat_x_mat = 0.0_jprb |
|---|
| 169 | |
|---|
| 170 | if (i_actual_matrix_pattern == IMatrixPatternShortwave) then |
|---|
| 171 | ! Matrix has a sparsity pattern |
|---|
| 172 | ! (C D E) |
|---|
| 173 | ! A = (F G H) |
|---|
| 174 | ! (0 0 I) |
|---|
| 175 | mblock = m/3 |
|---|
| 176 | m2block = 2*mblock |
|---|
| 177 | ! Do the top-left (C, D, F, G) |
|---|
| 178 | do j2 = 1,m2block |
|---|
| 179 | do j1 = 1,m2block |
|---|
| 180 | do j3 = 1,m2block |
|---|
| 181 | mat_x_mat(1:iend,j1,j2) = mat_x_mat(1:iend,j1,j2) & |
|---|
| 182 | & + A(1:iend,j1,j3)*B(1:iend,j3,j2) |
|---|
| 183 | end do |
|---|
| 184 | end do |
|---|
| 185 | end do |
|---|
| 186 | do j2 = m2block+1,m |
|---|
| 187 | ! Do the top-right (E & H) |
|---|
| 188 | do j1 = 1,m2block |
|---|
| 189 | do j3 = 1,m |
|---|
| 190 | mat_x_mat(1:iend,j1,j2) = mat_x_mat(1:iend,j1,j2) & |
|---|
| 191 | & + A(1:iend,j1,j3)*B(1:iend,j3,j2) |
|---|
| 192 | end do |
|---|
| 193 | end do |
|---|
| 194 | ! Do the bottom-right (I) |
|---|
| 195 | do j1 = m2block+1,m |
|---|
| 196 | do j3 = m2block+1,m |
|---|
| 197 | mat_x_mat(1:iend,j1,j2) = mat_x_mat(1:iend,j1,j2) & |
|---|
| 198 | & + A(1:iend,j1,j3)*B(1:iend,j3,j2) |
|---|
| 199 | end do |
|---|
| 200 | end do |
|---|
| 201 | end do |
|---|
| 202 | else |
|---|
| 203 | ! Ordinary dense matrix |
|---|
| 204 | do j2 = 1,m |
|---|
| 205 | do j1 = 1,m |
|---|
| 206 | do j3 = 1,m |
|---|
| 207 | mat_x_mat(1:iend,j1,j2) = mat_x_mat(1:iend,j1,j2) & |
|---|
| 208 | & + A(1:iend,j1,j3)*B(1:iend,j3,j2) |
|---|
| 209 | end do |
|---|
| 210 | end do |
|---|
| 211 | end do |
|---|
| 212 | end if |
|---|
| 213 | |
|---|
| 214 | if (lhook) call dr_hook('radiation_matrix:mat_x_mat',1,hook_handle) |
|---|
| 215 | |
|---|
| 216 | end function mat_x_mat |
|---|
| 217 | |
|---|
| 218 | |
|---|
| 219 | !--------------------------------------------------------------------- |
|---|
| 220 | ! Treat A as an m-by-m matrix and B as n m-by-m square matrices |
|---|
| 221 | ! (with the n dimension varying fastest) and perform matrix |
|---|
| 222 | ! multiplications on the first iend matrix pairs |
|---|
| 223 | function singlemat_x_mat(n,iend,m,A,B) |
|---|
| 224 | |
|---|
| 225 | use yomhook, only : lhook, dr_hook, jphook |
|---|
| 226 | |
|---|
| 227 | integer, intent(in) :: n, m, iend |
|---|
| 228 | real(jprb), intent(in), dimension(m,m) :: A |
|---|
| 229 | real(jprb), intent(in), dimension(:,:,:) :: B |
|---|
| 230 | real(jprb), dimension(iend,m,m) :: singlemat_x_mat |
|---|
| 231 | |
|---|
| 232 | integer :: j1, j2, j3 |
|---|
| 233 | real(jphook) :: hook_handle |
|---|
| 234 | |
|---|
| 235 | if (lhook) call dr_hook('radiation_matrix:singlemat_x_mat',0,hook_handle) |
|---|
| 236 | |
|---|
| 237 | ! Array-wise assignment |
|---|
| 238 | singlemat_x_mat = 0.0_jprb |
|---|
| 239 | |
|---|
| 240 | do j2 = 1,m |
|---|
| 241 | do j1 = 1,m |
|---|
| 242 | do j3 = 1,m |
|---|
| 243 | singlemat_x_mat(1:iend,j1,j2) = singlemat_x_mat(1:iend,j1,j2) & |
|---|
| 244 | & + A(j1,j3)*B(1:iend,j3,j2) |
|---|
| 245 | end do |
|---|
| 246 | end do |
|---|
| 247 | end do |
|---|
| 248 | |
|---|
| 249 | if (lhook) call dr_hook('radiation_matrix:singlemat_x_mat',1,hook_handle) |
|---|
| 250 | |
|---|
| 251 | end function singlemat_x_mat |
|---|
| 252 | |
|---|
| 253 | |
|---|
| 254 | !--------------------------------------------------------------------- |
|---|
| 255 | ! Treat B as an m-by-m matrix and A as n m-by-m square matrices |
|---|
| 256 | ! (with the n dimension varying fastest) and perform matrix |
|---|
| 257 | ! multiplications on the first iend matrix pairs |
|---|
| 258 | function mat_x_singlemat(n,iend,m,A,B) |
|---|
| 259 | |
|---|
| 260 | use yomhook, only : lhook, dr_hook, jphook |
|---|
| 261 | |
|---|
| 262 | integer, intent(in) :: n, m, iend |
|---|
| 263 | real(jprb), intent(in), dimension(:,:,:) :: A |
|---|
| 264 | real(jprb), intent(in), dimension(m,m) :: B |
|---|
| 265 | |
|---|
| 266 | real(jprb), dimension(iend,m,m) :: mat_x_singlemat |
|---|
| 267 | integer :: j1, j2, j3 |
|---|
| 268 | real(jphook) :: hook_handle |
|---|
| 269 | |
|---|
| 270 | if (lhook) call dr_hook('radiation_matrix:mat_x_singlemat',0,hook_handle) |
|---|
| 271 | |
|---|
| 272 | ! Array-wise assignment |
|---|
| 273 | mat_x_singlemat = 0.0_jprb |
|---|
| 274 | |
|---|
| 275 | do j2 = 1,m |
|---|
| 276 | do j1 = 1,m |
|---|
| 277 | do j3 = 1,m |
|---|
| 278 | mat_x_singlemat(1:iend,j1,j2) = mat_x_singlemat(1:iend,j1,j2) & |
|---|
| 279 | & + A(1:iend,j1,j3)*B(j3,j2) |
|---|
| 280 | end do |
|---|
| 281 | end do |
|---|
| 282 | end do |
|---|
| 283 | |
|---|
| 284 | if (lhook) call dr_hook('radiation_matrix:mat_x_singlemat',1,hook_handle) |
|---|
| 285 | |
|---|
| 286 | end function mat_x_singlemat |
|---|
| 287 | |
|---|
| 288 | |
|---|
| 289 | !--------------------------------------------------------------------- |
|---|
| 290 | ! Compute I-A*B where I is the identity matrix and A & B are n |
|---|
| 291 | ! m-by-m square matrices |
|---|
| 292 | function identity_minus_mat_x_mat(n,iend,m,A,B,i_matrix_pattern) |
|---|
| 293 | |
|---|
| 294 | use yomhook, only : lhook, dr_hook, jphook |
|---|
| 295 | |
|---|
| 296 | integer, intent(in) :: n, m, iend |
|---|
| 297 | integer, intent(in), optional :: i_matrix_pattern |
|---|
| 298 | real(jprb), intent(in), dimension(:,:,:) :: A, B |
|---|
| 299 | real(jprb), dimension(iend,m,m) :: identity_minus_mat_x_mat |
|---|
| 300 | |
|---|
| 301 | integer :: j |
|---|
| 302 | real(jphook) :: hook_handle |
|---|
| 303 | |
|---|
| 304 | if (lhook) call dr_hook('radiation_matrix:identity_mat_x_mat',0,hook_handle) |
|---|
| 305 | |
|---|
| 306 | if (present(i_matrix_pattern)) then |
|---|
| 307 | identity_minus_mat_x_mat = mat_x_mat(n,iend,m,A,B,i_matrix_pattern) |
|---|
| 308 | else |
|---|
| 309 | identity_minus_mat_x_mat = mat_x_mat(n,iend,m,A,B) |
|---|
| 310 | end if |
|---|
| 311 | |
|---|
| 312 | identity_minus_mat_x_mat = - identity_minus_mat_x_mat |
|---|
| 313 | do j = 1,m |
|---|
| 314 | identity_minus_mat_x_mat(1:iend,j,j) & |
|---|
| 315 | & = 1.0_jprb + identity_minus_mat_x_mat(1:iend,j,j) |
|---|
| 316 | end do |
|---|
| 317 | |
|---|
| 318 | if (lhook) call dr_hook('radiation_matrix:identity_mat_x_mat',1,hook_handle) |
|---|
| 319 | |
|---|
| 320 | end function identity_minus_mat_x_mat |
|---|
| 321 | |
|---|
| 322 | |
|---|
| 323 | |
|---|
| 324 | !--------------------------------------------------------------------- |
|---|
| 325 | ! Replacement for matmul in the case that the first matrix is sparse |
|---|
| 326 | function sparse_x_dense(sparse, dense) |
|---|
| 327 | |
|---|
| 328 | real(jprb), intent(in) :: sparse(:,:), dense(:,:) |
|---|
| 329 | real(jprb) :: sparse_x_dense(size(sparse,1),size(dense,2)) |
|---|
| 330 | |
|---|
| 331 | integer :: j1, j2, j3 ! Loop indices |
|---|
| 332 | integer :: n1, n2, n3 ! Array sizes |
|---|
| 333 | |
|---|
| 334 | n1 = size(sparse,1) |
|---|
| 335 | n2 = size(sparse,2) |
|---|
| 336 | n3 = size(dense,2) |
|---|
| 337 | |
|---|
| 338 | sparse_x_dense = 0.0_jprb |
|---|
| 339 | do j2 = 1,n2 |
|---|
| 340 | do j1 = 1,n1 |
|---|
| 341 | if (sparse(j1,j2) /= 0.0_jprb) then |
|---|
| 342 | sparse_x_dense(j1,:) = sparse_x_dense(j1,:) + sparse(j1,j2)*dense(j2,:) |
|---|
| 343 | end if |
|---|
| 344 | end do |
|---|
| 345 | end do |
|---|
| 346 | |
|---|
| 347 | end function sparse_x_dense |
|---|
| 348 | |
|---|
| 349 | |
|---|
| 350 | ! --- REPEATEDLY SQUARE A MATRIX --- |
|---|
| 351 | |
|---|
| 352 | !--------------------------------------------------------------------- |
|---|
| 353 | ! Square m-by-m matrix "A" nrepeat times. A will be corrupted by |
|---|
| 354 | ! this function. |
|---|
| 355 | function repeated_square(m,A,nrepeat,i_matrix_pattern) |
|---|
| 356 | integer, intent(in) :: m, nrepeat |
|---|
| 357 | real(jprb), intent(inout) :: A(m,m) |
|---|
| 358 | integer, intent(in), optional :: i_matrix_pattern |
|---|
| 359 | real(jprb) :: repeated_square(m,m) |
|---|
| 360 | |
|---|
| 361 | integer :: j1, j2, j3, j4 |
|---|
| 362 | integer :: mblock, m2block |
|---|
| 363 | integer :: i_actual_matrix_pattern |
|---|
| 364 | |
|---|
| 365 | if (present(i_matrix_pattern)) then |
|---|
| 366 | i_actual_matrix_pattern = i_matrix_pattern |
|---|
| 367 | else |
|---|
| 368 | i_actual_matrix_pattern = IMatrixPatternDense |
|---|
| 369 | end if |
|---|
| 370 | |
|---|
| 371 | if (i_actual_matrix_pattern == IMatrixPatternShortwave) then |
|---|
| 372 | ! Matrix has a sparsity pattern |
|---|
| 373 | ! (C D E) |
|---|
| 374 | ! A = (F G H) |
|---|
| 375 | ! (0 0 I) |
|---|
| 376 | mblock = m/3 |
|---|
| 377 | m2block = 2*mblock |
|---|
| 378 | do j4 = 1,nrepeat |
|---|
| 379 | repeated_square = 0.0_jprb |
|---|
| 380 | ! Do the top-left (C, D, F & G) |
|---|
| 381 | do j2 = 1,m2block |
|---|
| 382 | do j1 = 1,m2block |
|---|
| 383 | do j3 = 1,m2block |
|---|
| 384 | repeated_square(j1,j2) = repeated_square(j1,j2) & |
|---|
| 385 | & + A(j1,j3)*A(j3,j2) |
|---|
| 386 | end do |
|---|
| 387 | end do |
|---|
| 388 | end do |
|---|
| 389 | do j2 = m2block+1, m |
|---|
| 390 | ! Do the top-right (E & H) |
|---|
| 391 | do j1 = 1,m2block |
|---|
| 392 | do j3 = 1,m |
|---|
| 393 | repeated_square(j1,j2) = repeated_square(j1,j2) & |
|---|
| 394 | & + A(j1,j3)*A(j3,j2) |
|---|
| 395 | end do |
|---|
| 396 | end do |
|---|
| 397 | ! Do the bottom-right (I) |
|---|
| 398 | do j1 = m2block+1, m |
|---|
| 399 | do j3 = m2block+1,m |
|---|
| 400 | repeated_square(j1,j2) = repeated_square(j1,j2) & |
|---|
| 401 | & + A(j1,j3)*A(j3,j2) |
|---|
| 402 | end do |
|---|
| 403 | end do |
|---|
| 404 | end do |
|---|
| 405 | if (j4 < nrepeat) then |
|---|
| 406 | A = repeated_square |
|---|
| 407 | end if |
|---|
| 408 | end do |
|---|
| 409 | else |
|---|
| 410 | ! Ordinary dense matrix |
|---|
| 411 | do j4 = 1,nrepeat |
|---|
| 412 | repeated_square = 0.0_jprb |
|---|
| 413 | do j2 = 1,m |
|---|
| 414 | do j1 = 1,m |
|---|
| 415 | do j3 = 1,m |
|---|
| 416 | repeated_square(j1,j2) = repeated_square(j1,j2) & |
|---|
| 417 | & + A(j1,j3)*A(j3,j2) |
|---|
| 418 | end do |
|---|
| 419 | end do |
|---|
| 420 | end do |
|---|
| 421 | if (j4 < nrepeat) then |
|---|
| 422 | A = repeated_square |
|---|
| 423 | end if |
|---|
| 424 | end do |
|---|
| 425 | end if |
|---|
| 426 | |
|---|
| 427 | end function repeated_square |
|---|
| 428 | |
|---|
| 429 | |
|---|
| 430 | ! --- SOLVE LINEAR EQUATIONS --- |
|---|
| 431 | |
|---|
| 432 | !--------------------------------------------------------------------- |
|---|
| 433 | ! Solve Ax=b to obtain x. Version optimized for 2x2 matrices using |
|---|
| 434 | ! Cramer's method: "A" contains n 2x2 matrices and "b" contains n |
|---|
| 435 | ! 2-element vectors; returns A^-1 b. |
|---|
| 436 | pure subroutine solve_vec_2(n,iend,A,b,x) |
|---|
| 437 | |
|---|
| 438 | integer, intent(in) :: n, iend |
|---|
| 439 | real(jprb), intent(in) :: A(:,:,:) |
|---|
| 440 | real(jprb), intent(in) :: b(:,:) |
|---|
| 441 | real(jprb), intent(out) :: x(:,:) |
|---|
| 442 | |
|---|
| 443 | real(jprb) :: inv_det(iend) |
|---|
| 444 | |
|---|
| 445 | inv_det = 1.0_jprb / ( A(1:iend,1,1)*A(1:iend,2,2) & |
|---|
| 446 | & - A(1:iend,1,2)*A(1:iend,2,1)) |
|---|
| 447 | |
|---|
| 448 | x(1:iend,1) = inv_det*(A(1:iend,2,2)*b(1:iend,1)-A(1:iend,1,2)*b(1:iend,2)) |
|---|
| 449 | x(1:iend,2) = inv_det*(A(1:iend,1,1)*b(1:iend,2)-A(1:iend,2,1)*b(1:iend,1)) |
|---|
| 450 | |
|---|
| 451 | end subroutine solve_vec_2 |
|---|
| 452 | |
|---|
| 453 | |
|---|
| 454 | !--------------------------------------------------------------------- |
|---|
| 455 | ! Solve AX=B to obtain X, i.e. the matrix right-hand-side version of |
|---|
| 456 | ! solve_vec_2, with A, X and B all containing n 2x2 matrices; |
|---|
| 457 | ! returns A^-1 B using Cramer's method. |
|---|
| 458 | pure subroutine solve_mat_2(n,iend,A,B,X) |
|---|
| 459 | integer, intent(in) :: n, iend |
|---|
| 460 | real(jprb), intent(in) :: A(:,:,:) |
|---|
| 461 | real(jprb), intent(in) :: B(:,:,:) |
|---|
| 462 | real(jprb), intent(out) :: X(:,:,:) |
|---|
| 463 | |
|---|
| 464 | real(jprb) :: inv_det(iend) |
|---|
| 465 | |
|---|
| 466 | inv_det = 1.0_jprb / ( A(1:iend,1,1)*A(1:iend,2,2) & |
|---|
| 467 | & - A(1:iend,1,2)*A(1:iend,2,1)) |
|---|
| 468 | |
|---|
| 469 | X(1:iend,1,1) = inv_det*( A(1:iend,2,2)*B(1:iend,1,1) & |
|---|
| 470 | & -A(1:iend,1,2)*B(1:iend,2,1)) |
|---|
| 471 | X(1:iend,2,1) = inv_det*( A(1:iend,1,1)*B(1:iend,2,1) & |
|---|
| 472 | & -A(1:iend,2,1)*B(1:iend,1,1)) |
|---|
| 473 | X(1:iend,1,2) = inv_det*( A(1:iend,2,2)*B(1:iend,1,2) & |
|---|
| 474 | & -A(1:iend,1,2)*B(1:iend,2,2)) |
|---|
| 475 | X(1:iend,2,2) = inv_det*( A(1:iend,1,1)*B(1:iend,2,2) & |
|---|
| 476 | & -A(1:iend,2,1)*B(1:iend,1,2)) |
|---|
| 477 | |
|---|
| 478 | end subroutine solve_mat_2 |
|---|
| 479 | |
|---|
| 480 | |
|---|
| 481 | !--------------------------------------------------------------------- |
|---|
| 482 | ! Solve Ax=b optimized for 3x3 matrices, using LU |
|---|
| 483 | ! factorization and substitution without pivoting. |
|---|
| 484 | pure subroutine solve_vec_3(n,iend,A,b,x) |
|---|
| 485 | integer, intent(in) :: n, iend |
|---|
| 486 | real(jprb), intent(in) :: A(:,:,:) |
|---|
| 487 | real(jprb), intent(in) :: b(:,:) |
|---|
| 488 | real(jprb), intent(out) :: x(:,:) |
|---|
| 489 | |
|---|
| 490 | real(jprb), dimension(iend) :: L21, L31, L32 |
|---|
| 491 | real(jprb), dimension(iend) :: U22, U23, U33 |
|---|
| 492 | real(jprb), dimension(iend) :: y2, y3 |
|---|
| 493 | |
|---|
| 494 | ! Some compilers unfortunately don't support assocate |
|---|
| 495 | ! associate (U11 => A(:,1,1), U12 => A(:,1,2), U13 => A(1,3), & |
|---|
| 496 | ! y1 => b(:,1), x1 => solve_vec3(:,1), & |
|---|
| 497 | ! x2 => solve_vec3(:,2), x3 => solve_vec3(:,3)) |
|---|
| 498 | |
|---|
| 499 | ! LU decomposition: |
|---|
| 500 | ! ( 1 ) (U11 U12 U13) |
|---|
| 501 | ! A = (L21 1 ) * ( U22 U23) |
|---|
| 502 | ! (L31 L32 1) ( U33) |
|---|
| 503 | L21 = A(1:iend,2,1) / A(1:iend,1,1) |
|---|
| 504 | L31 = A(1:iend,3,1) / A(1:iend,1,1) |
|---|
| 505 | U22 = A(1:iend,2,2) - L21*A(1:iend,1,2) |
|---|
| 506 | U23 = A(1:iend,2,3) - L21*A(1:iend,1,3) |
|---|
| 507 | L32 =(A(1:iend,3,2) - L31*A(1:iend,1,2)) / U22 |
|---|
| 508 | U33 = A(1:iend,3,3) - L31*A(1:iend,1,3) - L32*U23 |
|---|
| 509 | |
|---|
| 510 | ! Solve Ly = b by forward substitution |
|---|
| 511 | y2 = b(1:iend,2) - L21*b(1:iend,1) |
|---|
| 512 | y3 = b(1:iend,3) - L31*b(1:iend,1) - L32*y2 |
|---|
| 513 | |
|---|
| 514 | ! Solve Ux = y by back substitution |
|---|
| 515 | x(1:iend,3) = y3/U33 |
|---|
| 516 | x(1:iend,2) = (y2 - U23*x(1:iend,3)) / U22 |
|---|
| 517 | x(1:iend,1) = (b(1:iend,1) - A(1:iend,1,2)*x(1:iend,2) & |
|---|
| 518 | & - A(1:iend,1,3)*x(1:iend,3)) / A(1:iend,1,1) |
|---|
| 519 | ! end associate |
|---|
| 520 | |
|---|
| 521 | end subroutine solve_vec_3 |
|---|
| 522 | |
|---|
| 523 | |
|---|
| 524 | !--------------------------------------------------------------------- |
|---|
| 525 | ! Solve AX=B optimized for 3x3 matrices, using LU factorization and |
|---|
| 526 | ! substitution with no pivoting. |
|---|
| 527 | pure subroutine solve_mat_3(n,iend,A,B,X) |
|---|
| 528 | integer, intent(in) :: n, iend |
|---|
| 529 | real(jprb), intent(in) :: A(:,:,:) |
|---|
| 530 | real(jprb), intent(in) :: B(:,:,:) |
|---|
| 531 | real(jprb), intent(out) :: X(:,:,:) |
|---|
| 532 | |
|---|
| 533 | real(jprb), dimension(iend) :: L21, L31, L32 |
|---|
| 534 | real(jprb), dimension(iend) :: U22, U23, U33 |
|---|
| 535 | real(jprb), dimension(iend) :: y2, y3 |
|---|
| 536 | |
|---|
| 537 | integer :: j |
|---|
| 538 | |
|---|
| 539 | ! associate (U11 => A(:,1,1), U12 => A(:,1,2), U13 => A(1,3)) |
|---|
| 540 | ! LU decomposition: |
|---|
| 541 | ! ( 1 ) (U11 U12 U13) |
|---|
| 542 | ! A = (L21 1 ) * ( U22 U23) |
|---|
| 543 | ! (L31 L32 1) ( U33) |
|---|
| 544 | L21 = A(1:iend,2,1) / A(1:iend,1,1) |
|---|
| 545 | L31 = A(1:iend,3,1) / A(1:iend,1,1) |
|---|
| 546 | U22 = A(1:iend,2,2) - L21*A(1:iend,1,2) |
|---|
| 547 | U23 = A(1:iend,2,3) - L21*A(1:iend,1,3) |
|---|
| 548 | L32 =(A(1:iend,3,2) - L31*A(1:iend,1,2)) / U22 |
|---|
| 549 | U33 = A(1:iend,3,3) - L31*A(1:iend,1,3) - L32*U23 |
|---|
| 550 | |
|---|
| 551 | do j = 1,3 |
|---|
| 552 | ! Solve Ly = B(:,:,j) by forward substitution |
|---|
| 553 | ! y1 = B(:,1,j) |
|---|
| 554 | y2 = B(1:iend,2,j) - L21*B(1:iend,1,j) |
|---|
| 555 | y3 = B(1:iend,3,j) - L31*B(1:iend,1,j) - L32*y2 |
|---|
| 556 | ! Solve UX(:,:,j) = y by back substitution |
|---|
| 557 | X(1:iend,3,j) = y3 / U33 |
|---|
| 558 | X(1:iend,2,j) = (y2 - U23*X(1:iend,3,j)) / U22 |
|---|
| 559 | X(1:iend,1,j) = (B(1:iend,1,j) - A(1:iend,1,2)*X(1:iend,2,j) & |
|---|
| 560 | & - A(1:iend,1,3)*X(1:iend,3,j)) / A(1:iend,1,1) |
|---|
| 561 | end do |
|---|
| 562 | |
|---|
| 563 | end subroutine solve_mat_3 |
|---|
| 564 | |
|---|
| 565 | |
|---|
| 566 | !--------------------------------------------------------------------- |
|---|
| 567 | ! Return X = B A^-1 = (A^-T B)^T optimized for 3x3 matrices, where B |
|---|
| 568 | ! is a diagonal matrix, using LU factorization and substitution with |
|---|
| 569 | ! no pivoting. |
|---|
| 570 | pure subroutine diag_mat_right_divide_3(n,iend,A,B,X) |
|---|
| 571 | integer, intent(in) :: n, iend |
|---|
| 572 | real(jprb), intent(in) :: A(iend,3,3) |
|---|
| 573 | real(jprb), intent(in) :: B(iend,3) |
|---|
| 574 | real(jprb), intent(out) :: X(n,3,3) |
|---|
| 575 | |
|---|
| 576 | real(jprb), dimension(iend) :: L21, L31, L32 |
|---|
| 577 | real(jprb), dimension(iend) :: U22, U23, U33 |
|---|
| 578 | real(jprb), dimension(iend) :: y2, y3 |
|---|
| 579 | |
|---|
| 580 | ! associate (U11 => A(:,1,1), U12 => A(:,1,2), U13 => A(1,3)) |
|---|
| 581 | ! LU decomposition of the *transpose* of A: |
|---|
| 582 | ! ( 1 ) (U11 U12 U13) |
|---|
| 583 | ! A^T = (L21 1 ) * ( U22 U23) |
|---|
| 584 | ! (L31 L32 1) ( U33) |
|---|
| 585 | L21 = A(1:iend,1,2) / A(1:iend,1,1) |
|---|
| 586 | L31 = A(1:iend,1,3) / A(1:iend,1,1) |
|---|
| 587 | U22 = A(1:iend,2,2) - L21*A(1:iend,2,1) |
|---|
| 588 | U23 = A(1:iend,3,2) - L21*A(1:iend,3,1) |
|---|
| 589 | L32 =(A(1:iend,2,3) - L31*A(1:iend,2,1)) / U22 |
|---|
| 590 | U33 = A(1:iend,3,3) - L31*A(1:iend,3,1) - L32*U23 |
|---|
| 591 | |
|---|
| 592 | ! Solve X(1,:) = A^-T ( B(1) ) |
|---|
| 593 | ! ( 0 ) |
|---|
| 594 | ! ( 0 ) |
|---|
| 595 | ! Solve Ly = B(:,:,j) by forward substitution |
|---|
| 596 | ! y1 = B(:,1) |
|---|
| 597 | y2 = - L21*B(1:iend,1) |
|---|
| 598 | y3 = - L31*B(1:iend,1) - L32*y2 |
|---|
| 599 | ! Solve UX(:,:,j) = y by back substitution |
|---|
| 600 | X(1:iend,1,3) = y3 / U33 |
|---|
| 601 | X(1:iend,1,2) = (y2 - U23*X(1:iend,1,3)) / U22 |
|---|
| 602 | X(1:iend,1,1) = (B(1:iend,1) - A(1:iend,2,1)*X(1:iend,1,2) & |
|---|
| 603 | & - A(1:iend,3,1)*X(1:iend,1,3)) / A(1:iend,1,1) |
|---|
| 604 | |
|---|
| 605 | ! Solve X(2,:) = A^-T ( 0 ) |
|---|
| 606 | ! ( B(2) ) |
|---|
| 607 | ! ( 0 ) |
|---|
| 608 | ! Solve Ly = B(:,:,j) by forward substitution |
|---|
| 609 | ! y1 = 0 |
|---|
| 610 | ! y2 = B(1:iend,2) |
|---|
| 611 | y3 = - L32*B(1:iend,2) |
|---|
| 612 | ! Solve UX(:,:,j) = y by back substitution |
|---|
| 613 | X(1:iend,2,3) = y3 / U33 |
|---|
| 614 | X(1:iend,2,2) = (B(1:iend,2) - U23*X(1:iend,2,3)) / U22 |
|---|
| 615 | X(1:iend,2,1) = (-A(1:iend,2,1)*X(1:iend,2,2) & |
|---|
| 616 | & -A(1:iend,3,1)*X(1:iend,2,3)) / A(1:iend,1,1) |
|---|
| 617 | |
|---|
| 618 | ! Solve X(3,:) = A^-T ( 0 ) |
|---|
| 619 | ! ( 0 ) |
|---|
| 620 | ! ( B(3) ) |
|---|
| 621 | ! Solve Ly = B(:,:,j) by forward substitution |
|---|
| 622 | ! y1 = 0 |
|---|
| 623 | ! y2 = 0 |
|---|
| 624 | ! y3 = B(1:iend,3) |
|---|
| 625 | ! Solve UX(:,:,j) = y by back substitution |
|---|
| 626 | X(1:iend,3,3) = B(1:iend,3) / U33 |
|---|
| 627 | X(1:iend,3,2) = -U23*X(1:iend,3,3) / U22 |
|---|
| 628 | X(1:iend,3,1) = (-A(1:iend,2,1)*X(1:iend,3,2) & |
|---|
| 629 | & - A(1:iend,3,1)*X(1:iend,3,3)) / A(1:iend,1,1) |
|---|
| 630 | |
|---|
| 631 | end subroutine diag_mat_right_divide_3 |
|---|
| 632 | |
|---|
| 633 | |
|---|
| 634 | !--------------------------------------------------------------------- |
|---|
| 635 | ! Treat A as n m-by-m matrices and return the LU factorization of A |
|---|
| 636 | ! compressed into a single matrice (with L below the diagonal and U |
|---|
| 637 | ! on and above the diagonal; the diagonal elements of L are 1). No |
|---|
| 638 | ! pivoting is performed. |
|---|
| 639 | pure subroutine lu_factorization(n, iend, m, A, LU) |
|---|
| 640 | integer, intent(in) :: n, m, iend |
|---|
| 641 | real(jprb), intent(in) :: A(:,:,:) |
|---|
| 642 | real(jprb), intent(out) :: LU(iend,m,m) |
|---|
| 643 | |
|---|
| 644 | real(jprb) :: s(iend) |
|---|
| 645 | integer :: j1, j2, j3 |
|---|
| 646 | |
|---|
| 647 | ! This routine is adapted from an in-place one, so we first copy |
|---|
| 648 | ! the input into the output. |
|---|
| 649 | LU(1:iend,1:m,1:m) = A(1:iend,1:m,1:m) |
|---|
| 650 | |
|---|
| 651 | do j2 = 1, m |
|---|
| 652 | do j1 = 1, j2-1 |
|---|
| 653 | s = LU(1:iend,j1,j2) |
|---|
| 654 | do j3 = 1, j1-1 |
|---|
| 655 | s = s - LU(1:iend,j1,j3) * LU(1:iend,j3,j2) |
|---|
| 656 | end do |
|---|
| 657 | LU(1:iend,j1,j2) = s |
|---|
| 658 | end do |
|---|
| 659 | do j1 = j2, m |
|---|
| 660 | s = LU(1:iend,j1,j2) |
|---|
| 661 | do j3 = 1, j2-1 |
|---|
| 662 | s = s - LU(1:iend,j1,j3) * LU(1:iend,j3,j2) |
|---|
| 663 | end do |
|---|
| 664 | LU(1:iend,j1,j2) = s |
|---|
| 665 | end do |
|---|
| 666 | if (j2 /= m) then |
|---|
| 667 | s = 1.0_jprb / LU(1:iend,j2,j2) |
|---|
| 668 | do j1 = j2+1, m |
|---|
| 669 | LU(1:iend,j1,j2) = LU(1:iend,j1,j2) * s |
|---|
| 670 | end do |
|---|
| 671 | end if |
|---|
| 672 | end do |
|---|
| 673 | |
|---|
| 674 | end subroutine lu_factorization |
|---|
| 675 | |
|---|
| 676 | |
|---|
| 677 | !--------------------------------------------------------------------- |
|---|
| 678 | ! Treat LU as an LU-factorization of an original matrix A, and |
|---|
| 679 | ! return x where Ax=b. LU consists of n m-by-m matrices and b as n |
|---|
| 680 | ! m-element vectors. |
|---|
| 681 | pure subroutine lu_substitution(n,iend,m,LU,b,x) |
|---|
| 682 | ! CHECK: dimensions should be ":"? |
|---|
| 683 | integer, intent(in) :: n, m, iend |
|---|
| 684 | real(jprb), intent(in) :: LU(iend,m,m) |
|---|
| 685 | real(jprb), intent(in) :: b(:,:) |
|---|
| 686 | real(jprb), intent(out):: x(iend,m) |
|---|
| 687 | |
|---|
| 688 | integer :: j1, j2 |
|---|
| 689 | |
|---|
| 690 | x(1:iend,1:m) = b(1:iend,1:m) |
|---|
| 691 | |
|---|
| 692 | ! First solve Ly=b |
|---|
| 693 | do j2 = 2, m |
|---|
| 694 | do j1 = 1, j2-1 |
|---|
| 695 | x(1:iend,j2) = x(1:iend,j2) - x(1:iend,j1)*LU(1:iend,j2,j1) |
|---|
| 696 | end do |
|---|
| 697 | end do |
|---|
| 698 | ! Now solve Ux=y |
|---|
| 699 | do j2 = m, 1, -1 |
|---|
| 700 | do j1 = j2+1, m |
|---|
| 701 | x(1:iend,j2) = x(1:iend,j2) - x(1:iend,j1)*LU(1:iend,j2,j1) |
|---|
| 702 | end do |
|---|
| 703 | x(1:iend,j2) = x(1:iend,j2) / LU(1:iend,j2,j2) |
|---|
| 704 | end do |
|---|
| 705 | |
|---|
| 706 | end subroutine lu_substitution |
|---|
| 707 | |
|---|
| 708 | |
|---|
| 709 | !--------------------------------------------------------------------- |
|---|
| 710 | ! Return matrix X where AX=B. LU, A, X, B all consist of n m-by-m |
|---|
| 711 | ! matrices. |
|---|
| 712 | pure subroutine solve_mat_n(n,iend,m,A,B,X) |
|---|
| 713 | integer, intent(in) :: n, m, iend |
|---|
| 714 | real(jprb), intent(in) :: A(:,:,:) |
|---|
| 715 | real(jprb), intent(in) :: B(:,:,:) |
|---|
| 716 | real(jprb), intent(out):: X(iend,m,m) |
|---|
| 717 | |
|---|
| 718 | real(jprb) :: LU(iend,m,m) |
|---|
| 719 | |
|---|
| 720 | integer :: j |
|---|
| 721 | |
|---|
| 722 | call lu_factorization(n,iend,m,A,LU) |
|---|
| 723 | |
|---|
| 724 | do j = 1, m |
|---|
| 725 | call lu_substitution(n,iend,m,LU,B(1:,1:m,j),X(1:iend,1:m,j)) |
|---|
| 726 | ! call lu_substitution(n,iend,m,LU,B(1:n,1:m,j),X(1:iend,1:m,j)) |
|---|
| 727 | end do |
|---|
| 728 | |
|---|
| 729 | end subroutine solve_mat_n |
|---|
| 730 | |
|---|
| 731 | |
|---|
| 732 | !--------------------------------------------------------------------- |
|---|
| 733 | ! Solve Ax=b, where A consists of n m-by-m matrices and x and b |
|---|
| 734 | ! consist of n m-element vectors. For m=2 or m=3, this function |
|---|
| 735 | ! calls optimized versions, otherwise it uses general LU |
|---|
| 736 | ! decomposition without pivoting. |
|---|
| 737 | function solve_vec(n,iend,m,A,b) |
|---|
| 738 | |
|---|
| 739 | use yomhook, only : lhook, dr_hook, jphook |
|---|
| 740 | |
|---|
| 741 | integer, intent(in) :: n, m, iend |
|---|
| 742 | real(jprb), intent(in) :: A(:,:,:) |
|---|
| 743 | real(jprb), intent(in) :: b(:,:) |
|---|
| 744 | |
|---|
| 745 | real(jprb) :: solve_vec(iend,m) |
|---|
| 746 | real(jprb) :: LU(iend,m,m) |
|---|
| 747 | real(jphook) :: hook_handle |
|---|
| 748 | |
|---|
| 749 | if (lhook) call dr_hook('radiation_matrix:solve_vec',0,hook_handle) |
|---|
| 750 | |
|---|
| 751 | if (m == 2) then |
|---|
| 752 | call solve_vec_2(n,iend,A,b,solve_vec) |
|---|
| 753 | elseif (m == 3) then |
|---|
| 754 | call solve_vec_3(n,iend,A,b,solve_vec) |
|---|
| 755 | else |
|---|
| 756 | call lu_factorization(n,iend,m,A,LU) |
|---|
| 757 | call lu_substitution(n,iend,m,LU,b,solve_vec) |
|---|
| 758 | end if |
|---|
| 759 | |
|---|
| 760 | if (lhook) call dr_hook('radiation_matrix:solve_vec',1,hook_handle) |
|---|
| 761 | |
|---|
| 762 | end function solve_vec |
|---|
| 763 | |
|---|
| 764 | |
|---|
| 765 | !--------------------------------------------------------------------- |
|---|
| 766 | ! Solve AX=B, where A, X and B consist of n m-by-m matrices. For m=2 |
|---|
| 767 | ! or m=3, this function calls optimized versions, otherwise it uses |
|---|
| 768 | ! general LU decomposition without pivoting. |
|---|
| 769 | function solve_mat(n,iend,m,A,B) |
|---|
| 770 | |
|---|
| 771 | use yomhook, only : lhook, dr_hook, jphook |
|---|
| 772 | |
|---|
| 773 | integer, intent(in) :: n, m, iend |
|---|
| 774 | real(jprb), intent(in) :: A(:,:,:) |
|---|
| 775 | real(jprb), intent(in) :: B(:,:,:) |
|---|
| 776 | |
|---|
| 777 | real(jprb) :: solve_mat(iend,m,m) |
|---|
| 778 | real(jphook) :: hook_handle |
|---|
| 779 | |
|---|
| 780 | if (lhook) call dr_hook('radiation_matrix:solve_mat',0,hook_handle) |
|---|
| 781 | |
|---|
| 782 | if (m == 2) then |
|---|
| 783 | call solve_mat_2(n,iend,A,B,solve_mat) |
|---|
| 784 | elseif (m == 3) then |
|---|
| 785 | call solve_mat_3(n,iend,A,B,solve_mat) |
|---|
| 786 | else |
|---|
| 787 | call solve_mat_n(n,iend,m,A,B,solve_mat) |
|---|
| 788 | end if |
|---|
| 789 | |
|---|
| 790 | if (lhook) call dr_hook('radiation_matrix:solve_mat',1,hook_handle) |
|---|
| 791 | |
|---|
| 792 | end function solve_mat |
|---|
| 793 | |
|---|
| 794 | |
|---|
| 795 | ! --- MATRIX EXPONENTIATION --- |
|---|
| 796 | |
|---|
| 797 | !--------------------------------------------------------------------- |
|---|
| 798 | ! Perform matrix exponential of n m-by-m matrices stored in A (where |
|---|
| 799 | ! index n varies fastest) using the Higham scaling and squaring |
|---|
| 800 | ! method. The result is placed in A. This routine is intended for |
|---|
| 801 | ! speed so is accurate only to single precision. For simplicity and |
|---|
| 802 | ! to aid vectorization, the Pade approximant of order 7 is used for |
|---|
| 803 | ! all input matrices, perhaps leading to a few too many |
|---|
| 804 | ! multiplications for matrices with a small norm. |
|---|
| 805 | subroutine expm(n,iend,m,A,i_matrix_pattern) |
|---|
| 806 | |
|---|
| 807 | use yomhook, only : lhook, dr_hook, jphook |
|---|
| 808 | |
|---|
| 809 | integer, intent(in) :: n, m, iend |
|---|
| 810 | real(jprb), intent(inout) :: A(n,m,m) |
|---|
| 811 | integer, intent(in) :: i_matrix_pattern |
|---|
| 812 | |
|---|
| 813 | real(jprb), parameter :: theta(3) = (/4.258730016922831e-01_jprb, & |
|---|
| 814 | & 1.880152677804762e+00_jprb, & |
|---|
| 815 | & 3.925724783138660e+00_jprb/) |
|---|
| 816 | real(jprb), parameter :: c(8) = (/17297280.0_jprb, 8648640.0_jprb, & |
|---|
| 817 | & 1995840.0_jprb, 277200.0_jprb, 25200.0_jprb, & |
|---|
| 818 | & 1512.0_jprb, 56.0_jprb, 1.0_jprb/) |
|---|
| 819 | |
|---|
| 820 | real(jprb), dimension(iend,m,m) :: A2, A4, A6 |
|---|
| 821 | real(jprb), dimension(iend,m,m) :: U, V |
|---|
| 822 | |
|---|
| 823 | real(jprb) :: normA(iend), sum_column(iend) |
|---|
| 824 | |
|---|
| 825 | integer :: j1, j2, j3 |
|---|
| 826 | real(jprb) :: frac(iend) |
|---|
| 827 | integer :: expo(iend) |
|---|
| 828 | real(jprb) :: scaling(iend) |
|---|
| 829 | |
|---|
| 830 | real(jphook) :: hook_handle |
|---|
| 831 | |
|---|
| 832 | if (lhook) call dr_hook('radiation_matrix:expm',0,hook_handle) |
|---|
| 833 | |
|---|
| 834 | normA = 0.0_jprb |
|---|
| 835 | |
|---|
| 836 | ! Compute the 1-norms of A |
|---|
| 837 | do j3 = 1,m |
|---|
| 838 | sum_column(:) = 0.0_jprb |
|---|
| 839 | do j2 = 1,m |
|---|
| 840 | do j1 = 1,iend |
|---|
| 841 | sum_column(j1) = sum_column(j1) + abs(A(j1,j2,j3)) |
|---|
| 842 | end do |
|---|
| 843 | end do |
|---|
| 844 | do j1 = 1,iend |
|---|
| 845 | if (sum_column(j1) > normA(j1)) then |
|---|
| 846 | normA(j1) = sum_column(j1) |
|---|
| 847 | end if |
|---|
| 848 | end do |
|---|
| 849 | end do |
|---|
| 850 | |
|---|
| 851 | frac = fraction(normA/theta(3)) |
|---|
| 852 | expo = exponent(normA/theta(3)) |
|---|
| 853 | where (frac == 0.5_jprb) |
|---|
| 854 | expo = expo - 1 |
|---|
| 855 | end where |
|---|
| 856 | |
|---|
| 857 | where (expo < 0) |
|---|
| 858 | expo = 0 |
|---|
| 859 | end where |
|---|
| 860 | |
|---|
| 861 | ! Scale the input matrices by a power of 2 |
|---|
| 862 | scaling = 2.0_jprb**(-expo) |
|---|
| 863 | do j3 = 1,m |
|---|
| 864 | do j2 = 1,m |
|---|
| 865 | A(1:iend,j2,j3) = A(1:iend,j2,j3) * scaling |
|---|
| 866 | end do |
|---|
| 867 | end do |
|---|
| 868 | ! Pade approximant of degree 7 |
|---|
| 869 | A2 = mat_x_mat(n,iend,m,A, A, i_matrix_pattern) |
|---|
| 870 | A4 = mat_x_mat(n,iend,m,A2,A2,i_matrix_pattern) |
|---|
| 871 | A6 = mat_x_mat(n,iend,m,A2,A4,i_matrix_pattern) |
|---|
| 872 | |
|---|
| 873 | V = c(8)*A6 + c(6)*A4 + c(4)*A2 |
|---|
| 874 | do j3 = 1,m |
|---|
| 875 | V(:,j3,j3) = V(:,j3,j3) + c(2) |
|---|
| 876 | end do |
|---|
| 877 | U = mat_x_mat(n,iend,m,A,V,i_matrix_pattern) |
|---|
| 878 | V = c(7)*A6 + c(5)*A4 + c(3)*A2 |
|---|
| 879 | ! Add a multiple of the identity matrix |
|---|
| 880 | do j3 = 1,m |
|---|
| 881 | V(:,j3,j3) = V(:,j3,j3) + c(1) |
|---|
| 882 | end do |
|---|
| 883 | |
|---|
| 884 | V = V-U |
|---|
| 885 | U = 2.0_jprb*U |
|---|
| 886 | A(1:iend,1:m,1:m) = solve_mat(n,iend,m,V,U) |
|---|
| 887 | |
|---|
| 888 | ! Add the identity matrix |
|---|
| 889 | do j3 = 1,m |
|---|
| 890 | A(1:iend,j3,j3) = A(1:iend,j3,j3) + 1.0_jprb |
|---|
| 891 | end do |
|---|
| 892 | |
|---|
| 893 | ! Loop through the matrices |
|---|
| 894 | do j1 = 1,iend |
|---|
| 895 | if (expo(j1) > 0) then |
|---|
| 896 | ! Square matrix j1 expo(j1) times |
|---|
| 897 | A(j1,:,:) = repeated_square(m,A(j1,:,:),expo(j1),i_matrix_pattern) |
|---|
| 898 | end if |
|---|
| 899 | end do |
|---|
| 900 | |
|---|
| 901 | if (lhook) call dr_hook('radiation_matrix:expm',1,hook_handle) |
|---|
| 902 | |
|---|
| 903 | end subroutine expm |
|---|
| 904 | |
|---|
| 905 | |
|---|
| 906 | !--------------------------------------------------------------------- |
|---|
| 907 | ! Return the matrix exponential of n 2x2 matrices representing |
|---|
| 908 | ! conservative exchange between SPARTACUS regions, where the |
|---|
| 909 | ! matrices have the structure |
|---|
| 910 | ! (-a b) |
|---|
| 911 | ! ( a -b) |
|---|
| 912 | ! and a and b are assumed to be positive or zero. The solution uses |
|---|
| 913 | ! Putzer's algorithm - see the appendix of Hogan et al. (GMD 2018) |
|---|
| 914 | subroutine fast_expm_exchange_2(n,iend,a,b,R) |
|---|
| 915 | |
|---|
| 916 | use yomhook, only : lhook, dr_hook, jphook |
|---|
| 917 | |
|---|
| 918 | integer, intent(in) :: n, iend |
|---|
| 919 | real(jprb), dimension(n), intent(in) :: a, b |
|---|
| 920 | real(jprb), dimension(n,2,2), intent(out) :: R |
|---|
| 921 | |
|---|
| 922 | real(jprb), dimension(iend) :: factor |
|---|
| 923 | |
|---|
| 924 | real(jphook) :: hook_handle |
|---|
| 925 | |
|---|
| 926 | if (lhook) call dr_hook('radiation_matrix:fast_expm_exchange_2',0,hook_handle) |
|---|
| 927 | |
|---|
| 928 | ! Security to ensure that if a==b==0 then the identity matrix is returned |
|---|
| 929 | factor = (1.0_jprb - exp(-(a(1:iend)+b(1:iend))))/max(1.0e-12_jprb,a(1:iend)+b(1:iend)) |
|---|
| 930 | |
|---|
| 931 | R(1:iend,1,1) = 1.0_jprb - factor*a(1:iend) |
|---|
| 932 | R(1:iend,2,1) = factor*a(1:iend) |
|---|
| 933 | R(1:iend,1,2) = factor*b(1:iend) |
|---|
| 934 | R(1:iend,2,2) = 1.0_jprb - factor*b(1:iend) |
|---|
| 935 | |
|---|
| 936 | if (lhook) call dr_hook('radiation_matrix:fast_expm_exchange_2',1,hook_handle) |
|---|
| 937 | |
|---|
| 938 | end subroutine fast_expm_exchange_2 |
|---|
| 939 | |
|---|
| 940 | |
|---|
| 941 | !--------------------------------------------------------------------- |
|---|
| 942 | ! Return the matrix exponential of n 3x3 matrices representing |
|---|
| 943 | ! conservative exchange between SPARTACUS regions, where the |
|---|
| 944 | ! matrices have the structure |
|---|
| 945 | ! (-a b 0) |
|---|
| 946 | ! ( a -b-c d) |
|---|
| 947 | ! ( 0 c -d) |
|---|
| 948 | ! and a-d are assumed to be positive or zero. The solution uses the |
|---|
| 949 | ! diagonalization method and is a slight generalization of the |
|---|
| 950 | ! solution provided in the appendix of Hogan et al. (GMD 2018), |
|---|
| 951 | ! which assumed c==d. |
|---|
| 952 | subroutine fast_expm_exchange_3(n,iend,a,b,c,d,R) |
|---|
| 953 | |
|---|
| 954 | use yomhook, only : lhook, dr_hook, jphook |
|---|
| 955 | |
|---|
| 956 | real(jprb), parameter :: my_epsilon = 1.0e-12_jprb |
|---|
| 957 | |
|---|
| 958 | integer, intent(in) :: n, iend |
|---|
| 959 | real(jprb), dimension(n), intent(in) :: a, b, c, d |
|---|
| 960 | real(jprb), dimension(n,3,3), intent(out) :: R |
|---|
| 961 | |
|---|
| 962 | ! Eigenvectors |
|---|
| 963 | real(jprb), dimension(iend,3,3) :: V |
|---|
| 964 | |
|---|
| 965 | ! Non-zero Eigenvalues |
|---|
| 966 | real(jprb), dimension(iend) :: lambda1, lambda2 |
|---|
| 967 | |
|---|
| 968 | ! Diagonal matrix of the exponential of the eigenvalues |
|---|
| 969 | real(jprb), dimension(iend,3) :: diag |
|---|
| 970 | |
|---|
| 971 | ! Result of diag right-divided by V |
|---|
| 972 | real(jprb), dimension(iend,3,3) :: diag_rdivide_V |
|---|
| 973 | |
|---|
| 974 | ! Intermediate arrays |
|---|
| 975 | real(jprb), dimension(iend) :: tmp1, tmp2 |
|---|
| 976 | |
|---|
| 977 | integer :: j1, j2 |
|---|
| 978 | |
|---|
| 979 | real(jphook) :: hook_handle |
|---|
| 980 | |
|---|
| 981 | if (lhook) call dr_hook('radiation_matrix:fast_expm_exchange_3',0,hook_handle) |
|---|
| 982 | |
|---|
| 983 | ! Eigenvalues lambda1 and lambda2 |
|---|
| 984 | tmp1 = 0.5_jprb * (a(1:iend)+b(1:iend)+c(1:iend)+d(1:iend)) |
|---|
| 985 | tmp2 = sqrt(max(0.0_jprb, tmp1*tmp1 - (a(1:iend)*c(1:iend) & |
|---|
| 986 | & + a(1:iend)*d(1:iend) + b(1:iend)*d(1:iend)))) |
|---|
| 987 | ! The eigenvalues must not be the same or the LU decomposition |
|---|
| 988 | ! fails; this can occur occasionally in single precision, which we |
|---|
| 989 | ! avoid by limiting the minimum value of tmp2 |
|---|
| 990 | tmp2 = max(tmp2, epsilon(1.0_jprb) * tmp1) |
|---|
| 991 | lambda1 = -tmp1 + tmp2 |
|---|
| 992 | lambda2 = -tmp1 - tmp2 |
|---|
| 993 | |
|---|
| 994 | ! Eigenvectors, with securities such that if a--d are all zero |
|---|
| 995 | ! then V is non-singular and the identity matrix is returned in R; |
|---|
| 996 | ! note that lambdaX is typically negative so we need a |
|---|
| 997 | ! sign-preserving security |
|---|
| 998 | V(1:iend,1,1) = max(my_epsilon, b(1:iend)) & |
|---|
| 999 | & / sign(max(my_epsilon, abs(a(1:iend) + lambda1)), a(1:iend) + lambda1) |
|---|
| 1000 | V(1:iend,1,2) = b(1:iend) & |
|---|
| 1001 | & / sign(max(my_epsilon, abs(a(1:iend) + lambda2)), a(1:iend) + lambda2) |
|---|
| 1002 | V(1:iend,1,3) = b(1:iend) / max(my_epsilon, a(1:iend)) |
|---|
| 1003 | V(1:iend,2,:) = 1.0_jprb |
|---|
| 1004 | V(1:iend,3,1) = c(1:iend) & |
|---|
| 1005 | & / sign(max(my_epsilon, abs(d(1:iend) + lambda1)), d(1:iend) + lambda1) |
|---|
| 1006 | V(1:iend,3,2) = c(1:iend) & |
|---|
| 1007 | & / sign(max(my_epsilon, abs(d(1:iend) + lambda2)), d(1:iend) + lambda2) |
|---|
| 1008 | V(1:iend,3,3) = max(my_epsilon, c(1:iend)) / max(my_epsilon, d(1:iend)) |
|---|
| 1009 | |
|---|
| 1010 | diag(:,1) = exp(lambda1) |
|---|
| 1011 | diag(:,2) = exp(lambda2) |
|---|
| 1012 | diag(:,3) = 1.0_jprb |
|---|
| 1013 | |
|---|
| 1014 | ! Compute diag_rdivide_V = diag * V^-1 |
|---|
| 1015 | call diag_mat_right_divide_3(iend,iend,V,diag,diag_rdivide_V) |
|---|
| 1016 | |
|---|
| 1017 | ! Compute V * diag_rdivide_V |
|---|
| 1018 | do j1 = 1,3 |
|---|
| 1019 | do j2 = 1,3 |
|---|
| 1020 | R(1:iend,j2,j1) = V(1:iend,j2,1)*diag_rdivide_V(1:iend,1,j1) & |
|---|
| 1021 | & + V(1:iend,j2,2)*diag_rdivide_V(1:iend,2,j1) & |
|---|
| 1022 | & + V(1:iend,j2,3)*diag_rdivide_V(1:iend,3,j1) |
|---|
| 1023 | end do |
|---|
| 1024 | end do |
|---|
| 1025 | |
|---|
| 1026 | if (lhook) call dr_hook('radiation_matrix:fast_expm_exchange_3',1,hook_handle) |
|---|
| 1027 | |
|---|
| 1028 | end subroutine fast_expm_exchange_3 |
|---|
| 1029 | |
|---|
| 1030 | ! generic :: fast_expm_exchange => fast_expm_exchange_2, fast_expm_exchange_3 |
|---|
| 1031 | |
|---|
| 1032 | |
|---|
| 1033 | end module radiation_matrix |
|---|
