complete cg with dirichlet flag, add example, fix di filter

Author: jalvesz

example/linalg/CMakeLists.txt

@@ -41,6 +41,7 @@ ADD_EXAMPLE(get_norm)
 ADD_EXAMPLE(solve1)
 ADD_EXAMPLE(solve2)
 ADD_EXAMPLE(solve3)
+ADD_EXAMPLE(solve_cg)
 ADD_EXAMPLE(solve_pccg)
 ADD_EXAMPLE(sparse_from_ijv)
 ADD_EXAMPLE(sparse_data_accessors)

example/linalg/example_solve_cg.f90

@@ -0,0 +1,17 @@
+program example_solve_pccg
+    use stdlib_kinds, only: dp
+    use stdlib_linalg_iterative_solvers, only: solve_cg
+
+    real(dp) :: matrix(2,2)
+    real(dp) :: x(2), load(2)
+
+    matrix = reshape( [4, 1,&
+                       1, 3] , [2,2])
+
+    x    = dble( [2,1] ) !> initial guess
+    load = dble( [1,2] ) !> load vector
+
+    call solve_cg(matrix, load, x, restart=.false.)
+    print *, x !> solution: [0.0909, 0.6364]
+
+end program

example/linalg/example_solve_pccg.f90

@@ -24,9 +24,9 @@ program example_solve_pccg
     dirichlet([1,5]) = .true._1
 
     call solve_pccg(laplacian, load, x, tol=1.d-6, di=dirichlet)
-    print *, x
+    print *, x !> solution: [0.0, 2.5, 5.0, 2.5, 0.0]
     x = 0._dp
 
     call solve_pccg(laplacian_csr, load, x, tol=1.d-6, di=dirichlet)
-    print *, x
+    print *, x !> solution: [0.0, 2.5, 5.0, 2.5, 0.0]
 end program

src/stdlib_linalg_iterative_solvers.fypp

@@ -42,12 +42,14 @@ module stdlib_linalg_iterative_solvers
 
     interface solve_cg_generic
         #:for k, t, s in R_KINDS_TYPES
-        module subroutine solve_cg_generic_${s}$(A,b,x,tol,maxiter,workspace)
+        module subroutine solve_cg_generic_${s}$(A,b,x,di,tol,maxiter,restart,workspace)
             class(linop_${s}$), intent(in) :: A
             ${t}$, intent(in) :: b(:)
             ${t}$, intent(inout) :: x(:)
             ${t}$, intent(in) :: tol
+            logical(1), intent(in) :: di(:)
             integer, intent(in) :: maxiter
+            logical, intent(in) :: restart
             type(solver_workspace_${s}$), intent(inout) :: workspace
         end subroutine
         #:endfor
@@ -57,7 +59,7 @@ module stdlib_linalg_iterative_solvers
     interface solve_cg
         #:for matrix in MATRIX_TYPES
         #:for k, t, s in R_KINDS_TYPES
-        module subroutine solve_cg_${matrix}$_${s}$(A,b,x,tol,maxiter,workspace)
+        module subroutine solve_cg_${matrix}$_${s}$(A,b,x,di,tol,maxiter,restart,workspace)
             #:if matrix == "dense"
             ${t}$, intent(in) :: A(:,:)
             #:else 
@@ -66,7 +68,9 @@ module stdlib_linalg_iterative_solvers
             ${t}$, intent(in) :: b(:)
             ${t}$, intent(inout) :: x(:)
             ${t}$, intent(in), optional :: tol
+            logical(1), intent(in), optional, target  :: di(:)
             integer, intent(in), optional :: maxiter
+            logical, intent(in), optional :: restart
             type(solver_workspace_${s}$), optional, intent(inout), target :: workspace
         end subroutine
         #:endfor

src/stdlib_linalg_iterative_solvers_cg.fypp

@@ -15,43 +15,63 @@ submodule(stdlib_linalg_iterative_solvers) stdlib_linalg_iterative_cg
 contains
 
     #:for k, t, s in R_KINDS_TYPES
-    module subroutine solve_cg_generic_${s}$(A,b,x,tol,maxiter,workspace)
+    module subroutine solve_cg_generic_${s}$(A,b,x,di,tol,maxiter,restart,workspace)
         class(linop_${s}$), intent(in) :: A
         ${t}$, intent(in) :: b(:), tol
         ${t}$, intent(inout) :: x(:)
+        logical(1), intent(in) :: di(:)
         integer, intent(in) :: maxiter
+        logical, intent(in) :: restart 
         type(solver_workspace_${s}$), intent(inout) :: workspace
         !-------------------------
         integer :: iter
