Skip to content

Commit

Permalink
feat: update the vartype of kwargs rel. to symbolic-array
Browse files Browse the repository at this point in the history 
While here, removes passing `indices` around to generate_var, update_kwargs_and_metadata! and parse_variable_def! as these no longer handles symbolic-arrays.
  • Loading branch information
@ven-k
ven-k committed Oct 8, 2024
1 parent 78f0967 commit 5239e4f
Showing 1 changed file with 36 additions and 46 deletions.
82 changes: 36 additions & 46 deletions src/systems/model_parsing.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -148,31 +148,17 @@ end

pop_structure_dict!(dict, key) = length(dict[key]) == 0 && pop!(dict, key)

function update_kwargs_and_metadata!(dict, kwargs, a, def, indices, type, var,
function update_kwargs_and_metadata!(dict, kwargs, a, def, type,
varclass, where_types, meta)
if indices isa Nothing
if !isnothing(meta) && haskey(meta, VariableUnit)
uvar = gensym()
push!(where_types, uvar)
push!(kwargs, Expr(:kw, :($a::Union{Nothing, $uvar}), nothing))
else
push!(kwargs, Expr(:kw, :($a::Union{Nothing, $type}), nothing))
end
dict[:kwargs][a] = Dict(:value => def, :type => type)
if !isnothing(meta) && haskey(meta, VariableUnit)
uvar = gensym()
push!(where_types, uvar)
push!(kwargs, Expr(:kw, :($a::Union{Nothing, $uvar}), nothing))
else
vartype = gensym(:T)
push!(kwargs,
Expr(:kw,
Expr(:(::), a,
Expr(:curly, :Union, :Nothing, Expr(:curly, :AbstractArray, vartype))),
nothing))
if !isnothing(meta) && haskey(meta, VariableUnit)
push!(where_types, vartype)
else
push!(where_types, :($vartype <: $type))
end
dict[:kwargs][a] = Dict(:value => def, :type => AbstractArray{type})
push!(kwargs, Expr(:kw, :($a::Union{Nothing, $type}), nothing))
end
dict[:kwargs][a] = Dict(:value => def, :type => type)

if dict[varclass] isa Vector
dict[varclass][1][a][:type] = AbstractArray{type}
else
Expand Down Expand Up @@ -207,8 +193,8 @@ function update_readable_metadata!(varclass_dict, meta::Dict, varname)
end
end

function push_array_kwargs_and_metadata!(
dict, indices, meta, type, varclass, varname, varval)
function update_array_kwargs_and_metadata!(
dict, indices, kwargs, meta, type, varclass, varname, varval, where_types)
dict[varclass] = get!(dict, varclass) do
Dict{Symbol, Dict{Symbol, Any}}()
end
Expand All @@ -221,6 +207,18 @@ function push_array_kwargs_and_metadata!(
:type => type
)))

vartype = gensym(:T)
push!(kwargs,
Expr(:kw,
Expr(:(::), varname,
Expr(:curly, :Union, :Nothing, Expr(:curly, :AbstractArray, vartype))),
nothing))
if !isnothing(meta) && haskey(meta, VariableUnit)
push!(where_types, vartype)
else
push!(where_types, :($vartype <: $type))
end

# Useful keys for kwargs entry are: value, type and size.
dict[:kwargs][varname] = varclass_dict[][varname]

Expand All @@ -237,13 +235,12 @@ function unit_handled_variable_value(meta, varname)
end

