Merge pull request #2384 from chriselrod/cachehash

Cache hash of ConnectionElement

Author: ChrisRackauckas

src/systems/abstractsystem.jl

@@ -187,11 +187,14 @@ function SymbolicIndexingInterface.is_variable(sys::AbstractSystem, sym)
     if unwrap(sym) isa Int    # [x, 1] coerces 1 to a Num
         return unwrap(sym) in 1:length(unknown_states(sys))
     end
-    return any(isequal(sym), unknown_states(sys)) || hasname(sym) && is_variable(sys, getname(sym))
+    return any(isequal(sym), unknown_states(sys)) ||
+           hasname(sym) && is_variable(sys, getname(sym))
 end
 
 function SymbolicIndexingInterface.is_variable(sys::AbstractSystem, sym::Symbol)
-    return any(isequal(sym), getname.(unknown_states(sys))) || count('₊', string(sym)) == 1 && count(isequal(sym), Symbol.(sys.name, :₊, getname.(unknown_states(sys)))) == 1
+    return any(isequal(sym), getname.(unknown_states(sys))) ||
+           count('₊', string(sym)) == 1 &&
+           count(isequal(sym), Symbol.(sys.name, :₊, getname.(unknown_states(sys)))) == 1
 end
 
 function SymbolicIndexingInterface.variable_index(sys::AbstractSystem, sym)
@@ -224,12 +227,14 @@ function SymbolicIndexingInterface.is_parameter(sys::AbstractSystem, sym)
         return unwrap(sym) in 1:length(parameters(sys))
     end
 
-    return any(isequal(sym), parameters(sys)) || hasname(sym) && is_parameter(sys, getname(sym))
+    return any(isequal(sym), parameters(sys)) ||
+           hasname(sym) && is_parameter(sys, getname(sym))
 end
 
 function SymbolicIndexingInterface.is_parameter(sys::AbstractSystem, sym::Symbol)
     return any(isequal(sym), getname.(parameters(sys))) ||
-        count('₊', string(sym)) == 1 && count(isequal(sym), Symbol.(sys.name, :₊, getname.(parameters(sys)))) == 1
+           count('₊', string(sym)) == 1 &&
+           count(isequal(sym), Symbol.(sys.name, :₊, getname.(parameters(sys)))) == 1
 end
 
 function SymbolicIndexingInterface.parameter_index(sys::AbstractSystem, sym)
@@ -270,7 +275,8 @@ function SymbolicIndexingInterface.independent_variable_symbols(sys::AbstractSys
 end
 
 function SymbolicIndexingInterface.is_observed(sys::AbstractSystem, sym)
-    return !is_variable(sys, sym) && !is_parameter(sys, sym) && !is_independent_variable(sys, sym) && symbolic_type(sym) != NotSymbolic()
+    return !is_variable(sys, sym) && !is_parameter(sys, sym) &&
+           !is_independent_variable(sys, sym) && symbolic_type(sym) != NotSymbolic()
 end
 
 SymbolicIndexingInterface.is_time_dependent(::AbstractTimeDependentSystem) = true
@@ -621,13 +627,25 @@ function namespace_expr(O, sys, n = nameof(sys); ivs = independent_variables(sys
         O
     end
 end
-
+_nonum(@nospecialize x) = x isa Num ? x.val : x
 function states(sys::AbstractSystem)
     sts = get_states(sys)
     systems = get_systems(sys)
-    unique(isempty(systems) ?
-           sts :
-           [sts; reduce(vcat, namespace_variables.(systems))])
+    nonunique_states = if isempty(systems)
+        sts
+    else
+        system_states = reduce(vcat, namespace_variables.(systems))
+        isempty(sts) ? system_states : [sts; system_states]
+    end
+    isempty(nonunique_states) && return nonunique_states
+    # `Vector{Any}` is incompatible with the `SymbolicIndexingInterface`, which uses
+    # `elsymtype = symbolic_type(eltype(_arg))` 
+    # which inappropriately returns `NotSymbolic()`
+    if nonunique_states isa Vector{Any}
+        nonunique_states = _nonum.(nonunique_states)
+    end
+    @assert typeof(nonunique_states) !== Vector{Any}
+    unique(nonunique_states)
 end
 
 function parameters(sys::AbstractSystem)

src/systems/connectors.jl

@@ -147,18 +147,39 @@ struct ConnectionElement
     sys::LazyNamespace
     v::Any
     isouter::Bool
+    h::UInt
 end
-Base.nameof(l::ConnectionElement) = renamespace(nameof(l.sys), getname(l.v))
-function Base.hash(l::ConnectionElement, salt::UInt)
-    hash(nameof(l.sys)) ⊻ hash(l.v) ⊻ hash(l.isouter) ⊻ salt
+function _hash_impl(sys, v, isouter)
+    hashcore = hash(nameof(sys)) ⊻ hash(getname(v))
+    hashouter = isouter ? hash(true) : hash(false)
+    hashcore ⊻ hashouter
+end
+function ConnectionElement(sys::LazyNamespace, v, isouter::Bool)
+    ConnectionElement(sys, v, isouter, _hash_impl(sys, v, isouter))
 end
+Base.nameof(l::ConnectionElement) = renamespace(nameof(l.sys), getname(l.v))
 Base.isequal(l1::ConnectionElement, l2::ConnectionElement) = l1 == l2
 function Base.:(==)(l1::ConnectionElement, l2::ConnectionElement)
     nameof(l1.sys) == nameof(l2.sys) && isequal(l1.v, l2.v) && l1.isouter == l2.isouter
 end
+function Base.hash(e::ConnectionElement, salt::UInt)
+    @inbounds begin
+        @boundscheck begin
+            @assert e.h === _hash_impl(e.sys, e.v, e.isouter)
+        end
+    end
+    e.h ⊻ salt
+end
 namespaced_var(l::ConnectionElement) = states(l, l.v)
 states(l::ConnectionElement, v) = states(copy(l.sys), v)
 
+function withtrueouter(e::ConnectionElement)
+    e.isouter && return e
+    # we undo the xor
+    newhash = (e.h ⊻ hash(false)) ⊻ hash(true)
+    ConnectionElement(e.sys, e.v, true, newhash)
+end
+
 struct ConnectionSet
     set::Vector{ConnectionElement} # namespace.sys, var, isouter
 end
@@ -353,9 +374,7 @@ function partial_merge(csets::AbstractVector{<:ConnectionSet}, allouter = false)
     merged = false
     for (j, cset) in enumerate(csets)
         if allouter
-            cset = ConnectionSet(map(cset.set) do e
-                @set! e.isouter = true
-            end)
+            cset = ConnectionSet(map(withtrueouter, cset.set))
         end
         idx = nothing
         for e in cset.set
@@ -390,7 +409,7 @@ end
 
 function generate_connection_equations_and_stream_connections(csets::AbstractVector{
         <:ConnectionSet,
-    })
+})
     eqs = Equation[]
     stream_connections = ConnectionSet[]
 

src/systems/jumps/jumpsystem.jl

@@ -491,7 +491,7 @@ function (ratemap::JumpSysMajParamMapper{
         U,
         V,
         W,
-    })(params) where {U <: AbstractArray,
+})(params) where {U <: AbstractArray,
         V <: AbstractArray, W}
     updateparams!(ratemap, params)
     [convert(W, value(substitute(paramexpr, ratemap.subdict)))

src/systems/nonlinear/initializesystem.jl

@@ -52,4 +52,4 @@ function initializesystem(sys::ODESystem; name = nameof(sys), kwargs...)
         kwargs...)
 
     return sys_nl
