Fix after rebase

Author: vyudu

.github/workflows/Downstream.yml

@@ -39,7 +39,7 @@ jobs:
           - {user: SciML, repo: MethodOfLines.jl, group: DAE}
           - {user: SciML, repo: ModelingToolkitNeuralNets.jl, group: All}
 
-          - {user: Neuroblox, repo: Neuroblox.jl, group: NNPDE}
+          - {user: Neuroblox, repo: Neuroblox.jl, group: All}
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v1

docs/src/basics/InputOutput.md

@@ -70,7 +70,7 @@ Now we can test the generated function `f` with random input and state values
 p = [1]
 x = [rand()]
 u = [rand()]
-@test f[1](x, u, p, 1) ≈ -p[] * (x + u) # Test that the function computes what we expect D(x) = -k*(x + u)
+@test f(x, u, p, 1) ≈ -p[] * (x + u) # Test that the function computes what we expect D(x) = -k*(x + u)
 ```
 
 ## Generating an output function, ``g``

docs/src/tutorials/disturbance_modeling.md

@@ -184,7 +184,7 @@ disturbance_inputs = [ssys.d1, ssys.d2]
 P = ssys.system_model
 outputs = [P.inertia1.phi, P.inertia2.phi, P.inertia1.w, P.inertia2.w]
 
-(f_oop, f_ip), x_sym, p_sym, io_sys = ModelingToolkit.generate_control_function(
+f, x_sym, p_sym, io_sys = ModelingToolkit.generate_control_function(
     model_with_disturbance, inputs, disturbance_inputs; disturbance_argument = true)
 
 g = ModelingToolkit.build_explicit_observed_function(
@@ -195,12 +195,12 @@ x0, _ = ModelingToolkit.get_u0_p(io_sys, op, op)
 p = MTKParameters(io_sys, op)
 u = zeros(1) # Control input
 w = zeros(length(disturbance_inputs)) # Disturbance input
-@test f_oop(x0, u, p, t, w) == zeros(5)
+@test f(x0, u, p, t, w) == zeros(5)
 @test g(x0, u, p, 0.0) == [0, 0, 0, 0]
 
 # Non-zero disturbance inputs should result in non-zero state derivatives. We call `sort` since we do not generally know the order of the state variables
 w = [1.0, 2.0]
-@test sort(f_oop(x0, u, p, t, w)) == [0, 0, 0, 1, 2]
+@test sort(f(x0, u, p, t, w)) == [0, 0, 0, 1, 2]
 ```
 
 ## Input signal library

ext/MTKInfiniteOptExt.jl

@@ -340,7 +340,7 @@ function constructDefault(T::Type = Float64)
     A = map(T, A)
     α = map(T, α)
     c = map(T, c)
-    
+
     DiffEqBase.ImplicitRKTableau(A, c, α, 5)
 end
 
@@ -422,7 +422,6 @@ function _solve(prob::AbstractDynamicOptProblem, jump_solver, solver)
     DynamicOptSolution(model, sol, input_sol)
 end
 
-
 import InfiniteOpt: JuMP, GeneralVariableRef
 
 for ff in [acos, log1p, acosh, log2, asin, tan, atanh, cos, log, sin, log10, sqrt]

src/discretedomain.jl

@@ -39,9 +39,13 @@ SymbolicUtils.isbinop(::Shift) = false
 
 function (D::Shift)(x, allow_zero = false)
     !allow_zero && D.steps == 0 && return x
-    Term{symtype(x)}(D, Any[x])
+    if Symbolics.isarraysymbolic(x)
+        Symbolics.array_term(D, x)
+    else
+        term(D, x)
+    end
 end
-function (D::Shift)(x::Num, allow_zero = false)
+function (D::Shift)(x::Union{Num, Symbolics.Arr}, allow_zero = false)
     !allow_zero && D.steps == 0 && return x
     vt = value(x)
     if iscall(vt)
@@ -52,11 +56,11 @@ function (D::Shift)(x::Num, allow_zero = false)
             if D.t === nothing || isequal(D.t, op.t)
                 arg = arguments(vt)[1]
                 newsteps = D.steps + op.steps
-                return Num(newsteps == 0 ? arg : Shift(D.t, newsteps)(arg))
+                return wrap(newsteps == 0 ? arg : Shift(D.t, newsteps)(arg))
             end
         end
     end
