A Modular Arbitrary-Order Ocean-Atmosphere Model -- Stochastic implementation
sqrt_mod.f90
Go to the documentation of this file.
1 
2 ! sqrt_mod.f90
3 !
4 !> Utility module with various routine to compute matrix square root.
5 !
6 !> @copyright
7 !> 2018 Jonathan Demaeyer.
8 !> See LICENSE.txt for license information.
9 !
10 !---------------------------------------------------------------------------!
11 !
12 !> @remark
13 !> Mainly based on the numerical recipes and from:
14 !> Edvin Deadman, Nicholas J. Higham, Rui Ralha (2013)
15 !> "Blocked Schur Algorithms for Computing the Matrix Square Root",
16 !> Lecture Notes in Computer Science, 7782. pp. 171-182.
17 !
18 !---------------------------------------------------------------------------!
19 
20 MODULE sqrt_mod
21  USE params, only:ndim
22  USE util, only: triu,diag,cdiag,floordiv
23  IMPLICIT NONE
24 
25  PRIVATE
26 
27  REAL(KIND=8), DIMENSION(:), ALLOCATABLE :: work
28  ! COMPLEX(KIND=16), DIMENSION(:), ALLOCATABLE :: zwork
29 
30  INTEGER :: lwork
31 
32  REAL(KIND=8), PARAMETER :: real_eps = 2.2204460492503131e-16
33 
34  PUBLIC :: sqrtm,init_sqrt,chol,sqrtm_svd
35 
36 CONTAINS
37 
38  SUBROUTINE init_sqrt
39  INTEGER :: AllocStat
40  lwork=10
42 
43  ! print*, lwork
44 
45  IF (ALLOCATED(work)) THEN
46  DEALLOCATE(work, stat=allocstat)
47  IF (allocstat /= 0) stop "*** Deallocation problem ! ***"
48  ENDIF
49  ALLOCATE(work(lwork), stat=allocstat)
50  IF (allocstat /= 0) stop "*** Not enough memory ! ***"
51 
52  ! ALLOCATE(zwork(lwork), STAT=AllocStat)
53  ! IF (AllocStat /= 0) STOP "*** Not enough memory ! ***"
54  END SUBROUTINE init_sqrt
55 
56  !> Routine to compute a real square-root of a matrix
57  !> @param A Matrix whose square root to evaluate.
58  !> @param sqA Square root of `A`.
59  !> @param info Information code returned by the Lapack routines.
60  !> @param info_triu Information code returned by the triangular matrix Lapack routines.
61  !> @param bs Optional blocksize specification variable.
62  SUBROUTINE sqrtm(A,sqA,info,info_triu,bs)
63  REAL(KIND=8), DIMENSION(:,:), INTENT(IN) :: A
64  REAL(KIND=8), DIMENSION(:,:), INTENT(OUT) :: sqA
65  INTEGER, INTENT(IN), OPTIONAL :: bs
66  INTEGER, INTENT(OUT) :: info,info_triu
67  REAL(KIND=8), DIMENSION(SIZE(A,1),SIZE(A,1)) :: T,Z,R
68  COMPLEX(KIND=16), DIMENSION(SIZE(A,1),SIZE(A,1)) :: Tz,Zz,Rz
69  REAL(KIND=8), DIMENSION(SIZE(A,1)) :: wr,wi
70  LOGICAL, DIMENSION(SIZE(A,1)) :: bwork
71  LOGICAL :: selectev
72  INTEGER :: n
73  INTEGER :: sdim=0
74  n=SIZE(a,1)
75  t=a
76  ! print*, n,size(work,1)
77  CALL dgees('v','n',selectev,n,t,n,sdim,wr,wi,z,n,work,lwork,bwork,info)
78  ! print*, 'Z'
79  ! CALL printmat(Z)
80  ! print*, 'T'
81  ! CALL printmat(T)
82  ! CALL DGEES('V','N',SIZE(T,1),T,SIZE(T,1),0,wr,wi,Z,SIZE(Z,1),work,lwork,info)
83  ! print*, info
84  CALL triu(t,r)
85  IF (any(t /= r)) THEN
86  ! print*, 'T'
87  ! CALL printmat(T)
88  ! print*, 'Z'
89  ! CALL printmat(Z)
90  CALL rsf2csf(t,z,tz,zz)
91  ! print*, 'Tz'
92  ! CALL printmat(dble(Tz))
93  ! print*, 'iTz'
94  ! CALL printmat(dble(aimag(Tz)))
95  ! print*, 'Zz'
96  ! CALL printmat(dble(Zz))
97  ! print*, 'iZz'
98  ! CALL printmat(dble(aimag(Zz)))
99  IF (PRESENT(bs)) THEN
100  CALL csqrtm_triu(tz,rz,info_triu,bs)
101  ELSE
102  CALL csqrtm_triu(tz,rz,info_triu)
103  END IF
104  rz=matmul(zz,matmul(rz,conjg(transpose(zz))))
105  ! print*, 'sqAz'
106  ! CALL printmat(dble(Rz))
107  ! print*, 'isqAz'
108  ! CALL printmat(dble(aimag(Rz)))
109  sqa=dble(rz)
110  ELSE
111  IF (PRESENT(bs)) THEN
112  CALL sqrtm_triu(t,r,info_triu,bs)
113  ELSE
114  CALL sqrtm_triu(t,r,info_triu)
115  END IF
116  sqa=matmul(z,matmul(r,transpose(z)))
117  ENDIF
118 
119  END SUBROUTINE sqrtm
120 
121  FUNCTION selectev(a,b)
122  REAL(KIND=8) :: a,b
123  LOGICAL selectev
124  selectev=.false.
125  ! IF (a>b) selectev=.true.
126  RETURN
127  END FUNCTION selectev
128 
129 
130  SUBROUTINE sqrtm_triu(A,sqA,info,bs)
131  REAL(KIND=8), DIMENSION(:,:), INTENT(IN) :: A
132  INTEGER, INTENT(IN), OPTIONAL :: bs
133  REAL(KIND=8), DIMENSION(:,:), INTENT(OUT) :: sqA
134  INTEGER, INTENT(OUT) :: info
135  REAL(KIND=8), DIMENSION(SIZE(A,1)) :: A_diag
136  REAL(KIND=8), DIMENSION(SIZE(A,1),SIZE(A,1)) :: R,Sm,Rii,Rjj
137  INTEGER, DIMENSION(2*SIZE(A,1),2) :: start_stop_pairs
138  REAL(KIND=8) :: s,denom,scale
139  INTEGER :: i,j,k,start,n,sstop,m
140  INTEGER :: istart,istop,jstart,jstop
141  INTEGER :: nblocks,blocksize
142  INTEGER :: bsmall,blarge,nlarge,nsmall
143 
144  blocksize=64
145  IF (PRESENT(bs)) blocksize=bs
146  n=SIZE(a,1)
147  ! print*, blocksize
148 
149  CALL diag(a,a_diag)
150  r=0.d0
151  DO i=1,n
152  r(i,i)=sqrt(a_diag(i))
153  ENDDO
154 
155 
156  nblocks=max(floordiv(n,blocksize),1)
157  bsmall=floordiv(n,nblocks)
158  nlarge=mod(n,nblocks)
159  blarge=bsmall+1
160  nsmall=nblocks-nlarge
161  IF (nsmall*bsmall + nlarge*blarge /= n) stop 'Sqrtm: Internal inconsistency'
162 
163  ! print*, nblocks,bsmall,nsmall,blarge,nlarge
164 
165  start=1
166  DO i=1,nsmall
167  start_stop_pairs(i,1)=start
168  start_stop_pairs(i,2)=start+bsmall-1
169  start=start+bsmall
170  ENDDO
171  DO i=nsmall+1,nsmall+nlarge
172  start_stop_pairs(i,1)=start
173  start_stop_pairs(i,2)=start+blarge-1
174  start=start+blarge
175  ENDDO
176 
177  ! DO i=1,SIZE(start_stop_pairs,1)
178  ! print*, i
179  ! print*, start_stop_pairs(i,1),start_stop_pairs(i,2)
180  ! END DO
181 
182  DO k=1,nsmall+nlarge
183  start=start_stop_pairs(k,1)
184  sstop=start_stop_pairs(k,2)
185  DO j=start,sstop
186  DO i=j-1,start,-1
187  s=0.d0
188  IF (j-i>1) s= dot_product(r(i,i+1:j-1),r(i+1:j-1,j))
189  denom= r(i,i)+r(j,j)
190  IF (denom==0.d0) stop 'Sqrtm: Failed to find the matrix square root'
191  r(i,j)=(a(i,j)-s)/denom
192  END DO
193  END DO
194  END DO
195 
196  ! print*, 'R'
197  ! CALL printmat(R)
198 
199  DO j=1,nblocks
200  jstart=start_stop_pairs(j,1)
201  jstop=start_stop_pairs(j,2)
202  DO i=j-1,1,-1
203  istart=start_stop_pairs(i,1)
204  istop=start_stop_pairs(i,2)
205  sm=0.d0
206  sm(istart:istop,jstart:jstop)=a(istart:istop,jstart:jstop)
207  IF (j-i>1) sm(istart:istop,jstart:jstop) = sm(istart:istop&
208  &,jstart:jstop) - matmul(r(istart:istop,istop:jstart)&
209  &,r(istop:jstart,jstart:jstop))
210  rii=0.d0
211  rii = r(istart:istop, istart:istop)
212  rjj=0.d0
213  rjj = r(jstart:jstop, jstart:jstop)
214  m=istop-istart+1
215  n=jstop-jstart+1
216  k=1
217  ! print*, m,n
218  ! print*, istart,istop
219  ! print*, jstart,jstop
220 
221  ! print*, 'Rii',Rii(istart:istop, istart:istop)
222  ! print*, 'Rjj',Rjj(jstart:jstop,jstart:jstop)
223  ! print*, 'Sm',Sm(istart:istop,jstart:jstop)
224 
225  CALL dtrsyl('N','N',k,m,n,rii(istart:istop, istart:istop),m&
226  &,rjj(jstart:jstop,jstart:jstop),n,sm(istart:istop&
227  &,jstart:jstop),m,scale,info)
228  r(istart:istop,jstart:jstop)=sm(istart:istop,jstart:jstop)*scale
229  ENDDO
230  ENDDO
231  sqa=r
232  END SUBROUTINE sqrtm_triu
233 
234  SUBROUTINE csqrtm_triu(A,sqA,info,bs)
235  COMPLEX(KIND=16), DIMENSION(:,:), INTENT(IN) :: A
236  INTEGER, INTENT(IN), OPTIONAL :: bs
237  COMPLEX(KIND=16), DIMENSION(:,:), INTENT(OUT) :: sqA
238  INTEGER, INTENT(OUT) :: info
239  COMPLEX(KIND=16), DIMENSION(SIZE(A,1)) :: A_diag
240  COMPLEX(KIND=16), DIMENSION(SIZE(A,1),SIZE(A,1)) :: R,Sm,Rii,Rjj
241  INTEGER, DIMENSION(2*SIZE(A,1),2) :: start_stop_pairs
242  COMPLEX(KIND=16) :: s,denom,scale
243  INTEGER :: i,j,k,start,n,sstop,m
244  INTEGER :: istart,istop,jstart,jstop
245  INTEGER :: nblocks,blocksize
246  INTEGER :: bsmall,blarge,nlarge,nsmall
247 
248  blocksize=64
249  IF (PRESENT(bs)) blocksize=bs
250  n=SIZE(a,1)
251  ! print*, blocksize
252 
253  CALL cdiag(a,a_diag)
254  r=0.d0
255  DO i=1,n
256  r(i,i)=sqrt(a_diag(i))
257  ENDDO
258 
259 
260  nblocks=max(floordiv(n,blocksize),1)
261  bsmall=floordiv(n,nblocks)
262  nlarge=mod(n,nblocks)
263  blarge=bsmall+1
264  nsmall=nblocks-nlarge
265  IF (nsmall*bsmall + nlarge*blarge /= n) stop 'Sqrtm: Internal inconsistency'
266 
267  ! print*, nblocks,bsmall,nsmall,blarge,nlarge
268 
269  start=1
270  DO i=1,nsmall
271  start_stop_pairs(i,1)=start
272  start_stop_pairs(i,2)=start+bsmall-1
273  start=start+bsmall
274  ENDDO
275  DO i=nsmall+1,nsmall+nlarge
276  start_stop_pairs(i,1)=start
277  start_stop_pairs(i,2)=start+blarge-1
278  start=start+blarge
279  ENDDO
280 
281  ! DO i=1,SIZE(start_stop_pairs,1)
282  ! print*, i
283  ! print*, start_stop_pairs(i,1),start_stop_pairs(i,2)
284  ! END DO
285 
286  DO k=1,nsmall+nlarge
287  start=start_stop_pairs(k,1)
288  sstop=start_stop_pairs(k,2)
289  DO j=start,sstop
290  DO i=j-1,start,-1
291  s=0.d0
292  IF (j-i>1) s= dot_product(r(i,i+1:j-1),r(i+1:j-1,j))
293  denom= r(i,i)+r(j,j)
294  IF (denom==0.d0) stop 'Sqrtm: Failed to find the matrix square root'
295  r(i,j)=(a(i,j)-s)/denom
296  END DO
297  END DO
298  END DO
299 
300  ! print*, 'R'
301  ! CALL printmat(R)
302 
303  DO j=1,nblocks
304  jstart=start_stop_pairs(j,1)
305  jstop=start_stop_pairs(j,2)
306  DO i=j-1,1,-1
307  istart=start_stop_pairs(i,1)
308  istop=start_stop_pairs(i,2)
309  sm=0.d0
310  sm(istart:istop,jstart:jstop)=a(istart:istop,jstart:jstop)
311  IF (j-i>1) sm(istart:istop,jstart:jstop) = sm(istart:istop&
312  &,jstart:jstop) - matmul(r(istart:istop,istop:jstart)&
313  &,r(istop:jstart,jstart:jstop))
314  rii=0.d0
315  rii = r(istart:istop, istart:istop)
316  rjj=0.d0
317  rjj = r(jstart:jstop, jstart:jstop)
318  m=istop-istart+1
319  n=jstop-jstart+1
320  k=1
321  ! print*, m,n
322  ! print*, istart,istop
323  ! print*, jstart,jstop
324 
325  ! print*, 'Rii',Rii(istart:istop, istart:istop)
326  ! print*, 'Rjj',Rjj(jstart:jstop,jstart:jstop)
327  ! print*, 'Sm',Sm(istart:istop,jstart:jstop)
328 
329  CALL ztrsyl('N','N',k,m,n,rii(istart:istop, istart:istop),m&
330  &,rjj(jstart:jstop,jstart:jstop),n,sm(istart:istop&
331  &,jstart:jstop),m,scale,info)
332  r(istart:istop,jstart:jstop)=sm(istart:istop,jstart:jstop)*scale
333  ENDDO
334  ENDDO
335  sqa=r
336  END SUBROUTINE csqrtm_triu
337 
338  SUBROUTINE rsf2csf(T,Z,Tz,Zz)
339  REAL(KIND=8), DIMENSION(:,:), INTENT(IN) :: T,Z
340  COMPLEX(KIND=16), DIMENSION(:,:), INTENT(OUT) :: Tz,Zz
341  INTEGER, PARAMETER :: nb=2
342  COMPLEX(KIND=16), DIMENSION(nb) :: w
343  !COMPLEX(KIND=16), DIMENSION(nb,nb) :: vl,vr
344  COMPLEX(KIND=16) :: r,c,s,mu
345  COMPLEX(KIND=16), DIMENSION(nb,nb) :: G,Gc
346  INTEGER :: N,m!,info
347  !REAL(KIND=8), DIMENSION(2*nb) :: rwork
348  !REAL(KIND=8), DIMENSION(2*nb) :: ztwork
349 
350  ! print*, lwork
351  tz=cmplx(t,kind=16)
352  zz=cmplx(z,kind=16)
353  n=SIZE(t,1)
354  DO m=n,2,-1
355  IF (abs(tz(m,m-1)) > real_eps*(abs(tz(m-1,m-1)) + abs(tz(m,m)))) THEN
356  g=tz(m-1:m,m-1:m)
357  ! CALL printmat(dble(G))
358  ! CALL zgeev('N','N',nb,G,nb,w,vl,nb,vr,nb,ztwork,2*nb,rwork,info)
359  ! CALL cprintmat(G)
360  ! print*, m,w,info
361  s=g(1,1)+g(2,2)
362  c=g(1,1)*g(2,2)-g(1,2)*g(2,1)
363  w(1)=s/2+sqrt(s**2/4-c)
364  mu=w(1)-tz(m,m)
365  r=sqrt(mu*conjg(mu)+tz(m,m-1)*conjg(tz(m,m-1)))
366  c=mu/r
367  s=tz(m,m-1)/r
368  g(1,1)=conjg(c)
369  g(1,2)=s
370  g(2,1)=-s
371  g(2,2)=c
372  gc=conjg(transpose(g))
373  tz(m-1:m,m-1:n)=matmul(g,tz(m-1:m,m-1:n))
374  tz(1:m,m-1:m)=matmul(tz(1:m,m-1:m),gc)
375  zz(:,m-1:m)=matmul(zz(:,m-1:m),gc)
376  END IF
377  tz(m,m-1)=cmplx(0.d0,kind=16)
378  END DO
379  END SUBROUTINE rsf2csf
380 
381  !> Routine to perform a Cholesky decomposition
382  !> @param A Matrix whose decomposition is evaluated.
383  !> @param sqA Cholesky decomposition of `A`.
384  !> @param info Information code returned by the Lapack routines.
385  SUBROUTINE chol(A,sqA,info)
386  REAL(KIND=8), DIMENSION(:,:), INTENT(IN) :: A
387  REAL(KIND=8), DIMENSION(:,:), INTENT(OUT) :: sqA
388  INTEGER, INTENT(OUT) :: info
389 
390  sqa=a
391  CALL dpotrf('L',SIZE(sqa,1),sqa,SIZE(sqa,1),info)
392  END SUBROUTINE chol
393 
394  !> Routine to compute a real square-root of a matrix via a SVD decomposition
395  !> @param A Matrix whose square root to evaluate.
396  !> @param sqA Square root of `A`.
397  !> @param info Information code returned by the Lapack routines.
398  !> @param info_triu Not used (present for compatibility).
399  !> @param bs Not used (present for compatibility).
400  SUBROUTINE sqrtm_svd(A,sqA,info,info_triu,bs)
401  REAL(KIND=8), DIMENSION(:,:), INTENT(IN) :: A
402  REAL(KIND=8), DIMENSION(:,:), INTENT(OUT) :: sqA
403  INTEGER, INTENT(IN), OPTIONAL :: bs
404  INTEGER, INTENT(OUT) :: info,info_triu
405  REAL(KIND=8), DIMENSION(SIZE(A,1)) :: S
406  REAL(KIND=8), DIMENSION(SIZE(A,1),SIZE(A,1)) :: Sq,U,VT
407  INTEGER :: i,n
408 
409  sqa=a
410  n=SIZE(sqa,1)
411  CALL dgesvd('A','A',n,n,sqa,n,s,u,n,vt,n,work,lwork,info)
412  sq=0.d0
413  DO i=1,n
414  sq(i,i)=sqrt(s(i))
415  ENDDO
416  sqa=matmul(u,matmul(sq,vt))
417  END SUBROUTINE sqrtm_svd
418 
419 END MODULE sqrt_mod
420 
421 
422 
423 
424 
integer ndim
Number of variables (dimension of the model)
Definition: params.f90:85
subroutine, public chol(A, sqA, info)
Routine to perform a Cholesky decomposition.
Definition: sqrt_mod.f90:386
Utility module.
Definition: util.f90:12
Statistics accumulators.
Definition: stat.f90:14
integer function, public floordiv(i, j)
Definition: util.f90:280
subroutine, public cdiag(A, d)
Definition: util.f90:269
subroutine, public sqrtm(A, sqA, info, info_triu, bs)
Routine to compute a real square-root of a matrix.
Definition: sqrt_mod.f90:63
real(kind=8), dimension(:), allocatable work
Definition: sqrt_mod.f90:27
subroutine, public diag(A, d)
Definition: util.f90:259
integer lwork
Definition: sqrt_mod.f90:30
subroutine, public triu(A, T)
Definition: util.f90:247
subroutine sqrtm_triu(A, sqA, info, bs)
Definition: sqrt_mod.f90:131
subroutine rsf2csf(T, Z, Tz, Zz)
Definition: sqrt_mod.f90:339
subroutine, public init_sqrt
Definition: sqrt_mod.f90:39
logical function selectev(a, b)
Definition: sqrt_mod.f90:122
The model parameters module.
Definition: params.f90:18
real(kind=8), parameter real_eps
Definition: sqrt_mod.f90:32
Utility module with various routine to compute matrix square root.
Definition: sqrt_mod.f90:20
subroutine csqrtm_triu(A, sqA, info, bs)
Definition: sqrt_mod.f90:235
subroutine, public sqrtm_svd(A, sqA, info, info_triu, bs)
Routine to compute a real square-root of a matrix via a SVD decomposition.
Definition: sqrt_mod.f90:401