refactor: use TaylorDiff.jl instead of TaylorSeries.jl

Author: AayushSabharwal

lib/NonlinearSolveHomotopyContinuation/Project.toml

@@ -14,7 +14,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NonlinearSolveBase = "be0214bd-f91f-a760-ac4e-3421ce2b2da0"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
-TaylorSeries = "6aa5eb33-94cf-58f4-a9d0-e4b2c4fc25ea"
+TaylorDiff = "b36ab563-344f-407b-a36a-4f200bebf99c"
 
 [compat]
 ADTypes = "1.11.0"
@@ -29,7 +29,7 @@ NonlinearSolve = "4"
 NonlinearSolveBase = "1.3.3"
 SciMLBase = "2.71"
 SymbolicIndexingInterface = "0.3.36"
-TaylorSeries = "0.18.2"
+TaylorDiff = "0.3.1"
 Test = "1.10"
 julia = "1.10"
 

lib/NonlinearSolveHomotopyContinuation/src/NonlinearSolveHomotopyContinuation.jl

@@ -6,7 +6,7 @@ using NonlinearSolveBase
 using SymbolicIndexingInterface
 using LinearAlgebra
 using ADTypes
-using TaylorSeries
+using TaylorDiff
 using DocStringExtensions
 import CommonSolve
 import HomotopyContinuation as HC

lib/NonlinearSolveHomotopyContinuation/src/interface_types.jl

@@ -161,47 +161,74 @@ for V in (Inplace, OutOfPlace, Scalar)
 end
 
 function update_taylorvars_from_taylorvector!(
-        vars, x::HC.ModelKit.TaylorVector{M}) where {M}
-    for i in eachindex(vars)
-        for j in 0:(M - 1)
-            vars[i][j] = x[i, j + 1]
+        vars, x::HC.ModelKit.TaylorVector)
+    for i in eachindex(x)
+        xvar = x[i]
+        realx = ntuple(Val(5)) do j
+            j <= length(xvar) ? real(xvar[j - 1]) : 0.0
         end
-        for j in M:4
-            vars[i][j] = zero(vars[i][j])
+        imagx = ntuple(Val(5)) do j
+            j <= length(xvar) ? imag(xvar[j - 1]) : 0.0
         end
+
+        vars[2i - 1] = TaylorScalar(realx)
+        vars[2i] = TaylorScalar(imagx)
     end
 end
 
 function update_taylorvars_from_taylorvector!(vars, x::AbstractVector)
-    for i in eachindex(vars)
-        vars[i][0] = x[i]
-        for j in 1:4
-            vars[i][j] = zero(vars[i][j])
-        end
+    for i in eachindex(x)
+        vars[2i - 1] = TaylorScalar(real(x[i]), ntuple(Returns(0.0), Val(4)))
+        vars[2i] = TaylorScalar(imag(x[i]), ntuple(Returns(0.0), Val(4)))
+    end
+end
+
+function check_taylor_equality(vars, x::HC.ModelKit.TaylorVector)
+    for i in eachindex(x)
+        TaylorDiff.flatten(vars[2i-1]) == map(real, x[i]) || return false
+        TaylorDiff.flatten(vars[2i]) == map(imag, x[i]) || return false
+    end
+    return true
+end
+function check_taylor_equality(vars, x::AbstractVector)
+    for i in eachindex(x)
+        TaylorDiff.value(vars[2i-1]) != real(x[i]) && return false
+        TaylorDiff.value(vars[2i]) != imag(x[i]) && return false
     end
+    return true
 end
 
 function update_maybe_taylorvector_from_taylorvars!(
         u::Vector, vars, buffer, ::Val{N}) where {N}
     for i in eachindex(vars)
-        u[i] = buffer[i][N]
+        rval = TaylorDiff.flatten(real(buffer[i]))
+        ival = TaylorDiff.flatten(imag(buffer[i]))
+        u[i] = rval[N] + im * ival[N]
     end
 end
 
 function update_maybe_taylorvector_from_taylorvars!(
-        u::HC.ModelKit.TaylorVector, vars, buffer, ::Val{N}) where {N}
+        u::HC.ModelKit.TaylorVector{M}, vars, buffer, ::Val{N}) where {M, N}
     for i in eachindex(vars)