-    Num(D(vt, allow_zero))
+    wrap(D(vt, allow_zero))
 end
 SymbolicUtils.promote_symtype(::Shift, t) = t
 
@@ -202,11 +206,19 @@ function (xn::Num)(k::ShiftIndex)
     x = value(xn)
     # Verify that the independent variables of k and x match and that the expression doesn't have multiple variables
     vars = Symbolics.get_variables(x)
-    length(vars) == 1 ||
+    if length(vars) != 1
         error("Cannot shift a multivariate expression $x. Either create a new unknown and shift this, or shift the individual variables in the expression.")
-    args = Symbolics.arguments(vars[]) # args should be one element vector with the t in x(t)
-    length(args) == 1 ||
+    end
+    var = only(vars)
+    if !iscall(var)
+        throw(ArgumentError("Cannot shift time-independent variable $var"))
+    end
+    if operation(var) == getindex
+        var = first(arguments(var))
+    end
+    if length(arguments(var)) != 1
         error("Cannot shift an expression with multiple independent variables $x.")
+    end
 
     # d, _ = propagate_time_domain(xn)
     # if d != clock # this is only required if the variable has another clock
@@ -220,6 +232,34 @@ function (xn::Num)(k::ShiftIndex)
     Shift(t, steps)(xn) # a shift of k steps
 end
 
+function (xn::Symbolics.Arr)(k::ShiftIndex)
+    @unpack clock, steps = k
+    x = value(xn)
+    # Verify that the independent variables of k and x match and that the expression doesn't have multiple variables
+    vars = ModelingToolkit.vars(x)
+    if length(vars) != 1
+        error("Cannot shift a multivariate expression $x. Either create a new unknown and shift this, or shift the individual variables in the expression.")
+    end
+    var = only(vars)
+    if !iscall(var)
+        throw(ArgumentError("Cannot shift time-independent variable $var"))
+    end
+    if length(arguments(var)) != 1
+        error("Cannot shift an expression with multiple independent variables $x.")
+    end
+
+    # d, _ = propagate_time_domain(xn)
+    # if d != clock # this is only required if the variable has another clock
+    #     xn = Sample(t, clock)(xn)
+    # end
+    # QUESTION: should we return a variable with time domain set to k.clock?
+    xn = wrap(setmetadata(unwrap(xn), VariableTimeDomain, k.clock))
+    if steps == 0
+        return xn # x(k) needs no shift operator if the step of k is 0
+    end
+    Shift(t, steps)(xn) # a shift of k steps
+end
+
 Base.:+(k::ShiftIndex, i::Int) = ShiftIndex(k.clock, k.steps + i)
 Base.:-(k::ShiftIndex, i::Int) = k + (-i)
 

src/inputoutput.jl

