Context Navigation

radiation_matrix.F90 @ 5450

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

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

Download in other formats:

Original Format