feat: initial implementation of HomotopyContinuation interface

Project.toml

@@ -8,6 +8,7 @@ AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
+CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
@@ -61,12 +62,14 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 BifurcationKit = "0f109fa4-8a5d-4b75-95aa-f515264e7665"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DeepDiffs = "ab62b9b5-e342-54a8-a765-a90f495de1a6"
+HomotopyContinuation = "f213a82b-91d6-5c5d-acf7-10f1c761b327"
 LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 
 [extensions]
 MTKBifurcationKitExt = "BifurcationKit"
 MTKChainRulesCoreExt = "ChainRulesCore"
 MTKDeepDiffsExt = "DeepDiffs"
+MTKHomotopyContinuationExt = "HomotopyContinuation"
 MTKLabelledArraysExt = "LabelledArrays"
 
 [compat]
@@ -76,6 +79,7 @@ BifurcationKit = "0.3"
 BlockArrays = "1.1"
 ChainRulesCore = "1"
 Combinatorics = "1"
+CommonSolve = "0.2.4"
 Compat = "3.42, 4"
 ConstructionBase = "1"
 DataInterpolations = "6.4"
@@ -97,6 +101,7 @@ ForwardDiff = "0.10.3"
 FunctionWrappers = "1.1"
 FunctionWrappersWrappers = "0.1"
 Graphs = "1.5.2"
+HomotopyContinuation = "2.11"
 InteractiveUtils = "1"
 JuliaFormatter = "1.0.47"
 JumpProcesses = "9.13.1"

ext/MTKHomotopyContinuationExt.jl