@@ -160,16 +160,17 @@ has_var(ex, x) = x ∈ Set(get_variables(ex))
 # Build control function
 
 """
-    (f_oop, f_ip), x_sym, p_sym, io_sys = generate_control_function(
+    f, x_sym, p_sym, io_sys = generate_control_function(
             sys::AbstractODESystem,
             inputs             = unbound_inputs(sys),
             disturbance_inputs = disturbances(sys);
             implicit_dae       = false,
             simplify           = false,
         )
 
-For a system `sys` with inputs (as determined by [`unbound_inputs`](@ref) or user specified), generate a function with additional input argument `in`
+For a system `sys` with inputs (as determined by [`unbound_inputs`](@ref) or user specified), generate functions with additional input argument `u`
 
+The returned function `f` can be called in the out-of-place or in-place form:
 ```
 f_oop : (x,u,p,t)      -> rhs
 f_ip  : (xout,x,u,p,t) -> nothing
@@ -187,7 +188,7 @@ f, x_sym, ps = generate_control_function(sys, expression=Val{false}, simplify=fa
 p = varmap_to_vars(defaults(sys), ps)
 x = varmap_to_vars(defaults(sys), x_sym)
 t = 0
-f[1](x, inputs, p, t)
+f(x, inputs, p, t)
 ```
 """
 function generate_control_function(sys::AbstractODESystem, inputs = unbound_inputs(sys),
@@ -249,9 +250,9 @@ function generate_control_function(sys::AbstractODESystem, inputs = unbound_inpu
     f = build_function_wrapper(sys, rhss, args...; p_start = 3 + implicit_dae,
         p_end = length(p) + 2 + implicit_dae, kwargs...)
     f = eval_or_rgf.(f; eval_expression, eval_module)
-    f = GeneratedFunctionWrapper{(3, length(args) - length(p) + 1, is_split(sys))}(f...)
+    f = GeneratedFunctionWrapper{(3 + implicit_dae, length(args) - length(p) + 1, is_split(sys))}(f...)
     ps = setdiff(parameters(sys), inputs, disturbance_inputs)
-    (; f, dvs, ps, io_sys = sys)
+    (; f = (f, f), dvs, ps, io_sys = sys)
 end
 
 """
@@ -418,7 +419,7 @@ function add_input_disturbance(sys, dist::DisturbanceModel, inputs = Any[]; kwar
     augmented_sys = extend(augmented_sys, sys)
     ssys = structural_simplify(augmented_sys, inputs = all_inputs, disturbance_inputs = [d])
 
-    f, dvs, p, io_sys = generate_control_function(ssys, all_inputs,
+    (f_oop, f_ip), dvs, p, io_sys = generate_control_function(ssys, all_inputs,
         [d]; kwargs...)
-    f, augmented_sys, dvs, p, io_sys
+    (f_oop, f_ip), augmented_sys, dvs, p, io_sys
 end

src/linearization.jl

@@ -143,6 +143,17 @@ function eq_idxs(sys::AbstractSystem)
     diff_idxs, alge_idxs
 end
 
+"""
+Return the set of indexes of differential equations and algebraic equations in the simplified system.
+"""
+function eq_idxs(sys::AbstractSystem)
+    eqs = equations(sys)
+    alge_idxs = findall(!isdiffeq, eqs)
+    diff_idxs = setdiff(1:length(eqs), alge_idxs)
+
+    diff_idxs, alge_idxs
+end
+
 """
     $(TYPEDEF)
 
@@ -607,7 +618,7 @@ function markio!(state, orig_inputs, inputs, outputs, disturbances; check = true
         end
         (all(values(outputset)) || error(
             "Some specified outputs were not found in system. The following Dict indicates the found variables ",
-        outputset))
+            outputset))
     end
     state, orig_inputs
 end

src/structural_transformation/utils.jl

@@ -472,6 +472,13 @@ function shift2term(var)
     op = operation(var)
     iv = op.t
     arg = only(arguments(var))
+    if operation(arg) === getindex
+        idxs = arguments(arg)[2:end]
+        newvar = shift2term(op(first(arguments(arg))))[idxs...]
+        unshifted = ModelingToolkit.getunshifted(newvar)[idxs...]
+        newvar = setmetadata(newvar, ModelingToolkit.VariableUnshifted, unshifted)
+        return newvar
+    end
     is_lowered = !isnothing(ModelingToolkit.getunshifted(arg))
 
     backshift = is_lowered ? op.steps + ModelingToolkit.getshift(arg) : op.steps

src/systems/discrete_system/discrete_system.jl

@@ -319,6 +319,7 @@ function SciMLBase.DiscreteProblem(
     iv = get_iv(sys)
 
     u0map = to_varmap(u0map, dvs)
+    scalarize_varmap!(u0map)
     u0map = shift_u0map_forward(sys, u0map, defaults(sys))
     f, u0, p = process_SciMLProblem(
         DiscreteFunction, sys, u0map, parammap; eval_expression, eval_module, build_initializeprob = false)

src/systems/if_lifting.jl

@@ -111,8 +111,8 @@ function (cw::CondRewriter)(expr, dep)
         # and ELSE branch is true
         # similarly for expression being false
         return (ifelse(rw_cond, rw_conda, rw_condb),
-            implies(ctrue, truea) | implies(cfalse, trueb),
-            implies(ctrue, falsea) | implies(cfalse, falseb))
+            ctrue & truea | cfalse & trueb,
+            ctrue & falsea | cfalse & falseb)
     elseif operation(expr) == Base.:(!) # NOT of expression
         (a,) = arguments(expr)
         (rw, ctrue, cfalse) = cw(a, dep)