Skip to content

Commit 5450d39

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 3c47904 commit 5450d39

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
@@ -239,20 +239,6 @@ end
239239
function parse_variable_def!(dict, mod, arg, varclass, kwargs, where_types;
240240
def = nothing, indices::Union{Vector{UnitRange{Int}}, Nothing} = nothing,
241241
type::Type = Real, meta = Dict{DataType, Expr}())
242-
metatypes = [(:connection_type, VariableConnectType),
243-
(:description, VariableDescription),
244-
(:unit, VariableUnit),
245-
(:bounds, VariableBounds),
246-
(:noise, VariableNoiseType),
247-
(:input, VariableInput),
248-
(:output, VariableOutput),
249-
(:irreducible, VariableIrreducible),
250-
(:state_priority, VariableStatePriority),
251-
(:misc, VariableMisc),
252-
(:disturbance, VariableDisturbance),
253-
(:tunable, VariableTunable),
254-
(:dist, VariableDistribution)]
255-
256242
arg isa LineNumberNode && return
257243
MLStyle.@match arg begin
258244
a::Symbol => begin
@@ -278,27 +264,86 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs, where_types;
278264
varclass, where_types, meta)
279265
return var, def, Dict()
280266
end
267+
Expr(:tuple, Expr(:(::), Expr(:ref, a, indices...), type), meta_val) ||
268+
Expr(:tuple, Expr(:ref, a, indices...), meta_val) => begin
269+
(@isdefined type) || (type = Real)
270+
varname = Meta.isexpr(a, :call) ? a.args[1] : a
271+
push!(kwargs, Expr(:kw, varname, nothing))
272+
meta = parse_metadata(mod, meta_val)
273+
varval = (@isdefined default_val) ? default_val :
274+
unit_handled_variable_value(meta, varname)
275+
if varclass == :parameters
276+
var = :($varname = $first(@parameters ($a[$(indices...)]::$type = $varval),
277+
$meta_val))
278+
else
279+
var = :($varname = $first(@variables ($a[$(indices)]::$type = $varval),
280+
$meta_val))
281+
end
282+
push_array_kwargs_and_metadata!(
283+
dict, indices, meta, type, varclass, varname, varval)
284+
(:($varname...), var), nothing, Dict()
285+
end
286+
Expr(:(=), Expr(:(::), Expr(:ref, a, indices...), type), def_n_meta) ||
287+
Expr(:(=), Expr(:ref, a, indices...), def_n_meta) => begin
288+
(@isdefined type) || (type = Real)
289+
varname = Meta.isexpr(a, :call) ? a.args[1] : a
290+
if Meta.isexpr(def_n_meta, :tuple)
291+
meta = parse_metadata(mod, def_n_meta)
292+
varval = unit_handled_variable_value(meta, varname)
293+
val, def_n_meta = (def_n_meta.args[1], def_n_meta.args[2:end])
294+
push!(kwargs, Expr(:kw, varname, nothing))
295+
if varclass == :parameters
296+
var = :($varname = $varname === nothing ? $val : $varname;
297+
$varname = $first(@parameters ($a[$(indices...)]::$type = $varval),
298+
$(def_n_meta...)))
299+
else
300+
var = :($varname = $varname === nothing ? $val : $varname;
301+
$varname = $first(@variables $a[$(indices...)]::$type = (
302+
$varval),
303+
$(def_n_meta...)))
304+
end
305+
else
306+
push!(kwargs, Expr(:kw, varname, nothing))
307+
if varclass == :parameters
308+
var = :($varname = $varname === nothing ? $def_n_meta : $varname;
309+
$varname = $first(@parameters $a[$(indices...)]::$type = $varname))
310+
else
311+
var = :($varname = $varname === nothing ? $def_n_meta : $varname;
312+
$varname = $first(@variables $a[$(indices...)]::$type = $varname))
313+
end
314+
varval, meta = def_n_meta, nothing
315+
end
316+
push_array_kwargs_and_metadata!(
317+
dict, indices, meta, type, varclass, varname, varval)
318+
(:($varname...), var), nothing, Dict()
319+
end
320+
Expr(:(::), Expr(:ref, a, indices...), type) ||
321+
Expr(:ref, a, indices...) => begin
322+
(@isdefined type) || (type = Real)
323+
varname = a isa Expr && a.head == :call ? a.args[1] : a
324+
push!(kwargs, Expr(:kw, varname, nothing))
325+
if varclass == :parameters
326+
var = :($varname = $first(@parameters $a[$(indices...)]::$type = $varname))
327+
elseif varclass == :variables
328+
var = :($varname = $first(@variables $a[$(indices...)]::$type = $varname))
329+
else
330+
throw("Symbolic array with arbitrary length is not handled for $varclass.
331+
Please open an issue with an example.")
332+
end
333+
push_array_kwargs_and_metadata!(
334+
dict, indices, nothing, type, varclass, varname, nothing)
335+
(:($varname...), var), nothing, Dict()
336+
end
281337
Expr(:(=), a, b) => begin
282338
Base.remove_linenums!(b)
283339
def, meta = parse_default(mod, b)
284340
var, def, _ = parse_variable_def!(
285341
dict, mod, a, varclass, kwargs, where_types; def, type, meta)
286-
if dict[varclass] isa Vector
287-
dict[varclass][1][getname(var)][:default] = def
288-
else
289-
dict[varclass][getname(var)][:default] = def
290-
end
342+
varclass_dict = dict[varclass] isa Vector ? Ref(dict[varclass][1]) :
343+
Ref(dict[varclass])
344+
varclass_dict[][getname(var)][:default] = def
291345
if meta !== nothing
292-
for (type, key) in metatypes
293-
if (mt = get(meta, key, nothing)) !== nothing
294-
key == VariableConnectType && (mt = nameof(mt))
295-
if dict[varclass] isa Vector
296-
dict[varclass][1][getname(var)][type] = mt
297-
else
298-
dict[varclass][getname(var)][type] = mt
299-
end
300-
end
301-
end
346+
update_readable_metadata!(varclass_dict[], meta, getname(var))
302347
var, metadata_with_exprs = set_var_metadata(var, meta)
303348
return var, def, metadata_with_exprs
304349
end
@@ -308,27 +353,15 @@ function parse_variable_def!(dict, mod, arg, varclass, kwargs, where_types;
308353
meta = parse_metadata(mod, b)
309354
var, def, _ = parse_variable_def!(
310355
dict, mod, a, varclass, kwargs, where_types; type, meta)
356+
varclass_dict = dict[varclass] isa Vector ? Ref(dict[varclass][1]) :
357+
Ref(dict[varclass])
311358
if meta !== nothing
312-
for (type, key) in metatypes
313-
if (mt = get(meta, key, nothing)) !== nothing
314-
key == VariableConnectType && (mt = nameof(mt))
315-
if dict[varclass] isa Vector
316-
dict[varclass][1][getname(var)][type] = mt
317-
else
318-
dict[varclass][getname(var)][type] = mt
319-
end
320-
end
321-
end
359+
update_readable_metadata!(varclass_dict[], meta, getname(var))
322360
var, metadata_with_exprs = set_var_metadata(var, meta)
323361
return var, def, metadata_with_exprs
324362
end
325363
return var, def, Dict()
326364
end
327-
Expr(:ref, a, b...) => begin
328-
indices = map(i -> UnitRange(i.args[2], i.args[end]), b)
329-
parse_variable_def!(dict, mod, a, varclass, kwargs, where_types;
330-
def, indices, type, meta)
331-
end
332365
_ => error("$arg cannot be parsed")
333366
end
334367
end
@@ -436,14 +469,23 @@ function parse_default(mod, a)
436469
end
437470
end
438471

439-
function parse_metadata(mod, a)
472+
function parse_metadata(mod, a::Expr)
440473
MLStyle.@match a begin
441-
Expr(:vect, eles...) => Dict(parse_metadata(mod, e) for e in eles)
474+
Expr(:vect, b...) => Dict(parse_metadata(mod, m) for m in b)
475+
Expr(:tuple, a, b...) => parse_metadata(mod, b)
442476
Expr(:(=), a, b) => Symbolics.option_to_metadata_type(Val(a)) => get_var(mod, b)
443477
_ => error("Cannot parse metadata $a")
444478
end
445479
end
446480

481+
function parse_metadata(mod, metadata::AbstractArray)
482+
ret = Dict()
483+
for m in metadata
484+
merge!(ret, parse_metadata(mod, m))
485+
end
486+
ret
487+
end
488+
447489
function _set_var_metadata!(metadata_with_exprs, a, m, v::Expr)
448490
push!(metadata_with_exprs, m => v)
449491
a
@@ -701,6 +743,7 @@ function parse_variable_arg!(exprs, vs, dict, mod, arg, varclass, kwargs, where_
701743
end
702744

703745
function convert_units(varunits::DynamicQuantities.Quantity, value)
746+
value isa Nothing && return nothing
704747
DynamicQuantities.ustrip(DynamicQuantities.uconvert(
705748
DynamicQuantities.SymbolicUnits.as_quantity(varunits), value))
706749
end
@@ -712,6 +755,7 @@ function convert_units(
712755
end
713756

714757
function convert_units(varunits::Unitful.FreeUnits, value)
758+
value isa Nothing && return nothing
715759
Unitful.ustrip(varunits, value)
716760
end
717761

@@ -730,47 +774,50 @@ end
730774
function parse_variable_arg(dict, mod, arg, varclass, kwargs, where_types)
731775
vv, def, metadata_with_exprs = parse_variable_def!(
732776
dict, mod, arg, varclass, kwargs, where_types)
733-
name = getname(vv)
734-
735-
varexpr = if haskey(metadata_with_exprs, VariableUnit)
736-
unit = metadata_with_exprs[VariableUnit]
737-
quote
738-
$name = if $name === nothing
739-
$setdefault($vv, $def)
740-
else
741-
try
742-
$setdefault($vv, $convert_units($unit, $name))
743-
catch e
744-
if isa(e, $(DynamicQuantities.DimensionError)) ||
745-
isa(e, $(Unitful.DimensionError))
746-
error("Unable to convert units for \'" * string(:($$vv)) * "\'")
747-
elseif isa(e, MethodError)
748-
error("No or invalid units provided for \'" * string(:($$vv)) *
749-
"\'")
750-
else
751-
rethrow(e)
777+
if !(vv isa Tuple)
778+
name = getname(vv)
779+
varexpr = if haskey(metadata_with_exprs, VariableUnit)
780+
unit = metadata_with_exprs[VariableUnit]
781+
quote
782+
$name = if $name === nothing
783+
$setdefault($vv, $def)
784+
else
785+
try
786+
$setdefault($vv, $convert_units($unit, $name))
787+
catch e
788+
if isa(e, $(DynamicQuantities.DimensionError)) ||
789+
isa(e, $(Unitful.DimensionError))
790+
error("Unable to convert units for \'" * string(:($$vv)) * "\'")
791+
elseif isa(e, MethodError)
792+
error("No or invalid units provided for \'" * string(:($$vv)) *
793+
"\'")
794+
else
795+
rethrow(e)
796+
end
752797
end
753798
end
754799
end
755-
end
756-
else
757-
quote
758-
$name = if $name === nothing
759-
$setdefault($vv, $def)
760-
else
761-
$setdefault($vv, $name)
800+
else
801+
quote
802+
$name = if $name === nothing
803+
$setdefault($vv, $def)
804+
else
805+
$setdefault($vv, $name)
806+
end
762807
end
763808
end
764-
end
765809

766-
metadata_expr = Expr(:block)
767-
for (k, v) in metadata_with_exprs
768-
push!(metadata_expr.args,
769-
:($name = $wrap($set_scalar_metadata($unwrap($name), $k, $v))))
770-
end
810+
metadata_expr = Expr(:block)
811+
for (k, v) in metadata_with_exprs
812+
push!(metadata_expr.args,
813+
:($name = $wrap($set_scalar_metadata($unwrap($name), $k, $v))))
814+
end
771815

772-
push!(varexpr.args, metadata_expr)
773-
return vv isa Num ? name : :($name...), varexpr
816+
push!(varexpr.args, metadata_expr)
817+
return vv isa Num ? name : :($name...), varexpr
818+
else
819+
return vv
820+
end
774821
end
775822

776823
function handle_conditional_vars!(

0 commit comments

Comments
 (0)