Pivoting QR decomposition

doc/specs/stdlib_linalg.md

@@ -953,7 +953,7 @@ the full problem is solved. On reduced matrices (`shape(Q)==[m,k]`, `shape(R)==[
 
 ### Syntax
 
-`call ` [[stdlib_linalg(module):qr(interface)]] `(a, q, r, [, storage] [, overwrite_a] [, err])`
+`call ` [[stdlib_linalg(module):qr(interface)]] `(a, q, r [, pivots] [, overwrite_a] [, storage] [, err])`
 
 ### Arguments
 
@@ -963,15 +963,17 @@ the full problem is solved. On reduced matrices (`shape(Q)==[m,k]`, `shape(R)==[
 
 `r`: Shall be a rank-2 array of the same kind as `a`, containing the upper triangular matrix `r`. It is an `intent(out)` argument. It should have a shape equal to either `[m,n]` or `[k,n]`, whether the full or the reduced problem is sought for.
 
-`storage` (optional): Shall be a rank-1 array of the same type and kind as `a`, providing working storage for the solver. Its minimum size can be determined with a call to [[stdlib_linalg(module):qr_space(interface)]]. It is an `intent(out)` argument.
+`pivots` (optional): Shall be an `integer` array of size `n`. If provided, QR factorization with column-pivoting is being computed.
 
 `overwrite_a` (optional): Shall be an input `logical` flag (default: `.false.`). If `.true.`, input matrix `a` will be used as temporary storage and overwritten. This avoids internal data allocation. It is an `intent(in)` argument.
 
+`storage` (optional): Shall be a rank-1 array of the same type and kind as `a`, providing working storage for the solver. Its minimum size can be determined with a call to [[stdlib_linalg(module):qr_space(interface)]]. It is an `intent(out)` argument.
+
 `err` (optional): Shall be a `type(linalg_state_type)` value. It is an `intent(out)` argument.
 
 ### Return value
 
-Returns the QR factorization matrices into the \( Q \) and \( R \) arguments. 
+Returns the QR factorization matrices into the \( Q \) and \( R \) arguments and the optional pivots in `pivots`.
 
 Raises `LINALG_VALUE_ERROR` if any of the matrices has invalid or unsuitable size for the full/reduced problem.
 Raises `LINALG_ERROR` on insufficient user storage space.
@@ -981,6 +983,8 @@ If the state argument `err` is not present, exceptions trigger an `error stop`.
 
 ```fortran
 {!example/linalg/example_qr.f90!}
+
+{!example/linalg/example_pivoting_qr.f90!}
 ```
 
 ## `qr_space` - Compute internal working space requirements for the QR factorization.
@@ -995,20 +999,24 @@ This subroutine computes the internal working space requirements for the QR fact
 
 ### Syntax
 
-`call ` [[stdlib_linalg(module):qr_space(interface)]] `(a, lwork, [, err])`
+`call ` [[stdlib_linalg(module):qr_space(interface)]] `(a, lwork, [, pivoting] [, err])`
 
 ### Arguments
 
 `a`: Shall be a rank-2 `real` or `complex` array containing the coefficient matrix. It is an `intent(in)` argument.
 
 `lwork`: Shall be an `integer` scalar, that returns the minimum array size required for the working storage in [[stdlib_linalg(module):qr(interface)]] to factorize `a`.
 
+`pivoting` (optional): Shall a `logical` flag (default: `.false.`). If `.true.`, on exit `lwork` is the optimal workspace size for the QR factorization with column pivoting. If `.false.`, `lwork` is the optimal workspace size for the standard QR factorization.
+
 `err` (optional): Shall be a `type(linalg_state_type)` value. This is an `intent(out)` argument.
 
 ### Example
 
 ```fortran
 {!example/linalg/example_qr_space.f90!}
+
+{!example/linalg/example_pivoting_qr_space.f90!}
 ```
 
 ## `schur` - Compute the Schur decomposition of a matrix

example/linalg/CMakeLists.txt

@@ -54,7 +54,9 @@ ADD_EXAMPLE(svdvals)
 ADD_EXAMPLE(determinant)
 ADD_EXAMPLE(determinant2)
 ADD_EXAMPLE(qr)
+ADD_EXAMPLE(pivoting_qr)
 ADD_EXAMPLE(qr_space)
+ADD_EXAMPLE(pivoting_qr_space)
 ADD_EXAMPLE(cholesky)
 ADD_EXAMPLE(chol)
 ADD_EXAMPLE(expm)

example/linalg/example_pivoting_qr.f90

@@ -0,0 +1,16 @@
+program example_pivoting_qr
+   use stdlib_linalg, only: qr
+   implicit none
+   real :: A(104, 32), Q(104, 32), R(32, 32)
+   integer :: pivots(32)
+
+   ! Create a random matrix
+   call random_number(A)
+
+   ! Compute its QR factorization (reduced)
+   call qr(A, Q, R, pivots)
+
+   ! Test factorization: Q*R = A
+   print *, maxval(abs(matmul(Q, R) - A(:, pivots)))
+
+end program example_pivoting_qr

example/linalg/example_pivoting_qr_space.f90

@@ -0,0 +1,25 @@
+! QR example with pre-allocated storage
+program example_pivoting_qr_space
+   use stdlib_linalg_constants, only: ilp
+   use stdlib_linalg, only: qr, qr_space, linalg_state_type
+   implicit none
+   real :: A(104, 32), Q(104, 32), R(32, 32)
+   real, allocatable :: work(:)
+   integer(ilp) :: lwork, pivots(32)
+   type(linalg_state_type) :: err
+
+   ! Create a random matrix
+   call random_number(A)
+
+   ! Prepare QR workspace
+   call qr_space(A, lwork, pivoting=.true.)
+   allocate (work(lwork))
+
+   ! Compute its QR factorization (reduced)
+   call qr(A, Q, R, pivots, storage=work, err=err)
+
+   ! Test factorization: Q*R = A
+   print *, maxval(abs(matmul(Q, R) - A(:, pivots)))
+   print *, err%print()
+
+end program example_pivoting_qr_space

src/lapack/stdlib_linalg_lapack_aux.fypp

@@ -46,6 +46,7 @@ module stdlib_linalg_lapack_aux
      public :: handle_gesv_info
      public :: handle_gees_info
      public :: handle_geqrf_info
+     public :: handle_geqp3_info
      public :: handle_orgqr_info
      public :: handle_gelsd_info
      public :: handle_geev_info
@@ -1462,6 +1463,27 @@ module stdlib_linalg_lapack_aux
 
      end subroutine handle_geqrf_info
 
+    elemental subroutine handle_geqp3_info(this, info, m, n, lwork, err)
+        character(len=*), intent(in) :: this
+        integer(ilp), intent(in) :: info, m, n, lwork
+        type(linalg_state_type), intent(out) :: err
+        ! Process output
+        select case (info)
+            case(0)
+                ! Success
+            case(-1)
+                err = linalg_state_type(this, LINALG_VALUE_ERROR, 'invalid matrix size m=', m)
+            case(-2)
+                err = linalg_state_type(this, LINALG_VALUE_ERROR, 'invalid matrix size n=', n)
+            case(-4)
+                err = linalg_state_type(this, LINALG_VALUE_ERROR, 'invalid matrix size a=', [m, n])
+            case(-7)
+                err = linalg_state_type(this, LINALG_VALUE_ERROR, 'invalid input for lwork=', lwork)
+            case default
+                err = linalg_state_type(this, LINALG_VALUE_ERROR, 'catastrophic error')
+        end select
+    end subroutine handle_geqp3_info
+
      elemental subroutine handle_orgqr_info(this,info,m,n,k,lwork,err)
          character(len=*), intent(in) :: this
          integer(ilp), intent(in) :: info,m,n,k,lwork

src/stdlib_linalg.fypp

@@ -614,6 +614,23 @@ module stdlib_linalg
          !> [optional] state return flag. On error if not requested, the code will stop
          type(linalg_state_type), optional, intent(out) :: err
       end subroutine stdlib_linalg_${ri}$_qr
+
+      pure module subroutine stdlib_linalg_${ri}$_pivoting_qr(a, q, r, pivots, overwrite_a, storage, err)
+          !> Input matrix a[m, n]
+          ${rt}$, intent(inout), target :: a(:, :)
+          !> Orthogonal matrix Q ([m, m] or [m, k] if reduced)
+          ${rt}$, intent(out), contiguous, target :: q(:, :)
+          !> Upper triangular matrix R ([m, n] or [k, n] if reduced)
+          ${rt}$, intent(out), contiguous, target :: r(:, :)
+          !> Pivots.
+          integer(ilp), intent(out) :: pivots(:)
+          !> [optional] Can A data be overwritten and destroyed?
+          logical(lk), optional, intent(in) :: overwrite_a
+          !> [optional] Provide pre-allocated workspace, size to be checked with pivoting_qr_space.
+          ${rt}$, intent(out), optional, target :: storage(:)
+          !> [optional] state return flag. On error if not requested, the code will stop.
+          type(linalg_state_type), optional, intent(out) :: err
+      end subroutine stdlib_linalg_${ri}$_pivoting_qr
     #:endfor    
   end interface qr
 
@@ -641,6 +658,17 @@ module stdlib_linalg
          !> State return flag. Returns an error if the query failed
          type(linalg_state_type), optional, intent(out) :: err
       end subroutine get_qr_${ri}$_workspace
+
+      pure module subroutine get_pivoting_qr_${ri}$_workspace(a, lwork, pivoting, err)
+        !> Input matrix a[m, n]
+        ${rt}$, intent(in), target :: a(:, :)
+        !> Minimum workspace size for both operations.
+        integer(ilp), intent(out) :: lwork
+        !> Pivoting flag.
+        logical(lk), intent(in) :: pivoting
+        !> State return flag. Returns an error if the query failed.
+        type(linalg_state_type), optional, intent(out) :: err
+      end subroutine get_pivoting_qr_${ri}$_workspace
     #:endfor
   end interface qr_space
 

src/stdlib_linalg_lapack.fypp

@@ -3230,6 +3230,77 @@ module stdlib_linalg_lapack
 #:endfor
           end interface geqrt3
 
+          interface geqp3
+          !! GEQP3 computes a QR factorization with column pivoting of a real or complex
+          !! M-by-N matrix A:
+          !!
+          !!    A * P = Q * R,
+          !!
+          !! where:
+          !! Q is an M-by-min(M, N) orthogonal matrix
+          !! R is an min(M, N)-by-N upper triangular matrix;
+#:for ik, it, ii in LINALG_INT_KINDS_TYPES
+#ifdef STDLIB_EXTERNAL_LAPACK${ii}$
+            pure subroutine sgeqp3(m, n, a, lda, jpvt, tau, work, lwork, info)
+                import sp, dp, qp, ${ik}$, lk
+                implicit none
+                integer(${ik}$), intent(in) :: m, n, lda, lwork
+                integer(${ik}$), intent(out) :: info
+                integer(${ik}$), intent(inout) :: jpvt(*)
+                real(sp), intent(inout) :: a(lda, *)
+                real(sp), intent(out) :: tau(*), work(*)
+            end subroutine sgeqp3
+
+            pure subroutine dgeqp3(m, n, a, lda, jpvt, tau, work, lwork, info)
+                import sp, dp, qp, ${ik}$, lk
+                implicit none
+                integer(${ik}$), intent(in) :: m, n, lda, lwork
+                integer(${ik}$), intent(out) :: info
+                integer(${ik}$), intent(inout) :: jpvt(*)
+                real(dp), intent(inout) :: a(lda, *)
+                real(dp), intent(out) :: tau(*), work(*)
+            end subroutine dgeqp3
+
+            pure subroutine cgeqp3(m, n, a, lda, jpvt, tau, work, lwork, rwork, info)
+                import sp, dp, qp, ${ik}$, lk
+                implicit none
+                integer(${ik}$), intent(in) :: m, n, lda, lwork
+                integer(${ik}$), intent(out) :: info
+                integer(${ik}$), intent(inout) :: jpvt(*)
+                complex(sp), intent(inout) :: a(lda, *)
+                complex(sp), intent(out) :: tau(*), work(*)
+                real(sp), intent(out) :: rwork(*)
+            end subroutine cgeqp3
+
+            pure subroutine zgeqp3(m, n, a, lda, jpvt, tau, work, lwork, rwork, info)
+                import sp, dp, qp, ${ik}$, lk
+                implicit none
+                integer(${ik}$), intent(in) :: m, n, lda, lwork
+                integer(${ik}$), intent(out) :: info
+                integer(${ik}$), intent(inout) :: jpvt(*)
+                complex(dp), intent(inout) :: a(lda, *)
+                complex(dp), intent(out) :: tau(*), work(*)
+                real(dp), intent(out) :: rwork(*)
+            end subroutine zgeqp3
+#else
+            module procedure stdlib${ii}$_sgeqp3
+            module procedure stdlib${ii}$_dgeqp3
+            module procedure stdlib${ii}$_cgeqp3
+            module procedure stdlib${ii}$_zgeqp3
+#endif
+#:for rk, rt, ri in REAL_KINDS_TYPES
+#:if not rk in ["sp", "dp"]
+            module procedure stdlib${ii}$_${ri}$geqp3
+#:endif
+#:endfor
+#:for rk, rt, ri in CMPLX_KINDS_TYPES
+#:if not rk in ["sp", "dp"]
+            module procedure stdlib${ii}$_${ri}$geqp3
+#:endif
+#:endfor
+#:endfor
+          end interface geqp3
+
           interface gerfs
           !! GERFS improves the computed solution to a system of linear
           !! equations and provides error bounds and backward error estimates for