Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add Symbolics hashconsing benchmarks #1164

Draft
wants to merge 1 commit into
base: master
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 24 additions & 0 deletions benchmarks/Symbolics/Project.toml
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
[deps]
CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0"
JuliaSimCompiler = "8391cb6b-4921-5777-4e45-fd9aab8cb88d"
JuliaSimCompilerRuntime = "9cbdfd5a-25c0-4dde-8b55-206f91b28bd9"
ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
ModelingToolkitStandardLibrary = "16a59e39-deab-5bd0-87e4-056b12336739"
Multibody = "e1cad5d1-98ef-44f9-a79a-9ca4547f95b9"
OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
Polynomials = "f27b6e38-b328-58d1-80ce-0feddd5e7a45"
PreferenceTools = "ba661fbb-e901-4445-b070-854aec6bfbc5"
SciMLBenchmarks = "31c91b34-3c75-11e9-0341-95557aab0344"
SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
XSteam = "95ff35a0-be81-11e9-2ca3-5b4e338e8476"

[compat]
CairoMakie = "0.11.11"
ModelingToolkit = "9.20.0"
ModelingToolkitStandardLibrary = "2"
Polynomials = "4.0.8"
PreferenceTools = "0.1.2"
SciMLBenchmarks = "0.1.3"
SymbolicUtils = "3.11"
Symbolics = "6.23"
332 changes: 332 additions & 0 deletions benchmarks/Symbolics/ThermalFluid.jmd
Original file line number Diff line number Diff line change
@@ -0,0 +1,332 @@
---
title: Symbolic Manipulation - Large Jacobians
author: Aayush Sabharwal, Bowen Zhu, Chris Rackauckas
---

This benchmark utilizes a thermal fluid ODE setup to benchmark the scaling of
symbolic jacobian computation with and without hashconsing.

```julia
using Pkg
# Rev fixes precompilation https://github.com/hzgzh/XSteam.jl/pull/2
Pkg.add(Pkg.PackageSpec(;name="XSteam", rev="f2a1c589054cfd6bba307985a3a534b6f5a1863b"))

using ModelingToolkit, JuliaSimCompiler, Symbolics, SymbolicUtils, XSteam, Polynomials, BenchmarkTools, OrdinaryDiffEq

using ModelingToolkit: t_nounits as t, D_nounits as D
```

## Setup Julia Code

