Merge pull request #223 from JuliaDynamics/hw/despecialize

despecialize VertexModel/EdgeModel

Author: hexaeder

.github/workflows/tests.yml

@@ -8,7 +8,7 @@ on:
 concurrency:
   # Skip intermediate builds: always.
   group: ${{ github.workflow }}-${{ github.ref }}
-  cancel-in-progress: true # allways cancle old runs on new push
+  cancel-in-progress: true # always cancle old runs on new push
 jobs:
   test:
     name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }}
@@ -38,6 +38,28 @@ jobs:
       - uses: julia-actions/cache@v2
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
+        with:
+          coverage: false
+
+  test-coverage:
+    name: Test Coverage
+    runs-on: ubuntu-latest
+    permissions: # needed to allow julia-actions/cache to proactively delete old caches that it has created
+      actions: write
+      contents: read
+    steps:
+      - uses: actions/checkout@v4
+      - uses: julia-actions/setup-julia@v2
+        with:
+          version: 1
+          arch: x64
+      - uses: julia-actions/cache@v2
+      - uses: julia-actions/julia-buildpkg@v1
+      - uses: julia-actions/julia-runtest@v1
+        env:
+          JULIA_COVERAGE: true
+        with:
+          coverage: true
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v4
         with:

Project.toml

@@ -59,7 +59,7 @@ KernelAbstractions = "0.9.18"
 LinearAlgebra = "1"
 MacroTools = "0.5.15"
 Mixers = "0.1.2"
-ModelingToolkit = "9.61"
+ModelingToolkit = "9.67"
 NNlib = "0.9.13"
 NonlinearSolve = "4"
 Polyester = "0.7.12"

docs/examples/cascading_failure.jl

@@ -245,8 +245,8 @@ prob = ODEProblem(nw, uflat(u0), (0,6), pflat(u0); callback=cbset)
 Main.test_execution_styles(prob) # testing all ex styles #src
 sol2 = solve(prob, Tsit5());
 ## we want to test the reconstruction of the observables # hide
-@test all(!iszero, sol2(sol2.t; idxs=eidxs(sol2,:,:P))[begin]) # hide
-@test all(iszero, sol2(sol2.t; idxs=eidxs(sol2,:,:P))[end][[1:5...,7]]) # hide
+@test all(!iszero, sol2(sol2.t; idxs=eidxs(sol2,:,:P)).u[begin]) # hide
+@test all(iszero, sol2(sol2.t; idxs=eidxs(sol2,:,:P)).u[end][[1:5...,7]]) # hide
 nothing #hide
 
 # Then again we plot the solution:

ext/MTKExt.jl

@@ -475,11 +475,9 @@ end
 
 _all_rhs_symbols(eqs) = mapreduce(eq->get_variables(eq.rhs), ∪, eqs, init=Set{Symbolic}())
 
-using PrecompileTools: @setup_workload, @compile_workload
-@setup_workload begin
-    @compile_workload begin
-        # include("precompile_workload.jl")
-    end
+using PrecompileTools: @compile_workload
+@compile_workload begin
+    include("precompile_workload.jl")
 end
 
 end

ext/precompile_workload.jl

@@ -1,30 +1,28 @@
-is_precompiling() = ccall(:jl_generating_output, Cint, ()) == 1
-
 using NetworkDynamics
 using ModelingToolkit
 using ModelingToolkit: t_nounits as t, D_nounits as Dt
 
 @connector Terminal begin
-    u_r(t), [description="d-voltage"]
-    u_i(t), [description="q-voltage"]
-    i_r(t), [description="d-current", connect=Flow]
-    i_i(t), [description="q-current", connect=Flow]
+    u_r(t)
+    u_i(t)
+    i_r(t), [connect=Flow]
+    i_i(t), [connect=Flow]
 end
 @mtkmodel Swing begin
     @components begin
         terminal = Terminal()
     end
     @variables begin
-        ω(t)=0.0, [description="Rotor frequency"]
-        θ(t)=0.0, [description="Rotor angle"]
-        Pel(t), [description="Electrical Power injected into the grid"]
+        ω(t), [description="Rotor frequency", guess=0]
+        θ(t), [description="Rotor angle", guess=0]
+        Pel(t)
     end
     @parameters begin