function parse_variable_def!(dict, mod, arg, varclass, kwargs, where_types;
def = nothing, indices::Union{Vector{UnitRange{Int}}, Nothing} = nothing,
type::Type = Real, meta = Dict{DataType, Expr}())
def = nothing, type::Type = Real, meta = Dict{DataType, Expr}())
arg isa LineNumberNode && return
MLStyle.@match arg begin
a::Symbol => begin
var = generate_var!(dict, a, varclass; indices, type)
update_kwargs_and_metadata!(dict, kwargs, a, def, indices, type, var,
var = generate_var!(dict, a, varclass; type)
update_kwargs_and_metadata!(dict, kwargs, a, def, type,
varclass, where_types, meta)
return var, def, Dict()
end
Expand All @@ -259,8 +256,8 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs, where_types;
dict, mod, a, varclass, kwargs, where_types; def, type, meta)
end
Expr(:call, a, b) => begin
var = generate_var!(dict, a, b, varclass, mod; indices, type)
update_kwargs_and_metadata!(dict, kwargs, a, def, indices, type, var,
var = generate_var!(dict, a, b, varclass, mod; type)
update_kwargs_and_metadata!(dict, kwargs, a, def, type,
varclass, where_types, meta)
return var, def, Dict()
end
Expand All @@ -270,7 +267,6 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs, where_types;
Expr(:tuple, Expr(:ref, a, indices...), meta_val) => begin
(@isdefined type) || (type = Real)
varname = Meta.isexpr(a, :call) ? a.args[1] : a
push!(kwargs, Expr(:kw, varname, nothing))
meta = parse_metadata(mod, meta_val)
varval = (@isdefined default_val) ? default_val :
unit_handled_variable_value(meta, varname)
Expand All @@ -285,8 +281,8 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs, where_types;
var = :($varname = $first(@variables ($a[$(indices)]::$type = $varval),
$meta_val))
end
push_array_kwargs_and_metadata!(
dict, indices, meta, type, varclass, varname, varval)
update_array_kwargs_and_metadata!(
dict, indices, kwargs, meta, type, varclass, varname, varval, where_types)
(:($varname...), var), nothing, Dict()
end
Expr(:(=), Expr(:(::), Expr(:ref, a, indices...), type), def_n_meta) ||
Expand All @@ -297,7 +293,6 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs, where_types;
meta = parse_metadata(mod, def_n_meta)
varval = unit_handled_variable_value(meta, varname)
val, def_n_meta = (def_n_meta.args[1], def_n_meta.args[2:end])
push!(kwargs, Expr(:kw, varname, nothing))
if varclass == :parameters
Meta.isexpr(a, :call) &&
assert_unique_independent_var(dict, a.args[end])
Expand All @@ -314,7 +309,6 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs, where_types;
$(def_n_meta...)))
end
else
push!(kwargs, Expr(:kw, varname, nothing))
if varclass == :parameters
Meta.isexpr(a, :call) &&
assert_unique_independent_var(dict, a.args[end])
Expand All @@ -329,15 +323,14 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs, where_types;
end
varval, meta = def_n_meta, nothing
end
push_array_kwargs_and_metadata!(
dict, indices, meta, type, varclass, varname, varval)
update_array_kwargs_and_metadata!(
dict, indices, kwargs, meta, type, varclass, varname, varval, where_types)
(:($varname...), var), nothing, Dict()
end
Expr(:(::), Expr(:ref, a, indices...), type) ||
Expr(:ref, a, indices...) => begin
(@isdefined type) || (type = Real)
varname = a isa Expr && a.head == :call ? a.args[1] : a
push!(kwargs, Expr(:kw, varname, nothing))
if varclass == :parameters
Meta.isexpr(a, :call) && assert_unique_independent_var(dict, a.args[end])
var = :($varname = $first(@parameters $a[$(indices...)]::$type = $varname))
Expand All @@ -350,8 +343,8 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs, where_types;
throw("Symbolic array with arbitrary length is not handled for $varclass.
Please open an issue with an example.")
end
push_array_kwargs_and_metadata!(
dict, indices, nothing, type, varclass, varname, nothing)
update_array_kwargs_and_metadata!(
dict, indices, kwargs, nothing, type, varclass, varname, nothing, where_types)
(:($varname...), var), nothing, Dict()
end
Expr(:(=), a, b) => begin
Expand Down Expand Up @@ -391,11 +384,8 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs, where_types;
end
end

function generate_var(a, varclass;
indices::Union{Vector{UnitRange{Int}}, Nothing} = nothing,
type = Real)
var = indices === nothing ? Symbolics.variable(a; T = type) :
first(@variables $a[indices...]::type)
function generate_var(a, varclass; type = Real)
var = Symbolics.variable(a; T = type)
if varclass == :parameters
var = toparam(var)
elseif varclass == :independent_variables
Expand Down Expand Up @@ -425,7 +415,7 @@ function generate_var!(dict, a, varclass;
vd isa Vector && (vd = first(vd))
vd[a] = Dict{Symbol, Any}()
indices !== nothing && (vd[a][:size] = Tuple(lastindex.(indices)))
generate_var(a, varclass; indices, type)
generate_var(a, varclass; type)
end

function assert_unique_independent_var(dict, iv::Num)
Expand Down

0 comments on commit 5239e4f

Please sign in to comment.