Skip to content

Commit 3912f73

Browse files
loiseaujcloiseaujc
committed
Added the Tridiagonal type and associated spmv kernel.
1 parent 4ebe266 commit 3912f73

File tree

2 files changed

+201
-3
lines changed

2 files changed

+201
-3
lines changed

‎src/stdlib_sparse_kinds.fypp

Lines changed: 149 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ module stdlib_sparse_kinds
1313
private
1414
public :: sparse_full, sparse_lower, sparse_upper
1515
public :: sparse_op_none, sparse_op_transpose, sparse_op_hermitian
16+
public :: Tridiagonal, dense
1617
!! version: experimental
1718
!!
1819
!! Base sparse type holding the meta data related to the storage capacity of a matrix.
@@ -128,6 +129,81 @@ module stdlib_sparse_kinds
128129
end type
129130
#:endfor
130131

132+
!------------------------------------------------------
133+
!----- -----
134+
!----- Tridiagonal matrix implementations -----
135+
!----- -----
136+
!------------------------------------------------------
137+
138+
!! version: experimental
139+
!!
140+
!! Tridiagonal matrix
141+
#:for k1, t1, s1 in (KINDS_TYPES)
142+
type, public, extends(sparse_type) :: Tridiagonal_${s1}$_type
143+
${t1}$, allocatable :: dl(:), dv(:), du(:)
144+
end type
145+
#:endfor
146+
147+
interface Tridiagonal
148+
!! This interface provides different methods to construct a
149+
!! `Tridiagonal` matrix. Only the non-zero elements of \( A \) are
150+
!! stored, i.e.
151+
!!
152+
!! \[
153+
!! A
154+
!! =
155+
!! \begin{bmatrix}
156+
!! a_1 & b_1 \\
157+
!! c_1 & a_2 & b_2 \\
158+
!! & \ddots & \ddots & \ddots \\
159+
!! & & c_{n-2} & a_{n-1} & b_{n-1} \\
160+
!! & & & c_{n-1} & a_n
161+
!! \end{bmatrix}.
162+
!! \]
163+
!!
164+
!! #### Syntax
165+
!!
166+
!! - Construct a real `Tridiagonal` matrix from rank-1 arrays:
167+
!!
168+
!! ```fortran
169+
!! integer, parameter :: n
170+
!! real(dp), allocatable :: dl(:), dv(:), du(:)
171+
!! type(Tridiagonal_rdp_type) :: A
172+
!! integer :: i
173+
!!
174+
!! dl = [(i, i=1, n-1)]; dv = [(2*i, i=1, n)]; du = [(3*i, i=1, n)]
175+
!! A = Tridiagonal(dl, dv, du)
176+
!! ```
177+
!!
178+
!! - Construct a real `Tridiagonal` matrix with constant diagonals:
179+
!!
180+
!! ```fortran
181+
!! integer, parameter :: n
182+
!! real(dp), parameter :: a = 1.0_dp, b = 1.0_dp, c = 2.0_dp
183+
!! type(Tridiagonal_rdp_type) :: A
184+
!!
185+
!! A = Tridiagonal(a, b, c, n)
186+
!! ```
187+
#:for k1, t1, s1 in (KINDS_TYPES)
188+
module procedure initialize_tridiagonal_${s1}$
189+
module procedure initialize_constant_tridiagonal_${s1}$
190+
#:endfor
191+
end interface
192+
193+
interface dense
194+
!! This interface provides methods to convert a `Tridiagonal` matrix
195+
!! to a regular rank-2 array.
196+
!!
197+
!! #### Syntax
198+
!!
199+
!! ```fortran
200+
!! B = dense(A)
201+
!! ```
202+
#:for k1, t1, s1 in (KINDS_TYPES)
203+
module procedure tridiagonal_to_dense_${s1}$
204+
#:endfor
205+
end interface
206+
131207
contains
132208

133209
!! (re)Allocate matrix memory for the COO type
@@ -589,5 +665,77 @@ contains
589665
end subroutine
590666