@@ -0,0 +1,192 @@
+module MTKHomotopyContinuationExt
+
+using ModelingToolkit
+using ModelingToolkit.SciMLBase
+using ModelingToolkit.Symbolics: unwrap, symtype
+using ModelingToolkit.SymbolicIndexingInterface
+using ModelingToolkit.DocStringExtensions
+using HomotopyContinuation
+using ModelingToolkit: iscomplete, parameters, has_index_cache, get_index_cache, get_u0,
+                       get_u0_p, check_eqs_u0, CommonSolve
+
+const MTK = ModelingToolkit
+
+function contains_variable(x, wrt)
+    any(y -> occursin(y, x), wrt)
+end
+
+"""
+$(TYPEDSIGNATURES)
+
+Check if `x` is polynomial with respect to the variables in `wrt`.
+"""
+function is_polynomial(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))
+    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
+        end
+        exponent_has_unknowns = contains_variable(p, wrt)
+        if exponent_has_unknowns
+            @warn "In $x: Exponent $p cannot contain unknowns of the system."
+        end
+        base_polynomial = is_polynomial(b, wrt)
+        if !base_polynomial
+            @warn "In $x: Base is not a polynomial"
+        end
+        return base_polynomial && !exponent_has_unknowns && is_pow_integer
+    end
+    @warn "In $x: Unrecognized operation $(operation(x)). Allowed polynomial operations are `*, +, -, ^`"
+    return false
+end
+
+"""
+$(TYPEDSIGNATURES)
+
+Convert `expr` from a symbolics expression to one that uses `HomotopyContinuation.ModelKit`.
+"""
+function symbolics_to_hc(expr)
+    if iscall(expr)
+        if operation(expr) == getindex
+            args = arguments(expr)
+            return ModelKit.Variable(getname(args[1]), args[2:end]...)
+        else
+            return operation(expr)(symbolics_to_hc.(arguments(expr))...)
+        end
+    elseif symbolic_type(expr) == NotSymbolic()
+        return expr
+    else
+        return ModelKit.Variable(getname(expr))
+    end
+end
+
+"""
+$(TYPEDEF)
+
+A subtype of `HomotopyContinuation.AbstractSystem` used to solve `HomotopyContinuationProblem`s.
+"""
+struct MTKHomotopySystem{F, P, J, V} <: HomotopyContinuation.AbstractSystem
+    """
+    The generated function for the residual of the polynomial system. In-place.
+    """
+    f::F
+    """
+    The parameter object.
+    """
+    p::P
+    """
+    The generated function for the jacobian of the polynomial system. In-place.
+    """
+    jac::J
+    """
+    The `HomotopyContinuation.ModelKit.Variable` representation of the unknowns of
+    the system.
+    """
+    vars::V
+    """
+    The number of polynomials in the system. Must also be equal to `length(vars)`.
+    """
+    nexprs::Int
+end
+
+Base.size(sys::MTKHomotopySystem) = (sys.nexprs, length(sys.vars))
+ModelKit.variables(sys::MTKHomotopySystem) = sys.vars
+
+function (sys::MTKHomotopySystem)(x, p = nothing)
+    sys.f(x, sys.p)
+end
+
+function ModelKit.evaluate!(u, sys::MTKHomotopySystem, x, p = nothing)
+    sys.f(u, x, sys.p)
+end
+
+function ModelKit.evaluate_and_jacobian!(u, U, sys::MTKHomotopySystem, x, p = nothing)
+    sys.f(u, x, sys.p)
+    sys.jac(U, x, sys.p)
+end
+
+SymbolicIndexingInterface.parameter_values(s::MTKHomotopySystem) = s.p
+
+"""
+    $(TYPEDSIGNATURES)
+
+Create a `HomotopyContinuationProblem` from a `NonlinearSystem` with polynomial equations.
+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).
+"""
+function MTK.HomotopyContinuationProblem(
+        sys::NonlinearSystem, u0map, parammap = nothing; eval_expression = false,
+        eval_module = ModelingToolkit, 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)
+
+    for eq in eqs
+        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.")
+        end
+    end
+
+    nlfn, u0, p = MTK.process_SciMLProblem(NonlinearFunction{true}, sys, u0map, parammap;
+        jac = true, eval_expression, eval_module)
+
+    hvars = symbolics_to_hc.(dvs)
+    mtkhsys = MTKHomotopySystem(nlfn.f, p, nlfn.jac, hvars, length(eqs))
+
+    obsfn = MTK.ObservedFunctionCache(sys; eval_expression, eval_module)
+
+    return MTK.HomotopyContinuationProblem(u0, mtkhsys, sys, obsfn)
+end
+
+"""
+$(TYPEDSIGNATURES)
+
+Solve a `HomotopyContinuationProblem`. Ignores the algorithm passed to it, and always
+uses `HomotopyContinuation.jl`. All keyword arguments are forwarded to
+`HomotopyContinuation.solve`. 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`.
+"""
+function CommonSolve.solve(prob::MTK.HomotopyContinuationProblem,
+        alg = nothing; show_progress = false, kwargs...)
+    sol = HomotopyContinuation.solve(
+        prob.homotopy_continuation_system; show_progress, kwargs...)
+    realsols = HomotopyContinuation.results(sol; only_real = true)
+    if isempty(realsols)
+        u = state_values(prob)
+        resid = prob.homotopy_continuation_system(u)
+        retcode = SciMLBase.ReturnCode.ConvergenceFailure
+    else
+        distance, idx = findmin(realsols) do result
+            norm(result.solution - state_values(prob))
+        end
+        u = real.(realsols[idx].solution)
+        resid = prob.homotopy_continuation_system(u)
+        retcode = SciMLBase.ReturnCode.Success
+    end
+
+    return SciMLBase.build_solution(
+        prob, :HomotopyContinuation, u, resid; retcode, original = sol)
+end
+
+end

src/ModelingToolkit.jl

@@ -54,6 +54,7 @@ using Reexport
 using RecursiveArrayTools
 import Graphs: SimpleDiGraph, add_edge!, incidence_matrix
 import BlockArrays: BlockedArray, Block, blocksize, blocksizes
+import CommonSolve
 
 using RuntimeGeneratedFunctions
 using RuntimeGeneratedFunctions: drop_expr
@@ -281,4 +282,6 @@ export Clock, SolverStepClock, TimeDomain
 
 export MTKParameters, reorder_dimension_by_tunables!, reorder_dimension_by_tunables
 
