Skip to content

Commit b0ab722

Browse files
ven-kven-k
committed
feat: support arbitrary length arrays with metadata and default
Update metadata dict and kwargs for symbolic arrays
1 parent ae22c71 commit b0ab722

File tree

1 file changed

+127
-80
lines changed

1 file changed

+127
-80
lines changed

‎src/systems/model_parsing.jl

Lines changed: 127 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -245,20 +245,6 @@ end
245245
function parse_variable_def!(dict, mod, arg, varclass, kwargs, where_types;
246246
def = nothing, indices::Union{Vector{UnitRange{Int}}, Nothing} = nothing,
247247
type::Type = Real, meta = Dict{DataType, Expr}())
248-
metatypes = [(:connection_type, VariableConnectType),
249-
(:description, VariableDescription),
250-
(:unit, VariableUnit),
251-
(:bounds, VariableBounds),
252-
(:noise, VariableNoiseType),
253-
(:input, VariableInput),
254-
(:output, VariableOutput),
255-
(:irreducible, VariableIrreducible),
256-
(:state_priority, VariableStatePriority),
257-
(:misc, VariableMisc),
258-
(:disturbance, VariableDisturbance),
259-
(:tunable, VariableTunable),
260-
(:dist, VariableDistribution)]
261-
262248
arg isa LineNumberNode && return
263249
MLStyle.@match arg begin
264250
a::Symbol => begin
@@ -284,27 +270,86 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs, where_types;
284270
varclass, where_types, meta)
285271
return var, def, Dict()
286272
end
273+
Expr(:tuple, Expr(:(::), Expr(:ref, a, indices...), type), meta_val) ||
274+
Expr(:tuple, Expr(:ref, a, indices...), meta_val) => begin
275+
(@isdefined type) || (type = Real)
276+
varname = Meta.isexpr(a, :call) ? a.args[1] : a
277+
push!(kwargs, Expr(:kw, varname, nothing))
278+
meta = parse_metadata(mod, meta_val)
279+
varval = (@isdefined default_val) ? default_val :
280+
unit_handled_variable_value(meta, varname)
281+
if varclass == :parameters
282+
var = :($varname = $first(@parameters ($a[$(indices...)]::$type = $varval),
283+
$meta_val))
284+
else
285+
var = :($varname = $first(@variables ($a[$(indices)]::$type = $varval),
286+
$meta_val))
287+
end
288+
push_array_kwargs_and_metadata!(
289+
dict, indices, meta, type, varclass, varname, varval)
290+
(:($varname...), var), nothing, Dict()
291+
end
292+
Expr(:(=), Expr(:(::), Expr(:ref, a, indices...), type), def_n_meta) ||
293+
Expr(:(=), Expr(:ref, a, indices...), def_n_meta) => begin
294+
(@isdefined type) || (type = Real)
295+
varname = Meta.isexpr(a, :call) ? a.args[1] : a
296+
if Meta.isexpr(def_n_meta, :tuple)
297+
meta = parse_metadata(mod, def_n_meta)
298+
varval = unit_handled_variable_value(meta, varname)
299+
val, def_n_meta = (def_n_meta.args[1], def_n_meta.args[2:end])
300+
push!(kwargs, Expr(:kw, varname, nothing))
301+
if varclass == :parameters
302+
var = :($varname = $varname === nothing ? $val : $varname;
303+
$varname = $first(@parameters ($a[$(indices...)]::$type = $varval),
304+
$(def_n_meta...)))
305+
else
306+
var = :($varname = $varname === nothing ? $val : $varname;
307+
$varname = $first(@variables $a[$(indices...)]::$type = (
308+
$varval),
309+
$(def_n_meta...)))
310+
end
311+
else
312+
push!(kwargs, Expr(:kw, varname, nothing))
313+
if varclass == :parameters
314+
var = :($varname = $varname === nothing ? $def_n_meta : $varname;
315+
$varname = $first(@parameters $a[$(indices...)]::$type = $varname))
316+
else
317+
var = :($varname = $varname === nothing ? $def_n_meta : $varname;
318+
$varname = $first(@variables $a[$(indices...)]::$type = $varname))
319+
end
320+
varval, meta = def_n_meta, nothing
321+
end
322+
push_array_kwargs_and_metadata!(
323+
dict, indices, meta, type, varclass, varname, varval)
324+
(:($varname...), var), nothing, Dict()
325+
end
326+
Expr(:(::), Expr(:ref, a, indices...), type) ||
327+
Expr(:ref, a, indices...) => begin
328+
(@isdefined type) || (type = Real)
329+
varname = a isa Expr && a.head == :call ? a.args[1] : a
330+
push!(kwargs, Expr(:kw, varname, nothing))
331+
if varclass == :parameters
332+
var = :($varname = $first(@parameters $a[$(indices...)]::$type = $varname))
333+
elseif varclass == :variables
334+
var = :($varname = $first(@variables $a[$(indices...)]::$type = $varname))
335+
else
336+
throw("Symbolic array with arbitrary length is not handled for $varclass.
337+
Please open an issue with an example.")
338+
end
339+
push_array_kwargs_and_metadata!(
340+
dict, indices, nothing, type, varclass, varname, nothing)
341+
(:($varname...), var), nothing, Dict()
342+
end
287343
Expr(:(=), a, b) => begin
288344
Base.remove_linenums!(b)
289345
def, meta = parse_default(mod, b)
290346
var, def, _ = parse_variable_def!(
291347
dict, mod, a, varclass, kwargs, where_types; def, type, meta)
292-
if dict[varclass] isa Vector
293-
dict[varclass][1][getname(var)][:default] = def
294-
else
295-
dict[varclass][getname(var)][:default] = def
296-
end
348+
varclass_dict = dict[varclass] isa Vector ? Ref(dict[varclass][1]) :
349+
Ref(dict[varclass])
350+
varclass_dict[][getname(var)][:default] = def
297351
if meta !== nothing
298-
for (type, key) in metatypes
299-
if (mt = get(meta, key, nothing)) !== nothing
300-
key == VariableConnectType && (mt = nameof(mt))
301-
if dict[varclass] isa Vector
302-
dict[varclass][1][getname(var)][type] = mt
303-
else
304-
dict[varclass][getname(var)][type] = mt
305-
end
306-
end
307-
end
352+
update_readable_metadata!(varclass_dict[], meta, getname(var))
308353
var, metadata_with_exprs = set_var_metadata(var, meta)
309354
return var, def, metadata_with_exprs
310355
end
@@ -314,27 +359,15 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs, where_types;
314359
meta = parse_metadata(mod, b)
315360
var, def, _ = parse_variable_def!(
316361
dict, mod, a, varclass, kwargs, where_types; type, meta)
362+
varclass_dict = dict[varclass] isa Vector ? Ref(dict[varclass][1]) :
363+
Ref(dict[varclass])
317364
if meta !== nothing
318-
for (type, key) in metatypes
319-
if (mt = get(meta, key, nothing)) !== nothing
320-
key == VariableConnectType && (mt = nameof(mt))
321-
if dict[varclass] isa Vector
322-
dict[varclass][1][getname(var)][type] = mt
323-
else
324-
dict[varclass][getname(var)][type] = mt
325-
end
326-
end
327-
end
365+
update_readable_metadata!(varclass_dict[], meta, getname(var))
328366
var, metadata_with_exprs = set_var_metadata(var, meta)
329367
return var, def, metadata_with_exprs
330368
end
331369
return var, def, Dict()
332370
end
333-
Expr(:ref, a, b...) => begin
334-
indices = map(i -> UnitRange(i.args[2], i.args[end]), b)
335-
parse_variable_def!(dict, mod, a, varclass, kwargs, where_types;
336-
def, indices, type, meta)
337-
end
338371
_ => error("$arg cannot be parsed")
339372
end
340373
end
@@ -442,14 +475,23 @@ function parse_default(mod, a)
442475
end
443476
end
444477

445-
function parse_metadata(mod, a)
478+
function parse_metadata(mod, a::Expr)
446479
MLStyle.@match a begin
447-
Expr(:vect, eles...) => Dict(parse_metadata(mod, e) for e in eles)
480+
Expr(:vect, b...) => Dict(parse_metadata(mod, m) for m in b)
481+
Expr(:tuple, a, b...) => parse_metadata(mod, b)
448482
Expr(:(=), a, b) => Symbolics.option_to_metadata_type(Val(a)) => get_var(mod, b)
449483
_ => error("Cannot parse metadata $a")
450484
end
451485
end
452486

487+
function parse_metadata(mod, metadata::AbstractArray)
488+
ret = Dict()
489+
for m in metadata
490+
merge!(ret, parse_metadata(mod, m))
491+
end
492+
ret
493+
end
494+
453495
function _set_var_metadata!(metadata_with_exprs, a, m, v::Expr)
454496
push!(metadata_with_exprs, m => v)
455497
a
@@ -707,6 +749,7 @@ function parse_variable_arg!(exprs, vs, dict, mod, arg, varclass, kwargs, where_
707749
end
708750

709751
function convert_units(varunits::DynamicQuantities.Quantity, value)
752+
value isa Nothing && return nothing
710753
DynamicQuantities.ustrip(DynamicQuantities.uconvert(
711754
DynamicQuantities.SymbolicUnits.as_quantity(varunits), value))
712755
end
@@ -718,6 +761,7 @@ function convert_units(
718761
end
719762

720763
function convert_units(varunits::Unitful.FreeUnits, value)
764+
value isa Nothing && return nothing
721765
Unitful.ustrip(varunits, value)
722766
end
723767

@@ -736,47 +780,50 @@ end
736780
function parse_variable_arg(dict, mod, arg, varclass, kwargs, where_types)
737781
vv, def, metadata_with_exprs = parse_variable_def!(
738782
dict, mod, arg, varclass, kwargs, where_types)
739-
name = getname(vv)
740-
741-
varexpr = if haskey(metadata_with_exprs, VariableUnit)
742-
unit = metadata_with_exprs[VariableUnit]
743-
quote
744-
$name = if $name === $NO_VALUE
745-
$setdefault($vv, $def)
746-
else
747-
try
748-
$setdefault($vv, $convert_units($unit, $name))
749-
catch e
750-
if isa(e, $(DynamicQuantities.DimensionError)) ||
751-
isa(e, $(Unitful.DimensionError))
752-
error("Unable to convert units for \'" * string(:($$vv)) * "\'")
753-
elseif isa(e, MethodError)
754-
error("No or invalid units provided for \'" * string(:($$vv)) *
755-
"\'")
756-
else
757-
rethrow(e)
783+
if !(vv isa Tuple)
784+
name = getname(vv)
785+
varexpr = if haskey(metadata_with_exprs, VariableUnit)
786+
unit = metadata_with_exprs[VariableUnit]
787+
quote
788+
$name = if $name === $NO_VALUE
789+
$setdefault($vv, $def)
790+
else
791+
try
792+
$setdefault($vv, $convert_units($unit, $name))
793+
catch e
794+
if isa(e, $(DynamicQuantities.DimensionError)) ||
795+
isa(e, $(Unitful.DimensionError))
796+
error("Unable to convert units for \'" * string(:($$vv)) * "\'")
797+
elseif isa(e, MethodError)
798+
error("No or invalid units provided for \'" * string(:($$vv)) *
799+
"\'")
800+
else
801+
rethrow(e)
802+
end
758803
end
759804
end
760805
end
761-
end
762-
else
763-
quote
764-
$name = if $name === $NO_VALUE
765-
$setdefault($vv, $def)
766-
else
767-
$setdefault($vv, $name)
806+
else
807+
quote
808+
$name = if $name === $NO_VALUE
809+
$setdefault($vv, $def)
810+
else
811+
$setdefault($vv, $name)
812+
end
768813
end
769814
end
770-
end
771815

772-
metadata_expr = Expr(:block)
773-
for (k, v) in metadata_with_exprs
774-
push!(metadata_expr.args,
775-
:($name = $wrap($set_scalar_metadata($unwrap($name), $k, $v))))
776-
end
816+
metadata_expr = Expr(:block)
817+
for (k, v) in metadata_with_exprs
818+
push!(metadata_expr.args,
819+
:($name = $wrap($set_scalar_metadata($unwrap($name), $k, $v))))
820+
end
777821

778-
push!(varexpr.args, metadata_expr)
779-
return vv isa Num ? name : :($name...), varexpr
822+
push!(varexpr.args, metadata_expr)
823+
return vv isa Num ? name : :($name...), varexpr
824+
else
825+
return vv
826+
end
780827
end
781828

782829
function handle_conditional_vars!(

0 commit comments

Comments
 (0)