add csc/coo conversions and diagonal extraction

jalvesz · jalvesz · commit bc0021b0885f · 2024-09-20T23:28:16.000+02:00
diff --git a/doc/specs/stdlib_sparse.md b/doc/specs/stdlib_sparse.md
@@ -259,6 +259,18 @@ This module provides facility functions for converting between storage formats.
 ```fortran
 {!example/linalg/example_sparse_from_ijv.f90!}
 ```
+### Syntax
+
+`call ` [[stdlib_sparse_conversion(module):diag(interface)]] `(matrix,diagonal)`
+
+### Arguments
+
+`matrix` : Shall be a `dense`, `COO`, `CSR` or `ELL` type. It is an `intent(in)` argument.
+
+`diagonal` : A rank-1 array of the same type as the `matrix`. It is an `intent(inout)` and `allocatable` argument.
+
+#### Note
+If the `diagonal` array has not been previously allocated, the `diag` subroutine will allocate it using the `nrows` of the `matrix`.
 
 ### Syntax
 
@@ -292,6 +304,16 @@ This module provides facility functions for converting between storage formats.
 
 ### Syntax
 
+`call ` [[stdlib_sparse_conversion(module):coo2csc(interface)]] `(coo,csc)`
+
+### Arguments
+
+`coo` : Shall be a `COO` type of `real` or `complex` type. It is an `intent(in)` argument.
+
+`csc` : Shall be a `CSC` type of `real` or `complex` type. It is an `intent(out)` argument.
+
+### Syntax
+
 `call ` [[stdlib_sparse_conversion(module):csr2coo(interface)]] `(csr,coo)`
 
 ### Arguments
@@ -312,6 +334,16 @@ This module provides facility functions for converting between storage formats.
 
 `chunk`, `optional`: chunk size for the Sliced ELLPACK format. It is an `intent(in)` argument.
 
