<commit_before>#lang racket

(require racket/match)
(require math/bigfloat)
(require casio/common)
(require casio/points)
(require casio/programs)
(require casio/alternative)
(require casio/redgreen)

(provide improve-by-analysis)

(struct annotation (expr exact-value approx-value local-error total-error loc) #:transparent)

(define (real-op->bigfloat-op op)
  (list-ref (hash-ref operations op) mode:bf))

(define (real-op->float-op op)
  (list-ref (hash-ref operations op) mode:fl))

(define (analyze-expressions prog inputs)
  (define varmap (map cons (program-variables prog) inputs))

  (location-induct
   prog

   #:constant
   (λ (c loc)
      (let* ([exact (bf c)] [approx ((*precision*) c)]
             [error (relative-error (->flonum exact) approx)])
        (annotation c exact approx error error loc)))

   #:variable
   (λ (v loc)
      (let* ([var (cdr (assoc v varmap))]
             [exact (bf var)] [approx ((*precision*) var)]
             [error (relative-error (->flonum exact) approx)])
        (annotation v exact approx error error loc)))

   #:primitive
   (λ (expr loc)
      (let* ([exact-op (real-op->bigfloat-op (car expr))]
             [approx-op (real-op->float-op (car expr))]
             [exact-inputs (map annotation-exact-value (cdr expr))]
             [semiapprox-inputs (map ->flonum exact-inputs)]
             [approx-inputs (map annotation-approx-value (cdr expr))]
             [exact-ans (apply exact-op exact-inputs)]
             [semiapprox-ans (apply approx-op semiapprox-inputs)]
             [approx-ans (apply approx-op approx-inputs)]
             [local-error (relative-error (->flonum exact-ans)
                                          semiapprox-ans)]
             [cumulative-error (relative-error (->flonum exact-ans)
                                               approx-ans)])
        (annotation expr exact-ans approx-ans local-error cumulative-error loc)))))

(define (find-most-local-error annot-prog)
  (define (search-expression expr)
    (if (list? expr)
        (let continue ([expr (cdr expr)] [err 0] [loc #f])
          (cond
           [(null? expr) (values err loc)]
           [#t
            (let-values ([(err* loc*) (search-annot (car expr))])
              (if (> err* err)
                  (continue (cdr expr) err* loc*)
                  (continue (cdr expr) err loc)))]))
        (values 0 #f)))

  (define (search-annot annot)
    (let-values ([(err* loc*)
                  (search-expression (annotation-expr annot))])
      ; Why >=? Because all else equal, we're more interested in small subexpressions
      (if (>= err* (annotation-local-error annot))
          (values err* loc*)
          (values (annotation-local-error annot) (annotation-loc annot)))))

  (let-values ([(err loc) (search-annot (program-body annot-prog))])
    (if (= err 0) #f loc)))

(define (pick-bad-input alt)
  (let* ([lst (enumerate (alt-errors alt) (*points*) (*exacts*))]
         ; We filter out very large errors.
         ; Those usually result in overfitting.
         ; TODO: Eliminate this ad-hoc solution.
         [lst* (filter (λ (x) (< (cadr x) 1)) lst)])
    (argmax cadr (if (null? lst*) lst lst*))))

(define (step alt input)
  (let* ([annot (analyze-expressions (alt-program alt) input)]
         [loc (find-most-local-error annot)])
    (if loc
        (alt-rewrite-expression alt #:destruct #t #:root loc)
        '())))

(define (alt-error-at alt idx)
  (list-ref (alt-errors alt) idx))

(define (alt<-at? idx alt1 alt2)
  "Compare two alternatives.
   Compares first by a lattice order on points, then by program cost."
  (eq? (list-ref (errors-compare (alt-errors alt1) (alt-errors alt2))
                 idx)
       '<))

(define (improve-by-analysis alt0 iters)
  (when (*debug*) (println "> iba " alt0 " for " iters))
  (define input (pick-bad-input alt0))

  (let loop ([altn alt0] [left iters])
    (cond
     [(<= left 0) altn]
     [(green? altn) altn]
     [#t
      (let* ([alts (step altn (caddr input))])
        (when (*debug*) (println "; iba(" left ") " altn))
        (if (null? alts)
            altn
            (loop (argmin (curryr alt-error-at (car input)) alts)
                  (- left 1))))])))

<commit_msg>Add Zach's crazy try-to-correlate-error idea.<commit_after>#lang racket

(require racket/match)
(require math/bigfloat)
(require casio/common)
(require casio/points)
(require casio/programs)
(require casio/alternative)
(require casio/redgreen)

(provide improve-by-analysis)

(struct annotation (expr exact-value approx-value local-error total-error loc) #:transparent)

(define (real-op->bigfloat-op op)
  (list-ref (hash-ref operations op) mode:bf))

(define (real-op->float-op op)
  (list-ref (hash-ref operations op) mode:fl))

(define (analyze-expressions prog inputs)
  (define varmap (map cons (program-variables prog) inputs))

  (location-induct
   prog

   #:constant
   (λ (c loc)
      (let* ([exact (bf c)] [approx ((*precision*) c)]
             [error (relative-error (->flonum exact) approx)])
        (annotation c exact approx error error loc)))

   #:variable
   (λ (v loc)
      (let* ([var (cdr (assoc v varmap))]
             [exact (bf var)] [approx ((*precision*) var)]
             [error (relative-error (->flonum exact) approx)])
        (annotation v exact approx error error loc)))

   #:primitive
   (λ (expr loc)
      (let* ([exact-op (real-op->bigfloat-op (car expr))]
             [approx-op (real-op->float-op (car expr))]
             [exact-inputs (map annotation-exact-value (cdr expr))]
             [semiapprox-inputs (map ->flonum exact-inputs)]
             [approx-inputs (map annotation-approx-value (cdr expr))]
             [exact-ans (apply exact-op exact-inputs)]
             [semiapprox-ans (apply approx-op semiapprox-inputs)]
             [approx-ans (apply approx-op approx-inputs)]
             [local-error (relative-error (->flonum exact-ans)
                                          semiapprox-ans)]
             [cumulative-error (relative-error (->flonum exact-ans)
                                               approx-ans)])
        (annotation expr exact-ans approx-ans local-error cumulative-error loc)))))

(define (find-most-local-error annot-prog)
  (define (search-expression expr)
    (if (list? expr)
        (let continue ([expr (cdr expr)] [err 0] [loc #f])
          (cond
           [(null? expr) (values err loc)]
           [#t
            (let-values ([(err* loc*) (search-annot (car expr))])
              (if (> err* err)
                  (continue (cdr expr) err* loc*)
                  (continue (cdr expr) err loc)))]))
        (values 0 #f)))

  (define (search-annot annot)
    (let-values ([(err* loc*)
                  (search-expression (annotation-expr annot))])
      ; Why >=? Because all else equal, we're more interested in small subexpressions
      (if (>= err* (annotation-local-error annot))
          (values err* loc*)
          (values (annotation-local-error annot) (annotation-loc annot)))))

  (let-values ([(err loc) (search-annot (program-body annot-prog))])
    (if (= err 0) #f loc)))

(define (pick-bad-input alt)
  (let* ([lst (enumerate (alt-errors alt) (*points*) (*exacts*))]
         ; We filter out very large errors.
         ; Those usually result in overfitting.
         ; TODO: Eliminate this ad-hoc solution.
         [lst* (filter (λ (x) (< (cadr x) 1)) lst)])
    (argmax cadr (if (null? lst*) lst lst*))))

(define (step alt input)
  (let* ([annot (analyze-expressions (alt-program alt) input)]
         [loc (find-most-local-error annot)])
    (if loc
        (alt-rewrite-expression alt #:destruct #t #:root loc)
        '())))

(define (alt-error-at alt idx)
  (list-ref (alt-errors alt) idx))

(define (alt<-at? idx alt1 alt2)
  "Compare two alternatives.
   Compares first by a lattice order on points, then by program cost."
  (eq? (list-ref (errors-compare (alt-errors alt1) (alt-errors alt2))
                 idx)
       '<))

(define (improve-by-analysis alt0 iters)
  (when (*debug*) (println "> iba " alt0 " for " iters))
  (define input (pick-bad-input alt0))

  (let loop ([altn alt0] [left iters])
    (cond
     [(<= left 0) altn]
     [(green? altn) altn]
     [#t
      (let* ([alts (step altn (caddr input))])
        (when (*debug*) (println "; iba(" left ") " altn))
        (if (null? alts)
            altn
            (loop (argmin (curryr alt-error-at (car input)) alts)
                  (- left 1))))])))

(define (strip-off-top loc)
  (let ([loc* (reverse loc)])
    (match loc*
      [`(car cdr car . ,rest)
       (values (reverse (cons 'car rest))
               (reverse (list* 'car 'cdr 'cdr 'car rest)))]
      [`(car cdr cdr car . ,rest)
       (values (reverse (cons 'car rest))
               (reverse (list* 'car 'cdr 'car rest)))]
      [#t
       (error "Wat?")])))

(define (correlate-binary-ops alt0)
  (define input (pick-bad-input alt0))

  (let* ([annot (analyze-expressions (alt-program alt0) (caddr input))]
         [loc (find-most-local-error annot)])
    (if loc
        (let-values ([(parent other) (strip-off-top loc)])
          (if (= (length (location-get parent (alt-program alt0))) 3)
              (alt-rewrite-tree alt0 #:root other)
              alt0))
        alt0)))

{'fraction_non_alphanumeric': 0.15537118844952053, 'fraction_numerical': 0.0038788923607369894, 'mean_word_length': 2.3472773169852146, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 29, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5264135', 'n_tokens_mistral': 3078, 'n_tokens_neox': 2965, 'n_words': 895}