Skip to content

Commit

Permalink
use can now provide a custom sparse linear solver
Browse files Browse the repository at this point in the history 
also added some params for the different sparsity modes
Fixes #20
  • Loading branch information
@jacobwilliams
jacobwilliams committed Feb 23, 2024
1 parent 6398e49 commit c309c20
Showing 1 changed file with 74 additions and 22 deletions.
96 changes: 74 additions & 22 deletions src/nlesolver_module.F90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,13 @@ module nlesolver_module
real(wp),parameter :: eps = epsilon(one) !! machine \( \epsilon \)
real(wp),parameter :: big = huge(one)

! options for sparsity_mode
integer,parameter,public :: NLESOLVER_SPARSITY_DENSE = 1 !! assume dense (use dense solver).
integer,parameter,public :: NLESOLVER_SPARSITY_LSQR = 2 !! assume sparse (use LSQR sparse solver).
integer,parameter,public :: NLESOLVER_SPARSITY_LUSOL = 3 !! assume sparse (use LUSOL sparse solver).
integer,parameter,public :: NLESOLVER_SPARSITY_LSMR = 4 !! assume sparse (use LSMR sparse solver).
integer,parameter,public :: NLESOLVER_SPARSITY_CUSTOM_SPARSE = 5 !! assume sparse (use a user provided sparse solver).

!*********************************************************
type,public :: nlesolver_type

Expand Down Expand Up @@ -105,12 +112,13 @@ module nlesolver_module
procedure(linesearch_func),pointer :: linesearch => null() !! line search method (determined by `step_mode` user input in [[nlesolver_type:initialize]])

! sparsity options:
integer :: sparsity_mode = 1 !! sparsity mode:
!!
!! * 1 - assume dense (use dense solver)
!! * 2 - assume sparse (use LSQR sparse solver).
!! * 3 - assume sparse (use LUSOL sparse solver).
!! * 4 - assume sparse (use LSMR sparse solver).
integer :: sparsity_mode = NLESOLVER_SPARSITY_DENSE !! sparsity mode:
!!
!! * 1 - assume dense (use dense solver).
!! * 2 - assume sparse (use LSQR sparse solver).
!! * 3 - assume sparse (use LUSOL sparse solver).
!! * 4 - assume sparse (use LSMR sparse solver).
!! * 5 - assume sparse (use a user provided sparse solver).
integer :: n_nonzeros = -1 !! number of nonzero Jacobian elements (used for `sparsity_mode > 1`)
integer,dimension(:),allocatable :: irow !! sparsity pattern nonzero elements row indices.
integer,dimension(:),allocatable :: icol !! sparsity pattern nonzero elements column indices
Expand Down Expand Up @@ -144,6 +152,9 @@ module nlesolver_module
! sparse version:
procedure(grad_func_sparse),pointer :: grad_sparse => null() !! user-supplied routine to compute the gradient of the function (sparse version)

! custom sparse solver:
procedure(sparse_solver_func),pointer :: custom_solver_sparse => null() !! user-supplied sparse linear solver routine (used for `sparsity_mode=5`)

contains

private
Expand All @@ -160,6 +171,22 @@ module nlesolver_module

abstract interface

subroutine sparse_solver_func(me,n_cols,n_rows,n_nonzero,irow,icol,a,b,x,istat)
!! for a custom user-provided linear solver (sparse version)
!! solve `Ax = b` for `x`, given `A` and `b`.
import :: wp,nlesolver_type
implicit none
class(nlesolver_type),intent(inout) :: me
integer,intent(in) :: n_cols !! `n`: number of columns in A.
integer,intent(in) :: n_rows !! `m`: number of rows in A.
integer,intent(in) :: n_nonzero !! number of nonzero elements of A.
integer,dimension(n_nonzero),intent(in) :: irow, icol !! sparsity pattern (size is `n_nonzero`)
real(wp),dimension(n_nonzero),intent(in) :: a !! matrix elements (size is `n_nonzero`)
real(wp),dimension(n_rows),intent(in) :: b !! right hand side (size is `m`)
real(wp),dimension(n_cols),intent(out) :: x !! solution (size is `n`)
integer,intent(out) :: istat !! status code (=0 for success)
end subroutine sparse_solver_func

