<commit_before>;; Copyright (c) 2014 <NAME>
;; Released under the terms of the 2-clause BSD license.
;; See the file COPYRIGHT for details.

#lang racket/base
(require racket/contract
         racket/math
         racket/flonum
         math/flonum
         (prefix-in m: math/distributions)
         (prefix-in m: math/special-functions)
         "../private/dist.rkt"
         "../private/dist-define.rkt"
         "../private/dist-impl.rkt"
         "../matrix.rkt")
(provide #| implicit in define-dist-type |#)

(define-dist-type multi-normal-dist
  ([mean col-matrix?] [cov square-matrix?])
  #:pdf multi-normal-pdf
  #:sample multi-normal-sample
  #:guard (lambda (mean cov _name)
            (define n (matrix-num-rows mean))
            (unless (= n (square-matrix-size cov))
              (error 'multi-normal-dist
                     "covariance matrix has wrong number of rows\n  expected: ~s\n value: ~e"
                     n cov))
            ;; check cov is symmetric, nonnegative-definite (?)
            (values mean cov))
  #:mean mean
  #:modes (list mean) ;; FIXME: degenerate cases?
  #:variance cov
  #:drift (lambda (value scale-factor) (cons (multi-normal-drift value cov scale-factor) 0)))

(define-dist-type wishart-dist
  ([n real?] [V square-matrix?])
  #:pdf wishart-pdf #:sample wishart-sample
  #:guard (lambda (n V _name)
            (unless (matrix-symmetric? V)
              (error 'wishart-dist "expected symmetric matrix\n  given: ~e" V))
            (define p (square-matrix-size V))
            (unless (> n (- p 1))
              (error 'wishart-dist "expected n > p - 1\n  n: ~e\n  p: ~e" n p))
            ;; FIXME: check V positive definite
            (values n V)))
;; DRIFT: draw another wishart, weighted avg with current value
;; NOTE: not symmetric!
;; NOTE: for weight eps, let eps = sample-factor / (1 + sample-factor)

(define-dist-type inverse-wishart-dist
  ([n real?] [Vinv square-matrix?])
  #:pdf inverse-wishart-pdf #:sample inverse-wishart-sample
  #:guard (lambda (n Vinv _name)
            (unless (matrix-symmetric? Vinv)
              (error 'wishart-dist "expected symmetric matrix\n  given: ~e" Vinv))
            (define p (square-matrix-size Vinv))
            (unless (> n (- p 1))
              (error 'wishart-dist "expected n > p - 1\n  n: ~e\n  p: ~e" n p))
            ;; FIXME: check V positive definite
            (values n Vinv)))
;; DRIFT: invert, do wishart drift, invert

;; ============================================================
;; Multivariate dist functions

(define-memoize1 (memo-matrix-inverse m) (matrix-inverse m))
(define-memoize1 (memo-matrix-cholesky m) (matrix-cholesky m))
(define-memoize1 (memo-matrix-determinant m) (matrix-determinant m))

;; ------------------------------------------------------------
;; Multivariate normal dist functions<commit_msg>drift proposal for wishart dist<commit_after>;; Copyright (c) 2014 <NAME>
;; Released under the terms of the 2-clause BSD license.
;; See the file COPYRIGHT for details.

#lang racket/base
(require racket/contract
         racket/math
         racket/flonum
         math/flonum
         (prefix-in m: math/distributions)
         (prefix-in m: math/special-functions)
         "../private/dist.rkt"
         "../private/dist-define.rkt"
         "../private/dist-impl.rkt"
         "../matrix.rkt")
(provide #| implicit in define-dist-type |#)

(define-dist-type multi-normal-dist
  ([mean col-matrix?] [cov square-matrix?])
  #:pdf multi-normal-pdf
  #:sample multi-normal-sample
  #:guard (lambda (mean cov _name)
            (define n (matrix-num-rows mean))
            (unless (= n (square-matrix-size cov))
              (error 'multi-normal-dist
                     "covariance matrix has wrong number of rows\n  expected: ~s\n value: ~e"
                     n cov))
            ;; check cov is symmetric, nonnegative-definite (?)
            (values mean cov))
  #:mean mean
  #:modes (list mean) ;; FIXME: degenerate cases?
  #:variance cov
  #:drift (lambda (value scale-factor) (cons (multi-normal-drift value cov scale-factor) 0)))

(define-dist-type wishart-dist
  ([n real?] [V square-matrix?])
  #:pdf wishart-pdf #:sample wishart-sample
  #:guard (lambda (n V _name)
            (unless (matrix-symmetric? V)
              (error 'wishart-dist "expected symmetric matrix\n  given: ~e" V))
            (define p (square-matrix-size V))
            (unless (> n (- p 1))
              (error 'wishart-dist "expected n > p - 1\n  n: ~e\n  p: ~e" n p))
            ;; FIXME: check V positive definite
            (values n V))
  #:drift (lambda (value scale-factor)
            (define w (/ scale-factor (+ 1 scale-factor)))
            (define F-sample (wishart-sample n V))
            (define value*
              (matrix+ (array* value (array (- 1 w)))
                       (array* (array w) F-sample)))
            (define F (wishart-pdf n V F-sample #t))
            ;; Reverse direction: solve value = (1 - w) * value* + w * R-sample
            (define R-sample
              (array* (matrix- value (array* (array (- 1 w)) value*))
                      (array (/ w))))
            (define R (wishart-pdf n V R-sample #t))
            (cons value* (- R F))))
;; DRIFT: draw another wishart, weighted avg with current value
;; NOTE: not symmetric!
;; NOTE: for weight eps, let eps = sample-factor / (1 + sample-factor)

(define-dist-type inverse-wishart-dist
  ([n real?] [Vinv square-matrix?])
  #:pdf inverse-wishart-pdf #:sample inverse-wishart-sample
  #:guard (lambda (n Vinv _name)
            (unless (matrix-symmetric? Vinv)
              (error 'wishart-dist "expected symmetric matrix\n  given: ~e" Vinv))
            (define p (square-matrix-size Vinv))
            (unless (> n (- p 1))
              (error 'wishart-dist "expected n > p - 1\n  n: ~e\n  p: ~e" n p))
            ;; FIXME: check V positive definite
            (values n Vinv)))
;; DRIFT: invert, do wishart drift, invert

;; ============================================================
;; Multivariate dist functions

(define-memoize1 (memo-matrix-inverse m) (matrix-inverse m))
(define-memoize1 (memo-matrix-cholesky m) (matrix-cholesky m))
(define-memoize1 (memo-matrix-determinant m) (matrix-determinant m))

;; ------------------------------------------------------------
;; Multivariate normal dist functions
{'fraction_non_alphanumeric': 0.1716616977875412, 'fraction_numerical': 0.005021183116271772, 'mean_word_length': 2.18063872255489, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5907799', 'n_tokens_mistral': 2053, 'n_tokens_neox': 1928, 'n_words': 596}