Merge pull request #191 from JuliaDynamics/hw/initbounds

add variable transformation to respect bounds in initialization problem

Author: hexaeder

docs/src/API.md

@@ -128,6 +128,8 @@ set_bounds!
 find_fixpoint
 initialize_component!
 init_residual
+get_initial_state
+dump_initial_state
 ```
 
 ## Execution Types

docs/src/initialization.md

@@ -67,3 +67,8 @@ As a quick test we can ensure that the angle indeed matches the voltag angel:
 ```@example compinit
 get_init(vf, :θ) ≈ atan(get_default(vf, :u_i), get_default(vf, :u_r))
 ```
+
+It is possible to inspect initial states (also for observed symbols) using [`get_initial_state`](@ref). You can print out the whole state using [`dump_initial_state`](@ref).
+```@example compinit
+dump_initial_state(vf)
+```

ext/MTKExt.jl

@@ -262,14 +262,13 @@ function generate_io_function(_sys, inputss::Tuple, outputss::Tuple;
         deepcopy(_sys)
     else
         _openinputs = setdiff(allinputs, Set(full_parameters(_sys)))
-        get_variables.(full_equations(_sys))
         all_eq_vars = mapreduce(get_variables, union, full_equations(_sys), init=Set{Symbolic}())
         if !(_openinputs ⊆ all_eq_vars)
             missing_inputs = setdiff(_openinputs, all_eq_vars)
             @warn "The specified inputs ($missing_inputs) do not appear in the equations of the system!"
             _openinputs = setdiff(_openinputs, missing_inputs)
         end
-        structural_simplify(_sys, (_openinputs, alloutputs); simplify=true)[1]
+        structural_simplify(_sys, (_openinputs, alloutputs); simplify=false)[1]
     end
 
     states = unknowns(sys)
@@ -334,6 +333,10 @@ function generate_io_function(_sys, inputss::Tuple, outputss::Tuple;
                 throw(ArgumentError("Output $out was neither foundin states nor in observed equations."))
             end
             eq = obseqs[idx]
+            if !isempty(rhs_differentials(eq))
+                println(obs_subs[out])
+                throw(ArgumentError("Algebraic FF equation for output $out contains differentials in the RHS: $(rhs_differentials(eq))"))
+            end
             deleteat!(obseqs, idx)
 
             if ff_to_constraint && !isempty(get_variables(eq.rhs) ∩ allinputs)
@@ -417,7 +420,7 @@ function generate_io_function(_sys, inputss::Tuple, outputss::Tuple;
             obsf = obsf_ip,
             equations=formulas,
             outputeqs=Dict(Iterators.flatten(outputss) .=> gformulas),
-            observed=Dict(obsstates .=> obsformulas),
+            observed=Dict(getname.(obsstates) .=> obsformulas),
             params)
 end
 

src/NetworkDynamics.jl

@@ -69,6 +69,7 @@ export has_guess, get_guess, set_guess!
 export has_init, get_init, set_init!
 export has_bounds, get_bounds, set_bounds!
 export has_graphelement, get_graphelement, set_graphelement!
+export get_initial_state, dump_initial_state
 include("metadata.jl")
 
 using NonlinearSolve: AbstractNonlinearSolveAlgorithm, NonlinearFunction

src/initialization.jl

@@ -48,7 +48,7 @@ function _solve_fixpoint(prob, alg::SteadyStateDiffEqAlgorithm; kwargs...)
     sol = SciMLBase.solve(prob, alg; kwargs...)
 end
 
-function initialization_problem(cf::T; t=NaN, verbose=true) where {T<:ComponentModel}
+function initialization_problem(cf::T; t=NaN, apply_bound_transformation, verbose=true) where {T<:ComponentModel}
     hasinsym(cf) || throw(ArgumentError("Component model musst have `insym`!"))
 
     outfree_ms = Tuple((!).(map(s -> has_default(cf, s), sv)) for sv in outsym_normalized(cf))
@@ -92,6 +92,48 @@ function initialization_problem(cf::T; t=NaN, verbose=true) where {T<:ComponentM
     unl_range_p = nextrange(pfree_m)
     @assert vcat(unl_range_outs..., unl_range_u, unl_range_ins..., unl_range_p) == 1:Nfree
 
+    # check for positivity and negativity constraints
+    bounds = map(freesym) do sym
+        if has_bounds(cf, sym)
+            bound = get_bounds(cf, sym)
+            if bound[1] >= 0 && bound[2] > bound[1]
+                return :pos
+            elseif bound[1] < bound[2]  && bound[2] <= 0
+                return :neg
+            end
+        end
+        :none
+    end
+    if !apply_bound_transformation || all(isequal(:none), bounds)
+        boundT! = identity
+        inv_boundT! = identity
+    else
+        if verbose
+            idxs = findall(!isequal(:none), bounds)
+            @info "Apply positivity/negativity conserving variable transformation on $(freesym[idxs]) to satisfy bounds."
+        end
+        boundT! = (u) -> begin
+            for i in eachindex(u, bounds)
+                if bounds[i] == :pos
+                    u[i] = u[i]^2
+                elseif bounds[i] == :neg
+                    u[i] = -u[i]^2
+                end
+            end
+            return u
+        end
+        inv_boundT! = (u) -> begin
+            for i in eachindex(u, bounds)
+                if bounds[i] == :pos
+                    u[i] = sqrt(u[i])
+                elseif bounds[i] == :neg
+                    u[i] = sqrt(-u[i])
+                end
+            end
+            return u
+        end
+    end
+
     # check for missin guesses
     missing_guesses = Symbol[]
     uguess = map(freesym) do s
@@ -101,67 +143,68 @@ function initialization_problem(cf::T; t=NaN, verbose=true) where {T<:ComponentM
             push!(missing_guesses, s)
         end
     end
+    isempty(missing_guesses) || throw(ArgumentError("Missing guesses for free variables $(missing_guesses)"))
 
-    for s in freesym
-        if has_bounds(cf, s)
-            @warn "Ignore bounds $(get_bounds(cf, s)) for $s. Not supported yet."
-        end
+    # apply bound conserving transformation to initial state
+    try
+        inv_boundT!(uguess)
+    catch
+        throw(ArgumentError("Initial guess violates bounds. Check the docstring on `NetworkDynamics.initialize_component!`\
+                             about bound satisfying transformations!"))
     end
 
-    isempty(missing_guesses) || throw(ArgumentError("Missing guesses for free variables $(missing_guesses)"))
-
     N = ForwardDiff.pickchunksize(Nfree)
-    fz = let fg = compfg(cf),
-        outcaches=map(d->DiffCache(zeros(d), N), outdim_normalized(cf)),
-        ucache=DiffCache(zeros(dim(cf)), N),
-        incaches=map(d->DiffCache(zeros(d), N), indim_normalized(cf)),
-        pcache=DiffCache(zeros(pdim(cf))),
-        t=t,
-        outfree_ms=outfree_ms, ufree_m=ufree_m, infree_ms=infree_ms, pfree_m=pfree_m,
-        outfixs=outfixs, ufix=ufix, infixs=infixs, pfix=pfix,
-        unl_range_outs=unl_range_outs, unl_range_u=unl_range_u,
-        unl_range_ins=unl_range_ins, unl_range_p=unl_range_p
-
-        (dunl, unl, _) -> begin
-            outbufs = PreallocationTools.get_tmp.(outcaches, Ref(dunl))
-            ubuf = PreallocationTools.get_tmp(ucache, dunl)
-            inbufs = PreallocationTools.get_tmp.(incaches, Ref(dunl))
-            pbuf = PreallocationTools.get_tmp(pcache, dunl)
-
-            # prefill buffers with fixed values
-            for (buf, fix) in zip(outbufs, outfixs)
-                buf .= fix
-            end
-            ubuf .= ufix
-            for (buf, fix) in zip(inbufs, infixs)
-                buf .= fix
-            end
-            pbuf .= pfix
+    fg = compfg(cf)
+    unlcache = map(d->DiffCache(zeros(d), N), length(freesym))
+    outcaches = map(d->DiffCache(zeros(d), N), outdim_normalized(cf))
+    ucache = DiffCache(zeros(dim(cf)), N)
+    incaches = map(d->DiffCache(zeros(d), N), indim_normalized(cf))
+    pcache = DiffCache(zeros(pdim(cf)))
+
+    fz = (dunl, unl, _) -> begin
+        # apply the bound conserving transformation
+        unlbuf = PreallocationTools.get_tmp(unlcache, unl)
+        unlbuf .= unl
+        boundT!(unlbuf)
+
+        outbufs = PreallocationTools.get_tmp.(outcaches, Ref(dunl))
+        ubuf = PreallocationTools.get_tmp(ucache, dunl)
+        inbufs = PreallocationTools.get_tmp.(incaches, Ref(dunl))
+        pbuf = PreallocationTools.get_tmp(pcache, dunl)
+
+        # prefill buffers with fixed values
+        for (buf, fix) in zip(outbufs, outfixs)
+            buf .= fix
+        end
+        ubuf .= ufix
+        for (buf, fix) in zip(inbufs, infixs)
+            buf .= fix
+        end
+        pbuf .= pfix
 
-            # overwrite nonfixed values
-            for (buf, mask, range) in zip(outbufs, outfree_ms, unl_range_outs)
-                buf[mask] .= unl[range]
-            end
-            ubuf[ufree_m] .= unl[unl_range_u]
-            for (buf, mask, range) in zip(inbufs, infree_ms, unl_range_ins)
-                buf[mask] .= unl[range]
-            end
-            pbuf[pfree_m] .= unl[unl_range_p]
+        # overwrite nonfixed values
+        for (buf, mask, range) in zip(outbufs, outfree_ms, unl_range_outs)
+            _overwrite_at_mask!(buf, mask, unlbuf, range)
+        end
+        _overwrite_at_mask!(ubuf, ufree_m, unlbuf, unl_range_u)
+        for (buf, mask, range) in zip(inbufs, infree_ms, unl_range_ins)
+            _overwrite_at_mask!(buf, mask, unlbuf, range)
+        end
+        _overwrite_at_mask!(pbuf, pfree_m, unlbuf, unl_range_p)
 
-            # view into du buffer for the fg funtion
-            @views dunl_fg = dunl[1:dim(cf)]
-            # view into the output buffer for the outputs
-            @views dunl_out = dunl[dim(cf)+1:end]
+        # view into du buffer for the fg funtion
+        @views dunl_fg = dunl[1:dim(cf)]
+        # view into the output buffer for the outputs
+        @views dunl_out = dunl[dim(cf)+1:end]
 
-            # 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)
+        # 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)
 
-            # calculate the residual for the second half ov the dunl buf, the outputs
-            dunl_out .= dunl_out .- RecursiveArrayTools.ArrayPartition(outbufs...)
-            nothing
-        end
+        # calculate the residual for the second half ov the dunl buf, the outputs
+        dunl_out .= dunl_out .- RecursiveArrayTools.ArrayPartition(outbufs...)
+        nothing
     end
 
     nlf = NonlinearFunction(fz; resid_prototype=zeros(Neqs), sys=SII.SymbolCache(freesym))
@@ -171,11 +214,21 @@ function initialization_problem(cf::T; t=NaN, verbose=true) where {T<:ComponentM
     else
         verbose && @info "Initialization problem is overconstrained ($Nfree vars for $Neqs equations). Create NonlinearLeastSquaresProblem for $freesym."
     end
-    NonlinearLeastSquaresProblem(nlf, uguess)
+    (NonlinearLeastSquaresProblem(nlf, uguess), boundT!)
+end
+function _overwrite_at_mask!(target, mask, source, range)
+    src_v = view(source, range)
+    j = 1
+    for i in eachindex(target, mask)
+        if mask[i]
+            target[i] = src_v[j]
+            j += 1
+        end
+    end
 end
 
 """
-    initialize_component!(cf::ComponentModel; verbose=true, kwargs...)
+    initialize_component!(cf::ComponentModel; verbose=true, apply_bound_transformation=true, kwargs...)
 
 Initialize a `ComponentModel` by solving the corresponding `NonlinearLeastSquaresProblem`.
 During initialization, everyting which has a `default` value (see [Metadata](@ref)) is considered
@@ -186,9 +239,16 @@ The result is stored in the `ComponentModel` itself. The values of the free vari
 in the metadata field `init`.
 
 The `kwargs` are passed to the nonlinear solver.
+
+## Bounds of free variables
+When encountering any bounds in the free variables, NetworkDynamics will try to conserve them
+by applying a coordinate transforamtion. This behavior can be supressed by setting `apply_bound_transformation`.
+The following transformations are used:
+- (a, b) intervals where both a and b are positive are transformed to `u^2`/`sqrt(u)`
+- (a, b) intervals where both a and b are negative are transformed to `-u^2`/`sqrt(-u)`
 """
-function initialize_component!(cf; verbose=true, kwargs...)
-    prob = initialization_problem(cf; verbose)
+function initialize_component!(cf; verbose=true, apply_bound_transformation=true, kwargs...)
+    prob, boundT! = initialization_problem(cf; verbose, apply_bound_transformation)
 
     if !isempty(prob.u0)
         sol = SciMLBase.solve(prob; kwargs...)
@@ -202,14 +262,22 @@ function initialize_component!(cf; verbose=true, kwargs...)
         else
             verbose && @info "Initialization successful with residual $(LinearAlgebra.norm(sol.resid))"
         end
-        set_init!.(Ref(cf), SII.variable_symbols(sol), sol.u)
+        # transform back to original space
+        u = boundT!(copy(sol.u))
+        set_init!.(Ref(cf), SII.variable_symbols(sol), u)
         resid = sol.resid
     else
         resid = init_residual(cf; recalc=true)
         verbose && @info "No free variables! Residual $(LinearAlgebra.norm(resid))"
     end
 
     set_metadata!(cf, :init_residual, resid)
+
+    broken = broken_bounds(cf)
+    if !isempty(broken)
+        @warn "Initialized model has broken bounds for $(broken). Use `dump_initial_state(mode)` \
+        to inspect further and try to adapt the initial guesses!"
+    end
     cf
 end
 
@@ -250,3 +318,17 @@ function init_residual(cf::T; t=NaN, recalc=false) where {T<:ComponentModel}
         LinearAlgebra.norm(get_metadata(cf, :init_residual))
     end
 end
+
+function broken_bounds(cf)
+    allsyms = vcat(sym(cf), psym(cf), insym_all(cf), outsym_flat(cf), obssym(cf))
+    boundsyms = filter(s -> has_bounds(cf, s), allsyms)
+    bounds = get_bounds.(Ref(cf), boundsyms)
+    vals = get_initial_state(cf, boundsyms)
+    broken = filter(i -> !bounds_satisfied(vals[i], bounds[i]), 1:length(bounds))
+    boundsyms[broken]
+end
+
+function bounds_satisfied(val, bounds)
+    @assert length(bounds) == 2
+    !isnothing(val) && !isnan(val) && first(bounds) ≤ val ≤ last(bounds)
+end