feat: add symbolic jacobian benchmark to ThermalFluid

Author: AayushSabharwal

benchmarks/ModelingToolkit/ThermalFluid.jmd

@@ -10,8 +10,10 @@ 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, XSteam, Polynomials, BenchmarkTools, CairoMakie, OrdinaryDiffEq
+using ModelingToolkit, JuliaSimCompiler, Symbolics, XSteam, Polynomials, BenchmarkTools, OrdinaryDiffEq
 using OMJulia
+
+using ModelingToolkit: t_nounits as t, D_nounits as D
 ```
 
 ## Setup Julia Code
@@ -24,8 +26,6 @@ using OMJulia
 #Source -> o--o--o--o--o--o--o -> Sink
 #       advection diff source PDE
 
-@variables t
-D = Differential(t)
 m_flow_source(t) = 2.75
 T_source(t) = (t > 12 * 3600) * 56.0 + 12.0
 @register_symbolic m_flow_source(t)
@@ -410,6 +410,72 @@ Legend(f2[1,2], _lines, names)
 f2
 ```
 
+## Symbolic jacobian timings
+
+To avoid scaling issues with `structural_simplify` on ModelingToolkit.jl models, we will
+compile functions using JuliaSimCompiler.jl and trace through the resultant functions.
+
+```julia
+function run_and_time_jacobian!(mtk_times, mtk_no_hashconsing_times, N)
+  @named testbench = TestBenchPreinsulated(L=470, N=N, dn=0.3127, t_layer=[0.0056, 0.058])
+  sys = structural_simplify(IRSystem(testbench), loop=false)
+  # tspan doesn't matter
+  prob = ODEProblem(sys, [], (0.0, 1.0); sparse = true)
+
+  t = ModelingToolkit.t_nounits
+  n_states = length(prob.u0)
+  @variables x(t)[1:n_states]
+  n_params = length(prob.p)
+  @variables p[1:n_params]
+
+  x = collect(x)
+  p = collect(p)
+  expr = zeros(Num, n_states)
+  prob.f.f(expr, x, p, t)
+
+  expr = map(Symbolics.unwrap, expr)
+  x = map(Symbolics.unwrap, x)
+  mtk_time = @elapsed jac1 = Symbolics.jacobian(expr, x)
+  push!(mtk_times, mtk_time)
+
+  SymbolicUtils.ENABLE_HASHCONSING[] = false
+  @variables x(t)[1:n_states]
+  @variables p[1:n_params]
+  x = collect(x)
+  p = collect(p)
+  expr = zeros(Num, n_states)
+  prob.f.f(expr, x, p, t)
+  expr = map(Symbolics.unwrap, expr)
+  no_hashconsing_time = @elapsed jac2 = Symbolics.jacobian(expr, x)
+  push!(mtk_no_hashconsing_times, no_hashconsing_time)
+  SymbolicUtils.ENABLE_HASHCONSING[] = true
+
+  @assert all(isequal.(jac1, jac2))
+end
+```
+
+
+```julia
+@time run_and_time_jacobian!(Float64[], Float64[], 4) # precompile
+mtk_times = Float64[]
+mtk_no_hashconsing_times = Float64[]
+for n in N
+  run_and_time_jacobian!(Float64[], Float64[], n)
+end
+```
+
+## Generate Plots
+
+```julia
+f3 = Figure(size=(800,1200));
+ss_names = ["MTK", "JSIR"];
+let ax = Axis(f3[1, 1]; yscale = log10, xscale = log10, title = "Jacobian Time")
+  _lines = [lines!(N, mtk_times), lines!(N, mtk_no_hashconsing_times)]
+  Legend(f3[1, 2], _lines, ["With Hashconsing", "Without Hashconsing"])
+end
+f3
+```
+
 ## Appendix
 
 ```julia, echo = false