feat: support polynomials of invertible functions in HomotopyContinuationExt

Author: AayushSabharwal

ext/MTKHomotopyContinuationExt.jl

@@ -54,16 +54,73 @@ end
 
 PolynomialData() = PolynomialData(BasicSymbolic[], NonPolynomialReason.T[], false)
 
+abstract type PolynomialTransformationError <: Exception end
+
+struct MultivarTerm <: PolynomialTransformationError
+    term::Any
+    vars::Any
+end
+
+function Base.showerror(io::IO, err::MultivarTerm)
+    println(io,
+        "Cannot convert system to polynomial: Found term $(err.term) which is a function of multiple unknowns $(err.vars).")
+end
+
+struct MultipleTermsOfSameVar <: PolynomialTransformationError
+    terms::Any
+    var::Any
+end
+
+function Base.showerror(io::IO, err::MultipleTermsOfSameVar)
+    println(io,
+        "Cannot convert system to polynomial: Found multiple non-polynomial terms $(err.terms) involving the same unknown $(err.var).")
+end
+
+struct SymbolicSolveFailure <: PolynomialTransformationError
+    term::Any
+    var::Any
+end
+
+function Base.showerror(io::IO, err::SymbolicSolveFailure)
+    println(io,
+        "Cannot convert system to polynomial: Unable to symbolically solve $(err.term) for $(err.var).")
+end
+
+struct NemoNotLoaded <: PolynomialTransformationError end
+
+function Base.showerror(io::IO, err::NemoNotLoaded)
+    println(io,
+        "ModelingToolkit may be able to solve this system as a polynomial system if `Nemo` is loaded. Run `import Nemo` and try again.")
+end
+
+struct VariablesAsPolyAndNonPoly <: PolynomialTransformationError
+    vars::Any
+end
+
+function Base.showerror(io::IO, err::VariablesAsPolyAndNonPoly)
+    println(io,
+        "Cannot convert convert system to polynomial: Variables $(err.vars) occur in both polynomial and non-polynomial terms in the system.")
+end
+
 struct NotPolynomialError <: Exception
-    eq::Equation
-    data::PolynomialData
+    transformation_err::Union{PolynomialTransformationError, Nothing}
+    eq::Vector{Equation}
+    data::Vector{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))
+    if err.transformation_err !== nothing
+        Base.showerror(io, err.transformation_err)
+    end
+    for (eq, data) in zip(err.eq, err.data)
+        if isempty(data.non_polynomial_terms)
+            continue
+        end
+        println(io,
+            "Equation $(eq) is not a polynomial in the unknowns for the following reasons:")
+        for (term, reason) in zip(data.non_polynomial_terms, data.reasons)
+            println(io, display_reason(reason, term))
+        end
     end
 end
 
@@ -86,7 +143,9 @@ function process_polynomial!(data::PolynomialData, x, wrt)
     any(isequal(x), wrt) && return true
 
     if operation(x) in (*, +, -, /)
-        return all(y -> is_polynomial!(data, y, wrt), arguments(x))
+        # `map` because `all` will early exit, but we want to search
+        # through everything to get all the non-polynomial terms
+        return all(map(y -> is_polynomial!(data, y, wrt), arguments(x)))
     end
     if operation(x) == (^)
         b, p = arguments(x)
@@ -105,10 +164,6 @@ function process_polynomial!(data::PolynomialData, x, wrt)
             push!(data.reasons, NonPolynomialReason.ExponentContainsUnknowns)
         end
         base_polynomial = is_polynomial!(data, b, wrt)
-        if !base_polynomial
-            push!(data.non_polynomial_terms, x)
-            push!(data.reasons, NonPolynomialReason.BaseNotPolynomial)
-        end
         return base_polynomial && !exponent_has_unknowns && is_pow_integer
     end
     push!(data.non_polynomial_terms, x)
@@ -234,6 +289,12 @@ end
 
 SymbolicIndexingInterface.parameter_values(s::MTKHomotopySystem) = s.p
 
