Symbolics exporation

[marius311]

After talking with Jamie and randomly scanning the repo / thinking about things, I was kind of curious if Symbolics could help. In theory it could be used to get analytic derivatives and sparsity detection for the hierarchy, which could maybe be helpful somewhere. I coded up the following tiny example which shows how it works, if you guys had never tried, TLDR doesnt seem too bad on first glance and the science code would in theory be unchaged, and derivative here is 5X faster. Maybe worth exploring further, maybe not :)

using Symbolics, BenchmarkTools, ForwardDiff
@generated fieldvalues(x) = Expr(:tuple, (:(x.$f) for f=fieldnames(x))...)


## science code

struct Params{T}
    ργ₀ :: T
    ρc₀ :: T
end
ργ(𝕡, z) = 𝕡.ργ₀ * (1+z)^4
ρm(𝕡, z) = 𝕡.ρc₀ * (1+z)^3
Hubble(𝕡, z) = √(8π/3 * (ργ(𝕡, z) + ρm(𝕡, z)))


## running with numbers

𝕡 = Params(1e-1, 1e-2)
@btime Hubble(𝕡, 1) # 15ns
@btime ForwardDiff.derivative(z -> Hubble(𝕡, z), 1) # 120ns
@btime ForwardDiff.hessian(((z,),) -> Hubble(𝕡, z), [1]) # 1.5us


## building symbolic expression then running that

@variables ργ₀ ρc₀ z
𝕡_sym = Params(ργ₀, ρc₀)

Hubble_from_sym = let
    f = build_function(Hubble(𝕡_sym, z), ργ₀, ρc₀, z, expression=Val(false))
    (𝕡, z) -> f(fieldvalues(𝕡)..., z)
end

dz_Hubble_from_sym = let
    f = build_function(simplify(expand_derivatives(Differential(z)(Hubble(𝕡_sym, z)))), ργ₀, ρc₀, z, expression=Val(false))
    (𝕡, z) -> f(fieldvalues(𝕡)..., z)
end

d²z_Hubble_from_sym = let
    f = build_function(simplify(expand_derivatives((Differential(z)^2)(Hubble(𝕡_sym, z)))), ργ₀, ρc₀, z, expression=Val(false))
    (𝕡, z) -> f(fieldvalues(𝕡)..., z)
end

@btime Hubble_from_sym(𝕡, 1) # 15ns
@btime dz_Hubble_from_sym(𝕡, 1) # 17ns
@btime d²z_Hubble_from_sym(𝕡, 1) # 20ns