Update max-w-in-R algorithm.

1 files changed, 7 insertions, 13 deletions

diff --git a/services/nnls.rkt b/services/nnls.rkt
index 2c5738f..898072c 100644
--- a/services/nnls.rkt
+++ b/services/nnls.rkt
@@ -34,7 +34,7 @@
   (lawson-hanson-1974 fertilizer-product-matrix deficits error-threshold))
 
 ;; Algorithm lifted from the Wikipedia article on NNLS
-(define/contract (lawson-hanson-1974 A y ε)
+(define (lawson-hanson-1974 A y ε)
   ;;;;;;;;;
   ;; Inputs
   ;;;;;;;;;
@@ -63,15 +63,13 @@
   (define (compute-w x)
     (matrix* (matrix-transpose A) (matrix- y (matrix* A x))))
 
-  ;; max over j in R of w_j, returning (values max-val j*)
+  ;; max over j in R of w_j.
   (define (max-w-in-R w)
-    (for/fold ([max-val -inf.0]
-               [max-j #f])
-              ([j (in-set R)])
-      (define v (colv-ref w j))
-      (if (> v max-val)
-          (values v j)
-          (values max-val max-j))))
+    (if (set-empty? R)
+        (values -inf.0 #f)
+        (let ([max-j (argmax (λ (j) (colv-ref w j)) (set->list R))])
+          (values (colv-ref w max-j) max-j))))
+
 
   ;; Build full candidate vector s from current P:
   ;; s_P = (A_Pᵀ A_P)⁻¹ A_Pᵀ y, s_R = 0
@@ -124,20 +122,16 @@
           ;; Inner loop: adjust until s_P > 0
           (let inner-loop ()
             (define s (make-s-from-P))
-
-            ;; min(s_P)
             (define min-sP
               (if (set-empty? P)
                   +inf.0
                   (for/fold ([mn +inf.0]) ([j (in-set P)])
                     (min mn (colv-ref s j)))))
-
             (cond
               ;; If all s_P > 0 (or P empty), accept s as new x and go back to outer loop
               [(or (set-empty? P) (> min-sP 0))
                (set! x s)
                (outer-loop)]
-
               [else
                ;; Compute α = min_{i in P, s_i <= 0} x_i / (x_i - s_i)
                (define α