feat: cache subexpressions dependent only on previous SCCs

Author: AayushSabharwal

src/systems/nonlinear/nonlinearsystem.jl

@@ -579,9 +579,11 @@ function SCCNonlinearFunction{iip}(
     f(resid, u, p) = f_iip(resid, u, p)
     f(resid, u, p::MTKParameters) = f_iip(resid, u, p...)
 
-    subsys = NonlinearSystem(_eqs, _dvs, ps; observed = _obs, parameter_dependencies = parameter_dependencies(sys), name = nameof(sys))
+    subsys = NonlinearSystem(_eqs, _dvs, ps; observed = _obs,
+        parameter_dependencies = parameter_dependencies(sys), name = nameof(sys))
     if get_index_cache(sys) !== nothing
-        @set! subsys.index_cache = subset_unknowns_observed(get_index_cache(sys), sys, _dvs, getproperty.(_obs, (:lhs,)))
+        @set! subsys.index_cache = subset_unknowns_observed(
+            get_index_cache(sys), sys, _dvs, getproperty.(_obs, (:lhs,)))
         @set! subsys.complete = true
     end
 
@@ -620,8 +622,7 @@ function SciMLBase.SCCNonlinearProblem{iip}(sys::NonlinearSystem, u0map,
     explicitfuns = []
     nlfuns = []
     prevobsidxs = Int[]
-    cachevars = []
-    cacheexprs = []
+    cachesize = 0
     for (i, (escc, vscc)) in enumerate(zip(eq_sccs, var_sccs))
         # subset unknowns and equations
         _dvs = dvs[vscc]
@@ -632,6 +633,26 @@ function SciMLBase.SCCNonlinearProblem{iip}(sys::NonlinearSystem, u0map,
         setdiff!(obsidxs, prevobsidxs)
         _obs = obs[obsidxs]
 
+        # get all subexpressions in the RHS which we can precompute in the cache
+        state = Dict()
+        for i in eachindex(_obs)
+            _obs[i] = _obs[i].lhs ~ subexpressions_not_involving_vars!(
+                _obs[i].rhs, _dvs, state)
+        end
+        for i in eachindex(_eqs)
+            _eqs[i] = _eqs[i].lhs ~ subexpressions_not_involving_vars!(
+                _eqs[i].rhs, _dvs, state)
+        end
+
+        # cached variables and their corresponding expressions
+        cachevars = [obs[i].lhs for i in prevobsidxs]
+        cacheexprs = [obs[i].rhs for i in prevobsidxs]
+        for (k, v) in state
+            push!(cachevars, v)
+            push!(cacheexprs, k)
+        end
+        cachesize = max(cachesize, length(cachevars))
+
         if isempty(cachevars)
             push!(explicitfuns, Returns(nothing))
         else
@@ -651,7 +672,9 @@ function SciMLBase.SCCNonlinearProblem{iip}(sys::NonlinearSystem, u0map,
         append!(prevobsidxs, obsidxs)
     end
 
-    p = rebuild_with_caches(p, BufferTemplate(eltype(u0), length(cachevars)))
+    if cachesize != 0
+        p = rebuild_with_caches(p, BufferTemplate(eltype(u0), cachesize))
+    end
 
     subprobs = []
     for (f, vscc) in zip(nlfuns, var_sccs)

src/utils.jl

@@ -1042,3 +1042,78 @@ function observed_equations_used_by(sys::AbstractSystem, exprs)
     sort!(obsidxs)
     return obsidxs
 end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Given an expression `expr`, return a dictionary mapping subexpressions of `expr` that do
+not involve variables in `vars` to anonymous symbolic variables. Also return the modified
+`expr` with the substitutions indicated by the dictionary. If `expr` is a function
+of only `vars`, then all of the returned subexpressions can be precomputed.
+"""
+function subexpressions_not_involving_vars(expr, vars)
+    expr = unwrap(expr)
+    vars = map(unwrap, vars)
+    state = Dict()
+    newexpr = subexpressions_not_involving_vars!(expr, vars, state)
+    return state, newexpr
+end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Mutating version of `subexpressions_not_involving_vars` which writes to `state`. Only
+returns the modified `expr`.
+"""
+function subexpressions_not_involving_vars!(expr, vars, state::Dict{Any, Any})
+    expr = unwrap(expr)
+    symbolic_type(expr) == NotSymbolic() && return expr
+    iscall(expr) || return expr
+    Symbolics.shape(expr) == Symbolics.Unknown() && return expr
+    haskey(state, expr) && return state[expr]
+    if !any(x -> occursin(x, expr), vars)
+        sym = gensym(:subexpr)
+        stype = symtype(expr)
+        var = similar_variable(expr, sym)
+        state[expr] = var
+        return var
+    end
+    op = operation(expr)
+    args = arguments(expr)
+    if (op == (+) || op == (*)) && symbolic_type(expr) !== ArraySymbolic()
+        indep_args = []
+        dep_args = []
+        for arg in args
+            if any(x -> occursin(x, arg), vars)
+                push!(dep_args, subexpressions_not_involving_vars!(arg, vars, state))
+            else
+                push!(indep_args, arg)
+            end
+        end
+        indep_term = reduce(op, indep_args; init = Int(op == (*)))
+        indep_term = subexpressions_not_involving_vars!(indep_term, vars, state)
+        dep_term = reduce(op, dep_args; init = Int(op == (*)))
+        return op(indep_term, dep_term)
+    end
+    newargs = map(args) do arg
+        symbolic_type(arg) != NotSymbolic() || is_array_of_symbolics(arg) || return arg
+        subexpressions_not_involving_vars!(arg, vars, state)
+    end
+    return maketerm(typeof(expr), op, newargs, metadata(expr))
+end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Create an anonymous symbolic variable of the same shape, size and symtype as `var`, with
+name `gensym(name)`. Does not support unsized array symbolics.
+"""
+function similar_variable(var::BasicSymbolic, name = :anon)
+    name = gensym(name)
+    stype = symtype(var)
+    sym = Symbolics.variable(name; T = stype)
+    if size(var) !== ()
+        sym = setmetadata(sym, Symbolics.ArrayShapeCtx, map(Base.OneTo, size(var)))
+    end
+    return sym
+end

test/scc_nonlinear_problem.jl

@@ -142,3 +142,20 @@ end
     @test sol.u≈sccsol.u atol=1e-10
 end
 
+@testset "Expression caching" begin
+    @variables x[1:4] = rand(4)
+    val = Ref(0)
+    function func(x, y)
+        val[] += 1
+        x + y
+    end
+    @register_symbolic func(x, y)
+    @mtkbuild sys = NonlinearSystem([0 ~ x[1]^3 + x[2]^3 - 5
+                                     0 ~ sin(x[1] - x[2]) - 0.5
+                                     0 ~ func(x[1], x[2]) * exp(x[3]) - x[4]^3 - 5
+                                     0 ~ func(x[1], x[2]) * exp(x[4]) - x[3]^3 - 4])
+    sccprob = SCCNonlinearProblem(sys, [])
+    sccsol = solve(sccprob, NewtonRaphson())
+    @test SciMLBase.successful_retcode(sccsol)
+    @test val[] == 1
+end