diff --git a/benchmark/run_benchmarks.jl b/benchmark/run_benchmarks.jl index 453e11a1..78ec9818 100755 --- a/benchmark/run_benchmarks.jl +++ b/benchmark/run_benchmarks.jl @@ -116,8 +116,8 @@ if bench_target || bench_baseline @info "Dirty directory, add everything to new commit!" @assert realpath(pwd()) == realpath(ndpath_tmp) "Julia is in $(pwd()) not it $ndpath_tmp" run(`git status`) - run(`git config --global user.email "benchmarkbot@example.com"`) - run(`git config --global user.name "Benchmark Bot"`) + run(`git config --local user.email "benchmarkbot@example.com"`) + run(`git config --local user.name "Benchmark Bot"`) run(`git checkout -b $(randstring(15))`) run(`git add -A`) run(`git commit -m "tmp commit for benchmarking"`) diff --git a/docs/make.jl b/docs/make.jl index 1551f397..7a91a7a3 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -36,6 +36,7 @@ makedocs(; "metadata.md", "initialization.md", "mtk_integration.md", + "external_inputs.md", ], "API.md", "Tutorials" => [ diff --git a/docs/src/external_inputs.md b/docs/src/external_inputs.md new file mode 100644 index 00000000..5d66bb1b --- /dev/null +++ b/docs/src/external_inputs.md @@ -0,0 +1,210 @@ +# External Inputs +External inputs for components are way to pass signals between components outside of the network structure. +The most common usecase for that are control systems: make your vertex `i` depend on some state of vertex `j`. + +If used, this will essentially wides the number of received inputs of a component function. I.e. the baseline mathematical models for vertex models + +```math +\begin{aligned} +M^{\mathrm v}\,\frac{\mathrm{d}}{\mathrm{d}t}x^{\mathrm v} &= f^{\mathrm v}(u^{\mathrm v}, i^{\mathrm v}, i^{\mathrm{ext}}, p^{\mathrm v}, t)\\ +y^{\mathrm v} &= g^{\mathrm v}(u^{\mathrm v}, i^{\mathrm v}, i^{\mathrm{ext}}, p^{\mathrm v}, t) +\end{aligned} +``` +```julia +fᵥ(dxᵥ, xᵥ, e_aggr, ext, pᵥ, t) +gᵥ(yᵥ, xᵥ, [e_aggr, ext,] pᵥ, t) +``` +and edge models +```math +\begin{aligned} +M^{\mathrm e}\,\frac{\mathrm{d}}{\mathrm{d}t}x^{\mathrm e} &= f^{\mathrm e}(u^{\mathrm e}, y^{\mathrm v}_{\mathrm{src}}, y^{\mathrm v}_{\mathrm{dst}}, i^{\mathrm{ext}}, p^{\mathrm e}, t)\\ +y^{\mathrm e}_{\mathrm{dst}} &= g_\mathrm{dst}^{\mathrm e}(u^{\mathrm e}, y^{\mathrm v}_{\mathrm{src}}, y^{\mathrm v}_{\mathrm{dst}}, i^{\mathrm{ext}}, p^{\mathrm e}, t)\\ +y^{\mathrm e}_{\mathrm{src}} &= g_\mathrm{src}^{\mathrm e}(u^{\mathrm e}, y^{\mathrm v}_{\mathrm{src}}, y^{\mathrm v}_{\mathrm{dst}}, i^{\mathrm{ext}}, p^{\mathrm e}, t) +\end{aligned} +``` +```julia +fₑ(dxₑ, xₑ, v_src, v_dst, ext, pₑ, t) +gₑ(y_src, y_dst, xᵥ, [v_src, v_dst, ext,] pₑ, t) +``` +change. You may still oomit the input section from `g` according to the different [Feed Forward Behavior](@ref)s. However you either have to use *all* inputs (including `ext`) or none. + +## Usage +Vertex and Edge models with external inputs can be created using the `extin` keyword of the [`EdgeModel`](@ref) and [`VertexModel`](@ref) constructors. + +You need to pass a vector of `SymbolicIndices` ([`VIndex`](@ref) and [`EIndex`](@ref)), e.g. +```julia +VertexModel(... , extin=[VIndex(12,:x), VIndex(9, :x)], ...) +``` +means your vertex receives a 2 dimensional external input vector with the states `x` of vertices 12 and 9. +See below for hands on example. + +## Limitations +There are some limitations in place. You can only reference **states** (i.e. things that appear in `xᵥ` or `xₑ` of some component model) or **outputs of non-feed-forward components**, i.e. states `yᵥ` or `yₑ` of some component model which does not have feed forward behavior in their `g` function. + +## Example +As an example system, we'll consider two capacitors with a resistor between them. +Vertex 1 `v1` has a controllable current source. +Using a PI controller for the current source, it tries to keep the voltage at the +second vertex stable under the disturbance of some time periodic current sind at `v2`. + +``` + + v1 Resistor v2 +PI controlled ─→─o─←────MMM────→─o─→─ time dependent +current source ┴ ┴ current sink + ┬ ┬ + ⏚ ⏚ +``` + +The example will be implemented 2 times: in plain NetworkDynamcics and using MTK. + +### Plain NetworkDynamics + +First we define the resistor: +```@example extin +using NetworkDynamics +using OrdinaryDiffEqTsit5 +using CairoMakie + +function resistor_g(idst, vsrc, vdst, (R,), t) + idst[1] = (vsrc[1] - vdst[1])/R +end +resistor = EdgeModel(g=AntiSymmetric(resistor_g), + outsym=:i, insym=:v, psym=:R=>0.1, + src=:source, dst=:load, name=:resistor) +``` + +Then we define the "load" vertex with sinusoidial load profile: +```@example extin +function f_load(dv, v, iin, (C,), t) + dv[1] = 1/C*(iin[1] - (1 + 0.1*sin(t))) +end +load = VertexModel(f=f_load, g=1, + sym=:V=>0.5, insym=:i, psym=:C=>1, + vidx=2, name=:load) +``` +Lastly, we define the "source" vertex +```@example extin +function f_source(dv, v, iin, extin, (C,Vref,Ki,Kp), t) + Δ = Vref - extin[1] # tracking error + dv[2] = Δ # integrator state of PI + i_inj = Kp*Δ + Ki*v[2] # controller output + dv[1] = 1/C*(iin[1] + i_inj) +end +source = VertexModel(f=f_source, g=1, + sym=[:V=>0.5,:ξ=>0], insym=:i, + psym=[:C=>1, :Vref=>1, :Ki=>0.5, :Kp=>10], + extin=[VIndex(:load, :V)], + vidx=1, name=:source) +``` +Then we can create the network and simulate: +```@example extin +nw = Network([source, load], [resistor]) +u0 = NWState(nw) # everything has default values +prob = ODEProblem(nw, uflat(u0), (0,100), pflat(u0)) +sol = solve(prob, Tsit5()) + +fig, ax, p = lines(sol, idxs=VIndex(:load, :V), label="Voltage @ load"); +lines!(ax, sol, idxs=VPIndex(:source, :Vref), label="Reference", color=Cycled(2)); +axislegend(ax; position=:rb); +fig # hide +``` + +### MTK Models +First we define the resistor: +```@example extin_mtk +using NetworkDynamics +using OrdinaryDiffEqTsit5 +using ModelingToolkit +using ModelingToolkit: t_nounits as t, D_nounits as Dt +using CairoMakie + +@mtkmodel Resistor begin + @variables begin + i(t), [description="Current at dst end"] + V_src(t), [description="Voltage at src end"] + V_dst(t), [description="Voltage at dst end"] + end + @parameters begin + R=0.1, [description="Resistance"] + end + @equations begin + i ~ (V_src - V_dst)/R + end +end +@named resistor = Resistor() +resistor_edge = EdgeModel(resistor, [:V_src], [:V_dst], AntiSymmetric([:i]); src=:load, dst=:source) +``` + +Then we define the "load" vertex with sinusoidial load profile: +```@example extin_mtk +@mtkmodel Load begin + @variables begin + V(t)=0.5, [description="Voltage"] + i_load(t), [description="Load current"] + i_grid(t), [description="Current from grid"] + end + @parameters begin + C=1, [description="Capacitance"] + end + @equations begin + i_load ~ 1 + 0.1*sin(t) + Dt(V) ~ 1/C*(i_grid - i_load) + end +end +@named load = Load() +load_vertex = VertexModel(load, [:i_grid], [:V]; vidx=2) +``` +Lastly, we define the "source" vertex +```@example extin_mtk +@mtkmodel Source begin + @variables begin + V(t)=0.5, [description="Voltage"] + ξ(t)=0, [description="Integrator state"] + i_grid(t), [description="Current from grid"] + i_source(t), [description="Current from source"] + Δ(t), [description="Tracking Error"] + V_load(t), [description="Voltage at load"] + end + @parameters begin + C=1, [description="Capacitance"] + Vref=1, [description="Reference voltage"] + Ki=0.5, [description="Integral gain"] + Kp=10, [description="Proportional gain"] + end + @equations begin + Δ ~ Vref - V_load + Dt(ξ) ~ Δ + i_source ~ Kp*Δ + Ki*ξ + Dt(V) ~ 1/C*(i_grid + i_source) + end +end +@named source = Source() +source_vertex = VertexModel(source, [:i_grid], [:V]; + extin=[:V_load => VIndex(:load, :V)], vidx=1) +``` +Then we can create the network and simulate: +```@example extin_mtk +nw = Network([source_vertex, load_vertex], [resistor_edge]) +u0 = NWState(nw) # everything has default values +prob = ODEProblem(nw, uflat(u0), (0,100), pflat(u0)) +sol = solve(prob, Tsit5()) + +let + fig = Figure(); + ax1 = Axis(fig[1,1]); + lines!(ax1, sol, idxs=VIndex(:load, :V), label="Voltage @ load"); + lines!(ax1, sol, idxs=VPIndex(:source, :Vref), label="Reference", color=Cycled(2)); + axislegend(ax1; position=:rb); + ax2 = Axis(fig[2,1]); + lines!(ax2, sol, idxs=VIndex(:load, :i_load), label="load current"); + lines!(ax2, sol, idxs=VIndex(:source, :i_source), label="source current", color=Cycled(2)); + axislegend(ax2); + fig +end +``` +Using MTK for modeling, we can also inspect the currents `i_load` and `i_source` the MTK interface preserves the [Observables](@ref). + + + + diff --git a/ext/CUDAExt.jl b/ext/CUDAExt.jl index 21257356..f2f49c80 100644 --- a/ext/CUDAExt.jl +++ b/ext/CUDAExt.jl @@ -2,7 +2,7 @@ module CUDAExt using NetworkDynamics: Network, NetworkLayer, ComponentBatch, KAAggregator, AggregationMap, SparseAggregator, LazyGBufProvider, EagerGBufProvider, LazyGBuf, - dispatchT, iscudacompatible, executionstyle + dispatchT, iscudacompatible, executionstyle, ExtMap using NetworkDynamics.PreallocationTools: DiffCache using NetworkDynamics: KernelAbstractions as KA @@ -32,12 +32,14 @@ function Adapt.adapt_structure(to, n::Network) mm = adapt(to, n.mass_matrix) gbp = adapt(to, n.gbufprovider) caches = (;output = _adapt_diffcache(to, n.caches.output), - aggregation = _adapt_diffcache(to, n.caches.aggregation)) + aggregation = _adapt_diffcache(to, n.caches.aggregation), + external = _adapt_diffcache(to, n.caches.external)) exT = typeof(executionstyle(n)) gT = typeof(n.im.g) + extmap = adapt(to, n.extmap) - Network{exT,gT,typeof(layer),typeof(vb),typeof(mm),eltype(caches),typeof(gbp)}( - vb, layer, n.im, caches, mm, gbp) + Network{exT,gT,typeof(layer),typeof(vb),typeof(mm),eltype(caches),typeof(gbp),typeof(extmap)}( + vb, layer, n.im, caches, mm, gbp, extmap) end Adapt.@adapt_structure NetworkLayer @@ -80,6 +82,14 @@ function _adapt_eager_gbufp(mapto, cacheto, gbp) EagerGBufProvider(map, cache) end +#### +#### Adapt external input map +#### +Adapt.@adapt_structure ExtMap +function Adapt.adapt_structure(to::Type{<:CuArray{<:AbstractFloat}}, em::ExtMap) + adapt(CuArray, em) +end + #### #### Adapt VertexBatch/EdgeBatch @@ -90,7 +100,7 @@ end function Adapt.adapt_structure(to, b::ComponentBatch) indices = adapt(to, b.indices) ComponentBatch(dispatchT(b), indices, b.compf, b.compg, b.ff, - b.statestride, b.pstride, b.inbufstride, b.outbufstride) + b.statestride, b.pstride, b.inbufstride, b.outbufstride, b.extbufstride) end #### diff --git a/ext/MTKExt.jl b/ext/MTKExt.jl index 7d8e14df..10dcdb87 100644 --- a/ext/MTKExt.jl +++ b/ext/MTKExt.jl @@ -16,20 +16,36 @@ import NetworkDynamics: VertexModel, EdgeModel, AnnotatedSym include("MTKUtils.jl") """ - VertexModel(sys::ODESystem, inputs, outputs; ff_to_constraint=true, kwargs...) + VertexModel(sys::ODESystem, inputs, outputs; + verbose=false, name=getname(sys), extin=nothing, ff_to_constraint=true, kwargs...) Create a `VertexModel` object from a given `ODESystem` created with ModelingToolkit. You need to provide 2 lists of symbolic names (`Symbol` or `Vector{Symbols}`): - `inputs`: names of variables in you equation representing the aggregated edge states - `outputs`: names of variables in you equation representing the node output -`ff_to_constraint` controlls, whether output transformations `g` which depend on inputs should be +Additional kw arguments: +- `name`: Set name of the component model. Will be lifted from the ODESystem name. +- `extin=nothing`: Provide external inputs as pairs, i.e. `extin=[:extvar => VIndex(1, :a)]` + will bound the variable `extvar(t)` in the equations to the state `a` of the first vertex. +- `ff_to_constraint=true`: Controlls, whether output transformations `g` which depend on inputs should be + transformed into constraints. Defaults to true since ND.jl does not handle vertices with FF yet. """ -function VertexModel(sys::ODESystem, inputs, outputs; verbose=false, name=getname(sys), ff_to_constraint=true, kwargs...) +function VertexModel(sys::ODESystem, inputs, outputs; verbose=false, name=getname(sys), + ff_to_constraint=true, extin=nothing, kwargs...) warn_events(sys) inputs = inputs isa AbstractVector ? inputs : [inputs] outputs = outputs isa AbstractVector ? outputs : [outputs] - gen = generate_io_function(sys, (inputs,), (outputs,); verbose, ff_to_constraint) + + if isnothing(extin) + extin_nwidx = nothing + ins = (inputs, ) + else + extin_sym, extin_nwidx = _split_extin(extin) + ins = (inputs, extin_sym) + end + + gen = generate_io_function(sys, ins, (outputs,); verbose, ff_to_constraint) f = gen.f g = gen.g @@ -52,7 +68,8 @@ function VertexModel(sys::ODESystem, inputs, outputs; verbose=false, name=getnam mass_matrix = gen.mass_matrix c = VertexModel(;f, g, sym, insym, outsym, psym, obssym, - obsf, mass_matrix, ff=gen.fftype, name, allow_output_sym_clash=true, kwargs...) + obsf, mass_matrix, ff=gen.fftype, name, extin=extin_nwidx, + allow_output_sym_clash=true, kwargs...) set_metadata!(c, :observed, gen.observed) set_metadata!(c, :equations, gen.equations) set_metadata!(c, :outputeqs, gen.outputeqs) @@ -60,9 +77,9 @@ function VertexModel(sys::ODESystem, inputs, outputs; verbose=false, name=getnam end """ - EdgeModel(sys::ODESystem, srcin, dstin, AntiSymmetric(dstout); ff_to_constraint=false, kwargs...) + EdgeModel(sys::ODESystem, srcin, dstin, AntiSymmetric(dstout); kwargs...) -Create a `EdgeModel` object from a given `ODESystem` created with ModelingToolkit. +Create a `EdgeModel` object from a given `ODESystem` created with ModelingToolkit for **single sided models**. Here you only need to provide one list of output symbols: `dstout`. To make it clear how to handle the single-sided output definiton, you musst wrap @@ -71,13 +88,13 @@ the symbol vector in - `Symmetric(dstout)`, or - `Directed(dstout)`. -`ff_to_constraint` controlls, whether output transformations `g` which depend on inputs should be -transformed into constraints. +Additional `kwargs` are the same as for the double-sided EdgeModel MTK constructor. """ EdgeModel(sys::ODESystem, srcin, dstin, dstout; kwargs...) = EdgeModel(sys, srcin, dstin, nothing, dstout; kwargs...) """ - EdgeModel(sys::ODESystem, srcin, srcout, dstin, dstout; ff_to_constraint=false, kwargs...) + EdgeModel(sys::ODESystem, srcin, srcout, dstin, dstout; + verbose=false, name=getname(sys), extin=nothing, ff_to_constraint=false, kwargs...) Create a `EdgeModel` object from a given `ODESystem` created with ModelingToolkit. You need to provide 4 lists of symbolic names (`Symbol` or `Vector{Symbols}`): @@ -86,14 +103,27 @@ You need to provide 4 lists of symbolic names (`Symbol` or `Vector{Symbols}`): - `srcout`: names of variables in you equation representing the output at the source - `dstout`: names of variables in you equation representing the output at the destination -`ff_to_constraint` controlls, whether output transformations `g` which depend on inputs should be -transformed into constraints. +Additional kw arguments: +- `name`: Set name of the component model. Will be lifted from the ODESystem name. +- `extin=nothing`: Provide external inputs as pairs, i.e. `extin=[:extvar => VIndex(1, :a)]` + will bound the variable `extvar(t)` in the equations to the state `a` of the first vertex. +- `ff_to_constraint=false`: Controlls, whether output transformations `g` which depend on inputs should be + transformed into constraints. """ -function EdgeModel(sys::ODESystem, srcin, dstin, srcout, dstout; verbose=false, name=getname(sys), ff_to_constraint=false, kwargs...) +function EdgeModel(sys::ODESystem, srcin, dstin, srcout, dstout; verbose=false, name=getname(sys), + ff_to_constraint=false, extin=nothing, kwargs...) warn_events(sys) srcin = srcin isa AbstractVector ? srcin : [srcin] dstin = dstin isa AbstractVector ? dstin : [dstin] + if isnothing(extin) + extin_nwidx = nothing + ins = (srcin, dstin) + else + extin_sym, extin_nwidx = _split_extin(extin) + ins = (srcin, dstin, extin_sym) + end + singlesided = isnothing(srcout) if singlesided && !(dstout isa AnnotatedSym) throw(ArgumentError("If you only provide one output (single sided \ @@ -111,10 +141,10 @@ function EdgeModel(sys::ODESystem, srcin, dstin, srcout, dstout; verbose=false, outs = (srcout, dstout) end - gen = generate_io_function(sys, (srcin, dstin), outs; verbose, ff_to_constraint) + gen = generate_io_function(sys, ins, outs; verbose, ff_to_constraint) f = gen.f - g = singlesided ? gwrap(gen.g; ff=gen.fftype) : gen.g + g = singlesided ? gwrap(gen.g) : gen.g obsf = gen.obsf _sym = getname.(gen.states) @@ -145,7 +175,8 @@ function EdgeModel(sys::ODESystem, srcin, dstin, srcout, dstout; verbose=false, mass_matrix = gen.mass_matrix c = EdgeModel(;f, g, sym, insym, outsym, psym, obssym, - obsf, mass_matrix, ff=gen.fftype, name, allow_output_sym_clash=true, kwargs...) + obsf, mass_matrix, ff=gen.fftype, name, extin=extin_nwidx, + allow_output_sym_clash=true, kwargs...) set_metadata!(c, :observed, gen.observed) set_metadata!(c, :equations, gen.equations) set_metadata!(c, :outputeqs, gen.outputeqs) @@ -199,6 +230,18 @@ function _get_metadata(sys, name) nt end +function _split_extin(extin) + try + extin_sym = first.(extin) + extin_nwidx = last.(extin) + @assert extin_sym isa Vector{Symbol} + @assert extin_nwidx isa Vector{<:NetworkDynamics.SymbolicIndex} + return extin_sym, extin_nwidx + catch e + @error "Could not evaluate extin keyword argument!" + end +end + function generate_io_function(_sys, inputss::Tuple, outputss::Tuple; expression=Val{false}, verbose=false, ff_to_constraint) diff --git a/src/NetworkDynamics.jl b/src/NetworkDynamics.jl index ce2e5e48..5552f7ed 100644 --- a/src/NetworkDynamics.jl +++ b/src/NetworkDynamics.jl @@ -55,6 +55,7 @@ include("aggregators.jl") include("gbufs.jl") include("construction.jl") include("coreloop.jl") +include("external_inputs.jl") # XXX: have both, s[:] and uflat(s) ? export VIndex, EIndex, VPIndex, EPIndex, NWState, NWParameter, uflat, pflat diff --git a/src/component_functions.jl b/src/component_functions.jl index f7e387f2..06751882 100644 --- a/src/component_functions.jl +++ b/src/component_functions.jl @@ -99,9 +99,7 @@ fftype(::StateMask) = PureStateMap() end -abstract type OutputWrapper{FF} end -hasfftype(::OutputWrapper) = true -fftype(::OutputWrapper{FF}) where {FF} = FF() +abstract type SingleSidedOutputWrapper end """ AntiSymmetric(g_dst) @@ -116,14 +114,10 @@ output function which applies See also [`Symmetric`](@ref), [`Directed`](@ref), [`Fiducial`](@ref) and [`StateMask`](@ref). """ -struct AntiSymmetric{FF,G} <: OutputWrapper{FF} +struct AntiSymmetric{G} <: SingleSidedOutputWrapper g::G - function AntiSymmetric(g; ff=nothing) - ff = isnothing(ff) ? _infer_ss_fftype(g) : ff - new{typeof(ff), typeof(g)}(g) - end end -AntiSymmetric(g::Union{AbstractVector,Number}; ff=nothing) = AntiSymmetric(StateMask(g); ff) +AntiSymmetric(g::Union{AbstractVector,Number}) = AntiSymmetric(StateMask(g)) @inline function (c::AntiSymmetric)(osrc, odst, args...) @inline c.g(odst, args...) @inbounds for i in 1:length(osrc) @@ -145,14 +139,10 @@ output function which applies See also [`AntiSymmetric`](@ref), [`Directed`](@ref), [`Fiducial`](@ref) and [`StateMask`](@ref). """ -struct Symmetric{FF,G} <: OutputWrapper{FF} +struct Symmetric{G} <: SingleSidedOutputWrapper g::G - function Symmetric(g; ff=nothing) - ff = isnothing(ff) ? _infer_ss_fftype(g) : ff - new{typeof(ff), typeof(g)}(g) - end end -Symmetric(g::Union{AbstractVector,Number}; ff=nothing) = Symmetric(StateMask(g); ff) +Symmetric(g::Union{AbstractVector,Number}) = Symmetric(StateMask(g)) @inline function (c::Symmetric)(osrc, odst, args...) @inline c.g(odst, args...) @inbounds for i in 1:length(osrc) @@ -174,14 +164,10 @@ With `Directed` there is no output for the `src` side. See also [`AntiSymmetric`](@ref), [`Symmetric`](@ref), [`Fiducial`](@ref) and [`StateMask`](@ref). """ -struct Directed{FF,G} <: OutputWrapper{FF} +struct Directed{G} <: SingleSidedOutputWrapper g::G - function Directed(g; ff=nothing) - ff = isnothing(ff) ? _infer_ss_fftype(g) : ff - new{typeof(ff), typeof(g)}(g) - end end -Directed(g::Union{AbstractVector,Number}; ff=nothing) = Directed(StateMask(g); ff) +Directed(g::Union{AbstractVector,Number}) = Directed(StateMask(g)) @inline function (c::Directed)(osrc, odst, args...) @inline c.g(odst, args...) nothing @@ -200,26 +186,14 @@ into a double sided output function which applies See also [`AntiSymmetric`](@ref), [`Directed`](@ref), [`Fiducial`](@ref) and [`StateMask`](@ref). """ -struct Fiducial{FF,GS,GD} <: OutputWrapper{FF} +struct Fiducial{GS,GD} <: SingleSidedOutputWrapper src::GS dst::GD - function Fiducial(src, dst; ff=nothing) - if isnothing(ff) - ffsrc = _infer_ss_fftype(src) - ffdst = _infer_ss_fftype(dst) - if ffsrc == ffdst - ff = ffsrc - else - error("Both src and dst coupling functions have different ff types.") - end - end - new{typeof(ff), typeof(src), typeof(dst)}(src, dst) - end end -function Fiducial(src::Union{AbstractVector,Number}, dst::Union{AbstractVector,Number}; ff=nothing) - Fiducial(StateMask(src), StateMask(dst); ff) +function Fiducial(src::Union{AbstractVector,Number}, dst::Union{AbstractVector,Number}) + Fiducial(StateMask(src), StateMask(dst)) end -Fiducial(;src, dst, ff=nothing) = Fiducial(src, dst; ff) +Fiducial(;src, dst) = Fiducial(src, dst) @inline function (c::Fiducial)(osrc, odst, args...) @inline c.src(osrc, args...) @@ -273,7 +247,7 @@ Directed(s::AbstractVector{<:Symbol}) = AnnotatedSym(Directed, s) abstract type ComponentModel end -struct VertexModel{F,G,FFT,OF,MM} <: ComponentModel +struct VertexModel{F,G,FFT,OF,MM,EX<:Union{Nothing,Vector{<:SymbolicIndex}}} <: ComponentModel name::Symbol # main function f::F @@ -283,13 +257,14 @@ struct VertexModel{F,G,FFT,OF,MM} <: ComponentModel g::G outsym::Vector{Symbol} ff::FFT - # parameters and option input sym + # parameters, optional input sym and optional external inputs psym::Vector{Symbol} insym::Union{Nothing, Vector{Symbol}} + extin::EX # observed obsf::OF obssym::Vector{Symbol} - # optional insyms + # metadata symmetadata::Dict{Symbol,Dict{Symbol, Any}} metadata::Dict{Symbol,Any} # cached symbol collections @@ -318,13 +293,16 @@ Optional Arguments: - `ff`: `FeedForwardType` of component. Will be typically infered from `g` automaticially. - `obssym`/`obsf`: Define additional "observable" states. - `symmetadata`/`metadata`: Provide prefilled metadata dictionaries. +- `extin=nothing`: + Define "external" inputs for the model with Network indices, i.e. `extin=[VIndex(7,:x), ..]`. + Those inputs will be provided as another input vector `f(x, in, extin, p, t)` and `g(y, x, in, extin, p, t)`. 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)) -struct EdgeModel{F,G,FFT,OF,MM} <: ComponentModel +struct EdgeModel{F,G,FFT,OF,MM,EX<:Union{Nothing,Vector{<:SymbolicIndex}}} <: ComponentModel name::Symbol # main function f::F @@ -334,9 +312,10 @@ struct EdgeModel{F,G,FFT,OF,MM} <: ComponentModel g::G outsym::@NamedTuple{src::Vector{Symbol},dst::Vector{Symbol}} ff::FFT - # parameters and option input sym + # parameters, optional input sym and optional external inputs psym::Vector{Symbol} insym::Union{Nothing, @NamedTuple{src::Vector{Symbol},dst::Vector{Symbol}}} + extin::EX # observed obsf::OF obssym::Vector{Symbol} @@ -372,6 +351,9 @@ Optional Arguments: - `ff`: `FeedForwardType` of component. Will be typically infered from `g` automaticially. - `obssym`/`obsf`: Define additional "observable" states. - `symmetadata`/`metadata`: Provide prefilled metadata dictionaries. +- `extin=nothing`: + Define "external" inputs for the model with Network indices, i.e. `extin=[VIndex(7,:x), ..]`. + Those inputs will be provided as another input vector `f(x, insrc, indst, extin, p, t)` and `g(ysrc, ydst, x, insrc, indst, extin, p, t)`. All Symbol arguments can be used to set default values, i.e. `psym=[:K=>1, :p]`. """ @@ -515,6 +497,20 @@ Gives the input dimension(s). indim(c::VertexModel)::Int = length(insym(c)) indim(c::EdgeModel)::@NamedTuple{src::Int,dst::Int} = (; src=length(insym(c).src), dst=length(insym(c).dst)) +""" + extin(c::ComponentModel) + +Retrieve the external input symbols of the component. +""" +extin(c::ComponentModel) = c.extin + +""" + extdim(c::ComponentModel)::Int + +Retrieve the external input dimension of the component. +""" +extdim(c::ComponentModel) = has_external_input(c) ? length(extin(c)) : 0 + # return both "observed" outputs (those that do not shadow states) and true observed outsym_flat(c::ComponentModel) = c._outsym_flat obssym_all(c::ComponentModel) = c._obssym_all @@ -527,10 +523,21 @@ outsym_normalized(c::EdgeModel) = values(outsym(c)) outsym_normalized(c::VertexModel) = (outsym(c),) outdim_normalized(c::ComponentModel) = map(length, outsym_normalized(c)) -_infer_ss_fftype(g) = _infer_fftype(g, 1, 2, nothing) -_infer_fftype(::Type{<:VertexModel}, g, dim) = _infer_fftype(g, 1, 1, dim) -_infer_fftype(::Type{<:EdgeModel}, g, dim) = _infer_fftype(g, 2, 2, dim) - +infer_fftype(::Type{<:VertexModel}, g, dim, hasext) = _infer_fftype(g, 1, 1+hasext, dim) +infer_fftype(::Type{<:EdgeModel}, g, dim, hasext) = _infer_fftype(g, 2, 2+hasext, dim) +# special cases for wrapped output functions +function infer_fftype(::Type{<:EdgeModel}, g::Union{Symmetric, AntiSymmetric, Directed}, dim, hasext) + _infer_fftype(g.g, 1, 2+hasext, dim) +end +function infer_fftype(::Type{<:EdgeModel}, g::Fiducial, dim, hasext) + ffsrc = _infer_fftype(g.src, 1, 2+hasext, dim) + ffdst = _infer_fftype(g.dst, 1, 2+hasext, dim) + if ffsrc == ffdst + return ffsrc + else + error("Both src and dst coupling functions have different ff types.") + end +end function _infer_fftype(g, nout, nin, dim) pureff = _takes_n_vecs_and_t(g, nout + nin + 1) # (outs..., ins..., p, t) ff = _takes_n_vecs_and_t(g, nout + 1 + nin + 1) # (outs..., u, ins..., p, t) @@ -564,18 +571,19 @@ 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} -dispatchT(T::Type{<:EdgeModel}) = EdgeModel{nothing,nothing,nothing,nothing,nothing} +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 + 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 @@ -735,11 +743,6 @@ function _fill_defaults(T, @nospecialize(kwargs)) end sym = dict[:sym] - # infer fftype (needs dim) - if !haskey(dict, :ff) - dict[:ff] = hasfftype(g) ? fftype(g) : _infer_fftype(T, g, dim) - end - #### #### parameter sym #### @@ -957,6 +960,23 @@ function _fill_defaults(T, @nospecialize(kwargs)) dict[:_outsym_flat] = _outsym_flat dict[:_obssym_all] = setdiff(_outsym_flat, sym) ∪ obssym + #### + #### External Inputs + #### + if haskey(dict, :extin) + @assert dict[:extin] isa Union{Nothing, Vector{<:SymbolicIndex}} + else + dict[:extin] = nothing + end + + # infer fftype (needs dim and extdim) + if !haskey(dict, :ff) + dict[:ff] = if hasfftype(g) + fftype(g) + else + infer_fftype(T, g, dim, !isnothing(dict[:extin])) + end + end # check for name clashes (at the end because only now sym, psym, obssym are initialized) _s = sym @@ -986,25 +1006,25 @@ end # define the symbolmapping to infer output symbols from state symbols _has_sym_to_outsym_mapping(::Any) = false _has_sym_to_outsym_mapping(::StateMask) = true -_has_sym_to_outsym_mapping(::Directed{<:Any, <:StateMask}) = true -_has_sym_to_outsym_mapping(::AntiSymmetric{<:Any, <:StateMask}) = true -_has_sym_to_outsym_mapping(::Symmetric{<:Any, <:StateMask}) = true -_has_sym_to_outsym_mapping(::Fiducial{<:Any, <:StateMask, <:StateMask}) = true +_has_sym_to_outsym_mapping(::Directed{<:StateMask}) = true +_has_sym_to_outsym_mapping(::AntiSymmetric{<:StateMask}) = true +_has_sym_to_outsym_mapping(::Symmetric{<:StateMask}) = true +_has_sym_to_outsym_mapping(::Fiducial{<:StateMask, <:StateMask}) = true _sym_to_outsym(g::StateMask, s::AbstractVector{Symbol}) = s[g.idxs] -function _sym_to_outsym(g::AntiSymmetric{<:Any, <:StateMask}, s::AbstractVector{Symbol}) +function _sym_to_outsym(g::AntiSymmetric{<:StateMask}, s::AbstractVector{Symbol}) s = _sym_to_outsym(g.g, s) _symvec_to_sym_tup(g, s) end -function _sym_to_outsym(g::Symmetric{<:Any, <:StateMask}, s::AbstractVector{Symbol}) +function _sym_to_outsym(g::Symmetric{<:StateMask}, s::AbstractVector{Symbol}) s = _sym_to_outsym(g.g, s) _symvec_to_sym_tup(g, s) end -function _sym_to_outsym(g::Directed{<:Any, <:StateMask}, s::AbstractVector{Symbol}) +function _sym_to_outsym(g::Directed{<:StateMask}, s::AbstractVector{Symbol}) s = _sym_to_outsym(g.g, s) _symvec_to_sym_tup(g, s) end -function _sym_to_outsym(g::Fiducial{<:Any, <:StateMask, <:StateMask}, s::AbstractVector{Symbol}) +function _sym_to_outsym(g::Fiducial{<:StateMask, <:StateMask}, s::AbstractVector{Symbol}) dst = _sym_to_outsym(g.dst, s) src = _sym_to_outsym(g.src, s) (; src, dst) @@ -1097,21 +1117,13 @@ function Base.:(==)(cf1::ComponentModel, cf2::ComponentModel) typeof(cf1) == typeof(cf2) && equal_fields(cf1, cf2) end -function compfg(c::VertexModel) - (out, du, u, in, p, t) -> begin +function compfg(c) + (outs, du, u, ins, p, t) -> begin f = compf(c) - isnothing(f) || f(du, u, in, p, t) - compg(c)(_gargs(fftype(c), (out,), du, u, (in,), p, t)...) + isnothing(f) || f(du, u, ins..., p, t) + compg(c)(_gargs(fftype(c), outs, du, u, ins, p, t)...) nothing - end -end -function compfg(c::EdgeModel) - (out1, out2, du, u, in1, in2, p, t) -> begin - f = compf(c) - isnothing(f) || f(du, u, in1, in2, p, t) - compg(c)(_gargs(fftype(c), (out1,out2), du, u, (in1,in2), p, t)...) - nothing - end + end end _gargs(::PureFeedForward, outs, du, u, ins, p, t) = (outs..., ins..., p, t) _gargs(::FeedForward, outs, du, u, ins, p, t) = (outs..., u, ins..., p, t) @@ -1130,6 +1142,8 @@ Returns the transformed `VertexModel`. """ function ff_to_constraint(v::VertexModel) hasff(v) || error("Vertex does not have feed forward property.") + # TODO: ff_to_constraint with ext einput + has_external_input(v) && error("ff_to_constraint for model with external input not implemented.") olddim = dim(v) odim = outdim(v) diff --git a/src/construction.jl b/src/construction.jl index a12525e1..8f962420 100644 --- a/src/construction.jl +++ b/src/construction.jl @@ -171,7 +171,8 @@ function Network(g::AbstractGraph, mass_matrix = construct_mass_matrix(im) N = ForwardDiff.pickchunksize(max(im.lastidx_dynamic, im.lastidx_p)) caches = (; output = DiffCache(zeros(im.lastidx_out), N), - aggregation = DiffCache(zeros(im.lastidx_aggr), N)) + aggregation = DiffCache(zeros(im.lastidx_aggr), N), + external = DiffCache(zeros(im.lastidx_extbuf), N)) gbufprovider = if usebuffer(execution) EagerGBufProvider(im, edgebatches) @@ -179,13 +180,17 @@ function Network(g::AbstractGraph, LazyGBufProvider(im, edgebatches) end + # create map for extenral inputs + extmap = has_external_input(im) ? ExtMap(im) : nothing + nw = Network{typeof(execution),typeof(g),typeof(nl), typeof(vertexbatches), - typeof(mass_matrix),eltype(caches),typeof(gbufprovider)}( + typeof(mass_matrix),eltype(caches),typeof(gbufprovider),typeof(extmap)}( vertexbatches, nl, im, caches, mass_matrix, - gbufprovider + gbufprovider, + extmap, ) end @@ -276,37 +281,39 @@ function aliasgroups(cfs::Vector{T}) where {T<:ComponentModel} end function VertexBatch(im::IndexManager, idxs::Vector{Int}; verbose) - components = @view im.vertexm[idxs] - _compT = dispatchT(first(components)) - _compf = compf(first(components)) - _compg = compg(first(components)) - _ff = fftype(first(components)) - _dim = dim(first(components)) - _pdim = pdim(first(components)) - _outdim = outdim(first(components)) - - strides = register_vertices!(im, _dim, _outdim, _pdim, idxs) + first_comp = im.vertexm[first(idxs)] + _compT = dispatchT(first_comp) + _compf = compf(first_comp) + _compg = compg(first_comp) + _ff = fftype(first_comp) + _dim = dim(first_comp) + _pdim = pdim(first_comp) + _outdim = outdim(first_comp) + _extdim = extdim(first_comp) + + strides = register_vertices!(im, idxs, _dim, _outdim, _pdim, _extdim) verbose && println(" - VertexBatch: dim=$(_dim), pdim=$(_pdim), length=$(length(idxs))") - ComponentBatch(_compT, idxs, _compf, _compg, _ff, strides.state, strides.p, strides.in, strides.out) + ComponentBatch(_compT, idxs, _compf, _compg, _ff, strides.state, strides.p, strides.in, strides.out, strides.ext) end function EdgeBatch(im::IndexManager, idxs::Vector{Int}; verbose) - components = @view im.edgem[idxs] - _compT = dispatchT(first(components)) - _compf = compf(first(components)) - _compg = compg(first(components)) - _ff = fftype(first(components)) - _dim = dim(first(components)) - _pdim = pdim(first(components)) - _outdim = outdim(first(components)) - - strides = register_edges!(im, _dim, _outdim, _pdim, idxs) + first_comp = im.edgem[first(idxs)] + _compT = dispatchT(first_comp) + _compf = compf(first_comp) + _compg = compg(first_comp) + _ff = fftype(first_comp) + _dim = dim(first_comp) + _pdim = pdim(first_comp) + _outdim = outdim(first_comp) + _extdim = extdim(first_comp) + + strides = register_edges!(im, idxs, _dim, _outdim, _pdim, _extdim) verbose && println(" - EdgeBatch: dim=$(_dim), pdim=$(_pdim), length=$(length(idxs))") - ComponentBatch(_compT, idxs, _compf, _compg, _ff, strides.state, strides.p, strides.in, strides.out) + ComponentBatch(_compT, idxs, _compf, _compg, _ff, strides.state, strides.p, strides.in, strides.out, strides.ext) end batch_by_idxs(v, idxs::Vector{Vector{Int}}) = [v for batch in idxs] diff --git a/src/coreloop.jl b/src/coreloop.jl index 861a221e..1162c152 100644 --- a/src/coreloop.jl +++ b/src/coreloop.jl @@ -2,34 +2,37 @@ function (nw::Network{A,B,C,D,E})(du::dT, u::T, p, t) where {A,B,C,D,E,dT,T} if !(eachindex(du) == eachindex(u) == 1:nw.im.lastidx_dynamic) throw(ArgumentError("du or u does not have expected size $(nw.im.lastidx_dynamic)")) end - if nw.im.lastidx_p > 0 && _indexable(p) && !(eachindex(p) == 1:nw.im.lastidx_p) + if pdim(nw) > 0 && !(eachindex(p) == 1:nw.im.lastidx_p) throw(ArgumentError("p does not has expecte size $(nw.im.lastidx_p)")) end ex = executionstyle(nw) fill!(du, zero(eltype(du))) o = get_output_cache(nw, du) - fill!(o, convert(eltype(o), NaN)) + extbuf = has_external_input(nw) ? get_extinput_cache(nw, du) : nothing duopt = (du, u, o, p, t) aggbuf = get_aggregation_cache(nw, du) gbuf = get_gbuf(nw.gbufprovider, o) # vg without ff - process_batches!(ex, Val{:g}(), !hasff, nw.vertexbatches, nothing, duopt) + process_batches!(ex, Val{:g}(), !hasff, nw.vertexbatches, (nothing, nothing), duopt) # eg without ff - process_batches!(ex, Val{:g}(), !hasff, nw.layer.edgebatches, nothing, duopt) + process_batches!(ex, Val{:g}(), !hasff, nw.layer.edgebatches, (nothing, nothing), duopt) # process batches might be async so sync before next step ex isa KAExecution && KernelAbstractions.synchronize(get_backend(du)) + # gather the external inputs + has_external_input(nw) && collect_externals!(nw.extmap, extbuf, u, o) + # gather the vertex results for edges with ff - gather!(nw.gbufprovider, gbuf, o) # 2.6ms + gather!(nw.gbufprovider, gbuf, o) # execute f for the edges without ff - process_batches!(ex, Val{:f}(), !hasff, nw.layer.edgebatches, gbuf, duopt) + process_batches!(ex, Val{:f}(), !hasff, nw.layer.edgebatches, (gbuf, extbuf), duopt) # execute f&g for edges with ff - process_batches!(ex, Val{:fg}(), hasff, nw.layer.edgebatches, gbuf, duopt) + process_batches!(ex, Val{:fg}(), hasff, nw.layer.edgebatches, (gbuf, extbuf), duopt) # process batches might be async so sync before next step ex isa KAExecution && KernelAbstractions.synchronize(get_backend(du)) @@ -38,7 +41,7 @@ function (nw::Network{A,B,C,D,E})(du::dT, u::T, p, t) where {A,B,C,D,E,dT,T} aggregate!(nw.layer.aggregator, aggbuf, o) # vf for vertices without ff - process_batches!(ex, Val{:f}(), !hasff, nw.vertexbatches, aggbuf, duopt) + process_batches!(ex, Val{:f}(), !hasff, nw.vertexbatches, (aggbuf, extbuf), duopt) # process batches might be async so sync before next step ex isa KAExecution && KernelAbstractions.synchronize(get_backend(du)) @@ -47,6 +50,7 @@ end function get_buffers(nw, u, p, t; initbufs) o = get_output_cache(nw, u) aggbuf = get_aggregation_cache(nw, u) + extbuf = has_external_input(nw) ? get_extinput_cache(nw, u) : nothing if initbufs du = nothing duopt = (du, u, o, p, t) @@ -55,54 +59,56 @@ function get_buffers(nw, u, p, t; initbufs) gbuf = get_gbuf(nw.gbufprovider, o) # vg without ff - process_batches!(ex, Val{:g}(), !hasff, nw.vertexbatches, nothing, duopt) + process_batches!(ex, Val{:g}(), !hasff, nw.vertexbatches, (nothing, nothing), duopt) # eg without ff - process_batches!(ex, Val{:g}(), !hasff, nw.layer.edgebatches, nothing, duopt) + process_batches!(ex, Val{:g}(), !hasff, nw.layer.edgebatches, (nothing, nothing), duopt) # process batches might be async so sync before next step ex isa KAExecution && KernelAbstractions.synchronize(get_backend(u)) + # gather the external inputs + has_external_input(nw) && collect_externals!(nw.extmap, extbuf, u, o) # gather the vertex results for edges with ff gather!(nw.gbufprovider, gbuf, o) # 2.6ms # execute g for edges with ff - process_batches!(ex, Val{:g}(), hasff, nw.layer.edgebatches, gbuf, duopt) + process_batches!(ex, Val{:g}(), hasff, nw.layer.edgebatches, (gbuf, nothing), duopt) # process batches might be async so sync before next step ex isa KAExecution && KernelAbstractions.synchronize(get_backend(u)) # aggegrate the results aggregate!(nw.layer.aggregator, aggbuf, o) end - o, aggbuf + o, aggbuf, extbuf end -@inline function process_batches!(::SequentialExecution, fg, filt::F, batches, inbuf, duopt) where {F} +@inline function process_batches!(::SequentialExecution, fg, filt::F, batches, inbufs, duopt) where {F} unrolled_foreach(filt, batches) do batch (du, u, o, p, t) = duopt for i in 1:length(batch) _type = dispatchT(batch) - apply_comp!(_type, fg, batch, i, du, u, o, inbuf, p, t) + apply_comp!(_type, fg, batch, i, du, u, o, inbufs, p, t) end end end -@inline function process_batches!(::ThreadedExecution, fg, filt::F, batches, inbuf, duopt) where {F} +@inline function process_batches!(::ThreadedExecution, fg, filt::F, batches, inbufs, duopt) where {F} unrolled_foreach(filt, batches) do batch (du, u, o, p, t) = duopt Threads.@threads for i in 1:length(batch) _type = dispatchT(batch) - apply_comp!(_type, fg, batch, i, du, u, o, inbuf, p, t) + apply_comp!(_type, fg, batch, i, du, u, o, inbufs, p, t) end end end -@inline function process_batches!(::PolyesterExecution, fg, filt::F, batches, inbuf, duopt) where {F} +@inline function process_batches!(::PolyesterExecution, fg, filt::F, batches, inbufs, duopt) where {F} unrolled_foreach(filt, batches) do batch (du, u, o, p, t) = duopt Polyester.@batch for i in 1:length(batch) _type = dispatchT(batch) - apply_comp!(_type, fg, batch, i, du, u, o, inbuf, p, t) + apply_comp!(_type, fg, batch, i, du, u, o, inbufs, p, t) end end end -@inline function process_batches!(::KAExecution, fg, filt::F, batches, inbuf, duopt) where {F} +@inline function process_batches!(::KAExecution, fg, filt::F, batches, inbufs, duopt) where {F} _backend = get_backend(duopt[1]) unrolled_foreach(filt, batches) do batch (du, u, o, p, t) = duopt @@ -114,50 +120,60 @@ end elseif evalg(fg, batch) compkernel_g!(_backend) end - isnothing(kernel) || kernel(_type, fg, batch, du, u, o, inbuf, p, t; ndrange=length(batch)) + isnothing(kernel) || kernel(_type, fg, batch, du, u, o, inbufs, p, t; ndrange=length(batch)) end end @kernel function compkernel_f!(::Type{T}, @Const(fg), @Const(batch), - du, @Const(u), @Const(o), @Const(inbuf), @Const(p), @Const(t)) where {T} + du, @Const(u), @Const(o), @Const(inbufs), @Const(p), @Const(t)) where {T} I = @index(Global) - apply_comp!(T, fg, batch, I, du, u, o, inbuf, p, t) + apply_comp!(T, fg, batch, I, du, u, o, inbufs, p, t) nothing end @kernel function compkernel_g!(::Type{T}, @Const(fg), @Const(batch), - @Const(du), @Const(u), o, @Const(inbuf), @Const(p), @Const(t)) where {T} + @Const(du), @Const(u), o, @Const(inbufs), @Const(p), @Const(t)) where {T} I = @index(Global) - apply_comp!(T, fg, batch, I, du, u, o, inbuf, p, t) + apply_comp!(T, fg, batch, I, du, u, o, inbufs, p, t) nothing end @kernel function compkernel_fg!(::Type{T}, @Const(fg), @Const(batch), du, @Const(u), o, @Const(inbuf), @Const(p), @Const(t)) where {T} I = @index(Global) - apply_comp!(T, fg, batch, I, du, u, o, inbuf, p, t) + apply_comp!(T, fg, batch, I, du, u, o, inbufs, p, t) nothing end -@inline function apply_comp!(::Type{<:VertexModel}, fg, batch, i, du, u, o, aggbuf, p, t) +@inline function apply_comp!(::Type{<:VertexModel}, fg, batch, i, du, u, o, inbufs, p, t) @inbounds begin + aggbuf, extbuf = inbufs _o = _needs_out(fg, batch) ? view(o, out_range(batch, i)) : nothing _du = _needs_du(fg, batch) ? view(du, state_range(batch, i)) : nothing _u = _needs_u(fg, batch) ? view(u, state_range(batch, i)) : nothing - _agg = _needs_in(fg, batch) ? view(aggbuf, in_range(batch, i)) : nothing - _p = _needs_p(fg, batch) && _indexable(p) ? view(p, parameter_range(batch, i)) : p - evalf(fg, batch) && apply_compf(compf(batch), _du, _u, (_agg,), _p, t) - evalg(fg, batch) && apply_compg(fftype(batch), compg(batch), (_o,), _u, (_agg,), _p, t) + _ins = _needs_in(fg, batch) ? (view(aggbuf, in_range(batch, i)),) : nothing + _p = _needs_p(fg, batch) ? view(p, parameter_range(batch, i)) : nothing + if has_external_input(batch) && _needs_in(fg, batch) + _ext = view(extbuf, extbuf_range(batch, i)) + _ins = (_ins..., _ext) + end + evalf(fg, batch) && apply_compf(compf(batch), _du, _u, _ins, _p, t) + evalg(fg, batch) && apply_compg(fftype(batch), compg(batch), (_o,), _u, _ins, _p, t) end nothing end -@inline function apply_comp!(::Type{<:EdgeModel}, fg, batch, i, du, u, o, gbuf, p, t) +@inline function apply_comp!(::Type{<:EdgeModel}, fg, batch, i, du, u, o, inbufs, p, t) @inbounds begin - _osrc = _needs_out(fg, batch) ? view(o, out_range(batch, i, 1)) : nothing - _odst = _needs_out(fg, batch) ? view(o, out_range(batch, i, 2)) : nothing - _du = _needs_du(fg, batch) ? view(du, state_range(batch, i)) : nothing - _u = _needs_u(fg, batch) ? view(u, state_range(batch, i)) : nothing - _ins = _needs_in(fg, batch) ? get_src_dst(gbuf, batch, i) : nothing - _p = _needs_p(fg, batch) && _indexable(p) ? view(p, parameter_range(batch, i)) : p + gbuf, extbuf = inbufs + _osrc = _needs_out(fg, batch) ? view(o, out_range(batch, i, :src)) : nothing + _odst = _needs_out(fg, batch) ? view(o, out_range(batch, i, :dst)) : nothing + _du = _needs_du(fg, batch) ? view(du, state_range(batch, i)) : nothing + _u = _needs_u(fg, batch) ? view(u, state_range(batch, i)) : nothing + _ins = _needs_in(fg, batch) ? get_src_dst(gbuf, batch, i) : nothing + _p = _needs_p(fg, batch) ? view(p, parameter_range(batch, i)) : nothing + if has_external_input(batch) && _needs_in(fg, batch) + _ext = view(extbuf, extbuf_range(batch, i)) + _ins = (_ins..., _ext) + end evalf(fg, batch) && apply_compf(compf(batch), _du, _u, _ins, _p, t) evalg(fg, batch) && apply_compg(fftype(batch), compg(batch), (_osrc, _odst), _u, _ins, _p, t) end @@ -187,11 +203,11 @@ end end # check if the function arguments are actually used -_needs_du(fg, batch) = evalf(fg, batch) -_needs_u(fg, batch) = evalf(fg, batch) || fftype(batch) != PureFeedForward() +_needs_du(fg, batch) = evalf(fg, batch) +_needs_u(fg, batch) = evalf(fg, batch) || fftype(batch) != PureFeedForward() _needs_out(fg, batch) = evalg(fg, batch) -_needs_in(fg, batch) = evalf(fg, batch) || hasff(batch) -_needs_p(fg, batch) = evalf(fg, batch) || fftype(batch) != PureStateMap() +_needs_in(fg, batch) = evalf(fg, batch) || hasff(batch) +_needs_p(fg, batch) = !iszero(pdim(batch)) && (evalf(fg, batch) || fftype(batch) != PureStateMap()) # check if eval of f or g is necessary evalf(::Val{:f}, batch) = !isnothing(compf(batch)) @@ -200,8 +216,3 @@ evalf(::Val{:fg}, batch) = !isnothing(compf(batch)) evalg(::Val{:f}, _) = false evalg(::Val{:g}, _) = true evalg(::Val{:fg}, _) = true - -# check if indexing into p is necessary -_indexable(::Nothing) = false -_indexable(::SciMLBase.NullParameters) = false -_indexable(::AbstractVector) = true diff --git a/src/doctor.jl b/src/doctor.jl index 8c73de60..c46422f6 100644 --- a/src/doctor.jl +++ b/src/doctor.jl @@ -99,10 +99,14 @@ function chk_component(c::ComponentModel) Tuple(AccessTracker(rand(indim_guess)) for _ in outdim_normalized(c)) end outs = Tuple(AccessTracker(rand(l)) for l in values(outdim(c))) + ext = AccessTracker(rand(extdim(c))) + if has_external_input(c) + ins = (ins..., ext) + end t = NaN try - compfg(c)(outs..., du, u, ins..., p, t) + compfg(c)(outs, du, u, ins, p, t) catch e if e isa MethodError @warn "Encountered MethodError. All arguments are AbstractArrays, make sure to allways index into them: $e" @@ -123,6 +127,7 @@ function chk_component(c::ComponentModel) has_oob(du) && @warn "There is out of bound acces to du: reads $(oob_reads(du)) and writes $(oob_writes(du))! Check dim/sym!" has_oob(u) && @warn "There is out of bound acces to u: reads $(oob_reads(u)) and writes $(oob_writes(u))! Check dim/sym!" has_oob(p) && @warn "There is out of bound acces to p: reads $(oob_reads(p)) and writes $(oob_writes(p))! Check pdim/psim!" + has_oob(ext) && @warn "There is out of bound acces to external input: reads $(oob_reads(ext)) and writes $(oob_writes(ext))!" for (j, o) in enumerate(outs) has_oob(o) && @warn "There is out of bound acces to output#$j: reads $(oob_reads(o)) and writes $(oob_writes(o))!" end @@ -148,6 +153,7 @@ function chk_component(c::ComponentModel) end has_writes(p) && @warn "There is write access to p: $(writes(p))!" + has_writes(ext) && @warn "There is write access to external inputs: $(writes(ext))!" similars = String[] has_similars(du) && push!(similars, "du") diff --git a/src/external_inputs.jl b/src/external_inputs.jl new file mode 100644 index 00000000..c4afda48 --- /dev/null +++ b/src/external_inputs.jl @@ -0,0 +1,86 @@ +struct StateBufIdx + idx::Int +end +struct OutBufIdx + idx::Int +end +struct ExtMap{M<:AbstractVector{<:Union{StateBufIdx, OutBufIdx}}} + map::M +end + +function ExtMap(im::IndexManager) + map = Vector{Union{StateBufIdx, OutBufIdx}}(undef, im.lastidx_extbuf) + isempty(map) && error("There are no external inputs.") + + for (vi, vm) in pairs(im.vertexm) + has_external_input(vm) || continue + for (i, si) in pairs(vm.extin) + map[im.v_ext[vi][i]] = _symidx_to_extidx(im, si) + end + end + for (ei, em) in pairs(im.edgem) + has_external_input(em) || continue + for (i, si) in pairs(em.extin) + map[im.e_ext[ei][i]] = _symidx_to_extidx(im, si) + end + end + + # narrow down type if all are output or state + if all(e -> typeof(e) == typeof(first(map)), map) + map = Vector{typeof(first(map))}(map) + end + ExtMap(map) +end + +function _symidx_to_extidx(im, sni) + cm = getcomp(im, sni) + if subsym_has_idx(sni.subidx, sym(cm)) + range = getcomprange(im, sni) + return StateBufIdx(range[subsym_to_idx(sni.subidx, sym(cm))]) + elseif subsym_has_idx(sni.subidx, outsym_flat(cm)) + range = getcompoutrange(im, sni) + if hasff(cm) + throw(ArgumentError("Cannot resolve external input $sni! Outputs of feed-forward components are not allowed as external inputs.")) + end + return OutBufIdx(range[subsym_to_idx(sni.subidx, sym(cm))]) + else + throw(ArgumentError("Cannot resolve external input $sni! External inputs musst be states or outputs of non-feed-forward components.")) + end +end + +# CPU version +function collect_externals!(map::ExtMap, extbuf::Vector, u::Vector, o::Vector) + # the performance of this function really depends on the fact + # that julia somehow inferes that if src isn't a statebuf it musst be a outbuf + # wasn't able to recreate something this fast in a more general way + @inbounds for (dst, src) in pairs(map.map) + if src isa StateBufIdx + extbuf[dst] = u[src.idx] + else + extbuf[dst] = o[src.idx] + end + end + extbuf +end + +# GPU Version +function collect_externals!(map::ExtMap, extbuf, u, o) + _backend = get_backend(extbuf) + kernel = collect_ext_kernel!(_backend) + kernel(map.map, extbuf, u, o; ndrange=length(extbuf)) +end +@kernel function collect_ext_kernel!(map, extbuf, @Const(u), @Const(o)) + dst = @index(Global) + src = map[dst] + if src isa StateBufIdx + extbuf[dst] = u[src.idx] + else + extbuf[dst] = o[src.idx] + end + nothing +end + +has_external_input(c::ComponentModel) = !isnothing(c.extin) +has_external_input(cb::ComponentBatch) = !iszero(cb.extbufstride.strides) +has_external_input(nw::Network) = !isnothing(nw.extmap) +has_external_input(im::IndexManager) = !iszero(im.lastidx_extbuf) diff --git a/src/initialization.jl b/src/initialization.jl index ae9a5739..272b6040 100644 --- a/src/initialization.jl +++ b/src/initialization.jl @@ -156,7 +156,7 @@ function initialization_problem(cf::T; t=NaN, verbose=true) where {T<:ComponentM # this fills the second half of the du buffer with the fixed and current outputs dunl_out .= RecursiveArrayTools.ArrayPartition(outbufs...) # execute fg to fill dunl and outputs - fg(outbufs..., dunl_fg, ubuf, inbufs..., pbuf, t) + fg(outbufs, dunl_fg, ubuf, inbufs, pbuf, t) # calculate the residual for the second half ov the dunl buf, the outputs dunl_out .= dunl_out .- RecursiveArrayTools.ArrayPartition(outbufs...) @@ -242,7 +242,7 @@ function init_residual(cf::T; t=NaN, recalc=false) where {T<:ComponentModel} res_fg = @views res[1:dim(cf)] res_out = @views res[dim(cf)+1:end] res_out .= RecursiveArrayTools.ArrayPartition(outs...) - compfg(cf)(outs..., res_fg, u, ins..., p, t) + compfg(cf)(outs, res_fg, u, ins, p, t) res_out .= res_out .- RecursiveArrayTools.ArrayPartition(outs...) LinearAlgebra.norm(res) diff --git a/src/network_structure.jl b/src/network_structure.jl index d998b3e7..5910f6f3 100644 --- a/src/network_structure.jl +++ b/src/network_structure.jl @@ -6,11 +6,13 @@ mutable struct IndexManager{G} v_out::Vector{UnitRange{Int}} # v range in output buf v_para::Vector{UnitRange{Int}} # v para in flat para v_aggr::Vector{UnitRange{Int}} # v input in aggbuf + v_ext::Vector{UnitRange{Int}} # v external inputs in ext buf # positions of edge data e_data::Vector{UnitRange{Int}} # e data in flat states e_out::Vector{@NamedTuple{src::UnitRange{Int},dst::UnitRange{Int}}} # e range in output buf e_para::Vector{UnitRange{Int}} # e para in flat para e_gbufr::Vector{@NamedTuple{src::UnitRange{Int},dst::UnitRange{Int}}} # e range in input buf + e_ext::Vector{UnitRange{Int}} # v external inputs in ext buf # metadata edepth::Int vdepth::Int @@ -19,6 +21,7 @@ mutable struct IndexManager{G} lastidx_p::Int lastidx_aggr::Int lastidx_gbuf::Int + lastidx_extbuf::Int vertexm::Vector{VertexModel} edgem::Vector{EdgeModel} aliased_vertexms::IdDict{VertexModel, @NamedTuple{idxs::Vector{Int}, hash::UInt}} @@ -31,13 +34,14 @@ mutable struct IndexManager{G} unique_vnames = unique_mappings(getproperty.(vertexm, :name), 1:nv(g)) unique_enames = unique_mappings(getproperty.(edgem, :name), 1:ne(g)) new{typeof(g)}(g, collect(edges(g)), - (Vector{UnitRange{Int}}(undef, nv(g)) for i in 1:4)..., + (Vector{UnitRange{Int}}(undef, nv(g)) for i in 1:5)..., Vector{UnitRange{Int}}(undef, ne(g)), Vector{@NamedTuple{src::UnitRange{Int},dst::UnitRange{Int}}}(undef, ne(g)), Vector{UnitRange{Int}}(undef, ne(g)), Vector{@NamedTuple{src::UnitRange{Int},dst::UnitRange{Int}}}(undef, ne(g)), + Vector{UnitRange{Int}}(undef, ne(g)), edepth, vdepth, - 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, vertexm, edgem, aliased_vertexm_hashes, aliased_edgem_hashes, @@ -62,10 +66,9 @@ end dim(im::IndexManager) = im.lastidx_dynamic pdim(im::IndexManager) = im.lastidx_p -sdim(im::IndexManager) = im.lastidx_static - im.lastidx_dynamic -struct Network{EX<:ExecutionStyle,G,NL,VTup,MM,CT,GBT} +struct Network{EX<:ExecutionStyle,G,NL,VTup,MM,CT,GBT,EM} "vertex batches of same function" vertexbatches::VTup "network layer" @@ -73,11 +76,13 @@ struct Network{EX<:ExecutionStyle,G,NL,VTup,MM,CT,GBT} "index manager" im::IndexManager{G} "lazy cache pool" - caches::@NamedTuple{output::CT,aggregation::CT} + caches::@NamedTuple{output::CT,aggregation::CT,external::CT} "mass matrix" mass_matrix::MM "Gather buffer provider (lazy or eager)" gbufprovider::GBT + "map to gather external inputs" + extmap::EM end executionstyle(::Network{ex}) where {ex} = ex() nvbatches(::Network) = length(vertexbatches) @@ -106,9 +111,14 @@ function get_output_cache(nw::Network, T) throw(ArgumentError("Network caches are initialized with $(eltype(nw.caches.output.du)) \ but is used for $(eltype(T)) data!")) end - get_tmp(nw.caches.output, T) + o = get_tmp(nw.caches.output, T) + fill!(o, convert(eltype(o), NaN)) end get_aggregation_cache(nw::Network, T) = get_tmp(nw.caches.aggregation, T) +function get_extinput_cache(nw::Network, T) + ext = get_tmp(nw.caches.external, T) + fill!(ext, convert(eltype(ext), NaN)) +end iscudacompatible(nw::Network) = iscudacompatible(executionstyle(nw)) && iscudacompatible(nw.layer.aggregator) @@ -125,7 +135,7 @@ struct NetworkLayer{GT,ETup,AF} vdepth::Int # potential becomes range for multilayer end -struct ComponentBatch{T,F,G,FFT,DIM,PDIM,INDIMS,OUTDIMS,IV} +struct ComponentBatch{T,F,G,FFT,DIM,PDIM,INDIMS,OUTDIMS,EXTDIM,IV} "indices contained in batch" indices::IV "internal function" @@ -141,9 +151,11 @@ struct ComponentBatch{T,F,G,FFT,DIM,PDIM,INDIMS,OUTDIMS,IV} inbufstride::BatchStride{INDIMS} "outputbuf: dimension(s) and first index" outbufstride::BatchStride{OUTDIMS} - function ComponentBatch(dT, i, f, g, ff, ss, ps, is, os) - new{dT,typeof.((f,g,ff))...,stridesT.((ss, ps, is, os))...,typeof(i)}( - i, f, g, ff, ss, ps, is, os) + "external inputs: dimension and first index" + extbufstride::BatchStride{EXTDIM} + function ComponentBatch(dT, i, f, g, ff, ss, ps, is, os, es) + new{dT,typeof.((f,g,ff))...,stridesT.((ss, ps, is, os, es))...,typeof(i)}( + i, f, g, ff, ss, ps, is, os, es) end end @@ -152,6 +164,8 @@ end @inline compf(b::ComponentBatch) = b.compf @inline compg(b::ComponentBatch) = b.compg @inline fftype(b::ComponentBatch) = b.ff +@inline pdim(b::ComponentBatch) = b.pstride.strides +@inline extdim(b::ComponentBatch) = b.extbufstride.strides @inline state_range(batch) = _fullrange(batch.statestride, length(batch)) @@ -161,22 +175,25 @@ end @inline out_range(batch, i, j) = _range(batch.outbufstride, i, j) @inline in_range(batch, i) = _range(batch.inbufstride, i) @inline in_range(batch, i, j) = _range(batch.inbufstride, i, j) +@inline extbuf_range(batch, i) = _range(batch.extbufstride, i) -function register_vertices!(im::IndexManager, dim, outdim, pdim, idxs) +function register_vertices!(im::IndexManager, idxs, dim, outdim, pdim, extdim) for i in idxs im.v_data[i] = _nexturange!(im, dim) im.v_out[i] = _nextoutrange!(im, outdim) im.v_para[i] = _nextprange!(im, pdim) im.v_aggr[i] = _nextaggrrange!(im, im.edepth) + im.v_ext[i] = _nextextrange!(im, extdim) end (; state = BatchStride(first(im.v_data[first(idxs)]), dim), p = BatchStride(first(im.v_para[first(idxs)]), pdim), in = BatchStride(first(im.v_aggr[first(idxs)]), im.edepth), out = BatchStride(first(im.v_out[first(idxs)]), outdim), + ext = BatchStride(first(im.v_ext[first(idxs)]), extdim), ) end -function register_edges!(im::IndexManager, dim, outdim, pdim, idxs) +function register_edges!(im::IndexManager, idxs, dim, outdim, pdim, extdim) for i in idxs e = im.edgevec[i] im.e_data[i] = _nexturange!(im, dim) @@ -185,12 +202,14 @@ function register_edges!(im::IndexManager, dim, outdim, pdim, idxs) im.e_para[i] = _nextprange!(im, pdim) im.e_gbufr[i] = (src = _nextgbufrange!(im, im.vdepth), dst = _nextgbufrange!(im, im.vdepth)) + im.e_ext[i] = _nextextrange!(im, extdim) end (; state = BatchStride(first(im.e_data[first(idxs)]), dim), p = BatchStride(first(im.e_para[first(idxs)]), pdim), in = BatchStride(first(flatrange(im.e_gbufr[first(idxs)])), (;src=im.vdepth, dst=im.vdepth)), out = BatchStride(first(flatrange(im.e_out[first(idxs)])), (;src=outdim.src, dst=outdim.dst)), + ext = BatchStride(first(im.e_ext[first(idxs)]), extdim), ) end function _nexturange!(im::IndexManager, N) @@ -218,12 +237,18 @@ function _nextgbufrange!(im::IndexManager, N) im.lastidx_gbuf = newlast return range end +function _nextextrange!(im::IndexManager, N) + newlast, range = _nextrange(im.lastidx_extbuf, N) + im.lastidx_extbuf = newlast + return range +end _nextrange(last, N) = last + N, last+1:last+N function isdense(im::IndexManager) - pidxs = Int[] - stateidxs = Int[] - outidxs = Int[] + stateidxs = sizehint!(Int[], im.lastidx_dynamic) + pidxs = sizehint!(Int[], im.lastidx_p) + outidxs = sizehint!(Int[], im.lastidx_out) + extidxs = sizehint!(Int[], im.lastidx_extbuf) for dataranges in (im.v_data, im.e_data) for range in dataranges append!(stateidxs, range) @@ -239,11 +264,18 @@ function isdense(im::IndexManager) append!(outidxs, flatrange(range)) end end + for extranges in (im.v_ext, im.e_ext) + for range in extranges + append!(extidxs, flatrange(range)) + end + end sort!(pidxs) sort!(stateidxs) sort!(outidxs) + sort!(outidxs) @assert pidxs == 1:im.lastidx_p @assert stateidxs == 1:im.lastidx_dynamic @assert outidxs == 1:im.lastidx_out + @assert extidxs == 1:im.lastidx_extbuf return true end diff --git a/src/show.jl b/src/show.jl index 48141210..547605fd 100644 --- a/src/show.jl +++ b/src/show.jl @@ -67,6 +67,12 @@ function print_states_params(io, @nospecialize(c::ComponentModel), styling) num, word = maybe_plural(pdim(c), "param") pdim(c) > 0 && push!(info, styled"$num &$word: &&$(stylesymbolarray(c.psym, pdef(c), pguess(c)))") + if has_external_input(c) + num = extdim(c) + arr = match(r"(\[.*\])", repr(extin(c)))[1] + push!(info, styled"$num &ext in: &&$arr") + end + print_treelike(io, align_strings(info)) end function _inout_string(@nospecialize(c::VertexModel), f, name) diff --git a/src/symbolicindexing.jl b/src/symbolicindexing.jl index fda4463e..1e0b17aa 100644 --- a/src/symbolicindexing.jl +++ b/src/symbolicindexing.jl @@ -1,6 +1,3 @@ -abstract type SymbolicIndex{C,S} end -abstract type SymbolicStateIndex{C,S} <: SymbolicIndex{C,S} end -abstract type SymbolicParameterIndex{C,S} <: SymbolicIndex{C,S} end """ VIndex{C,S} <: SymbolicStateIndex{C,S} idx = VIndex(comp, sub) @@ -105,21 +102,28 @@ function SII.getname(x::SymbolicEdgeIndex) Symbol(prefix, Symbol(x.compidx), :₊, Symbol(x.subidx)) end -resolvecompidx(nw::Network, sni::SymbolicIndex{Int}) = sni.compidx -function resolvecompidx(nw::Network, sni::SymbolicIndex{Symbol}) - dict = sni isa SymbolicVertexIndex ? nw.im.unique_vnames : nw.im.unique_enames +resolvecompidx(nw::Network, sni) = resolvecompidx(nw.im, sni) +resolvecompidx(::IndexManager, sni::SymbolicIndex{Int}) = sni.compidx +function resolvecompidx(im::IndexManager, sni::SymbolicIndex{Symbol}) + dict = sni isa SymbolicVertexIndex ? im.unique_vnames : im.unique_enames if haskey(dict, sni.compidx) return dict[sni.compidx] else throw(ArgumentError("Could not resolve component index for $sni, the name might not be unique?")) end end -getcomp(nw::Network, sni::SymbolicEdgeIndex) = nw.im.edgem[resolvecompidx(nw, sni)] -getcomp(nw::Network, sni::SymbolicVertexIndex) = nw.im.vertexm[resolvecompidx(nw, sni)] -getcomprange(nw::Network, sni::VIndex{<:Union{Symbol,Int}}) = nw.im.v_data[resolvecompidx(nw, sni)] -getcomprange(nw::Network, sni::EIndex{<:Union{Symbol,Int}}) = nw.im.e_data[resolvecompidx(nw, sni)] -getcompoutrange(nw::Network, sni::VIndex{<:Union{Symbol,Int}}) = nw.im.v_out[resolvecompidx(nw, sni)] -getcompoutrange(nw::Network, sni::EIndex{<:Union{Symbol,Int}}) = flatrange(nw.im.e_out[resolvecompidx(nw, sni)]) +getcomp(nw::Network, sni) = getcomp(nw.im, sni) +getcomp(im::IndexManager, sni::SymbolicEdgeIndex) = im.edgem[resolvecompidx(im, sni)] +getcomp(im::IndexManager, sni::SymbolicVertexIndex) = im.vertexm[resolvecompidx(im, sni)] + +getcomprange(nw::Network, sni) = getcomprange(nw.im, sni) +getcomprange(im::IndexManager, sni::VIndex{<:Union{Symbol,Int}}) = im.v_data[resolvecompidx(im, sni)] +getcomprange(im::IndexManager, sni::EIndex{<:Union{Symbol,Int}}) = im.e_data[resolvecompidx(im, sni)] + +getcompoutrange(nw::Network, sni) = getcompoutrange(nw.im, sni) +getcompoutrange(im::IndexManager, sni::VIndex{<:Union{Symbol,Int}}) = im.v_out[resolvecompidx(im, sni)] +getcompoutrange(im::IndexManager, sni::EIndex{<:Union{Symbol,Int}}) = flatrange(im.e_out[resolvecompidx(im, sni)]) + getcompprange(nw::Network, sni::VPIndex{<:Union{Symbol,Int}}) = nw.im.v_para[resolvecompidx(nw, sni)] getcompprange(nw::Network, sni::EPIndex{<:Union{Symbol,Int}}) = nw.im.e_para[resolvecompidx(nw, sni)] @@ -422,7 +426,7 @@ function SII.observed(nw::Network, snis) outidx[i] = _range[idx] elseif (idx=findfirst(isequal(sni.subidx), obssym(cf))) != nothing #found in observed _obsf = _get_observed_f(nw, cf, resolvecompidx(nw, sni)) - obsfuns[i] = (u, outbuf, aggbuf, p, t) -> _obsf(u, outbuf, aggbuf, p, t)[idx] + obsfuns[i] = (u, outbuf, aggbuf, extbuf, p, t) -> _obsf(u, outbuf, aggbuf, extbuf, p, t)[idx] else throw(ArgumentError("Cannot resolve observable $sni")) end @@ -432,7 +436,7 @@ function SII.observed(nw::Network, snis) if isscalar @closure (u, p, t) -> begin - outbuf, aggbuf = get_buffers(nw, u, p, t; initbufs) + outbuf, aggbuf, extbuf = get_buffers(nw, u, p, t; initbufs) if !isempty(stateidx) idx = only(stateidx).second u[idx] @@ -441,12 +445,12 @@ function SII.observed(nw::Network, snis) outbuf[idx] else obsf = only(obsfuns).second - obsf(u, outbuf, aggbuf, p, t)::eltype(u) + obsf(u, outbuf, aggbuf, extbuf, p, t)::eltype(u) end end else @closure (u, p, t) -> begin - outbuf, aggbuf = get_buffers(nw, u, p, t; initbufs) + outbuf, aggbuf, extbuf = get_buffers(nw, u, p, t; initbufs) out = similar(u, length(snis)) for (i, statei) in stateidx @@ -456,7 +460,7 @@ function SII.observed(nw::Network, snis) out[i] = outbuf[outi] end for (i, obsf) in obsfuns - out[i] = obsf(u, outbuf, aggbuf, p, t)::eltype(u) + out[i] = obsf(u, outbuf, aggbuf, extbuf, p, t)::eltype(u) end out end @@ -483,11 +487,17 @@ function _get_observed_f(nw::Network, cf::VertexModel, vidx) N = length(cf.obssym) ur = nw.im.v_data[vidx] aggr = nw.im.v_aggr[vidx] + extr = nw.im.v_out[vidx] pr = nw.im.v_para[vidx] ret = Vector{Float64}(undef, N) - @closure (u, outbuf, aggbuf, p, t) -> begin - cf.obsf(ret, view(u, ur), view(aggbuf, aggr), view(p, pr), t) + @closure (u, outbuf, aggbuf, extbuf, p, t) -> begin + ins = if has_external_input(cf) + (view(aggbuf, aggr), view(extbuf, extr)) + else + (view(aggbuf, aggr), ) + end + cf.obsf(ret, view(u, ur), ins..., view(p, pr), t) ret end end @@ -497,11 +507,16 @@ function _get_observed_f(nw::Network, cf::EdgeModel, eidx) ur = nw.im.e_data[eidx] esrcr = nw.im.v_out[nw.im.edgevec[eidx].src] edstr = nw.im.v_out[nw.im.edgevec[eidx].dst] + extr = nw.im.e_out[eidx] pr = nw.im.e_para[eidx] ret = Vector{Float64}(undef, N) - @closure (u, outbuf, aggbuf, p, t) -> begin - cf.obsf(ret, view(u, ur), view(outbuf, esrcr), view(outbuf, edstr), view(p, pr), t) + @closure (u, outbuf, aggbuf, extbuf, p, t) -> begin + ins = (view(outbuf, esrcr), view(outbuf, edstr)) + if has_external_input(cf) + (ins..., view(extbuf, extr)) + end + cf.obsf(ret, view(u, ur), ins..., view(p, pr), t) ret end end diff --git a/src/utils.jl b/src/utils.jl index 860785f8..1e184477 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -119,8 +119,18 @@ function rand_inputs_fg(rng, cf) u = rand(rng, dim(cf)) p = rand(rng, pdim(cf)) ins = Tuple(rand(rng, l) for l in values(indim(cf))) + if has_external_input(cf) + ext = rand(rng, extdim(cf)) + ins = (ins..., ext) + end outs = Tuple(rand(rng, l) for l in values(outdim(cf))) t = NaN - (outs..., du, u, ins..., p, t) + (outs, du, u, ins, p, t) end rand_inputs_fg(cf) = rand_inputs_fg(Random.default_rng(), cf) + + +# abstract symbolic index types +abstract type SymbolicIndex{C,S} end +abstract type SymbolicStateIndex{C,S} <: SymbolicIndex{C,S} end +abstract type SymbolicParameterIndex{C,S} <: SymbolicIndex{C,S} end diff --git a/test/construction_test.jl b/test/construction_test.jl index b4ebc18f..1a37e05b 100644 --- a/test/construction_test.jl +++ b/test/construction_test.jl @@ -149,7 +149,7 @@ end @testset "EdgeModel Constructor" begin using NetworkDynamics - using NetworkDynamics: Symmetric + using NetworkDynamics: Symmetric, infer_fftype f = (du, u, in1, in2, p, t) -> nothing g_pff = (o1, o2, in1, in2, p, t) -> nothing g_ff = (o1, o2, u, in1, in2, p, t) -> nothing @@ -161,11 +161,13 @@ end g_single_nff = (o2, u, p, t) -> nothing g_single_psm = (o2, u) -> nothing for c in (AntiSymmetric, Symmetric, Directed) - @test fftype(c(g_single_pff)) == PureFeedForward() - @test fftype(c(g_single_ff)) == FeedForward() - @test fftype(c(g_single_nff)) == NoFeedForward() - @test fftype(c(g_single_psm)) == PureStateMap() - @test fftype(c(1:2)) == PureStateMap() + T = EdgeModel + dim = nothing # not used + @test infer_fftype(T, c(g_single_pff), dim, false) == PureFeedForward() + @test infer_fftype(T, c(g_single_ff), dim, false) == FeedForward() + @test infer_fftype(T, c(g_single_nff), dim, false) == NoFeedForward() + @test infer_fftype(T, c(g_single_psm), dim, false) == PureStateMap() + @test infer_fftype(T, c(1:2), dim, false) == PureStateMap() end cf = EdgeModel(g=g_pff, outdim=1, insym=[:a]) @@ -390,16 +392,16 @@ end u = rand(dim(v)) p = rand(pdim(v)) in = rand(indim(v)) - # compfg(v)(out, du, u, in, p, NaN) - b = @b $(compfg(v))($out, $du, $u, $in, $p, NaN) + # compfg(v)((out,), du, u, (in,), p, NaN) + b = @b $(compfg(v))($(out,), $du, $u, $(in,), $p, NaN) @test b.allocs == 0 v2 = ff_to_constraint(v) out2 = zeros(outdim(v2)) du2 = zeros(dim(v2)) u2 = vcat(u, out) - # compfg(v2)(out2, du2, u2, in, p, NaN) - b = @b $(compfg(v2))($out2, $du2, $u2, $in, $p, NaN) + # compfg(v2)((out2,), du2, u2, (in,), p, NaN) + b = @b $(compfg(v2))($(out2,), $du2, $u2, $(in,), $p, NaN) @test b.allocs == 0 @test out ≈ out2 @test du2 ≈ vcat(du, zeros(length(out))) diff --git a/test/external_inputs_test.jl b/test/external_inputs_test.jl new file mode 100644 index 00000000..d6af3be4 --- /dev/null +++ b/test/external_inputs_test.jl @@ -0,0 +1,69 @@ +using NetworkDynamics, Graphs +using Chairmarks: @b + +using NetworkDynamics: StateBufIdx, OutBufIdx, ExtMap, collect_externals! +using ModelingToolkit +using ModelingToolkit: t_nounits as t, D_nounits as Dt +using Random +using CUDA, Adapt + +@testset "test extmap performance" begin + N = 10_000 + u = rand(N) + o = rand(N) + _map = Union{StateBufIdx,OutBufIdx}[rand((StateBufIdx(i), OutBufIdx(i))) for i in 1:N] + map = ExtMap(Random.shuffle(_map)) + + exbuf = zeros(N) + fill!(exbuf, NaN) + b = @b collect_externals!($map, $exbuf, $u, $o) # 6.583μs + @test iszero(b.allocs) + + if CUDA.functional() + for to in (CuArray{Float64}, CuArray{Float32}) + map_d = adapt(to, map) + exbuf_d = adapt(to, zeros(N)) + u_d = adapt(to, u) + o_d = adapt(to, o) + b = @b collect_externals!($map_d, $exbuf_d, $u_d, $o_d) + display(b) + @test Vector(exbuf_d) ≈ exbuf + end + end +end + +@testset "test construction of elements with external inputs" begin + function fvext(dv, v, ein, ext, p, t) + dv[1] = ext[1] + end + function gvext(out, v, ein, ext, p, t) + out .= ext + end + v_noff = VertexModel(f=fvext, g=1, dim=1, extin=[VIndex(2,:a)]) + v_ff = VertexModel(f=fvext, g=gvext, outdim=1, dim=1, extin=[VIndex(2,:a)]) + + function feext(de, e, vsrc, vdst, ext, p ,t) + de[1] = ext[1] + end + function geext(outsrc, outdst, e, vsrc, vdst, ext, p, t) + outsrc .= ext + outdst .= ext + end + e_noff = EdgeModel(f=feext, g=AntiSymmetric(1), dim=1, extin=[EIndex(2,:a)]) + e_ff = EdgeModel(f=feext, g=geext, outdim=1, dim=1, extin=[EIndex(2,:a)]) + + @mtkmodel EdgeExt begin + @variables begin + srcin(t) + dstin(t) + extin1(t) + extin2(t) + out(t) + end + @equations begin + out ~ extin1 + extin2 + end + end + @named edge = EdgeExt() + edgem = EdgeModel(edge, :srcin, :dstin, Directed(:out); extin=[:extin1 => VIndex(2,:a), :extin2 => VIndex(2,:a)]) +end diff --git a/test/runtests.jl b/test/runtests.jl index 3022c7da..6e09cd62 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -34,6 +34,8 @@ using ExplicitImports NetworkDynamics.CHECK_COMPONENT[] = true @safetestset "Symbolic Indexing Tests" begin include("symbolicindexing_test.jl") end + @safetestset "external input test" begin include("external_inputs_test.jl") end + @safetestset "doctor test" begin include("doctor_test.jl") end @safetestset "Diffusion test" begin include("diffusion_test.jl") end