```julia
# o o o o o o o < heat capacitors
# | | | | | | | < heat conductors
# o o o o o o o
# | | | | | | |
#Source -> o--o--o--o--o--o--o -> Sink
# advection diff source PDE

m_flow_source(t) = 2.75
T_source(t) = (t > 12 * 3600) * 56.0 + 12.0
@register_symbolic m_flow_source(t)
@register_symbolic T_source(t)

#build polynomial liquid-water property only dependent on Temperature
p_l = 5 #bar
T_vec = collect(1:1:150);
@generated kin_visc_T(t) = :(Base.evalpoly(t, $(fit(T_vec, my_pT.(p_l, T_vec) ./ rho_pT.(p_l, T_vec), 5).coeffs...,)))
@generated lambda_T(t) = :(Base.evalpoly(t, $(fit(T_vec, tc_pT.(p_l, T_vec), 3).coeffs...,)))
@generated Pr_T(t) = :(Base.evalpoly(t, $(fit(T_vec, 1e3 * Cp_pT.(p_l, T_vec) .* my_pT.(p_l, T_vec) ./ tc_pT.(p_l, T_vec), 5).coeffs...,)))
@generated rho_T(t) = :(Base.evalpoly(t, $(fit(T_vec, rho_pT.(p_l, T_vec), 4).coeffs...,)))
@generated rhocp_T(t) = :(Base.evalpoly(t, $(fit(T_vec, 1000 * rho_pT.(p_l, T_vec) .* Cp_pT.(p_l, T_vec), 5).coeffs...,)))
@register_symbolic kin_visc_T(t)
@register_symbolic lambda_T(t)
@register_symbolic Pr_T(t)
@register_symbolic rho_T(t)
@register_symbolic rhocp_T(t)

@connector function FluidPort(; name, p=101325.0, m=0.0, T=0.0)
sts = @variables p(t) = p m(t) = m [connect = Flow] T(t) = T [connect = Stream]
ODESystem(Equation[], t, sts, []; name=name)
end

@connector function VectorHeatPort(; name, N=100, T0=0.0, Q0=0.0)
sts = @variables (T(t))[1:N] = T0 (Q(t))[1:N] = Q0 [connect = Flow]
ODESystem(Equation[], t, [T; Q], []; name=name)
end

@register_symbolic Dxx_coeff(u, d, T)
#Taylor-aris dispersion model
function Dxx_coeff(u, d, T)
Re = abs(u) * d / kin_visc_T(T) + 0.1
if Re < 1000.0
(d^2 / 4) * u^2 / 48 / 0.14e-6
else
d * u * (1.17e9 * Re^(-2.5) + 0.41)
end
end

@register_symbolic Nusselt(Re, Pr, f)
#Nusselt number model
function Nusselt(Re, Pr, f)
if Re <= 2300.0
3.66
elseif Re <= 3100.0
3.5239 * (Re / 1000)^4 - 45.158 * (Re / 1000)^3 + 212.13 * (Re / 1000)^2 - 427.45 * (Re / 1000) + 316.08
else
f / 8 * ((Re - 1000) * Pr) / (1 + 12.7 * (f / 8)^(1 / 2) * (Pr^(2 / 3) - 1))
end
end

@register_symbolic Churchill_f(Re, epsilon, d)
#Darcy weisbach friction factor
function Churchill_f(Re, epsilon, d)
theta_1 = (-2.457 * log(((7 / Re)^0.9) + (0.27 * (epsilon / d))))^16
theta_2 = (37530 / Re)^16
8 * ((((8 / Re)^12) + (1 / ((theta_1 + theta_2)^1.5)))^(1 / 12))
end

function FluidRegion(; name, L=1.0, dn=0.05, N=100, T0=0.0,
lumped_T=50, diffusion=true, e=1e-4)
@named inlet = FluidPort()
@named outlet = FluidPort()
@named heatport = VectorHeatPort(N=N)

dx = L / N
c = [-1 / 8, -3 / 8, -3 / 8] # advection stencil coefficients
A = pi * dn^2 / 4

p = @parameters C_shift = 0.0 Rw = 0.0 # stuff for latter
@variables begin
(T(t))[1:N] = fill(T0, N)
Twall(t)[1:N] = fill(T0, N)
(S(t))[1:N] = fill(T0, N)
(C(t))[1:N] = fill(1.0, N)
u(t) = 1e-6
Re(t) = 1000.0
Dxx(t) = 0.0
Pr(t) = 1.0
alpha(t) = 1.0
f(t) = 1.0
end

sts = vcat(T, Twall, S, C, Num[u], Num[Re], Num[Dxx], Num[Pr], Num[alpha], Num[f])

eqs = Equation[
Re ~ 0.1 + dn * abs(u) / kin_visc_T(lumped_T)
Pr ~ Pr_T(lumped_T)
f ~ Churchill_f(Re, e, dn) #Darcy-weisbach
alpha ~ Nusselt(Re, Pr, f) * lambda_T(lumped_T) / dn
Dxx ~ diffusion * Dxx_coeff(u, dn, lumped_T)
inlet.m ~ -outlet.m
inlet.p ~ outlet.p
inlet.T ~ instream(inlet.T)
outlet.T ~ T[N]
u ~ inlet.m / rho_T(inlet.T) / A
[C[i] ~ dx * A * rhocp_T(T[i]) for i in 1:N]
[S[i] ~ heatport.Q[i] for i in 1:N]
[Twall[i] ~ heatport.T[i] for i in 1:N]

#source term
[S[i] ~ (1 / (1 / (alpha * dn * pi * dx) + abs(Rw / 1000))) * (Twall[i] - T[i]) for i in 1:N]

#second order upwind + diffusion + source
D(T[1]) ~ u / dx * (inlet.T - T[1]) + Dxx * (T[2] - T[1]) / dx^2 + S[1] / (C[1] - C_shift)
D(T[2]) ~ u / dx * (c[1] * inlet.T - sum(c) * T[1] + c[2] * T[2] + c[3] * T[3]) + Dxx * (T[1] - 2 * T[2] + T[3]) / dx^2 + S[2] / (C[2] - C_shift)
[D(T[i]) ~ u / dx * (c[1] * T[i-2] - sum(c) * T[i-1] + c[2] * T[i] + c[3] * T[i+1]) + Dxx * (T[i-1] - 2 * T[i] + T[i+1]) / dx^2 + S[i] / (C[i] - C_shift) for i in 3:N-1]
D(T[N]) ~ u / dx * (T[N-1] - T[N]) + Dxx * (T[N-1] - T[N]) / dx^2 + S[N] / (C[N] - C_shift)
]

ODESystem(eqs, t, sts, p; systems=[inlet, outlet, heatport], name=name)
end

@register_symbolic Cn_circular_wall_inner(d, D, cp, ρ)
function Cn_circular_wall_inner(d, D, cp, ρ)
C = pi / 4 * (D^2 - d^2) * cp * ρ
return C / 2
end

@register_symbolic Cn_circular_wall_outer(d, D, cp, ρ)
function Cn_circular_wall_outer(d, D, cp, ρ)
C = pi / 4 * (D^2 - d^2) * cp * ρ
return C / 2
end

@register_symbolic Ke_circular_wall(d, D, λ)
function Ke_circular_wall(d, D, λ)
2 * pi * λ / log(D / d)
end

function CircularWallFEM(; name, L=100, N=10, d=0.05, t_layer=[0.002],
λ=[50], cp=[500], ρ=[7850], T0=0.0)
@named inner_heatport = VectorHeatPort(N=N)
@named outer_heatport = VectorHeatPort(N=N)
dx = L / N
Ne = length(t_layer)
Nn = Ne + 1
dn = vcat(d, d .+ 2.0 .* cumsum(t_layer))
Cn = zeros(Nn)
Cn[1:Ne] += Cn_circular_wall_inner.(dn[1:Ne], dn[2:Nn], cp, ρ) .* dx
Cn[2:Nn] += Cn_circular_wall_outer.(dn[1:Ne], dn[2:Nn], cp, ρ) .* dx
p = @parameters C_shift = 0.0
Ke = Ke_circular_wall.(dn[1:Ne], dn[2:Nn], λ) .* dx
@variables begin
(Tn(t))[1:N, 1:Nn] = fill(T0, N, Nn)
(Qe(t))[1:N, 1:Ne] = fill(T0, N, Ne)
end
sts = [vec(Tn); vec(Qe)]
e0 = Equation[inner_heatport.T[i] ~ Tn[i, 1] for i in 1:N]
e1 = Equation[outer_heatport.T[i] ~ Tn[i, Nn] for i in 1:N]
e2 = Equation[Qe[i, j] ~ Ke[j] * (-Tn[i, j+1] + Tn[i, j]) for i in 1:N for j in 1:Ne]
e3 = Equation[D(Tn[i, 1]) * (Cn[1] + C_shift) ~ inner_heatport.Q[i] - Qe[i, 1] for i in 1:N]
e4 = Equation[D(Tn[i, j]) * Cn[j] ~ Qe[i, j-1] - Qe[i, j] for i in 1:N for j in 2:Nn-1]
e5 = Equation[D(Tn[i, Nn]) * Cn[Nn] ~ Qe[i, Ne] + outer_heatport.Q[i] for i in 1:N]
eqs = vcat(e0, e1, e2, e3, e4, e5)
ODESystem(eqs, t, sts, p; systems=[inner_heatport, outer_heatport], name=name)
end

function CylindricalSurfaceConvection(; name, L=100, N=100, d=1.0, α=5.0)
dx = L / N
S = pi * d * dx
@named heatport = VectorHeatPort(N=N)
sts = @variables Tenv(t) = 0.0
eqs = [
Tenv ~ 18.0
[heatport.Q[i] ~ α * S * (heatport.T[i] - Tenv) for i in 1:N]
]

ODESystem(eqs, t, sts, []; systems=[heatport], name=name)
end

function PreinsulatedPipe(; name, L=100.0, N=100.0, dn=0.05, T0=0.0, t_layer=[0.004, 0.013],
λ=[50, 0.04], cp=[500, 1200], ρ=[7800, 40], α=5.0,
e=1e-4, lumped_T=50, diffusion=true)
@named inlet = FluidPort()
@named outlet = FluidPort()
@named fluid_region = FluidRegion(L=L, N=N, dn=dn, e=e, lumped_T=lumped_T, diffusion=diffusion)
@named shell = CircularWallFEM(L=L, N=N, d=dn, t_layer=t_layer, λ=λ, cp=cp, ρ=ρ)
@named surfconv = CylindricalSurfaceConvection(L=L, N=N, d=dn + 2.0 * sum(t_layer), α=α)
systems = [fluid_region, shell, inlet, outlet, surfconv]
eqs = [
connect(fluid_region.inlet, inlet)
connect(fluid_region.outlet, outlet)
connect(fluid_region.heatport, shell.inner_heatport)
connect(shell.outer_heatport, surfconv.heatport)
]
ODESystem(eqs, t, [], []; systems=systems, name=name)
end

function Source(; name, p_feed=100000)
@named outlet = FluidPort()
sts = @variables m_flow(t) = 1e-6
eqs = [
m_flow ~ m_flow_source(t)
outlet.m ~ -m_flow
outlet.p ~ p_feed
outlet.T ~ T_source(t)
]
compose(ODESystem(eqs, t, sts, []; name=name), [outlet])
end

function Sink(; name)
@named inlet = FluidPort()
eqs = [
inlet.T ~ instream(inlet.T)
]
compose(ODESystem(eqs, t, [], []; name=name), [inlet])
end

function TestBenchPreinsulated(; name, L=1.0, dn=0.05, t_layer=[0.0056, 0.013], N=100, diffusion=true, lumped_T=20)
@named pipe = PreinsulatedPipe(L=L, dn=dn, N=N, diffusion=diffusion, t_layer=t_layer, lumped_T=lumped_T)
@named source = Source()
@named sink = Sink()
subs = [source, pipe, sink]
eqs = [
connect(source.outlet, pipe.inlet)
connect(pipe.outlet, sink.inlet)
]
compose(ODESystem(eqs, t, [], []; name=name), subs)
end
```

