refactor to remove hard dependency on dirichlet BCs

Author: jalvesz

example/linalg/example_solve_custom.f90

@@ -56,7 +56,7 @@ subroutine solve_pccg_custom(A,b,x,di,tol,maxiter,restart,workspace)
         where(abs(diagonal)>epsilon(0.d0)) diagonal = 1._dp/diagonal
         !-------------------------
         ! main call to the solver
-        call solve_pccg_generic(op,M,b,x,di_,tol_,maxiter_,restart_,workspace_)
+        call solve_pccg_generic(op,M,b,x,tol_,maxiter_,workspace_)
 
         !-------------------------
         ! internal memory cleanup
@@ -70,15 +70,20 @@ subroutine solve_pccg_custom(A,b,x,di,tol,maxiter,restart,workspace)
         workspace_ => null()
         contains
 
-        subroutine my_apply(x,y)
+        subroutine my_apply(x,y,alpha,beta)
             real(dp), intent(in)  :: x(:)
             real(dp), intent(inout) :: y(:)
-            call spmv( A , x, y )
+            real(dp), intent(in) :: alpha
+            real(dp), intent(in) :: beta
+            call spmv( A , x, y , alpha, beta )
+            y = merge( 0._dp, y, di_ )
         end subroutine
-        subroutine my_jacobi_preconditionner(x,y)
+        subroutine my_jacobi_preconditionner(x,y,alpha,beta)
             real(dp), intent(in)  :: x(:)
             real(dp), intent(inout) :: y(:)
-            y = diagonal * x
+            real(dp), intent(in) :: alpha
+            real(dp), intent(in) :: beta
+            y = merge( 0._dp, diagonal * x , di_ )
         end subroutine
         pure real(dp) function my_dot(x,y) result(r)
             real(dp), intent(in) :: x(:)

src/stdlib_linalg_iterative_solvers.fypp

@@ -36,10 +36,12 @@ module stdlib_linalg_iterative_solvers
 
     abstract interface
         #:for k, t, s in R_KINDS_TYPES
-        subroutine vector_sub_${s}$(x,y)
+        subroutine vector_sub_${s}$(x,y,alpha,beta)
             import :: ${k}$
             ${t}$, intent(in)  :: x(:)
             ${t}$, intent(inout) :: y(:)
+            ${t}$, intent(in) :: alpha
+            ${t}$, intent(in) :: beta
         end subroutine
         pure ${t}$ function reduction_sub_${s}$(x,y) result(r)
             import :: ${k}$
