Tridagonal matrices and associated spmv kernel

example/linalg/example_sparse_data_accessors.f90

@@ -1,49 +1,49 @@
 program example_sparse_data_accessors
-    use stdlib_linalg_constants, only: dp
-    use stdlib_sparse
-    implicit none
+   use stdlib_linalg_constants, only: dp
+   use stdlib_sparse
+   implicit none
 
-    real(dp) :: mat(2,2)
-    real(dp), allocatable :: dense(:,:)
-    type(CSR_dp_type) :: CSR
-    type(COO_dp_type) :: COO
-    integer :: i, j, locdof(2)
+   real(dp) :: mat(2, 2)
+   real(dp), allocatable :: dense_matrix(:, :)
+   type(CSR_dp_type) :: CSR
+   type(COO_dp_type) :: COO
+   integer :: i, j, locdof(2)
 
-    ! Initial data
-    mat(:,1) = [1._dp,2._dp]
-    mat(:,2) = [2._dp,1._dp]
-    allocate(dense(5,5) , source = 0._dp)
-    do i = 0, 3
-        dense(1+i:2+i,1+i:2+i) = dense(1+i:2+i,1+i:2+i) + mat
-    end do
+   ! Initial data
+   mat(:, 1) = [1._dp, 2._dp]
+   mat(:, 2) = [2._dp, 1._dp]
+   allocate (dense_matrix(5, 5), source=0._dp)
+   do i = 0, 3
+      dense_matrix(1 + i:2 + i, 1 + i:2 + i) = dense_matrix(1 + i:2 + i, 1 + i:2 + i) + mat
+   end do
 
-    print *, 'Original Matrix'
-    do j = 1 , 5
-        print '(5f8.1)',dense(j,:)
-    end do
+   print *, 'Original Matrix'
+   do j = 1, 5
+      print '(5f8.1)', dense_matrix(j, :)
+   end do
 
-    ! Initialize CSR data and reset dense reference matrix
-    call dense2coo(dense,COO)
-    call coo2csr(COO,CSR)
-    CSR%data = 0._dp
-    dense = 0._dp
+   ! Initialize CSR data and reset dense reference matrix
+   call dense2coo(dense_matrix, COO)
+   call coo2csr(COO, CSR)
+   CSR%data = 0._dp
+   dense_matrix = 0._dp
 
-    ! Iteratively add blocks of data
-    do i = 0, 3
-        locdof(1:2) = [1+i,2+i] 
-        call CSR%add(locdof,locdof,mat)
-        ! lets print a dense view of every step
-        call csr2dense(CSR,dense)
-        print '(A,I2)', 'Add block :', i+1
-        do j = 1 , 5
-            print '(5f8.1)',dense(j,:)
-        end do
-    end do
+   ! Iteratively add blocks of data
+   do i = 0, 3
+      locdof(1:2) = [1 + i, 2 + i]
+      call CSR%add(locdof, locdof, mat)
+      ! lets print a dense view of every step
+      call csr2dense(CSR, dense_matrix)
+      print '(A,I2)', 'Add block :', i + 1
+      do j = 1, 5
+         print '(5f8.1)', dense_matrix(j, :)
+      end do
+   end do
 
-    ! Request values from the matrix
-    print *, ''
-    print *, 'within sparse pattern  :',CSR%at(2,1)
-    print *, 'outside sparse pattern :',CSR%at(5,2)
-    print *, 'outside matrix pattern :',CSR%at(7,7)
-  
-end program example_sparse_data_accessors
+   ! Request values from the matrix
+   print *, ''
+   print *, 'within sparse pattern  :', CSR%at(2, 1)
+   print *, 'outside sparse pattern :', CSR%at(5, 2)
+   print *, 'outside matrix pattern :', CSR%at(7, 7)
+
+end program example_sparse_data_accessors

src/stdlib_sparse_kinds.fypp

@@ -13,6 +13,7 @@ module stdlib_sparse_kinds
     private
     public :: sparse_full, sparse_lower, sparse_upper
     public :: sparse_op_none, sparse_op_transpose, sparse_op_hermitian
+    public :: Tridiagonal, dense
     !! version: experimental
     !!
     !! Base sparse type holding the meta data related to the storage capacity of a matrix.
@@ -128,6 +129,81 @@ module stdlib_sparse_kinds
     end type
     #:endfor
 