To circumvent scaling issues with ModelingToolkit.jl, we will simplify and compile the system
using JuliaSimCompiler.jl and trace through the resultant function to obtain the expression
whose jacobian will be benchmarked.

```julia
function build_all_problems(N)
return map(N) do n
@named testbench = TestBenchPreinsulated(L=470, N=n, dn=0.3127, t_layer=[0.0056, 0.058])
sys_jsir = structural_simplify(IRSystem(testbench), loop=false)
return ODEProblem(sys_jsir, [], (0.0, 1.0))
end
end

function trace_and_jacobian!(trace_times, jac_times, prob, i; xname, pname, tname)
t = only(@independent_variables $tname)
n_states = length(prob.u0)
x = only(@variables $xname(t)[1:n_states])
n_params = length(prob.p)
p = only(@parameters $pname[1:n_params])
xvec = collect(Symbolics.unwrap(x))

expr = zeros(Num, n_states)
trace_times[i] = @elapsed prob.f.f(expr, x, p, t)

expr = Vector{Union{SymbolicUtils.BasicSymbolic{Real}, Float64}}(map(Symbolics.unwrap, expr))
jac_times[i] = @elapsed jac = Symbolics.jacobian(expr, xvec)

return jac
end
```

```julia
N = [5, 10, 20, 40, 60, 80, 160];
problems = build_all_problems(N)

hashconsing_jac_times = fill(NaN, length(N))
hashconsing_trace_times = fill(NaN, length(N))
no_hashconsing_jac_times = fill(NaN, length(N))
no_hashconsing_trace_times = fill(NaN, length(N))
```

