Add lascl2, fix weighted_lstsq error propagation

Author: aamrindersingh

src/core/stdlib_error.fypp

@@ -81,6 +81,9 @@ module stdlib_error
 
            !> Handle optional error message 
            procedure :: handle    => error_handling
+           
+           !> Update the location of the error message
+           procedure :: update_location => state_update_location
 
     end type state_type
 
@@ -252,6 +255,15 @@ contains
 
     end subroutine error_handling
 
+    !> Update the location of the error message
+    pure subroutine state_update_location(this, where_at)
+         class(state_type), intent(inout) :: this
+         character(len=*), intent(in) :: where_at
+         
+         if (len_trim(where_at) > 0) this%where_at = adjustl(where_at)
+         
+    end subroutine state_update_location
+
     !> Produce a nice error string
     pure function state_print(this) result(msg)
          class(state_type),intent(in) :: this

src/lapack/stdlib_lapack_base.fypp

@@ -1310,6 +1310,27 @@ interface
 #:endfor
 end interface 
 
+! LASCL2 performs diagonal scaling: X(i,:) = D(i) * X(i,:)
+interface 
+#:for ik,it,ii in LINALG_INT_KINDS_TYPES
+#:for rk,rt,ri in REAL_KINDS_TYPES
+     pure module subroutine stdlib${ii}$_${ri}$lascl2( m, n, d, x, ldx )
+           integer(${ik}$), intent(in) :: m, n, ldx
+           real(${rk}$), intent(in) :: d(*)
+           real(${rk}$), intent(inout) :: x(ldx,*)
+     end subroutine stdlib${ii}$_${ri}$lascl2
+
+#:endfor
+#:for ck,ct,ci in CMPLX_KINDS_TYPES
+     pure module subroutine stdlib${ii}$_${ci}$lascl2( m, n, d, x, ldx )
+           integer(${ik}$), intent(in) :: m, n, ldx
+           real(${ck}$), intent(in) :: d(*)
+           complex(${ck}$), intent(inout) :: x(ldx,*)
+     end subroutine stdlib${ii}$_${ci}$lascl2
+
+#:endfor
+#:endfor
+end interface 
 
 interface 
 #:for ik,it,ii in LINALG_INT_KINDS_TYPES

src/lapack/stdlib_lapack_blas_like_l2.fypp

@@ -5681,4 +5681,55 @@ submodule(stdlib_lapack_base) stdlib_lapack_blas_like_l2
 
 
 #:endfor
+
+! LASCL2 performs diagonal scaling on a matrix: X(i,:) = D(i) * X(i,:)
+#:for ik,it,ii in LINALG_INT_KINDS_TYPES
+#:for rk,rt,ri in REAL_KINDS_TYPES
+     pure module subroutine stdlib${ii}$_${ri}$lascl2( m, n, d, x, ldx )
+        !! LASCL2 performs a diagonal scaling on a matrix:
+        !! X(i,j) = D(i) * X(i,j), for i = 1,...,M and j = 1,...,N.
+        !! D is a vector of length M containing the diagonal scaling factors.
+           ! Scalar Arguments
+           integer(${ik}$), intent(in) :: m, n, ldx
+           ! Array Arguments
+           real(${rk}$), intent(in) :: d(*)
+           real(${rk}$), intent(inout) :: x(ldx,*)
+        ! =====================================================================
+           ! Local Scalars
+           integer(${ik}$) :: i, j
+           ! Executable Statements
+           do j = 1, n
+              do i = 1, m
+                 x(i,j) = d(i) * x(i,j)
+              end do
+           end do
+           return
+     end subroutine stdlib${ii}$_${ri}$lascl2
+
+#:endfor
+#:for ck,ct,ci in CMPLX_KINDS_TYPES
+     pure module subroutine stdlib${ii}$_${ci}$lascl2( m, n, d, x, ldx )
+        !! LASCL2 performs a diagonal scaling on a matrix:
+        !! X(i,j) = D(i) * X(i,j), for i = 1,...,M and j = 1,...,N.
+        !! D is a vector of length M containing the diagonal scaling factors.
+           ! Scalar Arguments
+           integer(${ik}$), intent(in) :: m, n, ldx
+           ! Array Arguments
+           real(${ck}$), intent(in) :: d(*)
+           complex(${ck}$), intent(inout) :: x(ldx,*)
+        ! =====================================================================
+           ! Local Scalars
+           integer(${ik}$) :: i, j
+           ! Executable Statements
+           do j = 1, n
+              do i = 1, m
+                 x(i,j) = d(i) * x(i,j)
+              end do
+           end do
+           return
+     end subroutine stdlib${ii}$_${ci}$lascl2
+
+#:endfor
+#:endfor
+
 end submodule stdlib_lapack_blas_like_l2

