Support symbolic arrays in @mtkmodel (#2282)

Author: ven-k

docs/src/basics/MTKModel_Connector.md

@@ -19,7 +19,7 @@ equations.
 `ModelingToolkit.Model`, which includes a constructor that returns the ODESystem, a
 `structure` dictionary with metadata, and flag `isconnector` which is set to `false`.
 
-## What can an MTK-Model definition have?
+### What can an MTK-Model definition have?
 
 `@mtkmodel` definition contains begin blocks of
 
@@ -37,8 +37,8 @@ using ModelingToolkit
 
 @mtkmodel ModelA begin
     @parameters begin
-        k1
-        k2
+        k
+        k_array[1:2]
     end
 end
 
@@ -59,17 +59,17 @@ end
     @variables begin
         v(t) = v_var
     end
-    @extend p1, p2 = model_b = ModelB(; p1)
+    @extend ModelB(; p1)
     @components begin
-        model_a = ModelA(; k1)
+        model_a = ModelA(; k_array)
     end
     @equations begin
-        model_a.k1 ~ f(v)
+        model_a.k ~ f(v)
     end
 end
 ```
 
-### `@parameters` and `@variables` begin block
+#### `@parameters` and `@variables` begin block
 
   - Parameters and variables are declared with respective begin blocks.
   - Variables must be functions of an independent variable.
@@ -78,49 +78,54 @@ end
   - Whenever a parameter or variable has initial value, for example `v(t) = 0.0`, a symbolic variable named `v` with initial value 0.0 and a keyword argument `v`, with default value `nothing` are created. <br> This way, users can optionally pass new value of `v` while creating a component.
 
 ```julia
-julia > @named model_c = ModelC(; v = 2.0);
+julia> @mtkbuild model_c1 = ModelC(; v = 2.0);
 
-julia > ModelingToolkit.getdefault(model_c.v)
+julia> ModelingToolkit.getdefault(model_c.v)
 2.0
 ```
 
-### `@structural_parameters` begin block
+#### `@structural_parameters` begin block
 
   - This block is for non symbolic input arguements. These are for inputs that usually are not meant to be part of components; but influence how they are defined. One can list inputs like boolean flags, functions etc... here.
   - Whenever default values are specified, unlike parameters/variables, they are reflected in the keyword argument list.
 
 #### `@extend` begin block
 
-To extend a partial system,
-
-  - List the variables to unpack. If there is a single variable, explicitly specify it as a tuple.
-  - Give a name to the base system
-  - List the kwargs of the base system that should be listed as kwargs of the main component.
-  - Note that in above example, `p1` is promoted as an argument of `ModelC`. Users can set the value of `p1` as
+  - Partial systems can be extended in a higher system as `@extend PartialSystem(; kwargs)`.
+  - Keyword arguments pf partial system in the `@extend` definition are added as the keyword arguments of the base system.
+  - Note that in above example, `p1` is promoted as an argument of `ModelC`. Users can set the value of `p1`. However, as `p2` isn't listed in the model definition, its initial guess can't be specified while creating an instance of `ModelC`.
 
 ```julia
-julia> @named model_c = ModelC(; p1 = 2.0)
+julia> @mtkbuild model_c2 = ModelC(; p1 = 2.0)
 
 ```
 
-However, as `p2` isn't listed in the model definition, its initial guess can't
-specified while creating an instance of `ModelC`.
-
-### `@components` begin block
+#### `@components` begin block
 
   - Declare the subcomponents within `@components` begin block.
   - The arguments in these subcomponents are promoted as keyword arguments as `subcomponent_name__argname` with `nothing` as default value.
   - Whenever components are created with `@named` macro, these can be accessed with `.` operator as `subcomponent_name.argname`
-  - In the above example, `k1` of `model_a` can be set in following ways:
+  - In the above example, as `k` of `model_a` isn't listed while defining the sub-component in `ModelC`, its default value can't be modified by users. While `k_array` can be set as:
 
-```julia
-julia> @named model_c1 = ModelC(; model_a.k1 = 1);
+```@example mtkmodel-example
+using ModelingToolkit: getdefault
 
-```
+@mtkbuild model_c3 = ModelC(; model_a.k_array = [1.0, 2.0])
+
+getdefault(model_c3.model_a.k_array[1])
+# 1.0
+getdefault(model_c3.model_a.k_array[2])
+# 2.0
 
-And as `k2` isn't listed in the sub-component definition of `ModelC`, its default value can't be modified by users.
+@mtkbuild model_c4 = ModelC(model_a.k_array = 3.0)
+
+getdefault(model_c4.model_a.k_array[1])
+# 3.0
+getdefault(model_c4.model_a.k_array[2])
+# 3.0
+```
 
-### `@equations` begin block
+#### `@equations` begin block
 
   - List all the equations here
 
@@ -192,8 +197,10 @@ end
   - `:components`: List of sub-components in the form of [[name, sub_component_name],...].
   - `:extend`: The list of extended states, name given to the base system, and name of the base system.
   - `:structural_parameters`: Dictionary of structural parameters mapped to their default values.
-  - `:parameters`: Dictionary of symbolic parameters mapped to their metadata.
-  - `:variables`: Dictionary of symbolic variables mapped to their metadata.
+  - `:parameters`: Dictionary of symbolic parameters mapped to their metadata. For
+    parameter arrays, length is added to the metadata as `:size`.
+  - `:variables`: Dictionary of symbolic variables mapped to their metadata. For
+    variable arrays, length is added to the metadata as `:size`.
   - `:kwargs`: Dictionary of keyword arguments mapped to their default values.
   - `:independent_variable`: Independent variable, which is added while generating the Model.
   - `:equations`: List of equations (represented as strings).
@@ -204,8 +211,8 @@ For example, the structure of `ModelC` is:
 julia> ModelC.structure
 Dict{Symbol, Any} with 6 entries:
   :components           => [[:model_a, :ModelA]]
-  :variables            => Dict{Symbol, Dict{Symbol, Any}}(:v=>Dict(:default=>:v_var))
-  :kwargs               => Dict{Symbol, Any}(:f=>:sin, :v=>:v_var, :p1=>nothing, :model_a__k1=>nothing)
+  :variables            => Dict{Symbol, Dict{Symbol, Any}}(:v=>Dict(:default=>:v_var), :v_array=>Dict(:size=>(4,)))
+  :kwargs               => Dict{Symbol, Any}(:f=>:sin, :v=>:v_var, :v_array=>nothing, :model_a__k_array=>nothing, :p1=>nothing)
   :independent_variable => t
   :extend               => Any[[:p1, :p2], :model_b, :ModelB]
   :equations            => ["model_a.k1 ~ f(v)"]

docs/src/systems/JumpSystem.md

@@ -26,6 +26,7 @@ structural_simplify
 ```@docs; canonical=false
 DiscreteProblem(sys::ModelingToolkit.DiscreteSystem, u0map, tspan)
 ```
+
 ```@docs
 JumpProblem(sys::JumpSystem, prob, aggregator)
 ```

docs/src/systems/SDESystem.md

@@ -26,6 +26,7 @@ sde = SDESystem(ode, noiseeqs)
 structural_simplify
 alias_elimination
 ```
+
 ```@docs
 ModelingToolkit.Girsanov_transform
 ```

docs/src/tutorials/bifurcation_diagram_computation.md

@@ -1,65 +1,91 @@
 # [Bifurcation Diagrams](@id bifurcation_diagrams)
+
 Bifurcation diagrams describes how, for a dynamic system, the quantity and quality of its steady states changes with a parameter's value. These can be computed through the [BifurcationKit.jl](https://github.com/bifurcationkit/BifurcationKit.jl) package. ModelingToolkit provides a simple interface for creating BifurcationKit compatible `BifurcationProblem`s from `NonlinearSystem`s and `ODESystem`s. All teh features provided by BifurcationKit can then be applied to these systems. This tutorial provides a brief introduction for these features, with BifurcationKit.jl providing [a more extensive documentation](https://bifurcationkit.github.io/BifurcationKitDocs.jl/stable/).
 
 ### Creating a `BifurcationProblem`
+
 Let us first consider a simple `NonlinearSystem`:
+
 ```@example Bif1
 using ModelingToolkit
 @variables t x(t) y(t)
 @parameters μ α
-eqs = [0 ~ μ*x - x^3 + α*y,
-        0 ~ -y]
+eqs = [0 ~ μ * x - x^3 + α * y,
+    0 ~ -y]
 @named nsys = NonlinearSystem(eqs, [x, y], [μ, α])
 ```
+
 we wish to compute a bifurcation diagram for this system as we vary the parameter `μ`. For this, we need to provide the following information:
-1. The system for which we wish to compute the bifurcation diagram (`nsys`).
-2. The parameter which we wish to vary (`μ`).
-3. The parameter set for which we want to compute the bifurcation diagram.
-4. An initial guess of the state of the system for which there is a steady state at our provided parameter value.
-5. The variable which value we wish to plot in the bifurcation diagram (this argument is optional, if not provided, BifurcationKit default plot functions are used).
+
+ 1. The system for which we wish to compute the bifurcation diagram (`nsys`).
+ 2. The parameter which we wish to vary (`μ`).
+ 3. The parameter set for which we want to compute the bifurcation diagram.
+ 4. An initial guess of the state of the system for which there is a steady state at our provided parameter value.
+ 5. The variable which value we wish to plot in the bifurcation diagram (this argument is optional, if not provided, BifurcationKit default plot functions are used).
 
 We declare this additional information:
+
 ```@example Bif1
 bif_par = μ
 p_start = [μ => -1.0, α => 1.0]
 u0_guess = [x => 1.0, y => 1.0]
 plot_var = x;
 ```
+
 For the initial state guess (`u0_guess`), typically any value can be provided, however, read BifurcatioKit's documentation for more details.
 
 We can now create our `BifurcationProblem`, which can be provided as input to BifurcationKit's various functions.
+
 ```@example Bif1
 using BifurcationKit
-bprob = BifurcationProblem(nsys, u0_guess, p_start, bif_par; plot_var=plot_var, jac=false)
+bprob = BifurcationProblem(nsys,
+    u0_guess,
+    p_start,
+    bif_par;
+    plot_var = plot_var,
+    jac = false)
 ```
-Here, the `jac` argument (by default set to `true`) sets whenever to provide BifurcationKit with a Jacobian or not.
 
+Here, the `jac` argument (by default set to `true`) sets whenever to provide BifurcationKit with a Jacobian or not.
 
 ### Computing a bifurcation diagram
 
-Let us consider the `BifurcationProblem` from the last section. If we wish to compute the corresponding bifurcation diagram we must first declare various settings used by BifurcationKit to compute the diagram. These are stored in a `ContinuationPar` structure (which also contain a `NewtonPar` structure). 
+Let us consider the `BifurcationProblem` from the last section. If we wish to compute the corresponding bifurcation diagram we must first declare various settings used by BifurcationKit to compute the diagram. These are stored in a `ContinuationPar` structure (which also contain a `NewtonPar` structure).
+
 ```@example Bif1
 p_span = (-4.0, 6.0)
 opt_newton = NewtonPar(tol = 1e-9, max_iterations = 20)
 opts_br = ContinuationPar(dsmin = 0.001, dsmax = 0.05, ds = 0.01,
-	max_steps = 100, nev = 2, newton_options = opt_newton,
-	p_min = p_span[1], p_max = p_span[2],
-	detect_bifurcation = 3, n_inversion = 4, tol_bisection_eigenvalue = 1e-8, dsmin_bisection = 1e-9);
+    max_steps = 100, nev = 2, newton_options = opt_newton,
+    p_min = p_span[1], p_max = p_span[2],
+    detect_bifurcation = 3, n_inversion = 4, tol_bisection_eigenvalue = 1e-8,
+    dsmin_bisection = 1e-9);
 ```
+
 Here, `p_span` sets the interval over which we wish to compute the diagram.
 
 Next, we can use this as input to our bifurcation diagram, and then plot it.
+
 ```@example Bif1
-bf = bifurcationdiagram(bprob, PALC(), 2, (args...) -> opts_br; bothside=true)
+bf = bifurcationdiagram(bprob, PALC(), 2, (args...) -> opts_br; bothside = true)
 ```
+
 Here, the value `2` sets how sub-branches of the diagram that BifurcationKit should compute. Generally, for bifurcation diagrams, it is recommended to use the `bothside=true` argument.
+
 ```@example Bif1
 using Plots
-plot(bf; putspecialptlegend=false, markersize=2, plotfold=false, xguide="μ", yguide = "x")
+plot(bf;
+    putspecialptlegend = false,
+    markersize = 2,
+    plotfold = false,
+    xguide = "μ",
+    yguide = "x")
 ```
+
 Here, the system exhibits a pitchfork bifurcation at *μ=0.0*.
 
 ### Using `ODESystem` inputs
+
 It is also possible to use `ODESystem`s (rather than `NonlinearSystem`s) as input to `BifurcationProblem`. Here follows a brief such example.
 
 ```@example Bif2
@@ -68,25 +94,38 @@ using BifurcationKit, ModelingToolkit, Plots
 @variables t x(t) y(t)
 @parameters μ
 D = Differential(t)
-eqs = [D(x) ~ μ*x - y - x*(x^2+y^2), 
-        D(y) ~ x + μ*y - y*(x^2+y^2)]
+eqs = [D(x) ~ μ * x - y - x * (x^2 + y^2),
+    D(y) ~ x + μ * y - y * (x^2 + y^2)]
 @named osys = ODESystem(eqs, t)
 
 bif_par = μ
 plot_var = x
 p_start = [μ => 1.0]
-u0_guess = [x => 0.0, y=> 0.0]
+u0_guess = [x => 0.0, y => 0.0]
 
-bprob = BifurcationProblem(osys, u0_guess, p_start, bif_par; plot_var=plot_var, jac=false)
+bprob = BifurcationProblem(osys,
+    u0_guess,
+    p_start,
+    bif_par;
+    plot_var = plot_var,
+    jac = false)
 
 p_span = (-3.0, 3.0)
 opt_newton = NewtonPar(tol = 1e-9, max_iterations = 20)
 opts_br = ContinuationPar(dsmin = 0.001, dsmax = 0.05, ds = 0.01,
-	max_steps = 100, nev = 2, newton_options = opt_newton,
-	p_max = p_span[2], p_min = p_span[1],
-	detect_bifurcation = 3, n_inversion = 4, tol_bisection_eigenvalue = 1e-8, dsmin_bisection = 1e-9)
+    max_steps = 100, nev = 2, newton_options = opt_newton,
+    p_max = p_span[2], p_min = p_span[1],
+    detect_bifurcation = 3, n_inversion = 4, tol_bisection_eigenvalue = 1e-8,
+    dsmin_bisection = 1e-9)
 
-bf = bifurcationdiagram(bprob, PALC(), 2, (args...) -> opts_br; bothside=true)
-plot(bf; putspecialptlegend=false, markersize=2, plotfold=false, xguide="μ", yguide = "x")
+bf = bifurcationdiagram(bprob, PALC(), 2, (args...) -> opts_br; bothside = true)
+using Plots
+plot(bf;
+    putspecialptlegend = false,
+    markersize = 2,
+    plotfold = false,
+    xguide = "μ",
+    yguide = "x")
 ```
+
 Here, the value of `x` in the steady state does not change, however, at `μ=0` a Hopf bifurcation occur and the steady state turn unstable.

ext/MTKBifurcationKitExt.jl

@@ -9,15 +9,23 @@ import BifurcationKit
 ### Creates BifurcationProblem Overloads ###
 
 # When input is a NonlinearSystem.
-function BifurcationKit.BifurcationProblem(nsys::NonlinearSystem, u0_bif, ps, bif_par, args...; plot_var=nothing, record_from_solution=BifurcationKit.record_sol_default, jac=true, kwargs...)
+function BifurcationKit.BifurcationProblem(nsys::NonlinearSystem,
+    u0_bif,
+    ps,
+    bif_par,
+    args...;
+    plot_var = nothing,
+    record_from_solution = BifurcationKit.record_sol_default,
+    jac = true,
+    kwargs...)
     # Creates F and J functions.
-    ofun = NonlinearFunction(nsys; jac=jac)
+    ofun = NonlinearFunction(nsys; jac = jac)
     F = ofun.f
     J = jac ? ofun.jac : nothing
 
     # Computes bifurcation parameter and plot var indexes.
     bif_idx = findfirst(isequal(bif_par), parameters(nsys))
-    if !isnothing(plot_var) 
+    if !isnothing(plot_var)
         plot_idx = findfirst(isequal(plot_var), states(nsys))
         record_from_solution = (x, p) -> x[plot_idx]
     end
@@ -26,12 +34,22 @@ function BifurcationKit.BifurcationProblem(nsys::NonlinearSystem, u0_bif, ps, bi
     u0_bif = ModelingToolkit.varmap_to_vars(u0_bif, states(nsys))
     ps = ModelingToolkit.varmap_to_vars(ps, parameters(nsys))
 
-    return BifurcationKit.BifurcationProblem(F, u0_bif, ps, (@lens _[bif_idx]), args...; record_from_solution = record_from_solution, J = J, kwargs...)
+    return BifurcationKit.BifurcationProblem(F,
+        u0_bif,
+        ps,
+        (@lens _[bif_idx]),
+        args...;
+        record_from_solution = record_from_solution,
+        J = J,
+        kwargs...)
 end
 
 # When input is a ODESystem.
 function BifurcationKit.BifurcationProblem(osys::ODESystem, args...; kwargs...)
-    nsys = NonlinearSystem([0 ~ eq.rhs for eq in equations(osys)], states(osys), parameters(osys); name=osys.name)
+    nsys = NonlinearSystem([0 ~ eq.rhs for eq in equations(osys)],
+        states(osys),
+        parameters(osys);
+        name = osys.name)
     return BifurcationKit.BifurcationProblem(nsys, args...; kwargs...)
 end
 

src/ModelingToolkit.jl

@@ -56,7 +56,7 @@ using PrecompileTools, Reexport
     using Symbolics: _parse_vars, value, @derivatives, get_variables,
         exprs_occur_in, solve_for, build_expr, unwrap, wrap,
         VariableSource, getname, variable, Connection, connect,
-        NAMESPACE_SEPARATOR
+        NAMESPACE_SEPARATOR, set_scalar_metadata, setdefaultval
     import Symbolics: rename, get_variables!, _solve, hessian_sparsity,
         jacobian_sparsity, isaffine, islinear, _iszero, _isone,
         tosymbol, lower_varname, diff2term, var_from_nested_derivative,