@@ -57,14 +59,12 @@ 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,di,tol,maxiter,restart,workspace)
+        module subroutine solve_cg_generic_${s}$(A,b,x,tol,maxiter,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
@@ -95,15 +95,13 @@ module stdlib_linalg_iterative_solvers
 
     interface solve_pccg_generic
         #:for k, t, s in R_KINDS_TYPES
-        module subroutine solve_pccg_generic_${s}$(A,M,b,x,di,tol,maxiter,restart,workspace)
+        module subroutine solve_pccg_generic_${s}$(A,M,b,x,tol,maxiter,workspace)
             class(linop_${s}$), intent(in) :: A
             class(linop_${s}$), intent(in) :: M !> preconditionner
             ${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
@@ -124,7 +122,7 @@ module stdlib_linalg_iterative_solvers
             ${t}$, intent(in), optional :: tol
             logical(1), intent(in), optional, target  :: di(:)
             integer, intent(in), optional  :: maxiter
-            logical, intent(in), optional  :: restart
+            logical, intent(in), optional :: restart
             class(linop_${s}$), optional , intent(in), target :: M !> preconditionner
             type(solver_workspace_${s}$), optional, intent(inout), target :: workspace
         end subroutine

src/stdlib_linalg_iterative_solvers_cg.fypp

@@ -14,13 +14,11 @@ 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,di,tol,maxiter,restart,workspace)
+    module subroutine solve_cg_generic_${s}$(A,b,x,tol,maxiter,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
@@ -35,11 +33,8 @@ contains
             norm_sq0 = A%inner_product(B, B)
             if(associated(workspace%callback)) call workspace%callback(x, norm_sq0, iter)
 
-            if(restart) X = zero_${s}$
-            X = merge( B, X, di ) !> copy dirichlet load conditions encoded in B and indicated by di
-
-            call A%apply(X, R)
-            R = merge( zero_${s}$, B - R , di ) !> R = B - A*X
+            R = B
+            call A%apply(X, R, alpha= -one_${s}$, beta=one_${s}$) !> R = B - A*X
             norm_sq = A%inner_product(R, R)
 
             P = R
@@ -48,8 +43,7 @@ contains
             beta = zero_${s}$
             if(associated(workspace%callback)) call workspace%callback(x, norm_sq, iter)
             do while( norm_sq > tolsq * norm_sq0 .and. iter < maxiter)
-                call A%apply(P,Ap) 
-                Ap = merge( zero_${s}$, Ap, di )
+                call A%apply(P,Ap, alpha= one_${s}$, beta=zero_${s}$) !> Ap = A*P
 
                 alpha = norm_sq / A%inner_product(P, Ap)
 
@@ -116,7 +110,9 @@ contains
         if(.not.allocated(workspace_%tmp)) allocate( workspace_%tmp(n,size_wksp_cg), source = zero_${s}$ )
         !-------------------------
         ! main call to the solver
-        call solve_cg_generic(op,b,x,di_,tol_,maxiter_,restart_,workspace_)
+        if(restart_) x = zero_${s}$
+        x = merge( b, x, di_ ) !> copy dirichlet load conditions encoded in B and indicated by di
+        call solve_cg_generic(op,b,x,tol_,maxiter_,workspace_)
 
         !-------------------------
         ! internal memory cleanup
@@ -130,14 +126,17 @@ contains
         workspace_ => null()
         contains
 
-        subroutine default_matvec(x,y)
+        subroutine default_matvec(x,y,alpha,beta)
             ${t}$, intent(in)  :: x(:)
             ${t}$, intent(inout) :: y(:)
+            ${t}$, intent(in) :: alpha
+            ${t}$, intent(in) :: beta
             #:if matrix == "dense"
-            y = matmul(A,x)
+            y = alpha * matmul(A,x) + beta * y
             #:else 
-            call spmv( A , x, y )
+            call spmv( A , x, y , alpha, beta )
             #:endif
+            y = merge( zero_${s}$, y, di_ )
         end subroutine
     end subroutine
 

src/stdlib_linalg_iterative_solvers_pccg.fypp

@@ -14,14 +14,12 @@ submodule(stdlib_linalg_iterative_solvers) stdlib_linalg_iterative_pccg
 contains
 
     #:for k, t, s in R_KINDS_TYPES
-    module subroutine solve_pccg_generic_${s}$(A,M,b,x,di,tol,maxiter,restart,workspace)
+    module subroutine solve_pccg_generic_${s}$(A,M,b,x,tol,maxiter,workspace)
         class(linop_${s}$), intent(in) :: A
         class(linop_${s}$), intent(in) :: M !> preconditionner
         ${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
@@ -38,40 +36,37 @@ contains
         if(associated(workspace%callback)) call workspace%callback(x, norm_sq, iter)
 
         if ( norm_sq0 > zero_${s}$ ) then
-            if(restart) X = zero_${s}$
-            X = merge( B, X, di ) !> copy dirichlet load conditions encoded in B and indicated by di
 
-            call A%apply(X, R)
-            R = merge( zero_${s}$, B - R , di ) !> R = B - A*X
+            R = B
+            call A%apply(X, R, alpha= -one_${s}$, beta=one_${s}$) !> R = B - A*X
 
-            call M%apply(R,P) !> P = M^{-1}*R
-            P = merge( zero_${s}$, P, di )
+            call M%apply(R,P, alpha= one_${s}$, beta=zero_${s}$) !> P = M^{-1}*R
 
             tolsq = tol*tol
 
             zr1 = zero_${s}$
             zr2 = one_${s}$
             do while ( (iter < maxiter) .AND. (norm_sq > tolsq * norm_sq0) )
 
-                call M%apply(R,S) !> S = M^{-1}*R
-                S = merge( zero_${s}$, S, di )
+                call M%apply(R,S, alpha= one_${s}$, beta=zero_${s}$) !> S = M^{-1}*R
                 zr2 = A%inner_product( R, S )
 
                 if (iter>0) then
                     beta = zr2 / zr1
                     P = S + beta * P
                 end if
 
-                call A%apply(P, Q)
-                Q = merge( zero_${s}$, Q, di )
+                call A%apply(P, Q, alpha= one_${s}$, beta=zero_${s}$) !> Q = A*P
                 zv2 = A%inner_product( P, Q )
 
                 alpha = zr2 / zv2
 
                 X = X + alpha * P
                 R = R - alpha * Q
+
                 norm_sq = A%inner_product( R, R )
                 norm_sq_old = norm_sq
+
                 zr1 = zr2
                 iter = iter + 1
                 if(associated(workspace%callback)) call workspace%callback(x, norm_sq, iter)
@@ -134,7 +129,9 @@ contains
         if(.not.allocated(workspace_%tmp)) allocate( workspace_%tmp(n,size_wksp_pccg) , source = zero_${s}$ )
         !-------------------------
         ! main call to the solver
-        call solve_pccg_generic(op,M_,b,x,di_,tol_,maxiter_,restart_,workspace_)
+        if(restart_) x = zero_${s}$
+        x = merge( b, x, di_ ) !> copy dirichlet load conditions encoded in B and indicated by di
+        call solve_pccg_generic(op,M_,b,x,tol_,maxiter_,workspace_)
 
         !-------------------------
         ! internal memory cleanup
@@ -149,19 +146,24 @@ contains
         workspace_ => null()
         contains
 
-        subroutine default_matvec(x,y)
+        subroutine default_matvec(x,y,alpha,beta)
             ${t}$, intent(in)  :: x(:)
             ${t}$, intent(inout) :: y(:)
+            ${t}$, intent(in) :: alpha
+            ${t}$, intent(in) :: beta
             #:if matrix == "dense"
-            y = matmul(A,x)
+            y = alpha * matmul(A,x) + beta * y
             #:else 
-            call spmv( A , x, y )
+            call spmv( A , x, y , alpha, beta )
             #:endif
+            y = merge( zero_${s}$, y, di_ )
         end subroutine
-        subroutine default_preconditionner(x,y)
+        subroutine default_preconditionner(x,y,alpha,beta)
             ${t}$, intent(in)  :: x(:)
             ${t}$, intent(inout) :: y(:)
-            y = x
+            ${t}$, intent(in) :: alpha
+            ${t}$, intent(in) :: beta
+            y = merge( zero_${s}$, x, di_ )
         end subroutine
     end subroutine