+### Syntax
+
+`call ` [[stdlib_sparse_conversion(module):csc2coo(interface)]] `(csc,coo)`
+
+### Arguments
+
+`csc` : Shall be a `CSC` type of `real` or `complex` type. It is an `intent(in)` argument.
+
+`coo` : Shall be a `COO` type of `real` or `complex` type. It is an `intent(out)` argument.
+
 ## Example
 ```fortran
 {!example/linalg/example_sparse_spmv.f90!}
diff --git a/src/stdlib_sparse_conversion.fypp b/src/stdlib_sparse_conversion.fypp
@@ -58,6 +58,19 @@ module stdlib_sparse_conversion
         #:endfor
     end interface
     public :: coo2csr
+
+    !! version: experimental
+    !!
+    !! Conversion from coo to csc
+    !! Enables transferring data from a COO matrix to a CSC matrix
+    !! under the hypothesis that the COO is already ordered.
+    !! [Specifications](../page/specs/stdlib_sparse.html#sparse_conversion)
+    interface coo2csc
+        #:for k1, t1, s1 in (KINDS_TYPES)
+        module procedure :: coo2csc_${s1}$
+        #:endfor
+    end interface
+    public :: coo2csc
     
     !! version: experimental
     !!
@@ -111,6 +124,34 @@ module stdlib_sparse_conversion
     end interface
     public :: csr2sellc
 
+    !! version: experimental
+    !!
+    !! Conversion from csc to coo
+    !! Enables transferring data from a CSC matrix to a COO matrix
+    !! under the hypothesis that the CSC is already ordered.
+    !! [Specifications](../page/specs/stdlib_sparse.html#sparse_conversion)
+    interface csc2coo
+        #:for k1, t1, s1 in (KINDS_TYPES)
+        module procedure :: csc2coo_${s1}$
+        #:endfor
+    end interface
+    public :: csc2coo
+
+    !! version: experimental
+    !!
+    !! Extraction of diagonal values
+    !! [Specifications](../page/specs/stdlib_sparse.html#sparse_conversion)
+    interface diag
+    #:for k1, t1, s1 in (KINDS_TYPES)
+        module procedure :: dense2diagonal_${s1}$
+        module procedure :: coo2diagonal_${s1}$
+        module procedure :: csr2diagonal_${s1}$
+        module procedure :: csc2diagonal_${s1}$
+        module procedure :: ell2diagonal_${s1}$
+    #:endfor
+    end interface
+    public :: diag
+
     !! version: experimental
     !!
     !! Enable creating a sparse matrix from ijv (row,col,data) triplet
@@ -202,6 +243,45 @@ contains
 
     #:endfor
 
+    #:for k1, t1, s1 in (KINDS_TYPES)
+    subroutine coo2csc_${s1}$(COO,CSC)
+        type(COO_${s1}$_type), intent(in)  :: COO
+        type(CSC_${s1}$_type), intent(out) :: CSC
+        ${t1}$, allocatable :: data(:)
+        integer(ilp), allocatable :: temp(:,:)
+        integer(ilp) :: i, nnz
+
+        CSC%nnz = COO%nnz; CSC%nrows = COO%nrows; CSC%ncols = COO%ncols
+        CSC%storage = COO%storage
+
+        allocate(temp(2,COO%nnz))
+        temp(1,1:COO%nnz) = COO%index(2,1:COO%nnz)
+        temp(2,1:COO%nnz) = COO%index(1,1:COO%nnz)
+        allocate(data, source = COO%data )
+        nnz = COO%nnz
+        call sort_coo_unique_${s1}$( temp, data, nnz, COO%nrows, COO%ncols )
+
+        if( allocated(CSC%row) ) then
+            CSC%row(1:COO%nnz)  = temp(2,1:COO%nnz)
+            CSC%colptr(1:CSC%ncols) = 0
+            CSC%data(1:CSC%nnz) = data(1:COO%nnz)
+        else 
+            allocate( CSC%row(CSC%nnz)  , source = temp(2,1:COO%nnz) )
+            allocate( CSC%colptr(CSC%ncols+1) , source = 0 )
+            allocate( CSC%data(CSC%nnz) , source = data(1:COO%nnz) )
+        end if
+
+        CSC%colptr(1) = 1
+        do i = 1, COO%nnz
+            CSC%colptr( temp(1,i)+1 ) = CSC%colptr( temp(1,i)+1 ) + 1
+        end do
+        do i = 1, CSC%ncols
+            CSC%colptr( i+1 ) = CSC%colptr( i+1 ) + CSC%colptr( i )
+        end do
+    end subroutine
+
+    #:endfor
+
     #:for k1, t1, s1 in (KINDS_TYPES)
     subroutine csr2dense_${s1}$(CSR,dense)
         type(CSR_${s1}$_type), intent(in) :: CSR
@@ -254,6 +334,33 @@ contains
 
     #:endfor
 
+    #:for k1, t1, s1 in (KINDS_TYPES)
+    subroutine csc2coo_${s1}$(CSC,COO)
+        type(CSC_${s1}$_type), intent(in)  :: CSC
+        type(COO_${s1}$_type), intent(out) :: COO
+        integer(ilp) :: i, j
+
+        COO%nnz = CSC%nnz; COO%nrows = CSC%nrows; COO%ncols = CSC%ncols
+        COO%storage = CSC%storage
+
+        if( .not.allocated(COO%data) ) then
+            allocate( COO%data(CSC%nnz) , source = CSC%data(1:CSC%nnz) )
+        else 
+            COO%data(1:CSC%nnz) = CSC%data(1:CSC%nnz)
+        end if
+
+        if( .not.allocated(COO%index) ) allocate( COO%index(2,CSC%nnz) )
+        
+        do j = 1, CSC%ncols
+            do i = CSC%colptr(j), CSC%colptr(j+1)-1
+                COO%index(1:2,i) = [CSC%row(i),j]
+            end do
+        end do
+        call sort_coo_unique_${s1}$( COO%index, COO%data, COO%nnz, COO%nrows, COO%ncols )
+    end subroutine
+
+    #:endfor
+
     #:for k1, t1, s1 in (KINDS_TYPES)
     subroutine csr2ell_${s1}$(CSR,ELL,num_nz_rows)
         type(CSR_${s1}$_type), intent(in)  :: CSR
@@ -712,4 +819,119 @@ contains
     end subroutine
     #:endfor
 
+    !! Diagonal extraction
+
+    #:for k1, t1, s1 in (KINDS_TYPES)
+    subroutine dense2diagonal_${s1}$(dense,diagonal)
+        ${t1}$, intent(in) :: dense(:,:)
+        ${t1}$, intent(inout), allocatable :: diagonal(:)
+        integer :: num_rows
+        integer :: i
+
+        num_rows = size(dense,dim=1)
+        if(.not.allocated(diagonal)) allocate(diagonal(num_rows))
+
+        do i = 1, num_rows
+            diagonal(i) = dense(i,i)
+        end do
+    end subroutine
+
+    #:endfor
+
+    #:for k1, t1, s1 in (KINDS_TYPES)
+    subroutine coo2diagonal_${s1}$(COO,diagonal)
+        type(COO_${s1}$_type), intent(in) :: COO
+        ${t1}$, intent(inout), allocatable :: diagonal(:)
+        integer :: idx
+
+        if(.not.allocated(diagonal)) allocate(diagonal(COO%nrows))
+
+        do concurrent(idx = 1:COO%nnz)
+            if(COO%index(1,idx)==COO%index(2,idx)) &
+            & diagonal( COO%index(1,idx) ) = COO%data(idx)
+        end do
+    end subroutine
+
+    #:endfor
+
+    #:for k1, t1, s1 in (KINDS_TYPES)
+    subroutine csr2diagonal_${s1}$(CSR,diagonal)
+        type(CSR_${s1}$_type), intent(in) :: CSR
+        ${t1}$, intent(inout), allocatable :: diagonal(:)
+        integer :: i, j
+
+        if(.not.allocated(diagonal)) allocate(diagonal(CSR%nrows))
+
+        select case(CSR%storage)
+        case(sparse_lower)
+            do i = 1, CSR%nrows
+                diagonal(i) = CSR%data( CSR%rowptr(i+1)-1 )
+            end do
+        case(sparse_upper)
+            do i = 1, CSR%nrows
+                diagonal(i) = CSR%data( CSR%rowptr(i) )
+            end do
+        case(sparse_full)
+            do i = 1, CSR%nrows
+                do j = CSR%rowptr(i), CSR%rowptr(i+1)-1
+                    if( CSR%col(j) == i ) then
+                        diagonal(i) = CSR%data(j)
+                        exit
+                    end if
+                end do
+            end do
+        end select
+    end subroutine
+
+    #:endfor
+
+    #:for k1, t1, s1 in (KINDS_TYPES)
+    subroutine csc2diagonal_${s1}$(CSC,diagonal)
+        type(CSC_${s1}$_type), intent(in) :: CSC
+        ${t1}$, intent(inout), allocatable :: diagonal(:)
+        integer :: i, j
+
+        if(.not.allocated(diagonal)) allocate(diagonal(CSC%nrows))
+
+        select case(CSC%storage)
+        case(sparse_lower)
+            do i = 1, CSC%ncols
+                diagonal(i) = CSC%data( CSC%colptr(i+1)-1 )
+            end do
+        case(sparse_upper)
+            do i = 1, CSC%ncols
+                diagonal(i) = CSC%data( CSC%colptr(i) )
+            end do
+        case(sparse_full)
+            do i = 1, CSC%ncols
+                do j = CSC%colptr(i), CSC%colptr(i+1)-1
+                    if( CSC%row(j) == i ) then
+                        diagonal(i) = CSC%data(j)
+                        exit
+                    end if
+                end do
+            end do
+        end select
+    end subroutine
+
+    #:endfor
+
+    #:for k1, t1, s1 in (KINDS_TYPES)
+    subroutine ell2diagonal_${s1}$(ELL,diagonal)
+        type(ELL_${s1}$_type), intent(in) :: ELL
+        ${t1}$, intent(inout), allocatable :: diagonal(:)
+        integer :: i, k
+
+        if(.not.allocated(diagonal)) allocate(diagonal(ELL%nrows))
+        if( ELL%storage == sparse_full) then
+            do i = 1, ELL%nrows
+                do k = 1, ELL%K
+                    if(ELL%index(i,k)==i) diagonal(i) = ELL%data(i,k)
+                end do
+            end do
+        end if
+    end subroutine
+
+    #:endfor
+
 end module
diff --git a/test/linalg/test_sparse_spmv.fypp b/test/linalg/test_sparse_spmv.fypp
@@ -24,6 +24,7 @@ contains
             new_unittest('ell', test_ell),  &
             new_unittest('sellc', test_sellc),  &
             new_unittest('symmetries', test_symmetries), &
+            new_unittest('diagonal', test_diagonal), &
             new_unittest('add_get_values', test_add_get_values) &
         ]
     end subroutine
@@ -253,6 +254,50 @@ contains
         #:endfor
     end subroutine
 
+    subroutine test_diagonal(error)
+        !> Error handling
+        type(error_type), allocatable, intent(out) :: error
+        #:for k1, t1, s1 in (KINDS_TYPES)
+        block
+            integer, parameter :: wp = ${k1}$
+            ${t1}$, allocatable :: dense(:,:)
+            type(COO_${s1}$_type) :: COO
+            type(CSR_${s1}$_type) :: CSR
+            type(CSC_${s1}$_type) :: CSC
+            type(ELL_${s1}$_type) :: ELL
+            ${t1}$, allocatable :: diagonal(:)
+
+            allocate( dense(4,4) , source = &
+                    reshape(real([1,0,0,5, &
+                                  0,2,0,0, &
+                                  0,6,3,0,&
+                                  0,0,7,4],kind=wp),[4,4]) )
+
+            call diag(dense,diagonal)
+            call check(error, all(diagonal == [1,2,3,4]) )
+            if (allocated(error)) return
+
+            diagonal = 0.0
+            call dense2coo( dense , COO )
+            call diag( COO , diagonal )
+            call check(error, all(diagonal == [1,2,3,4]) )
+            if (allocated(error)) return
+
+            diagonal = 0.0
+            call coo2csr( COO, CSR )
+            call diag( CSR , diagonal )
+            call check(error, all(diagonal == [1,2,3,4]) )
+            if (allocated(error)) return
+
+            diagonal = 0.0
+            call coo2csc( COO, CSC )
+            call diag( CSC , diagonal )
+            call check(error, all(diagonal == [1,2,3,4]) )
+            if (allocated(error)) return
+        end block
+        #:endfor
+    end subroutine
+
     subroutine test_add_get_values(error)
         !> Error handling
         type(error_type), allocatable, intent(out) :: error
