Merge remote-tracking branch 'origin/master' into add_odes_to_jumpsys

Author: isaacsas

Project.toml

@@ -21,6 +21,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 DomainSets = "5b8099bc-c8ec-5219-889f-1d9e522a28bf"
 DynamicQuantities = "06fc5a27-2a28-4c7c-a15d-362465fb6821"
+EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56"
 ExprTools = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
 Expronicon = "6b7a57c9-7cc1-4fdf-b7f5-e857abae3636"
 FindFirstFunctions = "64ca27bc-2ba2-4a57-88aa-44e436879224"
@@ -94,6 +95,7 @@ Distributions = "0.23, 0.24, 0.25"
 DocStringExtensions = "0.7, 0.8, 0.9"
 DomainSets = "0.6, 0.7"
 DynamicQuantities = "^0.11.2, 0.12, 0.13, 1"
+EnumX = "1.0.4"
 ExprTools = "0.1.10"
 Expronicon = "0.8"
 FindFirstFunctions = "1"
@@ -111,7 +113,7 @@ Libdl = "1"
 LinearAlgebra = "1"
 MLStyle = "0.4.17"
 NaNMath = "0.3, 1"
-NonlinearSolve = "3.14"
+NonlinearSolve = "3.14, 4"
 OffsetArrays = "1"
 OrderedCollections = "1"
 OrdinaryDiffEq = "6.82.0"
@@ -125,7 +127,7 @@ SciMLBase = "2.57.1"
 SciMLStructures = "1.0"
 Serialization = "1"
 Setfield = "0.7, 0.8, 1"
-SimpleNonlinearSolve = "0.1.0, 1"
+SimpleNonlinearSolve = "0.1.0, 1, 2"
 SparseArrays = "1"
 SpecialFunctions = "0.7, 0.8, 0.9, 0.10, 1.0, 2"
 StaticArrays = "0.10, 0.11, 0.12, 1.0"

docs/Project.toml

@@ -30,7 +30,7 @@ Distributions = "0.25"
 Documenter = "1"
 DynamicQuantities = "^0.11.2, 0.12, 1"
 ModelingToolkit = "8.33, 9"
-NonlinearSolve = "3"
+NonlinearSolve = "3, 4"
 Optim = "1.7"
 Optimization = "3.9"
 OptimizationOptimJL = "0.1"

docs/src/basics/MTKLanguage.md

@@ -147,9 +147,9 @@ julia> ModelingToolkit.getdefault(model_c1.v)
 2.0
 ```
 
-#### `@extend` begin block
+#### `@extend` statement
 
-Partial systems can be extended in a higher system in two ways:
+One or more partial systems can be extended in a higher system with `@extend` statements. This can be done in two ways:
 
   - `@extend PartialSystem(var1 = value1)`
 
@@ -313,7 +313,8 @@ end
   - `:components`: The list of sub-components in the form of [[name, sub_component_name],...].
   - `:constants`: Dictionary of constants mapped to its metadata.
   - `:defaults`: Dictionary of variables and default values specified in the `@defaults`.
-  - `:extend`: The list of extended unknowns, name given to the base system, and name of the base system.
+  - `:extend`: The list of extended unknowns, parameters and components, name given to the base system, and name of the base system.
+    When multiple extend statements are present, latter two are returned as lists.
   - `:structural_parameters`: Dictionary of structural parameters mapped to their metadata.
   - `:parameters`: Dictionary of symbolic parameters mapped to their metadata. For
     parameter arrays, length is added to the metadata as `:size`.

docs/src/tutorials/stochastic_diffeq.md

@@ -23,9 +23,9 @@ where the magnitude of the noise scales with (0.3 times) the magnitude of each o
 
 ```math
 \begin{aligned}
