path: root/services/nnls.rkt

#lang racket

(provide find-ferti-recipe)

(require math/array
         math/matrix
         "../models/nutrient.rkt"
         "../models/nutrient-measurement.rkt"
         "../models/nutrient-target.rkt"
         "../models/fertilizer-product.rkt")

(define (find-ferti-recipe)
  (define fertilizers (get-fertilizer-products))
  (define solution-array (solve-nnls fertilizers))
  (for/list ([fertilizer (in-list fertilizers)]
             [quantity (in-array solution-array)])
    (cons fertilizer quantity)))

(define (solve-nnls fertilizers)
  (define nutrients (get-nutrients))
  (define fertilizer-product-matrix (get-fertilizer-product-matrix nutrients fertilizers))
  (define deficits
    (->col-matrix (for/list ([n nutrients])
                    (define latest-measurement (get-latest-nutrient-measurement-value n))
                    (define latest-target (get-latest-nutrient-target-value n))
                    (define deficit
                      (cond
                        [(false? latest-target) 0]
                        [(or (false? latest-measurement) (zero? latest-measurement)) latest-target]
                        [(and (number? latest-measurement) (number? latest-target))
                         (* 100 (/ (- latest-target latest-measurement) latest-measurement))]))
                    deficit)))
  (define error-threshold 10e-4)
  (lawson-hanson-1974 fertilizer-product-matrix deficits error-threshold))

;; Algorithm lifted from the Wikipedia article on NNLS
(define (lawson-hanson-1974 A y ε)
  ;;;;;;;;;
  ;; Inputs
  ;;;;;;;;;

  (-> matrix? ; Real-valued matrix A of dimension m × n
      col-matrix? ; Real-valued column matrix (vector) y of dimension m
      real? ; Real-value ɛ, tolerance for the stopping criterion
      col-matrix?) ; Real-valued solution column matrix x
  (define-values (m ; Number of nutrients
                  n) ; Number of fertilizer products
    (matrix-shape A))

  ;;;;;;;;;;;;;
  ;; Initialize
  ;;;;;;;;;;;;;

  ;; The passive set P is initially empty.
  (define P (mutable-set))
  ;; The active set R initially contains the indexes to the nutrients allowed to be...
  (define R (list->mutable-set (range n)))

  (define (colv-ref v i)
    (matrix-ref v i 0))

  ;; Gradient-like vector for residual error.
  (define (compute-w x)
    (matrix* (matrix-transpose A) (matrix- y (matrix* A x))))

  ;; max over j in R of w_j.
  (define (max-w-in-R w)
    (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
  (define (make-s-from-P)
    (if (set-empty? P)
        (make-matrix n 1 0)
        (let* ([idxs (sort (set->list P) <)]
               [AP (submatrix A (::) idxs)]
               [sP (matrix* (matrix-inverse (matrix* (matrix-transpose AP) AP))
                            (matrix-transpose AP)
                            y)])
          ;; map: column index j in P -> corresponding sP entry
          (define mapping
            (for/list ([j idxs]
                       [k (in-naturals)])
              (cons j (colv-ref sP k))))
          (define (s-at i)
            (define p (assoc i mapping))
            (if p
                (cdr p)
                0))
          (build-matrix n 1 (λ (i j) (s-at i))))))

  ;; The "first try" x represents no addition of any fertilizer.
  (define x (make-matrix n 1 0))

  ;;;;;;;;;;;;;
  ;; Outer loop
  ;;;;;;;;;;;;;

  (let outer-loop ()
    (define w (compute-w x))

    (cond
      ;; If no remaining candidates in R, we're done.
      [(set-empty? R) x]

      [else
       (define-values (max-val j*) (max-w-in-R w))

       ;; Stopping criterion: max(w_R) <= ε
       (cond
         [(or (not j*) (<= max-val ε)) x]

         [else
          ;; Add j* to P, remove from R
          (set-remove! R j*)
          (set-add! P j*)

          ;; Inner loop: adjust until s_P > 0
          (let inner-loop ()
            (define s (make-s-from-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 α
                 (for/fold ([a +inf.0]) ([j (in-set P)])
                   (define sj (colv-ref s j))
                   (if (<= sj 0)
                       (let* ([xj (colv-ref x j)]
                              [den (- xj sj)])
                         (if (> den 0)
                             (min a (/ xj den))
                             a))
                       a)))

               (when (or (equal? α +inf.0) (<= α 0))
                 (error 'lawson-hanson-1974 "no valid α in inner loop"))

               ;; x ← x + α (s − x)
               (define new-x (matrix+ x (matrix-scale (matrix- s x) α)))

               ;; Move to R all indices j in P with x_j <= 0
               (define to-remove '())
               (for ([j (in-set P)])
                 (when (<= (colv-ref new-x j) 0)
                   (set! to-remove (cons j to-remove))))
               (for ([j to-remove])
                 (set-remove! P j)
                 (set-add! R j))

               (set! x new-x)
               (inner-loop)]))])])))

(define (get-fertilizer-product-matrix nutrients fertilizers)
  ;; Lines are nutrients, columns are fertilizers
  (build-matrix (length nutrients)
                (length fertilizers)
                (λ (i j)
                  (define selected-nutrient (list-ref nutrients i))
                  (define product (list-ref fertilizers j))
                  (hash-ref (fertilizer-product-values product) selected-nutrient 0))))

(module+ test
  (require rackunit
           rackunit/text-ui
           "../db/conn.rkt"
           "../db/migrations.rkt")

  (define test-date "2025-01-01")

  (run-tests (test-suite "Nutrient measurement model"
               #:before (λ ()
                          (connect! #:path 'memory)
                          ;; (connect! #:path "test.sqlite3")
                          (migrate-all!)

                          (define nitrogen (create-nutrient! "Nitrogen" "N"))
                          (define phosphorus (create-nutrient! "Phosphorus" "P"))

                          (create-nutrient-measurement! test-date (hash nitrogen 0 phosphorus 0))
                          (create-nutrient-target! test-date (hash nitrogen 100 phosphorus 50))

                          (create-fertilizer-product! "Nitrogen" "King Nitrogen" (hash nitrogen 100))
                          (create-fertilizer-product! "Phosphorus"
                                                      "Phosphorescent Baboon"
                                                      (hash nitrogen 10 phosphorus 100))
                          (create-fertilizer-product! "Diluted phosphorus"
                                                      "John's Phosphorus"
                                                      (hash nitrogen 3 phosphorus 30)))
               #:after (λ () (disconnect!))

               (test-case "Solve for NNLS"
                 (displayln (format "Final solution for fertilizers is combination ~a"
                                    (find-ferti-recipe)))))))