subroutine func_func(me,x,f)
!! compute the function
import :: wp,nlesolver_type
Expand Down Expand Up @@ -332,7 +359,7 @@ subroutine initialize_nlesolver_variables(me,&
verbose,iunit,n_uphill_max,n_intervals,&
sparsity_mode,irow,icol,&
atol,btol,conlim,damp,itnlim,nout,&
lusol_method,localSize )
lusol_method,localSize,custom_solver_sparse )

implicit none

Expand Down Expand Up @@ -378,6 +405,8 @@ subroutine initialize_nlesolver_variables(me,&
!! * 2 - assume sparse (use LSQR sparse solver).
!! * 3 - assume sparse (use LUSOL sparse solver).
!! * 4 - assume sparse (use LSMR sparse solver).
!! * 5 - assume sparse (use the user provided sparse
!! solver `custom_solver_sparse`).
!! Must also specify `grad_sparse` and the `irow` and `icol`
!! sparsity pattern for `mode>1`.
integer,dimension(:),intent(in),optional :: irow !! sparsity pattern nonzero elements row indices.
Expand All @@ -404,6 +433,8 @@ subroutine initialize_nlesolver_variables(me,&
!!
!! localSize need not be more than `min(m,n)`.
!! At most `min(m,n)` vectors will be allocated.
procedure(sparse_solver_func),optional :: custom_solver_sparse !! for `sparsity_mode=5`, this is the
!! user-provided linear solver.

logical :: status_ok !! true if there were no errors

Expand Down Expand Up @@ -493,9 +524,11 @@ subroutine initialize_nlesolver_variables(me,&
if (allocated(me%irow)) deallocate(me%irow)
if (allocated(me%icol)) deallocate(me%icol)

if (present(custom_solver_sparse)) me%custom_solver_sparse => custom_solver_sparse

if (present(sparsity_mode)) then
me%sparsity_mode = sparsity_mode
if (sparsity_mode>1) then ! sparse solver method
if (sparsity_mode>NLESOLVER_SPARSITY_DENSE) then ! sparse solver method
if (present(irow) .and. present(icol) .and. present(grad_sparse)) then
if (size(irow) == size(icol)) then
me%n_nonzeros = size(irow)
Expand Down Expand Up @@ -524,6 +557,12 @@ subroutine initialize_nlesolver_variables(me,&
if (present(localSize)) me%localSize = localSize

end if
if (sparsity_mode==NLESOLVER_SPARSITY_CUSTOM_SPARSE) then
if (.not. associated(me%custom_solver_sparse)) then
call me%set_status(istat = -16, string = 'Error: must specify custom_solver_sparse for sparsity_mode = 5')
return

Check warning on line 563 in src/nlesolver_module.F90

View check run for this annotation

Codecov / codecov/patch

src/nlesolver_module.F90#L561-L563 

Added lines #L561 - L563 were not covered by tests
end if
end if
end if

if (status_ok) then
Expand Down Expand Up @@ -581,11 +620,11 @@ subroutine nlesolver_solver(me,x)
call me%set_status(istat = -10, string = 'Error: function routine is not associated')
return
end if
if (me%sparsity_mode==1 .and. .not. associated(me%grad)) then
if (me%sparsity_mode==NLESOLVER_SPARSITY_DENSE .and. .not. associated(me%grad)) then
call me%set_status(istat = -11, string = 'Error: gradient routine is not associated')
return
end if
if (me%sparsity_mode>1 .and. .not. associated(me%grad_sparse)) then
if (me%sparsity_mode>NLESOLVER_SPARSITY_DENSE .and. .not. associated(me%grad_sparse)) then
call me%set_status(istat = -11, string = 'Error: gradient routine is not associated')
return
end if
Expand All @@ -600,7 +639,7 @@ subroutine nlesolver_solver(me,x)
alloc_stat = 0
if (alloc_stat==0) allocate(fvec (me%m) , stat=alloc_stat)

if (me%sparsity_mode==1) then
if (me%sparsity_mode==NLESOLVER_SPARSITY_DENSE) then
! dense
if (alloc_stat==0) allocate(fjac (me%m,me%n) , stat=alloc_stat)
else
Expand All @@ -615,7 +654,7 @@ subroutine nlesolver_solver(me,x)
if (me%use_broyden) then
! only need these for broyden:
if (alloc_stat==0) allocate(prev_fvec(me%m) , stat=alloc_stat)
if (me%sparsity_mode/=1 .and. alloc_stat==0) then
if (me%sparsity_mode/=NLESOLVER_SPARSITY_DENSE .and. alloc_stat==0) then
allocate(prev_fjac(me%m,me%n) , stat=alloc_stat)
index_array = [(i, i=1,me%n_nonzeros)] ! an array to index the sparse jacobian elements
end if
Expand Down Expand Up @@ -661,9 +700,9 @@ subroutine nlesolver_solver(me,x)
! always compute Jacobian on the first iteration
!call me%grad(x,fjac)
select case (me%sparsity_mode)
case (1)
case (NLESOLVER_SPARSITY_DENSE)
call me%grad(x,fjac)
case (2:)
case default ! sparse modes
call me%grad_sparse(x,fjac_sparse)
end select
broyden_counter = 0
Expand All @@ -675,7 +714,7 @@ subroutine nlesolver_solver(me,x)
delx(:,1) = x - xold
delf(:,1) = fvec - prev_fvec

if (me%sparsity_mode==1) then ! dense
if (me%sparsity_mode==NLESOLVER_SPARSITY_DENSE) then ! dense

delxmag2 = dot_product(delx(:,1),delx(:,1))
if (delxmag2 < eps) then
Expand Down Expand Up @@ -718,17 +757,17 @@ subroutine nlesolver_solver(me,x)
broyden_counter = broyden_counter + 1
end if
prev_fvec = fvec
if (me%sparsity_mode==1) then
if (me%sparsity_mode==NLESOLVER_SPARSITY_DENSE) then
prev_fjac = fjac
else
prev_fjac_sparse = fjac_sparse
end if
else
! compute the jacobian:
select case (me%sparsity_mode)
case (1)
case (NLESOLVER_SPARSITY_DENSE)
call me%grad(x,fjac)
case (2:)
case default ! sparse
call me%grad_sparse(x,fjac_sparse)
end select
recompute_jac = .false. ! for broyden
Expand All @@ -741,10 +780,12 @@ subroutine nlesolver_solver(me,x)
! compute the search direction p by solving linear system:
rhs = -fvec ! RHS of the linear system
select case (me%sparsity_mode)
case (1)

case (NLESOLVER_SPARSITY_DENSE)
! use dense solver
call linear_solver(me%m,me%n,fjac,rhs,p,info)
case (2)

case (NLESOLVER_SPARSITY_LSQR)
! initialize the LSQR sparse solver
call sparse_solver%initialize(me%m,me%n,fjac_sparse,me%irow,me%icol,&
atol = me%atol, &
Expand All @@ -766,12 +807,14 @@ subroutine nlesolver_solver(me,x)
case default
info = 0
end select
case (3)

case (NLESOLVER_SPARSITY_LUSOL)
! use lusol solver
lusol_options%method = me%lusol_method
call solve(me%n,me%m,me%n_nonzeros,me%irow,me%icol,fjac_sparse,rhs,p,info,&
settings=lusol_options)
case (4)

case (NLESOLVER_SPARSITY_LSMR)

! use LSMR solver:
!me%conlim = 1.0_wp/(100*epsilon(1.0_wp))
Expand All @@ -793,6 +836,15 @@ subroutine nlesolver_solver(me,x)
info = 0
end select

case (NLESOLVER_SPARSITY_CUSTOM_SPARSE)
if (associated(me%custom_solver_sparse)) then

Check warning on line 840 in src/nlesolver_module.F90

View check run for this annotation

Codecov / codecov/patch

src/nlesolver_module.F90#L840 

Added line #L840 was not covered by tests
! a user-provided sparse solver:
call me%custom_solver_sparse(me%n,me%m,me%n_nonzeros,me%irow,me%icol,fjac_sparse,rhs,p,info)

Check warning on line 842 in src/nlesolver_module.F90

View check run for this annotation

Codecov / codecov/patch

src/nlesolver_module.F90#L842 

Added line #L842 was not covered by tests
else
call me%set_status(istat = -1006, &
string = 'Error: The custom_solver_sparse procedure has not been set.')
exit

Check warning on line 846 in src/nlesolver_module.F90

View check run for this annotation

Codecov / codecov/patch

src/nlesolver_module.F90#L845-L846 

Added lines #L845 - L846 were not covered by tests
end if
case default
error stop 'invalid sparsity_mode'
end select
Expand Down

0 comments on commit c309c20

Please sign in to comment.