diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml index 8d8e5f123..455e52be9 100644 --- a/.github/workflows/Test.yml +++ b/.github/workflows/Test.yml @@ -42,6 +42,7 @@ jobs: - Back/FiniteDiff - Back/FiniteDifferences - Back/ForwardDiff + - Back/GTPSA - Back/Mooncake - Back/PolyesterForwardDiff - Back/ReverseDiff diff --git a/DifferentiationInterface/Project.toml b/DifferentiationInterface/Project.toml index 9eaf7fb73..6f48e2b9a 100644 --- a/DifferentiationInterface/Project.toml +++ b/DifferentiationInterface/Project.toml @@ -17,6 +17,7 @@ FastDifferentiation = "eb9bf01b-bf85-4b60-bf87-ee5de06c00be" FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +GTPSA = "b27dd330-f138-47c5-815b-40db9dd9b6e8" Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" @@ -35,6 +36,7 @@ DifferentiationInterfaceFastDifferentiationExt = "FastDifferentiation" DifferentiationInterfaceFiniteDiffExt = "FiniteDiff" DifferentiationInterfaceFiniteDifferencesExt = "FiniteDifferences" DifferentiationInterfaceForwardDiffExt = ["ForwardDiff", "DiffResults"] +DifferentiationInterfaceGTPSAExt = "GTPSA" DifferentiationInterfaceMooncakeExt = "Mooncake" DifferentiationInterfacePolyesterForwardDiffExt = "PolyesterForwardDiff" DifferentiationInterfaceReverseDiffExt = ["ReverseDiff", "DiffResults"] @@ -57,6 +59,7 @@ FastDifferentiation = "0.4.1" FiniteDiff = "2.23.1" FiniteDifferences = "0.12.31" ForwardDiff = "0.10.36" +GTPSA = "1.4.0" JuliaFormatter = "1" LinearAlgebra = "<0.0.1,1" Mooncake = "0.4.52" @@ -84,6 +87,7 @@ FastDifferentiation = "eb9bf01b-bf85-4b60-bf87-ee5de06c00be" FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +GTPSA = "b27dd330-f138-47c5-815b-40db9dd9b6e8" JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899" diff --git a/DifferentiationInterface/README.md b/DifferentiationInterface/README.md index 85586851b..4066195de 100644 --- a/DifferentiationInterface/README.md +++ b/DifferentiationInterface/README.md @@ -37,6 +37,7 @@ We support the following backends defined by [ADTypes.jl](https://github.com/Sci - [FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl) - [FiniteDifferences.jl](https://github.com/JuliaDiff/FiniteDifferences.jl) - [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) +- [GTPSA.jl](https://github.com/bmad-sim/GTPSA.jl) - [Mooncake.jl](https://github.com/compintell/Mooncake.jl) - [PolyesterForwardDiff.jl](https://github.com/JuliaDiff/PolyesterForwardDiff.jl) - [ReverseDiff.jl](https://github.com/JuliaDiff/ReverseDiff.jl) diff --git a/DifferentiationInterface/docs/src/explanation/backends.md b/DifferentiationInterface/docs/src/explanation/backends.md index 53ea81ffe..f206f90d9 100644 --- a/DifferentiationInterface/docs/src/explanation/backends.md +++ b/DifferentiationInterface/docs/src/explanation/backends.md @@ -11,6 +11,7 @@ We support the following dense backend choices from [ADTypes.jl](https://github. - [`AutoFiniteDiff`](@extref ADTypes.AutoFiniteDiff) - [`AutoFiniteDifferences`](@extref ADTypes.AutoFiniteDifferences) - [`AutoForwardDiff`](@extref ADTypes.AutoForwardDiff) +- [`AutoGTPSA`](@extref ADTypes.AutoGTPSA) - [`AutoMooncake`](@extref ADTypes.AutoMooncake) - [`AutoPolyesterForwardDiff`](@extref ADTypes.AutoPolyesterForwardDiff) - [`AutoReverseDiff`](@extref ADTypes.AutoReverseDiff) @@ -45,6 +46,7 @@ In practice, many AD backends have custom implementations for high-level operato | `AutoFiniteDiff` | 🔀 | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | | `AutoFiniteDifferences` | 🔀 | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | | `AutoForwardDiff` | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | + | `AutoGTPSA` | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | | `AutoMooncake` | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | `AutoPolyesterForwardDiff` | 🔀 | ❌ | 🔀 | ✅ | ✅ | 🔀 | 🔀 | 🔀 | | `AutoReverseDiff` | ❌ | 🔀 | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | @@ -64,6 +66,7 @@ Moreover, each context type is supported by a specific subset of backends: | `AutoFiniteDiff` | ✅ | | `AutoFiniteDifferences` | ✅ | | `AutoForwardDiff` | ✅ | +| `AutoGTPSA` | ✅ | | `AutoMooncake` | ✅ | | `AutoPolyesterForwardDiff` | ✅ | | `AutoReverseDiff` | ✅ | @@ -138,6 +141,15 @@ Nothing specific to mention. We implement [`pushforward`](@ref) directly using [`Dual` numbers](https://juliadiff.org/ForwardDiff.jl/stable/dev/how_it_works/), and preparation allocates the necessary space. For higher level operators, preparation creates a [config object](https://juliadiff.org/ForwardDiff.jl/stable/user/api/#Preallocating/Configuring-Work-Buffers), which can be type-unstable. +### GTPSA + +For all operators, preparation preallocates the input [`TPS`s](https://bmad-sim.github.io/GTPSA.jl/stable/man/c_tps/), and for in-place functions the output `TPS`s as well. For minimal allocations of `TPS` temporaries inside of a function, the [`@FastGTPSA`/`@FastGTPSA!`](https://bmad-sim.github.io/GTPSA.jl/stable/man/j_fastgtpsa/) macros are recommended. + +If a GTPSA [`Descriptor`](https://bmad-sim.github.io/GTPSA.jl/stable/man/b_descriptor/) is not provided to `AutoGTPSA`, then a `Descriptor` will be generated in preparation based on the context. + +!!! danger + When providing a custom GTPSA `Descriptor` to `AutoGTPSA`, it is the responsibility of the user to ensure that the number of [GTPSA "variables"](https://bmad-sim.github.io/GTPSA.jl/stable/quickstart/#Calculating-a-Truncated-Power-Series) specified in the `Descriptor` is consistent with the number of inputs of the provided function. Undefined behavior and crashes may occur if this is not the case. + ### PolyesterForwardDiff Most operators fall back on `AutoForwardDiff`. diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl new file mode 100644 index 000000000..b2d24dac3 --- /dev/null +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl @@ -0,0 +1,12 @@ +module DifferentiationInterfaceGTPSAExt + +import DifferentiationInterface as DI +using ADTypes: AutoGTPSA +using GTPSA: GTPSA, TPS, Descriptor + +DI.check_available(::AutoGTPSA) = true + +include("onearg.jl") +include("twoarg.jl") + +end diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl new file mode 100644 index 000000000..ca7937d63 --- /dev/null +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl @@ -0,0 +1,600 @@ +## Pushforward + +# Contains either a single pre-allocated initial TPS +# or a vector of pre-allocated TPSs. +struct GTPSAOneArgPushforwardPrep{X} <: DI.PushforwardPrep + xt::X +end + +function DI.prepare_pushforward( + ::F, backend::AutoGTPSA{D}, x, tx::NTuple, ::Vararg{DI.Constant,C} +) where {F,D,C} + + # For pushforward/JVP, we only actually need 1 single variable (in the GTPSA sense) + # because we even if we did multiple we will add up the derivatives of each at the end. + if D != Nothing + d = backend.descriptor + else + d = Descriptor(1, 1) # 1 variable to first order + end + if x isa Number + xt = TPS{promote_type(typeof(first(tx)), typeof(x), Float64)}(; use=d) + return GTPSAOneArgPushforwardPrep(xt) + else + xt = similar(x, TPS{promote_type(eltype(first(tx)), eltype(x), Float64)}) + for i in eachindex(xt) + xt[i] = TPS{promote_type(eltype(first(tx)), eltype(x), Float64)}(; use=d) + end + return GTPSAOneArgPushforwardPrep(xt) + end +end + +function DI.pushforward( + f, + prep::GTPSAOneArgPushforwardPrep, + ::AutoGTPSA, + x, + tx::NTuple, + contexts::Vararg{DI.Constant,C}, +) where {C} + fc = DI.with_contexts(f, contexts...) + ty = map(tx) do dx + foreach((t, xi, dxi) -> (t[0] = xi; t[1] = dxi), prep.xt, x, dx) + yt = fc(prep.xt) + if yt isa Number + return yt[1] + else + dy = map(t -> t[1], yt) + return dy + end + end + return ty +end + +function DI.pushforward!( + f, + ty::NTuple, + prep::GTPSAOneArgPushforwardPrep, + ::AutoGTPSA, + x, + tx::NTuple, + contexts::Vararg{DI.Constant,C}, +) where {C} + fc = DI.with_contexts(f, contexts...) + for b in eachindex(tx, ty) + dx, dy = tx[b], ty[b] + foreach((t, xi, dxi) -> (t[0] = xi; t[1] = dxi), prep.xt, x, dx) + yt = fc(prep.xt) + map!(t -> t[1], dy, yt) + end + return ty +end + +function DI.value_and_pushforward( + f, + prep::GTPSAOneArgPushforwardPrep, + backend::AutoGTPSA, + x, + tx::NTuple, + contexts::Vararg{DI.Constant,C}, +) where {C} + fc = DI.with_contexts(f, contexts...) + ty = DI.pushforward(fc, prep, backend, x, tx) + y = fc(x) # TO-DO: optimize + return y, ty +end + +function DI.value_and_pushforward!( + f, + ty::NTuple, + prep::GTPSAOneArgPushforwardPrep, + backend::AutoGTPSA, + x, + tx::NTuple, + contexts::Vararg{DI.Constant,C}, +) where {C} + fc = DI.with_contexts(f, contexts...) + DI.pushforward!(fc, ty, prep, backend, x, tx) + y = fc(x) # TO-DO: optimize + return y, ty +end + +## Gradient +# Contains a vector of pre-allocated TPSs. +struct GTPSAOneArgGradientPrep{X} <: DI.GradientPrep + xt::X +end + +# Unlike JVP, this requires us to use all variables +function DI.prepare_gradient( + f, backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Constant,C} +) where {D,C} + if D != Nothing + d = backend.descriptor + else + d = Descriptor(length(x), 1) # n variables to first order + end + xt = similar(x, TPS{promote_type(eltype(x), Float64)}) + j = 1 + for i in eachindex(xt) + xt[i] = TPS{promote_type(eltype(x), Float64)}(; use=d) + xt[i][j] = 1 + j += 1 + end + return GTPSAOneArgGradientPrep(xt) +end + +function DI.gradient( + f, prep::GTPSAOneArgGradientPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Constant,C} +) where {C} + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part (slopes set in prepare) + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) + grad = similar(x, GTPSA.numtype(yt)) + GTPSA.gradient!(grad, yt; include_params=true, unsafe_inbounds=true) + return grad +end + +function DI.gradient!( + f, grad, prep::GTPSAOneArgGradientPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Constant,C} +) where {C} + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) + GTPSA.gradient!(grad, yt; include_params=true, unsafe_inbounds=true) + return grad +end + +function DI.value_and_gradient( + f, prep::GTPSAOneArgGradientPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Constant,C} +) where {C} + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part (slopes set in prepare) + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) + grad = similar(x, GTPSA.numtype(yt)) + GTPSA.gradient!(grad, yt; include_params=true, unsafe_inbounds=true) + return yt[0], grad +end + +function DI.value_and_gradient!( + f, grad, prep::GTPSAOneArgGradientPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Constant,C} +) where {C} + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part (slopes set in prepare) + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) + GTPSA.gradient!(grad, yt; include_params=true, unsafe_inbounds=true) + return yt[0], grad +end + +## Jacobian +# Contains a vector of pre-allocated TPSs +struct GTPSAOneArgJacobianPrep{X} <: DI.JacobianPrep + xt::X +end + +# To materialize the entire Jacobian we use all variables +function DI.prepare_jacobian( + f, backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Constant,C} +) where {D,C} + if D != Nothing + d = backend.descriptor + else + d = Descriptor(length(x), 1) # n variables to first order + end + + # We set the slopes of each variable to 1 here, this will always be the case for Jacobian + xt = similar(x, TPS{promote_type(eltype(x), Float64)}) + j = 1 + for i in eachindex(xt) + xt[i] = TPS{promote_type(eltype(x), Float64)}(; use=d) + xt[i][j] = 1 + j += 1 + end + return GTPSAOneArgJacobianPrep(xt) +end + +function DI.jacobian( + f, prep::GTPSAOneArgJacobianPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Constant,C} +) where {C} + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) + jac = similar(x, GTPSA.numtype(eltype(yt)), (length(yt), length(x))) + GTPSA.jacobian!(jac, yt; include_params=true, unsafe_inbounds=true) + return jac +end + +function DI.jacobian!( + f, jac, prep::GTPSAOneArgJacobianPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Constant,C} +) where {C} + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) + GTPSA.jacobian!(jac, yt; include_params=true, unsafe_inbounds=true) + return jac +end + +function DI.value_and_jacobian( + f, prep::GTPSAOneArgJacobianPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Constant,C} +) where {C} + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) + jac = similar(x, GTPSA.numtype(eltype(yt)), (length(yt), length(x))) + GTPSA.jacobian!(jac, yt; include_params=true, unsafe_inbounds=true) + y = map(t -> t[0], yt) + return y, jac +end + +function DI.value_and_jacobian!( + f, jac, prep::GTPSAOneArgJacobianPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Constant,C} +) where {C} + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) + GTPSA.jacobian!(jac, yt; include_params=true, unsafe_inbounds=true) + y = map(t -> t[0], yt) + return y, jac +end + +## Second derivative +# Contains single pre-allocated TPS +struct GTPSAOneArgSecondDerivativePrep{X} <: DI.SecondDerivativePrep + xt::X +end + +function DI.prepare_second_derivative( + f, backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Constant,C} +) where {D,C} + if D != Nothing + d = backend.descriptor + else + d = Descriptor(1, 2) + end + xt = TPS{promote_type(typeof(x), Float64)}(; use=d) + xt[1] = 1 # Set slope + return GTPSAOneArgSecondDerivativePrep(xt) +end + +function DI.second_derivative( + f, + prep::GTPSAOneArgSecondDerivativePrep, + backend::AutoGTPSA{D}, + x, + contexts::Vararg{DI.Constant,C}, +) where {D,C} + prep.xt[0] = x + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) + if D == Nothing + idx2 = 2 + else + idx2 = GTPSA.numnn(backend.descriptor) + 1 # index of first second derivative + end + + if yt isa Number + return yt[idx2] * 2 + else + der2 = similar(yt, GTPSA.numtype(eltype(yt))) + for i in eachindex(yt) + der2[i] = yt[i][idx2] * 2 # *2 because monomial coefficient is 1/2 + end + return der2 + end +end + +function DI.second_derivative!( + f, + der2, + prep::GTPSAOneArgSecondDerivativePrep, + backend::AutoGTPSA{D}, + x, + contexts::Vararg{DI.Constant,C}, +) where {D,C} + prep.xt[0] = x + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) + if D == Nothing + idx2 = 2 + else + idx2 = GTPSA.numnn(backend.descriptor) + 1 # index of first second derivative + end + for i in eachindex(yt) + der2[i] = yt[i][idx2] * 2 + end + return der2 +end + +function DI.value_derivative_and_second_derivative( + f, + prep::GTPSAOneArgSecondDerivativePrep, + backend::AutoGTPSA{D}, + x, + contexts::Vararg{DI.Constant,C}, +) where {D,C} + prep.xt[0] = x + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) + if D == Nothing + idx2 = 2 + else + idx2 = GTPSA.numnn(backend.descriptor) + 1 # index of first second derivative + end + if yt isa Number + return yt[0], yt[1], yt[idx2] * 2 + else + y = map(t -> t[0], yt) + der = similar(yt, GTPSA.numtype(eltype(yt))) + der2 = similar(yt, GTPSA.numtype(eltype(yt))) + for i in eachindex(yt) + der[i] = yt[i][1] + der2[i] = yt[i][idx2] * 2 + end + return y, der, der2 + end +end + +function DI.value_derivative_and_second_derivative!( + f, + der, + der2, + prep::GTPSAOneArgSecondDerivativePrep, + backend::AutoGTPSA{D}, + x, + contexts::Vararg{DI.Constant,C}, +) where {D,C} + prep.xt[0] = x + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) + y = map(t -> t[0], yt) + if D == Nothing + idx2 = 2 + else + idx2 = GTPSA.numnn(backend.descriptor) + 1 # index of first second derivative + end + for i in eachindex(yt) + der[i] = yt[i][1] + der2[i] = yt[i][idx2] * 2 + end + return y, der, der2 +end + +## Hessian +# Stores allocated array of TPS and an array for the monomial coefficient +# indexing in GTPSA.cycle! (which is used if a Descriptor is specified) +struct GTPSAOneArgHessianPrep{X,M} <: DI.HessianPrep + xt::X + m::M +end + +function DI.prepare_hessian( + f, backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Constant,C} +) where {D,C} + if D != Nothing + d = backend.descriptor + m = Vector{UInt8}(undef, length(x)) + else + nn = length(x) + d = Descriptor(nn, 2) + # If all variables/variable+parameters have truncation order > 2, then + # the indexing is known beforehand and we can do it (very slightly) faster + m = nothing + end + xt = similar(x, TPS{promote_type(eltype(x), Float64)}) + + # xt and x have same indexing because of similar + # Setting the first derivatives must be 1-based + # linear with the variables. + j = 1 + for i in eachindex(xt) + xt[i] = TPS{promote_type(eltype(x), Float64)}(; use=d) + xt[i][j] = 1 + j += 1 + end + + return GTPSAOneArgHessianPrep(xt, m) +end + +function DI.hessian( + f, prep::GTPSAOneArgHessianPrep, ::AutoGTPSA{D}, x, contexts::Vararg{DI.Constant,C} +) where {D,C} + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) + hess = similar(x, GTPSA.numtype(yt), (length(x), length(x))) + unsafe_fast = D == Nothing ? true : false + GTPSA.hessian!( + hess, + yt; + include_params=true, + unsafe_inbounds=true, + unsafe_fast=unsafe_fast, + tmp_mono=prep.m, + ) + return hess +end + +function DI.hessian!( + f, + hess, + prep::GTPSAOneArgHessianPrep, + ::AutoGTPSA{D}, + x, + contexts::Vararg{DI.Constant,C}, +) where {D,C} + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) + unsafe_fast = D == Nothing ? true : false + GTPSA.hessian!( + hess, + yt; + include_params=true, + unsafe_inbounds=true, + unsafe_fast=unsafe_fast, + tmp_mono=prep.m, + ) + return hess +end + +function DI.value_gradient_and_hessian( + f, prep::GTPSAOneArgHessianPrep, ::AutoGTPSA{D}, x, contexts::Vararg{DI.Constant,C} +) where {D,C} + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) + grad = similar(x, GTPSA.numtype(yt)) + GTPSA.gradient!(grad, yt; include_params=true, unsafe_inbounds=true) + hess = similar(x, GTPSA.numtype(yt), (length(x), length(x))) + unsafe_fast = D == Nothing ? true : false + GTPSA.hessian!( + hess, + yt; + include_params=true, + unsafe_inbounds=true, + unsafe_fast=unsafe_fast, + tmp_mono=prep.m, + ) + return yt[0], grad, hess +end + +function DI.value_gradient_and_hessian!( + f, + grad, + hess, + prep::GTPSAOneArgHessianPrep, + ::AutoGTPSA{D}, + x, + contexts::Vararg{DI.Constant,C}, +) where {D,C} + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) + GTPSA.gradient!(grad, yt; include_params=true, unsafe_inbounds=true) + unsafe_fast = D == Nothing ? true : false + GTPSA.hessian!( + hess, + yt; + include_params=true, + unsafe_inbounds=true, + unsafe_fast=unsafe_fast, + tmp_mono=prep.m, + ) + return yt[0], grad, hess +end + +struct GTPSAOneArgHVPPrep{E,H} <: DI.HVPPrep + hessprep::E + hess::H +end + +function DI.prepare_hvp( + f, backend::AutoGTPSA, x, tx::NTuple, contexts::Vararg{DI.Constant,C} +) where {C} + hessprep = DI.prepare_hessian(f, backend, x) + fc = DI.with_contexts(f, contexts...) + hess = similar(x, typeof(fc(x)), (length(x), length(x))) + return GTPSAOneArgHVPPrep(hessprep, hess) +end + +function DI.hvp( + f, + prep::GTPSAOneArgHVPPrep, + backend::AutoGTPSA, + x, + tx::NTuple, + contexts::Vararg{DI.Constant,C}, +) where {C} + DI.hessian!(f, prep.hess, prep.hessprep, backend, x, contexts...) + tg = map(tx) do dx + dg = similar(x, eltype(prep.hess)) + dg .= 0 + j = 1 + for dxi in dx + for i in 1:size(prep.hess, 1) + dg[i] += prep.hess[i, j] * dxi + end + j += 1 + end + return dg + end + return tg +end + +function DI.hvp!( + f, + tg::NTuple, + prep::GTPSAOneArgHVPPrep, + backend::AutoGTPSA, + x, + tx::NTuple, + contexts::Vararg{DI.Constant,C}, +) where {C} + DI.hessian!(f, prep.hess, prep.hessprep, backend, x, contexts...) + for b in eachindex(tg) + dx, dg = tx[b], tg[b] + dg .= 0 + j = 1 + for dxi in dx + for i in 1:size(prep.hess, 1) + dg[i] += prep.hess[i, j] * dxi + end + j += 1 + end + end + return tg +end + +function DI.gradient_and_hvp( + f, + prep::GTPSAOneArgHVPPrep, + backend::AutoGTPSA{D}, + x, + tx::NTuple, + contexts::Vararg{DI.Constant,C}, +) where {D,C} + grad = similar(x, eltype(prep.hess)) + DI.value_gradient_and_hessian!( + f, grad, prep.hess, prep.hessprep, backend, x, contexts... + ) + tg = map(tx) do dx + dg = similar(x, eltype(prep.hess)) + dg .= 0 + j = 1 + for dxi in dx + for i in 1:size(prep.hess, 1) + dg[i] += prep.hess[i, j] * dxi + end + j += 1 + end + return dg + end + return grad, tg +end + +function DI.gradient_and_hvp!( + f, + grad, + tg, + prep::GTPSAOneArgHVPPrep, + backend::AutoGTPSA{D}, + x, + tx::NTuple, + contexts::Vararg{DI.Constant,C}, +) where {D,C} + DI.value_gradient_and_hessian!( + f, grad, prep.hess, prep.hessprep, backend, x, contexts... + ) + for b in eachindex(tg, tx) + dg, dx = tg[b], tx[b] + dg .= 0 + j = 1 + for dxi in dx + for i in 1:size(prep.hess, 1) + dg[i] += prep.hess[i, j] * dxi + end + j += 1 + end + end + return grad, tg +end diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl new file mode 100644 index 000000000..3e2e88de5 --- /dev/null +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl @@ -0,0 +1,195 @@ +## Pushforward + +# Input: Contains either a single pre-allocated initial TPS +# or a vector of pre-allocated TPSs. +# +# Output: Contains a vector of pre-allocated TPSs +struct GTPSATwoArgPushforwardPrep{X,Y} <: DI.PushforwardPrep + xt::X + yt::Y +end + +function DI.prepare_pushforward( + ::F, y, backend::AutoGTPSA{D}, x, tx::NTuple, ::Vararg{DI.Constant,C} +) where {F,D,C} + + # For pushforward/JVP, we only actually need 1 single variable (in the GTPSA sense) + # because we even if we did multiple we will add up the derivatives of each at the end. + if D != Nothing + d = backend.descriptor + else + d = Descriptor(1, 1) # 1 variable to first order + end + if x isa Number + xt = TPS{promote_type(typeof(first(tx)), typeof(x), Float64)}(; use=d) + else + xt = similar(x, TPS{promote_type(eltype(first(tx)), eltype(x), Float64)}) + for i in eachindex(xt) + xt[i] = TPS{promote_type(eltype(first(tx)), eltype(x), Float64)}(; use=d) + end + end + + yt = similar(y, TPS{promote_type(eltype(y), Float64)}) + + for i in eachindex(yt) + yt[i] = TPS{promote_type(eltype(y), Float64)}(; use=d) + end + return GTPSATwoArgPushforwardPrep(xt, yt) +end + +function DI.pushforward( + f!, + y, + prep::GTPSATwoArgPushforwardPrep, + ::AutoGTPSA, + x, + tx::NTuple, + contexts::Vararg{DI.Constant,C}, +) where {C} + fc! = DI.with_contexts(f!, contexts...) + ty = map(tx) do dx + foreach((t, xi, dxi) -> (t[0] = xi; t[1] = dxi), prep.xt, x, dx) + fc!(prep.yt, prep.xt) + dy = map(t -> t[1], prep.yt) + return dy + end + map!(t -> t[0], y, prep.yt) + return ty +end + +function DI.pushforward!( + f!, + y, + ty::NTuple, + prep::GTPSATwoArgPushforwardPrep, + ::AutoGTPSA, + x, + tx::NTuple, + contexts::Vararg{DI.Constant,C}, +) where {C} + fc! = DI.with_contexts(f!, contexts...) + for b in eachindex(tx, ty) + dx, dy = tx[b], ty[b] + foreach((t, xi, dxi) -> (t[0] = xi; t[1] = dxi), prep.xt, x, dx) + fc!(prep.yt, prep.xt) + map!(t -> t[1], dy, prep.yt) + end + map!(t -> t[0], y, prep.yt) + return ty +end + +function DI.value_and_pushforward( + f!, + y, + prep::GTPSATwoArgPushforwardPrep, + backend::AutoGTPSA, + x, + tx::NTuple, + contexts::Vararg{DI.Constant,C}, +) where {C} + ty = DI.pushforward(f!, y, prep, backend, x, tx, contexts...) + return y, ty +end + +function DI.value_and_pushforward!( + f!, + y, + ty::NTuple, + prep::GTPSATwoArgPushforwardPrep, + backend::AutoGTPSA, + x, + tx::NTuple, + contexts::Vararg{DI.Constant,C}, +) where {C} + DI.pushforward!(f!, y, ty, prep, backend, x, tx, contexts...) + return y, ty +end + +## Jacobian +# Input: Contains a vector of pre-allocated TPSs +# Output: Contains a vector of pre-allocated TPSs +struct GTPSATwoArgJacobianPrep{X,Y} <: DI.JacobianPrep + xt::X + yt::Y +end + +function DI.prepare_jacobian( + f, y, backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Constant,C} +) where {D,C} + if D != Nothing + d = backend.descriptor + else + d = Descriptor(length(x), 1) # n variables to first order + end + + # We set the slopes of each variable to 1 here, this will always be the case for Jacobian + xt = similar(x, TPS{promote_type(eltype(x), Float64)}) + j = 1 + for i in eachindex(xt) + xt[i] = TPS{promote_type(eltype(x), Float64)}(; use=d) + xt[i][j] = 1 + j += 1 + end + + yt = similar(y, TPS{promote_type(eltype(y), Float64)}) + + for i in eachindex(yt) + yt[i] = TPS{promote_type(eltype(y), Float64)}(; use=d) + end + + return GTPSATwoArgJacobianPrep(xt, yt) +end + +function DI.jacobian( + f!, y, prep::GTPSATwoArgJacobianPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Constant,C} +) where {C} + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part + fc! = DI.with_contexts(f!, contexts...) + fc!(prep.yt, prep.xt) + jac = similar(x, GTPSA.numtype(eltype(prep.yt)), (length(prep.yt), length(x))) + GTPSA.jacobian!(jac, prep.yt; include_params=true, unsafe_inbounds=true) + map!(t -> t[0], y, prep.yt) + return jac +end + +function DI.jacobian!( + f!, + y, + jac, + prep::GTPSATwoArgJacobianPrep, + ::AutoGTPSA, + x, + contexts::Vararg{DI.Constant,C}, +) where {C} + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part + fc! = DI.with_contexts(f!, contexts...) + fc!(prep.yt, prep.xt) + GTPSA.jacobian!(jac, prep.yt; include_params=true, unsafe_inbounds=true) + map!(t -> t[0], y, prep.yt) + return jac +end + +function DI.value_and_jacobian( + f!, + y, + prep::GTPSATwoArgJacobianPrep, + backend::AutoGTPSA, + x, + contexts::Vararg{DI.Constant,C}, +) where {C} + jac = DI.jacobian(f!, y, prep, backend, x, contexts...) # y set on line 151 + return y, jac +end + +function DI.value_and_jacobian!( + f!, + y, + jac, + prep::GTPSATwoArgJacobianPrep, + backend::AutoGTPSA, + x, + contexts::Vararg{DI.Constant,C}, +) where {C} + DI.jacobian!(f!, y, jac, prep, backend, x, contexts...) + return y, jac +end diff --git a/DifferentiationInterface/src/DifferentiationInterface.jl b/DifferentiationInterface/src/DifferentiationInterface.jl index 37997c220..99ecfc2d4 100644 --- a/DifferentiationInterface/src/DifferentiationInterface.jl +++ b/DifferentiationInterface/src/DifferentiationInterface.jl @@ -23,6 +23,7 @@ using ADTypes: AutoFiniteDiff, AutoFiniteDifferences, AutoForwardDiff, + AutoGTPSA, AutoMooncake, AutoPolyesterForwardDiff, AutoReverseDiff, @@ -110,6 +111,7 @@ export AutoFastDifferentiation export AutoFiniteDiff export AutoFiniteDifferences export AutoForwardDiff +export AutoGTPSA export AutoMooncake export AutoPolyesterForwardDiff export AutoReverseDiff diff --git a/DifferentiationInterface/test/Back/GTPSA/test.jl b/DifferentiationInterface/test/Back/GTPSA/test.jl new file mode 100644 index 000000000..3dd9599d3 --- /dev/null +++ b/DifferentiationInterface/test/Back/GTPSA/test.jl @@ -0,0 +1,29 @@ +using Pkg +Pkg.add("GTPSA") + +using DifferentiationInterface, DifferentiationInterfaceTest +using GTPSA: GTPSA +using Test + +using ExplicitImports +check_no_implicit_imports(DifferentiationInterface) + +LOGGING = get(ENV, "CI", "false") == "false" + +for backend in [AutoGTPSA()] + @test check_available(backend) + @test check_inplace(backend) +end + +# Test no Descriptor (use context) +test_differentiation(AutoGTPSA(); type_stability=:full, logging=LOGGING); + +# Test with Descriptor: +d1 = GTPSA.Descriptor(20, 2) # 20 variables to 2nd order +test_differentiation(AutoGTPSA(d1); type_stability=:full, logging=LOGGING); + +# Test with Descriptor using varying orders +vos = 2 * ones(Int, 20) +vos[1] = 3 +d2 = GTPSA.Descriptor(vos, 3) +test_differentiation(AutoGTPSA(d2); type_stability=:full, logging=LOGGING);