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

Last change on this file since 4489 was 4489, checked in by lguez, 15 months ago

Merge LMDZ_ECRad branch back into trunk!
File size: 32.2 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
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
55contains
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    !    associate (U11 => A(:,1,1), U12 => A(:,1,2), U13 => A(1,3))
553    ! LU decomposition of the *transpose* of A:
554    !       ( 1        )   (U11 U12 U13)
555    ! A^T = (L21  1    ) * (    U22 U23)
556    !       (L31 L32  1)   (        U33)
557    L21 = A(1:iend,1,2) / A(1:iend,1,1)
558    L31 = A(1:iend,1,3) / A(1:iend,1,1)
559    U22 = A(1:iend,2,2) - L21*A(1:iend,2,1)
560    U23 = A(1:iend,3,2) - L21*A(1:iend,3,1)
561    L32 =(A(1:iend,2,3) - L31*A(1:iend,2,1)) / U22
562    U33 = A(1:iend,3,3) - L31*A(1:iend,3,1) - L32*U23
563
564    ! Solve X(1,:) = A^-T ( B(1) )
565    !                     (  0   )
566    !                     (  0   )
567    ! Solve Ly = B(:,:,j) by forward substitution
568    ! y1 = B(:,1)
569    y2 = - L21*B(1:iend,1)
570    y3 = - L31*B(1:iend,1) - L32*y2
571    ! Solve UX(:,:,j) = y by back substitution
572    X(1:iend,1,3) = y3 / U33
573    X(1:iend,1,2) = (y2 - U23*X(1:iend,1,3)) / U22
574    X(1:iend,1,1) = (B(1:iend,1) - A(1:iend,2,1)*X(1:iend,1,2) &
575         &          - A(1:iend,3,1)*X(1:iend,1,3)) / A(1:iend,1,1)
576
577    ! Solve X(2,:) = A^-T (  0   )
578    !                     ( B(2) )
579    !                     (  0   )
580    ! Solve Ly = B(:,:,j) by forward substitution
581    ! y1 = 0
582    ! y2 = B(1:iend,2)
583    y3 = - L32*B(1:iend,2)
584    ! Solve UX(:,:,j) = y by back substitution
585    X(1:iend,2,3) = y3 / U33
586    X(1:iend,2,2) = (B(1:iend,2) - U23*X(1:iend,2,3)) / U22
587    X(1:iend,2,1) = (-A(1:iend,2,1)*X(1:iend,2,2) &
588         &           -A(1:iend,3,1)*X(1:iend,2,3)) / A(1:iend,1,1)
589
590    ! Solve X(3,:) = A^-T (  0   )
591    !                     (  0   )
592    !                     ( B(3) )
593    ! Solve Ly = B(:,:,j) by forward substitution
594    ! y1 = 0
595    ! y2 = 0
596    ! y3 = B(1:iend,3)
597    ! Solve UX(:,:,j) = y by back substitution
598    X(1:iend,3,3) = B(1:iend,3) / U33
599    X(1:iend,3,2) = -U23*X(1:iend,3,3) / U22
600    X(1:iend,3,1) = (-A(1:iend,2,1)*X(1:iend,3,2) &
601         &          - A(1:iend,3,1)*X(1:iend,3,3)) / A(1:iend,1,1)
602
603  end subroutine diag_mat_right_divide_3
604
605
606  !---------------------------------------------------------------------
607  ! Treat A as n m-by-m matrices and return the LU factorization of A
608  ! compressed into a single matrice (with L below the diagonal and U
609  ! on and above the diagonal; the diagonal elements of L are 1). No
610  ! pivoting is performed.
611  pure subroutine lu_factorization(n, iend, m, A, LU)
612    integer,    intent(in)  :: n, m, iend
613    real(jprb), intent(in)  :: A(:,:,:)
614    real(jprb), intent(out) :: LU(iend,m,m)
615
616    real(jprb) :: s(iend)
617    integer    :: j1, j2, j3
618
619    ! This routine is adapted from an in-place one, so we first copy
620    ! the input into the output.
621    LU(1:iend,1:m,1:m) = A(1:iend,1:m,1:m)
622
623    do j2 = 1, m
624      do j1 = 1, j2-1
625        s = LU(1:iend,j1,j2)
626        do j3 = 1, j1-1
627          s = s - LU(1:iend,j1,j3) * LU(1:iend,j3,j2)
628        end do
629        LU(1:iend,j1,j2) = s
630      end do
631      do j1 = j2, m
632        s = LU(1:iend,j1,j2)
633        do j3 = 1, j2-1
634          s = s - LU(1:iend,j1,j3) * LU(1:iend,j3,j2)
635        end do
636        LU(1:iend,j1,j2) = s
637      end do
638      if (j2 /= m) then
639        s = 1.0_jprb / LU(1:iend,j2,j2)
640        do j1 = j2+1, m
641          LU(1:iend,j1,j2) = LU(1:iend,j1,j2) * s
642        end do
643      end if
644    end do
645
646  end subroutine lu_factorization
647
648
649  !---------------------------------------------------------------------
650  ! Treat LU as an LU-factorization of an original matrix A, and
651  ! return x where Ax=b. LU consists of n m-by-m matrices and b as n
652  ! m-element vectors.
653  pure subroutine lu_substitution(n,iend,m,LU,b,x)
654    ! CHECK: dimensions should be ":"?
655    integer,    intent(in) :: n, m, iend
656    real(jprb), intent(in) :: LU(iend,m,m)
657    real(jprb), intent(in) :: b(:,:)
658    real(jprb), intent(out):: x(iend,m)
659
660    integer :: j1, j2
661
662    x(1:iend,1:m) = b(1:iend,1:m)
663
664    ! First solve Ly=b
665    do j2 = 2, m
666      do j1 = 1, j2-1
667        x(1:iend,j2) = x(1:iend,j2) - x(1:iend,j1)*LU(1:iend,j2,j1)
668      end do
669    end do
670    ! Now solve Ux=y
671    do j2 = m, 1, -1
672      do j1 = j2+1, m
673        x(1:iend,j2) = x(1:iend,j2) - x(1:iend,j1)*LU(1:iend,j2,j1)
674      end do
675      x(1:iend,j2) = x(1:iend,j2) / LU(1:iend,j2,j2)
676    end do
677
678  end subroutine lu_substitution
679
680
681  !---------------------------------------------------------------------
682  ! Return matrix X where AX=B. LU, A, X, B all consist of n m-by-m
683  ! matrices.
684  pure subroutine solve_mat_n(n,iend,m,A,B,X)
685    integer,    intent(in) :: n, m, iend
686    real(jprb), intent(in) :: A(:,:,:)
687    real(jprb), intent(in) :: B(:,:,:)
688    real(jprb), intent(out):: X(iend,m,m)
689
690    real(jprb) :: LU(iend,m,m)
691
692    integer :: j
693
694    call lu_factorization(n,iend,m,A,LU)
695
696    do j = 1, m
697      call lu_substitution(n,iend,m,LU,B(1:,1:m,j),X(1:iend,1:m,j))
698!      call lu_substitution(n,iend,m,LU,B(1:n,1:m,j),X(1:iend,1:m,j))
699    end do
700
701  end subroutine solve_mat_n
702
703
704  !---------------------------------------------------------------------
705  ! Solve Ax=b, where A consists of n m-by-m matrices and x and b
706  ! consist of n m-element vectors. For m=2 or m=3, this function
707  ! calls optimized versions, otherwise it uses general LU
708  ! decomposition without pivoting.
709  function solve_vec(n,iend,m,A,b)
710
711    use yomhook, only : lhook, dr_hook
712
713    integer,    intent(in) :: n, m, iend
714    real(jprb), intent(in) :: A(:,:,:)
715    real(jprb), intent(in) :: b(:,:)
716
717    real(jprb)             :: solve_vec(iend,m)
718    real(jprb)             :: LU(iend,m,m)
719    real(jprb)             :: hook_handle
720
721    if (lhook) call dr_hook('radiation_matrix:solve_vec',0,hook_handle)
722
723    if (m == 2) then
724      call solve_vec_2(n,iend,A,b,solve_vec)
725    elseif (m == 3) then
726      call solve_vec_3(n,iend,A,b,solve_vec)
727    else
728      call lu_factorization(n,iend,m,A,LU)
729      call lu_substitution(n,iend,m,LU,b,solve_vec)
730    end if
731
732    if (lhook) call dr_hook('radiation_matrix:solve_vec',1,hook_handle)
733
734  end function solve_vec
735
736
737  !---------------------------------------------------------------------
738  ! Solve AX=B, where A, X and B consist of n m-by-m matrices. For m=2
739  ! or m=3, this function calls optimized versions, otherwise it uses
740  ! general LU decomposition without pivoting.
741  function solve_mat(n,iend,m,A,B)
742
743    use yomhook, only : lhook, dr_hook
744
745    integer,    intent(in)  :: n, m, iend
746    real(jprb), intent(in)  :: A(:,:,:)
747    real(jprb), intent(in)  :: B(:,:,:)
748
749    real(jprb)              :: solve_mat(iend,m,m)
750    real(jprb)              :: hook_handle
751
752    if (lhook) call dr_hook('radiation_matrix:solve_mat',0,hook_handle)
753
754    if (m == 2) then
755      call solve_mat_2(n,iend,A,B,solve_mat)
756    elseif (m == 3) then
757      call solve_mat_3(n,iend,A,B,solve_mat)
758    else
759      call solve_mat_n(n,iend,m,A,B,solve_mat)
760    end if
761
762    if (lhook) call dr_hook('radiation_matrix:solve_mat',1,hook_handle)
763
764  end function solve_mat
765
766
767  ! --- MATRIX EXPONENTIATION ---
768
769  !---------------------------------------------------------------------
770  ! Perform matrix exponential of n m-by-m matrices stored in A (where
771  ! index n varies fastest) using the Higham scaling and squaring
772  ! method. The result is placed in A. This routine is intended for
773  ! speed so is accurate only to single precision.  For simplicity and
774  ! to aid vectorization, the Pade approximant of order 7 is used for
775  ! all input matrices, perhaps leading to a few too many
776  ! multiplications for matrices with a small norm.
777  subroutine expm(n,iend,m,A,i_matrix_pattern)
778
779    use yomhook, only : lhook, dr_hook
780
781    integer,    intent(in)      :: n, m, iend
782    real(jprb), intent(inout)   :: A(n,m,m)
783    integer,    intent(in)      :: i_matrix_pattern
784
785    real(jprb), parameter :: theta(3) = (/4.258730016922831e-01_jprb, &
786         &                                1.880152677804762e+00_jprb, &
787         &                                3.925724783138660e+00_jprb/)
788    real(jprb), parameter :: c(8) = (/17297280.0_jprb, 8648640.0_jprb, &
789         &                1995840.0_jprb, 277200.0_jprb, 25200.0_jprb, &
790         &                1512.0_jprb, 56.0_jprb, 1.0_jprb/)
791
792    real(jprb), dimension(iend,m,m) :: A2, A4, A6
793    real(jprb), dimension(iend,m,m) :: U, V
794
795    real(jprb) :: normA(iend), sum_column(iend)
796
797    integer    :: j1, j2, j3
798    real(jprb) :: frac(iend)
799    integer    :: expo(iend)
800    real(jprb) :: scaling(iend)
801
802    real(jprb) :: hook_handle
803
804    if (lhook) call dr_hook('radiation_matrix:expm',0,hook_handle)
805
806    normA = 0.0_jprb
807
808    ! Compute the 1-norms of A
809    do j3 = 1,m
810      sum_column(:) = 0.0_jprb
811      do j2 = 1,m
812        do j1 = 1,iend
813          sum_column(j1) = sum_column(j1) + abs(A(j1,j2,j3))
814        end do
815      end do
816      do j1 = 1,iend
817        if (sum_column(j1) > normA(j1)) then
818          normA(j1) = sum_column(j1)
819        end if
820      end do
821    end do
822
823    frac = fraction(normA/theta(3))
824    expo = exponent(normA/theta(3))
825    where (frac == 0.5_jprb)
826      expo = expo - 1
827    end where
828
829    where (expo < 0)
830      expo = 0
831    end where
832
833    ! Scale the input matrices by a power of 2
834    scaling = 2.0_jprb**(-expo)
835    do j3 = 1,m
836      do j2 = 1,m
837        A(1:iend,j2,j3) = A(1:iend,j2,j3) * scaling
838      end do
839    end do
840    ! Pade approximant of degree 7
841    A2 = mat_x_mat(n,iend,m,A, A, i_matrix_pattern)
842    A4 = mat_x_mat(n,iend,m,A2,A2,i_matrix_pattern)
843    A6 = mat_x_mat(n,iend,m,A2,A4,i_matrix_pattern)
844
845    V = c(8)*A6 + c(6)*A4 + c(4)*A2
846    do j3 = 1,m
847      V(:,j3,j3) = V(:,j3,j3) + c(2)
848    end do
849    U = mat_x_mat(n,iend,m,A,V,i_matrix_pattern)
850    V = c(7)*A6 + c(5)*A4 + c(3)*A2
851    ! Add a multiple of the identity matrix
852    do j3 = 1,m
853      V(:,j3,j3) = V(:,j3,j3) + c(1)
854    end do
855
856    V = V-U
857    U = 2.0_jprb*U
858    A(1:iend,1:m,1:m) = solve_mat(n,iend,m,V,U)
859
860    ! Add the identity matrix
861    do j3 = 1,m
862      A(1:iend,j3,j3) = A(1:iend,j3,j3) + 1.0_jprb
863    end do
864
865    ! Loop through the matrices
866    do j1 = 1,iend
867      if (expo(j1) > 0) then
868        ! Square matrix j1 expo(j1) times         
869        A(j1,:,:) = repeated_square(m,A(j1,:,:),expo(j1),i_matrix_pattern)
870      end if
871    end do
872
873    if (lhook) call dr_hook('radiation_matrix:expm',1,hook_handle)
874
875  end subroutine expm
876
877
878  !---------------------------------------------------------------------
879  ! Return the matrix exponential of n 2x2 matrices representing
880  ! conservative exchange between SPARTACUS regions, where the
881  ! matrices have the structure
882  !   (-a   b)
883  !   ( a  -b)
884  ! and a and b are assumed to be positive or zero.  The solution uses
885  ! Putzer's algorithm - see the appendix of Hogan et al. (GMD 2018)
886  subroutine fast_expm_exchange_2(n,iend,a,b,R)
887
888    use yomhook, only : lhook, dr_hook
889
890    integer,                      intent(in)  :: n, iend
891    real(jprb), dimension(n),     intent(in)  :: a, b
892    real(jprb), dimension(n,2,2), intent(out) :: R
893
894    real(jprb), dimension(iend) :: factor
895
896    real(jprb) :: hook_handle
897
898    if (lhook) call dr_hook('radiation_matrix:fast_expm_exchange_2',0,hook_handle)
899
900    ! Security to ensure that if a==b==0 then the identity matrix is returned
901    factor = (1.0_jprb - exp(-(a(1:iend)+b(1:iend))))/max(1.0e-12_jprb,a(1:iend)+b(1:iend))
902
903    R(1:iend,1,1) = 1.0_jprb - factor*a(1:iend)
904    R(1:iend,2,1) = factor*a(1:iend)
905    R(1:iend,1,2) = factor*b(1:iend)
906    R(1:iend,2,2) = 1.0_jprb - factor*b(1:iend)
907
908    if (lhook) call dr_hook('radiation_matrix:fast_expm_exchange_2',1,hook_handle)
909
910  end subroutine fast_expm_exchange_2
911
912
913  !---------------------------------------------------------------------
914  ! Return the matrix exponential of n 3x3 matrices representing
915  ! conservative exchange between SPARTACUS regions, where the
916  ! matrices have the structure
917  !   (-a   b   0)
918  !   ( a -b-c  d)
919  !   ( 0   c  -d)
920  ! and a-d are assumed to be positive or zero.  The solution uses the
921  ! diagonalization method and is a slight generalization of the
922  ! solution provided in the appendix of Hogan et al. (GMD 2018),
923  ! which assumed c==d.
924  subroutine fast_expm_exchange_3(n,iend,a,b,c,d,R)
925
926    use yomhook, only : lhook, dr_hook
927
928    real(jprb), parameter :: my_epsilon = 1.0e-12_jprb
929
930    integer,                      intent(in)  :: n, iend
931    real(jprb), dimension(n),     intent(in)  :: a, b, c, d
932    real(jprb), dimension(n,3,3), intent(out) :: R
933
934    ! Eigenvectors
935    real(jprb), dimension(iend,3,3) :: V
936
937    ! Non-zero Eigenvalues
938    real(jprb), dimension(iend) :: lambda1, lambda2
939
940    ! Diagonal matrix of the exponential of the eigenvalues
941    real(jprb), dimension(iend,3) :: diag
942
943    ! Result of diag right-divided by V
944    real(jprb), dimension(iend,3,3) :: diag_rdivide_V
945
946    ! Intermediate arrays
947    real(jprb), dimension(iend) :: tmp1, tmp2
948
949    integer :: j1, j2
950
951    real(jprb) :: hook_handle
952
953    if (lhook) call dr_hook('radiation_matrix:fast_expm_exchange_3',0,hook_handle)
954
955    ! Eigenvalues
956    tmp1 = 0.5_jprb * (a(1:iend)+b(1:iend)+c(1:iend)+d(1:iend))
957    tmp2 = sqrt(tmp1*tmp1 - (a(1:iend)*c(1:iend) + a(1:iend)*d(1:iend) + b(1:iend)*d(1:iend)))
958    lambda1 = -tmp1 + tmp2
959    lambda2 = -tmp1 - tmp2
960
961    ! Eigenvectors, with securities such taht if a--d are all zero
962    ! then V is non-singular and the identity matrix is returned in R;
963    ! note that lambdaX is typically negative so we need a
964    ! sign-preserving security
965    V(1:iend,1,1) = max(my_epsilon, b(1:iend)) &
966         &  / sign(max(my_epsilon, abs(a(1:iend) + lambda1)), a(1:iend) + lambda1)
967    V(1:iend,1,2) = b(1:iend) &
968         &  / sign(max(my_epsilon, abs(a(1:iend) + lambda2)), a(1:iend) + lambda2)
969    V(1:iend,1,3) = b(1:iend) / max(my_epsilon, a(1:iend))
970    V(1:iend,2,:) = 1.0_jprb
971    V(1:iend,3,1) = c(1:iend) &
972         &  / sign(max(my_epsilon, abs(d(1:iend) + lambda1)), d(1:iend) + lambda1)
973    V(1:iend,3,2) = c(1:iend) &
974         &  / sign(max(my_epsilon, abs(d(1:iend) + lambda2)), d(1:iend) + lambda2)
975    V(1:iend,3,3) = max(my_epsilon, c(1:iend)) / max(my_epsilon, d(1:iend))
976   
977    diag(:,1) = exp(lambda1)
978    diag(:,2) = exp(lambda2)
979    diag(:,3) = 1.0_jprb
980
981    ! Compute diag_rdivide_V = diag * V^-1
982    call diag_mat_right_divide_3(iend,iend,V,diag,diag_rdivide_V)
983
984    ! Compute V * diag_rdivide_V
985    do j1 = 1,3
986      do j2 = 1,3
987        R(1:iend,j2,j1) = V(1:iend,j2,1)*diag_rdivide_V(1:iend,1,j1) &
988             &          + V(1:iend,j2,2)*diag_rdivide_V(1:iend,2,j1) &
989             &          + V(1:iend,j2,3)*diag_rdivide_V(1:iend,3,j1)
990      end do
991    end do
992
993    if (lhook) call dr_hook('radiation_matrix:fast_expm_exchange_3',1,hook_handle)
994
995  end subroutine fast_expm_exchange_3
996
997!  generic :: fast_expm_exchange => fast_expm_exchange_2, fast_expm_exchange_3
998
999
1000end module radiation_matrix
Note: See TracBrowser for help on using the repository browser.