From bb68e04e067e771e709a46f544cffda5c9818e57 Mon Sep 17 00:00:00 2001
From: Aaron Steele <eightysteele@gmail.com>
Date: Fri, 31 May 2024 19:05:16 -0700
Subject: [PATCH] Gamma optimizations and benchmarks (#69)

* initial commit

* use test-runner, add baseline spec
---
 .dir-locals.el                                |  2 +-
 bb.edn                                        |  4 +
 .../distribution/math/gamma_benchmark.cljc    | 91 +++++++++++++++++++
 deps.edn                                      |  9 ++
 src/gen/distribution/math/gamma.cljc          | 19 ++--
 src/gen/distribution/math/utils.cljc          |  8 +-
 6 files changed, 121 insertions(+), 12 deletions(-)
 create mode 100644 benchmark/gen/distribution/math/gamma_benchmark.cljc

diff --git a/.dir-locals.el b/.dir-locals.el
index cfa3337..bb6fa96 100644
--- a/.dir-locals.el
+++ b/.dir-locals.el
@@ -1,7 +1,7 @@
 ((nil
   . ((cider-default-cljs-repl . node)
      (cider-preferred-build-tool . clojure-cli)
-     (cider-clojure-cli-aliases . ":test:nextjournal/clerk")
+     (cider-clojure-cli-aliases . ":benchmark:test:nextjournal/clerk")
 
      ;; Custom indentation:
      (eval . (put-clojure-indent 'gen 1)))))
diff --git a/bb.edn b/bb.edn
index ba690d7..402a778 100644
--- a/bb.edn
+++ b/bb.edn
@@ -27,6 +27,10 @@
   {:doc "Run CLJS tests."
    :task (shell "npm run test")}
 
+  benchmark:clj
+  {:doc "Run CLJ benchmarks."
+   :task (shell "clojure -X:benchmark:test")}
+
   build-static
   {:doc "Generate a fresh static build."
    :task
diff --git a/benchmark/gen/distribution/math/gamma_benchmark.cljc b/benchmark/gen/distribution/math/gamma_benchmark.cljc
new file mode 100644
index 0000000..93c2ff3
--- /dev/null
+++ b/benchmark/gen/distribution/math/gamma_benchmark.cljc
@@ -0,0 +1,91 @@
+(ns gen.distribution.math.gamma-benchmark
+  (:require [clojure.test :refer [deftest is testing]]
+            [same.core :refer [ish? with-comparator]]
+            [gen.generators :refer [within]]
+            [gen.distribution.math.gamma :as g]
+            [criterium.core :as crit]))
+
+(def ^:no-doc ^:const lanczos-approximation-baseline
+  [{:n 15
+    :x 9
+    :mean 3.127368933269932E-6
+    :variance 1.5780434242573877E-14
+    :lower-q 2.984515105404977E-6
+    :upper-q 3.320172412426692E-6}
+   {:n 15
+    :x 100
+    :mean 2.0160584455226908E-7
+    :variance 7.928671640576774E-18
+    :lower-q 1.9907849968797706E-7
+    :upper-q 2.0489675953499433E-7}
+   {:n 15
+    :x 170
+    :mean 2.040524636004644E-7
+    :variance 5.795520133553684E-18
+    :lower-q 2.0120967339244923E-7
+    :upper-q 2.0719692107157572E-7}])
+
+(deftest ^:benchmark lanczos-approximation-benchmark
+  []
+  (let [f (fn[b]
+            (let [n (:n b)
+                  x (:x b)
+                  mean (:mean b)
+                  variance (:variance b)
+                  lower-q (:lower-q b)
+                  upper-q (:upper-q b)
+                  stats (crit/quick-benchmark
+                         (dotimes [_ n]
+                           (g/lanczos-approximation x)) {})]
+              (testing "against lanczos approximation baseline"
+                (with-comparator (within 1e-3)
+                  (is (ish? mean (first (:mean stats))) x)
+                  (is (ish? variance (first (:variance stats))) x)
+                  (is (ish? lower-q (first (:lower-q stats))) x)
+                  (is (ish? upper-q (first (:upper-q stats))) x)))))]
+    (run! f lanczos-approximation-baseline)))
+
+(deftest ^:benchmark inv-gamma-1pm1-benchmark
+  []
+  (let [n 15
+        x [-0.05 1.5]
+        f (fn[x]
+            (println "inv-gamma-1pm1: " x)
+            (crit/quick-bench
+                (dotimes [_ n]
+                  (g/inv-gamma-1pm1 x))))]
+    (run! f x)))
+
+(deftest ^:benchmark log-gamma-1p-benchmark
+  []
+  (let [n 15
+        x [-0.05 1.5]
+        f (fn[x]
+            (println "log-gamma-1p: " x)
+            (crit/quick-bench
+                (dotimes [_ n]
+                  (g/log-gamma-1p x))))]
+    (run! f x)))
+
+
+(deftest ^:benchmark log-gamma-benchmark
+  []
+  (let [n 15
+        x [0.1 1.5 2.5 8.0 10.0]
+        f (fn[x]
+            (println "log-gamma: " x)
+            (crit/quick-bench
+                   (dotimes [_ n]
+                     (g/log-gamma x))))]
+    (run! f x)))
+
+(deftest ^:benchmark gamma-benchmark
+  []
+  (let [n 15
+        x [-0.001 0.1 1.5 2.5 8.0 10.0]
+        f (fn[x]
+            (println "gamma: " x)
+            (crit/quick-bench
+                (dotimes [_ n]
+                  (g/gamma x))))]
+    (run! f x)))
diff --git a/deps.edn b/deps.edn
index f2da2a2..b626877 100644
--- a/deps.edn
+++ b/deps.edn
@@ -33,6 +33,15 @@
                 org.clojure/test.check {:mvn/version "1.1.1"}
                 same/ish {:mvn/version "0.1.6"}}}
 
+  :benchmark
+  {:extra-paths ["benchmark"]
+   :extra-deps {io.github.cognitect-labs/test-runner
+                {:git/tag "v0.5.1" :git/sha "dfb30dd"}
+                criterium/criterium {:mvn/version "0.4.6"}}
+   :main-opts ["-m" "cognitect.test-runner"]
+   :exec-args {:includes [:benchmark] :dirs ["benchmark"] :patterns [".*-benchmark$"]}
+   :exec-fn cognitect.test-runner.api/test}
+
   ;; See https://github.com/cognitect-labs/test-runner for invocation
   ;; instructions, or call `clojure -X:test:runner`.
   :runner
diff --git a/src/gen/distribution/math/gamma.cljc b/src/gen/distribution/math/gamma.cljc
index 3951fc9..328eb44 100644
--- a/src/gen/distribution/math/gamma.cljc
+++ b/src/gen/distribution/math/gamma.cljc
@@ -27,7 +27,7 @@
 ;;
 ;; Apache Commons Numbers Gamma. Github.
 ;; https://github.com/apache/commons-numbers-gamma
-
+;;
 ;; Weisstein, Eric W. Lanczos Approximation. From MathWorld--A Wolfram Web
 ;; Resource. https://mathworld.wolfram.com/LanczosApproximation.html
 ;;
@@ -47,7 +47,7 @@
 ;; Software (TOMS), 12(4), pp.377-393.
 ;; https://dl.acm.org/doi/10.1145/22721.23109
 
-(def LANCZOS_G
+(def ^:const LANCZOS_G
   "The Lanczos constant.
 
   Based on Godfrey's research (Godfrey, P., 2001)."
@@ -59,7 +59,8 @@
   Based on Wolfram (Weisstein, Eric W., Lanczos Approximation). The coefficients
   were computed by Godfrey (Godfrey, P., 2001). See
   `LanczosApproximation` (Apache Commons Numbers Gamma)."
-  [x]
+  ^double
+  [^double x]
   {:pre [(> x 8)]}
   (let [c [[14 3.6899182659531625E-6]
            [13 -2.6190838401581408E-5]
@@ -84,7 +85,8 @@
   function (NSWC Mathematical Library, 1993) which defines the coefficients of
   the Maclaurin series expansion (A, B, C, P, Q). See also `InvGamma1pm` (Apache
   Commons Numbers Gamma)."
-  [x]
+  ^double
+  [^double x]
   {:pre [(>= x -0.5)
          (<= x 1.5)]}
   (let [t (if (<= x 0.5) x (- (- x 0.5) 0.5))
@@ -145,7 +147,8 @@
   Based on the NSWC implementation (Morris, A.H., 1993). See the `DGMLN1`
   function (NSWC Mathematical Library, 1993) and `LogGamma1p` (Apache Commons
   Numbers Gamma)."
-  [x]
+  ^double
+  [^double x]
   {:pre [(>= x -0.5)
          (<= x 1.5)]}
   (- (Math/log1p (inv-gamma-1pm1 x))))
@@ -158,7 +161,8 @@
   the NSWC implementation (Morris, A.H., 1993). See the `DGAMLN` function (NSWC
   Mathematical Library, 1993). For x > 8, follows the Lanczos approximation
   method. See `LogGamma` (Apache Commons Numbers Gamma)."
-  [x]
+  ^double
+  [^double x]
   {:pre [(pos? x)]}
   (cond
     (< x 0.5) (- (log-gamma-1p x) (Math/log x))
@@ -181,7 +185,8 @@
   Based on Wolfram (Weisstein, Eric W., Gamma Function). Follows the NSWC
   implementation (Morris, A.H., 1993). See the `DGAMMA` function (NSWC
   Mathematical Library, 1993) and `Gamma` (Apache Commons Numbers Gamma)."
-  [x]
+  ^double
+  [^double x]
   (let [abs-x (abs x)]
     (if (<= abs-x 20)
       (if (>= x 1)
diff --git a/src/gen/distribution/math/utils.cljc b/src/gen/distribution/math/utils.cljc
index 69bc38c..6a32f61 100644
--- a/src/gen/distribution/math/utils.cljc
+++ b/src/gen/distribution/math/utils.cljc
@@ -1,14 +1,14 @@
 (ns gen.distribution.math.utils
   "Math utilities and constants.")
 
-(def ^:no-doc log-pi
+(def ^:no-doc ^:const log-pi
   (Math/log Math/PI))
 
-(def ^:no-doc log-2pi
+(def ^:no-doc ^:const log-2pi
   (Math/log (* 2 Math/PI)))
 
-(def ^:no-doc sqrt-2pi
+(def ^:no-doc ^:const sqrt-2pi
   (Math/sqrt (* 2 Math/PI)))
 
-(def ^:no-doc half-log-2pi
+(def ^:no-doc ^:const half-log-2pi
   (* 0.5 (Math/log (* 2.0 Math/PI))))