From 8e183acbb8210d933b61539d740e66360b69a4c1 Mon Sep 17 00:00:00 2001 From: Aayush Sabharwal Date: Thu, 23 Jan 2025 21:57:51 +0530 Subject: [PATCH] feat: add new Symbolics jacobian benchmark --- benchmarks/Symbolics/Project.toml | 24 ++ benchmarks/Symbolics/ThermalFluid.jmd | 332 ++++++++++++++++++++++++++ 2 files changed, 356 insertions(+) create mode 100644 benchmarks/Symbolics/Project.toml create mode 100644 benchmarks/Symbolics/ThermalFluid.jmd diff --git a/benchmarks/Symbolics/Project.toml b/benchmarks/Symbolics/Project.toml new file mode 100644 index 000000000..e5e62b3c8 --- /dev/null +++ b/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" diff --git a/benchmarks/Symbolics/ThermalFluid.jmd b/benchmarks/Symbolics/ThermalFluid.jmd new file mode 100644 index 000000000..fd085be51 --- /dev/null +++ b/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]) +```