Merge pull request #339 from JuliaDynamics/hw/fixes

small fixes for pd tests

Author: hexaeder

NEWS.md

@@ -1,5 +1,8 @@
 # NetworkDynamics Release Notes
 
+## v0.10.15 Changelog
+- small fixes, mainly for PowerDynamics tests ([#339](https://github.com/JuliaDynamics/NetworkDynamics.jl/pull/339))
+
 ## v0.10.14 Changelog
 - Add `assume_io_coupling` parameter to MTK `VertexModel` and `EdgeModel` constructors ([#332](https://github.com/JuliaDynamics/NetworkDynamics.jl/pull/332))
   - New optional parameter forces MTK to consider direct dependency chains from outputs to inputs

Project.toml

@@ -7,6 +7,7 @@ authors = ["Frank Hellmann <[email protected]>, Michael Lindner <michaelli
 projects = ["test", "benchmark"]
 
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
 Atomix = "a9b6321e-bd34-4604-b9c9-b65b8de01458"
@@ -58,6 +59,7 @@ NetworkDynamicsSparsityExt = ["SparseConnectivityTracer"]
 NetworkDynamicsSymbolicsExt = ["Symbolics", "MacroTools"]
 
 [compat]
+ADTypes = "1"
 Accessors = "0.1.43"
 Adapt = "4.0.4"
 ArgCheck = "2.3.0"

ext/NetworkDynamicsMTKExt.jl

@@ -238,7 +238,7 @@ function _get_metadata(sys, name, alldefaults, sysguesses)
         if !endswith(string(name), "ˍt") # known for "internal" derivatives
             @warn "Could not extract metadata for $name $(e.msg)"
         end
-        return nt
+        return md
     end
     if haskey(alldefaults, sym)
         def = alldefaults[sym]

src/NetworkDynamics.jl

@@ -25,6 +25,7 @@ using DiffEqCallbacks: DiffEqCallbacks
 using MacroTools: postwalk, @capture
 using ConstructionBase: ConstructionBase
 using Accessors: Accessors
+using ADTypes: AutoForwardDiff
 
 @static if VERSION ≥ v"1.11.0-0"
     using Base: AnnotatedIOBuffer, AnnotatedString

src/initialization.jl

@@ -498,7 +498,10 @@ The function solves a nonlinear problem to find values for all free variables/pa
 - `t`: Time at which to solve for steady state. Only relevant for components with explicit time dependency.
 - `tol`: Tolerance for the residual of the initialized model (defaults to `1e-10`). Init throws error if resid ≥ tol.
 - `residual`: Optional `Ref{Float64}` which gets the final residual of the initialized model.
-- `alg=nothing`: Nonlinear solver algorithm (defaults to NonlinearSolve.jl default with QR factorization, since init problems tend to be ill-conditioned.)
+- `alg_kwargs=(;)`: Additional keyword arguments passed to the nonlinear solver algorithm constructor
+- `alg=FastShortcutNLLSPolyalg(; linsolve=QRFactorization(), autodiff=AutoForwardDiff(), alg_kwargs...)`
+
+   Nonlinear solver algorithm (defaults to NonlinearSolve.jl default with QR factorization, since init problems tend to be ill-conditioned.)
 - `solve_kwargs=(;)`: Additional keyword arguments passed to the SciML `solve` function
 - `io=stdout`: IO stream for printing information
 
@@ -530,7 +533,8 @@ function initialize_component(cf;
                              # internal keywords to "return" the final defaults/guesses after applying formulas
                              _final_defaults=nothing,
                              _final_guesses=nothing,
-                             alg=nothing,
+                             alg_kwargs=(;),
+                             alg=FastShortcutNLLSPolyalg(; linsolve=QRFactorization(), autodiff=AutoForwardDiff(), alg_kwargs...),
                              solve_kwargs=(;),
                              io=stdout,
                              kwargs...)
@@ -539,10 +543,6 @@ function initialize_component(cf;
         @warn "Passing `kwargs` to `initialize_component(!)` is deprecated. Use `alg` and `solve_kwargs=(; kw=val)` instead."
     end
 
-    if alg == nothing
-        alg = FastShortcutNLLSPolyalg(linsolve=QRFactorization())
-    end
-
     defaults = isnothing(default_overrides) ? defaults : merge(defaults, default_overrides)
     guesses  = isnothing(guess_overrides)   ? guesses  : merge(guesses, guess_overrides)
     bounds   = isnothing(bound_overrides)   ? bounds   : merge(bounds, bound_overrides)

src/symbolicindexing.jl

@@ -109,6 +109,7 @@ struct NWState{U,P,T,NW<:Network}
     t::T
     function NWState(thing, uflat::AbstractVector, p=nothing, t=nothing)
         nw = extract_nw(thing)
+        @argcheck length(uflat) == dim(nw) "Length of uflat ($(length(uflat))) does not match number of states in network ($(dim(nw)))."
         _p = p isa Union{NWParameter,Nothing} ? p : NWParameter(nw, p)
         s = new{typeof(uflat),typeof(_p),typeof(t),typeof(nw)}(nw,uflat,_p,t)
         @argcheck isnothing(p) || s.nw === s.p.nw

test/AD_test.jl

@@ -32,10 +32,10 @@ end
 # jacobian(fx, AutoReverseDiff(), x0)
 # jacobian(fx, AutoFiniteDiff(), x0)
 # jacobian(fx, AutoMooncake(), x0)
-@test_broken jacobian(fx, AutoEnzyme(; mode=EnzymeCore.Forward, function_annotation=EnzymeCore.Duplicated), x0)
-@test_broken jacobian(fx, AutoEnzyme(; mode=EnzymeCore.Reverse, function_annotation=EnzymeCore.Duplicated), x0)
-@test_broken jacobian(fx, AutoEnzyme(; mode=Enzyme.set_runtime_activity(EnzymeCore.Forward), function_annotation=EnzymeCore.Duplicated), x0)
-@test_broken jacobian(fx, AutoEnzyme(; mode=Enzyme.set_runtime_activity(EnzymeCore.Reverse), function_annotation=EnzymeCore.Duplicated), x0)
+# @test_broken jacobian(fx, AutoEnzyme(; mode=EnzymeCore.Forward, function_annotation=EnzymeCore.Duplicated), x0)
+# @test_broken jacobian(fx, AutoEnzyme(; mode=EnzymeCore.Reverse, function_annotation=EnzymeCore.Duplicated), x0)
+# @test_broken jacobian(fx, AutoEnzyme(; mode=Enzyme.set_runtime_activity(EnzymeCore.Forward), function_annotation=EnzymeCore.Duplicated), x0)
+# @test_broken jacobian(fx, AutoEnzyme(; mode=Enzyme.set_runtime_activity(EnzymeCore.Reverse), function_annotation=EnzymeCore.Duplicated), x0)
 
 fp = function(p)
     dx = similar(p,length(x0))
@@ -46,10 +46,10 @@ end
 # jacobian(fp, AutoReverseDiff(), pflat(p0))
 # jacobian(fp, AutoFiniteDiff(), pflat(p0))
 # jacobian(fp, AutoMooncake(), pflat(p0))
-@test_broken jacobian(fp, AutoEnzyme(; mode=EnzymeCore.Forward, function_annotation=EnzymeCore.Duplicated), pflat(p0))
-@test_broken jacobian(fp, AutoEnzyme(; mode=EnzymeCore.Reverse, function_annotation=EnzymeCore.Duplicated), pflat(p0))
-@test_broken jacobian(fp, AutoEnzyme(; mode=Enzyme.set_runtime_activity(EnzymeCore.Forward), function_annotation=EnzymeCore.Duplicated), pflat(p0))
-@test_broken jacobian(fp, AutoEnzyme(; mode=Enzyme.set_runtime_activity(EnzymeCore.Reverse), function_annotation=EnzymeCore.Duplicated), pflat(p0))
+# @test_broken jacobian(fp, AutoEnzyme(; mode=EnzymeCore.Forward, function_annotation=EnzymeCore.Duplicated), pflat(p0))
+# @test_broken jacobian(fp, AutoEnzyme(; mode=EnzymeCore.Reverse, function_annotation=EnzymeCore.Duplicated), pflat(p0))
+# @test_broken jacobian(fp, AutoEnzyme(; mode=Enzyme.set_runtime_activity(EnzymeCore.Forward), function_annotation=EnzymeCore.Duplicated), pflat(p0))
+# @test_broken jacobian(fp, AutoEnzyme(; mode=Enzyme.set_runtime_activity(EnzymeCore.Reverse), function_annotation=EnzymeCore.Duplicated), pflat(p0))
 
 scenarios = [Scenario{:jacobian, :in}(fx, x0; res1=jacobian(fx, AutoFiniteDiff(), x0)) ,
              Scenario{:jacobian, :in}(fp, pflat(p0); res1=jacobian(fp, AutoFiniteDiff(), pflat(p0)))]