-        ${t}$ :: rtr, rtrold, alpha, beta, norm0_sq
+        ${t}$ :: norm_sq, norm_sq_old, norm_sq0, residual0, residual
+        ${t}$ :: alpha, beta, tolsq
         !-------------------------
-        associate(  p  => workspace%tmp(:,1), &
-                    r  => workspace%tmp(:,2), &
+        associate(  P  => workspace%tmp(:,1), &
+                    R  => workspace%tmp(:,2), &
                     Ap => workspace%tmp(:,3))
-            x = zero_${s}$
-            rtr = A%inner_product(r, r)
-            norm0_sq = A%inner_product(b, b)
-            p = b
+            
+            norm_sq0 = A%inner_product(B, B)
+            residual0 = sqrt(norm_sq0)
+
+            if(restart) X = zero_${s}$
+            X = merge( B, X, di ) !> copy dirichlet load conditions encoded in B and indicated by di
+
+            call A%matvec(X, R)
+            R = merge( zero_${s}$, B - R , di ) !> R = B - A*X
+            norm_sq = A%inner_product(R, R)
+   
+            P = R
+            
+            tolsq = tol*tol
             beta = zero_${s}$
-            iter = 1
-            do while( rtr > tol**2 * norm0_sq .and. iter < maxiter)
-                p = r + beta * p
-                call A%matvec(p,Ap) 
-                alpha = rtr / A%inner_product(p, Ap)
-                x = x + alpha * p
-                r = r - alpha * Ap
-                rtrold = rtr
-                rtr = A%inner_product(r, r)
-                beta = rtr / rtrold
+            iter = 0
+            do while( norm_sq > tolsq * norm_sq0 .and. iter < maxiter)
+                call A%matvec(P,Ap) 
+                Ap = merge( zero_${s}$, Ap, di )
+                
+                alpha = norm_sq / A%inner_product(P, Ap)
+                
+                X = X + alpha * P
+                R = R - alpha * Ap
+
+                norm_sq_old = norm_sq
+                norm_sq = A%inner_product(R, R)
+                beta = norm_sq / norm_sq_old
+                
+                P = R + beta * P
+
                 iter = iter + 1
             end do
+            residual = sqrt(norm_sq)
         end associate
     end subroutine
     #:endfor
 
     #:for matrix in MATRIX_TYPES
     #:for k, t, s in R_KINDS_TYPES
-    module subroutine solve_cg_${matrix}$_${s}$(A,b,x,tol,maxiter,workspace)
+    module subroutine solve_cg_${matrix}$_${s}$(A,b,x,di,tol,maxiter,restart,workspace)
         #:if matrix == "dense"
         ${t}$, intent(in) :: A(:,:)
         #:else 
@@ -60,30 +80,57 @@ contains
         ${t}$, intent(in) :: b(:)
         ${t}$, intent(inout) :: x(:)
         ${t}$, intent(in), optional :: tol
+        logical(1), intent(in), optional, target  :: di(:)
         integer, intent(in), optional :: maxiter
+        logical, intent(in), optional  :: restart
         type(solver_workspace_${s}$), optional, intent(inout), target :: workspace
         !-------------------------
         type(linop_${s}$) :: op
         type(solver_workspace_${s}$), pointer :: workspace_
         integer :: n, maxiter_
         ${t}$ :: tol_
+        logical :: restart_
+        logical(1), pointer :: di_(:)
         !-------------------------
         n = size(b)
+        maxiter_ = n;       if(present(maxiter)) maxiter_ = maxiter
+        restart_ = .true.;  if(present(restart)) restart_ = restart
+        tol_ = 1.e-4_${s}$; if(present(tol)) tol_ = tol
+
+        !-------------------------
+        ! internal memory setup
         op%matvec => default_matvec
         op%inner_product => default_dot
-
-        maxiter_ = n
-        if(present(maxiter)) maxiter_ = maxiter
-        tol_ = 1.e-4_${s}$
-        if(present(tol)) tol_ = tol
+        if(present(di))then
+            di_ => di
+        else 
+            allocate(di_(n),source=.false._1)
+        end if
 
         if(present(workspace)) then
             if(.not.allocated(workspace_%tmp)) allocate( workspace_%tmp(n,3) )
             workspace_ => workspace
         else
-            allocate( workspace_%tmp(n,3) )
+            allocate( workspace_ )
+            allocate( workspace_%tmp(n,3), source = zero_${s}$ )
+        end if
+        !-------------------------
+        ! main call to the solver
+        call solve_cg_generic(op,b,x,di_,tol_,maxiter_,restart_,workspace_)
+
+        !-------------------------
+        ! internal memory cleanup
+        if(present(di))then
+            di_ => null()
+        else 
+            deallocate(di_)
+        end if
+        if(present(workspace)) then
+            workspace_ => null()
+        else
+            deallocate( workspace_%tmp )
+            deallocate( workspace_ )
         end if
-        call solve_cg_generic(op,b,x,tol,maxiter_,workspace_)
         contains
 
         subroutine default_matvec(x,y)

src/stdlib_linalg_iterative_solvers_pccg.fypp

@@ -25,7 +25,7 @@ contains
         type(solver_workspace_${s}$), intent(inout) :: workspace
         !-------------------------
         integer :: iter
-        ${t}$ :: norm_sq, norm_sq0, norm_sq_old, residual0
+        ${t}$ :: norm_sq, norm_sq0, norm_sq_old, residual0, residual
         ${t}$ :: zr1, zr2, zv2, alpha, beta, tolsq
         !-------------------------
         associate(  R => workspace%tmp(:,1), &
@@ -37,23 +37,22 @@ contains
         residual0 = sqrt(norm_sq0)
         if ( norm_sq0 > zero_${s}$ ) then
             if(restart) X = zero_${s}$
-            where( di ) X = B
+            X = merge( B, X, di ) !> copy dirichlet load conditions encoded in B and indicated by di
 
             call A%matvec(X, R)
-            R = B - R
-            where( di ) R = zero_${s}$
+            R = merge( zero_${s}$, B - R , di ) !> R = B - A*X
 
-            call A%preconditionner(R,P)
-            where( di ) P = zero_${s}$
+            call A%preconditionner(R,P) !> P = M^{-1}*R
+            P = merge( zero_${s}$, P, di )
 
             tolsq = tol*tol
             iter = 0
             zr1 = zero_${s}$
             zr2 = one_${s}$
             do while ( (iter < maxiter) .AND. (norm_sq > tolsq * norm_sq0) )
 
-                call A%preconditionner(R,S)
-                where ( di ) S = zero_${s}$
+                call A%preconditionner(R,S) !> S = M^{-1}*R
+                S = merge( zero_${s}$, S, di )
                 zr2 = A%inner_product( R, S )
 
                 if (iter>0) then
@@ -62,7 +61,7 @@ contains
                 end if
 
                 call A%matvec(P, Q)
-                where( di ) Q = zero_${s}$
+                Q = merge( zero_${s}$, Q, di )
                 zv2 = A%inner_product( P, Q )
 
                 alpha = zr2 / zv2
@@ -75,6 +74,7 @@ contains
                 iter = iter + 1
             end do
         end if
+        residual = sqrt(norm_sq)
         end associate
     end subroutine
     #:endfor
@@ -103,18 +103,15 @@ contains
         logical(1), pointer :: di_(:)
         !-------------------------
         n = size(b)
+        maxiter_ = n;       if(present(maxiter)) maxiter_ = maxiter
+        restart_ = .true.;  if(present(restart)) restart_ = restart
+        tol_ = 1.e-4_${s}$; if(present(tol)) tol_ = tol
 
+        !-------------------------
+        ! internal memory setup
         op%matvec => default_matvec
         op%inner_product => default_dot
         op%preconditionner => default_preconditionner
-
-        maxiter_ = n
-        if(present(maxiter)) maxiter_ = maxiter
-        restart_ = .true. 
-        if(present(restart)) restart_ = restart
-        tol_ = 1.e-4_${s}$
-        if(present(tol)) tol_ = tol
-
         if(present(di))then
             di_ => di
         else 
@@ -126,12 +123,25 @@ contains
             workspace_ => workspace
         else
             allocate( workspace_ )
-            allocate( workspace_%tmp(n,4) )
+            allocate( workspace_%tmp(n,4) , source = zero_${s}$ )
         end if
-        
-        call solve_pccg_generic(op,b,x,di_,tol,maxiter_,restart_,workspace_)
+        !-------------------------
+        ! main call to the solver
+        call solve_pccg_generic(op,b,x,di_,tol_,maxiter_,restart_,workspace_)
 
-        if(.not.present(di)) deallocate(di_)
+        !-------------------------
+        ! internal memory cleanup
+        if(present(di))then
+            di_ => null()
+        else 
+            deallocate(di_)
+        end if
+        if(present(workspace)) then
+            workspace_ => null()
+        else
+            deallocate( workspace_%tmp )
+            deallocate( workspace_ )
+        end if
         contains
 
         subroutine default_matvec(x,y)