# Timings

```julia
for (i, prob) in enumerate(problems)
# to prevent cached hashconsing entries from previously run jacobians
# affecting the timing of this jacobian, the names of the symbolic variables
# are gensym-ed
xname = gensym(:x)
pname = gensym(:p)
tname = gensym(:t)
jac1 = trace_and_jacobian!(hashconsing_trace_times, hashconsing_jac_times, prob, i; xname, pname, tname)
SymbolicUtils.ENABLE_HASHCONSING[] = false
jac2 = trace_and_jacobian!(no_hashconsing_trace_times, no_hashconsing_jac_times, prob, i; xname, pname, tname)
SymbolicUtils.ENABLE_HASHCONSING[] = true
@assert isequal(jac1, jac2)
end
```

# Generate plots

```julia
f = Figure(size=(800,1200));
names = ["Without hashconsing", "With hashconsing"]
let ax = Axis(f[1, 1]; yscale = log10, xscale = log10, title = "Tracing time")
_lines = [lines!(N, no_hashconsing_trace_times), lines!(N, hashconsing_trace_times)]
Legend(f[1, 2], _lines, names)
end
let ax = Axis(f[2, 1]; yscale = log10, xscale = log10, title = "Jacobian time")
_lines = [lines!(N, no_hashconsing_jac_times), lines!(N, hashconsing_jac_times)]
Legend(f[2, 2], _lines, names)
end
f
```

## Appendix

```julia, echo = false
using SciMLBenchmarks
SciMLBenchmarks.bench_footer(WEAVE_ARGS[:folder],WEAVE_ARGS[:file])
```
Loading