+struct PolynomialTransformationData
+    new_var::BasicSymbolic
+    term::BasicSymbolic
+    inv_term::Vector
+end
+
 """
     $(TYPEDSIGNATURES)
 
@@ -265,18 +326,95 @@ function MTK.HomotopyContinuationProblem(
     # CSE/hashconsing.
     eqs = full_equations(sys)
 
-    denoms = []
-    has_parametric_exponents = false
-    eqs2 = map(eqs) do eq
+    polydata = map(eqs) do eq
         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))
+        data
+    end
+
+    has_parametric_exponents = any(d -> d.has_parametric_exponent, polydata)
+
+    all_non_poly_terms = mapreduce(d -> d.non_polynomial_terms, vcat, polydata)
+    unique!(all_non_poly_terms)
+
+    var_to_nonpoly = Dict{BasicSymbolic, PolynomialTransformationData}()
+
+    is_poly = true
+    transformation_err = nothing
+    for t in all_non_poly_terms
+        # if the term involves multiple unknowns, we can't invert it
+        dvs_in_term = map(x -> occursin(x, t), dvs)
+        if count(dvs_in_term) > 1
+            transformation_err = MultivarTerm(t, dvs[dvs_in_term])
+            is_poly = false
+            break
+        end
+        # we already have a substitution solving for `var`
+        var = dvs[findfirst(dvs_in_term)]
+        if haskey(var_to_nonpoly, var) && !isequal(var_to_nonpoly[var].term, t)
+            transformation_err = MultipleTermsOfSameVar([t, var_to_nonpoly[var].term], var)
+            is_poly = false
+            break
+        end
+        # we want to solve `term - new_var` for `var`
+        new_var = gensym(Symbol(var))
+        new_var = unwrap(only(@variables $new_var))
+        invterm = Symbolics.ia_solve(
+            t - new_var, var; complex_roots = false, periodic_roots = false, warns = false)
+        # if we can't invert it, quit
+        if invterm === nothing || isempty(invterm)
+            transformation_err = SymbolicSolveFailure(t, var)
+            is_poly = false
+            break
+        end
+        # `ia_solve` returns lazy terms i.e. `asin(1.0)` instead of `pi/2`
+        # this just evaluates the constant expressions
+        invterm = Symbolics.substitute.(invterm, (Dict(),))
+        # RootsOf implies Symbolics couldn't solve the inner polynomial because
+        # `Nemo` wasn't loaded.
+        if any(x -> MTK.iscall(x) && MTK.operation(x) == Symbolics.RootsOf, invterm)
+            transformation_err = NemoNotLoaded()
+            is_poly = false
+            break
         end
-        num, den = handle_rational_polynomials(eq.rhs - eq.lhs, dvs)
+        var_to_nonpoly[var] = PolynomialTransformationData(new_var, t, invterm)
+    end
+
+    if !is_poly
+        throw(NotPolynomialError(transformation_err, eqs, polydata))
+    end
+
+    subrules = Dict()
+    combinations = Vector[]
+    new_dvs = []
+    for x in dvs
+        if haskey(var_to_nonpoly, x)
+            _data = var_to_nonpoly[x]
+            subrules[_data.term] = _data.new_var
+            push!(combinations, _data.inv_term)
+            push!(new_dvs, _data.new_var)
+        else
+            push!(combinations, [x])
+            push!(new_dvs, x)
+        end
+    end
+    all_solutions = collect.(collect(Iterators.product(combinations...)))
 
+    denoms = []
+    eqs2 = map(eqs) do eq
+        t = eq.rhs - eq.lhs
+        t = Symbolics.fixpoint_sub(t, subrules; maxiters = length(dvs))
+        # the substituted variable occurs outside the substituted term
+        poly_and_nonpoly = map(dvs) do x
+            haskey(var_to_nonpoly, x) && occursin(x, t)
+        end
+        if any(poly_and_nonpoly)
+            throw(NotPolynomialError(
+                VariablesAsPolyAndNonPoly(dvs[poly_and_nonpoly]), eqs, polydata))
+        end
+
+        num, den = handle_rational_polynomials(t, new_dvs)
         # make factors different elements, otherwise the nonzero factors artificially
         # inflate the error of the zero factor.
         if iscall(den) && operation(den) == *
@@ -292,16 +430,19 @@ function MTK.HomotopyContinuationProblem(
     end
 
     sys2 = MTK.@set sys.eqs = eqs2
+    MTK.@set! sys2.unknowns = new_dvs
     # 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)
+    _, u0, p = MTK.process_SciMLProblem(
+        MTK.EmptySciMLFunction, sys, u0map, parammap; eval_expression, eval_module)
+    nlfn = NonlinearFunction{true}(sys2; jac = true, eval_expression, eval_module)
 
-    denominator = MTK.build_explicit_observed_function(sys, denoms)
+    denominator = MTK.build_explicit_observed_function(sys2, denoms)
+    unpack_solution = MTK.build_explicit_observed_function(sys2, all_solutions)
 
-    hvars = symbolics_to_hc.(dvs)
+    hvars = symbolics_to_hc.(new_dvs)
     mtkhsys = MTKHomotopySystem(nlfn.f, p, nlfn.jac, hvars, length(eqs))
 
     obsfn = MTK.ObservedFunctionCache(sys; eval_expression, eval_module)
@@ -319,7 +460,7 @@ function MTK.HomotopyContinuationProblem(
         solver_and_starts = HomotopyContinuation.solver_startsolutions(mtkhsys; kwargs...)
     end
     return MTK.HomotopyContinuationProblem(
-        u0, mtkhsys, denominator, sys, obsfn, solver_and_starts)
+        u0, mtkhsys, denominator, sys, obsfn, solver_and_starts, unpack_solution)
 end
 
 """