+    !------------------------------------------------------
+    !-----                                            -----
+    !-----     Tridiagonal matrix implementations     -----
+    !-----                                            -----
+    !------------------------------------------------------
+
+    !! version: experimental
+    !!
+    !! Tridiagonal matrix
+    #:for k1, t1, s1 in (KINDS_TYPES)
+    type, public, extends(sparse_type) :: Tridiagonal_${s1}$_type
+        ${t1}$, allocatable :: dl(:), dv(:), du(:)
+    end type
+    #:endfor
+
+    interface Tridiagonal
+        !! This interface provides different methods to construct a
+        !! `Tridiagonal` matrix. Only the non-zero elements of \( A \) are
+        !! stored, i.e.
+        !!
+        !! \[
+        !!    A
+        !!    =
+        !!    \begin{bmatrix}
+        !!       a_1   &  b_1  \\
+        !!       c_1  &  a_2      &  b_2  \\
+        !!             &  \ddots   &  \ddots   &  \ddots   \\
+        !!             &           &  c_{n-2} &  a_{n-1}  &  b_{n-1} \\
+        !!             &           &           &  c_{n-1} &  a_n
+        !!    \end{bmatrix}.
+        !! \]
+        !!
+        !! #### Syntax
+        !!
+        !! - Construct a real `Tridiagonal` matrix from rank-1 arrays:
+        !!
+        !! ```fortran
+        !!    integer, parameter :: n
+        !!    real(dp), allocatable :: dl(:), dv(:), du(:)
+        !!    type(Tridiagonal_rdp_type) :: A
+        !!    integer :: i
+        !!
+        !!    dl = [(i, i=1, n-1)]; dv = [(2*i, i=1, n)]; du = [(3*i, i=1, n)]
+        !!    A = Tridiagonal(dl, dv, du)
+        !! ```
+        !!
+        !! - Construct a real `Tridiagonal` matrix with constant diagonals:
+        !!
+        !! ```fortran
+        !!    integer, parameter :: n
+        !!    real(dp), parameter :: a = 1.0_dp, b = 1.0_dp, c = 2.0_dp
+        !!    type(Tridiagonal_rdp_type) :: A
+        !!
+        !!    A = Tridiagonal(a, b, c, n)
+        !! ```
+        #:for k1, t1, s1 in (KINDS_TYPES)
+        module procedure initialize_tridiagonal_${s1}$
+        module procedure initialize_constant_tridiagonal_${s1}$
+        #:endfor
+    end interface
+
+    interface dense
+        !! This interface provides methods to convert a `Tridiagonal` matrix
+        !! to a regular rank-2 array.
+        !!
+        !! #### Syntax
+        !!
+        !! ```fortran
+        !!    B = dense(A)
+        !! ```
+        #:for k1, t1, s1 in (KINDS_TYPES)
+        module procedure tridiagonal_to_dense_${s1}$
+        #:endfor
+    end interface
+
 contains
 
     !! (re)Allocate matrix memory for the COO type
@@ -589,5 +665,77 @@ contains
     end subroutine
 
     #:endfor
+
+    !------------------------------------------------------
+    !-----                                            -----
+    !-----     Tridiagonal matrix implementations     -----
+    !-----                                            -----
+    !------------------------------------------------------
+
+    #:for k1, t1, s1 in (KINDS_TYPES)
+    pure function initialize_tridiagonal_${s1}$(dl, dv, du) result(A)
+        !! Construct a `Tridiagonal` matrix from the rank-1 arrays
+        !! `dl`, `dv` and `du`.
+        ${t1}$, intent(in) :: dl(:), dv(:), du(:)
+        !! Tridiagonal matrix elements.
+        type(Tridiagonal_${s1}$_type) :: A
+        !! Corresponding Tridiagonal matrix.
+
+        ! Internal variables.
+        integer(ilp) :: n
+
+        ! Sanity check.
+        n = size(dv)
+        if (size(dl) /= n-1) error stop "Vector dl does not have the correct length."
+        if (size(du) /= n-1) error stop "Vector du does not have the correct length."
+
+        ! Description of the matrix.
+        A%nrows = n ; A%ncols = n ; A%nnz = n + 2*(n-1)
+        ! Matrix elements.
+        A%dl = dl ; A%dv = dv ; A%du = du
+    end function
+
+    pure function initialize_constant_tridiagonal_${s1}$(dl, dv, du, n) result(A)
+        !! Construct a `Tridiagonal` matrix with constant elements.
+        ${t1}$, intent(in) :: dl, dv, du
+        !! Tridiagonal matrix elements.
+        integer(ilp), intent(in) :: n
+        !! Matrix dimension.
+        type(Tridiagonal_${s1}$_type) :: A
+        !! Corresponding Tridiagonal matrix.
+
+        ! Internal variables.
+        integer(ilp) :: i
+
+        ! Description of the matrix.
+        A%nrows = n ; A%ncols = n ; A%nnz = n + 2*(n-1)
+        ! Matrix elements.
+        A%dl = [(dl, i = 1, n-1)]
+        A%dv = [(dv, i = 1, n)]
+        A%du = [(du, i = 1, n-1)]
+    end function
+
+    pure function tridiagonal_to_dense_${s1}$(A) result(B)
+        !! Convert a `Tridiagonal` matrix to its dense representation.
+        type(Tridiagonal_${s1}$_type), intent(in) :: A
+        !! Input Tridiagonal matrix.
+        ${t1}$, allocatable :: B(:, :)
+        !! Corresponding dense matrix.
+
+        ! Internal variables.
+        integer(ilp) :: i
+
+        associate (n => A%nrows)
+        allocate(B(n, n)) ; B = 0.0_${k1}$
+        B(1, 1) = A%dv(1) ; B(1, 2) = A%du(1)
+        do concurrent (i=2:n-1)
+            B(i, i-1) = A%dl(i-1)
+            B(i, i) = A%dv(i)
+            B(i, i+1) = A%du(i)
+        enddo
+        B(n, n-1) = A%dl(n-1) ; B(n, n) = A%dv(n)
+        end associate
+    end function
+    #:endfor
 
