export working space interface

perazz · perazz · commit 2e370f45438a · 2024-04-29T10:34:03.000+02:00
diff --git a/src/stdlib_linalg.fypp b/src/stdlib_linalg.fypp
@@ -21,6 +21,8 @@ module stdlib_linalg
   public :: diag
   public :: eye
   public :: lstsq
+  public :: lstsq_space
+  public :: solve_lstsq
   public :: trace
   public :: outer_product
   public :: kronecker_product
@@ -273,7 +275,7 @@ module stdlib_linalg
     !! version: experimental 
     !!
     !! Computes the squares solution to system \( A \cdot x = b \). 
-    !! ([Specification](../page/specs/stdlib_linalg.html#XXXXXxXxxxxxxx))xxxxx
+    !! ([Specification](../page/specs/stdlib_linalg.html))
     !! 
     !!### Summary 
     !! Subroutine interface for computing least squares, i.e. the 2-norm \( || (b-A \cdot x ||_2 \) minimizing solution.
@@ -291,7 +293,7 @@ module stdlib_linalg
     #:for rk,rt,ri in RC_KINDS_TYPES
     #:if rk!="xdp"
       module subroutine stdlib_linalg_${ri}$_solve_lstsq_${ndsuf}$(a,b,x,real_storage,int_storage,&
-                        #{if rt.startswith('c')}#cmpl_storage,#{endif}#cond,overwrite_a,rank,err) 
+                        #{if rt.startswith('c')}#cmpl_storage,#{endif}#cond,singvals,overwrite_a,rank,err) 
          !> Input matrix a[n,n]
          ${rt}$, intent(inout), target :: a(:,:)
          !> Right hand side vector or array, b[n] or b[n,nrhs]
@@ -308,6 +310,8 @@ module stdlib_linalg
          #:endif
          !> [optional] cutoff for rank evaluation: singular values s(i)<=cond*maxval(s) are considered 0.
          real(${rk}$), optional, intent(in) :: cond
+         !> [optional] list of singular values [min(m,n)], in descending magnitude order, returned by the SVD
+         real(${rk}$), optional, intent(out), target :: singvals(:)                  
          !> [optional] Can A,b data be overwritten and destroyed?
          logical(lk), optional, intent(in) :: overwrite_a
          !> [optional] Return rank of A
@@ -320,6 +324,36 @@ module stdlib_linalg
     #:endfor
   end interface solve_lstsq
 
+  ! Return the working array space required by the least squares solver
+  interface lstsq_space
+    !! version: experimental 
+    !!
+    !! Computes the integer, real [, complex] working space required by the least-squares solver
+    !! ([Specification](../page/specs/stdlib_linalg.html))
+    !! 
+    !!### Description
+    !! 
+    !! This interface provides sizes of integer, real [, complex] working spaces required by the 
+    !! least-squares solver. These sizes can be used to pre-allocated working arrays in case several 
+    !! repeated least-squares solutions to a same system are sought. If pre-allocated working arrays 
+    !! are provided, no internal allocations will take place.
+    !! 
+    #:for nd,ndsuf,nde in ALL_RHS
+    #:for rk,rt,ri in RC_KINDS_TYPES
+    #:if rk!="xdp" 
+      pure module subroutine stdlib_linalg_${ri}$_lstsq_space_${ndsuf}$(a,b,lrwork,liwork#{if rt.startswith('c')}#,lcwork#{endif}#)
+         !> Input matrix a[m,n]
+         ${rt}$, intent(inout), target :: a(:,:)
+         !> Right hand side vector or array, b[n] or b[n,nrhs]
+         ${rt}$, intent(in) :: b${nd}$
+         !> Size of the working space arrays                
+         integer(ilp), intent(out) :: lrwork,liwork#{if rt.startswith('c')}#,lcwork#{endif}#         
+      end subroutine stdlib_linalg_${ri}$_lstsq_space_${ndsuf}$
+    #:endif
+    #:endfor
+    #:endfor
+  end interface lstsq_space
+
   interface det
     !! version: experimental 
     !!
diff --git a/src/stdlib_linalg_least_squares.fypp b/src/stdlib_linalg_least_squares.fypp
@@ -31,6 +31,7 @@ submodule (stdlib_linalg) stdlib_linalg_least_squares
                  err = linalg_state_type(this,LINALG_ERROR,'SVD did not converge.')
              case default
                  err = linalg_state_type(this,LINALG_INTERNAL_ERROR,'catastrophic error')
+          end select
 
       end subroutine handle_gelsd_info
 
@@ -80,6 +81,26 @@ submodule (stdlib_linalg) stdlib_linalg_least_squares
      #:for rk,rt,ri in RC_KINDS_TYPES
      #:if rk!="xdp"
 
+     ! Compute the integer, real, [complex] working space requested byu the least squares procedure
+     pure module subroutine stdlib_linalg_${ri}$_lstsq_space_${ndsuf}$(a,b,lrwork,liwork#{if rt.startswith('c')}#,lcwork#{endif}#)
+         !> Input matrix a[m,n]
+         ${rt}$, intent(inout), target :: a(:,:)
+         !> Right hand side vector or array, b[n] or b[n,nrhs]
+         ${rt}$, intent(in) :: b${nd}$
+         !> Size of the working space arrays       
+         integer(ilp), intent(out) :: lrwork,liwork
+         integer(ilp) #{if rt.startswith('c')}#, intent(out)#{endif}# :: lcwork
+         
+         integer(ilp) :: m,n,nrhs
+         
+         m    = size(a,1,kind=ilp)
+         n    = size(a,2,kind=ilp)
+         nrhs = size(b,kind=ilp)/size(b,1,kind=ilp)
+                  
+         call ${ri}$gelsd_space(m,n,nrhs,lrwork,liwork,lcwork)
+             
+     end subroutine stdlib_linalg_${ri}$_lstsq_space_${ndsuf}$
+
      ! Compute the least-squares solution to a real system of linear equations Ax = b
      module function stdlib_linalg_${ri}$_lstsq_${ndsuf}$(a,b,cond,overwrite_a,rank,err) result(x)
      !!### Summary
@@ -98,7 +119,7 @@ submodule (stdlib_linalg) stdlib_linalg_least_squares
      !! param: err [optional] State return flag. 
      !! return: x Solution vector of size [n] or solution matrix of size [n,nrhs].
      !!
-         !> Input matrix a[n,n]
+         !> Input matrix a[m,n]
          ${rt}$, intent(inout), target :: a(:,:)
          !> Right hand side vector or array, b[n] or b[n,nrhs]
          ${rt}$, intent(in) :: b${nd}$
@@ -122,8 +143,8 @@ submodule (stdlib_linalg) stdlib_linalg_least_squares
      end function stdlib_linalg_${ri}$_lstsq_${ndsuf}$
 
      ! Compute the least-squares solution to a real system of linear equations Ax = b
-     module subroutine stdlib_linalg_${ri}$_solve_lstsq_${ndsuf}$(a,b,x, &
-             real_storage,int_storage#{if rt.startswith('c')}#,cmpl_storage#{endif}#,cond,overwrite_a,rank,err) 
+     module subroutine stdlib_linalg_${ri}$_solve_lstsq_${ndsuf}$(a,b,x,real_storage,int_storage, &
+             #{if rt.startswith('c')}#cmpl_storage,#{endif}# cond,singvals,overwrite_a,rank,err) 
              
      !!### Summary
      !! Compute least-squares solution to a real system of linear equations \( Ax = b \) 
@@ -134,12 +155,18 @@ submodule (stdlib_linalg) stdlib_linalg_least_squares
      !!
      !! param: a Input matrix of size [m,n].
      !! param: b Right-hand-side vector of size [n] or matrix of size [n,nrhs].
+     !! param: x Solution vector of size [n] or solution matrix of size [n,nrhs].
+     !! param: real_storage [optional] Real working space
+     !! param: int_storage [optional] Integer working space
+     #:if rt.startswith('c')   
+     !! param: cmpl_storage [optional] Complex working space
+     #:endif
      !! param: cond [optional] Real input threshold indicating that singular values `s_i <= cond*maxval(s)` 
      !!        do not contribute to the matrix rank.
+     !! param: singvals [optional] Real array of size [min(m,n)] returning a list of singular values.
      !! param: overwrite_a [optional] Flag indicating if the input matrix can be overwritten.
      !! param: rank [optional] integer flag returning matrix rank. 
-     !! param: err [optional] State return flag. 
-     !! return: x Solution vector of size [n] or solution matrix of size [n,nrhs].
+     !! param: err [optional] State return flag.      
      !!
          !> Input matrix a[n,n]
          ${rt}$, intent(inout), target :: a(:,:)
@@ -157,6 +184,8 @@ submodule (stdlib_linalg) stdlib_linalg_least_squares
          #:endif
          !> [optional] cutoff for rank evaluation: singular values s(i)<=cond*maxval(s) are considered 0.
          real(${rk}$), optional, intent(in) :: cond
+         !> [optional] list of singular values [min(m,n)], in descending magnitude order, returned by the SVD
+         real(${rk}$), optional, intent(out), target :: singvals(:)         
          !> [optional] Can A,b data be overwritten and destroyed?
          logical(lk), optional, intent(in) :: overwrite_a
          !> [optional] Return rank of A
@@ -167,12 +196,11 @@ submodule (stdlib_linalg) stdlib_linalg_least_squares
          !! Local variables
          type(linalg_state_type) :: err0
          integer(ilp) :: m,n,lda,ldb,nrhs,ldx,nrhsx,info,mnmin,mnmax,arank,lrwork,liwork,lcwork
-         integer(ilp) :: nrs,nis,ncs
+         integer(ilp) :: nrs,nis,ncs,nsvd
          integer(ilp), pointer :: iwork(:)
          logical(lk) :: copy_a
          real(${rk}$) :: acond,rcond
-         real(${rk}$), allocatable :: singular(:)
-         real(${rk}$), pointer :: rwork(:)
+         real(${rk}$), pointer :: rwork(:),singular(:)
          ${rt}$, pointer :: xmat(:,:),amat(:,:),cwork(:)
 
          ! Problem sizes
@@ -213,8 +241,14 @@ submodule (stdlib_linalg) stdlib_linalg_least_squares
          xmat(1:n,1:nrhs) => x
 
          ! Singular values array (in decreasing order)
-         allocate(singular(mnmin))
-
+         if (present(singvals)) then 
+            singular => singvals            
+            nsvd = size(singular,kind=ilp)
+         else
+            allocate(singular(mnmin))
+            nsvd = mnmin
+         endif
+         
          ! rcond is used to determine the effective rank of A.
          ! Singular values S(i) <= RCOND*maxval(S) are treated as zero.
          ! Use same default value as NumPy
@@ -254,12 +288,14 @@ submodule (stdlib_linalg) stdlib_linalg_least_squares
          ncs = size(iwork,kind=ilp)
          #:endif         
                     
-         if (nrs<lrwork .or. nis<liwork#{if rt.startswith('c')}# .or. ncs<lcwork#{endif}#) then 
+         if (nrs<lrwork .or. nis<liwork#{if rt.startswith('c')}# .or. ncs<lcwork#{endif}# &
+             .or. nsvd<mnmin) then 
             ! Halt on insufficient space
             err0 = linalg_state_type(this,LINALG_VALUE_ERROR,'insufficient working space: ',&
                          'real=',nrs,' should be >=',lrwork, &
-                        ', int=',nis,' should be >=',liwork  &
-                        #{if rt.startswith('complex')}#,', cmplx=',ncs,' should be >=',lcwork#{endif}#)
+                        ', int=',nis,' should be >=',liwork, &
+                        #{if rt.startswith('complex')}#', cmplx=',ncs,' should be >=',lcwork, &#{endif}# 
+                      ', singv=',nsvd,' should be >=',mnmin)
          
          else
 
@@ -270,23 +306,24 @@ submodule (stdlib_linalg) stdlib_linalg_least_squares
             #:else
                        rwork,nrs,iwork,info)
             #:endif
+
+            ! The condition number of A in the 2-norm = S(1)/S(min(m,n)).
+            acond = singular(1)/singular(mnmin)
+
+            ! Process output
+            call handle_gelsd_info(info,lda,n,ldb,nrhs,err0)
          
          endif           
                   
-         ! The condition number of A in the 2-norm = S(1)/S(min(m,n)).
-         acond = singular(1)/singular(mnmin)
-
-         ! Process output
-         call handle_gelsd_info(info,lda,n,ldb,nrhs,err0)
-
          ! Process output and return
-         1 if (copy_a) deallocate(amat)         
+         if (copy_a) deallocate(amat)         
          if (present(rank)) rank = arank
          if (.not.present(real_storage)) deallocate(rwork)
          if (.not.present(int_storage))  deallocate(iwork)
          #:if rt.startswith('complex') 
          if (.not.present(cmpl_storage)) deallocate(cwork)
          #:endif
+         if (.not.present(singvals)) deallocate(singular)
          call linalg_error_handling(err0,err)
 
      end subroutine stdlib_linalg_${ri}$_solve_lstsq_${ndsuf}$
