add ordinary least square with lapack backend

Author: kchanqvq

src/bindings/lapack-cffi.lisp

@@ -28,6 +28,8 @@
                               matrix-class type (lapack-routine "getrf"))
                              (generate-lapack-lu-solve-for-type
                               matrix-class type (lapack-routine "getrs"))
+                             (generate-lapack-lsd-for-type
+                              matrix-class type (lapack-routine "gelsd"))
                              (generate-lapack-inv-for-type
                               matrix-class type
                               (lapack-routine "getrf") (lapack-routine "getri"))

src/extensions/lapack/lapack-bindings.lisp

@@ -17,6 +17,8 @@
 
 (magicl:extend-function (magicl:lu-solve lapack-lu-solve :lapack) (lu ipiv b))
 
+(magicl:extend-function (magicl:lsd lapack-lsd :lapack) (a b rcond))
+
 (magicl:extend-function (magicl:inv lapack-inv :lapack) (matrix))
 
 (magicl:extend-function (magicl:csd-blocks csd-blocks-extension :lapack) (matrix p q))

src/extensions/lapack/lapack-generics.lisp

@@ -149,6 +149,81 @@
         info)
        b-tensor)))
 
+(defun generate-lapack-lsd-for-type (class type lstsq-function)
+  `(defmethod lapack-lsd ((a ,class) (b-tensor ,class) rcond)
+     (let* ((m (nrows a))
+            (n (ncols a)))
+       (policy-cond:with-expectations (> speed safety)
+           ((assertion (cl:= m (nrows b-tensor))))
+         (let* ((nrhs (ncols b-tensor))
+                (a (magicl::storage (deep-copy-tensor (if (eql :row-major (layout a)) (transpose a) a))))
+                (lda m)
+                (ldb (max m n))
+                (b (make-array (* ldb nrhs) :element-type ',type))
+                (s (make-array (min m n) :element-type ',type))
+                (rcond (or rcond
+                           (* (max m n)
+                              ,(cond ((or (eq type 'single-float)
+                                          (equal type '(complex single-float)))
+                                      'single-float-epsilon)
+                                     ((or (eq type 'double-float)
+                                          (equal type '(complex double-float)))
+                                      'double-float-epsilon)
+                                     (t (error "Unknown type for lapack-lsd: ~a" type))))))
+                (rank 0)
+                (work1 (make-array 1 :element-type ',type))
+                (work nil)
+                (lwork -1)
+                (iwork1 (make-array 1 :element-type '(signed-byte 32)))
+                (iwork nil)
+                (info 0))
+           (setf (subseq b 0)
+                 (magicl::storage
+                  (if (eql :row-major (layout b-tensor)) (transpose b-tensor) b-tensor)))
+           ;; Perform work size query with work of length 1
+           (,lstsq-function
+            m
+            n
+            nrhs
+            a
+            lda
+            b
+            ldb
+            s
+            rcond
+            rank
+            work1
+            lwork
+            iwork1
+            info)
+           (setf lwork (round (realpart (aref work1 0))))
+           (setf work (make-array (max 1 lwork) :element-type ',type))
+           (setf iwork (make-array (max 1 (aref iwork1 0)) :element-type '(signed-byte 32)))
+           ;; Perform actual operation with correct work size
+           (,lstsq-function
+            m
+            n
+            nrhs
+            a
+            lda
+            b
+            ldb
+            s
+            rcond
+            rank
+            work
+            lwork
+            iwork
+            info)
+           (let* ((sol (magicl::from-storage b (list ldb nrhs) :layout :column-major))
+                  (x (slice sol (list 0 0) (list n nrhs))))
+             (cond ((> m n)
+                    ;; r^2 is given by sum of squares of n+1:m rows in B
+                    (let ((r (slice sol (list n 0) (list m nrhs))))
+                      (values x (row-matrix->vector (@ (ones (list 1 (- m n)) :type ',type) (.* r r))))))
+                   (t
+                    (values x (zeros (list nrhs) :type ',type))))))))))
+
 (defun generate-lapack-inv-for-type (class type lu-function inv-function)
   `(defmethod lapack-inv ((a ,class))
      (let ((a-tensor (deep-copy-tensor a)))

src/extensions/lapack/lapack-templates.lisp

@@ -16,6 +16,7 @@ gges
 gees
 getrf
 getrs
+gelsd
 getri
 gesvd
 geev

src/functions.txt

@@ -567,6 +567,9 @@ So assuming P is a permutation matrix representing IPIV, we have
 (define-backend-function lu-solve (lu ipiv b)
   "Solve the system AX=B, where A is a square matrix, B is a compatibly shaped matrix, and A has PLU factorization indicated by the permutation vector IPIV and lower & upper triangular portions of the argument LU.")
 
+(define-backend-function lsd (a b &optional rcond)
+  "Solve the linear least square problem: argmin_X ||B-AX||^2 using the singular value decomposition (SVD) of A. A is a M-by-N matrix, and B is a M-by-NRHS right hand side matrix. Return two values: the solution as a N-by-NRHS solution matrix, and the sum of squared residuals as a NRHS vector. The effective rank of A is determined by treating as zero those singular values which are less than RCOND times the largest singular value.")
+
 (define-extensible-function (csd-blocks csd-blocks-lisp) (matrix p q)
   (:documentation "Compute the cosine-sine decomposition of the matrix MATRIX and return the result as blocks. See LISP-CSD-BLOCKS for mathematical details.
 
@@ -677,3 +680,11 @@ NOTE: If H is not Hermitian, the behavior is undefined.")
                  :matrix A)))
       (let ((rmat (lu-solve lu ipiv bmat)))
         (from-storage (storage rmat) (shape b))))))