-        M=1, [description="Inertia"]
-        D=0.1, [description="Damping"]
-        V=1.0, [description="Voltage magnitude"]
-        ω_ref=0, [description="Reference frequency"]
-        Pm, [description="Mechanical Power"]
+        M=0.005
+        D=0.1
+        V, [guess=1]
+        ω_ref=0
+        Pm, [description="Mechanical Power", guess=-1]
     end
     @equations begin
         Dt(θ) ~ ω - ω_ref
@@ -36,26 +34,10 @@ end
 end
 @mtkmodel BusBase begin
     @variables begin
-        u_r(t)=1, [description="bus d-voltage", output=true]
-        u_i(t)=0, [description="bus q-voltage", output=true]
-        i_r(t), [description="bus d-current (flowing into bus)", input=true]
-        i_i(t), [description="bus d-current (flowing into bus)", input=true]
-        P(t), [description="bus active power (flowing into network)"]
-        Q(t), [description="bus reactive power (flowing into network)"]
-        u_mag(t), [description="bus voltage magnitude"]
-        u_arg(t), [description="bus voltage argument"]
-        i_mag(t), [description="bus current magnitude"]
-        i_arg(t), [description="bus current argument"]
-        ω(t), [description="bus angular frequency"]
-    end
-    @equations begin
-        P ~ u_r * (-i_r) + u_i * (-i_i)
-        Q ~ u_i * (-i_r) - u_r * (-i_i)
-        u_mag ~ sqrt(u_r^2 + u_i^2)
-        u_arg ~ atan(u_i, u_r)
-        i_mag ~ sqrt(i_r^2 + i_i^2)
-        i_arg ~ atan(i_i, i_r)
-        ω ~ Dt(u_arg)
+        u_r(t)=1.0
+        u_i(t)=0.0
+        i_r(t)=1.0
+        i_i(t)=0.0
     end
 end
 @mtkmodel BusBar begin
@@ -80,10 +62,5 @@ end
     end
 end
 @named swingbus = Bus()
-
-if is_precompiling()
-    VertexModel(swingbus, [:busbar₊i_r, :busbar₊i_i], [:busbar₊u_r, :busbar₊u_i]; verbose=false)
-else
-    @info "ODEVertex"
-    @time @eval VertexModel(swingbus, [:busbar₊i_r, :busbar₊i_i], [:busbar₊u_r, :busbar₊u_i]; verbose=false)
-end
+vm = VertexModel(swingbus, [:busbar₊i_r, :busbar₊i_i], [:busbar₊u_r, :busbar₊u_i]; verbose=false)
+NetworkDynamics.initialize_component!(vm; verbose=false)

src/NetworkDynamics.jl

@@ -125,11 +125,9 @@ function reloadfaces!()
 end
 # NetworkDynamics.reloadfaces!()
 
-using PrecompileTools: @setup_workload, @compile_workload
-@setup_workload begin
-    @compile_workload begin
-        # include("precompile_workload.jl")
-    end
+using PrecompileTools: @compile_workload
+@compile_workload begin
+    include("precompile_workload.jl")
 end
 
 end

src/callbacks.jl

@@ -235,7 +235,7 @@ function wrap_component_callbacks(nw)
     end
     # group the callbacks such that they are in groups which are "batchequal"
     # batchequal groups can be wrapped into a single callback
-    idx_per_type = _find_identical(callbacks, 1:length(components))
+    idx_per_type = find_identical(callbacks; equality=_batchequal)
     batches = []
     for typeidx in idx_per_type
         batchcomps = components[typeidx]
@@ -251,21 +251,22 @@ function wrap_component_callbacks(nw)
     end
     return batches
 end
-function batchequal(a::ContinousComponentCallback, b::ContinousComponentCallback)
-    batchequal(a.condition, b.condition) || return false
-    batchequal(a.kwargs, b.kwargs)       || return false
+_batchequal(a, b) = false
+function _batchequal(a::ContinousComponentCallback, b::ContinousComponentCallback)
+    _batchequal(a.condition, b.condition) || return false
+    _batchequal(a.kwargs, b.kwargs)       || return false
     return true
 end
