DISPATCH
poisson_mod.f90
1 !===============================================================================
2 !> Methods for solving the Poisson equation; partially derived from routines
3 !> developed by Troels Haugbølle for the course Computational Astrophysics.
4 !===============================================================================
5 MODULE poisson_mod
6  USE mpi_mod
7  USE io_mod
8  USE trace_mod
9  USE kinds_mod
10  USE patch_mod
11  USE mesh_mod
12  USE bits_mod
13  USE scaling_mod
14  USE omp_timer_mod
15  implicit none
16  private
17  type, public:: poisson_t
18  real(8):: omega_sor ! Overrelaxation parameter
19  real(8):: rjac2 ! Jacobi spectral radius squared
20  procedure(cg), pointer:: update => null()
21  contains
22  procedure:: init
23  procedure:: cg
24  procedure:: sor
25  procedure, nopass:: laplace
26  end type
27  real, save:: floor=1e-2, tolerance=1e-4, fourpig=1.0
28  integer, save:: verbose=0
29  integer, save:: max_iter=500
30  logical, save:: precondition=.true.
31  real(8), parameter:: pi = asin(1.0) * 2.0
32  logical, save:: chebyshev = .true. ! Activate Chebyshev acceleration of SOR?
33 
34 CONTAINS
35 
36 !===============================================================================
37 !> Initialize poisson_params
38 !===============================================================================
39 SUBROUTINE init (self)
40  implicit none
41  class(poisson_t) :: self
42  !.............................................................................
43  integer :: iostat
44  !.............................................................................
45  character(len=12), save :: solver='cg'
46  logical, save :: first_time = .true.
47  namelist /poisson_params/ tolerance, floor, max_iter, fourpig, precondition, &
48  verbose, solver, chebyshev
49  !-----------------------------------------------------------------------------
50  !$omp critical (poisson_cr)
51  if (first_time) then
52  first_time = .false.
53  fourpig = scaling%grav
54  rewind(io%input)
55  read (io%input, poisson_params, iostat=iostat)
56  if (io%master) write (io%output, poisson_params)
57  end if
58  !$omp end critical (poisson_cr)
59  select case(trim(solver))
60  case ('cg')
61  self%update => cg
62  case ('sor')
63  self%update => sor
64  case default
65  if (io%master) then
66  write(*,*) 'Poisson solver '//trim(solver)//' is unknown'
67  write(*,*) 'Valid solvers are: cg, sor'
68  call mpi%abort()
69  end if
70  end select
71 END SUBROUTINE init
72 
73 !===============================================================================
74 !> Conjugate gradient Poisson solver (Ghysels & Vanroose 2014, Algorithm 1)
75 !> http://www.sciencedirect.com/science/article/pii/S0167819113000719
76 !> With the "incomplete Poisson" preconditioner from:
77 !> http://www.vis.uni-stuttgart.de/~weiskopf/publications/pdp10.pdf
78 !> or http://ieeexplore.ieee.org/document/5452414/
79 !> See also the CG method in RAMSES: `ramses/poisson/phi_fine_cg.f90`.
80 !===============================================================================
81 SUBROUTINE cg (self, patch, phi1, d, d0, detailed_timer)
82  class(poisson_t) :: self
83  class(patch_t) :: patch
84  real(kind=KindScalarVar), dimension(:,:,:), pointer:: phi1
85  real(kind=KindScalarVar), dimension(:,:,:), pointer, intent(in):: d
86  real(8) :: d0
87  logical :: detailed_timer
88  !.............................................................................
89  integer :: niter
90  real :: error, error0, usppt
91  real(8) :: pap, ru, ru2, alpha, beta, ds(3), wc
92  integer :: iter, m(3), ione, jone, kone
93  integer, save :: itimer=0
94  real(8), dimension(:,:,:), allocatable:: source, phi, r, p, s, u
95  interface
96  subroutine laplace_interface(ds, phi, out)
97  real(8), intent(in) :: ds(3)
98  real(8), dimension(:,:,:), intent(in) :: phi
99  real(8), dimension(:,:,:), intent(out):: out
100  end subroutine laplace_interface
101  end interface
102  procedure(laplace_interface), pointer:: laplace
103  !-----------------------------------------------------------------------------
104  if (detailed_timer) call trace%begin('poisson_mod%cg', itimer=itimer)
105  ione = 1
106  jone = 1
107  kone = 1
108  if (size(phi1,1) <= 1) ione = 0
109  if (size(phi1,2) <= 1) jone = 0
110  if (size(phi1,3) <= 1) kone = 0
111  !
112  if (patch%mesh_type == mesh_types%Cartesian) then
113  laplace => laplace_cartesian
114  else
115  laplace => laplace_curvilinear
116  end if
117  !-----------------------------------------------------------------------------
118  ! Allocate double precision scratch arrays
119  !-----------------------------------------------------------------------------
120  m = shape(phi1)
121  allocate (r(m(1),m(2),m(3)), p(m(1),m(2),m(3)))
122  allocate (s(m(1),m(2),m(3)), u(m(1),m(2),m(3)))
123  allocate (source(m(1),m(2),m(3)), phi(m(1),m(2),m(3)))
124  !-----------------------------------------------------------------------------
125  ! Calculate the initial residual = source - Laplace(phi)
126  !-----------------------------------------------------------------------------
127  wc = wallclock()
128  ds = patch%ds
129  phi = phi1
130  source = fourpig*(d-d0)
131  call zero_boundaries (source)
132  p = 0d0
133  r = 0d0
134  s = 0d0
135  u = 0d0
136  call laplace (ds, phi, r)
137  r = source - r
138  error0 = patch%fmaxval(abs(r/(d+floor)))
139  if (verbose > 1) print "(i6,2x,'cg: iter = ',i4.1,3x,'error =',1p,g12.4)", &
140  patch%id, 0, error0
141  call precond (r, u)
142  ru = patch%fsum(r*u)
143  !-----------------------------------------------------------------------------
144  ! Iteratively improve the solution
145  !-----------------------------------------------------------------------------
146  p = u
147  do iter = 1,max_iter
148  call laplace (ds, p, s)
149  alpha = ru/patch%fsum(s*p)
150  phi = phi + alpha*p
151  r = r - alpha*s
152  error = patch%fmaxval(abs(r/(d+floor)))/fourpig
153  call precond (r, u)
154  ru2 = patch%fsum(r*u)
155  if (error < tolerance) exit
156  beta = ru2/ru
157  p = u + beta*p
158  ru = ru2
159  if (verbose > 1) print "(i6,2x,'cg: iter = ',i4.1,3x,'error =',1p,g12.4)", &
160  patch%id, iter, error
161  !!!write (11) phi
162  end do
163  if (iter==max_iter+1) then
164  print *, 'poisson_t%cg: NO CONVERGENCE! id, niter, error0, error =', &
165  patch%id, max_iter, error0, error
166  else if (verbose>0 .and. io%master) then
167  usppt = 1e6*(wallclock()-wc)/product(shape(phi))
168  print '(a,i7,i5,3x,2e12.3,f8.2)', &
169  'poisson_t%cg: id, niter, initial/final error, us/pt =', &
170  patch%id, iter, error0, error, usppt
171  end if
172  phi1 = phi
173 
174  deallocate (phi, source, r, p, s, u)
175  if (detailed_timer) call trace%end (itimer)
176 contains
177  !===============================================================================
178  !> Set scratch array outer boundary to zero
179  !===============================================================================
180  subroutine zero_boundaries (f)
181  real(8):: f(:,:,:)
182  class(mesh_t), pointer:: m(:)
183  !.............................................................................
184  if (patch%is_periodic()) return
185  m => patch%mesh
186  if (size(f,1) > 1) then
187  f(m(1)%lb:m(1)%lo,:,:) = 0d0
188  f(m(1)%uo:m(1)%ub,:,:) = 0d0
189  end if
190  if (size(f,2) > 1) then
191  f(:,m(2)%lb:m(2)%lo,:) = 0d0
192  f(:,m(2)%uo:m(2)%ub,:) = 0d0
193  end if
194  if (size(f,3) > 1) then
195  f(:,:,m(3)%lb:m(3)%lo) = 0d0
196  f(:,:,m(3)%uo:m(3)%ub) = 0d0
197  end if
198  end subroutine zero_boundaries
199  !===============================================================================
200  !> 2nd order Laplace operator
201  !===============================================================================
202  subroutine laplace_cartesian (ds, phi, out)
203  real(8), intent(in) :: ds(3)
204  real(8), dimension(:,:,:), intent(in) :: phi
205  real(8), dimension(:,:,:), intent(out):: out
206  !.............................................................................
207  integer :: i, j, k, im1, jm1, km1, ip1, jp1, kp1, li(3), n(3)
208  real(8) :: c(3)
209  !-----------------------------------------------------------------------------
210  c = 1d0/ds**2
211  n = patch%mesh%nc
212  where (n == 1)
213  c = 0.0
214  end where
215  li = patch%mesh%li
216  !-----------------------------------------------------------------------------
217  ! For periodic patches, wrap around
218  !-----------------------------------------------------------------------------
219  if (patch%is_periodic()) then
220  do k=patch%mesh(3)%lb,patch%mesh(3)%ub
221  kp1 = modulo(k-li(3)+1,n(3))+li(3)
222  km1 = modulo(k-li(3)-1,n(3))+li(3)
223  do j=patch%mesh(2)%lb,patch%mesh(2)%ub
224  jp1 = modulo(j-li(2)+1,n(2))+li(2)
225  jm1 = modulo(j-li(2)-1,n(2))+li(2)
226  do i=patch%mesh(1)%lb,patch%mesh(1)%lo
227  ip1 = modulo(i-li(1)+1,n(1))+li(1)
228  im1 = modulo(i-li(1)-1,n(1))+li(1)
229  out(i,j,k) = c(1)*(phi(ip1,j,k)+phi(im1,j,k)-2d0*phi(i,j,k)) &
230  + c(2)*(phi(i,jp1,k)+phi(i,jm1,k)-2d0*phi(i,j,k)) &
231  + c(3)*(phi(i,j,kp1)+phi(i,j,km1)-2d0*phi(i,j,k))
232  end do
233  do i=patch%mesh(1)%li,patch%mesh(1)%ui
234  out(i,j,k) = c(1)*(phi(i+ione,j,k)+phi(i-ione,j,k)-2d0*phi(i,j,k)) &
235  + c(2)*(phi(i ,jp1,k)+phi(i ,jm1,k)-2d0*phi(i,j,k)) &
236  + c(3)*(phi(i ,j,kp1)+phi(i ,j,km1)-2d0*phi(i,j,k))
237  end do
238  do i=patch%mesh(1)%uo,patch%mesh(1)%ub
239  ip1 = modulo(i-li(1)+1,n(1))+li(1)
240  im1 = modulo(i-li(1)-1,n(1))+li(1)
241  out(i,j,k) = c(1)*(phi(ip1,j,k)+phi(im1,j,k)-2d0*phi(i,j,k)) &
242  + c(2)*(phi(i,jp1,k)+phi(i,jm1,k)-2d0*phi(i,j,k)) &
243  + c(3)*(phi(i,j,kp1)+phi(i,j,km1)-2d0*phi(i,j,k))
244  end do
245  end do
246  end do
247  else
248  !-----------------------------------------------------------------------------
249  ! For non-periodic patches, loop from li to ui
250  !-----------------------------------------------------------------------------
251  call zero_boundaries (out)
252  do k=patch%mesh(3)%li,patch%mesh(3)%ui
253  do j=patch%mesh(2)%li,patch%mesh(2)%ui
254  do i=patch%mesh(1)%li,patch%mesh(1)%ui
255  out(i,j,k) = c(1)*(phi(i+ione,j,k)+phi(i-ione,j,k)-2d0*phi(i,j,k)) &
256  + c(2)*(phi(i,j+jone,k)+phi(i,j-jone,k)-2d0*phi(i,j,k)) &
257  + c(3)*(phi(i,j,k+kone)+phi(i,j,k-kone)-2d0*phi(i,j,k))
258  end do
259  end do
260  end do
261  end if
262  end subroutine laplace_cartesian
263  !===============================================================================
264  !> 2nd order Laplace operator for curvilinear coords.
265  !===============================================================================
266  subroutine laplace_curvilinear (ds, phi, out)
267  real(8), intent(in) :: ds(3)
268  real(8), dimension(:,:,:), intent(in) :: phi
269  real(8), dimension(:,:,:), intent(out):: out
270  !.............................................................................
271  class(mesh_t), pointer:: m1, m2
272  integer :: i, j, k
273  real(8) :: c(3), h2ci, h31ci, h32ci
274  !-----------------------------------------------------------------------------
275  c = 1d0/ds**2
276  where (patch%mesh%nc == 1)
277  c = 0.0
278  end where
279  m1 => patch%mesh(1)
280  m2 => patch%mesh(2)
281  !-----------------------------------------------------------------------------
282  ! For periodic patches, wrap around
283  ! NOTE FROM JPR: I'm not sure a curvilinear patch will ever be periodic in all
284  ! three directions, thus the abort.
285  !-----------------------------------------------------------------------------
286  if (patch%is_periodic()) then
287  call mpi%abort('NOT IMPLEMENTED ERROR: This curvilinear patch is periodic in all three directions. Really??')
288  else
289  !-----------------------------------------------------------------------------
290  ! For non-periodic patches, loop from li to ui
291  !-----------------------------------------------------------------------------
292  call zero_boundaries (out)
293  do k=patch%mesh(3)%li,patch%mesh(3)%ui
294  do j=patch%mesh(2)%li,patch%mesh(2)%ui
295  h32ci = 1.0d0 / m2%h32c(j)
296  do i=patch%mesh(1)%li,patch%mesh(1)%ui
297  h2ci = 1.0d0 / m1%h2c(i)
298  h31ci = 1.0d0 / m1%h31c(i)
299  out(i,j,k) = c(1) * h2ci * h31ci &
300  * (m1%h2f(i+ione) * m1%h31f(i+ione) * (phi(i+ione,j,k) - phi(i,j,k)) &
301  - m1%h2f(i ) * m1%h31f(i ) * (phi(i,j,k) - phi(i-ione,j,k))) &
302  + c(2) * h2ci**2 * h32ci &
303  * (m2%h32f(j+jone) * (phi(i,j+jone,k) - phi(i,j,k)) &
304  - m2%h32f(j ) * (phi(i,j,k) - phi(i,j-jone,k))) &
305  + c(3) * h31ci**2 * h32ci**2 &
306  * (phi(i,j,k+kone) + phi(i,j,k-kone) - 2.0d0 * phi(i,j,k))
307  end do
308  end do
309  end do
310  end if
311  end subroutine laplace_curvilinear
312  !===============================================================================
313  !> Preconditioner for conjugate gradient method
314  !===============================================================================
315  subroutine precond (res, out)
316  real(8), dimension(:,:,:) :: res, out
317  !.............................................................................
318  integer :: i, j, k, im1, jm1, km1, ip1, jp1, kp1, li(3), n(3), &
319  ndim
320  real(8) :: a, b(3)
321  !-----------------------------------------------------------------------------
322  if (precondition) then
323  n = patch%mesh%nc
324  ndim = sum(merge(1,0,n > 1))
325  b = 0.5d0 / ndim
326  a = 1d0 + 0.25d0 / ndim
327  where (n == 1)
328  b = 0.0d0
329  end where
330  li = patch%mesh%li
331  !---------------------------------------------------------------------------
332  ! For periodic patches, wrap around (cheap patch so no need to optimize)
333  !---------------------------------------------------------------------------
334  if (patch%is_periodic()) then
335  do k=patch%mesh(3)%lb,patch%mesh(3)%ub
336  kp1 = modulo(k-li(3)+1,n(3))+li(3)
337  km1 = modulo(k-li(3)-1,n(3))+li(3)
338  do j=patch%mesh(2)%lb,patch%mesh(2)%ub
339  jp1 = modulo(j-li(2)+1,n(2))+li(2)
340  jm1 = modulo(j-li(2)-1,n(2))+li(2)
341  do i=patch%mesh(1)%lb,patch%mesh(1)%ub
342  ip1 = modulo(i-li(1)+1,n(1))+li(1)
343  im1 = modulo(i-li(1)-1,n(1))+li(1)
344  out(i,j,k) = b(1)*(res(ip1,j,k)+res(im1,j,k)) &
345  + b(2)*(res(i,jp1,k)+res(i,jm1,k)) &
346  + b(3)*(res(i,j,kp1)+res(i,j,km1)) &
347  + a*res(i,j,k)
348  end do
349  end do
350  end do
351  !---------------------------------------------------------------------------
352  ! For non- periodic patches, loop over internal cells
353  !---------------------------------------------------------------------------
354  else
355  do k=patch%li(3),patch%ui(3)
356  do j=patch%li(2),patch%ui(2)
357  do i=patch%li(1),patch%ui(1)
358  out(i,j,k) = b(1)*(res(i+ione,j,k)+res(i-ione,j,k)) &
359  + b(2)*(res(i,j+jone,k)+res(i,j-jone,k)) &
360  + b(3)*(res(i,j,k+kone)+res(i,j,k-kone)) &
361  + a*res(i,j,k)
362  end do
363  end do
364  end do
365  call zero_boundaries (out)
366  end if
367  else
368  out = res
369  end if
370  end subroutine precond
371 END SUBROUTINE cg
372 
373 !===============================================================================
374 !> 2nd order successive over relaxation with red-black ordering
375 !===============================================================================
376 SUBROUTINE sor (self, patch, phi1, d, d0, detailed_timer)
377  class(poisson_t) :: self
378  class(patch_t) :: patch
379  real(kind=KindScalarVar), dimension(:,:,:), pointer:: phi1
380  real(kind=KindScalarVar), dimension(:,:,:), pointer, intent(in):: d
381  real(8):: d0
382  logical :: detailed_timer
383  !.............................................................................
384  integer :: i, j, k, sweep, ss, ione, jone, kone, m(3), iter, niter
385  real :: error, error0
386  real(8) :: a, b, b1, b2, b3, res, numer, denom, wc, usppt
387  integer, save :: itimer=0
388  real(8), dimension(:,:,:), allocatable:: phi, source
389 
390  !-----------------------------------------------------------------------------
391  if (detailed_timer) call trace%begin ('poisson_t%sor_iter', itimer=itimer)
392 
393  associate(n=>patch%n,ds=>patch%ds,l=>patch%li,u=>patch%ui,gn=>patch%gn)
394  allocate (phi(gn(1),gn(2),gn(3)), source(gn(1),gn(2),gn(3)))
395 
396  !-------------------------------------------------------------------
397  ! Calculate optimal guess for the SOR prefactor depending on
398  ! dimensionality, ds, and grid points n:
399  ! delta = (1/sum(1/ds^2)) sum(1/ds^2 * cos(pi/n))
400  ! omega_sor = 2 / (1 + (1 - delta^2)^0.5)
401  ! For an even dimensionality grid
402  ! delta = cos(pi / n)
403  ! which for n > 10 is well approximated by
404  ! delta ~ 1 - 1/2 (pi/n)**2
405  ! and then
406  ! omega_sor = 2 / (1 + pi/n)
407  !-------------------------------------------------------------------
408  numer = 0.0
409  denom = 0.0
410  do i=1,3
411  if (patch%n(i) > 1) then
412  numer = numer + cos(pi / patch%n(i))
413  denom = denom + 1d0
414  end if
415  end do
416  self%rjac2 = (numer / denom)**2
417  if (chebyshev) then
418  ! initial value for Chebyshev acceleration
419  self%omega_sor = 1.0
420  else
421  ! optimal value for the relaxation parameter
422  self%omega_sor = 2d0 / (1d0 + sqrt(1d0 - self%rjac2))
423  end if
424 
425  if (patch%is_periodic()) call patch%make_periodic (phi)
426  ione = 1
427  jone = 1
428  kone = 1
429  if (n(1) <= 1) ione = 0
430  if (n(2) <= 1) jone = 0
431  if (n(3) <= 1) kone = 0
432 
433  wc = wallclock()
434  phi = phi1
435  source = fourpig * (d - d0)
436  ! Calculate 1 / [(2/dx + 2/dy + 2/dz)] prefactor according to dimensionality
437  a = 0.0
438  if (n(1) > 1) a = a + 2. / ds(1)**2
439  if (n(2) > 1) a = a + 2. / ds(2)**2
440  if (n(3) > 1) a = a + 2. / ds(3)**2
441  a = 1. / a
442  b1 = a / ds(1)**2 * ione
443  b2 = a / ds(2)**2 * jone
444  b3 = a / ds(3)**2 * kone
445 
446  do iter=1,max_iter
447  error = 0.0
448  do sweep=0,1
449  do k=l(3),u(3)
450  do j=l(2),u(2)
451  ss = modulo(j + k + sweep,2)
452  do i=l(1)+ss,u(1),2
453  !---------------------------------------------------------------------
454  ! This form of the residual has no factor in front of phi, which makes
455  ! it clear how to apply the over-relaxation factor
456  !---------------------------------------------------------------------
457  res = b1 * (phi(i+ione,j,k) + phi(i-ione,j,k)) &
458  + b2 * (phi(i,j+jone,k) + phi(i,j-jone,k)) &
459  + b3 * (phi(i,j,k+kone) + phi(i,j,k-kone)) &
460  - phi(i,j,k)- a * source(i,j,k)
461  phi(i,j,k) = phi(i,j,k) + self%omega_sor * res
462  error = max(abs(res / (d(i,j,k) + floor)), error)
463  end do
464  end do
465  end do
466  if (chebyshev) then
467  if (iter == 1) then
468  self%omega_sor = 1.0 / (1.0 - self%rjac2 * 0.5)
469  else
470  self%omega_sor = 1.0 / (1.0 - self%rjac2 * 0.25 * self%omega_sor)
471  end if
472  end if
473  end do
474  !-----------------------------------------------------------------------------
475  ! We are supposed to return the residual as (Laplace(phi)-source)/d, so we
476  ! need to divide with the factor that source (FourPiGrho) was multiplied by
477  !-----------------------------------------------------------------------------
478  error = error / a
479  if (iter==1) error0 = error
480  if (verbose > 1) print *, 'poisson_t%sor: iter, error =', iter, error
481  if (error < tolerance) exit
482  !if (patch%is_periodic()) call patch%make_periodic (phi)
483  end do
484 
485  if (iter==max_iter+1) then
486  print *, 'poisson_t%sor: no convergence, niter, error =', max_iter, error
487  else if (verbose > 0 .and. io%master) then
488  usppt = 1e6 * (wallclock() - wc) / product(shape(phi))
489  print '(a,i7,i5,3x,2e12.3,f8.2)', &
490  'poisson_t%sor: id, niter, initial/final error, us/pt =', &
491  patch%id, iter, error0, error, usppt
492  end if
493  niter = iter
494  phi1 = phi
495  deallocate (phi)
496 
497  end associate
498  if (detailed_timer) call trace%end (itimer)
499 END SUBROUTINE sor
500 
501 !===============================================================================
502 !> 2nd order Laplace operator
503 !===============================================================================
504 SUBROUTINE laplace (ds, phi, source, res)
505  real(8) :: ds(3)
506  real, dimension(:,:,:) :: phi, source, res
507  !.............................................................................
508  integer :: ix, iy, iz, m(3), ione, jone, kone
509  real(8) :: c(3)
510  !-----------------------------------------------------------------------------
511  ione = 1
512  jone = 1
513  kone = 1
514  if (size(phi,1) <= 1) ione = 0
515  if (size(phi,2) <= 1) jone = 0
516  if (size(phi,3) <= 1) kone = 0
517 
518  c = 1d0/ds**2
519  m = shape(phi)
520  where (m == 1)
521  c = 0.0d0
522  end where
523  do iz=2,m(3)-kone
524  do iy=2,m(2)-jone
525  do ix=2,m(1)-ione
526  res(ix,iy,iz) = c(1)*(phi(ix+ione,iy,iz)+phi(ix-ione,iy,iz)-2d0*phi(ix,iy,iz)) &
527  + c(2)*(phi(ix,iy+jone,iz)+phi(ix,iy-jone,iz)-2d0*phi(ix,iy,iz)) &
528  + c(3)*(phi(ix,iy,iz+kone)+phi(ix,iy,iz-kone)-2d0*phi(ix,iy,iz)) &
529  - source(ix,iy,iz)
530  end do
531  end do
532  end do
533 END SUBROUTINE laplace
534 
535 END MODULE poisson_mod
Support tic/toc timing, as in MATLAB, and accurate wallclock() function. The timing is generally much...
Define code units, in terms of (here) CGS units.
Definition: scaling_mod.f90:4
Template module for patches, which adds pointers to memory and mesh, and number of dimensions and var...
Definition: patch_mod.f90:6
Methods for solving the Poisson equation; partially derived from routines developed by Troels Haugbøll...
Definition: poisson_mod.f90:5
Template module for mesh.
Definition: mesh_mod.f90:4
Definition: io_mod.f90:4