gistfile1.scm

;; -*- coding: utf-8 -*-

#!c-expr

;; I always use Scheme as a quick calculators, and a few such expressions
;; happened to be in my scratch buffer so I turned it to C-exprs.

(use srfi-42)

(define ^^ expt)

(define (S0 λ N)
  (sum-ec (: n 1 N)
          {λ * (exp (- {n * λ})) * (- {2 ^^ (ceiling (log n 2))}) * n}))

(define (S1 λ K)
  (sum-ec (: k K)
          {{2 ^^ k} * (exp (- {{2 ^^ k} * λ}))}))

(define (St λ K)
  {(- {(/ λ) * (exp (- λ))})
   + (sum-ec (: k 1 K)
             {{2 ^^ {k - 1}} * (exp (- {{2 ^^ k} * λ}))})})


;; Jeremy Gibbons's Unbounded Spigot algorithm.
;;   http://sourceforge.net/projects/readable/files/
;; Original code is in Haskell.  Here's my port to Gauche:
;;   http://blog.practical-scheme.net/shiro/20120713-spigot
;; Since it was one of the times that I envied Haskell's infix syntax,
;; let's try to use C-expr.

#!c-expr

(use gauche.lazy)
(use util.match)

(define-syntax define*
  (syntax-rules ()
    [(_ (fn . pats) . body) (define fn (match-lambda* [pats . body]))]))

(define (stream next safe prod kons seed xs)
  (^[] (let loop ([y (next seed)])
         (cond [(safe seed y) (set! seed (prod seed y)) y]
               [else (set! seed (kons seed (car xs)))
                     (set! xs (cdr xs))
                     (loop (next seed))]))))

(define (lft q r s t)
  (let1 f (gcd q r s t)
    (vector {q / f} {r / f} {s / f} {t / f})))

