path: root/services/nnls.rkt

#lang racket

(provide find-ferti-recipe)

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

(define (find-ferti-recipe date-string)
  (define fertilizers (get-fertilizer-products))
  (define measurement-values
    (nutrient-measurement-values
     (or (get-nutrient-measurement #:date date-string)
         (error 'nnls
                "Can't compute the Ferti recipe (missing nutrient measurement for ~a)"
                date-string))))
  (define rotation-values
    (average-crop-requirement-nutrient-values
     (crop-rotation-requirements
      (or
       (get-crop-rotation #:date date-string)
       (error 'nnls "Can't compute the Ferti recipe (missing crop rotation for ~a)" date-string)))))
  (define solution-array (solve-nnls fertilizers measurement-values rotation-values))
  (for/hash ([fertilizer (in-list fertilizers)]
             [quantity (in-array solution-array)])
    (values fertilizer quantity)))

(define (solve-nnls fertilizers measurement-values rotation-values)
  (define nutrients (get-nutrients))
  (define fertilizer-product-matrix (get-fertilizer-product-matrix nutrients fertilizers))
  (define deficits
    (->col-matrix (for/list ([n nutrients])
                    (define measured (hash-ref measurement-values n 0))
                    (define required (hash-ref rotation-values n 0))
                    (define deficit
                      (if (zero? required)
                          0
                          (* 100 (/ (- required measured) required))))
                    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")

  (run-tests
   (test-suite "NNLS"
     #:before (λ ()
                (connect! #:path 'memory)
                (migrate-all!)
                (create-nutrient! "Nitrogen" "Azote" "N")
                (create-nutrient! "Phosphorus" "Phosphore" "P")
                (create-nutrient! "Potassium" "Potassium" "K"))
     #:after (λ () (disconnect!))

     (test-case "Build fertilizer product matrix"
       (define n (get-nutrient #:name "Nitrogen"))
       (define p (get-nutrient #:name "Phosphorus"))

       (define nutrients (list n p))

       (define f1 (create-fertilizer-product! "F1" "F1" (hash n 10 p 20)))
       (define f2 (create-fertilizer-product! "F2" "F2" (hash n 30 p 5)))
       (define fertilizers (list f1 f2))

       (define matrix (get-fertilizer-product-matrix nutrients fertilizers))

       (check-= (matrix-ref matrix 0 0) 10 0 "N1 in F1")
       (check-= (matrix-ref matrix 0 1) 30 0 "N1 in F2")
       (check-= (matrix-ref matrix 1 0) 20 0 "N2 in F1")
       (check-= (matrix-ref matrix 1 1) 5 0 "N2 in F2"))

     (test-case "Single nutrient, single fertilizer"
       (define A (matrix [[2]]))
       (define y (col-matrix [10]))
       (define ε 1e-6)

       (define result (lawson-hanson-1974 A y ε))

       (check-= (matrix-ref result 0 0) 5.0 ε "Should give x = 5 since 2*5 = 10"))

     (test-case "Two variables, known solution"
       (define A (matrix [[1 0] [0 1]]))
       (define y (col-matrix [3 4]))
       (define ε 1e-6)

       (define result (lawson-hanson-1974 A y ε))

       (check-= (matrix-ref result 0 0) 3.0 ε "x1 should be 3")
       (check-= (matrix-ref result 1 0) 4.0 ε "x2 should be 4"))

     (test-case "Overdetermined system"
       (define A (matrix [[1 1] [2 1] [1 2]]))
       (define y (col-matrix [3 5 5]))
       (define ε 1e-4)

       (define result (lawson-hanson-1974 A y ε))

       ;; Solution should be approximately [1.636, 1.636] (least squares fit)
       (check-= (matrix-ref result 0 0) 1.636 0.01 "x1 approximately 1.636")
       (check-= (matrix-ref result 1 0) 1.636 0.01 "x2 approximately 1.636"))

     (test-case "Non-negativity enforcement"
       (define A (matrix [[1 -1] [1 1]]))
       (define y (col-matrix [1 3]))
       (define ε 1e-6)

       (define result (lawson-hanson-1974 A y ε))

       ;; All results should be non-negative
       (check-true (>= (matrix-ref result 0 0) 0) "x1 should be non-negative")
       (check-true (>= (matrix-ref result 1 0) 0) "x2 should be non-negative"))

     (test-case "Zero target"
       (define A (matrix [[1 2] [3 4]]))
       (define y (col-matrix [0 0]))
       (define ε 1e-6)

       (define result (lawson-hanson-1974 A y ε))

       (check-= (matrix-ref result 0 0) 0.0 ε "x1 should be 0")
       (check-= (matrix-ref result 1 0) 0.0 ε "x2 should be 0"))

     (test-case "Ferti recipe"
       (define test-date "2025-01-01")
       (define n (get-nutrient #:name "Nitrogen"))
       (define p (get-nutrient #:name "Phosphorus"))
       (define k (get-nutrient #:name "Potassium"))

       (create-nutrient-measurement! (nutrient-measurement #f test-date (hash n 0 p 0)))
       (define test-crop-requirement (create-crop-requirement! "Test requirement" (hash n 100 p 50)))
       (create-crop-rotation! test-date (hash test-crop-requirement 100))

       (create-fertilizer-product! "Nitrogen" "King Nitrogen" (hash n 100))
       (create-fertilizer-product! "Phosphorus" "Phosphorescent Baboon" (hash n 10 p 100))
       (create-fertilizer-product! "Diluted phosphorus" "John's Phosphorus" (hash n 3 p 30))

       (define recipe (find-ferti-recipe test-date))

       (check-equal? (hash-count recipe) 5 "Should have 5 fertilizer products")

       (for ([(fertilizer quantity) (in-hash recipe)])
         (check-true (>= quantity 0) "Fertilizer quantity should be non-negative")))

     (test-case "Recipe calculation with real-world scenario"
       (define test-date "2025-01-02")
       (define n (get-nutrient #:name "Nitrogen"))
       (define p (get-nutrient #:name "Phosphorus"))
       (define k (get-nutrient #:name "Potassium"))

       ;; Current levels
       (create-nutrient-measurement! (nutrient-measurement #f test-date (hash n 50 p 10 k 100)))

       ;; Target levels
       (define test-crop-requirement
         (create-crop-requirement! "Test requirement 2" (hash n 150 p 30 k 200)))
       (create-crop-rotation! test-date (hash test-crop-requirement 100))

       ;; Available fertilizer products
       (create-fertilizer-product! "" "Balanced" (hash n 100 p 100 k 100))
       (create-fertilizer-product! "Nitrogen blend" "High-N" (hash n 200 p 50 k 50))
       (create-fertilizer-product! "Phosphorus blend" "High-P" (hash n 50 p 200 k 50))
       (create-fertilizer-product! "Potassium blend" "High-K" (hash n 50 p 50 k 200))

       (define recipe (find-ferti-recipe test-date))

       (check-equal? (hash-count recipe) 9 "Recipe should have 9 fertilizers")

       ;; Verify solution is non-negative
       (for ([(fertilizer quantity) (in-hash recipe)])
         (check-true (>= quantity 0)
                     (format "~a quantity must be non-negative"
                             (fertilizer-product-name fertilizer))))))))