591667
#:endfor
668+
669+
!------------------------------------------------------
670+
!----- -----
671+
!----- Tridiagonal matrix implementations -----
672+
!----- -----
673+
!------------------------------------------------------
674+
675+
#:for k1, t1, s1 in (KINDS_TYPES)
676+
pure function initialize_tridiagonal_${s1}$(dl, dv, du) result(A)
677+
!! Construct a `Tridiagonal` matrix from the rank-1 arrays
678+
!! `dl`, `dv` and `du`.
679+
${t1}$, intent(in) :: dl(:), dv(:), du(:)
680+
!! Tridiagonal matrix elements.
681+
type(Tridiagonal_${s1}$_type) :: A
682+
!! Corresponding Tridiagonal matrix.
683+
684+
! Internal variables.
685+
integer(ilp) :: n
686+
687+
! Sanity check.
688+
n = size(dv)
689+
if (size(dl) /= n-1) error stop "Vector dl does not have the correct length."
690+
if (size(du) /= n-1) error stop "Vector du does not have the correct length."
691+
692+
! Description of the matrix.
693+
A%nrows = n ; A%ncols = n ; A%nnz = n + 2*(n-1)
694+
! Matrix elements.
695+
A%dl = dl ; A%dv = dv ; A%du = du
696+
end function
697+
698+
pure function initialize_constant_tridiagonal_${s1}$(dl, dv, du, n) result(A)
699+
!! Construct a `Tridiagonal` matrix with constant elements.
700+
${t1}$, intent(in) :: dl, dv, du
701+
!! Tridiagonal matrix elements.
702+
integer(ilp), intent(in) :: n
703+
!! Matrix dimension.
704+
type(Tridiagonal_${s1}$_type) :: A
705+
!! Corresponding Tridiagonal matrix.
706+
707+
! Internal variables.
708+
integer(ilp) :: i
709+
710+
! Description of the matrix.
711+
A%nrows = n ; A%ncols = n ; A%nnz = n + 2*(n-1)
712+
! Matrix elements.
713+
A%dl = [(dl, i = 1, n-1)]
714+
A%dv = [(dv, i = 1, n)]
715+
A%du = [(du, i = 1, n-1)]
716+
end function
717+
718+
pure function tridiagonal_to_dense_${s1}$(A) result(B)
719+
!! Convert a `Tridiagonal` matrix to its dense representation.
720+
type(Tridiagonal_${s1}$_type), intent(in) :: A
721+
!! Input Tridiagonal matrix.
722+
${t1}$, allocatable :: B(:, :)
723+
!! Corresponding dense matrix.
724+
725+
! Internal variables.
726+
integer(ilp) :: i
727+
728+
associate (n => A%nrows)
729+
allocate(B(n, n)) ; B = 0.0_${k1}$
730+
B(1, 1) = A%dv(1) ; B(1, 2) = A%du(1)
731+
do concurrent (i=2:n-1)
732+
B(i, i-1) = A%dl(i-1)
733+
B(i, i) = A%dv(i)
734+
B(i, i+1) = A%du(i)
735+
enddo
736+
B(n, n-1) = A%dl(n-1) ; B(n, n) = A%dv(n)
737+
end associate
738+
end function
739+
#:endfor
592740

593-
end module stdlib_sparse_kinds
741+
end module stdlib_sparse_kinds

‎src/stdlib_sparse_spmv.fypp

Lines changed: 52 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
module stdlib_sparse_spmv
1414
use stdlib_sparse_constants
1515
use stdlib_sparse_kinds
16+
use stdlib_linalg_lapack, only: lagtm
1617
implicit none
1718
private
1819

@@ -27,6 +28,9 @@ module stdlib_sparse_spmv
2728
module procedure :: spmv_csr_${rank}$d_${s1}$
2829
module procedure :: spmv_csc_${rank}$d_${s1}$
2930
module procedure :: spmv_ell_${rank}$d_${s1}$
31+
#:if k1 != "qp" and k1 != "xdp"
32+
module procedure :: spmv_tridiag_${rank}$d_${s1}$
33+
#:endif
3034
#:endfor
3135
module procedure :: spmv_sellc_${s1}$
3236
#:endfor
@@ -589,5 +593,51 @@ contains
589593
end subroutine
590594

591595
#:endfor
592-
593-
end module
596+
597+
!-----------------------------------------------------------
598+
!----- -----
599+
!----- TRIDIAGONAL MATRIX SPMV IMPLEMENTATIONS -----
600+
!----- -----
601+
!-----------------------------------------------------------
602+
!! spmv_tridiag
603+
#:for k1, t1, s1 in (KINDS_TYPES)
604+
#:for rank in RANKS
605+
#:if k1 != "qp" and k1 != "xdp"
606+
subroutine spmv_tridiag_${rank}$d_${s1}$(A, x, y, alpha, beta, op)
607+
type(Tridiagonal_${s1}$_type), intent(in) :: A
608+
${t1}$, intent(in), contiguous, target :: x${ranksuffix(rank)}$
609+
${t1}$, intent(inout), contiguous, target :: y${ranksuffix(rank)}$
610+
real(${k1}$), intent(in), optional :: alpha
611+
real(${k1}$), intent(in), optional :: beta
612+
character(1), intent(in), optional :: op
613+
614+
! Internal variables.
615+
real(${k1}$) :: alpha_, beta_
616+
integer(ilp) :: n, nrhs, ldx, ldy
617+
character(1) :: op_
618+
#:if rank == 1
619+
${t1}$, pointer :: xmat(:, :), ymat(:, :)
620+
#:endif
621+
622+
! Deal with optional arguments.
623+
alpha_ = 1.0_${k1}$ ; if (present(alpha)) alpha_ = alpha
624+
beta_ = 0.0_${k1}$ ; if (present(beta)) beta_ = beta
625+
op_ = "N" ; if (present(op)) op_ = op
626+
627+
! Prepare Lapack arguments.
628+
n = A%nrows ; ldx = n ; ldy = n ; y = 0.0_${k1}$
629+
nrhs = #{if rank==1}# 1 #{else}# size(x, 2) #{endif}#
630+
631+
#:if rank == 1
632+
! Pointer trick.
633+
xmat(1:n, 1:nrhs) => x ; ymat(1:n, 1:nrhs) => y
634+
call lagtm(op_, n, nrhs, alpha_, A%dl, A%dv, A%du, xmat, ldx, beta_, ymat, ldy)
635+
#:else
636+
call lagtm(op_, n, nrhs, alpha_, A%dl, A%dv, A%du, x, ldx, beta_, y, ldy)
637+
#:endif
638+
end subroutine
639+
#:endif
640+
#:endfor
641+
#:endfor
642+
643+
end module

0 commit comments

Comments
 (0)