Push more computations into node functions to improve performance

Project.toml

@@ -19,6 +19,7 @@ LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
 LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 MetaGraphsNext = "fa8bd995-216d-47f1-8a91-f3b68fbeb377"
+MistyClosures = "dbe65cb8-6be2-42dd-bbc5-4196aaced4f4"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -65,6 +66,7 @@ LogExpFunctions = "0.3"
 MCMCChains = "6"
 MacroTools = "0.5"
 MetaGraphsNext = "0.6, 0.7"
+MistyClosures = "2.0.0"
 OrderedCollections = "1"
 PDMats = "0.10, 0.11"
 SpecialFunctions = "2"

src/BUGSPrimitives/distributions.jl

@@ -8,12 +8,11 @@ with mean ``μ`` and standard deviation ``\\frac{1}{√τ}``.
 p(x|μ,τ) = \\sqrt{\\frac{τ}{2π}} e^{-τ \\frac{(x-μ)^2}{2}}
 ```
 """
-function dnorm(μ, τ)
-    if τ < 0
+function dnorm(μ, τ)::Normal
+    if τ <= 0
         throw(DomainError((μ, τ), "Requires τ > 0"))
     end
-    σ² = 1 / τ # variance
-    σ = √σ² # standard deviation
+    σ = inv(√τ) # standard deviation
     return Normal(μ, σ)
 end
 
@@ -27,9 +26,8 @@ with location parameter ``μ`` and scale parameter ``\\frac{1}{√τ}``.
 p(x|μ,τ) = \\frac{\\sqrt{τ} e^{-\\sqrt{τ}(x-μ)}}{(1+e^{-\\sqrt{τ}(x-μ)})^2}
 ```
 """
-function dlogis(μ, τ)
-    s = 1 / √τ
-    return Logistic(μ, s)
+function dlogis(μ, τ)::Logistic
+    return Logistic(μ, inv(√τ))
 end
 
 """
@@ -70,9 +68,9 @@ p(x|ν,μ,σ) = \\frac{Γ((ν+1)/2)}{Γ(ν/2) \\sqrt{νπσ}}
 \\left(1+\\frac{1}{ν}\\left(\\frac{x-μ}{σ}\\right)^2\\right)^{-\\frac{ν+1}{2}}
 ```
 """
-function dt(μ, τ, ν)
-    σ = sqrt(1 / τ)
-    if μ == 0 && σ == 1
+function dt(μ::Real, τ::Real, ν::Real)
+    σ = inv(√τ)
+    if iszero(μ) && isone(σ)
         return TDist(ν)
     else
         return TDistShiftedScaled(ν, μ, σ)
@@ -89,9 +87,9 @@ with location ``μ`` and scale ``\\frac{1}{\\sqrt{τ}}``.
 p(x|μ,τ) = \\frac{\\sqrt{τ}}{2} e^{-\\sqrt{τ} |x-μ|}
 ```
 """
-function ddexp(μ, τ)
-    b = 1 / √τ
-    return Laplace(μ, b)
+function ddexp(μ::Real, τ::Real)
+    b = one(τ) / sqrt(τ)
+    return Laplace(convert(typeof(b), μ), b)
 end
 
 """

src/JuliaBUGS.jl

@@ -10,13 +10,14 @@ using Graphs, MetaGraphsNext
 using LinearAlgebra
 using LogDensityProblems, LogDensityProblemsAD
 using MacroTools
+using MistyClosures
 using OrderedCollections: OrderedDict
 using Random
 using StaticArrays
 
 import Base: ==, hash, Symbol, size
 import Distributions: truncated
-import AbstractPPL: AbstractContext, evaluate!!
+import AbstractPPL: condition, decondition, evaluate!!
 
 export @bugs
 export compile, initialize!