@@ -353,25 +494,35 @@ function CommonSolve.solve(prob::MTK.HomotopyContinuationProblem,
     if isempty(realsols)
         u = state_values(prob)
         retcode = SciMLBase.ReturnCode.ConvergenceFailure
+        resid = prob.homotopy_continuation_system(u)
     else
         T = eltype(state_values(prob))
-        distance, idx = findmin(realsols) do result
+        distance = T(Inf)
+        u = state_values(prob)
+        resid = nothing
+        for result in realsols
             if any(<=(denominator_abstol),
                 prob.denominator(real.(result.solution), parameter_values(prob)))
-                return T(Inf)
+                continue
+            end
+            for truesol in prob.unpack_solution(result.solution, parameter_values(prob))
+                dist = norm(truesol - state_values(prob))
+                if dist < distance
+                    distance = dist
+                    u = T.(real.(truesol))
+                    resid = T.(real.(prob.homotopy_continuation_system(result.solution)))
+                end
             end
-            norm(result.solution - state_values(prob))
         end
         # all roots cause denominator to be zero
         if isinf(distance)
             u = state_values(prob)
+            resid = prob.homotopy_continuation_system(u)
             retcode = SciMLBase.ReturnCode.Infeasible
         else
-            u = real.(realsols[idx].solution)
             retcode = SciMLBase.ReturnCode.Success
         end
     end
-    resid = prob.homotopy_continuation_system(u)
 
     return SciMLBase.build_solution(
         prob, :HomotopyContinuation, u, resid; retcode, original = sol)

src/systems/nonlinear/nonlinearsystem.jl

@@ -690,7 +690,7 @@ 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, D, O, SS} <:
+struct HomotopyContinuationProblem{uType, H, D, O, SS, U} <:
        SciMLBase.AbstractNonlinearProblem{uType, true}
     """
     The initial values of states in the system. If there are multiple real roots of
@@ -721,6 +721,12 @@ struct HomotopyContinuationProblem{uType, H, D, O, SS} <:
     `HomotopyContinuation.solver_startsystems`.
     """
     solver_and_starts::SS