+export HomotopyContinuationProblem
+
 end # module

src/systems/nonlinear/nonlinearsystem.jl

@@ -565,3 +565,55 @@ function Base.:(==)(sys1::NonlinearSystem, sys2::NonlinearSystem)
         _eq_unordered(get_ps(sys1), get_ps(sys2)) &&
         all(s1 == s2 for (s1, s2) in zip(get_systems(sys1), get_systems(sys2)))
 end
+
+"""
+$(TYPEDEF)
+
+A type of Nonlinear problem which specializes on polynomial systems and uses
+HomotopyContinuation.jl to solve the system. Requires importing HomotopyContinuation.jl to
+create and solve.
+"""
+struct HomotopyContinuationProblem{uType, H, O} <:
+       SciMLBase.AbstractNonlinearProblem{uType, true}
+    """
+    The initial values of states in the system. If there are multiple real roots of
+    the system, the one closest to this point is returned.
+    """
+    u0::uType
+    """
+    A subtype of `HomotopyContinuation.AbstractSystem` to solve. Also contains the
+    parameter object.
+    """
+    homotopy_continuation_system::H
+    """
+    The `NonlinearSystem` used to create this problem. Used for symbolic indexing.
+    """
+    sys::NonlinearSystem
+    """
+    A function which generates and returns observed expressions for the given system.
+    """
+    obsfn::O
+end
+
+function HomotopyContinuationProblem(::AbstractSystem, _u0, _p; kwargs...)
+    error("HomotopyContinuation.jl is required to create and solve `HomotopyContinuationProblem`s. Please run `Pkg.add(\"HomotopyContinuation\")` to continue.")
+end
+
+SymbolicIndexingInterface.symbolic_container(p::HomotopyContinuationProblem) = p.sys
+SymbolicIndexingInterface.state_values(p::HomotopyContinuationProblem) = p.u0
+function SymbolicIndexingInterface.set_state!(p::HomotopyContinuationProblem, args...)
+    set_state!(p.u0, args...)
+end
+function SymbolicIndexingInterface.parameter_values(p::HomotopyContinuationProblem)
+    parameter_values(p.homotopy_continuation_system)
+end
+function SymbolicIndexingInterface.set_parameter!(p::HomotopyContinuationProblem, args...)
+    set_parameter!(parameter_values(p), args...)
+end
+function SymbolicIndexingInterface.observed(p::HomotopyContinuationProblem, sym)
+    if p.obsfn !== nothing
+        return p.obsfn(sym)
+    else
+        return SymbolicIndexingInterface.observed(p.sys, sym)
+    end
+end

test/downstream/linearize.jl

@@ -121,13 +121,13 @@ lsys = ModelingToolkit.reorder_unknowns(lsys0, unknowns(ssys), desired_order)
 lsyss, _ = ModelingToolkit.linearize_symbolic(pid, [reference.u, measurement.u],
     [ctr_output.u])
 
-@test substitute(
+@test ModelingToolkit.fixpoint_sub(
     lsyss.A, ModelingToolkit.defaults_and_guesses(pid)) == lsys.A
-@test substitute(
+@test ModelingToolkit.fixpoint_sub(
     lsyss.B, ModelingToolkit.defaults_and_guesses(pid)) == lsys.B
-@test substitute(
+@test ModelingToolkit.fixpoint_sub(
     lsyss.C, ModelingToolkit.defaults_and_guesses(pid)) == lsys.C
-@test substitute(
+@test ModelingToolkit.fixpoint_sub(
     lsyss.D, ModelingToolkit.defaults_and_guesses(pid)) == lsys.D
 
 # Test with the reverse desired unknown order as well to verify that similarity transform and reoreder_unknowns really works

test/extensions/Project.toml

@@ -3,10 +3,12 @@ BifurcationKit = "0f109fa4-8a5d-4b75-95aa-f515264e7665"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+HomotopyContinuation = "f213a82b-91d6-5c5d-acf7-10f1c761b327"
 LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1"
 SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
+Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"