feat: initial implementation of HomotopyContinuation interface

Author: AayushSabharwal

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,114 @@
+module MTKHomotopyContinuationExt
+
+using ModelingToolkit
+using ModelingToolkit.Symbolics: unwrap
+using ModelingToolkit.SymbolicIndexingInterface
+using HomotopyContinuation
+using ModelingToolkit: iscomplete, parameters, has_index_cache, get_index_cache, get_u0,
+                       get_u0_p, check_eqs_u0, CommonSolve
+
+function contains_variable(x, wrt)
+    any(isequal(x), wrt) && return true
+    istree(x) || return false
+    return any(y -> contains_variable(y, wrt), arguments(x))
+end
+
+function is_polynomial(x, wrt)
+    x = unwrap(x)
+    symbolic_type(x) == NotSymbolic() && return true
+    istree(x) || return true
+    contains_variable(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)
+        return is_polynomial(b, wrt) && !contains_variable(p, wrt)
+    end
+    return false
+end
+
+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
+
+struct MTKHomotopySystem{F, P, J, V} <: HomotopyContinuation.AbstractSystem
+    f::F
+    p::P
+    jac::J
+    vars::V
+    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
+
+function ModelingToolkit.HomotopyContinuationProblem(
+        sys::NonlinearSystem, u0map, parammap; compile = :all, 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)
+    ps = parameters(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")
+        end
+    end
+
+    nlfn = NonlinearFunction(sys; jac = true)
+    hvars = symbolics_to_hc.(dvs)
+
+    if has_index_cache(sys) && get_index_cache(sys) !== nothing
+        u0, defs = get_u0(sys, u0map, parammap)
+        check_eqs_u0(eqs, dvs, u0; kwargs...)
+        p = MTKParameters(sys, parammap, u0map)
+    else
+        u0, p, defs = get_u0_p(sys, u0map, parammap; tofloat, use_union)
+        check_eqs_u0(eqs, dvs, u0; kwargs...)
+    end
+
+    mtkhsys = MTKHomotopySystem(nlfn.f, p, nlfn.jac, hvars, length(eqs))
+
+    prob = ModelingToolkit.HomotopyContinuationProblem{typeof(mtkhsys), typeof(u0)}(
+        sys, mtkhsys, u0)
+end
+
+function CommonSolve.solve(prob::ModelingToolkit.HomotopyContinuationProblem; kwargs...)
+    sol = HomotopyContinuation.solve(prob.hcsys; kwargs...)
+    rsols = HomotopyContinuation.real_solutions(sol)
+    rsol = findmin(rsols) do val
+        norm(prob.u0 - val)
+    end
+    return rsol
+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

src/systems/nonlinear/nonlinearsystem.jl

@@ -598,3 +598,25 @@ 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
+
+struct HomotopyContinuationProblem{H, U}
+    sys::NonlinearSystem
+    hcsys::H
+    u0::U
+end
+
+function HomotopyContinuationProblem(args...; kwargs...)
+    error("Requires HomotopyContinuationExt")
+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.hcsys)
+end
+function SymbolicIndexingInterface.set_parameter!(p::HomotopyContinuationProblem, args...)
+    set_parameter!(p.hcsys, args...)
+end