+    """
+    A function which takes a solution of the transformed system, and returns a vector
+    of solutions for the original system. This is utilized when converting systems
+    to polynomials.
+    """
+    unpack_solution::U
 end
 
 function HomotopyContinuationProblem(::AbstractSystem, _u0, _p; kwargs...)

test/extensions/homotopy_continuation.jl

@@ -75,19 +75,59 @@ end
 @testset "Polynomial check and warnings" begin
     @variables x = 1.0
     @mtkbuild sys = NonlinearSystem([x^1.5 + x^2 - 1 ~ 0])
-    @test_throws ["Exponent", "not an integer", "not a polynomial"] HomotopyContinuationProblem(
+    @test_throws ["Cannot convert", "Unable", "symbolically solve",
+        "Exponent", "not an integer", "not a polynomial"] HomotopyContinuationProblem(
         sys, [])
     @mtkbuild sys = NonlinearSystem([x^x - x ~ 0])
-    @test_throws ["Exponent", "unknowns", "not a polynomial"] HomotopyContinuationProblem(
+    @test_throws ["Cannot convert", "Unable", "symbolically solve",
+        "Exponent", "unknowns", "not a polynomial"] HomotopyContinuationProblem(
         sys, [])
     @mtkbuild sys = NonlinearSystem([((x^2) / sin(x))^2 + x ~ 0])
-    @test_throws ["recognized", "sin", "not a polynomial"] HomotopyContinuationProblem(
+    @test_throws ["Cannot convert", "both polynomial", "non-polynomial",
+        "recognized", "sin", "not a polynomial"] HomotopyContinuationProblem(
         sys, [])
 
     @variables y = 2.0
     @mtkbuild sys = NonlinearSystem([x^2 + y^2 + 2 ~ 0, y ~ sin(x)])
-    @test_throws ["recognized", "sin", "not a polynomial"] HomotopyContinuationProblem(
+    @test_throws ["Cannot convert", "recognized", "sin", "not a polynomial"] HomotopyContinuationProblem(
         sys, [])
+
+    @mtkbuild sys = NonlinearSystem([x^2 + y^2 - 2 ~ 0, sin(x + y) ~ 0])
+    @test_throws ["Cannot convert", "function of multiple unknowns"] HomotopyContinuationProblem(
+        sys, [])
+
+    @mtkbuild sys = NonlinearSystem([sin(x)^2 + 1 ~ 0, cos(y) - cos(x) - 1 ~ 0])
+    @test_throws ["Cannot convert", "multiple non-polynomial terms", "same unknown"] HomotopyContinuationProblem(
+        sys, [])
+
+    @mtkbuild sys = NonlinearSystem([sin(x^2)^2 + sin(x^2) - 1 ~ 0])
+    @test_throws ["import Nemo"] HomotopyContinuationProblem(sys, [])
+end
+
+import Nemo
+
+@testset "With Nemo" begin
+    @variables x = 2.0
+    @mtkbuild sys = NonlinearSystem([sin(x^2)^2 + sin(x^2) - 1 ~ 0])
+    prob = HomotopyContinuationProblem(sys, [])
+    @test prob[1] ≈ 2.0
+    sol = solve(prob; threading = false)
+    _x = sol[1]
+    @test sin(_x^2)^2 + sin(_x^2) - 1≈0.0 atol=1e-12
+end
+
+@testset "Function of polynomial" begin
+    @variables x=0.25 y=0.125
+    a = sin(x^2 - 4x + 1)
+    b = cos(3log(y) + 4)
+    @mtkbuild sys = NonlinearSystem([(a^2 - 4a * b + 4b^2) / (a - 0.25) ~ 0
+                                     (a^2 - 0.75a + 0.125) ~ 0])
+    prob = HomotopyContinuationProblem(sys, [])
+    @test prob[x] ≈ 0.25
+    @test prob[y] ≈ 0.125
+    sol = solve(prob; threading = false)
+    @test sol[a]≈0.5 atol=1e-6
+    @test sol[b]≈0.25 atol=1e-6
 end
 
 @testset "Rational functions" begin