+
+(defun ols (a b &optional rcond)
+  "Attempt to solve the ordinary least square problem argmin_X ||B-AX||^2. Returns X and sum of squared residuals."
+  (cond ((vector-p b)
+         (multiple-value-bind (x r) (lsd a (vector->column-matrix b) rcond)
+           (values (column-matrix->vector x) (tref r 0))))
+        (t
+         (lsd a b rcond))))

src/high-level/matrix.lisp

@@ -173,6 +173,7 @@
            #:inv
            #:lu
            #:lu-solve
+           #:lsd
            #:csd-blocks
            #:csd
            #:svd
@@ -186,6 +187,7 @@
            #:logm
            #:expih
            #:linear-solve
+           #:ols
 
            #:polynomial
            #:make-polynomial

src/packages.lisp

@@ -292,3 +292,19 @@
 
   (signals magicl::rank-deficiency-error
     (magicl:linear-solve (magicl:ones '(3 3)) (magicl:ones '(3)))))
+
+(deftest test-ols ()
+  "Check that we can solve ordinary least square."
+  (let ((A (magicl:from-list '((1d0 1d0) (2d0 1d0) (3d0 1d0) (4d0 1d0)) '(4 2)))
+        (b (magicl:from-list '(0d0 1d0 5d0 6d0) '(4))))
+    (multiple-value-bind (x r) (magicl:ols A b)
+      (is (magicl:= x (magicl:from-list '(2.2d0 -2.5d0) '(2))))
+      (is (< (abs (- r 1.8d0))) (* 1.0d2 double-float-epsilon)))
+
+    (signals error (magicl:ols A (magicl:from-list '(0d0 1d0 5d0) '(3)))))
+
+  (let ((A (magicl:from-list '((1d0 1d0) (2d0 1d0) (3d0 1d0) (4d0 1d0)) '(4 2)))
+        (b (magicl:from-list '((0d0 0d0) (1d0 1d0) (5d0 2d0) (6d0 3d0)) '(4 2))))
+    (multiple-value-bind (x r) (magicl:ols A b)
+      (is (magicl:= x (magicl:from-list '((2.2d0 1d0) (-2.5d0 -1d0)) '(2 2))))
+      (is (magicl:= r (magicl:from-list '(1.8d0 0.0d0) '(2)))))))