-end
+end

src/systems/optimization/modelingtoolkitize.jl

@@ -27,7 +27,7 @@ function modelingtoolkitize(prob::DiffEqBase.OptimizationProblem; kwargs...)
             for i in 1:num_cons
                 if !isinf(prob.lcons[i])
                     if prob.lcons[i] != prob.ucons[i]
-                       push!(cons, prob.lcons[i] ≲ lhs[i])
+                        push!(cons, prob.lcons[i] ≲ lhs[i])
                     else
                         push!(cons, lhs[i] ~ prob.ucons[i])
                     end

src/systems/optimization/optimizationsystem.jl

@@ -269,7 +269,7 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
         expression = Val{false})
 
     obj_expr = subs_constants(objective(sys))
-    
+
     if grad
         grad_oop, grad_iip = generate_gradient(sys, checkbounds = checkbounds,
             linenumbers = linenumbers,

src/variables.jl

@@ -463,4 +463,4 @@ function get_default_or_guess(x)
     else
         return getguess(x)
     end
-end
+end

test/nonlinearsystem.jl

@@ -237,7 +237,7 @@ end
 @testset "Initialization System" begin
     # Define the Lotka Volterra system which begins at steady state
     @parameters t
-    pars = @parameters a=1.5 b=1.0 c=3.0 d=1.0 dx_ss = 1e-5
+    pars = @parameters a=1.5 b=1.0 c=3.0 d=1.0 dx_ss=1e-5
 
     vars = @variables begin
         dx(t),
@@ -274,4 +274,4 @@ end
 
     # Confirm for all the states of the simplified system
     @test all(.≈(sol[states(sys_simple)], [1e-5 / 1.5, 0]; atol = 1e-8))
-end
+end

test/optimizationsystem.jl

@@ -289,10 +289,17 @@ end
     @parameters a b
     loss = (a - x)^2 + b * (y - x^2)^2
     @named sys = OptimizationSystem(loss, [x, y], [a, b], constraints = [x^2 + y^2 ≲ 0.0])
-    @test_throws ArgumentError OptimizationProblem(sys, [x => 0.0, y => 0.0], [a => 1.0, b => 100.0], lcons = [0.0])
-    @test_throws ArgumentError OptimizationProblem(sys, [x => 0.0, y => 0.0], [a => 1.0, b => 100.0], ucons = [0.0])
+    @test_throws ArgumentError OptimizationProblem(sys,
+        [x => 0.0, y => 0.0],
+        [a => 1.0, b => 100.0],
+        lcons = [0.0])
+    @test_throws ArgumentError OptimizationProblem(sys,
+        [x => 0.0, y => 0.0],
+        [a => 1.0, b => 100.0],
+        ucons = [0.0])
 
     prob = OptimizationProblem(sys, [x => 0.0, y => 0.0], [a => 1.0, b => 100.0])
     @test prob.f.expr isa Symbolics.Symbolic
-    @test all(prob.f.cons_expr[i].lhs isa Symbolics.Symbolic for i in 1:length(prob.f.cons_expr)) 
-end
+    @test all(prob.f.cons_expr[i].lhs isa Symbolics.Symbolic
+              for i in 1:length(prob.f.cons_expr))
+end