Context Navigation

← Previous Revision
Latest Revision
Next Revision →
Blame
Revision Log

radiation_matrix.F90 @ 4791

Last change on this file since 4791 was 4773, checked in by idelkadi, 7 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
26	module radiation_matrix
27
28	use parkind1, only : jprb
29
30	implicit none
31	public
32
33	! Codes to describe sparseness pattern, where the SHORTWAVE
34	! pattern is of the form:
35	! (x x x)
36	! (x x x)
37	! (0 0 x)
38	! where each element may itself be a square matrix.
39	integer, parameter :: IMatrixPatternDense = 0
40	integer, parameter :: IMatrixPatternShortwave = 1
41
42	public :: mat_x_vec, singlemat_x_vec, mat_x_mat, &
43	& singlemat_x_mat, mat_x_singlemat, &
44	& identity_minus_mat_x_mat, solve_vec, solve_mat, expm, &
45	& fast_expm_exchange_2, fast_expm_exchange_3, &
46	& sparse_x_dense
47
48	private :: solve_vec_2, solve_vec_3, solve_mat_2, &
49	& solve_mat_3, lu_factorization, lu_substitution, solve_mat_n, &
50	& diag_mat_right_divide_3
51
52	interface fast_expm_exchange
53	module procedure fast_expm_exchange_2, fast_expm_exchange_3
54	end interface fast_expm_exchange
55
56	contains
57
58	! --- MATRIX-VECTOR MULTIPLICATION ---
59
60	!---------------------------------------------------------------------
61	! Treat A as n m-by-m square matrices (with the n dimension varying
62	! fastest) and b as n m-element vectors, and perform matrix-vector
63	! multiplications on first iend pairs
64	function mat_x_vec(n,iend,m,A,b,do_top_left_only_in)
65
66	use yomhook, only : lhook, dr_hook, jphook
67
68	integer, intent(in) :: n, m, iend
69	real(jprb), intent(in), dimension(:,:,:) :: A
70	real(jprb), intent(in), dimension(:,:) :: b
71	logical, intent(in), optional :: do_top_left_only_in
72	real(jprb), dimension(iend,m):: mat_x_vec
73
74	integer :: j1, j2
75	logical :: do_top_left_only
76
77	real(jphook) :: hook_handle
78
79	if (lhook) call dr_hook('radiation_matrix:mat_x_vec',0,hook_handle)
80
81	if (present(do_top_left_only_in)) then
82	do_top_left_only = do_top_left_only_in
83	else
84	do_top_left_only = .false.
85	end if
86
87	! Array-wise assignment
88	mat_x_vec = 0.0_jprb
89
90	if (do_top_left_only) then
91	mat_x_vec(1:iend,1) = A(1:iend,1,1)*b(1:iend,1)
92	else
93	do j1 = 1,m
94	do j2 = 1,m
95	mat_x_vec(1:iend,j1) = mat_x_vec(1:iend,j1) &
96	& + A(1:iend,j1,j2)*b(1:iend,j2)
97	end do
98	end do
99	end if
100
101	if (lhook) call dr_hook('radiation_matrix:mat_x_vec',1,hook_handle)
102
103	end function mat_x_vec
104
105
106	!---------------------------------------------------------------------
107	! Treat A as an m-by-m square matrix and b as n m-element vectors
108	! (with the n dimension varying fastest), and perform matrix-vector
109	! multiplications on first iend pairs
110	function singlemat_x_vec(n,iend,m,A,b)
111
112	! use yomhook, only : lhook, dr_hook, jphook
113
114	integer, intent(in) :: n, m, iend
115	real(jprb), intent(in), dimension(m,m) :: A
116	real(jprb), intent(in), dimension(:,:) :: b
117	real(jprb), dimension(iend,m) :: singlemat_x_vec
118
119	integer :: j1, j2
120	! real(jphook) :: hook_handle
121
122	! if (lhook) call dr_hook('radiation_matrix:single_mat_x_vec',0,hook_handle)
123
124	! Array-wise assignment
125	singlemat_x_vec = 0.0_jprb
126
127	do j1 = 1,m
128	do j2 = 1,m
129	singlemat_x_vec(1:iend,j1) = singlemat_x_vec(1:iend,j1) &
130	& + A(j1,j2)*b(1:iend,j2)
131	end do
132	end do
133
134	! if (lhook) call dr_hook('radiation_matrix:single_mat_x_vec',1,hook_handle)
135
136	end function singlemat_x_vec
137
138
139	! --- SQUARE MATRIX-MATRIX MULTIPLICATION ---
140
141	!---------------------------------------------------------------------
142	! Treat A and B each as n m-by-m square matrices (with the n
143	! dimension varying fastest) and perform matrix multiplications on
144	! all n matrix pairs
145	function mat_x_mat(n,iend,m,A,B,i_matrix_pattern)
146
147	use yomhook, only : lhook, dr_hook, jphook
148
149	integer, intent(in) :: n, m, iend
150	integer, intent(in), optional :: i_matrix_pattern
151	real(jprb), intent(in), dimension(:,:,:) :: A, B
152
153	real(jprb), dimension(iend,m,m) :: mat_x_mat
154	integer :: j1, j2, j3
155	integer :: mblock, m2block
156	integer :: i_actual_matrix_pattern
157	real(jphook) :: hook_handle
158
159	if (lhook) call dr_hook('radiation_matrix:mat_x_mat',0,hook_handle)
160
161	if (present(i_matrix_pattern)) then
162	i_actual_matrix_pattern = i_matrix_pattern
163	else
164	i_actual_matrix_pattern = IMatrixPatternDense
165	end if
166
167	! Array-wise assignment
168	mat_x_mat = 0.0_jprb
169
170	if (i_actual_matrix_pattern == IMatrixPatternShortwave) then
171	! Matrix has a sparsity pattern
172	! (C D E)
173	! A = (F G H)
174	! (0 0 I)
175	mblock = m/3
176	m2block = 2*mblock
177	! Do the top-left (C, D, F, G)
178	do j2 = 1,m2block
179	do j1 = 1,m2block
180	do j3 = 1,m2block
181	mat_x_mat(1:iend,j1,j2) = mat_x_mat(1:iend,j1,j2) &
182	& + A(1:iend,j1,j3)*B(1:iend,j3,j2)
183	end do
184	end do
185	end do
186	do j2 = m2block+1,m
187	! Do the top-right (E & H)
188	do j1 = 1,m2block
189	do j3 = 1,m
190	mat_x_mat(1:iend,j1,j2) = mat_x_mat(1:iend,j1,j2) &
191	& + A(1:iend,j1,j3)*B(1:iend,j3,j2)
192	end do
193	end do
194	! Do the bottom-right (I)
195	do j1 = m2block+1,m
196	do j3 = m2block+1,m
197	mat_x_mat(1:iend,j1,j2) = mat_x_mat(1:iend,j1,j2) &
198	& + A(1:iend,j1,j3)*B(1:iend,j3,j2)
199	end do
200	end do
201	end do
202	else
203	! Ordinary dense matrix
204	do j2 = 1,m
205	do j1 = 1,m
206	do j3 = 1,m
207	mat_x_mat(1:iend,j1,j2) = mat_x_mat(1:iend,j1,j2) &
208	& + A(1:iend,j1,j3)*B(1:iend,j3,j2)
209	end do
210	end do
211	end do
212	end if
213
214	if (lhook) call dr_hook('radiation_matrix:mat_x_mat',1,hook_handle)
215
216	end function mat_x_mat
217
218
219	!---------------------------------------------------------------------
220	! Treat A as an m-by-m matrix and B as n m-by-m square matrices
221	! (with the n dimension varying fastest) and perform matrix
222	! multiplications on the first iend matrix pairs
223	function singlemat_x_mat(n,iend,m,A,B)
224
225	use yomhook, only : lhook, dr_hook, jphook
226
227	integer, intent(in) :: n, m, iend
228	real(jprb), intent(in), dimension(m,m) :: A
229	real(jprb), intent(in), dimension(:,:,:) :: B
230	real(jprb), dimension(iend,m,m) :: singlemat_x_mat
231
232	integer :: j1, j2, j3
233	real(jphook) :: hook_handle
234
235	if (lhook) call dr_hook('radiation_matrix:singlemat_x_mat',0,hook_handle)
236
237	! Array-wise assignment
238	singlemat_x_mat = 0.0_jprb
239
240	do j2 = 1,m
241	do j1 = 1,m
242	do j3 = 1,m
243	singlemat_x_mat(1:iend,j1,j2) = singlemat_x_mat(1:iend,j1,j2) &
244	& + A(j1,j3)*B(1:iend,j3,j2)
245	end do
246	end do
247	end do
248
249	if (lhook) call dr_hook('radiation_matrix:singlemat_x_mat',1,hook_handle)
250
251	end function singlemat_x_mat
252
253
254	!---------------------------------------------------------------------
255	! Treat B as an m-by-m matrix and A as n m-by-m square matrices
256	! (with the n dimension varying fastest) and perform matrix
257	! multiplications on the first iend matrix pairs
258	function mat_x_singlemat(n,iend,m,A,B)
259
260	use yomhook, only : lhook, dr_hook, jphook
261
262	integer, intent(in) :: n, m, iend
263	real(jprb), intent(in), dimension(:,:,:) :: A
264	real(jprb), intent(in), dimension(m,m) :: B
265
266	real(jprb), dimension(iend,m,m) :: mat_x_singlemat
267	integer :: j1, j2, j3
268	real(jphook) :: hook_handle
269
270	if (lhook) call dr_hook('radiation_matrix:mat_x_singlemat',0,hook_handle)
271
272	! Array-wise assignment
273	mat_x_singlemat = 0.0_jprb
274
275	do j2 = 1,m
276	do j1 = 1,m
277	do j3 = 1,m
278	mat_x_singlemat(1:iend,j1,j2) = mat_x_singlemat(1:iend,j1,j2) &
279	& + A(1:iend,j1,j3)*B(j3,j2)
280	end do
281	end do
282	end do
283
284	if (lhook) call dr_hook('radiation_matrix:mat_x_singlemat',1,hook_handle)
285
286	end function mat_x_singlemat
287
288
289	!---------------------------------------------------------------------
290	! Compute I-A*B where I is the identity matrix and A & B are n
291	! m-by-m square matrices
292	function identity_minus_mat_x_mat(n,iend,m,A,B,i_matrix_pattern)
293
294	use yomhook, only : lhook, dr_hook, jphook
295
296	integer, intent(in) :: n, m, iend
297	integer, intent(in), optional :: i_matrix_pattern
298	real(jprb), intent(in), dimension(:,:,:) :: A, B
299	real(jprb), dimension(iend,m,m) :: identity_minus_mat_x_mat
300
301	integer :: j
302	real(jphook) :: hook_handle
303
304	if (lhook) call dr_hook('radiation_matrix:identity_mat_x_mat',0,hook_handle)
305
306	if (present(i_matrix_pattern)) then
307	identity_minus_mat_x_mat = mat_x_mat(n,iend,m,A,B,i_matrix_pattern)
308	else
309	identity_minus_mat_x_mat = mat_x_mat(n,iend,m,A,B)
310	end if
311
312	identity_minus_mat_x_mat = - identity_minus_mat_x_mat
313	do j = 1,m
314	identity_minus_mat_x_mat(1:iend,j,j) &
315	& = 1.0_jprb + identity_minus_mat_x_mat(1:iend,j,j)
316	end do
317
318	if (lhook) call dr_hook('radiation_matrix:identity_mat_x_mat',1,hook_handle)
319
320	end function identity_minus_mat_x_mat
321
322
323
324	!---------------------------------------------------------------------
325	! Replacement for matmul in the case that the first matrix is sparse
326	function sparse_x_dense(sparse, dense)
327
328	real(jprb), intent(in) :: sparse(:,:), dense(:,:)
329	real(jprb) :: sparse_x_dense(size(sparse,1),size(dense,2))
330
331	integer :: j1, j2, j3 ! Loop indices
332	integer :: n1, n2, n3 ! Array sizes
333
334	n1 = size(sparse,1)
335	n2 = size(sparse,2)
336	n3 = size(dense,2)
337
338	sparse_x_dense = 0.0_jprb
339	do j2 = 1,n2
340	do j1 = 1,n1
341	if (sparse(j1,j2) /= 0.0_jprb) then
342	sparse_x_dense(j1,:) = sparse_x_dense(j1,:) + sparse(j1,j2)*dense(j2,:)
343	end if
344	end do
345	end do
346
347	end function sparse_x_dense
348
349
350	! --- REPEATEDLY SQUARE A MATRIX ---
351
352	!---------------------------------------------------------------------
353	! Square m-by-m matrix "A" nrepeat times. A will be corrupted by
354	! this function.
355	function repeated_square(m,A,nrepeat,i_matrix_pattern)
356	integer, intent(in) :: m, nrepeat
357	real(jprb), intent(inout) :: A(m,m)
358	integer, intent(in), optional :: i_matrix_pattern
359	real(jprb) :: repeated_square(m,m)
360
361	integer :: j1, j2, j3, j4
362	integer :: mblock, m2block
363	integer :: i_actual_matrix_pattern
364
365	if (present(i_matrix_pattern)) then
366	i_actual_matrix_pattern = i_matrix_pattern
367	else
368	i_actual_matrix_pattern = IMatrixPatternDense
369	end if
370
371	if (i_actual_matrix_pattern == IMatrixPatternShortwave) then
372	! Matrix has a sparsity pattern
373	! (C D E)
374	! A = (F G H)
375	! (0 0 I)
376	mblock = m/3
377	m2block = 2*mblock
378	do j4 = 1,nrepeat
379	repeated_square = 0.0_jprb
380	! Do the top-left (C, D, F & G)
381	do j2 = 1,m2block
382	do j1 = 1,m2block
383	do j3 = 1,m2block
384	repeated_square(j1,j2) = repeated_square(j1,j2) &
385	& + A(j1,j3)*A(j3,j2)
386	end do
387	end do
388	end do
389	do j2 = m2block+1, m
390	! Do the top-right (E & H)
391	do j1 = 1,m2block
392	do j3 = 1,m
393	repeated_square(j1,j2) = repeated_square(j1,j2) &
394	& + A(j1,j3)*A(j3,j2)
395	end do
396	end do
397	! Do the bottom-right (I)
398	do j1 = m2block+1, m
399	do j3 = m2block+1,m
400	repeated_square(j1,j2) = repeated_square(j1,j2) &
401	& + A(j1,j3)*A(j3,j2)
402	end do
403	end do
404	end do
405	if (j4 < nrepeat) then
406	A = repeated_square
407	end if
408	end do
409	else
410	! Ordinary dense matrix
411	do j4 = 1,nrepeat
412	repeated_square = 0.0_jprb
413	do j2 = 1,m
414	do j1 = 1,m
415	do j3 = 1,m
416	repeated_square(j1,j2) = repeated_square(j1,j2) &
417	& + A(j1,j3)*A(j3,j2)
418	end do
419	end do
420	end do
421	if (j4 < nrepeat) then
422	A = repeated_square
423	end if
424	end do
425	end if
426
427	end function repeated_square
428
429
430	! --- SOLVE LINEAR EQUATIONS ---
431
432	!---------------------------------------------------------------------
433	! Solve Ax=b to obtain x. Version optimized for 2x2 matrices using
434	! Cramer's method: "A" contains n 2x2 matrices and "b" contains n
435	! 2-element vectors; returns A^-1 b.
436	pure subroutine solve_vec_2(n,iend,A,b,x)
437
438	integer, intent(in) :: n, iend
439	real(jprb), intent(in) :: A(:,:,:)
440	real(jprb), intent(in) :: b(:,:)
441	real(jprb), intent(out) :: x(:,:)
442
443	real(jprb) :: inv_det(iend)
444
445	inv_det = 1.0_jprb / ( A(1:iend,1,1)*A(1:iend,2,2) &
446	& - A(1:iend,1,2)*A(1:iend,2,1))
447
448	x(1:iend,1) = inv_det(A(1:iend,2,2)b(1:iend,1)-A(1:iend,1,2)*b(1:iend,2))
449	x(1:iend,2) = inv_det(A(1:iend,1,1)b(1:iend,2)-A(1:iend,2,1)*b(1:iend,1))
450
451	end subroutine solve_vec_2
452
453
454	!---------------------------------------------------------------------
455	! Solve AX=B to obtain X, i.e. the matrix right-hand-side version of
456	! solve_vec_2, with A, X and B all containing n 2x2 matrices;
457	! returns A^-1 B using Cramer's method.
458	pure subroutine solve_mat_2(n,iend,A,B,X)
459	integer, intent(in) :: n, iend
460	real(jprb), intent(in) :: A(:,:,:)
461	real(jprb), intent(in) :: B(:,:,:)
462	real(jprb), intent(out) :: X(:,:,:)
463
464	real(jprb) :: inv_det(iend)
465
466	inv_det = 1.0_jprb / ( A(1:iend,1,1)*A(1:iend,2,2) &
467	& - A(1:iend,1,2)*A(1:iend,2,1))
468
469	X(1:iend,1,1) = inv_det( A(1:iend,2,2)B(1:iend,1,1) &
470	& -A(1:iend,1,2)*B(1:iend,2,1))
471	X(1:iend,2,1) = inv_det( A(1:iend,1,1)B(1:iend,2,1) &
472	& -A(1:iend,2,1)*B(1:iend,1,1))
473	X(1:iend,1,2) = inv_det( A(1:iend,2,2)B(1:iend,1,2) &
474	& -A(1:iend,1,2)*B(1:iend,2,2))
475	X(1:iend,2,2) = inv_det( A(1:iend,1,1)B(1:iend,2,2) &
476	& -A(1:iend,2,1)*B(1:iend,1,2))
477
478	end subroutine solve_mat_2
479
480
481	!---------------------------------------------------------------------
482	! Solve Ax=b optimized for 3x3 matrices, using LU
483	! factorization and substitution without pivoting.
484	pure subroutine solve_vec_3(n,iend,A,b,x)
485	integer, intent(in) :: n, iend
486	real(jprb), intent(in) :: A(:,:,:)
487	real(jprb), intent(in) :: b(:,:)
488	real(jprb), intent(out) :: x(:,:)
489
490	real(jprb), dimension(iend) :: L21, L31, L32
491	real(jprb), dimension(iend) :: U22, U23, U33
492	real(jprb), dimension(iend) :: y2, y3
493
494	! Some compilers unfortunately don't support assocate
495	! associate (U11 => A(:,1,1), U12 => A(:,1,2), U13 => A(1,3), &
496	! y1 => b(:,1), x1 => solve_vec3(:,1), &
497	! x2 => solve_vec3(:,2), x3 => solve_vec3(:,3))
498
499	! LU decomposition:
500	! ( 1 ) (U11 U12 U13)
501	! A = (L21 1 ) * ( U22 U23)
502	! (L31 L32 1) ( U33)
503	L21 = A(1:iend,2,1) / A(1:iend,1,1)
504	L31 = A(1:iend,3,1) / A(1:iend,1,1)
505	U22 = A(1:iend,2,2) - L21*A(1:iend,1,2)
506	U23 = A(1:iend,2,3) - L21*A(1:iend,1,3)
507	L32 =(A(1:iend,3,2) - L31*A(1:iend,1,2)) / U22
508	U33 = A(1:iend,3,3) - L31A(1:iend,1,3) - L32U23
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) - L31b(1:iend,1) - L32y2
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) - L31A(1:iend,1,3) - L32U23
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) - L31B(1:iend,1,j) - L32y2
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) - L31A(1:iend,3,1) - L32U23
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 = - L31B(1:iend,1) - L32y2
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, tmp1tmp1 - (a(1:iend)c(1:iend) &
986	& + a(1:iend)d(1:iend) + b(1:iend)d(1:iend))))
987	! The eigenvalues must not be the same or the LU decomposition
988	! fails; this can occur occasionally in single precision, which we
989	! avoid by limiting the minimum value of tmp2
990	tmp2 = max(tmp2, epsilon(1.0_jprb) * tmp1)
991	lambda1 = -tmp1 + tmp2
992	lambda2 = -tmp1 - tmp2
993
994	! Eigenvectors, with securities such that if a--d are all zero
995	! then V is non-singular and the identity matrix is returned in R;
996	! note that lambdaX is typically negative so we need a
997	! sign-preserving security
998	V(1:iend,1,1) = max(my_epsilon, b(1:iend)) &
999	& / sign(max(my_epsilon, abs(a(1:iend) + lambda1)), a(1:iend) + lambda1)
1000	V(1:iend,1,2) = b(1:iend) &
1001	& / sign(max(my_epsilon, abs(a(1:iend) + lambda2)), a(1:iend) + lambda2)
1002	V(1:iend,1,3) = b(1:iend) / max(my_epsilon, a(1:iend))
1003	V(1:iend,2,:) = 1.0_jprb
1004	V(1:iend,3,1) = c(1:iend) &
1005	& / sign(max(my_epsilon, abs(d(1:iend) + lambda1)), d(1:iend) + lambda1)
1006	V(1:iend,3,2) = c(1:iend) &
1007	& / sign(max(my_epsilon, abs(d(1:iend) + lambda2)), d(1:iend) + lambda2)
1008	V(1:iend,3,3) = max(my_epsilon, c(1:iend)) / max(my_epsilon, d(1:iend))
1009
1010	diag(:,1) = exp(lambda1)
1011	diag(:,2) = exp(lambda2)
1012	diag(:,3) = 1.0_jprb
1013
1014	! Compute diag_rdivide_V = diag * V^-1
1015	call diag_mat_right_divide_3(iend,iend,V,diag,diag_rdivide_V)
1016
1017	! Compute V * diag_rdivide_V
1018	do j1 = 1,3
1019	do j2 = 1,3
1020	R(1:iend,j2,j1) = V(1:iend,j2,1)*diag_rdivide_V(1:iend,1,j1) &
1021	& + V(1:iend,j2,2)*diag_rdivide_V(1:iend,2,j1) &
1022	& + V(1:iend,j2,3)*diag_rdivide_V(1:iend,3,j1)
1023	end do
1024	end do
1025
1026	if (lhook) call dr_hook('radiation_matrix:fast_expm_exchange_3',1,hook_handle)
1027
1028	end subroutine fast_expm_exchange_3
1029
1030	! generic :: fast_expm_exchange => fast_expm_exchange_2, fast_expm_exchange_3
1031
1032
1033	end module radiation_matrix

Note: See TracBrowser for help on using the repository browser.

Download in other formats:

Original Format