(define *unit-lft* '#(1 0 0 1))

(define* (lft*v #(q r s t) x y)
  (floor {{{q * x} + {r * y}} / {{s * x} + {t * y}}}))

(define* (lft*lft #(q r s t)  #(u v w x))
  (lft {{q * u} + {r * w}} {{q * v} + {r * x}}
       {{s * u} + {t * w}} {{s * v} + {t * x}}))

(define (pi)
  (define lfts (lmap (^k (lft k {{4 * k} + 2} 0 {{2 * k} + 1})) (lrange 1)))
  (define (next z)   (lft*v z 3 1))
  (define (safe z n) {n = (lft*v z 4 1)}
  (define (prod z n) (lft*lft (lft 10 {-10 * n} 0 1) z))
  (stream next safe prod lft*lft *unit-lft* lfts))

(define (piL)
  (define lfts (lmap (^i (lft {{2 * i} - 1} {i * i} 1 0)) (lrange 1)))
  (define* (next (z . i))    (lft*v z {{2 * i} - 1} 1))
  (define* (safe (z . i) n)  {n = (lft*v z {{5 * i} - 2} 2)})
  (define* (prod (z . i) n)  (cons (lft*lft (lft 10 {-10 * n} 0 1) z) i))
  (define* (kons (z . i) z_) (cons (lft*lft z z_) {i + 1}))
  (stream next safe prod kons `(,(lft 0 4 1 0) . 1) lfts))

;; Lazy infinite sequence of prime numbers.
;;  http://blog.practical-scheme.net/shiro/20110929-primes

(define (prime? n)
  (not (any zero? (lmap (pa$ mod n) (take-while (^k {{k * k} <= n}) *primes*)))))

(define (primes-from k)
  (if (prime? k)
    (lcons k (primes-from {k + 2}))
    (primes-from {k + 2})))

(define *primes* (list* 2 3 (lcons 5 (primes-from 7))))

;; Some math-heavy code excerpts from https://github.com/shirok/mpm-experiment

(define-constant √2π (sqrt {pi * 2}))

(define (φ x) {(exp {-.5 * x * x}) / √2π})

(define (Φ x)
  (define (Φ+ x)
    (let* ([t {/ {1 + {0.2316419 * x}}}]
           [t2 {t * t}]
           [t3 {t * t2}]
           [t4 {t * t3}]
           [t5 {t * t4}])
      {1 - {(φ x) * {{0.319381530 * t} + {-0.356563782 * t2}
                     + {1.781477937 * t3} + {-1.821255978 * t4}
                     + {1.330274429 * t5}}}}))
  (cond [{x = 0} .5]
        [{x > 0} (Φ+ x)]
        [else {1 - (Φ+ (- x))}]))

(define (cpick μ σ)
  (let ([min {μ - {4 * σ}}]
        [max {μ + {4 * σ}}]
        [u (random-real)])
    (let loop ([x0 μ])
      (let1 Fx {(Φ {{x0 - μ} / σ}) - u}
        (if {(abs Fx) < 10e-6}
          x0
          (let* ([fx {(φ {{x0 - μ} / σ}) / σ}]
                 [x1 {x0 - {Fx / fx}}])
            (cond [{x1 <= min} min]
                  [{x1 >= max} max]
                  [else (loop x1)])))))))

(define (derive! tree X)
  (define (derive-1 Xs)
    (if (null? Xs) (finalize!) (grow! (car Xs) (cdr Xs))))
  (define (finalize!)
    (let ([λs '()] [αs '()] [βs '()] [Xs '()])
      (define (rec tree)
        (match tree
          [('F <λ>) (push! λs <λ>)]
          [('X _ subtree) (push! Xs tree) (rec subtree)]
          [((or 'L 'R) <α> <β>) (push! αs <α>) (push! βs <β>)]
          [('Tree nodes ...) (for-each rec nodes)]
          [_ #f]))
      (rec (Tree-ω tree))
      (Tree-λs-set! tree (list->vector λs))
      (Tree-αs-set! tree (list->vector αs))
      (Tree-βs-set! tree (list->vector βs))
      (Tree-num-params-set! tree {(vector-length (Tree-λs tree))
                                  + {(vector-length (Tree-αs tree)) * 2}})
      (Tree-Xs-set! tree Xs)
      (recalculate-weights! tree)))
  (define (grow! X Xs)
    (let* ([l (cadr X)] [u (random-real)])
      (if {u < {l / M}}
        (let1 <λ.> (list (cpick μ_λ σ_λ))
          (set! (caddr X) `(Tree (F ,<λ.>) Z))
          (derive-1 Xs))
        (let ([<λ.> (list (cpick μ_λ σ_λ))]
              [<α.> (list (cpick μ_α σ_α))]
              [<β.> (list (cpick μ_β σ_β))]
              [XL (list 'X {l + 1} #f)]
              [XR (list 'X {l + 1} #f)])
          (set! (caddr X) `(Tree (F ,<λ.>)
                                 (Tree (L ,<α.> ,<β.>) ,XL)
                                 (Tree (R ,<α.> ,<β.>) ,XR)))
          (derive-1 (list* XL XR Xs))))))
  (derive-1 `(,X)))

(define (recalculate-weights! tree)
  (let1 maxdepth (apply max (map cadr (Tree-Xs tree)))
    (receive (ws sum)
        (map-accum (^[Xi sum]
                     (let1 wi {2 ^^ {maxdepth - (cadr Xi)}}
                       (values wi {wi + sum})))
                   0 (Tree-Xs tree))
      (set! (Tree-ws tree) ws)
      (set! (Tree-Σw tree) sum))))

(define (walk-node node twig leaf)
  (define (rec node level x y θ)
    (match node
      [('F (λ))
       (let ([x. {x + {λ * (sin θ)}}]
             [y. {y + {λ * (cos θ)}}])
         (when twig (twig x y x. y. level))
         (values x. y. θ))]
      ['Z (when leaf (leaf x y level)) (values x y θ)]
      [('X level subtree) (rec subtree level x y θ)]
      [('Tree nodes ...)
       (let loop ([ns nodes] [x. x] [y. y] [θ. θ])
         (match ns
           [() (values x y θ)]
           [(n . ns) (receive (x. y. θ.) (rec n level x. y. θ.)
                       (loop ns x. y. θ.))]))]
      [((and (or 'L 'R) z) (α) (β))
       (values x y {θ + (if {z eq? 'L} {β + α} {β - α})})]))
  (rec node 0 0 0 0))

(define kernel-center {2 * (round->exact σ_z)})
(define kernel-size {{2 * kernel-center} - 1})
(define I0 (make-f32vector {I_size * I_size} 1.0))

(define kernel
  (rlet1 v (make-f32vector {kernel-size * kernel-size})
    (do-ec (: y kernel-size)
           (: x kernel-size)
           (f32vector-set! v {x + {y * kernel-size}}
                           (exp (- {{{{y - kernel-center - 1} ^^ 2}
                                     + {{x - kernel-center - 1} ^^ 2}}
                                    / σ_z}))))))

(define (likelihood tree I)             ;I :: <f32vector>
  (define buf (likelihood-buf tree))
  (define penalty 0)  ; if we're outside of canvas, we penalize it a lot.
  (define (roundI a min max)
    (round->exact (clamp {{I_size - 1} * {{a - min} /. {max - min}}}
                         0 {I_size - 1})))
  (define (add-Z x y)
    (if {{x < canvas-x0} or {canvas-x1 < x}
         or {y < canvas-y0} or {canvas-y1 < y}}
      (inc! penalty 5.0)
      (let ([ix (roundI x canvas-x0 canvas-x1)]
            [iy (roundI y canvas-y0 canvas-y1)])
        (do-ec (: y kernel-size)
               (: x kernel-size)
               (let ([jy {iy + y + (- kernel-center)}]
                     [jx {ix + x + (- kernel-center)}])
                 (when {{-1 < jy < I_size} and {-1 < jx < I_size}}
                   (f32vector-set! buf {jx + {jy * I_size}}
                                   {(f32vector-ref kernel
                                                   {x + {y * kernel-size}})
                                    +
                                    (f32vector-ref buf
                                                   {jx + {jy * I_size}})}))))
        )))

  (f32vector-fill! buf 0)
  (walk-node (Tree-ω tree) #f (^[x y level] (add-Z x y)))
  (f32vector-clamp! buf 0 1)
  (f32vector-sub! buf I)
  ;; Just to calculate likelihood, we can do (exp (- (f32vector-dot buf buf)))
  ;; but we want the squared image for visualization.
  (f32vector-mul! buf buf)
  (exp (- {penalty + (f32vector-dot buf I0)})))