-function batchequal(a::VectorContinousComponentCallback, b::VectorContinousComponentCallback)
-    batchequal(a.condition, b.condition) || return false
-    batchequal(a.kwargs, b.kwargs)       || return false
+function _batchequal(a::VectorContinousComponentCallback, b::VectorContinousComponentCallback)
+    _batchequal(a.condition, b.condition) || return false
+    _batchequal(a.kwargs, b.kwargs)       || return false
     a.len == b.len                       || return false
     return true
 end
-function batchequal(a::T, b::T) where {T <: Union{ComponentCondition, ComponentAffect}}
+function _batchequal(a::T, b::T) where {T <: Union{ComponentCondition, ComponentAffect}}
     typeof(a) == typeof(b)
 end
-function batchequal(a::NamedTuple, b::NamedTuple)
+function _batchequal(a::NamedTuple, b::NamedTuple)
     length(a) == length(b) || return false
     for (k, v) in a
         haskey(b, k) || return false

src/component_functions.jl

@@ -247,22 +247,22 @@ Directed(s::AbstractVector{<:Symbol}) = AnnotatedSym(Directed, s)
 
 abstract type ComponentModel end
 
-struct VertexModel{F,G,FFT,OF,MM,EX<:Union{Nothing,Vector{<:SymbolicIndex}}} <: ComponentModel
+struct VertexModel <: ComponentModel
     name::Symbol
     # main function
-    f::F
+    f::Any
     sym::Vector{Symbol}
-    mass_matrix::MM
+    mass_matrix::Any
     # outputs
-    g::G
+    g::Any
     outsym::Vector{Symbol}
-    ff::FFT
+    ff::FeedForwardType
     # parameters, optional input sym and optional external inputs
     psym::Vector{Symbol}
     insym::Union{Nothing, Vector{Symbol}}
-    extin::EX
+    extin::Union{Nothing, Vector{<:SymbolicIndex}}
     # observed
-    obsf::OF
+    obsf::Any
     obssym::Vector{Symbol}
     # metadata
     symmetadata::Dict{Symbol,Dict{Symbol, Any}}
@@ -299,25 +299,25 @@ Optional Arguments:
 
 All Symbol arguments can be used to set default values, i.e. `psym=[:K=>1, :p]`.
 """
-VertexModel(; kwargs...) = _construct_comp(VertexModel, Base.inferencebarrier(kwargs))
-VertexModel(v::VertexModel; kwargs...) = _reconstruct_comp(VertexModel, v, Base.inferencebarrier(kwargs))
+VertexModel(; kwargs...) = _construct_comp(VertexModel, kwargs)
+VertexModel(v::VertexModel; kwargs...) = _reconstruct_comp(VertexModel, v, kwargs)
 
-struct EdgeModel{F,G,FFT,OF,MM,EX<:Union{Nothing,Vector{<:SymbolicIndex}}} <: ComponentModel
+struct EdgeModel <: ComponentModel
     name::Symbol
     # main function
-    f::F
+    f::Any
     sym::Vector{Symbol}
-    mass_matrix::MM
+    mass_matrix::Any
     # outputs
-    g::G
+    g::Any
     outsym::@NamedTuple{src::Vector{Symbol},dst::Vector{Symbol}}
-    ff::FFT
+    ff::FeedForwardType
     # parameters, optional input sym and optional external inputs
     psym::Vector{Symbol}
     insym::Union{Nothing, @NamedTuple{src::Vector{Symbol},dst::Vector{Symbol}}}
-    extin::EX
+    extin::Union{Nothing, Vector{<:SymbolicIndex}}
     # observed
-    obsf::OF
+    obsf::Any
     obssym::Vector{Symbol}
     # metadata
     symmetadata::Dict{Symbol,Dict{Symbol, Any}}
@@ -357,8 +357,8 @@ Optional Arguments:
 
 All Symbol arguments can be used to set default values, i.e. `psym=[:K=>1, :p]`.
 """
