Skip to content

Commit

Permalink
refactor: use TaylorDiff.jl instead of TaylorSeries.jl
Browse files Browse the repository at this point in the history
  • Loading branch information
AayushSabharwal committed Jan 29, 2025
1 parent 5426e99 commit d2e8ec6
Show file tree
Hide file tree
Showing 4 changed files with 71 additions and 30 deletions.
4 changes: 2 additions & 2 deletions lib/NonlinearSolveHomotopyContinuation/Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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"

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ using NonlinearSolveBase
using SymbolicIndexingInterface
using LinearAlgebra
using ADTypes
using TaylorSeries
using TaylorDiff
using DocStringExtensions
import CommonSolve
import HomotopyContinuation as HC
Expand Down
85 changes: 62 additions & 23 deletions lib/NonlinearSolveHomotopyContinuation/src/interface_types.jl
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,7 @@ $(FIELDS)
"""
vars
"""
The `TaylorSeries.Taylor1` objects used to compute the taylor series of `f`.
The `TaylorDiff.TaylorScalar` objects used to compute the taylor series of `f`.
"""
taylorvars
"""
Expand Down Expand Up @@ -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(4)) 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(4)) 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(3)))
vars[2i] = TaylorScalar(imag(x[i]), ntuple(Returns(0.0), Val(3)))
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
Expand All @@ -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
Expand All @@ -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

Expand Down
10 changes: 6 additions & 4 deletions lib/NonlinearSolveHomotopyContinuation/src/solve.jl
Original file line number Diff line number Diff line change
Expand Up @@ -44,12 +44,14 @@ function homotopy_continuation_preprocessing(
end

taylorvars = if isscalar
Taylor1(zeros(ComplexF64, 5), 4)
[TaylorScalar(ntuple(Returns(0.0), 4)), TaylorScalar(ntuple(Returns(0.0), 4))]
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
Expand Down

0 comments on commit d2e8ec6

Please sign in to comment.