Refactor to better support type promotion due to thinks like Dual Numbers (#30)

* Refactor to better support type promotion due to thinks like ForwardDiff.Dual

* bump version

Author: MasonProtter

Project.toml

@@ -1,6 +1,6 @@
 name = "GraphDynamics"
 uuid = "bcd5d0fe-e6b7-4ef1-9848-780c183c7f4c"
-version = "0.4.1"
+version = "0.4.2"
 
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"

README.md

@@ -39,8 +39,7 @@ function GraphDynamics.to_subsystem(p::Particle)
     # Use `name`, `m`, and `q` as parameters
     # Every subsystem should have a unique name symbol.
     params = SubsystemParams{Particle}(
-        ;name,
-        m,
+        ;m,
         q,
     )
     # Assemble a Subsystem
@@ -82,8 +81,7 @@ function GraphDynamics.to_subsystem(p::Oscillator)
     # Use `name`, `m`, `k`, `x₀`, and `q` as parameters
     # Every subsystem should have a unique name symbol.
     params = SubsystemParams{Oscillator}(
-        ;name,
-        m,
+        ;m,
         k,
         x₀,
         q,

src/GraphDynamics.jl

@@ -300,11 +300,21 @@ abstract type ConnectionRule end
 (c::ConnectionRule)(src, dst, t) = c(src, dst)
 Base.zero(::T) where {T <: ConnectionRule} = zero(T)
 
-struct NotConnected <: ConnectionRule end
-(::NotConnected)(l, r) = zero(promote_type(eltype(l), eltype(r)))
-struct ConnectionMatrix{N, CR, Tup <: NTuple{N, NTuple{N, Union{NotConnected, AbstractMatrix{CR}}}}}
+struct NotConnected{CR <: ConnectionRule} end
+Base.getindex(::NotConnected{CR}, inds...) where {CR} = zero(CR)
+Base.copy(c::NotConnected) = c
+struct ConnectionMatrix{N, CR, Tup <: NTuple{N, NTuple{N, Union{NotConnected{CR}, AbstractMatrix{CR}}}}}
     data::Tup
 end
+function Base.copy(c::ConnectionMatrix)
+    data′ = map(c.data) do col
+        map(col) do mat
+            copy(mat)
+        end
+    end
+    ConnectionMatrix(data′)
+end
+
 struct ConnectionMatrices{NConn, Tup <: NTuple{NConn, ConnectionMatrix}}
     matrices::Tup
 end
@@ -314,7 +324,7 @@ Base.getindex(m::ConnectionMatrix, ::Val{i}, ::Val{j}) where {i, j} = m.data[i][
 @inline Base.getindex(m::ConnectionMatrices, i) = m.matrices[i]
 Base.length(m::ConnectionMatrices) = length(m.matrices)
 Base.size(m::ConnectionMatrix{N}) where {N} = (N, N)
-
+Base.copy(c::ConnectionMatrices) = ConnectionMatrices(map(copy, c.matrices))
 
 #----------------------------------------------------------
 # Infrastructure for subsystems

src/graph_system.jl

@@ -60,8 +60,11 @@ end
 function system_wiring_rule!(g, node)
     add_node!(g, node)
 end
-function system_wiring_rule!(g, src, dst; conn, kwargs...)
-    add_connection!(g, src, dst; conn, kwargs...)
+function system_wiring_rule!(g, src, dst; kwargs...)
+    if !haskey(kwargs, :conn)
+        error("conn keyword argument not specified for connection between $src and $dst")
+    end
+    add_connection!(g, src, dst; conn=kwargs[:conn], kwargs...)
 end
 
 @kwdef struct PartitionedGraphSystem{CM <: ConnectionMatrices, S, P, EVT, Ns, EP, SNM, PNM, CNM}
@@ -140,10 +143,10 @@ function PartitionedGraphSystem(g::GraphSystem)
 	@named n2 = SysType1(x=1, y=3)
 	@named n3 = SysType2(a=1, b=2, c=3)
     
-    add_connection!(g, n1, n2; conn=C1(1))
-    add_connection!(g, n2, n3; conn=C1(2))
-    add_connection!(g, n3, n1; conn=C2(3))
-    add_connection!(g, n3, n2; conn=C3(4))
+    add_connection!(g, n1, n2; conn=Conn1(1))
+    add_connection!(g, n2, n3; conn=Conn1(2))
+    add_connection!(g, n3, n1; conn=Conn2(3))
+    add_connection!(g, n3, n2; conn=Conn2(4))
     
     we'd get
     connection_matrix_1 = Conn1[⎡. 1⎤⎡.⎤
@@ -226,7 +229,7 @@ function make_connection_matrices(g_flat, nodes_partitioned=make_partitioned_nod
                     end
                 end
                 rule_matrix = if isempty(conns)
-                    NotConnected()#{CT}(length(nodeks), length(nodeis))
+                    NotConnected{CT}() #{CT}(length(nodeks), length(nodeis))
                 else
                     sparse(ls, js, conns, length(nodeks), length(nodeis))
                 end

src/problems.jl

@@ -27,20 +27,10 @@ function  SciMLBase.ODEProblem(g::PartitionedGraphSystem, u0map, tspan, param_ma
     end
     tstops = vcat(tstops, nt.tstops)
     prob = ODEProblem{true, SciMLBase.FullSpecialize}(f, u, tspan, p; callback, tstops, kwargs...)
-    # let ukeys = map(first, u0map),
-    #     uvals = map(last, u0map)
-    #     setu(prob, ukeys)(prob, uvals)
-    # end
-    # let pkeys = map(first, param_map),
-    #     pvals = map(last, param_map)
-    #     setp(prob, pkeys)(prob, pvals)
-    # end
-     for (k, v) ∈ u0map
-         setu(prob, k)(prob, v)
-     end
-    for (k, v) ∈ param_map
-        setp(prob, k)(prob, v)
+    for (k, v) ∈ u0map
+        setu(prob, k)(prob, v)
     end
+    @reset prob.p = set_params!!(prob.p, param_map)
     prob
 end
 
@@ -58,27 +48,32 @@ function SciMLBase.SDEProblem(g::PartitionedGraphSystem, u0map, tspan, param_map
         uvals = map(last, u0map)
         setu(prob, ukeys)(prob, uvals)
     end
-    let pkeys = map(first, param_map),
-        pvals = map(last, param_map)
-        setp(prob, pkeys)(prob, pvals)
-    end
+    @reset prob.p = set_params!!(prob.p, param_map)
     prob
 end
 
-Base.@kwdef struct GraphSystemParameters{PP, CM, S, PAP, DEC, EP<:NamedTuple}
+Base.@kwdef struct GraphSystemParameters{PP, CM, S, PAP, DEC, NP, SNM, PNM, CNM, EP<:NamedTuple}
     params_partitioned::PP
     connection_matrices::CM
     scheduler::S
     partition_plan::PAP
     discrete_event_cache::DEC
+    names_partitioned::NP
+    state_namemap::SNM
+    param_namemap::PNM
+    compu_namemap::CNM
     extra_params::EP=(;)
 end
 
 function _problem(g::PartitionedGraphSystem, tspan; scheduler, allow_nonconcrete, u0map, param_map, global_events)
     (; states_partitioned,
      params_partitioned,
      connection_matrices,
-     tstops) = g
+     tstops,
+     names_partitioned,
+     state_namemap,
+     param_namemap,
+     compu_namemap) = g
 
     total_eltype = let
         states_eltype = mapreduce(promote_type, states_partitioned) do v
@@ -87,12 +82,24 @@ function _problem(g::PartitionedGraphSystem, tspan; scheduler, allow_nonconcrete
         u0map_eltype = mapreduce(promote_type, u0map; init=Union{}) do (k, v)
             typeof(v)
         end
-        promote_type(states_eltype, u0map_eltype)
+        numeric_params_eltype = mapreduce(promote_type, params_partitioned) do v
+            if isconcretetype(eltype(v))
+                promote_numeric_param_eltype(eltype(v))
+            else
+                mapreduce(promote_type, v) do params
+                    promote_numeric_param_eltype(typeof(params))
+                end
+            end
+        end
+        numeric_param_map_eltype = let numeric_params_from_map = [v for (_, v) in param_map if v isa Number]
+            mapreduce(typeof, promote_type, numeric_params_from_map; init=Union{})
+        end
+        promote_type(states_eltype, u0map_eltype, numeric_params_eltype, numeric_param_map_eltype)
     end
 
     re_eltype(s::SubsystemStates{T}) where {T} = convert(SubsystemStates{T, total_eltype}, s) 
     states_partitioned = map(states_partitioned) do v
-        if eltype(eltype(v)) <: total_eltype && eltype(eltype(v)) !== Union{}
+        if eltype(eltype(v)) <: total_eltype
             v
         else
             re_eltype.(v)
@@ -159,7 +166,11 @@ function _problem(g::PartitionedGraphSystem, tspan; scheduler, allow_nonconcrete
                               connection_matrices,
                               scheduler,
                               partition_plan,
-                              discrete_event_cache)
+                              discrete_event_cache,
+                              names_partitioned,
+                              state_namemap,
+                              param_namemap,
+                              compu_namemap)
 
     (; f, u, tspan, p, callback, tstops)
 end

src/subsystems.jl

@@ -14,16 +14,23 @@ function ConstructionBase.getproperties(s::SubsystemParams)
 end
 
 function ConstructionBase.setproperties(s::SubsystemParams{T}, patch::NamedTuple) where {T}
+    set_param_prop(s, patch; allow_typechange=false)
+end
+function set_param_prop(s::SubsystemParams{T}, key, val; allow_typechange=false) where {T}
+    set_param_prop(s, NamedTuple{(key,)}(val); allow_typechange)
+end
+function set_param_prop(s::SubsystemParams{T}, patch; allow_typechange=false) where {T}
     props = NamedTuple(s)
     props′ = merge(props, patch)
-    if typeof(props) != typeof(props′)
+    if typeof(props) != typeof(props′) && !allow_typechange
         param_setproperror(props, props′)
     end
     SubsystemParams{T}(props′)
 end
+
 @noinline function param_setproperror(props, props′)
     error("Type unstable change to subsystem params! Expected properties of type\n  $(typeof(props))\nbut got\n  $(typeof(props′))")
-end 
+end
 
 get_tag(::SubsystemParams{Name}) where {Name} = Name
 get_tag(::Type{<:SubsystemParams{Name}}) where {Name} = Name
@@ -37,6 +44,9 @@ end
 function Base.convert(::Type{SubsystemParams{Name, NT}}, p::SubsystemParams{Name}) where {Name, NT}
     SubsystemParams{Name}(convert(NT, NamedTuple(p)))
 end
+@generated function promote_numeric_param_eltype(::Type{SubsystemParams{Name, NamedTuple{props, Tup}}}) where {Name, props, Tup}
+    :(promote_type($(param for param in Tup.parameters if param <: Number)...))
+end
 
 #------------------------------------------------------------
 # Subsystem states
@@ -156,8 +166,9 @@ function Base.convert(::Type{Subsystem{Name, Eltype}}, s::Subsystem{Name}) where
 end
 
 @generated function promote_nt_type(::Type{NamedTuple{names, T1}},
-                                    ::Type{NamedTuple{names, T2}}) where {names, T1, T2}
-    NamedTuple{names, Tuple{(promote_type(T1.parameters[i], T2.parameters[i]) for i ∈ eachindex(names))...}}
+                         ::Type{NamedTuple{names, T2}}) where {names, T1, T2}
+    proms = [:(promote_type($(T1.parameters[i]), $(T2.parameters[i]))) for i in eachindex(names)]
+    :(NamedTuple{names, Tuple{$(proms...)}})
 end
 
 function Base.promote_rule(::Type{SubsystemParams{Name, NT1}},

src/symbolic_indexing.jl

@@ -200,3 +200,36 @@ end
 #         :t
 #     )
 # end
+
+function SymbolicIndexingInterface.remake_buffer(sys, oldbuffer::GraphSystemParameters, idxs, vals)
+    newbuffer = @set oldbuffer.params_partitioned = copy.(oldbuffer.params_partitioned)
+    set_params!!(newbuffer, zip(idxs, vals))
+end
+
+function set_params!!(buffer::GraphSystemParameters, param_map)
+    (; params_partitioned, param_namemap) = buffer
+    for (key, val) ∈ param_map
+        let (;tup_index, v_index, prop) = param_namemap[key]
+            params = params_partitioned[tup_index][v_index]
+            params_new = set_param_prop(params, prop, val; allow_typechange=true)
+            peltype = eltype(params_partitioned[tup_index])
+            if !(typeof(params_new) <: peltype)
+                new_eltype = promote_type(typeof(params_new), peltype)
+                @reset params_partitioned[tup_index] = convert.(new_eltype, params_partitioned[tup_index])
+            end
+            params_partitioned[tup_index][v_index] = params_new
+        end
+    end
+    @reset buffer.params_partitioned = params_partitioned#re_eltype_params(params_partitioned)
+end
+
+function re_eltype_params(params_partitioned)
+    map(params_partitioned) do v
+        ptype = mapreduce(typeof, promote_type, v)
+        if ptype == eltype(v)
+            v
+        else
+            convert.(ptype, v)
+        end
+    end
+end

test/Project.toml

@@ -1,5 +1,6 @@
 [deps]
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 GraphDynamics = "bcd5d0fe-e6b7-4ef1-9848-780c183c7f4c"
 OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"