feat: add new Symbolics jacobian benchmark

Author: AayushSabharwal

benchmarks/Symbolics/Project.toml

@@ -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"

benchmarks/Symbolics/ThermalFluid.jmd

@@ -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])
+```