-        u[i] = ntuple(j -> buffer[i][j - 1], Val(N + 1))
+        rval = TaylorDiff.flatten(real(buffer[i]))
+        ival = TaylorDiff.flatten(imag(buffer[i]))
+        u[i] = ntuple(i -> rval[i]  + im * ival[i], Val(length(rval)))
     end
 end
 
 function HC.ModelKit.taylor!(u::AbstractVector, ::Val{N},
         sys::HomotopySystemWrapper{Inplace}, x, p = nothing) where {N}
     f = sys.f
     p = sys.p
-    buffer, vars = sys.taylorvars
-    update_taylorvars_from_taylorvector!(vars, x)
-    f(buffer, vars, p)
+    vars, buffer = sys.taylorvars
+    if !check_taylor_equality(vars, x)
+        update_taylorvars_from_taylorvector!(vars, x)
+        vars = reinterpret(Complex{eltype(vars)}, vars)
+        buffer = reinterpret(Complex{eltype(buffer)}, buffer)
+        f(buffer, vars, p)
+    else
+        vars = reinterpret(Complex{eltype(vars)}, vars)
+    end
     update_maybe_taylorvector_from_taylorvars!(u, vars, buffer, Val(N))
     return u
 end
@@ -211,8 +238,14 @@ function HC.ModelKit.taylor!(u::AbstractVector, ::Val{N},
     f = sys.f
     p = sys.p
     vars = sys.taylorvars
-    update_taylorvars_from_taylorvector!(vars, x)
-    buffer = f(vars, p)
+    if !check_taylor_equality(vars, x)
+        update_taylorvars_from_taylorvector!(vars, x)
+        vars = reinterpret(Complex{eltype(vars)}, vars)
+        buffer = f(vars, p)
+        copyto!(vars, buffer)
+    else
+        vars = buffer = reinterpret(Complex{eltype(vars)}, vars)
+    end
     update_maybe_taylorvector_from_taylorvars!(u, vars, buffer, Val(N))
     return u
 end
@@ -222,9 +255,15 @@ function HC.ModelKit.taylor!(u::AbstractVector, ::Val{N},
     f = sys.f
     p = sys.p
     var = sys.taylorvars
-    update_taylorvars_from_taylorvector!((var,), x)
-    buffer = f(var, p)
-    update_maybe_taylorvector_from_taylorvars!(u, (var,), (buffer,), Val(N))
+    if !check_taylor_equality(var, x)
+        update_taylorvars_from_taylorvector!(var, x)
+        var = reinterpret(Complex{eltype(var)}, var)
+        buffer = f(var[1], p)
+        var[1] = buffer
+    else
+        var = buffer = reinterpret(Complex{eltype(var)}, var)
+    end
+    update_maybe_taylorvector_from_taylorvars!(u, var, buffer, Val(N))
     return u
 end
 

lib/NonlinearSolveHomotopyContinuation/src/solve.jl

@@ -44,12 +44,14 @@ function homotopy_continuation_preprocessing(
     end
 
     taylorvars = if isscalar
-        Taylor1(zeros(ComplexF64, 5), 4)
+        [TaylorScalar(ntuple(Returns(0.0), 5)), TaylorScalar(ntuple(Returns(0.0), 5))]
     elseif iip
-        ([Taylor1(zeros(ComplexF64, 5), 4) for _ in u0],
-            [Taylor1(zeros(ComplexF64, 5), 4) for _ in u0])
+        (
+            [TaylorScalar(ntuple(Returns(0.0), 4)) for _ in 1:2length(u0)],
+            [TaylorScalar(ntuple(Returns(0.0), 4)) for _ in 1:2length(u0)]
+        )
     else
-        [Taylor1(zeros(ComplexF64, 5), 4) for _ in u0]
+        [TaylorScalar(ntuple(Returns(0.0), 4)) for _ in 1:2length(u0)]
     end
 
     jacobian_buffers = if isscalar