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