-EdgeModel(; kwargs...) = _construct_comp(EdgeModel, Base.inferencebarrier(kwargs))
-EdgeModel(v::EdgeModel; kwargs...) = _reconstruct_comp(EdgeModel, v, Base.inferencebarrier(kwargs))
+EdgeModel(; kwargs...) = _construct_comp(EdgeModel, kwargs)
+EdgeModel(v::EdgeModel; kwargs...) = _reconstruct_comp(EdgeModel, v, kwargs)
 
 """
     compf(c::ComponentModel)
@@ -576,30 +576,17 @@ Fills up type parameters with `nothing` to ensure `Core.Compiler.isconstType`
 for GPU compatibility.
 """
 dispatchT(::T) where {T<:ComponentModel} = dispatchT(T)
-dispatchT(T::Type{<:VertexModel}) = VertexModel{nothing,nothing,nothing,nothing,nothing,Nothing}
-dispatchT(T::Type{<:EdgeModel}) = EdgeModel{nothing,nothing,nothing,nothing,nothing,Nothing}
-
-# TODO: introduce batchequal hash for faster batching of component models
-batchequal(a, b) = false
-function batchequal(a::ComponentModel, b::ComponentModel)
-    compf(a)  == compf(b)  || return false
-    compg(a)  == compg(b)  || return false
-    fftype(a) == fftype(b) || return false
-    dim(a)    == dim(b)    || return false
-    outdim(a) == outdim(b) || return false
-    pdim(a)   == pdim(b)   || return false
-    extdim(a) == extdim(b) || return false
-    return true
-end
+dispatchT(T::Type{<:VertexModel}) = VertexModel
+dispatchT(T::Type{<:EdgeModel}) = EdgeModel
 
 """
     _construct_comp(::Type{T}, kwargs) where {T}
 
 Internal function to construct a component model from keyword arguments.
 Fills up kw arguments with default values and performs sanity checks.
 """
-function _construct_comp(::Type{T}, @nospecialize(kwargs)) where {T}
-    dict = _fill_defaults(T, Base.inferencebarrier(kwargs))
+Base.@nospecializeinfer function _construct_comp(::Type{T}, @nospecialize(kwargs)) where {T}
+    dict = _fill_defaults(T, kwargs)
 
     # check signature of f
     # if !_valid_signature(T, dict[:f])
@@ -625,7 +612,7 @@ function _construct_comp(::Type{T}, @nospecialize(kwargs)) where {T}
     return c
 end
 
-function _reconstruct_comp(::Type{T}, cf::ComponentModel, kwargs) where {T}
+Base.@nospecializeinfer function _reconstruct_comp(::Type{T}, cf::ComponentModel, @nospecialize(kwargs)) where {T}
     fields = fieldnames(T)
     dict = Dict{Symbol, Any}()
     for f in fields
@@ -643,7 +630,7 @@ end
 Fill up keyword arguments `kwargs` for type T with default values.
 Also perfoms sanity check some properties like mass matrix, depth, ...
 """
-function _fill_defaults(T, @nospecialize(kwargs))
+Base.@nospecializeinfer function _fill_defaults(T, @nospecialize(kwargs))
     dict = Dict{Symbol, Any}(kwargs)
     allow_output_sym_clash = pop!(dict, :allow_output_sym_clash, false)
 
@@ -1121,11 +1108,12 @@ function Base.:(==)(cf1::ComponentModel, cf2::ComponentModel)
     typeof(cf1) == typeof(cf2) && equal_fields(cf1, cf2)
 end
 
-function compfg(c)
+# force specialization on f, g, fft
+compfg(c) = _compfg(compf(c), compg(c), fftype(c))
+function _compfg(f::F, g::G, fft::FFT) where {F, G, FFT}
     (outs, du, u, ins, p, t) -> begin
-        f = compf(c)
         isnothing(f) || f(du, u, ins..., p, t)
-        compg(c)(_gargs(fftype(c), outs, du, u, ins, p, t)...)
+        g(_gargs(fft, outs, du, u, ins, p, t)...)
         nothing
     end
 end