-end module stdlib_sparse_kinds
+end module stdlib_sparse_kinds

src/stdlib_sparse_spmv.fypp

@@ -13,6 +13,7 @@
 module stdlib_sparse_spmv
     use stdlib_sparse_constants
     use stdlib_sparse_kinds
+    use stdlib_linalg_lapack, only: lagtm
     implicit none
     private
 
@@ -27,6 +28,9 @@ module stdlib_sparse_spmv
         module procedure :: spmv_csr_${rank}$d_${s1}$
         module procedure :: spmv_csc_${rank}$d_${s1}$
         module procedure :: spmv_ell_${rank}$d_${s1}$
+        #:if k1 != "qp" and k1 != "xdp"
+        module procedure :: spmv_tridiag_${rank}$d_${s1}$
+        #:endif
         #:endfor
         module procedure :: spmv_sellc_${s1}$
         #:endfor
@@ -589,5 +593,51 @@ contains
     end subroutine
 
     #:endfor
-
-end module
+
+    !-----------------------------------------------------------
+    !-----                                                 -----
+    !-----     TRIDIAGONAL MATRIX SPMV IMPLEMENTATIONS     -----
+    !-----                                                 -----
+    !-----------------------------------------------------------
+    !! spmv_tridiag
+    #:for k1, t1, s1 in (KINDS_TYPES)
+    #:for rank in RANKS
+    #:if k1 != "qp" and k1 != "xdp"
+    subroutine spmv_tridiag_${rank}$d_${s1}$(A, x, y, alpha, beta, op)
+        type(Tridiagonal_${s1}$_type), intent(in) :: A
+        ${t1}$, intent(in), contiguous, target :: x${ranksuffix(rank)}$
+        ${t1}$, intent(inout), contiguous, target :: y${ranksuffix(rank)}$
+        real(${k1}$), intent(in), optional :: alpha
+        real(${k1}$), intent(in), optional :: beta
+        character(1), intent(in), optional :: op
+
+        ! Internal variables.
+        real(${k1}$) :: alpha_, beta_
+        integer(ilp) :: n, nrhs, ldx, ldy
+        character(1) :: op_
+        #:if rank == 1
+        ${t1}$, pointer :: xmat(:, :), ymat(:, :)
+        #:endif
+
+        ! Deal with optional arguments.
+        alpha_ = 1.0_${k1}$     ; if (present(alpha)) alpha_ = alpha
+        beta_  = 0.0_${k1}$     ; if (present(beta))  beta_  = beta
+        op_    = sparse_op_none ; if (present(op))    op_    = op
+
+        ! Prepare Lapack arguments.
+        n = A%nrows ; ldx = n ; ldy = n ; y = 0.0_${k1}$
+        nrhs = #{if rank==1}# 1 #{else}# size(x, 2) #{endif}#
+
+        #:if rank == 1
+        ! Pointer trick.
+        xmat(1:n, 1:nrhs) => x ; ymat(1:n, 1:nrhs) => y
+        call lagtm(op_, n, nrhs, alpha_, A%dl, A%dv, A%du, xmat, ldx, beta_, ymat, ldy)
+        #:else
+        call lagtm(op_, n, nrhs, alpha_, A%dl, A%dv, A%du, x, ldx, beta_, y, ldy)
+        #:endif
+    end subroutine
+    #:endif
+    #:endfor
+    #:endfor
+
+end module