Skip to content

Commit

Permalink
Merge pull request #2836 from AayushSabharwal/as/more-metadata
Browse files Browse the repository at this point in the history
fix: fix edge cases with metadata dumping, add tests
  • Loading branch information
ChrisRackauckas authored Jul 23, 2024
2 parents 6311360 + 105bdbe commit 2c2e914
Show file tree
Hide file tree
Showing 3 changed files with 24 additions and 9 deletions.
6 changes: 5 additions & 1 deletion src/systems/abstractsystem.jl
Original file line number Diff line number Diff line change
Expand Up @@ -2776,11 +2776,15 @@ ModelingToolkit.dump_unknowns(sys)
See also: [`ModelingToolkit.dump_variable_metadata`](@ref), [`ModelingToolkit.dump_parameters`](@ref)
"""
function dump_unknowns(sys::AbstractSystem)
defs = defaults(sys)
defs = varmap_with_toterm(defaults(sys))
gs = varmap_with_toterm(guesses(sys))
map(dump_variable_metadata.(unknowns(sys))) do meta
if haskey(defs, meta.var)
meta = merge(meta, (; default = defs[meta.var]))
end
if haskey(gs, meta.var)
meta = merge(meta, (; guess = gs[meta.var]))
end
meta
end
end
Expand Down
13 changes: 11 additions & 2 deletions src/variables.jl
Original file line number Diff line number Diff line change
Expand Up @@ -30,8 +30,13 @@ ModelingToolkit.dump_variable_metadata(p)
function dump_variable_metadata(var)
uvar = unwrap(var)
vartype, name = get(uvar.metadata, VariableSource, (:unknown, :unknown))
shape = Symbolics.shape(var)
if shape == ()
type = symtype(uvar)
if type <: AbstractArray
shape = Symbolics.shape(var)
if shape == ()
shape = nothing
end
else
shape = nothing
end
unit = get(uvar.metadata, VariableUnit, nothing)
Expand Down Expand Up @@ -208,6 +213,10 @@ function _varmap_to_vars(varmap::Dict, varlist; defaults = Dict(), check = false
return [values[unwrap(var)] for var in varlist]
end

function varmap_with_toterm(varmap; toterm = Symbolics.diff2term)
return merge(todict(varmap), Dict(toterm(unwrap(k)) => v for (k, v) in varmap))
end

function canonicalize_varmap(varmap; toterm = Symbolics.diff2term)
new_varmap = Dict()
for (k, v) in varmap
Expand Down
14 changes: 8 additions & 6 deletions test/test_variable_metadata.jl
Original file line number Diff line number Diff line change
Expand Up @@ -131,13 +131,15 @@ sp = Set(p)

@test_nowarn show(stdout, "text/plain", sys)

# Defaults overridden by system, parameter dependencies
@variables x(t) = 1.0
# Defaults, guesses overridden by system, parameter dependencies
@variables x(t)=1.0 y(t) [guess = 1.0]
@parameters p=2.0 q
@named sys = ODESystem(Equation[], t, [x], [p]; defaults = Dict(x => 2.0, p => 3.0),
parameter_dependencies = [q => 2p])
x_meta = ModelingToolkit.dump_unknowns(sys)[]
@test x_meta.default == 2.0
@named sys = ODESystem(Equation[], t, [x, y], [p]; defaults = Dict(x => 2.0, p => 3.0),
guesses = Dict(y => 2.0), parameter_dependencies = [q => 2p])
unks_meta = ModelingToolkit.dump_unknowns(sys)
unks_meta = Dict([ModelingToolkit.getname(meta.var) => meta for meta in unks_meta])
@test unks_meta[:x].default == 2.0
@test unks_meta[:y].guess == 2.0
params_meta = ModelingToolkit.dump_parameters(sys)
params_meta = Dict([ModelingToolkit.getname(meta.var) => meta for meta in params_meta])
@test params_meta[:p].default == 3.0
Expand Down

0 comments on commit 2c2e914

Please sign in to comment.