source: LMDZ6/trunk/libf/phylmd/ecrad/radiation/radiation_matrix.F90 @ 5408

Last change on this file since 5408 was 4773, checked in by idelkadi, 12 months ago
  • Update of Ecrad in LMDZ The same organization of the Ecrad offline version is retained in order to facilitate the updating of Ecrad in LMDZ and the comparison between online and offline results. version 1.6.1 of Ecrad (https://github.com/lguez/ecrad.git)
  • Implementation of the double call of Ecrad in LMDZ


File size: 33.4 KB
Line 
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
26module 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
56contains
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 
1033end module radiation_matrix
Note: See TracBrowser for help on using the repository browser.