src/lapack/stdlib_linalg_lapack.fypp

@@ -16228,6 +16228,20 @@ module stdlib_linalg_lapack
 #:endfor
           end interface lascl
 
+          interface lascl2
+          !! LASCL2 performs a diagonal scaling on a matrix:
+          !! X(i,j) = D(i) * X(i,j), for i = 1,...,M and j = 1,...,N.
+          !! D is a vector of length M containing the diagonal scaling factors.
+#:for ik,it,ii in LINALG_INT_KINDS_TYPES
+#:for rk,rt,ri in REAL_KINDS_TYPES
+               module procedure stdlib${ii}$_${ri}$lascl2
+#:endfor
+#:for ck,ct,ci in CMPLX_KINDS_TYPES
+               module procedure stdlib${ii}$_${ci}$lascl2
+#:endfor
+#:endfor
+          end interface lascl2
+
           interface lasd0
           !! Using a divide and conquer approach, LASD0: computes the singular
           !! value decomposition (SVD) of a real upper bidiagonal N-by-M

src/linalg/stdlib_linalg_least_squares.fypp

@@ -7,7 +7,7 @@
 submodule (stdlib_linalg) stdlib_linalg_least_squares
 !! Least-squares solution to Ax=b
      use stdlib_linalg_constants
-     use stdlib_linalg_lapack, only: gelsd, gglse, stdlib_ilaenv
+     use stdlib_linalg_lapack, only: gelsd, gglse, stdlib_ilaenv, lascl2
      use stdlib_linalg_lapack_aux, only: handle_gelsd_info, handle_gglse_info
      use stdlib_linalg_state, only: linalg_state_type, linalg_error_handling, LINALG_ERROR, &
          LINALG_INTERNAL_ERROR, LINALG_VALUE_ERROR
@@ -590,7 +590,7 @@ submodule (stdlib_linalg) stdlib_linalg_least_squares
 
         ! Local variables
         type(linalg_state_type) :: err0
-        integer(ilp) :: m, n, j
+        integer(ilp) :: m, n
         logical(lk) :: copy_a
         ${rt}$, pointer :: amat(:,:)
         ${rt}$, allocatable, target :: amat_alloc(:,:)
@@ -644,22 +644,23 @@ submodule (stdlib_linalg) stdlib_linalg_least_squares
 
         ! Handle A matrix: either copy or use original
         if (copy_a) then
-            allocate(amat_alloc(m, n))
+            allocate(amat_alloc(m, n), source=a)
             amat => amat_alloc
         else
             amat => a
         end if
 
-        ! Scale A column-wise (cache-friendly: column-major order)
-        do j = 1, n
-            amat(:, j) = sqrt_w(:) * a(:, j)
-        end do
+        ! Scale rows of A by sqrt(w) using LAPACK's lascl2
+        call lascl2(m, n, sqrt_w, amat, m)
 
         ! Scale b
         b_scaled = sqrt_w * b
 
-        ! Solve transformed OLS problem
-        call stdlib_linalg_${ri}$_solve_lstsq_one(amat, b_scaled, x, cond=cond, overwrite_a=.true., rank=rank, err=err)
+        ! Solve transformed OLS problem using local error state
+        call stdlib_linalg_${ri}$_solve_lstsq_one(amat, b_scaled, x, cond=cond, overwrite_a=.true., rank=rank, err=err0)
+        
+        ! Propagate error with updated location
+        call linalg_error_handling(err0, err, where_at=this)
 
         ! Cleanup
         if (copy_a) deallocate(amat_alloc)

src/linalg_core/stdlib_linalg_state.fypp

@@ -73,17 +73,25 @@ module stdlib_linalg_state
      end function linalg_message
 
      !> Flow control: on output flag present, return it; otherwise, halt on error
-     pure subroutine linalg_error_handling(ierr,ierr_out)
+     !> If where_at is provided, update the error location before returning/stopping
+     pure subroutine linalg_error_handling(ierr,ierr_out,where_at)
          type(linalg_state_type),intent(in) :: ierr
          type(linalg_state_type),optional,intent(out) :: ierr_out
+         character(len=*),optional,intent(in) :: where_at
 
          character(len=:),allocatable :: err_msg
 
          if (present(ierr_out)) then
              ! Return error flag
              ierr_out = ierr
+             ! Update location if requested
+             if (present(where_at)) call ierr_out%update_location(where_at)
          elseif (ierr%error()) then
-             err_msg = ierr%print()
+             if (present(where_at)) then
+                 err_msg = '['//trim(where_at)//'] returned '//ierr%print_msg()
+             else
+                 err_msg = ierr%print()
+             end if
              error stop err_msg
          end if