Fun with SymbolicUtils

[cscherrer]

Check it out:

julia> using MeasureTheory, SymbolicUtils

julia> @syms μ σ x
(μ, σ, x)

julia> d = Normal(μ,σ)
Normal(μ = μ, σ = σ)

julia> typeof(d)
Normal{(:μ, :σ), Tuple{SymbolicUtils.Sym{Number}, SymbolicUtils.Sym{Number}}}

julia> logdensity(d, x)
(-log(σ)) - (((x - μ) ^ 2) / (2 * (σ ^ 2)))

julia> simplify(ans, polynorm=true)
(-1 * log(σ)) + (x * μ * (σ ^ -2)) + (-0.5 * (x ^ 2) * (σ ^ -2)) + (-0.5 * (μ ^ 2) * (σ ^ -2))

julia> rand(d)
(1.0183271329944876 * σ) + μ

In Soss, I'm using the simplify result (after a few other rules as well). It's something like

1. Get the symbolic logdensity
2. Simplify
3. Statically evaluate any <: Number subexpression where all values are known at compile time
4. Common subexpression elimination
5. Codegen with GeneralizedGenerated

[cscherrer]

(4) is here

julia> using SymbolicCodegen

julia> simplify(logdensity(d,x), polynorm=true) |> cse
9-element Vector{Pair{Symbol, SymbolicUtils.Term{Number}}}:
 Symbol("##286") => log(σ)
 Symbol("##287") => -1 * ##286
 Symbol("##288") => σ ^ -2
 Symbol("##289") => x * μ * ##288
 Symbol("##290") => x ^ 2
 Symbol("##291") => -0.5 * ##290 * ##288
 Symbol("##292") => μ ^ 2
 Symbol("##293") => -0.5 * ##292 * ##288
 Symbol("##294") => ##287 + ##289 + ##291 + ##293

julia> simplify(logdensity(d,x), polynorm=true) |> codegen
quote
    var"##303" = (log)(σ)
    var"##304" = (*)(-1, var"##303")
    var"##305" = (^)(σ, -2)
    var"##306" = (*)(x, μ, var"##305")
    var"##307" = (^)(x, 2)
    var"##308" = (*)(-0.5, var"##307", var"##305")
    var"##309" = (^)(μ, 2)
    var"##310" = (*)(-0.5, var"##309", var"##305")
    var"##311" = (+)(var"##304", var"##306", var"##308", var"##310")
end

[sethaxen]

Very nice! This really showcases the benefit of a more lightweight and less constrained implementation of distributions. What does the polynorm flag do?

[cscherrer]

polynorm=true expands polynomials. There might be other differences, but here's the big one I see:

julia> using SymbolicUtils
[ Info: Precompiling SymbolicUtils [d1185830-fcd6-423d-90d6-eec64667417b]

julia> @syms x
(x,)

julia> simplify((x+1)^2)
(1 + x)^2

julia> simplify((x+1)^2; polynorm=true)
1 + (x^2) + 2x

[cscherrer]

Just a note: The polynorm flag is now broken for this use case (see JuliaSymbolics/SymbolicUtils.jl#175). This root of this seems to be the use of AbstractAlgebra.jl, which apparently lacks type safety. The details of this aren't yet very clear to me.