add example for custom solver extending the generic interface

Author: jalvesz

example/linalg/CMakeLists.txt

@@ -43,6 +43,7 @@ ADD_EXAMPLE(solve2)
 ADD_EXAMPLE(solve3)
 ADD_EXAMPLE(solve_cg)
 ADD_EXAMPLE(solve_pccg)
+ADD_EXAMPLE(solve_custom)
 ADD_EXAMPLE(sparse_from_ijv)
 ADD_EXAMPLE(sparse_data_accessors)
 ADD_EXAMPLE(sparse_spmv)

example/linalg/example_solve_custom.f90

@@ -0,0 +1,116 @@
+module custom_solver
+    use stdlib_kinds, only: dp
+    use stdlib_sparse
+    use stdlib_linalg_iterative_solvers, only: linop_dp, solver_workspace_dp, solve_pccg_generic
+    implicit none
+contains
+    subroutine solve_pccg_custom(A,b,x,di,tol,maxiter,restart,workspace)
+        type(CSR_dp_type), intent(in) :: A
+        real(dp), intent(in) :: b(:)
+        real(dp), intent(inout) :: x(:)
+        real(dp), intent(in), optional :: tol
+        logical(1), intent(in), optional, target  :: di(:)
+        integer, intent(in), optional  :: maxiter
+        logical, intent(in), optional  :: restart
+        type(solver_workspace_dp), optional, intent(inout), target :: workspace
+        !-------------------------
+        type(linop_dp) :: op
+        type(solver_workspace_dp), pointer :: workspace_
+        integer :: n, maxiter_
+        real(dp) :: tol_
+        logical :: restart_
+        logical(1), pointer :: di_(:)
+        real(dp), allocatable :: diagonal(:)
+        !-------------------------
+        n = size(b)
+        maxiter_ = n;       if(present(maxiter)) maxiter_ = maxiter
+        restart_ = .true.;  if(present(restart)) restart_ = restart
+        tol_ = 1.e-4_dp;    if(present(tol)) tol_ = tol
+
+        !-------------------------
+        ! internal memory setup
+        op%matvec => my_matvec
+        op%inner_product => my_dot
+        op%preconditionner => jacobi_preconditionner
+        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,4) )
+            workspace_ => workspace
+        else
+            allocate( workspace_ )
+            allocate( workspace_%tmp(n,4) , source = 0._dp )
+        end if
+        call diag(A,diagonal)
+        where(abs(diagonal)>epsilon(0.d0)) diagonal = 1._dp/diagonal
+        !-------------------------
+        ! main call to the solver
+        call solve_pccg_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
+        contains
+        
+        subroutine my_matvec(x,y)
+            real(dp), intent(in)  :: x(:)
+            real(dp), intent(inout) :: y(:)
+            call spmv( A , x, y )
+        end subroutine
+        subroutine jacobi_preconditionner(x,y)
+            real(dp), intent(in)  :: x(:)
+            real(dp), intent(inout) :: y(:)
+            y = diagonal * x
+        end subroutine
+        pure real(dp) function my_dot(x,y) result(r)
+            real(dp), intent(in) :: x(:)
+            real(dp), intent(in) :: y(:)
+            r = dot_product(x,y)
+        end function
+    end subroutine
+    
+end module custom_solver
+
+
+program example_solve_custom
+    use custom_solver
+    implicit none
+
+    type(CSR_dp_type) :: laplacian_csr
+    type(COO_dp_type) :: COO
+    real(dp) :: laplacian(5,5)
+    real(dp) :: x(5), load(5)
+    logical(1) :: dirichlet(5)
+
+    laplacian = reshape( [1, -1,  0,  0,  0,&
+                         -1,  2, -1,  0,  0,&
+                          0, -1,  2, -1,  0,&
+                          0,  0, -1,  2, -1,&
+                          0,  0,  0, -1,  1] , [5,5])
+    call dense2coo(laplacian,COO)
+    call coo2csr(COO,laplacian_csr)
+
+    x = 0._dp
+    load = dble( [0,0,5,0,0] )
+
+    dirichlet = .false._1 
+    dirichlet([1,5]) = .true._1
+
+    call solve_pccg_custom(laplacian_csr, load, x, tol=1.d-6, di=dirichlet)
+    print *, x !> solution: [0.0, 2.5, 5.0, 2.5, 0.0]
+    
+end program example_solve_custom