-\frac{dx}{dt} &= (\sigma (y-x))  &+ 0.1x\frac{dB}{dt} \\
-\frac{dy}{dt} &= (x(\rho-z) - y) &+ 0.1y\frac{dB}{dt}  \\
-\frac{dz}{dt} &= (xy - \beta z)  &+ 0.1z\frac{dB}{dt}  \\
+\frac{dx}{dt} &= (\sigma (y-x))  &+ 0.3x\frac{dB}{dt} \\
+\frac{dy}{dt} &= (x(\rho-z) - y) &+ 0.3y\frac{dB}{dt}  \\
+\frac{dz}{dt} &= (xy - \beta z)  &+ 0.3z\frac{dB}{dt}  \\
 \end{aligned}
 ```
 

ext/MTKHomotopyContinuationExt.jl

@@ -15,42 +15,104 @@ function contains_variable(x, wrt)
     any(y -> occursin(y, x), wrt)
 end
 
+"""
+Possible reasons why a term is not polynomial
+"""
+MTK.EnumX.@enumx NonPolynomialReason begin
+    NonIntegerExponent
+    ExponentContainsUnknowns
+    BaseNotPolynomial
+    UnrecognizedOperation
+end
+
+function display_reason(reason::NonPolynomialReason.T, sym)
+    if reason == NonPolynomialReason.NonIntegerExponent
+        pow = arguments(sym)[2]
+        "In $sym: Exponent $pow is not an integer"
+    elseif reason == NonPolynomialReason.ExponentContainsUnknowns
+        pow = arguments(sym)[2]
+        "In $sym: Exponent $pow contains unknowns of the system"
+    elseif reason == NonPolynomialReason.BaseNotPolynomial
+        base = arguments(sym)[1]
+        "In $sym: Base $base is not a polynomial in the unknowns"
+    elseif reason == NonPolynomialReason.UnrecognizedOperation
+        op = operation(sym)
+        """
+        In $sym: Operation $op is not recognized. Allowed polynomial operations are \
+        `*, /, +, -, ^`.
+        """
+    else
+        error("This should never happen. Please open an issue in ModelingToolkit.jl.")
+    end
+end
+
+mutable struct PolynomialData
+    non_polynomial_terms::Vector{BasicSymbolic}
+    reasons::Vector{NonPolynomialReason.T}
+    has_parametric_exponent::Bool
+end
+
+PolynomialData() = PolynomialData(BasicSymbolic[], NonPolynomialReason.T[], false)
+
+struct NotPolynomialError <: Exception
+    eq::Equation
+    data::PolynomialData
+end
+
+function Base.showerror(io::IO, err::NotPolynomialError)
+    println(io,
+        "Equation $(err.eq) is not a polynomial in the unknowns for the following reasons:")
+    for (term, reason) in zip(err.data.non_polynomial_terms, err.data.reasons)
+        println(io, display_reason(reason, term))
+    end
+end
+
+function is_polynomial!(data, y, wrt)
+    process_polynomial!(data, y, wrt)
+    isempty(data.reasons)
+end
+
 """
 $(TYPEDSIGNATURES)
 
-Check if `x` is polynomial with respect to the variables in `wrt`.
+Return information about the polynmial `x` with respect to variables in `wrt`,
+writing said information to `data`.
 """
-function is_polynomial(x, wrt)
+function process_polynomial!(data::PolynomialData, x, wrt)
     x = unwrap(x)
     symbolic_type(x) == NotSymbolic() && return true
     iscall(x) || return true
     contains_variable(x, wrt) || return true
     any(isequal(x), wrt) && return true
 
     if operation(x) in (*, +, -, /)
-        return all(y -> is_polynomial(y, wrt), arguments(x))
+        return all(y -> is_polynomial!(data, y, wrt), arguments(x))
     end
     if operation(x) == (^)
         b, p = arguments(x)
         is_pow_integer = symtype(p) <: Integer
         if !is_pow_integer
-            if symbolic_type(p) == NotSymbolic()
-                @warn "In $x: Exponent $p is not an integer"
-            else
-                @warn "In $x: Exponent $p is not an integer. Use `@parameters p::Integer` to declare integer parameters."
-            end
+            push!(data.non_polynomial_terms, x)
+            push!(data.reasons, NonPolynomialReason.NonIntegerExponent)
+        end
+        if symbolic_type(p) != NotSymbolic()
+            data.has_parametric_exponent = true
         end
+
         exponent_has_unknowns = contains_variable(p, wrt)
         if exponent_has_unknowns
-            @warn "In $x: Exponent $p cannot contain unknowns of the system."
+            push!(data.non_polynomial_terms, x)
+            push!(data.reasons, NonPolynomialReason.ExponentContainsUnknowns)
         end
-        base_polynomial = is_polynomial(b, wrt)
+        base_polynomial = is_polynomial!(data, b, wrt)
         if !base_polynomial
-            @warn "In $x: Base is not a polynomial"
+            push!(data.non_polynomial_terms, x)
+            push!(data.reasons, NonPolynomialReason.BaseNotPolynomial)
         end
         return base_polynomial && !exponent_has_unknowns && is_pow_integer
     end
-    @warn "In $x: Unrecognized operation $(operation(x)). Allowed polynomial operations are `*, +, -, ^`"
+    push!(data.non_polynomial_terms, x)
+    push!(data.reasons, NonPolynomialReason.UnrecognizedOperation)
     return false
 end
 
@@ -179,21 +241,39 @@ Create a `HomotopyContinuationProblem` from a `NonlinearSystem` with polynomial
 The problem will be solved by HomotopyContinuation.jl. The resultant `NonlinearSolution`
 will contain the polynomial root closest to the point specified by `u0map` (if real roots
 exist for the system).
+
+Keyword arguments:
+- `eval_expression`: Whether to `eval` the generated functions or use a `RuntimeGeneratedFunction`.
+- `eval_module`: The module to use for `eval`/`@RuntimeGeneratedFunction`
+- `warn_parametric_exponent`: Whether to warn if the system contains a parametric
+  exponent preventing the homotopy from being cached.
+
+All other keyword arguments are forwarded to `HomotopyContinuation.solver_startsystems`.
 """
 function MTK.HomotopyContinuationProblem(
         sys::NonlinearSystem, u0map, parammap = nothing; eval_expression = false,
-        eval_module = ModelingToolkit, kwargs...)
+        eval_module = ModelingToolkit, warn_parametric_exponent = true, kwargs...)
     if !iscomplete(sys)
         error("A completed `NonlinearSystem` is required. Call `complete` or `structural_simplify` on the system before creating a `HomotopyContinuationProblem`")
     end
 
     dvs = unknowns(sys)
-    eqs = equations(sys)
+    # we need to consider `full_equations` because observed also should be
+    # polynomials (if used in equations) and we don't know if observed is used
+    # in denominator.
+    # This is not the most efficient, and would be improved significantly with
+    # CSE/hashconsing.
+    eqs = full_equations(sys)
 
     denoms = []
+    has_parametric_exponents = false
     eqs2 = map(eqs) do eq
-        if !is_polynomial(eq.lhs, dvs) || !is_polynomial(eq.rhs, dvs)
-            error("Equation $eq is not a polynomial in the unknowns. See warnings for further details.")
+        data = PolynomialData()
+        process_polynomial!(data, eq.lhs, dvs)
+        process_polynomial!(data, eq.rhs, dvs)
+        has_parametric_exponents |= data.has_parametric_exponent
+        if !isempty(data.non_polynomial_terms)
+            throw(NotPolynomialError(eq, data))
         end
         num, den = handle_rational_polynomials(eq.rhs - eq.lhs, dvs)
 
@@ -212,6 +292,9 @@ function MTK.HomotopyContinuationProblem(
     end
 
     sys2 = MTK.@set sys.eqs = eqs2
+    # remove observed equations to avoid adding them in codegen
+    MTK.@set! sys2.observed = Equation[]
+    MTK.@set! sys2.substitutions = nothing
 
     nlfn, u0, p = MTK.process_SciMLProblem(NonlinearFunction{true}, sys2, u0map, parammap;
         jac = true, eval_expression, eval_module)
@@ -223,29 +306,49 @@ function MTK.HomotopyContinuationProblem(
 
     obsfn = MTK.ObservedFunctionCache(sys; eval_expression, eval_module)
 
-    return MTK.HomotopyContinuationProblem(u0, mtkhsys, denominator, sys, obsfn)
+    if has_parametric_exponents
+        if warn_parametric_exponent
+            @warn """
+            The system has parametric exponents, preventing caching of the homotopy. \
+            This will cause `solve` to be slower. Pass `warn_parametric_exponent \
+            = false` to turn off this warning
+            """
+        end
+        solver_and_starts = nothing
+    else
+        solver_and_starts = HomotopyContinuation.solver_startsolutions(mtkhsys; kwargs...)
+    end
+    return MTK.HomotopyContinuationProblem(
+        u0, mtkhsys, denominator, sys, obsfn, solver_and_starts)
 end
 
 """
 $(TYPEDSIGNATURES)
 
 Solve a `HomotopyContinuationProblem`. Ignores the algorithm passed to it, and always
-uses `HomotopyContinuation.jl`. All keyword arguments except the ones listed below are
-forwarded to `HomotopyContinuation.solve`. The original solution as returned by
+uses `HomotopyContinuation.jl`. The original solution as returned by
 `HomotopyContinuation.jl` will be available in the `.original` field of the returned
 `NonlinearSolution`.
 
-All keyword arguments have their default values in HomotopyContinuation.jl, except
-`show_progress` which defaults to `false`.
+All keyword arguments except the ones listed below are forwarded to
+`HomotopyContinuation.solve`. Note that the solver and start solutions are precomputed,
+and only keyword arguments related to the solve process are valid. All keyword
+arguments have their default values in HomotopyContinuation.jl, except `show_progress`
+which defaults to `false`.
 
 Extra keyword arguments:
 - `denominator_abstol`: In case `prob` is solving a rational function, roots which cause
   the denominator to be below `denominator_abstol` will be discarded.
 """
 function CommonSolve.solve(prob::MTK.HomotopyContinuationProblem,
         alg = nothing; show_progress = false, denominator_abstol = 1e-7, kwargs...)
-    sol = HomotopyContinuation.solve(
-        prob.homotopy_continuation_system; show_progress, kwargs...)
+    if prob.solver_and_starts === nothing
+        sol = HomotopyContinuation.solve(
+            prob.homotopy_continuation_system; show_progress, kwargs...)
+    else
+        solver, starts = prob.solver_and_starts
+        sol = HomotopyContinuation.solve(solver, starts; show_progress, kwargs...)
+    end
     realsols = HomotopyContinuation.results(sol; only_real = true)
     if isempty(realsols)
         u = state_values(prob)

src/ModelingToolkit.jl

@@ -55,6 +55,7 @@ using RecursiveArrayTools
 import Graphs: SimpleDiGraph, add_edge!, incidence_matrix
 import BlockArrays: BlockedArray, Block, blocksize, blocksizes
 import CommonSolve
+import EnumX
 
 using RuntimeGeneratedFunctions
 using RuntimeGeneratedFunctions: drop_expr

src/systems/abstractsystem.jl

@@ -918,7 +918,7 @@ Mark a system as completed. A completed system is a system which is done being
 defined/modified and is ready for structural analysis or other transformations.
 This allows for analyses and optimizations to be performed which require knowing
 the global structure of the system.
-        
+
 One property to note is that if a system is complete, the system will no longer
 namespace its subsystems or variables, i.e. `isequal(complete(sys).v.i, v.i)`.
 """
@@ -1933,7 +1933,7 @@ function Base.show(
         end
     end
     limited = nrows < nsubs
-    limited && print(io, "\n  ⋮") # too many to print 
+    limited && print(io, "\n  ⋮") # too many to print
 
     # Print equations
     eqs = equations(sys)
@@ -3043,10 +3043,19 @@ function extend(sys::AbstractSystem, basesys::AbstractSystem;
     return T(args...; kwargs...)
 end
 
+function extend(sys, basesys::Vector{T}) where {T <: AbstractSystem}
+    foldl(extend, basesys, init = sys)
+end
+
 function Base.:(&)(sys::AbstractSystem, basesys::AbstractSystem; kwargs...)
     extend(sys, basesys; kwargs...)
 end
 
+function Base.:(&)(
+        sys::AbstractSystem, basesys::Vector{T}; kwargs...) where {T <: AbstractSystem}
+    extend(sys, basesys; kwargs...)
+end
+
 """
 $(SIGNATURES)
 

src/systems/callbacks.jl

@@ -552,7 +552,7 @@ end
 """
     compile_condition(cb::SymbolicDiscreteCallback, sys, dvs, ps; expression, kwargs...)
 
-Returns a function `condition(u,p,t)` returning the `condition(cb)`.
+Returns a function `condition(u,t,integrator)` returning the `condition(cb)`.
 
 Notes
 
@@ -573,7 +573,8 @@ function compile_condition(cb::SymbolicDiscreteCallback, sys, dvs, ps;
     end
     expr = build_function(
         condit, u, t, p...; expression = Val{true},
-        wrap_code = condition_header(sys) .∘ wrap_array_vars(sys, condit; dvs, ps) .∘
+        wrap_code = condition_header(sys) .∘
+                    wrap_array_vars(sys, condit; dvs, ps, inputs = true) .∘
                     wrap_parameter_dependencies(sys, !(condit isa AbstractArray)),
         kwargs...)
     if expression == Val{true}