From 5a69e77be405ebd20b06c056fdf5e42289106b74 Mon Sep 17 00:00:00 2001 From: mattsignorelli Date: Mon, 17 Jun 2024 19:15:29 -0400 Subject: [PATCH 01/44] add extension directory --- .../DifferentiationInterfaceGTPSAExt.jl | 98 +++++++++++++++++++ 1 file changed, 98 insertions(+) create mode 100644 DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl new file mode 100644 index 000000000..06ebe41d9 --- /dev/null +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl @@ -0,0 +1,98 @@ +module DifferentiationInterfaceGTPSAExt + +using ADTypes: AbstractADType, ForwardMode +import ADTypes: mode +import DifferentiationInterface as DI +using DifferentiationInterface: + DerivativeExtras, + GradientExtras, + HessianExtras, + JacobianExtras, + NoDerivativeExtras, + NoSecondDerivativeExtras, + PushforwardExtras +using GTPSA + +struct AutoGTPSA <: AbstractADType end + +mode(::AutoGTPSA) = ForwardMode() + +DI.check_available(::AutoGTPSA) = true + +# --- pushforward --- + +""" + GTPSAPushforwardExtras{T} <: PushforwardExtras + +This struct contains pre-allocated TPS(s) corresponding to `x` with seed +`dx` for the derivative. If taking derivative of single-variable function, +then T is a TPS, else T is some vector of TPSs. +""" +struct GTPSAPushforwardExtras{T} <: PushforwardExtras + v::T +end + +function DI.prepare_pushforward(f, backend::AutoGTPSA, x, dx) + if x isa Number + d = Descriptor(1,1) + t = TPS(use=d) + t[0] = x + t[1] = dx + return GTPSAPushforwardExtras(t) + else + NV = length(x) + d = Descriptor(NV,1) # only first order + v = similar(x, TPS) + + # v 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(v) + v[i] = TPS(use=d) # allocate + v[i][0] = x[i] + v[i][j] = dx[i] + j += 1 + end + return GTPSAPushforwardExtras(v) + end +end + +function DI.pushforward(f, ::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras{T}) where {T} + w = f(extras.v) + + if w isa Number + return w[1] + else + dy = similar(extras.v, GTPSA.numtype(eltype(w))) + # once again, derivatives in TPSs must be + # 1-based linear indexing + j = 1 + for i in eachindex(dy) + dy[i] = w[i][j] + j += 1 + end + return dy + end +end + +function DI.pushforward!(f, dy, ::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras{T}) where {T} + w = f(extras.v) + + if w isa Number + return w[1] # this should never be reached when this ! function is called + else + # once again, derivatives in TPSs must be + # 1-based linear indexing + j = 1 + for i in eachindex(dy) + dy[i] = w[i][j] + j += 1 + end + return dy + end +end + + + +end # module From aeedb5dbd91eda7f0aa7f5446bcad65c91b12afa Mon Sep 17 00:00:00 2001 From: mattsignorelli Date: Tue, 18 Jun 2024 08:31:56 -0400 Subject: [PATCH 02/44] derivative,second derivative, gradient --- DifferentiationInterface/Project.toml | 17 +- .../DifferentiationInterfaceGTPSAExt.jl | 93 +---- .../onearg.jl | 346 ++++++++++++++++++ DifferentiationInterfaceTest/Project.toml | 1 + 4 files changed, 369 insertions(+), 88 deletions(-) create mode 100644 DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl diff --git a/DifferentiationInterface/Project.toml b/DifferentiationInterface/Project.toml index 9199ac2db..a1c41eb4c 100644 --- a/DifferentiationInterface/Project.toml +++ b/DifferentiationInterface/Project.toml @@ -8,6 +8,7 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" Compat = "34da2185-b29b-5c13-b0c7-acf172513d20" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" +GTPSA = "b27dd330-f138-47c5-815b-40db9dd9b6e8" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" @@ -73,7 +74,6 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" -# DifferentiationInterfaceTest = "a82114a7-5aa3-49a8-9643-716bb13727a3" Diffractor = "9f5e2b26-1114-432f-b630-d3fe2085c51c" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" FastDifferentiation = "eb9bf01b-bf85-4b60-bf87-ee5de06c00be" @@ -96,17 +96,4 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [targets] -test = [ - "ADTypes", - "Aqua", - "DataFrames", - # "DifferentiationInterfaceTest", - "JET", - "JuliaFormatter", - "Pkg", - "SparseArrays", - "SparseConnectivityTracer", - "SparseMatrixColorings", - "StableRNGs", - "Test", -] +test = ["ADTypes", "Aqua", "DataFrames", "JET", "JuliaFormatter", "Pkg", "SparseArrays", "SparseConnectivityTracer", "SparseMatrixColorings", "StableRNGs", "Test"] diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl index 06ebe41d9..0bda8a3e7 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl @@ -5,94 +5,41 @@ import ADTypes: mode import DifferentiationInterface as DI using DifferentiationInterface: DerivativeExtras, + SecondDerivativeExtras, GradientExtras, - HessianExtras, JacobianExtras, - NoDerivativeExtras, - NoSecondDerivativeExtras, + HessianExtras, PushforwardExtras using GTPSA -struct AutoGTPSA <: AbstractADType end - -mode(::AutoGTPSA) = ForwardMode() - -DI.check_available(::AutoGTPSA) = true - -# --- pushforward --- - -""" - GTPSAPushforwardExtras{T} <: PushforwardExtras -This struct contains pre-allocated TPS(s) corresponding to `x` with seed -`dx` for the derivative. If taking derivative of single-variable function, -then T is a TPS, else T is some vector of TPSs. """ -struct GTPSAPushforwardExtras{T} <: PushforwardExtras - v::T -end + AutoGTPSA{D} -function DI.prepare_pushforward(f, backend::AutoGTPSA, x, dx) - if x isa Number - d = Descriptor(1,1) - t = TPS(use=d) - t[0] = x - t[1] = dx - return GTPSAPushforwardExtras(t) - else - NV = length(x) - d = Descriptor(NV,1) # only first order - v = similar(x, TPS) +Struct used to select the [GTPSA.jl](https://github.com/bmad-sim/GTPSA.jl) backend for automatic differentiation. - # v 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(v) - v[i] = TPS(use=d) # allocate - v[i][0] = x[i] - v[i][j] = dx[i] - j += 1 - end - return GTPSAPushforwardExtras(v) - end -end +# Constructors -function DI.pushforward(f, ::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras{T}) where {T} - w = f(extras.v) + AutoGTPSA(; descriptor=nothing) - if w isa Number - return w[1] - else - dy = similar(extras.v, GTPSA.numtype(eltype(w))) - # once again, derivatives in TPSs must be - # 1-based linear indexing - j = 1 - for i in eachindex(dy) - dy[i] = w[i][j] - j += 1 - end - return dy - end -end +# Fields -function DI.pushforward!(f, dy, ::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras{T}) where {T} - w = f(extras.v) + - `descriptor::D`: can be either - if w isa Number - return w[1] # this should never be reached when this ! function is called - else - # once again, derivatives in TPSs must be - # 1-based linear indexing - j = 1 - for i in eachindex(dy) - dy[i] = w[i][j] - j += 1 - end - return dy - end + + a GTPSA `Descriptor` specifying the number of variables/parameters, parameter + order, individual variable/parameter truncation orders, and maximum order. See + [here](https://bmad-sim.github.io/GTPSA.jl/stable/man/c_descriptor/) for more details. + + `nothing` to automatically use a `Descriptor` given the context + +""" +Base.@kwdef struct AutoGTPSA{D} <: AbstractADType + descriptor::D = nothing end +mode(::AutoGTPSA) = ForwardMode() + +DI.check_available(::AutoGTPSA) = true +include("onearg.jl") end # module diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl new file mode 100644 index 000000000..fd4a4d838 --- /dev/null +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl @@ -0,0 +1,346 @@ +# --- derivative --- +struct GTPSADerivativeExtras <: DerivativeExtras + t::TPS +end + +function DI.prepare_derivative(f, backend::AutoGTPSA{D}, x) where {D} + if D != Nothing + d = backend.descriptor + else + d = Descriptor(1,1) + end + t = TPS(use=d) + t[1] = 1 + return GTPSADerivativeExtras(t) +end + +function DI.derivative(f, backend::AutoGTPSA, x, extras::GTPSADerivativeExtras) + extras.t[0] = x + yt = f(extras.t) + if yt isa Number + return yt[1] + else + der = similar(yt, GTPSA.numtype(eltype(yt))) + for i in eachindex(yt) + der[i] = yt[i][1] + end + return der + end +end + +function DI.derivative!(f, der, backend::AutoGTPSA, x, extras::GTPSADerivativeExtras) + extras.t[0] = x + yt = f(extras.t) + if yt isa Number # This should never be reached + return yt[1] + else + for i in eachindex(yt) + der[i] = yt[i][1] + end + return der + end +end + +function DI.value_and_derivative(f, backend::AutoGTPSA, x, extras::GTPSADerivativeExtras) + extras.t[0] = x + yt = f(extras.t) + if yt isa Number + return yt[0], yt[1] + else + y = map(t->t[0], yt) + der = similar(yt, GTPSA.numtype(eltype(yt))) + for i in eachindex(yt) + der[i] = yt[i][1] + end + return y, der + end +end + +function DI.value_and_derivative!(f, der, backend::AutoGTPSA, x, extras::GTPSADerivativeExtras) + extras.t[0] = x + yt = f(extras.t) + if yt isa Number # This should never be reached + return yt[0], yt[1] + else + y = map(t->t[0], yt) + for i in eachindex(yt) + der[i] = yt[i][1] + end + return y, der + end +end + +# --- second_derivative --- +struct GTPSASecondDerivativeExtras <: SecondDerivativeExtras + t::TPS +end + +function DI.prepare_second_derivative(f, backend::AutoGTPSA{D}, x) where {D} + if D != Nothing + d = backend.descriptor + else + d = Descriptor(1,2) + end + t = TPS(use=d) + t[1] = 1 + return GTPSASecondDerivativeExtras(t) +end + +function DI.second_derivative(f, backend::AutoGTPSA, x, extras::GTPSADerivativeExtras) + extras.t[0] = x + yt = f(extras.t) + if yt isa Number + return yt[2] + else + der2 = similar(yt, GTPSA.numtype(eltype(yt))) + for i in eachindex(yt) + der2[i] = yt[i][2] + end + return der2 + end +end + +function DI.second_derivative!(f, der2, backend::AutoGTPSA, x, extras::GTPSADerivativeExtras) + extras.t[0] = x + yt = f(extras.t) + if yt isa Number # This should never be reached + return yt[2] + else + for i in eachindex(yt) + der2[i] = yt[i][2] + end + return der2 + end +end + +function DI.value_derivative_and_second_derivative(f, backend::AutoGTPSA, x, extras::GTPSADerivativeExtras) + extras.t[0] = x + yt = f(extras.t) + if yt isa Number + return yt[0], yt[1], yt[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][2] + end + return y, der, der2 + end +end + +function DI.value_derivative_and_second_derivative!(f, der, der2, backend::AutoGTPSA, x, extras::GTPSADerivativeExtras) + extras.t[0] = x + yt = f(extras.t) + if yt isa Number # This should never be reached + return yt[0], yt[1], yt[2] + else + y = map(t->t[0], yt) + for i in eachindex(yt) + der[i] = yt[i][1] + der2[i] = yt[i][2] + end + return y, der, der2 + end +end + +# --- gradient --- +struct GTPSAGradientExtras <: GradientExtras + v::Vector{TPS} +end + +function DI.prepare_gradient(f, backend::AutoGTPSA{D}, x) where {D} + if D != Nothing + d = backend.descriptor + nn = GTPSA.numnn(d) + else + nn = length(x) + d = Descriptor(nn,1) + end + v = similar(x, TPS) + + # v 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(v) + v[i] = TPS(use=d) + v[i][j] = 1 + j += 1 + end + + return GTPSAGradientExtras(v) +end + +function DI.gradient(f, backend::AutoGTPSA, x, extras::GTPSAGradientExtras) + foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + grad = similar(x, GTPSA.numtype(eltype(yt))) + GTPSA.gradient!(grad, yt, include_params=true) + return grad +end + +function DI.gradient!(f, grad, backend::AutoGTPSA, x, extras::GTPSAGradientExtras) + foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + GTPSA.gradient!(grad, yt, include_params=true) + return grad +end + +function DI.value_and_gradient(f, backend::AutoGTPSA, x, extras::GTPSAGradientExtras) + foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + grad = similar(x, GTPSA.numtype(eltype(yt))) + GTPSA.gradient!(grad, yt, include_params=true) + y = map(t->t[0], yt) + return y, grad +end + +function DI.value_and_gradient!(f, grad, backend::AutoGTPSA, x, extras::GTPSAGradientExtras) + foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + GTPSA.gradient!(grad, yt, include_params=true) + y = map(t->t[0], yt) + return y, grad +end + +# --- jacobian --- +struct GTPSAJacobianExtras <: JacobianExtras + v::Vector{TPS} +end + + + + +# --- pushforward --- + +""" + GTPSAPushforwardExtras{T} <: PushforwardExtras + +This struct contains pre-allocated TPS(s) corresponding to `x` with seed +`dx` for the derivative. If taking derivative of single-variable function, +then T is a TPS, else T is some vector of TPSs. +""" +struct GTPSAPushforwardExtras{T} <: PushforwardExtras + v::T +end + +function DI.prepare_pushforward(f, backend::AutoGTPSA, x, dx) + if x isa Number + d = Descriptor(1,1) + t = TPS(use=d) + t[0] = x + t[1] = dx + return GTPSAPushforwardExtras(t) + else + NV = length(x) + d = Descriptor(NV,1) # only first order + v = similar(x, TPS) + + # v 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(v) + v[i] = TPS(use=d) # allocate + v[i][0] = x[i] + v[i][j] = dx[i] + j += 1 + end + return GTPSAPushforwardExtras(v) + end +end + +function DI.pushforward(f, ::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras{T}) where {T} + w = f(extras.v) + + if w isa Number + if x isa Number # Single variable scalar function derivative + return w[1] + else # Multivariable scalar function derivative (gradient) + dy = similar(w, GTPSA.numtype(eltype(w))) + return GTPSA.gradient!(dy, w, include_params=true) + end + else + if x isa Number # Single variable vector function derivative + + else # Multivariable vector function derivative (Jacobian) + dy = similar(w, GTPSA.numtype(eltype(w)), Size(length(w),length(x))) + end + dy = similar(extras.v, GTPSA.numtype(eltype(w))) + # once again, derivatives in TPSs must be + # 1-based linear indexing + j = 1 + for i in eachindex(dy) + dy[i] = w[i][j] + j += 1 + end + return dy + + end +end + +function DI.value_and_pushforward(f, ::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras{T}) where {T} + w = f(extras.v) + + if w isa Number + return w[0], w[1] + else + dy = similar(extras.v, GTPSA.numtype(eltype(w))) + # once again, derivatives in TPSs must be + # 1-based linear indexing + j = 1 + for i in eachindex(dy) + dy[i] = w[i][j] + j += 1 + end + return map(t->t[0], w), dy + end +end + +function DI.pushforward!(f, dy, ::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras{T}) where {T} + w = f(extras.v) + + if w isa Number + return w[1] # this should never be reached when this ! function is called + else + # once again, derivatives in TPSs must be + # 1-based linear indexing + j = 1 + for i in eachindex(dy) + dy[i] = w[i][j] + j += 1 + end + return dy + end +end + +function DI.value_and_pushforward!(f, dy, ::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras{T}) where {T} + w = f(extras.v) + + if w isa Number + return w[0], w[1] # this should never be reached when this ! function is called + else + # once again, derivatives in TPSs must be + # 1-based linear indexing + j = 1 + for i in eachindex(dy) + dy[i] = w[i][j] + j += 1 + end + return map(t->t[0], w), dy + end +end + +# --- derivative --- +""" + GTPSADerivativeExtras{T} <: PushforwardExtras + +This struct contains pre-allocated TPS(s) corresponding to `x` with seed +`dx` for the derivative. If taking derivative of single-variable function, +then T is a TPS, else T is some vector of TPSs. +""" +struct GTPSAPushforwardExtras{T} <: PushforwardExtras + v::T +end \ No newline at end of file diff --git a/DifferentiationInterfaceTest/Project.toml b/DifferentiationInterfaceTest/Project.toml index b0877332c..61516b001 100644 --- a/DifferentiationInterfaceTest/Project.toml +++ b/DifferentiationInterfaceTest/Project.toml @@ -10,6 +10,7 @@ Compat = "34da2185-b29b-5c13-b0c7-acf172513d20" ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" +GTPSA = "b27dd330-f138-47c5-815b-40db9dd9b6e8" JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" From 9572d57fc1904a03a0a1e474d5a0020c68367a97 Mon Sep 17 00:00:00 2001 From: mattsignorelli Date: Tue, 18 Jun 2024 11:45:10 -0400 Subject: [PATCH 03/44] jacobian, hessian --- .../onearg.jl | 238 +++++++++--------- 1 file changed, 113 insertions(+), 125 deletions(-) diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl index fd4a4d838..4a92af488 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl @@ -14,7 +14,7 @@ function DI.prepare_derivative(f, backend::AutoGTPSA{D}, x) where {D} return GTPSADerivativeExtras(t) end -function DI.derivative(f, backend::AutoGTPSA, x, extras::GTPSADerivativeExtras) +function DI.derivative(f, ::AutoGTPSA, x, extras::GTPSADerivativeExtras) extras.t[0] = x yt = f(extras.t) if yt isa Number @@ -28,7 +28,7 @@ function DI.derivative(f, backend::AutoGTPSA, x, extras::GTPSADerivativeExtras) end end -function DI.derivative!(f, der, backend::AutoGTPSA, x, extras::GTPSADerivativeExtras) +function DI.derivative!(f, der, ::AutoGTPSA, x, extras::GTPSADerivativeExtras) extras.t[0] = x yt = f(extras.t) if yt isa Number # This should never be reached @@ -41,7 +41,7 @@ function DI.derivative!(f, der, backend::AutoGTPSA, x, extras::GTPSADerivativeEx end end -function DI.value_and_derivative(f, backend::AutoGTPSA, x, extras::GTPSADerivativeExtras) +function DI.value_and_derivative(f, ::AutoGTPSA, x, extras::GTPSADerivativeExtras) extras.t[0] = x yt = f(extras.t) if yt isa Number @@ -56,7 +56,7 @@ function DI.value_and_derivative(f, backend::AutoGTPSA, x, extras::GTPSADerivati end end -function DI.value_and_derivative!(f, der, backend::AutoGTPSA, x, extras::GTPSADerivativeExtras) +function DI.value_and_derivative!(f, der, ::AutoGTPSA, x, extras::GTPSADerivativeExtras) extras.t[0] = x yt = f(extras.t) if yt isa Number # This should never be reached @@ -86,7 +86,7 @@ function DI.prepare_second_derivative(f, backend::AutoGTPSA{D}, x) where {D} return GTPSASecondDerivativeExtras(t) end -function DI.second_derivative(f, backend::AutoGTPSA, x, extras::GTPSADerivativeExtras) +function DI.second_derivative(f, ::AutoGTPSA, x, extras::GTPSASecondDerivativeExtras) extras.t[0] = x yt = f(extras.t) if yt isa Number @@ -100,7 +100,7 @@ function DI.second_derivative(f, backend::AutoGTPSA, x, extras::GTPSADerivativeE end end -function DI.second_derivative!(f, der2, backend::AutoGTPSA, x, extras::GTPSADerivativeExtras) +function DI.second_derivative!(f, der2, ::AutoGTPSA, x, extras::GTPSASecondDerivativeExtras) extras.t[0] = x yt = f(extras.t) if yt isa Number # This should never be reached @@ -113,7 +113,7 @@ function DI.second_derivative!(f, der2, backend::AutoGTPSA, x, extras::GTPSADeri end end -function DI.value_derivative_and_second_derivative(f, backend::AutoGTPSA, x, extras::GTPSADerivativeExtras) +function DI.value_derivative_and_second_derivative(f, ::AutoGTPSA, x, extras::GTPSASecondDerivativeExtras) extras.t[0] = x yt = f(extras.t) if yt isa Number @@ -130,7 +130,7 @@ function DI.value_derivative_and_second_derivative(f, backend::AutoGTPSA, x, ext end end -function DI.value_derivative_and_second_derivative!(f, der, der2, backend::AutoGTPSA, x, extras::GTPSADerivativeExtras) +function DI.value_derivative_and_second_derivative!(f, der, der2, ::AutoGTPSA, x, extras::GTPSASecondDerivativeExtras) extras.t[0] = x yt = f(extras.t) if yt isa Number # This should never be reached @@ -173,7 +173,7 @@ function DI.prepare_gradient(f, backend::AutoGTPSA{D}, x) where {D} return GTPSAGradientExtras(v) end -function DI.gradient(f, backend::AutoGTPSA, x, extras::GTPSAGradientExtras) +function DI.gradient(f, ::AutoGTPSA, x, extras::GTPSAGradientExtras) foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part yt = f(extras.v) grad = similar(x, GTPSA.numtype(eltype(yt))) @@ -181,14 +181,14 @@ function DI.gradient(f, backend::AutoGTPSA, x, extras::GTPSAGradientExtras) return grad end -function DI.gradient!(f, grad, backend::AutoGTPSA, x, extras::GTPSAGradientExtras) +function DI.gradient!(f, grad, ::AutoGTPSA, x, extras::GTPSAGradientExtras) foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part yt = f(extras.v) GTPSA.gradient!(grad, yt, include_params=true) return grad end -function DI.value_and_gradient(f, backend::AutoGTPSA, x, extras::GTPSAGradientExtras) +function DI.value_and_gradient(f, ::AutoGTPSA, x, extras::GTPSAGradientExtras) foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part yt = f(extras.v) grad = similar(x, GTPSA.numtype(eltype(yt))) @@ -197,7 +197,7 @@ function DI.value_and_gradient(f, backend::AutoGTPSA, x, extras::GTPSAGradientEx return y, grad end -function DI.value_and_gradient!(f, grad, backend::AutoGTPSA, x, extras::GTPSAGradientExtras) +function DI.value_and_gradient!(f, grad, ::AutoGTPSA, x, extras::GTPSAGradientExtras) foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part yt = f(extras.v) GTPSA.gradient!(grad, yt, include_params=true) @@ -210,137 +210,125 @@ struct GTPSAJacobianExtras <: JacobianExtras v::Vector{TPS} end +function DI.prepare_jacobian(f, backend::AutoGTPSA{D}, x) where {D} + if D != Nothing + d = backend.descriptor + nn = GTPSA.numnn(d) + else + nn = length(x) + d = Descriptor(nn,1) + end + v = similar(x, TPS) + # v 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(v) + v[i] = TPS(use=d) + v[i][j] = 1 + j += 1 + end + return GTPSAJacobianExtras(v) +end -# --- pushforward --- +function DI.jacobian(f, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) + foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + jac = similar(x, GTPSA.numtype(eltype(yt)), (length(yt),length(x))) + GTPSA.jacobian!(jac, yt, include_params=true) + return jac +end -""" - GTPSAPushforwardExtras{T} <: PushforwardExtras +function DI.jacobian!(f, jac, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) + foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + GTPSA.jacobian!(jac, yt, include_params=true) + return jac +end -This struct contains pre-allocated TPS(s) corresponding to `x` with seed -`dx` for the derivative. If taking derivative of single-variable function, -then T is a TPS, else T is some vector of TPSs. -""" -struct GTPSAPushforwardExtras{T} <: PushforwardExtras - v::T +function DI.value_and_jacobian(f, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) + foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + jac = similar(x, GTPSA.numtype(eltype(yt)), (length(yt),length(x))) + GTPSA.jacobian!(jac, yt, include_params=true) + y = map(t->t[0], yt) + return y, jac end -function DI.prepare_pushforward(f, backend::AutoGTPSA, x, dx) - if x isa Number - d = Descriptor(1,1) - t = TPS(use=d) - t[0] = x - t[1] = dx - return GTPSAPushforwardExtras(t) - else - NV = length(x) - d = Descriptor(NV,1) # only first order - v = similar(x, TPS) - - # v 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(v) - v[i] = TPS(use=d) # allocate - v[i][0] = x[i] - v[i][j] = dx[i] - j += 1 - end - return GTPSAPushforwardExtras(v) - end +function DI.value_and_jacobian!(f, jac, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) + foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + GTPSA.jacobian!(jac, yt, include_params=true) + y = map(t->t[0], yt) + return y, jac end -function DI.pushforward(f, ::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras{T}) where {T} - w = f(extras.v) - if w isa Number - if x isa Number # Single variable scalar function derivative - return w[1] - else # Multivariable scalar function derivative (gradient) - dy = similar(w, GTPSA.numtype(eltype(w))) - return GTPSA.gradient!(dy, w, include_params=true) - end - else - if x isa Number # Single variable vector function derivative - - else # Multivariable vector function derivative (Jacobian) - dy = similar(w, GTPSA.numtype(eltype(w)), Size(length(w),length(x))) - end - dy = similar(extras.v, GTPSA.numtype(eltype(w))) - # once again, derivatives in TPSs must be - # 1-based linear indexing - j = 1 - for i in eachindex(dy) - dy[i] = w[i][j] - j += 1 - end - return dy - - end +# --- hessian --- +struct GTPSAHessianExtras <: HessianExtras + v::Vector{TPS} end -function DI.value_and_pushforward(f, ::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras{T}) where {T} - w = f(extras.v) - - if w isa Number - return w[0], w[1] - else - dy = similar(extras.v, GTPSA.numtype(eltype(w))) - # once again, derivatives in TPSs must be - # 1-based linear indexing - j = 1 - for i in eachindex(dy) - dy[i] = w[i][j] - j += 1 - end - return map(t->t[0], w), dy +function DI.prepare_hessian(f, backend::AutoGTPSA{D}, x) where {D} + if D != Nothing + d = backend.descriptor + nn = GTPSA.numnn(d) + else + nn = length(x) + d = Descriptor(nn,2) end + v = similar(x, TPS) + + # v 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(v) + v[i] = TPS(use=d) + v[i][j] = 1 + j += 1 + end + + return GTPSAHessianExtras(v) end -function DI.pushforward!(f, dy, ::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras{T}) where {T} - w = f(extras.v) +function DI.hessian(f, ::AutoGTPSA, x, extras::GTPSAHessianExtras) + foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + hess = similar(x, GTPSA.numtype(eltype(yt)), (length(x),length(x))) + GTPSA.hessian!(hess, yt, include_params=true) + return hess +end - if w isa Number - return w[1] # this should never be reached when this ! function is called - else - # once again, derivatives in TPSs must be - # 1-based linear indexing - j = 1 - for i in eachindex(dy) - dy[i] = w[i][j] - j += 1 - end - return dy - end +function DI.hessian!(f, hess, ::AutoGTPSA, x, extras::GTPSAHessianExtras) + foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + GTPSA.hessian!(hess, yt, include_params=true) + return hess end -function DI.value_and_pushforward!(f, dy, ::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras{T}) where {T} - w = f(extras.v) +function DI.value_gradient_and_hessian(f, ::AutoGTPSA, x, extras::GTPSAHessianExtras) + foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + y = map(t->t[0], yt) + grad = similar(x, GTPSA.numtype(eltype(yt))) + GTPSA.gradient!(grad, yt, include_params=true) + hess = similar(x, GTPSA.numtype(eltype(yt)), (length(x),length(x))) + GTPSA.hessian!(hess, yt, include_params=true) + return y, grad, hess +end - if w isa Number - return w[0], w[1] # this should never be reached when this ! function is called - else - # once again, derivatives in TPSs must be - # 1-based linear indexing - j = 1 - for i in eachindex(dy) - dy[i] = w[i][j] - j += 1 - end - return map(t->t[0], w), dy - end +function DI.value_gradient_and_hessian!(f, grad, hess, ::AutoGTPSA, x, extras::GTPSAHessianExtras) + foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + y = map(t->t[0], yt) + GTPSA.gradient!(grad, yt, include_params=true) + GTPSA.hessian!(hess, yt, include_params=true) + return y, grad, hess end -# --- derivative --- -""" - GTPSADerivativeExtras{T} <: PushforwardExtras - -This struct contains pre-allocated TPS(s) corresponding to `x` with seed -`dx` for the derivative. If taking derivative of single-variable function, -then T is a TPS, else T is some vector of TPSs. -""" -struct GTPSAPushforwardExtras{T} <: PushforwardExtras - v::T -end \ No newline at end of file + + + From 6e69e6c0b7d904c0a045a4a3fae6ba72f4b34b9c Mon Sep 17 00:00:00 2001 From: mattsignorelli Date: Mon, 24 Jun 2024 11:45:27 -0400 Subject: [PATCH 04/44] add GTPSA to Project.toml, move AutoGTPSA to ADTypes --- .github/workflows/Test.yml | 2 ++ DifferentiationInterface/Project.toml | 5 ++- .../DifferentiationInterfaceGTPSAExt.jl | 32 ++----------------- .../onearg.jl | 6 ++++ .../src/DifferentiationInterface.jl | 2 ++ .../test/Single/GTPSA/test.jl | 13 ++++++++ DifferentiationInterface/test/runtests.jl | 3 +- 7 files changed, 32 insertions(+), 31 deletions(-) create mode 100644 DifferentiationInterface/test/Single/GTPSA/test.jl diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml index cc8b93911..5f72b84ee 100644 --- a/.github/workflows/Test.yml +++ b/.github/workflows/Test.yml @@ -65,6 +65,8 @@ jobs: - version: '1.6' group: Single/FastDifferentiation - version: '1.6' + group: Single/GTPSA + - version: '1.9' group: Single/PolyesterForwardDiff - version: '1.6' group: Single/Symbolics diff --git a/DifferentiationInterface/Project.toml b/DifferentiationInterface/Project.toml index a1c41eb4c..d4b76057a 100644 --- a/DifferentiationInterface/Project.toml +++ b/DifferentiationInterface/Project.toml @@ -8,7 +8,6 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" Compat = "34da2185-b29b-5c13-b0c7-acf172513d20" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" -GTPSA = "b27dd330-f138-47c5-815b-40db9dd9b6e8" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" @@ -22,6 +21,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" PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7" @@ -37,6 +37,7 @@ DifferentiationInterfaceFastDifferentiationExt = "FastDifferentiation" DifferentiationInterfaceFiniteDiffExt = "FiniteDiff" DifferentiationInterfaceFiniteDifferencesExt = "FiniteDifferences" DifferentiationInterfaceForwardDiffExt = "ForwardDiff" +DifferentiationInterfaceGTPSAExt = "GTPSA" DifferentiationInterfacePolyesterForwardDiffExt = "PolyesterForwardDiff" DifferentiationInterfaceReverseDiffExt = "ReverseDiff" DifferentiationInterfaceSymbolicsExt = "Symbolics" @@ -56,6 +57,7 @@ FillArrays = "1.7.0" FiniteDiff = "2.23.1" FiniteDifferences = "0.12.31" ForwardDiff = "0.10.36" +GTPSA = "0.6.1" LinearAlgebra = "<0.0.1,1" PackageExtensionCompat = "1.0.2" PolyesterForwardDiff = "0.1.1" @@ -80,6 +82,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" JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl index 0bda8a3e7..afd895879 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl @@ -1,7 +1,6 @@ module DifferentiationInterfaceGTPSAExt -using ADTypes: AbstractADType, ForwardMode -import ADTypes: mode +using ADTypes: AutoGTPSA import DifferentiationInterface as DI using DifferentiationInterface: DerivativeExtras, @@ -9,35 +8,10 @@ using DifferentiationInterface: GradientExtras, JacobianExtras, HessianExtras, - PushforwardExtras + PushforwardExtras, + NoPushforwardExtras using GTPSA - -""" - AutoGTPSA{D} - -Struct used to select the [GTPSA.jl](https://github.com/bmad-sim/GTPSA.jl) backend for automatic differentiation. - -# Constructors - - AutoGTPSA(; descriptor=nothing) - -# Fields - - - `descriptor::D`: can be either - - + a GTPSA `Descriptor` specifying the number of variables/parameters, parameter - order, individual variable/parameter truncation orders, and maximum order. See - [here](https://bmad-sim.github.io/GTPSA.jl/stable/man/c_descriptor/) for more details. - + `nothing` to automatically use a `Descriptor` given the context - -""" -Base.@kwdef struct AutoGTPSA{D} <: AbstractADType - descriptor::D = nothing -end - -mode(::AutoGTPSA) = ForwardMode() - DI.check_available(::AutoGTPSA) = true include("onearg.jl") diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl index 4a92af488..7d5f09184 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl @@ -1,3 +1,9 @@ +# --- pushforward --- +function DI.pushforward(f, backend::AutoGTPSA{D}, x, dx, extras::NoPushforwardExtras) + println("hi!") +end + + # --- derivative --- struct GTPSADerivativeExtras <: DerivativeExtras t::TPS diff --git a/DifferentiationInterface/src/DifferentiationInterface.jl b/DifferentiationInterface/src/DifferentiationInterface.jl index 504171bbe..d2d19bc45 100644 --- a/DifferentiationInterface/src/DifferentiationInterface.jl +++ b/DifferentiationInterface/src/DifferentiationInterface.jl @@ -22,6 +22,7 @@ using ADTypes: AutoFiniteDiff, AutoFiniteDifferences, AutoForwardDiff, + AutoGTPSA, AutoPolyesterForwardDiff, AutoReverseDiff, AutoSymbolics, @@ -122,6 +123,7 @@ export AutoFastDifferentiation export AutoFiniteDiff export AutoFiniteDifferences export AutoForwardDiff +export AutoGTPSA export AutoPolyesterForwardDiff export AutoReverseDiff export AutoSymbolics diff --git a/DifferentiationInterface/test/Single/GTPSA/test.jl b/DifferentiationInterface/test/Single/GTPSA/test.jl new file mode 100644 index 000000000..4a25be82b --- /dev/null +++ b/DifferentiationInterface/test/Single/GTPSA/test.jl @@ -0,0 +1,13 @@ +using DifferentiationInterface, DifferentiationInterfaceTest +using GTPSA: GTPSA +using Test + +for backend in [AutoGTPSA()] + @test check_available(backend) + #@test check_twoarg(backend) + #@test check_hessian(backend) +end + +test_differentiation( + AutoGTPSA(); excluded=[SecondDerivativeScenario, HVPScenario], logging=LOGGING +); diff --git a/DifferentiationInterface/test/runtests.jl b/DifferentiationInterface/test/runtests.jl index a4910e175..18a723e33 100644 --- a/DifferentiationInterface/test/runtests.jl +++ b/DifferentiationInterface/test/runtests.jl @@ -11,7 +11,7 @@ end LOGGING = get(ENV, "CI", "false") == "false" -GROUP = get(ENV, "JULIA_DI_TEST_GROUP", "All") +GROUP = get(ENV, "JULIA_DI_TEST_GROUP", "Single/GTPSA") ALL_BACKENDS = [ "Diffractor", @@ -20,6 +20,7 @@ ALL_BACKENDS = [ "FiniteDifferences", "FastDifferentiation", "ForwardDiff", + "GTPSA", "PolyesterForwardDiff", "ReverseDiff", "Symbolics", From a39fd4ea8abab343036298bcc07fe9227ab6484d Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Tue, 25 Jun 2024 08:18:00 +0200 Subject: [PATCH 05/44] Fix tests --- .github/workflows/Test.yml | 3 +- .../DifferentiationInterfaceGTPSAExt.jl | 3 +- .../onearg.jl | 500 +++++++++--------- .../src/DifferentiationInterface.jl | 6 +- .../src/misc/temp_backends.jl | 27 + .../test/Single/GTPSA/test.jl | 4 +- 6 files changed, 280 insertions(+), 263 deletions(-) create mode 100644 DifferentiationInterface/src/misc/temp_backends.jl diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml index 5f72b84ee..ab6aca229 100644 --- a/.github/workflows/Test.yml +++ b/.github/workflows/Test.yml @@ -39,6 +39,7 @@ jobs: - Single/FiniteDiff - Single/FiniteDifferences - Single/ForwardDiff + - Single/GTPSA - Single/PolyesterForwardDiff - Single/ReverseDiff - Single/Symbolics @@ -66,7 +67,7 @@ jobs: group: Single/FastDifferentiation - version: '1.6' group: Single/GTPSA - - version: '1.9' + - version: '1.6' group: Single/PolyesterForwardDiff - version: '1.6' group: Single/Symbolics diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl index afd895879..689c926a1 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl @@ -1,7 +1,7 @@ module DifferentiationInterfaceGTPSAExt -using ADTypes: AutoGTPSA import DifferentiationInterface as DI +using DifferentiationInterface: AutoGTPSA # TODO: replace with ADTypes using DifferentiationInterface: DerivativeExtras, SecondDerivativeExtras, @@ -13,6 +13,7 @@ using DifferentiationInterface: using GTPSA DI.check_available(::AutoGTPSA) = true +DI.twoarg_support(::AutoGTPSA) = DI.TwoArgNotSupported() include("onearg.jl") diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl index 7d5f09184..3ef85dcf1 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl @@ -1,340 +1,326 @@ -# --- pushforward --- -function DI.pushforward(f, backend::AutoGTPSA{D}, x, dx, extras::NoPushforwardExtras) - println("hi!") -end +## Pushforward +## Derivative -# --- derivative --- struct GTPSADerivativeExtras <: DerivativeExtras - t::TPS + t::TPS end function DI.prepare_derivative(f, backend::AutoGTPSA{D}, x) where {D} - if D != Nothing - d = backend.descriptor - else - d = Descriptor(1,1) - end - t = TPS(use=d) - t[1] = 1 - return GTPSADerivativeExtras(t) + if D != Nothing + d = backend.descriptor + else + d = Descriptor(1, 1) + end + t = TPS(; use=d) + t[1] = 1 + return GTPSADerivativeExtras(t) end function DI.derivative(f, ::AutoGTPSA, x, extras::GTPSADerivativeExtras) - extras.t[0] = x - yt = f(extras.t) - if yt isa Number - return yt[1] - else - der = similar(yt, GTPSA.numtype(eltype(yt))) - for i in eachindex(yt) - der[i] = yt[i][1] + extras.t[0] = x + yt = f(extras.t) + if yt isa Number + return yt[1] + else + der = similar(yt, GTPSA.numtype(eltype(yt))) + for i in eachindex(yt) + der[i] = yt[i][1] + end + return der end - return der - end end function DI.derivative!(f, der, ::AutoGTPSA, x, extras::GTPSADerivativeExtras) - extras.t[0] = x - yt = f(extras.t) - if yt isa Number # This should never be reached - return yt[1] - else + extras.t[0] = x + yt = f(extras.t) for i in eachindex(yt) - der[i] = yt[i][1] + der[i] = yt[i][1] end return der - end end function DI.value_and_derivative(f, ::AutoGTPSA, x, extras::GTPSADerivativeExtras) - extras.t[0] = x - yt = f(extras.t) - if yt isa Number - return yt[0], yt[1] - else - y = map(t->t[0], yt) - der = similar(yt, GTPSA.numtype(eltype(yt))) - for i in eachindex(yt) - der[i] = yt[i][1] + extras.t[0] = x + yt = f(extras.t) + if yt isa Number + return yt[0], yt[1] + else + y = map(t -> t[0], yt) + der = similar(yt, GTPSA.numtype(eltype(yt))) + for i in eachindex(yt) + der[i] = yt[i][1] + end + return y, der end - return y, der - end end function DI.value_and_derivative!(f, der, ::AutoGTPSA, x, extras::GTPSADerivativeExtras) - extras.t[0] = x - yt = f(extras.t) - if yt isa Number # This should never be reached - return yt[0], yt[1] - else - y = map(t->t[0], yt) + extras.t[0] = x + yt = f(extras.t) + y = map(t -> t[0], yt) for i in eachindex(yt) - der[i] = yt[i][1] + der[i] = yt[i][1] end return y, der - end -end - -# --- second_derivative --- -struct GTPSASecondDerivativeExtras <: SecondDerivativeExtras - t::TPS end -function DI.prepare_second_derivative(f, backend::AutoGTPSA{D}, x) where {D} - if D != Nothing - d = backend.descriptor - else - d = Descriptor(1,2) - end - t = TPS(use=d) - t[1] = 1 - return GTPSASecondDerivativeExtras(t) -end +## Gradient -function DI.second_derivative(f, ::AutoGTPSA, x, extras::GTPSASecondDerivativeExtras) - extras.t[0] = x - yt = f(extras.t) - if yt isa Number - return yt[2] - else - der2 = similar(yt, GTPSA.numtype(eltype(yt))) - for i in eachindex(yt) - der2[i] = yt[i][2] - end - return der2 - end -end - -function DI.second_derivative!(f, der2, ::AutoGTPSA, x, extras::GTPSASecondDerivativeExtras) - extras.t[0] = x - yt = f(extras.t) - if yt isa Number # This should never be reached - return yt[2] - else - for i in eachindex(yt) - der2[i] = yt[i][2] - end - return der2 - end +struct GTPSAGradientExtras <: GradientExtras + v::Vector{TPS} end -function DI.value_derivative_and_second_derivative(f, ::AutoGTPSA, x, extras::GTPSASecondDerivativeExtras) - extras.t[0] = x - yt = f(extras.t) - if yt isa Number - return yt[0], yt[1], yt[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][2] +function DI.prepare_gradient(f, backend::AutoGTPSA{D}, x) where {D} + if D != Nothing + d = backend.descriptor + nn = GTPSA.numnn(d) + else + nn = length(x) + d = Descriptor(nn, 1) end - return y, der, der2 - end -end - -function DI.value_derivative_and_second_derivative!(f, der, der2, ::AutoGTPSA, x, extras::GTPSASecondDerivativeExtras) - extras.t[0] = x - yt = f(extras.t) - if yt isa Number # This should never be reached - return yt[0], yt[1], yt[2] - else - y = map(t->t[0], yt) - for i in eachindex(yt) - der[i] = yt[i][1] - der2[i] = yt[i][2] + v = similar(x, TPS) + + # v 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(v) + v[i] = TPS(; use=d) + v[i][j] = 1 + j += 1 end - return y, der, der2 - end -end -# --- gradient --- -struct GTPSAGradientExtras <: GradientExtras - v::Vector{TPS} -end - -function DI.prepare_gradient(f, backend::AutoGTPSA{D}, x) where {D} - if D != Nothing - d = backend.descriptor - nn = GTPSA.numnn(d) - else - nn = length(x) - d = Descriptor(nn,1) - end - v = similar(x, TPS) - - # v 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(v) - v[i] = TPS(use=d) - v[i][j] = 1 - j += 1 - end - - return GTPSAGradientExtras(v) + return GTPSAGradientExtras(v) end function DI.gradient(f, ::AutoGTPSA, x, extras::GTPSAGradientExtras) - foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part - yt = f(extras.v) - grad = similar(x, GTPSA.numtype(eltype(yt))) - GTPSA.gradient!(grad, yt, include_params=true) - return grad + foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + grad = similar(x, GTPSA.numtype(eltype(yt))) + GTPSA.gradient!(grad, yt; include_params=true) + return grad end function DI.gradient!(f, grad, ::AutoGTPSA, x, extras::GTPSAGradientExtras) - foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part - yt = f(extras.v) - GTPSA.gradient!(grad, yt, include_params=true) - return grad + foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + GTPSA.gradient!(grad, yt; include_params=true) + return grad end function DI.value_and_gradient(f, ::AutoGTPSA, x, extras::GTPSAGradientExtras) - foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part - yt = f(extras.v) - grad = similar(x, GTPSA.numtype(eltype(yt))) - GTPSA.gradient!(grad, yt, include_params=true) - y = map(t->t[0], yt) - return y, grad + foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + grad = similar(x, GTPSA.numtype(eltype(yt))) + GTPSA.gradient!(grad, yt; include_params=true) + y = map(t -> t[0], yt) + return y, grad end function DI.value_and_gradient!(f, grad, ::AutoGTPSA, x, extras::GTPSAGradientExtras) - foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part - yt = f(extras.v) - GTPSA.gradient!(grad, yt, include_params=true) - y = map(t->t[0], yt) - return y, grad + foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + GTPSA.gradient!(grad, yt; include_params=true) + y = map(t -> t[0], yt) + return y, grad end -# --- jacobian --- +## Jacobian + struct GTPSAJacobianExtras <: JacobianExtras - v::Vector{TPS} + v::Vector{TPS} end function DI.prepare_jacobian(f, backend::AutoGTPSA{D}, x) where {D} - if D != Nothing - d = backend.descriptor - nn = GTPSA.numnn(d) - else - nn = length(x) - d = Descriptor(nn,1) - end - v = similar(x, TPS) - - # v 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(v) - v[i] = TPS(use=d) - v[i][j] = 1 - j += 1 - end - - return GTPSAJacobianExtras(v) + if D != Nothing + d = backend.descriptor + nn = GTPSA.numnn(d) + else + nn = length(x) + d = Descriptor(nn, 1) + end + v = similar(x, TPS) + + # v 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(v) + v[i] = TPS(; use=d) + v[i][j] = 1 + j += 1 + end + + return GTPSAJacobianExtras(v) end function DI.jacobian(f, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) - foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part - yt = f(extras.v) - jac = similar(x, GTPSA.numtype(eltype(yt)), (length(yt),length(x))) - GTPSA.jacobian!(jac, yt, include_params=true) - return jac + foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + jac = similar(x, GTPSA.numtype(eltype(yt)), (length(yt), length(x))) + GTPSA.jacobian!(jac, yt; include_params=true) + return jac end function DI.jacobian!(f, jac, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) - foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part - yt = f(extras.v) - GTPSA.jacobian!(jac, yt, include_params=true) - return jac + foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + GTPSA.jacobian!(jac, yt; include_params=true) + return jac end function DI.value_and_jacobian(f, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) - foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part - yt = f(extras.v) - jac = similar(x, GTPSA.numtype(eltype(yt)), (length(yt),length(x))) - GTPSA.jacobian!(jac, yt, include_params=true) - y = map(t->t[0], yt) - return y, jac + foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + jac = similar(x, GTPSA.numtype(eltype(yt)), (length(yt), length(x))) + GTPSA.jacobian!(jac, yt; include_params=true) + y = map(t -> t[0], yt) + return y, jac end function DI.value_and_jacobian!(f, jac, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) - foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part - yt = f(extras.v) - GTPSA.jacobian!(jac, yt, include_params=true) - y = map(t->t[0], yt) - return y, jac + foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + GTPSA.jacobian!(jac, yt; include_params=true) + y = map(t -> t[0], yt) + return y, jac +end + +## Second derivative + +struct GTPSASecondDerivativeExtras <: SecondDerivativeExtras + t::TPS +end + +function DI.prepare_second_derivative(f, backend::AutoGTPSA{D}, x) where {D} + if D != Nothing + d = backend.descriptor + else + d = Descriptor(1, 2) + end + t = TPS(; use=d) + t[1] = 1 + return GTPSASecondDerivativeExtras(t) end +function DI.second_derivative(f, ::AutoGTPSA, x, extras::GTPSASecondDerivativeExtras) + extras.t[0] = x + yt = f(extras.t) + if yt isa Number + return yt[2] + else + der2 = similar(yt, GTPSA.numtype(eltype(yt))) + for i in eachindex(yt) + der2[i] = yt[i][2] + end + return der2 + end +end + +function DI.second_derivative!(f, der2, ::AutoGTPSA, x, extras::GTPSASecondDerivativeExtras) + extras.t[0] = x + yt = f(extras.t) + for i in eachindex(yt) + der2[i] = yt[i][2] + end + return der2 +end + +function DI.value_derivative_and_second_derivative( + f, ::AutoGTPSA, x, extras::GTPSASecondDerivativeExtras +) + extras.t[0] = x + yt = f(extras.t) + if yt isa Number + return yt[0], yt[1], yt[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][2] + end + return y, der, der2 + end +end + +function DI.value_derivative_and_second_derivative!( + f, der, der2, ::AutoGTPSA, x, extras::GTPSASecondDerivativeExtras +) + extras.t[0] = x + yt = f(extras.t) + y = map(t -> t[0], yt) + for i in eachindex(yt) + der[i] = yt[i][1] + der2[i] = yt[i][2] + end + return y, der, der2 +end + +## Hessian -# --- hessian --- struct GTPSAHessianExtras <: HessianExtras - v::Vector{TPS} + v::Vector{TPS} end function DI.prepare_hessian(f, backend::AutoGTPSA{D}, x) where {D} - if D != Nothing - d = backend.descriptor - nn = GTPSA.numnn(d) - else - nn = length(x) - d = Descriptor(nn,2) - end - v = similar(x, TPS) - - # v 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(v) - v[i] = TPS(use=d) - v[i][j] = 1 - j += 1 - end - - return GTPSAHessianExtras(v) + if D != Nothing + d = backend.descriptor + nn = GTPSA.numnn(d) + else + nn = length(x) + d = Descriptor(nn, 2) + end + v = similar(x, TPS) + + # v 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(v) + v[i] = TPS(; use=d) + v[i][j] = 1 + j += 1 + end + + return GTPSAHessianExtras(v) end function DI.hessian(f, ::AutoGTPSA, x, extras::GTPSAHessianExtras) - foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part - yt = f(extras.v) - hess = similar(x, GTPSA.numtype(eltype(yt)), (length(x),length(x))) - GTPSA.hessian!(hess, yt, include_params=true) - return hess + foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + hess = similar(x, GTPSA.numtype(eltype(yt)), (length(x), length(x))) + GTPSA.hessian!(hess, yt; include_params=true) + return hess end function DI.hessian!(f, hess, ::AutoGTPSA, x, extras::GTPSAHessianExtras) - foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part - yt = f(extras.v) - GTPSA.hessian!(hess, yt, include_params=true) - return hess + foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + GTPSA.hessian!(hess, yt; include_params=true) + return hess end function DI.value_gradient_and_hessian(f, ::AutoGTPSA, x, extras::GTPSAHessianExtras) - foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part - yt = f(extras.v) - y = map(t->t[0], yt) - grad = similar(x, GTPSA.numtype(eltype(yt))) - GTPSA.gradient!(grad, yt, include_params=true) - hess = similar(x, GTPSA.numtype(eltype(yt)), (length(x),length(x))) - GTPSA.hessian!(hess, yt, include_params=true) - return y, grad, hess + foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + y = map(t -> t[0], yt) + grad = similar(x, GTPSA.numtype(eltype(yt))) + GTPSA.gradient!(grad, yt; include_params=true) + hess = similar(x, GTPSA.numtype(eltype(yt)), (length(x), length(x))) + GTPSA.hessian!(hess, yt; include_params=true) + return y, grad, hess end -function DI.value_gradient_and_hessian!(f, grad, hess, ::AutoGTPSA, x, extras::GTPSAHessianExtras) - foreach((t,xi)->t[0]=xi, extras.v, x) # Set the scalar part - yt = f(extras.v) - y = map(t->t[0], yt) - GTPSA.gradient!(grad, yt, include_params=true) - GTPSA.hessian!(hess, yt, include_params=true) - return y, grad, hess +function DI.value_gradient_and_hessian!( + f, grad, hess, ::AutoGTPSA, x, extras::GTPSAHessianExtras +) + foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part + yt = f(extras.v) + y = map(t -> t[0], yt) + GTPSA.gradient!(grad, yt; include_params=true) + GTPSA.hessian!(hess, yt; include_params=true) + return y, grad, hess end - - - - diff --git a/DifferentiationInterface/src/DifferentiationInterface.jl b/DifferentiationInterface/src/DifferentiationInterface.jl index d2d19bc45..7653514f9 100644 --- a/DifferentiationInterface/src/DifferentiationInterface.jl +++ b/DifferentiationInterface/src/DifferentiationInterface.jl @@ -22,7 +22,6 @@ using ADTypes: AutoFiniteDiff, AutoFiniteDifferences, AutoForwardDiff, - AutoGTPSA, AutoPolyesterForwardDiff, AutoReverseDiff, AutoSymbolics, @@ -74,6 +73,7 @@ include("sparse/hessian.jl") include("misc/differentiate_with.jl") include("misc/sparsity_detector.jl") +include("misc/temp_backends.jl") function __init__() @require_extensions @@ -133,6 +133,10 @@ export AutoZygote export AutoSparse +## Temporary export + +export AutoGTPSA + ## Re-exported from SparseMatrixColorings export GreedyColoringAlgorithm diff --git a/DifferentiationInterface/src/misc/temp_backends.jl b/DifferentiationInterface/src/misc/temp_backends.jl new file mode 100644 index 000000000..2763db91e --- /dev/null +++ b/DifferentiationInterface/src/misc/temp_backends.jl @@ -0,0 +1,27 @@ +""" + AutoGTPSA{D} + +Struct used to select the [GTPSA.jl](https://github.com/bmad-sim/GTPSA.jl) backend for automatic differentiation. + +!!! warning + This type is not part of the public API, it is only used for testing before being moved to ADTypes.jl. + +# Constructors + + AutoGTPSA(; descriptor=nothing) + +# Fields + + - `descriptor::D`: can be either + + + a GTPSA `Descriptor` specifying the number of variables/parameters, parameter + order, individual variable/parameter truncation orders, and maximum order. See + the [GTPSA.jl documentation](https://bmad-sim.github.io/GTPSA.jl/stable/man/c_descriptor/) for more details. + + `nothing` to automatically use a `Descriptor` given the context. + +""" +Base.@kwdef struct AutoGTPSA{D} <: AbstractADType + descriptor::D = nothing +end + +ADTypes.mode(::AutoGTPSA) = ForwardMode() diff --git a/DifferentiationInterface/test/Single/GTPSA/test.jl b/DifferentiationInterface/test/Single/GTPSA/test.jl index 4a25be82b..3f1c5febf 100644 --- a/DifferentiationInterface/test/Single/GTPSA/test.jl +++ b/DifferentiationInterface/test/Single/GTPSA/test.jl @@ -8,6 +8,4 @@ for backend in [AutoGTPSA()] #@test check_hessian(backend) end -test_differentiation( - AutoGTPSA(); excluded=[SecondDerivativeScenario, HVPScenario], logging=LOGGING -); +test_differentiation(AutoGTPSA(); logging=LOGGING); From 79d6ee172276a4b089bff4dfe9e63382046cb391 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Tue, 25 Jun 2024 11:06:24 +0200 Subject: [PATCH 06/44] Lighter workflow --- .github/workflows/Test.yml | 134 ++++++++++++++++++------------------- 1 file changed, 67 insertions(+), 67 deletions(-) diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml index 311b9ba8a..d4621b74f 100644 --- a/.github/workflows/Test.yml +++ b/.github/workflows/Test.yml @@ -28,28 +28,28 @@ jobs: matrix: version: - '1' - - '1.6' + # - '1.6' - '~1.11.0-0' group: - Formalities - - Internals - - Single/Diffractor - - Single/Enzyme - - Single/FastDifferentiation - - Single/FiniteDiff - - Single/FiniteDifferences - - Single/ForwardDiff + # - Internals + # - Single/Diffractor + # - Single/Enzyme + # - Single/FastDifferentiation + # - Single/FiniteDiff + # - Single/FiniteDifferences + # - Single/ForwardDiff - Single/GTPSA - - Single/PolyesterForwardDiff - - Single/ReverseDiff - - Single/Symbolics - - Single/Tapir - - Single/Tracker - - Single/Zygote - - Double/ChainRulesCore-Zygote - - Double/Enzyme-ForwardDiff - - Double/ForwardDiff-Zygote - - Double/ReverseDiff-Zygote + # - Single/PolyesterForwardDiff + # - Single/ReverseDiff + # - Single/Symbolics + # - Single/Tapir + # - Single/Tracker + # - Single/Zygote + # - Double/ChainRulesCore-Zygote + # - Double/Enzyme-ForwardDiff + # - Double/ForwardDiff-Zygote + # - Double/ReverseDiff-Zygote exclude: - version: '1.6' group: Formalities @@ -108,53 +108,53 @@ jobs: token: ${{ secrets.CODECOV_TOKEN }} fail_ci_if_error: true - test-DIT: - name: DIT (${{ matrix.group }}) - ${{ matrix.version }} - runs-on: ubuntu-latest - timeout-minutes: 60 - permissions: # needed to allow julia-actions/cache to proactively delete old caches that it has created - actions: write - contents: read - strategy: - fail-fast: false - matrix: - version: - - '1' - - '1.6' - - '~1.11.0-0' - group: - - Formalities - - Zero - - ForwardDiff - - Zygote - exclude: - - version: '1.6' - group: Formalities + # test-DIT: + # name: DIT (${{ matrix.group }}) - ${{ matrix.version }} + # runs-on: ubuntu-latest + # timeout-minutes: 60 + # permissions: # needed to allow julia-actions/cache to proactively delete old caches that it has created + # actions: write + # contents: read + # strategy: + # fail-fast: false + # matrix: + # version: + # - '1' + # - '1.6' + # - '~1.11.0-0' + # group: + # - Formalities + # - Zero + # - ForwardDiff + # - Zygote + # exclude: + # - version: '1.6' + # group: Formalities - steps: - - uses: actions/checkout@v4 - - uses: julia-actions/setup-julia@v2 - with: - version: ${{ matrix.version }} - arch: x64 - - uses: julia-actions/cache@v2 - - uses: julia-actions/julia-buildpkg@v1 - - name: Dev dependencies (temporary) - run: julia --project='./DifferentiationInterfaceTest' -e ' - using Pkg; - Pkg.Registry.update(); - Pkg.develop(PackageSpec(path="./DifferentiationInterface")); - Pkg.instantiate();' - - uses: julia-actions/julia-runtest@v1 - with: - project: ./DifferentiationInterfaceTest - env: - JULIA_DI_TEST_GROUP: ${{ matrix.group }} - - uses: julia-actions/julia-processcoverage@v1 - with: - directories: ./DifferentiationInterfaceTest/src,./DifferentiationInterfaceTest/ext,./DifferentiationInterfaceTest/test - - uses: codecov/codecov-action@v4 - with: - files: lcov.info - token: ${{ secrets.CODECOV_TOKEN }} - fail_ci_if_error: true \ No newline at end of file + # steps: + # - uses: actions/checkout@v4 + # - uses: julia-actions/setup-julia@v2 + # with: + # version: ${{ matrix.version }} + # arch: x64 + # - uses: julia-actions/cache@v2 + # - uses: julia-actions/julia-buildpkg@v1 + # - name: Dev dependencies (temporary) + # run: julia --project='./DifferentiationInterfaceTest' -e ' + # using Pkg; + # Pkg.Registry.update(); + # Pkg.develop(PackageSpec(path="./DifferentiationInterface")); + # Pkg.instantiate();' + # - uses: julia-actions/julia-runtest@v1 + # with: + # project: ./DifferentiationInterfaceTest + # env: + # JULIA_DI_TEST_GROUP: ${{ matrix.group }} + # - uses: julia-actions/julia-processcoverage@v1 + # with: + # directories: ./DifferentiationInterfaceTest/src,./DifferentiationInterfaceTest/ext,./DifferentiationInterfaceTest/test + # - uses: codecov/codecov-action@v4 + # with: + # files: lcov.info + # token: ${{ secrets.CODECOV_TOKEN }} + # fail_ci_if_error: true \ No newline at end of file From cee4b7ec899884530af4db9a8cdc89eb8b29169e Mon Sep 17 00:00:00 2001 From: mattsignorelli Date: Sun, 7 Jul 2024 15:06:50 -0400 Subject: [PATCH 07/44] pushforward --- DifferentiationInterface/Project.toml | 2 +- .../DifferentiationInterfaceGTPSAExt.jl | 3 +- .../onearg.jl | 152 ++++++++++++++++-- 3 files changed, 144 insertions(+), 13 deletions(-) diff --git a/DifferentiationInterface/Project.toml b/DifferentiationInterface/Project.toml index d4b76057a..b8e13e935 100644 --- a/DifferentiationInterface/Project.toml +++ b/DifferentiationInterface/Project.toml @@ -57,7 +57,7 @@ FillArrays = "1.7.0" FiniteDiff = "2.23.1" FiniteDifferences = "0.12.31" ForwardDiff = "0.10.36" -GTPSA = "0.6.1" +GTPSA = "0.7" LinearAlgebra = "<0.0.1,1" PackageExtensionCompat = "1.0.2" PolyesterForwardDiff = "0.1.1" diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl index 689c926a1..db797a922 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl @@ -8,8 +8,7 @@ using DifferentiationInterface: GradientExtras, JacobianExtras, HessianExtras, - PushforwardExtras, - NoPushforwardExtras + PushforwardExtras using GTPSA DI.check_available(::AutoGTPSA) = true diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl index 3ef85dcf1..095ac3f8c 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl @@ -1,9 +1,141 @@ ## Pushforward +struct GTPSAPushforwardExtras{X} <: PushforwardExtras + xd::X +end + +function DI.prepare_pushforward(f, backend::AutoGTPSA{D}, x, dx) where {D} + if D != Nothing + d = backend.descriptor + else + d = Descriptor(length(x), 1) + end + + if x isa Number + t = TPS(; use=d) + return GTPSAPushforwardExtras(t) + else + v = similar(x, TPS) + + # v and x have same indexing because of similar + for i in eachindex(v) + v[i] = TPS(; use=d) + end + return GTPSAPushforwardExtras(v) + end +end + +function DI.pushforward(f, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras) + if x isa Number + extras.xd[0] = x + extras.xd[1] = dx + else + j = 1 + for i in eachindex(x) + extras.xd[i][0] = x[i] + extras.xd[i][j] = dx[i] + j += 1 + end + end + + yt = f(extras.xd) + if yt isa Number + return yt[1] + else + dy = similar(yt, GTPSA.numtype(eltype(yt))) + j = 1 + for i in eachindex(yt) + dy[i] = yt[i][j] + j += 1 + end + return dy + end +end + +function DI.pushforward!(f, dy, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras) + if x isa Number + extras.xd[0] = x + extras.xd[1] = dx + else + j = 1 + for i in eachindex(x) + extras.xd[i][0] = x[i] + extras.xd[i][j] = dx[i] + j += 1 + end + end + + yt = f(extras.xd) + if yt isa Number + return yt[1] + else + j = 1 + for i in eachindex(yt) + dy[i] = yt[i][j] + j += 1 + end + return dy + end +end + +function DI.value_and_pushforward(f, backend::AutoGTPSA, x, dx, extras:GTPSAPushforwardExtras) + if x isa Number + extras.xd[0] = x + extras.xd[1] = dx + else + j = 1 + for i in eachindex(x) + extras.xd[i][0] = x[i] + extras.xd[i][j] = dx[i] + j += 1 + end + end + + yt = f(extras.xd) + if yt isa Number + return yt[1] + else + dy = similar(yt, GTPSA.numtype(eltype(yt))) + j = 1 + for i in eachindex(yt) + dy[i] = yt[i][j] + j += 1 + end + y = map(t->t[0], yt) + return y, dy + end +end + +function DI.value_and_pushforward!(f, dy, backend::AutoGTPSA, x, dx, extras:GTPSAPushforwardExtras) + if x isa Number + extras.xd[0] = x + extras.xd[1] = dx + else + j = 1 + for i in eachindex(x) + extras.xd[i][0] = x[i] + extras.xd[i][j] = dx[i] + j += 1 + end + end + + yt = f(extras.xd) + if yt isa Number + return yt[1] + else + j = 1 + for i in eachindex(yt) + dy[i] = yt[i][j] + j += 1 + end + y = map(t->t[0], yt) + return y, dy + end +end ## Derivative -struct GTPSADerivativeExtras <: DerivativeExtras - t::TPS +struct GTPSADerivativeExtras{T} <: DerivativeExtras + t::T end function DI.prepare_derivative(f, backend::AutoGTPSA{D}, x) where {D} @@ -67,8 +199,8 @@ end ## Gradient -struct GTPSAGradientExtras <: GradientExtras - v::Vector{TPS} +struct GTPSAGradientExtras{V} <: GradientExtras + v::V end function DI.prepare_gradient(f, backend::AutoGTPSA{D}, x) where {D} @@ -128,8 +260,8 @@ end ## Jacobian -struct GTPSAJacobianExtras <: JacobianExtras - v::Vector{TPS} +struct GTPSAJacobianExtras{V} <: JacobianExtras + v::V end function DI.prepare_jacobian(f, backend::AutoGTPSA{D}, x) where {D} @@ -189,8 +321,8 @@ end ## Second derivative -struct GTPSASecondDerivativeExtras <: SecondDerivativeExtras - t::TPS +struct GTPSASecondDerivativeExtras{T} <: SecondDerivativeExtras + t::T end function DI.prepare_second_derivative(f, backend::AutoGTPSA{D}, x) where {D} @@ -261,8 +393,8 @@ end ## Hessian -struct GTPSAHessianExtras <: HessianExtras - v::Vector{TPS} +struct GTPSAHessianExtras{V} <: HessianExtras + v::V end function DI.prepare_hessian(f, backend::AutoGTPSA{D}, x) where {D} From b362dd974cdef6040f7e857e8bc0cdde9be1e534 Mon Sep 17 00:00:00 2001 From: mattsignorelli Date: Sun, 7 Jul 2024 15:10:16 -0400 Subject: [PATCH 08/44] trigger workflow From 147385871a29c5e061d917f4ebdfe82e682db92c Mon Sep 17 00:00:00 2001 From: mattsignorelli Date: Sun, 7 Jul 2024 17:41:43 -0400 Subject: [PATCH 09/44] fix tests? --- .github/workflows/Test.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml index d4621b74f..1bdbafdfb 100644 --- a/.github/workflows/Test.yml +++ b/.github/workflows/Test.yml @@ -4,6 +4,7 @@ on: push: branches: - main + - addgtpsa tags: ['*'] pull_request: workflow_dispatch: From 20c026a3f9582df86ebea8ae940083484051a4d1 Mon Sep 17 00:00:00 2001 From: mattsignorelli Date: Sun, 7 Jul 2024 17:45:39 -0400 Subject: [PATCH 10/44] fix typo --- .../ext/DifferentiationInterfaceGTPSAExt/onearg.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl index 095ac3f8c..36f274b0a 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl @@ -77,7 +77,7 @@ function DI.pushforward!(f, dy, backend::AutoGTPSA, x, dx, extras::GTPSAPushforw end end -function DI.value_and_pushforward(f, backend::AutoGTPSA, x, dx, extras:GTPSAPushforwardExtras) +function DI.value_and_pushforward(f, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras) if x isa Number extras.xd[0] = x extras.xd[1] = dx @@ -105,7 +105,7 @@ function DI.value_and_pushforward(f, backend::AutoGTPSA, x, dx, extras:GTPSAPush end end -function DI.value_and_pushforward!(f, dy, backend::AutoGTPSA, x, dx, extras:GTPSAPushforwardExtras) +function DI.value_and_pushforward!(f, dy, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras) if x isa Number extras.xd[0] = x extras.xd[1] = dx From 45012c75b0d13a299a748905bea6042bc13c76c8 Mon Sep 17 00:00:00 2001 From: mattsignorelli Date: Tue, 6 Aug 2024 21:00:31 -0400 Subject: [PATCH 11/44] hvp done and two args done --- DifferentiationInterface/Project.toml | 2 +- .../DifferentiationInterfaceGTPSAExt.jl | 7 +- .../onearg.jl | 353 ++++++++++++------ .../twoarg.jl | 293 +++++++++++++++ 4 files changed, 531 insertions(+), 124 deletions(-) create mode 100644 DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl diff --git a/DifferentiationInterface/Project.toml b/DifferentiationInterface/Project.toml index 497f3a7a7..f648703d1 100644 --- a/DifferentiationInterface/Project.toml +++ b/DifferentiationInterface/Project.toml @@ -57,7 +57,7 @@ FillArrays = "1.7.0" FiniteDiff = "2.23.1" FiniteDifferences = "0.12.31" ForwardDiff = "0.10.36" -GTPSA = "0.7" +GTPSA = "1.0.1" LinearAlgebra = "<0.0.1,1" PackageExtensionCompat = "1.0.2" PolyesterForwardDiff = "0.1.1" diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl index db797a922..74432c3e7 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl @@ -8,12 +8,13 @@ using DifferentiationInterface: GradientExtras, JacobianExtras, HessianExtras, - PushforwardExtras + PushforwardExtras, + HVPExtras using GTPSA DI.check_available(::AutoGTPSA) = true -DI.twoarg_support(::AutoGTPSA) = DI.TwoArgNotSupported() include("onearg.jl") +include("twoarg.jl") -end # module +end diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl index 36f274b0a..4531cca21 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl @@ -1,6 +1,6 @@ ## Pushforward struct GTPSAPushforwardExtras{X} <: PushforwardExtras - xd::X + xt::X end function DI.prepare_pushforward(f, backend::AutoGTPSA{D}, x, dx) where {D} @@ -11,37 +11,37 @@ function DI.prepare_pushforward(f, backend::AutoGTPSA{D}, x, dx) where {D} end if x isa Number - t = TPS(; use=d) - return GTPSAPushforwardExtras(t) + xt = TPS{promote_type(typeof(dx),typeof(x),Float64)}(; use=d) + return GTPSAPushforwardExtras(xt) else - v = similar(x, TPS) + xt = similar(x, TPS{promote_type(eltype(dx),eltype(x),Float64)}) - # v and x have same indexing because of similar - for i in eachindex(v) - v[i] = TPS(; use=d) + # xt and x have same indexing because of similar + for i in eachindex(xt) + xt[i] = TPS{promote_type(eltype(dx),eltype(x),Float64)}(; use=d) end - return GTPSAPushforwardExtras(v) + return GTPSAPushforwardExtras(xt) end end function DI.pushforward(f, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras) if x isa Number - extras.xd[0] = x - extras.xd[1] = dx + extras.xt[0] = x + extras.xt[1] = dx else j = 1 for i in eachindex(x) - extras.xd[i][0] = x[i] - extras.xd[i][j] = dx[i] + extras.xt[i][0] = x[i] + extras.xt[i][j] = dx[i] j += 1 end end - yt = f(extras.xd) + yt = f(extras.xt) if yt isa Number return yt[1] else - dy = similar(yt, GTPSA.numtype(eltype(yt))) + dy = similar(yt, eltype(eltype(yt))) j = 1 for i in eachindex(yt) dy[i] = yt[i][j] @@ -53,18 +53,18 @@ end function DI.pushforward!(f, dy, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras) if x isa Number - extras.xd[0] = x - extras.xd[1] = dx + extras.xt[0] = x + extras.xt[1] = dx else j = 1 for i in eachindex(x) - extras.xd[i][0] = x[i] - extras.xd[i][j] = dx[i] + extras.xt[i][0] = x[i] + extras.xt[i][j] = dx[i] j += 1 end end - yt = f(extras.xd) + yt = f(extras.xt) if yt isa Number return yt[1] else @@ -79,22 +79,22 @@ end function DI.value_and_pushforward(f, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras) if x isa Number - extras.xd[0] = x - extras.xd[1] = dx + extras.xt[0] = x + extras.xt[1] = dx else j = 1 for i in eachindex(x) - extras.xd[i][0] = x[i] - extras.xd[i][j] = dx[i] + extras.xt[i][0] = x[i] + extras.xt[i][j] = dx[i] j += 1 end end - yt = f(extras.xd) + yt = f(extras.xt) if yt isa Number return yt[1] else - dy = similar(yt, GTPSA.numtype(eltype(yt))) + dy = similar(yt, eltype(eltype(yt))) j = 1 for i in eachindex(yt) dy[i] = yt[i][j] @@ -107,18 +107,18 @@ end function DI.value_and_pushforward!(f, dy, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras) if x isa Number - extras.xd[0] = x - extras.xd[1] = dx + extras.xt[0] = x + extras.xt[1] = dx else j = 1 for i in eachindex(x) - extras.xd[i][0] = x[i] - extras.xd[i][j] = dx[i] + extras.xt[i][0] = x[i] + extras.xt[i][j] = dx[i] j += 1 end end - yt = f(extras.xd) + yt = f(extras.xt) if yt isa Number return yt[1] else @@ -134,8 +134,8 @@ end ## Derivative -struct GTPSADerivativeExtras{T} <: DerivativeExtras - t::T +struct GTPSADerivativeExtras{X} <: DerivativeExtras + xt::X end function DI.prepare_derivative(f, backend::AutoGTPSA{D}, x) where {D} @@ -144,18 +144,18 @@ function DI.prepare_derivative(f, backend::AutoGTPSA{D}, x) where {D} else d = Descriptor(1, 1) end - t = TPS(; use=d) - t[1] = 1 - return GTPSADerivativeExtras(t) + xt = TPS{promote_type(typeof(x),Float64)}(; use=d) + xt[1] = 1 + return GTPSADerivativeExtras(xt) end function DI.derivative(f, ::AutoGTPSA, x, extras::GTPSADerivativeExtras) - extras.t[0] = x - yt = f(extras.t) + extras.xt[0] = x + yt = f(extras.xt) if yt isa Number return yt[1] else - der = similar(yt, GTPSA.numtype(eltype(yt))) + der = similar(yt, eltype(eltype(yt))) for i in eachindex(yt) der[i] = yt[i][1] end @@ -164,8 +164,8 @@ function DI.derivative(f, ::AutoGTPSA, x, extras::GTPSADerivativeExtras) end function DI.derivative!(f, der, ::AutoGTPSA, x, extras::GTPSADerivativeExtras) - extras.t[0] = x - yt = f(extras.t) + extras.xt[0] = x + yt = f(extras.xt) for i in eachindex(yt) der[i] = yt[i][1] end @@ -173,13 +173,13 @@ function DI.derivative!(f, der, ::AutoGTPSA, x, extras::GTPSADerivativeExtras) end function DI.value_and_derivative(f, ::AutoGTPSA, x, extras::GTPSADerivativeExtras) - extras.t[0] = x - yt = f(extras.t) + extras.xt[0] = x + yt = f(extras.xt) if yt isa Number return yt[0], yt[1] else y = map(t -> t[0], yt) - der = similar(yt, GTPSA.numtype(eltype(yt))) + der = similar(yt, eltype(eltype(yt))) for i in eachindex(yt) der[i] = yt[i][1] end @@ -188,8 +188,8 @@ function DI.value_and_derivative(f, ::AutoGTPSA, x, extras::GTPSADerivativeExtra end function DI.value_and_derivative!(f, der, ::AutoGTPSA, x, extras::GTPSADerivativeExtras) - extras.t[0] = x - yt = f(extras.t) + extras.xt[0] = x + yt = f(extras.xt) y = map(t -> t[0], yt) for i in eachindex(yt) der[i] = yt[i][1] @@ -199,8 +199,8 @@ end ## Gradient -struct GTPSAGradientExtras{V} <: GradientExtras - v::V +struct GTPSAGradientExtras{X} <: GradientExtras + xt::X end function DI.prepare_gradient(f, backend::AutoGTPSA{D}, x) where {D} @@ -211,48 +211,48 @@ function DI.prepare_gradient(f, backend::AutoGTPSA{D}, x) where {D} nn = length(x) d = Descriptor(nn, 1) end - v = similar(x, TPS) + xt = similar(x, TPS{promote_type(eltype(x),Float64)}) - # v and x have same indexing because of similar + # 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(v) - v[i] = TPS(; use=d) - v[i][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 GTPSAGradientExtras(v) + return GTPSAGradientExtras(xt) end function DI.gradient(f, ::AutoGTPSA, x, extras::GTPSAGradientExtras) - foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part - yt = f(extras.v) - grad = similar(x, GTPSA.numtype(eltype(yt))) + foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part + yt = f(extras.xt) + grad = similar(x, eltype(eltype(yt))) GTPSA.gradient!(grad, yt; include_params=true) return grad end function DI.gradient!(f, grad, ::AutoGTPSA, x, extras::GTPSAGradientExtras) - foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part - yt = f(extras.v) + foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part + yt = f(extras.xt) GTPSA.gradient!(grad, yt; include_params=true) return grad end function DI.value_and_gradient(f, ::AutoGTPSA, x, extras::GTPSAGradientExtras) - foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part - yt = f(extras.v) - grad = similar(x, GTPSA.numtype(eltype(yt))) + foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part + yt = f(extras.xt) + grad = similar(x, eltype(eltype(yt))) GTPSA.gradient!(grad, yt; include_params=true) y = map(t -> t[0], yt) return y, grad end function DI.value_and_gradient!(f, grad, ::AutoGTPSA, x, extras::GTPSAGradientExtras) - foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part - yt = f(extras.v) + foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part + yt = f(extras.xt) GTPSA.gradient!(grad, yt; include_params=true) y = map(t -> t[0], yt) return y, grad @@ -260,8 +260,8 @@ end ## Jacobian -struct GTPSAJacobianExtras{V} <: JacobianExtras - v::V +struct GTPSAJacobianExtras{X} <: JacobianExtras + xt::X end function DI.prepare_jacobian(f, backend::AutoGTPSA{D}, x) where {D} @@ -272,48 +272,48 @@ function DI.prepare_jacobian(f, backend::AutoGTPSA{D}, x) where {D} nn = length(x) d = Descriptor(nn, 1) end - v = similar(x, TPS) + xt = similar(x, TPS{promote_type(eltype(x),Float64)}) - # v and x have same indexing because of similar + # 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(v) - v[i] = TPS(; use=d) - v[i][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 GTPSAJacobianExtras(v) + return GTPSAJacobianExtras(xt) end function DI.jacobian(f, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) - foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part - yt = f(extras.v) - jac = similar(x, GTPSA.numtype(eltype(yt)), (length(yt), length(x))) + foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part + yt = f(extras.xt) + jac = similar(x, eltype(eltype(yt)), (length(yt), length(x))) GTPSA.jacobian!(jac, yt; include_params=true) return jac end function DI.jacobian!(f, jac, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) - foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part - yt = f(extras.v) + foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part + yt = f(extras.xt) GTPSA.jacobian!(jac, yt; include_params=true) return jac end function DI.value_and_jacobian(f, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) - foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part - yt = f(extras.v) - jac = similar(x, GTPSA.numtype(eltype(yt)), (length(yt), length(x))) + foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part + yt = f(extras.xt) + jac = similar(x, eltype(eltype(yt)), (length(yt), length(x))) GTPSA.jacobian!(jac, yt; include_params=true) y = map(t -> t[0], yt) return y, jac end function DI.value_and_jacobian!(f, jac, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) - foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part - yt = f(extras.v) + foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part + yt = f(extras.xt) GTPSA.jacobian!(jac, yt; include_params=true) y = map(t -> t[0], yt) return y, jac @@ -321,8 +321,8 @@ end ## Second derivative -struct GTPSASecondDerivativeExtras{T} <: SecondDerivativeExtras - t::T +struct GTPSASecondDerivativeExtras{X} <: SecondDerivativeExtras + xt::X end function DI.prepare_second_derivative(f, backend::AutoGTPSA{D}, x) where {D} @@ -331,30 +331,30 @@ function DI.prepare_second_derivative(f, backend::AutoGTPSA{D}, x) where {D} else d = Descriptor(1, 2) end - t = TPS(; use=d) - t[1] = 1 - return GTPSASecondDerivativeExtras(t) + xt = TPS{promote_type(typeof(x),Float64)}(; use=d) + xt[1] = 1 + return GTPSASecondDerivativeExtras(xt) end function DI.second_derivative(f, ::AutoGTPSA, x, extras::GTPSASecondDerivativeExtras) - extras.t[0] = x - yt = f(extras.t) + extras.xt[0] = x + yt = f(extras.xt) if yt isa Number return yt[2] else - der2 = similar(yt, GTPSA.numtype(eltype(yt))) + der2 = similar(yt, eltype(eltype(yt))) for i in eachindex(yt) - der2[i] = yt[i][2] + der2[i] = yt[i][2]/2 end return der2 end end function DI.second_derivative!(f, der2, ::AutoGTPSA, x, extras::GTPSASecondDerivativeExtras) - extras.t[0] = x - yt = f(extras.t) + extras.xt[0] = x + yt = f(extras.xt) for i in eachindex(yt) - der2[i] = yt[i][2] + der2[i] = yt[i][2]/2 end return der2 end @@ -362,17 +362,17 @@ end function DI.value_derivative_and_second_derivative( f, ::AutoGTPSA, x, extras::GTPSASecondDerivativeExtras ) - extras.t[0] = x - yt = f(extras.t) + extras.xt[0] = x + yt = f(extras.xt) if yt isa Number return yt[0], yt[1], yt[2] else y = map(t -> t[0], yt) - der = similar(yt, GTPSA.numtype(eltype(yt))) - der2 = similar(yt, GTPSA.numtype(eltype(yt))) + der = similar(yt, eltype(eltype(yt))) + der2 = similar(yt, eltype(eltype(yt))) for i in eachindex(yt) der[i] = yt[i][1] - der2[i] = yt[i][2] + der2[i] = yt[i][2]/2 end return y, der, der2 end @@ -381,67 +381,72 @@ end function DI.value_derivative_and_second_derivative!( f, der, der2, ::AutoGTPSA, x, extras::GTPSASecondDerivativeExtras ) - extras.t[0] = x - yt = f(extras.t) + extras.xt[0] = x + yt = f(extras.xt) y = map(t -> t[0], yt) for i in eachindex(yt) der[i] = yt[i][1] - der2[i] = yt[i][2] + der2[i] = yt[i][2]/2 end return y, der, der2 end ## Hessian -struct GTPSAHessianExtras{V} <: HessianExtras - v::V +struct GTPSAHessianExtras{X,M} <: HessianExtras + xt::X + m::M end function DI.prepare_hessian(f, backend::AutoGTPSA{D}, x) where {D} if D != Nothing d = backend.descriptor nn = GTPSA.numnn(d) + m = Vector{UInt8}(undef, nn) 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 - v = similar(x, TPS) + xt = similar(x, TPS{promote_type(eltype(x),Float64)}) - # v and x have same indexing because of similar + # 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(v) - v[i] = TPS(; use=d) - v[i][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 GTPSAHessianExtras(v) + return GTPSAHessianExtras(xt, m) end function DI.hessian(f, ::AutoGTPSA, x, extras::GTPSAHessianExtras) - foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part - yt = f(extras.v) - hess = similar(x, GTPSA.numtype(eltype(yt)), (length(x), length(x))) + foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part + yt = f(extras.xt) + hess = similar(x, eltype(eltype(yt)), (length(x), length(x))) GTPSA.hessian!(hess, yt; include_params=true) return hess end function DI.hessian!(f, hess, ::AutoGTPSA, x, extras::GTPSAHessianExtras) - foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part - yt = f(extras.v) + foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part + yt = f(extras.xt) GTPSA.hessian!(hess, yt; include_params=true) return hess end function DI.value_gradient_and_hessian(f, ::AutoGTPSA, x, extras::GTPSAHessianExtras) - foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part - yt = f(extras.v) + foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part + yt = f(extras.xt) y = map(t -> t[0], yt) - grad = similar(x, GTPSA.numtype(eltype(yt))) + grad = similar(x, eltype(eltype(yt))) GTPSA.gradient!(grad, yt; include_params=true) - hess = similar(x, GTPSA.numtype(eltype(yt)), (length(x), length(x))) + hess = similar(x, eltype(eltype(yt)), (length(x), length(x))) GTPSA.hessian!(hess, yt; include_params=true) return y, grad, hess end @@ -449,10 +454,118 @@ end function DI.value_gradient_and_hessian!( f, grad, hess, ::AutoGTPSA, x, extras::GTPSAHessianExtras ) - foreach((t, xi) -> t[0] = xi, extras.v, x) # Set the scalar part - yt = f(extras.v) + foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part + yt = f(extras.xt) y = map(t -> t[0], yt) GTPSA.gradient!(grad, yt; include_params=true) GTPSA.hessian!(hess, yt; include_params=true) return y, grad, hess end + +# HVP +struct GTPSAHVPExtras{X,M} <: HVPExtras + xt::X + m::M +end + +function DI.prepare_hvp(f, backend::AutoGTPSA{D}, x, dx) where {D} + if D != Nothing + d = backend.descriptor + nn = GTPSA.numnn(d) + m = Vector{UInt8}(undef, nn) + 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(dx),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(dx),eltype(x),Float64)}(; use=d) + xt[i][j] = 1 + j += 1 + end + + return GTPSAHVPExtras(xt, mono) +end + +function DI.hvp(f, backend::AutoGTPSA{D}, x, dx, extras::GTPSAHVPExtras) where {D} + foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part + + yt = f(extras.xt) + + + dg = similar(x, eltype(eltype(yt))) + dg .= 0 + + d = GTPSA.getdesc(yt) + desc = unsafe_load(d.desc) + nn = desc.nn + + + if D == Nothing + idx = desc.nv+desc.np + endidx = floor(nn*(nn+1)/2)+nn + curdiag = 1 + col = 1 + xt = Ref{eltype(yt)}() + idx = cycle!(yt, idx, 0, C_NULL, xt) + while idx <= endidx && idx > 0 + h_idx = idx-nn + while h_idx > curdiag + col += 1 + curdiag += col + end + row = col-(curdiag-h_idx) + #println("row = ", row, ", col = ", col) + if row==col + dg[row] += 2*xt[]*dx[col] + else + dg[row] += xt[]*dx[col] + dg[col] += xt[]*dx[row] + end + idx = cycle!(yt, idx, 0, C_NULL, xt) + end + else + # If there are some variables/parameters with TO == 1, we have to do it "slow" + # because the indexing of TPSA index -> hessian index can be very complicated. + # I saw slow in quotes because it is likely still much faster than the calculation + # of the Hessian itself (this is just a getter) + idx = desc.nv+desc.np # start at 2nd order + xt = Ref{eltype(yt)}() + mono = extras.mono + idx = cycle!(yt, idx, nn, mono, xt) + while idx > 0 + if sum(mono) > 0x2 + return dg + end + i = findfirst(x->x==0x1, mono) + if isnothing(i) + i = findfirst(x->x==0x2, mono) + if isnothing(i) + return dg + end + if i <= nn + dg[i] += 2*xt[]*dx[i] # Multiply by 2 because taylor coefficient on diagonal is 1/2!*d2f/dx2 + end + else + j = findlast(x->x==0x1, mono) + if isnothing(j) + return dg + end + if i <= nn && j <= nn + dg[i] += xt[]*dx[j] + dg[j] += xt[]*dx[i] + end + end + idx = cycle!(yt, idx, nn, mono, xt) + end + end + return dg +end \ No newline at end of file diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl new file mode 100644 index 000000000..1d64d3c83 --- /dev/null +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl @@ -0,0 +1,293 @@ +## Pushforward +struct GTPSATwoArgPushforwardExtras{X,Y} <: PushforwardExtras + xt::X + yt::Y +end + +function DI.prepare_pushforward(f!, y, backend::AutoGTPSA{D}, x, dx) where {D} + if D != Nothing + d = backend.descriptor + else + d = Descriptor(length(x), 1) + end + + if x isa Number + xt = TPS{promote_type(typeof(dx),typeof(x),Float64)}(; use=d) + else + xt = similar(x, TPS{promote_type(eltype(dx),eltype(x),Float64)}) + + # xt and x have same indexing because of similar + for i in eachindex(xt) + xt[i] = TPS{promote_type(eltype(dx),eltype(x),Float64)}(; use=d) + end + end + + if y isa Number + yt = TPS{promote_type(typeof(y),Float64)}(; use=d) + else + 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 + end + + return GTPSATwoArgPushforwardExtras(xt, yt) +end + +function DI.pushforward(f!, y, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPushforwardExtras) + if x isa Number + extras.xt[0] = x + extras.xt[1] = dx + else + j = 1 + for i in eachindex(x) + extras.xt[i][0] = x[i] + extras.xt[i][j] = dx[i] + j += 1 + end + end + + f!(extras.yt, extras.xt) + if extras.yt isa Number + return extras.yt[1] + else + dy = similar(extras.yt, eltype(eltype(extras.yt))) + j = 1 + for i in eachindex(extras.yt) + dy[i] = extras.yt[i][j] + j += 1 + end + return dy + end +end + +function DI.pushforward!(f!, y, dy, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPushforwardExtras) + if x isa Number + extras.xt[0] = x + extras.xt[1] = dx + else + j = 1 + for i in eachindex(x) + extras.xt[i][0] = x[i] + extras.xt[i][j] = dx[i] + j += 1 + end + end + + f!(extras.yt, extras.xt) + if extras.yt isa Number + return extras.yt[1] + else + j = 1 + for i in eachindex(extras.yt) + dy[i] = extras.yt[i][j] + j += 1 + end + return dy + end +end + +function DI.value_and_pushforward(f!, y, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPushforwardExtras) + if x isa Number + extras.xt[0] = x + extras.xt[1] = dx + else + j = 1 + for i in eachindex(x) + extras.xt[i][0] = x[i] + extras.xt[i][j] = dx[i] + j += 1 + end + end + + f!(extras.yt, extras.xt) + if extras.yt isa Number + return extras.yt[1] + else + dy = similar(extras.yt, eltype(eltype(extras.yt))) + j = 1 + for i in eachindex(extras.yt) + dy[i] = extras.yt[i][j] + j += 1 + end + y = map(t->t[0], extras.yt) + return y, dy + end +end + +function DI.value_and_pushforward!(f!, y, dy, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPushforwardExtras) + if x isa Number + extras.xt[0] = x + extras.xt[1] = dx + else + j = 1 + for i in eachindex(x) + extras.xt[i][0] = x[i] + extras.xt[i][j] = dx[i] + j += 1 + end + end + + f!(extras.yt, extras.xt) + if extras.yt isa Number + return extras.yt[1] + else + j = 1 + for i in eachindex(extras.yt) + dy[i] = extras.yt[i][j] + j += 1 + end + y = map(t->t[0], extras.yt) + return y, dy + end +end + +## Derivative + +struct GTPSATwoArgDerivativeExtras{X,Y} <: TwoArgDerivativeExtras + xt::X + yt::Y +end + +function DI.prepare_derivative(f!, y, backend::AutoGTPSA{D}, x) where {D} + if D != Nothing + d = backend.descriptor + else + d = Descriptor(1, 1) + end + xt = TPS{promote_type(typeof(x),Float64)}(; use=d) + xt[1] = 1 + + if y isa Number + yt = TPS{promote_type(typeof(y),Float64)}(; use=d) + else + 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 + end + + return GTPSATwoArgDerivativeExtras(xt, yt) +end + +function DI.derivative(f!, y, ::AutoGTPSA, x, extras::GTPSATwoArgDerivativeExtras) + extras.xt[0] = x + f!(extras.yt, extras.xt) + if extras.yt isa Number + return extras.yt[1] + else + der = similar(extras.yt, eltype(eltype(extras.yt))) + for i in eachindex(extras.yt) + der[i] = extras.yt[i][1] + end + return der + end +end + +function DI.derivative!(f!, y, der, ::AutoGTPSA, x, extras::GTPSATwoArgDerivativeExtras) + extras.xt[0] = x + f!(extras.yt, extras.xt) + for i in eachindex(extras.yt) + der[i] = extras.yt[i][1] + end + return der +end + +function DI.value_and_derivative(f!, y, ::AutoGTPSA, x, extras::GTPSATwoArgDerivativeExtras) + extras.xt[0] = x + f!(extras.yt, extras.xt) + if extras.yt isa Number + return extras.yt[0], extras.yt[1] + else + y = map(t -> t[0], extras.yt) + der = similar(extras.yt, eltype(eltype(extras.yt))) + for i in eachindex(extras.yt) + der[i] = extras.yt[i][1] + end + return y, der + end +end + +function DI.value_and_derivative!(f!, y, der, ::AutoGTPSA, x, extras::GTPSATwoArgDerivativeExtras) + extras.xt[0] = x + f!(extras.yt, extras.xt) + y = map(t -> t[0], extras.yt) + for i in eachindex(extras.yt) + der[i] = extras.yt[i][1] + end + return y, der +end + +## Jacobian + +struct GTPSATwoArgJacobianExtras{X,Y} <: JacobianExtras + xt::X + yt::Y +end + +function DI.prepare_jacobian(f!, y, backend::AutoGTPSA{D}, x) where {D} + if D != Nothing + d = backend.descriptor + nn = GTPSA.numnn(d) + else + nn = length(x) + d = Descriptor(nn, 1) + 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 + + if y isa Number + yt = TPS{promote_type(typeof(y),Float64)}(; use=d) + else + 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 + end + + return GTPSATwoArgJacobianExtras(xt, yt) +end + +function DI.jacobian(f!, y, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) + foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part + f!(extras.yt, extras.xt) + jac = similar(x, eltype(eltype(extras.yt)), (length(extras.yt), length(x))) + GTPSA.jacobian!(jac, extras.yt; include_params=true) + return jac +end + +function DI.jacobian!(f!, y, jac, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) + foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part + f!(extras.yt, extras.xt) + GTPSA.jacobian!(jac, extras.yt; include_params=true) + return jac +end + +function DI.value_and_jacobian(f!, y, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) + foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part + f!(extras.yt, extras.xt) + jac = similar(x, eltype(eltype(extras.yt)), (length(extras.yt), length(x))) + GTPSA.jacobian!(jac, extras.yt; include_params=true) + y = map(t -> t[0], extras.yt) + return y, jac +end + +function DI.value_and_jacobian!(f!, y, jac, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) + foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part + f!(extras.yt, extras.xt) + GTPSA.jacobian!(jac, extras.yt; include_params=true) + y = map(t -> t[0], extras.yt) + return y, jac +end From 3afbbb9978d4e9e374c23031a94cf6706d682ba9 Mon Sep 17 00:00:00 2001 From: mattsignorelli Date: Tue, 6 Aug 2024 21:03:18 -0400 Subject: [PATCH 12/44] fix typo --- .../ext/DifferentiationInterfaceGTPSAExt/twoarg.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl index 1d64d3c83..cd4d1d44d 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl @@ -145,7 +145,7 @@ end ## Derivative -struct GTPSATwoArgDerivativeExtras{X,Y} <: TwoArgDerivativeExtras +struct GTPSATwoArgDerivativeExtras{X,Y} <: DerivativeExtras xt::X yt::Y end From 3039528d84564cf76c735af0d06dde21f23711fd Mon Sep 17 00:00:00 2001 From: mattsignorelli <113714234+mattsignorelli@users.noreply.github.com> Date: Wed, 7 Aug 2024 10:03:45 -0400 Subject: [PATCH 13/44] Update Test.yml --- .github/workflows/Test.yml | 36 ++++++++++++++++++------------------ 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml index 93c5d11cc..92952e1cb 100644 --- a/.github/workflows/Test.yml +++ b/.github/workflows/Test.yml @@ -35,24 +35,24 @@ jobs: group: - Formalities - Internals - - Back/Diffractor - - Back/Enzyme - - Back/FastDifferentiation - - Back/FiniteDiff - - Back/FiniteDifferences - - Back/ForwardDiff + #- Back/Diffractor + #- Back/Enzyme + #- Back/FastDifferentiation + #- Back/FiniteDiff + #- Back/FiniteDifferences + #- Back/ForwardDiff - Back/GTPSA - - Back/PolyesterForwardDiff - - Back/ReverseDiff - - Back/SecondOrder - - Back/Symbolics - - Back/Tapir - - Back/Tracker - - Back/Zygote - - Down/Detector - - Down/DifferentiateWith - - Down/Flux - - Down/Lux + #- Back/PolyesterForwardDiff + #- Back/ReverseDiff + #- Back/SecondOrder + #- Back/Symbolics + #- Back/Tapir + #- Back/Tracker + #- Back/Zygote + #- Down/Detector + #- Down/DifferentiateWith + #- Down/Flux + #- Down/Lux exclude: # lts - version: 'lts' @@ -178,4 +178,4 @@ jobs: # with: # files: lcov.info # token: ${{ secrets.CODECOV_TOKEN }} -# fail_ci_if_error: true \ No newline at end of file +# fail_ci_if_error: true From ecee93eee5bdad32559301e84fd38ce5cb82ceaf Mon Sep 17 00:00:00 2001 From: mattsignorelli <113714234+mattsignorelli@users.noreply.github.com> Date: Wed, 7 Aug 2024 11:36:43 -0400 Subject: [PATCH 14/44] no code coverage --- .github/workflows/Test.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml index 92952e1cb..57d3d9b99 100644 --- a/.github/workflows/Test.yml +++ b/.github/workflows/Test.yml @@ -111,7 +111,7 @@ jobs: using Pkg; Pkg.Registry.update(); Pkg.update(); - Pkg.test("DifferentiationInterface"; coverage=true);' + Pkg.test("DifferentiationInterface"; coverage=false);' - uses: julia-actions/julia-processcoverage@v1 if: ${{ env.SHOULDRUN == 'true' }} with: From f0a0d79d9dbb1e87f4f936e480e47b80c1675bcb Mon Sep 17 00:00:00 2001 From: mattsignorelli Date: Wed, 7 Aug 2024 11:45:39 -0400 Subject: [PATCH 15/44] add GTPSA test --- .../test/Back/GTPSA/test.jl | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) create mode 100644 DifferentiationInterface/test/Back/GTPSA/test.jl diff --git a/DifferentiationInterface/test/Back/GTPSA/test.jl b/DifferentiationInterface/test/Back/GTPSA/test.jl new file mode 100644 index 000000000..87b7a0910 --- /dev/null +++ b/DifferentiationInterface/test/Back/GTPSA/test.jl @@ -0,0 +1,19 @@ +using Pkg +Pkg.add("GTPSA") + +using DifferentiationInterface, DifferentiationInterfaceTest +using GTPSA: GTPSA +using Test + +LOGGING = get(ENV, "CI", "false") == "false" + +for backend in [AutoGTPSA()] + @test check_available(backend) + @test check_twoarg(backend) + @test check_hessian(backend) +end + +test_differentiation( + AutoGTPSA(); + type_stability=true, + logging=LOGGING); \ No newline at end of file From 8997b932d243804680b562be2e5e3ec97107dd6b Mon Sep 17 00:00:00 2001 From: mattsignorelli Date: Wed, 7 Aug 2024 20:47:41 -0400 Subject: [PATCH 16/44] almost all tests passing, bugs in pushforward --- .../onearg.jl | 172 +++++++++++++++--- .../twoarg.jl | 110 ++++++++--- 2 files changed, 234 insertions(+), 48 deletions(-) diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl index 4531cca21..dc5d90435 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl @@ -25,6 +25,7 @@ function DI.prepare_pushforward(f, backend::AutoGTPSA{D}, x, dx) where {D} end function DI.pushforward(f, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras) + if x isa Number extras.xt[0] = x extras.xt[1] = dx @@ -39,19 +40,32 @@ function DI.pushforward(f, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardEx yt = f(extras.xt) if yt isa Number - return yt[1] + if dx isa Number + + return yt[1] + else + dy = 0 + for j=1:length(dx) + dy += yt[j] + end + + return dy + end else dy = similar(yt, eltype(eltype(yt))) - j = 1 + dy .= 0 for i in eachindex(yt) - dy[i] = yt[i][j] - j += 1 + for j=1:length(dx) + dy[i] += yt[i][j] + end end + return dy end end function DI.pushforward!(f, dy, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras) + if x isa Number extras.xt[0] = x extras.xt[1] = dx @@ -66,18 +80,31 @@ function DI.pushforward!(f, dy, backend::AutoGTPSA, x, dx, extras::GTPSAPushforw yt = f(extras.xt) if yt isa Number - return yt[1] + if dx isa Number + + return yt[1] + else + dy = 0 + for j=1:length(dx) + dy += yt[j] + end + + return dy + end else - j = 1 + dy .= 0 for i in eachindex(yt) - dy[i] = yt[i][j] - j += 1 + for j=1:length(dx) + dy[i] += yt[i][j] + end end + return dy end end function DI.value_and_pushforward(f, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras) + if x isa Number extras.xt[0] = x extras.xt[1] = dx @@ -92,20 +119,33 @@ function DI.value_and_pushforward(f, backend::AutoGTPSA, x, dx, extras::GTPSAPus yt = f(extras.xt) if yt isa Number - return yt[1] + if dx isa Number + + return yt[0], yt[1] + else + dy = 0 + for j=1:length(dx) + dy += yt[j] + end + + return yt[0], dy + end else dy = similar(yt, eltype(eltype(yt))) - j = 1 + dy .= 0 for i in eachindex(yt) - dy[i] = yt[i][j] - j += 1 + for j=1:length(dx) + dy[i] += yt[i][j] + end end y = map(t->t[0], yt) + return y, dy end end function DI.value_and_pushforward!(f, dy, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras) + if x isa Number extras.xt[0] = x extras.xt[1] = dx @@ -120,16 +160,29 @@ function DI.value_and_pushforward!(f, dy, backend::AutoGTPSA, x, dx, extras::GTP yt = f(extras.xt) if yt isa Number - return yt[1] + if dx isa Number + return yt[0], yt[1] + + else + dy = 0 + for j=1:length(dx) + dy += yt[j] + end + + return yt[0], dy + end else - j = 1 + dy .= 0 for i in eachindex(yt) - dy[i] = yt[i][j] - j += 1 + for j=1:length(dx) + dy[i] += yt[i][j] + end end y = map(t->t[0], yt) + return y, dy end + end ## Derivative @@ -340,11 +393,11 @@ function DI.second_derivative(f, ::AutoGTPSA, x, extras::GTPSASecondDerivativeEx extras.xt[0] = x yt = f(extras.xt) if yt isa Number - return yt[2] + return yt[2]*2 else der2 = similar(yt, eltype(eltype(yt))) for i in eachindex(yt) - der2[i] = yt[i][2]/2 + der2[i] = yt[i][2]*2 # *2 because monomial coefficient is 1/2 end return der2 end @@ -354,7 +407,7 @@ function DI.second_derivative!(f, der2, ::AutoGTPSA, x, extras::GTPSASecondDeriv extras.xt[0] = x yt = f(extras.xt) for i in eachindex(yt) - der2[i] = yt[i][2]/2 + der2[i] = yt[i][2]*2 end return der2 end @@ -365,14 +418,14 @@ function DI.value_derivative_and_second_derivative( extras.xt[0] = x yt = f(extras.xt) if yt isa Number - return yt[0], yt[1], yt[2] + return yt[0], yt[1], yt[2]*2 else y = map(t -> t[0], yt) der = similar(yt, eltype(eltype(yt))) der2 = similar(yt, eltype(eltype(yt))) for i in eachindex(yt) der[i] = yt[i][1] - der2[i] = yt[i][2]/2 + der2[i] = yt[i][2]*2 end return y, der, der2 end @@ -386,7 +439,7 @@ function DI.value_derivative_and_second_derivative!( y = map(t -> t[0], yt) for i in eachindex(yt) der[i] = yt[i][1] - der2[i] = yt[i][2]/2 + der2[i] = yt[i][2]*2 end return y, der, der2 end @@ -509,6 +562,81 @@ function DI.hvp(f, backend::AutoGTPSA{D}, x, dx, extras::GTPSAHVPExtras) where { nn = desc.nn + if D == Nothing + idx = desc.nv+desc.np + endidx = floor(nn*(nn+1)/2)+nn + curdiag = 1 + col = 1 + xt = Ref{eltype(yt)}() + idx = cycle!(yt, idx, 0, C_NULL, xt) + while idx <= endidx && idx > 0 + h_idx = idx-nn + while h_idx > curdiag + col += 1 + curdiag += col + end + row = col-(curdiag-h_idx) + #println("row = ", row, ", col = ", col) + if row==col + dg[row] += 2*xt[]*dx[col] + else + dg[row] += xt[]*dx[col] + dg[col] += xt[]*dx[row] + end + idx = cycle!(yt, idx, 0, C_NULL, xt) + end + else + # If there are some variables/parameters with TO == 1, we have to do it "slow" + # because the indexing of TPSA index -> hessian index can be very complicated. + # I saw slow in quotes because it is likely still much faster than the calculation + # of the Hessian itself (this is just a getter) + idx = desc.nv+desc.np # start at 2nd order + xt = Ref{eltype(yt)}() + mono = extras.mono + idx = cycle!(yt, idx, nn, mono, xt) + while idx > 0 + if sum(mono) > 0x2 + return dg + end + i = findfirst(x->x==0x1, mono) + if isnothing(i) + i = findfirst(x->x==0x2, mono) + if isnothing(i) + return dg + end + if i <= nn + dg[i] += 2*xt[]*dx[i] # Multiply by 2 because taylor coefficient on diagonal is 1/2!*d2f/dx2 + end + else + j = findlast(x->x==0x1, mono) + if isnothing(j) + return dg + end + if i <= nn && j <= nn + dg[i] += xt[]*dx[j] + dg[j] += xt[]*dx[i] + end + end + idx = cycle!(yt, idx, nn, mono, xt) + end + end + return dg +end + +function DI.hvp!(f, dg, backend::AutoGTPSA{D}, x, dx, extras::GTPSAHVPExtras) where {D} + foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part + + yt = f(extras.xt) + + + #dg = similar(x, eltype(eltype(yt))) + dg .= 0 + + d = GTPSA.getdesc(yt) + desc = unsafe_load(d.desc) + nn = desc.nn + + if D == Nothing idx = desc.nv+desc.np endidx = floor(nn*(nn+1)/2)+nn diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl index cd4d1d44d..038bf50a9 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl @@ -36,6 +36,7 @@ function DI.prepare_pushforward(f!, y, backend::AutoGTPSA{D}, x, dx) where {D} end function DI.pushforward(f!, y, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPushforwardExtras) + if x isa Number extras.xt[0] = x extras.xt[1] = dx @@ -50,19 +51,33 @@ function DI.pushforward(f!, y, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPus f!(extras.yt, extras.xt) if extras.yt isa Number - return extras.yt[1] + if dx isa Number + + return extras.yt[1] + else + dy = 0 + for j=1:length(dx) + dy += extras.yt[j] + end + + return dy + end else + map!(t -> t[0], y, extras.yt) dy = similar(extras.yt, eltype(eltype(extras.yt))) - j = 1 + dy .= 0 for i in eachindex(extras.yt) - dy[i] = extras.yt[i][j] - j += 1 + for j=1:length(dx) + dy[i] = extras.yt[i][j] + end end + return dy end end function DI.pushforward!(f!, y, dy, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPushforwardExtras) + if x isa Number extras.xt[0] = x extras.xt[1] = dx @@ -77,18 +92,32 @@ function DI.pushforward!(f!, y, dy, backend::AutoGTPSA, x, dx, extras::GTPSATwoA f!(extras.yt, extras.xt) if extras.yt isa Number - return extras.yt[1] + if dx isa Number + + return extras.yt[1] + else + dy = 0 + for j=1:length(dx) + dy += extras.yt[j] + end + + return extras.dy + end else - j = 1 + map!(t -> t[0], y, extras.yt) + dy .= 0 for i in eachindex(extras.yt) - dy[i] = extras.yt[i][j] - j += 1 + for j=1:length(dx) + dy[i] += extras.yt[i][j] + end end + return dy end end function DI.value_and_pushforward(f!, y, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPushforwardExtras) + if x isa Number extras.xt[0] = x extras.xt[1] = dx @@ -103,20 +132,33 @@ function DI.value_and_pushforward(f!, y, backend::AutoGTPSA, x, dx, extras::GTPS f!(extras.yt, extras.xt) if extras.yt isa Number - return extras.yt[1] + if dx isa Number + + return extras.yt[0], extras.yt[1] + else + dy = 0 + for j=1:length(dx) + dy += extras.yt[j] + end + + return extras.yt[0], dy + end else dy = similar(extras.yt, eltype(eltype(extras.yt))) - j = 1 + dy .= 0 for i in eachindex(extras.yt) - dy[i] = extras.yt[i][j] - j += 1 + for j=1:length(dx) + dy[i] += extras.yt[i][j] + end end - y = map(t->t[0], extras.yt) + map!(t -> t[0], y, extras.yt) + return y, dy end end function DI.value_and_pushforward!(f!, y, dy, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPushforwardExtras) + if x isa Number extras.xt[0] = x extras.xt[1] = dx @@ -131,14 +173,26 @@ function DI.value_and_pushforward!(f!, y, dy, backend::AutoGTPSA, x, dx, extras: f!(extras.yt, extras.xt) if extras.yt isa Number - return extras.yt[1] + if dx isa Number + return extras.yt[0], extras.yt[1] + + else + dy = 0 + for j=1:length(dx) + dy += yt[j] + end + + return extras.yt[0], dy + end else - j = 1 + map!(t -> t[0], y, extras.yt) + dy .= 0 for i in eachindex(extras.yt) - dy[i] = extras.yt[i][j] - j += 1 + for j=1:length(dx) + dy[i] += extras.yt[i][j] + end end - y = map(t->t[0], extras.yt) + return y, dy end end @@ -178,6 +232,7 @@ function DI.derivative(f!, y, ::AutoGTPSA, x, extras::GTPSATwoArgDerivativeExtra if extras.yt isa Number return extras.yt[1] else + map!(t -> t[0], y, extras.yt) der = similar(extras.yt, eltype(eltype(extras.yt))) for i in eachindex(extras.yt) der[i] = extras.yt[i][1] @@ -192,6 +247,7 @@ function DI.derivative!(f!, y, der, ::AutoGTPSA, x, extras::GTPSATwoArgDerivativ for i in eachindex(extras.yt) der[i] = extras.yt[i][1] end + map!(t -> t[0], y, extras.yt) return der end @@ -201,7 +257,7 @@ function DI.value_and_derivative(f!, y, ::AutoGTPSA, x, extras::GTPSATwoArgDeriv if extras.yt isa Number return extras.yt[0], extras.yt[1] else - y = map(t -> t[0], extras.yt) + map!(t -> t[0], y, extras.yt) der = similar(extras.yt, eltype(eltype(extras.yt))) for i in eachindex(extras.yt) der[i] = extras.yt[i][1] @@ -213,7 +269,7 @@ end function DI.value_and_derivative!(f!, y, der, ::AutoGTPSA, x, extras::GTPSATwoArgDerivativeExtras) extras.xt[0] = x f!(extras.yt, extras.xt) - y = map(t -> t[0], extras.yt) + map!(t -> t[0], y, extras.yt) for i in eachindex(extras.yt) der[i] = extras.yt[i][1] end @@ -260,34 +316,36 @@ function DI.prepare_jacobian(f!, y, backend::AutoGTPSA{D}, x) where {D} return GTPSATwoArgJacobianExtras(xt, yt) end -function DI.jacobian(f!, y, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) +function DI.jacobian(f!, y, ::AutoGTPSA, x, extras::GTPSATwoArgJacobianExtras) foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part f!(extras.yt, extras.xt) jac = similar(x, eltype(eltype(extras.yt)), (length(extras.yt), length(x))) GTPSA.jacobian!(jac, extras.yt; include_params=true) + map!(t -> t[0], y, extras.yt) return jac end -function DI.jacobian!(f!, y, jac, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) +function DI.jacobian!(f!, y, jac, ::AutoGTPSA, x, extras::GTPSATwoArgJacobianExtras) foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part f!(extras.yt, extras.xt) GTPSA.jacobian!(jac, extras.yt; include_params=true) + map!(t -> t[0], y, extras.yt) return jac end -function DI.value_and_jacobian(f!, y, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) +function DI.value_and_jacobian(f!, y, ::AutoGTPSA, x, extras::GTPSATwoArgJacobianExtras) foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part f!(extras.yt, extras.xt) jac = similar(x, eltype(eltype(extras.yt)), (length(extras.yt), length(x))) GTPSA.jacobian!(jac, extras.yt; include_params=true) - y = map(t -> t[0], extras.yt) + map!(t -> t[0], y, extras.yt) return y, jac end -function DI.value_and_jacobian!(f!, y, jac, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) +function DI.value_and_jacobian!(f!, y, jac, ::AutoGTPSA, x, extras::GTPSATwoArgJacobianExtras) foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part f!(extras.yt, extras.xt) GTPSA.jacobian!(jac, extras.yt; include_params=true) - y = map(t -> t[0], extras.yt) + map!(t -> t[0], y, extras.yt) return y, jac end From b9ae5e063e27e32ce7d4163fb6458b0b234d77dd Mon Sep 17 00:00:00 2001 From: mattsignorelli Date: Wed, 7 Aug 2024 20:52:51 -0400 Subject: [PATCH 17/44] all tests passing --- .../ext/DifferentiationInterfaceGTPSAExt/twoarg.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl index 038bf50a9..1d1c28a2c 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl @@ -68,7 +68,7 @@ function DI.pushforward(f!, y, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPus dy .= 0 for i in eachindex(extras.yt) for j=1:length(dx) - dy[i] = extras.yt[i][j] + dy[i] += extras.yt[i][j] end end From 12ac9d9d2e105ffe5f42458f550fd716486bb101 Mon Sep 17 00:00:00 2001 From: mattsignorelli Date: Wed, 7 Aug 2024 21:21:07 -0400 Subject: [PATCH 18/44] latest GTPSA version specified, docs updated --- DifferentiationInterface/Project.toml | 2 +- DifferentiationInterface/README.md | 1 + DifferentiationInterface/docs/Project.toml | 1 + DifferentiationInterface/docs/src/backends.md | 3 +++ DifferentiationInterface/docs/src/implementations.md | 9 +++++++++ DifferentiationInterfaceTest/Project.toml | 1 - 6 files changed, 15 insertions(+), 2 deletions(-) diff --git a/DifferentiationInterface/Project.toml b/DifferentiationInterface/Project.toml index d222862ab..5615d4944 100644 --- a/DifferentiationInterface/Project.toml +++ b/DifferentiationInterface/Project.toml @@ -57,7 +57,7 @@ FillArrays = "1.7.0" FiniteDiff = "2.23.1" FiniteDifferences = "0.12.31" ForwardDiff = "0.10.36" -GTPSA = "1.0.1" +GTPSA = "1.1.0" LinearAlgebra = "<0.0.1,1" PackageExtensionCompat = "1.0.2" PolyesterForwardDiff = "0.1.1" diff --git a/DifferentiationInterface/README.md b/DifferentiationInterface/README.md index 9156d7cae..f466a3e02 100644 --- a/DifferentiationInterface/README.md +++ b/DifferentiationInterface/README.md @@ -42,6 +42,7 @@ We support all of the backends defined by [ADTypes.jl](https://github.com/SciML/ - [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) - [PolyesterForwardDiff.jl](https://github.com/JuliaDiff/PolyesterForwardDiff.jl) - [ReverseDiff.jl](https://github.com/JuliaDiff/ReverseDiff.jl) - [Symbolics.jl](https://github.com/JuliaSymbolics/Symbolics.jl) diff --git a/DifferentiationInterface/docs/Project.toml b/DifferentiationInterface/docs/Project.toml index 00d1960b3..9d7909fbd 100644 --- a/DifferentiationInterface/docs/Project.toml +++ b/DifferentiationInterface/docs/Project.toml @@ -13,6 +13,7 @@ FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +GTPSA = "b27dd330-f138-47c5-815b-40db9dd9b6e8" Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a" PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b" PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d" diff --git a/DifferentiationInterface/docs/src/backends.md b/DifferentiationInterface/docs/src/backends.md index 408b90700..da52afa6c 100644 --- a/DifferentiationInterface/docs/src/backends.md +++ b/DifferentiationInterface/docs/src/backends.md @@ -13,6 +13,7 @@ We support all dense backend choices from [ADTypes.jl](https://github.com/SciML/ - [`AutoFiniteDiff`](@extref ADTypes.AutoFiniteDiff) - [`AutoFiniteDifferences`](@extref ADTypes.AutoFiniteDifferences) - [`AutoForwardDiff`](@extref ADTypes.AutoForwardDiff) +- [`AutoGTPSA](@extref ADTypes.AutoGTPSA) - [`AutoPolyesterForwardDiff`](@extref ADTypes.AutoPolyesterForwardDiff) - [`AutoReverseDiff`](@extref ADTypes.AutoReverseDiff) - [`AutoSymbolics`](@extref ADTypes.AutoSymbolics) @@ -48,6 +49,7 @@ import FastDifferentiation import FiniteDiff import FiniteDifferences import ForwardDiff +import GTPSA import PolyesterForwardDiff import ReverseDiff import Symbolics @@ -62,6 +64,7 @@ backend_examples = [ AutoFiniteDiff(), AutoFiniteDifferences(; fdm=FiniteDifferences.central_fdm(3, 1)), AutoForwardDiff(), + AutoGTPSA(), AutoPolyesterForwardDiff(; chunksize=1), AutoReverseDiff(), AutoSymbolics(), diff --git a/DifferentiationInterface/docs/src/implementations.md b/DifferentiationInterface/docs/src/implementations.md index c65b02eec..caea2f4cc 100644 --- a/DifferentiationInterface/docs/src/implementations.md +++ b/DifferentiationInterface/docs/src/implementations.md @@ -33,6 +33,7 @@ using FastDifferentiation: FastDifferentiation using FiniteDiff: FiniteDiff using FiniteDifferences: FiniteDifferences using ForwardDiff: ForwardDiff +using GTPSA: GTPSA using PolyesterForwardDiff: PolyesterForwardDiff using ReverseDiff: ReverseDiff using Symbolics: Symbolics @@ -228,6 +229,14 @@ For [`pushforward`](@ref), preparation allocates the necessary space for `Dual` print_overloads(AutoForwardDiff(), :DifferentiationInterfaceForwardDiffExt) # hide ``` +## GTPSA + +Preparation constructs a [GTPSA Descriptor](https://bmad-sim.github.io/GTPSA.jl/stable/man/c_descriptor/) if none is provided or does not already exist, and will pre-allocate the variables as [TPSs](https://bmad-sim.github.io/GTPSA.jl/stable/man/d_tps/). + +```@example overloads +print_overloads(AutoGTPSA(), :DifferentiationInterfaceGTPSAfExt) # hide +``` + ## PolyesterForwardDiff ```@example overloads diff --git a/DifferentiationInterfaceTest/Project.toml b/DifferentiationInterfaceTest/Project.toml index 0c02f36e5..db69147fe 100644 --- a/DifferentiationInterfaceTest/Project.toml +++ b/DifferentiationInterfaceTest/Project.toml @@ -11,7 +11,6 @@ DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196" -GTPSA = "b27dd330-f138-47c5-815b-40db9dd9b6e8" JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930" From 5c40296ce8974ade2ca40712fe14b342d618b6ca Mon Sep 17 00:00:00 2001 From: mattsignorelli <113714234+mattsignorelli@users.noreply.github.com> Date: Wed, 7 Aug 2024 21:32:55 -0400 Subject: [PATCH 19/44] no coverage test --- .github/workflows/Test.yml | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml index 57d3d9b99..e89790e51 100644 --- a/.github/workflows/Test.yml +++ b/.github/workflows/Test.yml @@ -112,16 +112,16 @@ jobs: Pkg.Registry.update(); Pkg.update(); Pkg.test("DifferentiationInterface"; coverage=false);' - - uses: julia-actions/julia-processcoverage@v1 - if: ${{ env.SHOULDRUN == 'true' }} - with: - directories: ./DifferentiationInterface/src,./DifferentiationInterface/ext,./DifferentiationInterface/test - - uses: codecov/codecov-action@v4 - if: ${{ env.SHOULDRUN == 'true' }} - with: - files: lcov.info - token: ${{ secrets.CODECOV_TOKEN }} - fail_ci_if_error: true + #- uses: julia-actions/julia-processcoverage@v1 + # if: ${{ env.SHOULDRUN == 'true' }} + # with: + # directories: ./DifferentiationInterface/src,./DifferentiationInterface/ext,./DifferentiationInterface/test + #- uses: codecov/codecov-action@v4 + # if: ${{ env.SHOULDRUN == 'true' }} + # with: + # files: lcov.info + # token: ${{ secrets.CODECOV_TOKEN }} + # fail_ci_if_error: true # test-DIT: # name: ${{ matrix.version }} - DIT (${{ matrix.group }}) From 3ed2f0821a98e2935d6a2a906142ce33582ea300 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Thu, 8 Aug 2024 08:27:35 +0200 Subject: [PATCH 20/44] Fix formatting and docs --- DifferentiationInterface/docs/src/api.md | 4 + DifferentiationInterface/docs/src/backends.md | 2 +- .../docs/src/implementations.md | 2 +- .../onearg.jl | 276 +++++++++--------- .../twoarg.jl | 95 +++--- .../test/Back/GTPSA/test.jl | 5 +- 6 files changed, 191 insertions(+), 193 deletions(-) diff --git a/DifferentiationInterface/docs/src/api.md b/DifferentiationInterface/docs/src/api.md index c37435907..acde2edc9 100644 --- a/DifferentiationInterface/docs/src/api.md +++ b/DifferentiationInterface/docs/src/api.md @@ -130,3 +130,7 @@ Modules = [DifferentiationInterface] Public = false Filter = t -> !(Symbol(t) in [:outer, :inner]) ``` + +```@docs +AutoGTPSA +``` \ No newline at end of file diff --git a/DifferentiationInterface/docs/src/backends.md b/DifferentiationInterface/docs/src/backends.md index da52afa6c..aceb16687 100644 --- a/DifferentiationInterface/docs/src/backends.md +++ b/DifferentiationInterface/docs/src/backends.md @@ -13,7 +13,7 @@ We support all dense backend choices from [ADTypes.jl](https://github.com/SciML/ - [`AutoFiniteDiff`](@extref ADTypes.AutoFiniteDiff) - [`AutoFiniteDifferences`](@extref ADTypes.AutoFiniteDifferences) - [`AutoForwardDiff`](@extref ADTypes.AutoForwardDiff) -- [`AutoGTPSA](@extref ADTypes.AutoGTPSA) +- `AutoGTPSA` - [`AutoPolyesterForwardDiff`](@extref ADTypes.AutoPolyesterForwardDiff) - [`AutoReverseDiff`](@extref ADTypes.AutoReverseDiff) - [`AutoSymbolics`](@extref ADTypes.AutoSymbolics) diff --git a/DifferentiationInterface/docs/src/implementations.md b/DifferentiationInterface/docs/src/implementations.md index caea2f4cc..5070bafa3 100644 --- a/DifferentiationInterface/docs/src/implementations.md +++ b/DifferentiationInterface/docs/src/implementations.md @@ -234,7 +234,7 @@ print_overloads(AutoForwardDiff(), :DifferentiationInterfaceForwardDiffExt) # hi Preparation constructs a [GTPSA Descriptor](https://bmad-sim.github.io/GTPSA.jl/stable/man/c_descriptor/) if none is provided or does not already exist, and will pre-allocate the variables as [TPSs](https://bmad-sim.github.io/GTPSA.jl/stable/man/d_tps/). ```@example overloads -print_overloads(AutoGTPSA(), :DifferentiationInterfaceGTPSAfExt) # hide +print_overloads(AutoGTPSA(), :DifferentiationInterfaceGTPSAExt) # hide ``` ## PolyesterForwardDiff diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl index dc5d90435..3cebee074 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl @@ -11,21 +11,20 @@ function DI.prepare_pushforward(f, backend::AutoGTPSA{D}, x, dx) where {D} end if x isa Number - xt = TPS{promote_type(typeof(dx),typeof(x),Float64)}(; use=d) + xt = TPS{promote_type(typeof(dx), typeof(x), Float64)}(; use=d) return GTPSAPushforwardExtras(xt) else - xt = similar(x, TPS{promote_type(eltype(dx),eltype(x),Float64)}) + xt = similar(x, TPS{promote_type(eltype(dx), eltype(x), Float64)}) # xt and x have same indexing because of similar for i in eachindex(xt) - xt[i] = TPS{promote_type(eltype(dx),eltype(x),Float64)}(; use=d) + xt[i] = TPS{promote_type(eltype(dx), eltype(x), Float64)}(; use=d) end return GTPSAPushforwardExtras(xt) end end function DI.pushforward(f, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras) - if x isa Number extras.xt[0] = x extras.xt[1] = dx @@ -41,31 +40,29 @@ function DI.pushforward(f, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardEx yt = f(extras.xt) if yt isa Number if dx isa Number - return yt[1] else dy = 0 - for j=1:length(dx) + for j in 1:length(dx) dy += yt[j] end - + return dy end else dy = similar(yt, eltype(eltype(yt))) dy .= 0 for i in eachindex(yt) - for j=1:length(dx) + for j in 1:length(dx) dy[i] += yt[i][j] end end - + return dy end end function DI.pushforward!(f, dy, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras) - if x isa Number extras.xt[0] = x extras.xt[1] = dx @@ -81,30 +78,30 @@ function DI.pushforward!(f, dy, backend::AutoGTPSA, x, dx, extras::GTPSAPushforw yt = f(extras.xt) if yt isa Number if dx isa Number - return yt[1] else dy = 0 - for j=1:length(dx) + for j in 1:length(dx) dy += yt[j] end - + return dy end else dy .= 0 for i in eachindex(yt) - for j=1:length(dx) + for j in 1:length(dx) dy[i] += yt[i][j] end end - + return dy end end -function DI.value_and_pushforward(f, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras) - +function DI.value_and_pushforward( + f, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras +) if x isa Number extras.xt[0] = x extras.xt[1] = dx @@ -120,32 +117,32 @@ function DI.value_and_pushforward(f, backend::AutoGTPSA, x, dx, extras::GTPSAPus yt = f(extras.xt) if yt isa Number if dx isa Number - return yt[0], yt[1] else dy = 0 - for j=1:length(dx) + for j in 1:length(dx) dy += yt[j] end - + return yt[0], dy end else dy = similar(yt, eltype(eltype(yt))) dy .= 0 for i in eachindex(yt) - for j=1:length(dx) + for j in 1:length(dx) dy[i] += yt[i][j] end end - y = map(t->t[0], yt) - + y = map(t -> t[0], yt) + return y, dy end end -function DI.value_and_pushforward!(f, dy, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras) - +function DI.value_and_pushforward!( + f, dy, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras +) if x isa Number extras.xt[0] = x extras.xt[1] = dx @@ -162,27 +159,26 @@ function DI.value_and_pushforward!(f, dy, backend::AutoGTPSA, x, dx, extras::GTP if yt isa Number if dx isa Number return yt[0], yt[1] - + else dy = 0 - for j=1:length(dx) + for j in 1:length(dx) dy += yt[j] end - + return yt[0], dy end else dy .= 0 for i in eachindex(yt) - for j=1:length(dx) + for j in 1:length(dx) dy[i] += yt[i][j] end end - y = map(t->t[0], yt) - + y = map(t -> t[0], yt) + return y, dy end - end ## Derivative @@ -197,7 +193,7 @@ function DI.prepare_derivative(f, backend::AutoGTPSA{D}, x) where {D} else d = Descriptor(1, 1) end - xt = TPS{promote_type(typeof(x),Float64)}(; use=d) + xt = TPS{promote_type(typeof(x), Float64)}(; use=d) xt[1] = 1 return GTPSADerivativeExtras(xt) end @@ -264,14 +260,14 @@ function DI.prepare_gradient(f, backend::AutoGTPSA{D}, x) where {D} nn = length(x) d = Descriptor(nn, 1) end - xt = similar(x, TPS{promote_type(eltype(x),Float64)}) + 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] = TPS{promote_type(eltype(x), Float64)}(; use=d) xt[i][j] = 1 j += 1 end @@ -325,14 +321,14 @@ function DI.prepare_jacobian(f, backend::AutoGTPSA{D}, x) where {D} nn = length(x) d = Descriptor(nn, 1) end - xt = similar(x, TPS{promote_type(eltype(x),Float64)}) + 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] = TPS{promote_type(eltype(x), Float64)}(; use=d) xt[i][j] = 1 j += 1 end @@ -384,7 +380,7 @@ function DI.prepare_second_derivative(f, backend::AutoGTPSA{D}, x) where {D} else d = Descriptor(1, 2) end - xt = TPS{promote_type(typeof(x),Float64)}(; use=d) + xt = TPS{promote_type(typeof(x), Float64)}(; use=d) xt[1] = 1 return GTPSASecondDerivativeExtras(xt) end @@ -393,11 +389,11 @@ function DI.second_derivative(f, ::AutoGTPSA, x, extras::GTPSASecondDerivativeEx extras.xt[0] = x yt = f(extras.xt) if yt isa Number - return yt[2]*2 + return yt[2] * 2 else der2 = similar(yt, eltype(eltype(yt))) for i in eachindex(yt) - der2[i] = yt[i][2]*2 # *2 because monomial coefficient is 1/2 + der2[i] = yt[i][2] * 2 # *2 because monomial coefficient is 1/2 end return der2 end @@ -407,7 +403,7 @@ function DI.second_derivative!(f, der2, ::AutoGTPSA, x, extras::GTPSASecondDeriv extras.xt[0] = x yt = f(extras.xt) for i in eachindex(yt) - der2[i] = yt[i][2]*2 + der2[i] = yt[i][2] * 2 end return der2 end @@ -418,14 +414,14 @@ function DI.value_derivative_and_second_derivative( extras.xt[0] = x yt = f(extras.xt) if yt isa Number - return yt[0], yt[1], yt[2]*2 + return yt[0], yt[1], yt[2] * 2 else y = map(t -> t[0], yt) der = similar(yt, eltype(eltype(yt))) der2 = similar(yt, eltype(eltype(yt))) for i in eachindex(yt) der[i] = yt[i][1] - der2[i] = yt[i][2]*2 + der2[i] = yt[i][2] * 2 end return y, der, der2 end @@ -439,7 +435,7 @@ function DI.value_derivative_and_second_derivative!( y = map(t -> t[0], yt) for i in eachindex(yt) der[i] = yt[i][1] - der2[i] = yt[i][2]*2 + der2[i] = yt[i][2] * 2 end return y, der, der2 end @@ -463,14 +459,14 @@ function DI.prepare_hessian(f, backend::AutoGTPSA{D}, x) where {D} # 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 = 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] = TPS{promote_type(eltype(x), Float64)}(; use=d) xt[i][j] = 1 j += 1 end @@ -533,167 +529,163 @@ function DI.prepare_hvp(f, backend::AutoGTPSA{D}, x, dx) where {D} # the indexing is known beforehand and we can do it (very slightly) faster m = nothing end - xt = similar(x, TPS{promote_type(eltype(dx),eltype(x),Float64)}) + xt = similar(x, TPS{promote_type(eltype(dx), 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(dx),eltype(x),Float64)}(; use=d) + xt[i] = TPS{promote_type(eltype(dx), eltype(x), Float64)}(; use=d) xt[i][j] = 1 j += 1 end - return GTPSAHVPExtras(xt, mono) + return GTPSAHVPExtras(xt, mono) end function DI.hvp(f, backend::AutoGTPSA{D}, x, dx, extras::GTPSAHVPExtras) where {D} foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part - + yt = f(extras.xt) - dg = similar(x, eltype(eltype(yt))) dg .= 0 - d = GTPSA.getdesc(yt) + d = GTPSA.getdesc(yt) desc = unsafe_load(d.desc) nn = desc.nn - if D == Nothing - idx = desc.nv+desc.np - endidx = floor(nn*(nn+1)/2)+nn + idx = desc.nv + desc.np + endidx = floor(nn * (nn + 1) / 2) + nn curdiag = 1 col = 1 xt = Ref{eltype(yt)}() idx = cycle!(yt, idx, 0, C_NULL, xt) while idx <= endidx && idx > 0 - h_idx = idx-nn + h_idx = idx - nn while h_idx > curdiag - col += 1 - curdiag += col + col += 1 + curdiag += col end - row = col-(curdiag-h_idx) + row = col - (curdiag - h_idx) #println("row = ", row, ", col = ", col) - if row==col - dg[row] += 2*xt[]*dx[col] + if row == col + dg[row] += 2 * xt[] * dx[col] else - dg[row] += xt[]*dx[col] - dg[col] += xt[]*dx[row] + dg[row] += xt[] * dx[col] + dg[col] += xt[] * dx[row] end idx = cycle!(yt, idx, 0, C_NULL, xt) end else - # If there are some variables/parameters with TO == 1, we have to do it "slow" - # because the indexing of TPSA index -> hessian index can be very complicated. - # I saw slow in quotes because it is likely still much faster than the calculation - # of the Hessian itself (this is just a getter) - idx = desc.nv+desc.np # start at 2nd order - xt = Ref{eltype(yt)}() - mono = extras.mono - idx = cycle!(yt, idx, nn, mono, xt) - while idx > 0 - if sum(mono) > 0x2 - return dg - end - i = findfirst(x->x==0x1, mono) - if isnothing(i) - i = findfirst(x->x==0x2, mono) - if isnothing(i) - return dg - end - if i <= nn - dg[i] += 2*xt[]*dx[i] # Multiply by 2 because taylor coefficient on diagonal is 1/2!*d2f/dx2 - end - else - j = findlast(x->x==0x1, mono) - if isnothing(j) - return dg - end - if i <= nn && j <= nn - dg[i] += xt[]*dx[j] - dg[j] += xt[]*dx[i] - end - end + # If there are some variables/parameters with TO == 1, we have to do it "slow" + # because the indexing of TPSA index -> hessian index can be very complicated. + # I saw slow in quotes because it is likely still much faster than the calculation + # of the Hessian itself (this is just a getter) + idx = desc.nv + desc.np # start at 2nd order + xt = Ref{eltype(yt)}() + mono = extras.mono idx = cycle!(yt, idx, nn, mono, xt) - end + while idx > 0 + if sum(mono) > 0x2 + return dg + end + i = findfirst(x -> x == 0x1, mono) + if isnothing(i) + i = findfirst(x -> x == 0x2, mono) + if isnothing(i) + return dg + end + if i <= nn + dg[i] += 2 * xt[] * dx[i] # Multiply by 2 because taylor coefficient on diagonal is 1/2!*d2f/dx2 + end + else + j = findlast(x -> x == 0x1, mono) + if isnothing(j) + return dg + end + if i <= nn && j <= nn + dg[i] += xt[] * dx[j] + dg[j] += xt[] * dx[i] + end + end + idx = cycle!(yt, idx, nn, mono, xt) + end end return dg end function DI.hvp!(f, dg, backend::AutoGTPSA{D}, x, dx, extras::GTPSAHVPExtras) where {D} foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part - + yt = f(extras.xt) - #dg = similar(x, eltype(eltype(yt))) dg .= 0 - d = GTPSA.getdesc(yt) + d = GTPSA.getdesc(yt) desc = unsafe_load(d.desc) nn = desc.nn - if D == Nothing - idx = desc.nv+desc.np - endidx = floor(nn*(nn+1)/2)+nn + idx = desc.nv + desc.np + endidx = floor(nn * (nn + 1) / 2) + nn curdiag = 1 col = 1 xt = Ref{eltype(yt)}() idx = cycle!(yt, idx, 0, C_NULL, xt) while idx <= endidx && idx > 0 - h_idx = idx-nn + h_idx = idx - nn while h_idx > curdiag - col += 1 - curdiag += col + col += 1 + curdiag += col end - row = col-(curdiag-h_idx) + row = col - (curdiag - h_idx) #println("row = ", row, ", col = ", col) - if row==col - dg[row] += 2*xt[]*dx[col] + if row == col + dg[row] += 2 * xt[] * dx[col] else - dg[row] += xt[]*dx[col] - dg[col] += xt[]*dx[row] + dg[row] += xt[] * dx[col] + dg[col] += xt[] * dx[row] end idx = cycle!(yt, idx, 0, C_NULL, xt) end else - # If there are some variables/parameters with TO == 1, we have to do it "slow" - # because the indexing of TPSA index -> hessian index can be very complicated. - # I saw slow in quotes because it is likely still much faster than the calculation - # of the Hessian itself (this is just a getter) - idx = desc.nv+desc.np # start at 2nd order - xt = Ref{eltype(yt)}() - mono = extras.mono - idx = cycle!(yt, idx, nn, mono, xt) - while idx > 0 - if sum(mono) > 0x2 - return dg - end - i = findfirst(x->x==0x1, mono) - if isnothing(i) - i = findfirst(x->x==0x2, mono) - if isnothing(i) - return dg - end - if i <= nn - dg[i] += 2*xt[]*dx[i] # Multiply by 2 because taylor coefficient on diagonal is 1/2!*d2f/dx2 - end - else - j = findlast(x->x==0x1, mono) - if isnothing(j) - return dg - end - if i <= nn && j <= nn - dg[i] += xt[]*dx[j] - dg[j] += xt[]*dx[i] - end - end + # If there are some variables/parameters with TO == 1, we have to do it "slow" + # because the indexing of TPSA index -> hessian index can be very complicated. + # I saw slow in quotes because it is likely still much faster than the calculation + # of the Hessian itself (this is just a getter) + idx = desc.nv + desc.np # start at 2nd order + xt = Ref{eltype(yt)}() + mono = extras.mono idx = cycle!(yt, idx, nn, mono, xt) - end + while idx > 0 + if sum(mono) > 0x2 + return dg + end + i = findfirst(x -> x == 0x1, mono) + if isnothing(i) + i = findfirst(x -> x == 0x2, mono) + if isnothing(i) + return dg + end + if i <= nn + dg[i] += 2 * xt[] * dx[i] # Multiply by 2 because taylor coefficient on diagonal is 1/2!*d2f/dx2 + end + else + j = findlast(x -> x == 0x1, mono) + if isnothing(j) + return dg + end + if i <= nn && j <= nn + dg[i] += xt[] * dx[j] + dg[j] += xt[] * dx[i] + end + end + idx = cycle!(yt, idx, nn, mono, xt) + end end return dg -end \ No newline at end of file +end diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl index 1d1c28a2c..43eed68dc 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl @@ -12,31 +12,32 @@ function DI.prepare_pushforward(f!, y, backend::AutoGTPSA{D}, x, dx) where {D} end if x isa Number - xt = TPS{promote_type(typeof(dx),typeof(x),Float64)}(; use=d) + xt = TPS{promote_type(typeof(dx), typeof(x), Float64)}(; use=d) else - xt = similar(x, TPS{promote_type(eltype(dx),eltype(x),Float64)}) + xt = similar(x, TPS{promote_type(eltype(dx), eltype(x), Float64)}) # xt and x have same indexing because of similar for i in eachindex(xt) - xt[i] = TPS{promote_type(eltype(dx),eltype(x),Float64)}(; use=d) + xt[i] = TPS{promote_type(eltype(dx), eltype(x), Float64)}(; use=d) end end if y isa Number - yt = TPS{promote_type(typeof(y),Float64)}(; use=d) + yt = TPS{promote_type(typeof(y), Float64)}(; use=d) else - yt = similar(y, TPS{promote_type(eltype(y),Float64)}) + yt = similar(y, TPS{promote_type(eltype(y), Float64)}) for i in eachindex(yt) - yt[i] = TPS{promote_type(eltype(y),Float64)}(; use=d) + yt[i] = TPS{promote_type(eltype(y), Float64)}(; use=d) end end return GTPSATwoArgPushforwardExtras(xt, yt) end -function DI.pushforward(f!, y, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPushforwardExtras) - +function DI.pushforward( + f!, y, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPushforwardExtras +) if x isa Number extras.xt[0] = x extras.xt[1] = dx @@ -52,14 +53,13 @@ function DI.pushforward(f!, y, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPus f!(extras.yt, extras.xt) if extras.yt isa Number if dx isa Number - return extras.yt[1] else dy = 0 - for j=1:length(dx) + for j in 1:length(dx) dy += extras.yt[j] end - + return dy end else @@ -67,17 +67,18 @@ function DI.pushforward(f!, y, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPus dy = similar(extras.yt, eltype(eltype(extras.yt))) dy .= 0 for i in eachindex(extras.yt) - for j=1:length(dx) + for j in 1:length(dx) dy[i] += extras.yt[i][j] end end - + return dy end end -function DI.pushforward!(f!, y, dy, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPushforwardExtras) - +function DI.pushforward!( + f!, y, dy, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPushforwardExtras +) if x isa Number extras.xt[0] = x extras.xt[1] = dx @@ -93,31 +94,31 @@ function DI.pushforward!(f!, y, dy, backend::AutoGTPSA, x, dx, extras::GTPSATwoA f!(extras.yt, extras.xt) if extras.yt isa Number if dx isa Number - return extras.yt[1] else dy = 0 - for j=1:length(dx) + for j in 1:length(dx) dy += extras.yt[j] end - + return extras.dy end else map!(t -> t[0], y, extras.yt) dy .= 0 for i in eachindex(extras.yt) - for j=1:length(dx) + for j in 1:length(dx) dy[i] += extras.yt[i][j] end end - + return dy end end -function DI.value_and_pushforward(f!, y, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPushforwardExtras) - +function DI.value_and_pushforward( + f!, y, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPushforwardExtras +) if x isa Number extras.xt[0] = x extras.xt[1] = dx @@ -133,32 +134,32 @@ function DI.value_and_pushforward(f!, y, backend::AutoGTPSA, x, dx, extras::GTPS f!(extras.yt, extras.xt) if extras.yt isa Number if dx isa Number - return extras.yt[0], extras.yt[1] else dy = 0 - for j=1:length(dx) + for j in 1:length(dx) dy += extras.yt[j] end - + return extras.yt[0], dy end else dy = similar(extras.yt, eltype(eltype(extras.yt))) dy .= 0 for i in eachindex(extras.yt) - for j=1:length(dx) + for j in 1:length(dx) dy[i] += extras.yt[i][j] end end map!(t -> t[0], y, extras.yt) - + return y, dy end end -function DI.value_and_pushforward!(f!, y, dy, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPushforwardExtras) - +function DI.value_and_pushforward!( + f!, y, dy, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPushforwardExtras +) if x isa Number extras.xt[0] = x extras.xt[1] = dx @@ -175,24 +176,24 @@ function DI.value_and_pushforward!(f!, y, dy, backend::AutoGTPSA, x, dx, extras: if extras.yt isa Number if dx isa Number return extras.yt[0], extras.yt[1] - + else dy = 0 - for j=1:length(dx) + for j in 1:length(dx) dy += yt[j] end - + return extras.yt[0], dy end else map!(t -> t[0], y, extras.yt) dy .= 0 for i in eachindex(extras.yt) - for j=1:length(dx) + for j in 1:length(dx) dy[i] += extras.yt[i][j] end end - + return y, dy end end @@ -210,16 +211,16 @@ function DI.prepare_derivative(f!, y, backend::AutoGTPSA{D}, x) where {D} else d = Descriptor(1, 1) end - xt = TPS{promote_type(typeof(x),Float64)}(; use=d) + xt = TPS{promote_type(typeof(x), Float64)}(; use=d) xt[1] = 1 if y isa Number - yt = TPS{promote_type(typeof(y),Float64)}(; use=d) + yt = TPS{promote_type(typeof(y), Float64)}(; use=d) else - yt = similar(y, TPS{promote_type(eltype(y),Float64)}) + yt = similar(y, TPS{promote_type(eltype(y), Float64)}) for i in eachindex(yt) - yt[i] = TPS{promote_type(eltype(y),Float64)}(; use=d) + yt[i] = TPS{promote_type(eltype(y), Float64)}(; use=d) end end @@ -266,7 +267,9 @@ function DI.value_and_derivative(f!, y, ::AutoGTPSA, x, extras::GTPSATwoArgDeriv end end -function DI.value_and_derivative!(f!, y, der, ::AutoGTPSA, x, extras::GTPSATwoArgDerivativeExtras) +function DI.value_and_derivative!( + f!, y, der, ::AutoGTPSA, x, extras::GTPSATwoArgDerivativeExtras +) extras.xt[0] = x f!(extras.yt, extras.xt) map!(t -> t[0], y, extras.yt) @@ -291,25 +294,25 @@ function DI.prepare_jacobian(f!, y, backend::AutoGTPSA{D}, x) where {D} nn = length(x) d = Descriptor(nn, 1) end - xt = similar(x, TPS{promote_type(eltype(x),Float64)}) + 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] = TPS{promote_type(eltype(x), Float64)}(; use=d) xt[i][j] = 1 j += 1 end if y isa Number - yt = TPS{promote_type(typeof(y),Float64)}(; use=d) + yt = TPS{promote_type(typeof(y), Float64)}(; use=d) else - yt = similar(y, TPS{promote_type(eltype(y),Float64)}) + yt = similar(y, TPS{promote_type(eltype(y), Float64)}) for i in eachindex(yt) - yt[i] = TPS{promote_type(eltype(y),Float64)}(; use=d) + yt[i] = TPS{promote_type(eltype(y), Float64)}(; use=d) end end @@ -342,7 +345,9 @@ function DI.value_and_jacobian(f!, y, ::AutoGTPSA, x, extras::GTPSATwoArgJacobia return y, jac end -function DI.value_and_jacobian!(f!, y, jac, ::AutoGTPSA, x, extras::GTPSATwoArgJacobianExtras) +function DI.value_and_jacobian!( + f!, y, jac, ::AutoGTPSA, x, extras::GTPSATwoArgJacobianExtras +) foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part f!(extras.yt, extras.xt) GTPSA.jacobian!(jac, extras.yt; include_params=true) diff --git a/DifferentiationInterface/test/Back/GTPSA/test.jl b/DifferentiationInterface/test/Back/GTPSA/test.jl index 87b7a0910..da98887ae 100644 --- a/DifferentiationInterface/test/Back/GTPSA/test.jl +++ b/DifferentiationInterface/test/Back/GTPSA/test.jl @@ -13,7 +13,4 @@ for backend in [AutoGTPSA()] @test check_hessian(backend) end -test_differentiation( - AutoGTPSA(); - type_stability=true, - logging=LOGGING); \ No newline at end of file +test_differentiation(AutoGTPSA(); type_stability=true, logging=LOGGING); From 5022e2e24fe7530a8ef584873a7a5e91af388fb5 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Thu, 8 Aug 2024 08:42:06 +0200 Subject: [PATCH 21/44] Exclude GTPSA from LTS testing --- .github/workflows/Test.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml index e89790e51..f734ea81c 100644 --- a/.github/workflows/Test.yml +++ b/.github/workflows/Test.yml @@ -65,6 +65,8 @@ jobs: group: Back/FiniteDiff - version: 'lts' group: Back/FastDifferentiation + - version: 'lts' + group: Back/GTPSA - version: 'lts' group: Back/PolyesterForwardDiff - version: 'lts' From 5f4c34ca587ef1667d9c66f2a84efe3bff070e79 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Thu, 8 Aug 2024 17:13:21 +0200 Subject: [PATCH 22/44] Reactivate coverage --- .github/workflows/Test.yml | 135 +++++++++--------- DifferentiationInterface/Project.toml | 2 +- .../test/Single/GTPSA/test.jl | 11 -- 3 files changed, 68 insertions(+), 80 deletions(-) delete mode 100644 DifferentiationInterface/test/Single/GTPSA/test.jl diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml index f734ea81c..fee322277 100644 --- a/.github/workflows/Test.yml +++ b/.github/workflows/Test.yml @@ -4,7 +4,6 @@ on: push: branches: - main - - addgtpsa tags: ['*'] pull_request: types: [opened, reopened, synchronize, ready_for_review] @@ -113,71 +112,71 @@ jobs: using Pkg; Pkg.Registry.update(); Pkg.update(); - Pkg.test("DifferentiationInterface"; coverage=false);' - #- uses: julia-actions/julia-processcoverage@v1 - # if: ${{ env.SHOULDRUN == 'true' }} - # with: - # directories: ./DifferentiationInterface/src,./DifferentiationInterface/ext,./DifferentiationInterface/test - #- uses: codecov/codecov-action@v4 - # if: ${{ env.SHOULDRUN == 'true' }} - # with: - # files: lcov.info - # token: ${{ secrets.CODECOV_TOKEN }} - # fail_ci_if_error: true + Pkg.test("DifferentiationInterface"; coverage=true);' + - uses: julia-actions/julia-processcoverage@v1 + if: ${{ env.SHOULDRUN == 'true' }} + with: + directories: ./DifferentiationInterface/src,./DifferentiationInterface/ext,./DifferentiationInterface/test + - uses: codecov/codecov-action@v4 + if: ${{ env.SHOULDRUN == 'true' }} + with: + files: lcov.info + token: ${{ secrets.CODECOV_TOKEN }} + fail_ci_if_error: true -# test-DIT: -# name: ${{ matrix.version }} - DIT (${{ matrix.group }}) -# runs-on: ubuntu-latest -# timeout-minutes: 60 -# permissions: # needed to allow julia-actions/cache to proactively delete old caches that it has created -# actions: write -# contents: read -# strategy: -# fail-fast: false -# matrix: -# version: -# - '1' -# - 'lts' -# - 'pre' -# group: -# - Formalities -# - Zero -# - ForwardDiff -# - Weird -# exclude: -# - version: 'lts' -# group: Formalities -# - version: 'lts' -# group: Weird -# env: -# SHOULDRUN: ${{ matrix.version == '1' || !github.event.pull_request.draft }} -# JULIA_DIT_TEST_GROUP: ${{ matrix.group }} -# steps: -# - run: echo "$SHOULDRUN" -# - uses: actions/checkout@v4 -# if: ${{ env.SHOULDRUN == 'true' }} -# - uses: julia-actions/setup-julia@v2 -# if: ${{ env.SHOULDRUN == 'true' }} -# with: -# version: ${{ matrix.version }} -# arch: x64 -# - uses: julia-actions/cache@v2 -# if: ${{ env.SHOULDRUN == 'true' }} -# - name: Install dependencies & run tests -# if: ${{ env.SHOULDRUN == 'true' }} -# run: julia --project=./DifferentiationInterfaceTest -e ' -# using Pkg; -# Pkg.Registry.update(); -# Pkg.update(); -# Pkg.develop(path="./DifferentiationInterface"); -# Pkg.test("DifferentiationInterfaceTest"; coverage=true);' -# - uses: julia-actions/julia-processcoverage@v1 -# if: ${{ env.SHOULDRUN == 'true' }} -# with: -# directories: ./DifferentiationInterfaceTest/src,./DifferentiationInterfaceTest/ext,./DifferentiationInterfaceTest/test -# - uses: codecov/codecov-action@v4 -# if: ${{ env.SHOULDRUN == 'true' }} -# with: -# files: lcov.info -# token: ${{ secrets.CODECOV_TOKEN }} -# fail_ci_if_error: true + test-DIT: + name: ${{ matrix.version }} - DIT (${{ matrix.group }}) + runs-on: ubuntu-latest + timeout-minutes: 60 + permissions: # needed to allow julia-actions/cache to proactively delete old caches that it has created + actions: write + contents: read + strategy: + fail-fast: false + matrix: + version: + - '1' + - 'lts' + - 'pre' + group: + - Formalities + - Zero + - ForwardDiff + - Weird + exclude: + - version: 'lts' + group: Formalities + - version: 'lts' + group: Weird + env: + SHOULDRUN: ${{ matrix.version == '1' || !github.event.pull_request.draft }} + JULIA_DIT_TEST_GROUP: ${{ matrix.group }} + steps: + - run: echo "$SHOULDRUN" + - uses: actions/checkout@v4 + if: ${{ env.SHOULDRUN == 'true' }} + - uses: julia-actions/setup-julia@v2 + if: ${{ env.SHOULDRUN == 'true' }} + with: + version: ${{ matrix.version }} + arch: x64 + - uses: julia-actions/cache@v2 + if: ${{ env.SHOULDRUN == 'true' }} + - name: Install dependencies & run tests + if: ${{ env.SHOULDRUN == 'true' }} + run: julia --project=./DifferentiationInterfaceTest -e ' + using Pkg; + Pkg.Registry.update(); + Pkg.update(); + Pkg.develop(path="./DifferentiationInterface"); + Pkg.test("DifferentiationInterfaceTest"; coverage=true);' + - uses: julia-actions/julia-processcoverage@v1 + if: ${{ env.SHOULDRUN == 'true' }} + with: + directories: ./DifferentiationInterfaceTest/src,./DifferentiationInterfaceTest/ext,./DifferentiationInterfaceTest/test + - uses: codecov/codecov-action@v4 + if: ${{ env.SHOULDRUN == 'true' }} + with: + files: lcov.info + token: ${{ secrets.CODECOV_TOKEN }} + fail_ci_if_error: true \ No newline at end of file diff --git a/DifferentiationInterface/Project.toml b/DifferentiationInterface/Project.toml index 5615d4944..a79b55764 100644 --- a/DifferentiationInterface/Project.toml +++ b/DifferentiationInterface/Project.toml @@ -77,6 +77,7 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" +# DifferentiationInterfaceTest = "a82114a7-5aa3-49a8-9643-716bb13727a3" Diffractor = "9f5e2b26-1114-432f-b630-d3fe2085c51c" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" FastDifferentiation = "eb9bf01b-bf85-4b60-bf87-ee5de06c00be" @@ -122,4 +123,3 @@ test = [ "StaticArrays", "Test", ] - diff --git a/DifferentiationInterface/test/Single/GTPSA/test.jl b/DifferentiationInterface/test/Single/GTPSA/test.jl deleted file mode 100644 index 3f1c5febf..000000000 --- a/DifferentiationInterface/test/Single/GTPSA/test.jl +++ /dev/null @@ -1,11 +0,0 @@ -using DifferentiationInterface, DifferentiationInterfaceTest -using GTPSA: GTPSA -using Test - -for backend in [AutoGTPSA()] - @test check_available(backend) - #@test check_twoarg(backend) - #@test check_hessian(backend) -end - -test_differentiation(AutoGTPSA(); logging=LOGGING); From 037df4015addc0dc5e46cec90fffa82bbe44ea9d Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Thu, 8 Aug 2024 17:15:04 +0200 Subject: [PATCH 23/44] Reactivate all tests --- .github/workflows/Test.yml | 34 +++++++++++++++++----------------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml index fee322277..f9e59470e 100644 --- a/.github/workflows/Test.yml +++ b/.github/workflows/Test.yml @@ -34,24 +34,24 @@ jobs: group: - Formalities - Internals - #- Back/Diffractor - #- Back/Enzyme - #- Back/FastDifferentiation - #- Back/FiniteDiff - #- Back/FiniteDifferences - #- Back/ForwardDiff + - Back/Diffractor + - Back/Enzyme + - Back/FastDifferentiation + - Back/FiniteDiff + - Back/FiniteDifferences + - Back/ForwardDiff - Back/GTPSA - #- Back/PolyesterForwardDiff - #- Back/ReverseDiff - #- Back/SecondOrder - #- Back/Symbolics - #- Back/Tapir - #- Back/Tracker - #- Back/Zygote - #- Down/Detector - #- Down/DifferentiateWith - #- Down/Flux - #- Down/Lux + - Back/PolyesterForwardDiff + - Back/ReverseDiff + - Back/SecondOrder + - Back/Symbolics + - Back/Tapir + - Back/Tracker + - Back/Zygote + - Down/Detector + - Down/DifferentiateWith + - Down/Flux + - Down/Lux exclude: # lts - version: 'lts' From c7e1382a0b3e805beaa564817a207b2f0ed89169 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Thu, 8 Aug 2024 17:17:13 +0200 Subject: [PATCH 24/44] Test group --- DifferentiationInterface/test/runtests.jl | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/DifferentiationInterface/test/runtests.jl b/DifferentiationInterface/test/runtests.jl index ebe50a323..a4097054c 100644 --- a/DifferentiationInterface/test/runtests.jl +++ b/DifferentiationInterface/test/runtests.jl @@ -9,8 +9,7 @@ else Pkg.add("DifferentiationInterfaceTest") end -#GROUP = get(ENV, "JULIA_DI_TEST_GROUP", "All") -GROUP = get(ENV, "JULIA_DI_TEST_GROUP", "Back/GTPSA") +GROUP = get(ENV, "JULIA_DI_TEST_GROUP", "All") ## Main tests From 65d5a4ea8556f6a90f4eb42afa27a5ef4fd7f7c1 Mon Sep 17 00:00:00 2001 From: mattsignorelli Date: Fri, 9 Aug 2024 15:29:31 -0400 Subject: [PATCH 25/44] simplified code and removed derivative --- DifferentiationInterface/Project.toml | 2 +- .../onearg.jl | 315 +++--------------- .../twoarg.jl | 211 ++---------- 3 files changed, 79 insertions(+), 449 deletions(-) diff --git a/DifferentiationInterface/Project.toml b/DifferentiationInterface/Project.toml index a79b55764..0296fbb74 100644 --- a/DifferentiationInterface/Project.toml +++ b/DifferentiationInterface/Project.toml @@ -57,7 +57,7 @@ FillArrays = "1.7.0" FiniteDiff = "2.23.1" FiniteDifferences = "0.12.31" ForwardDiff = "0.10.36" -GTPSA = "1.1.0" +GTPSA = "1.1.1" LinearAlgebra = "<0.0.1,1" PackageExtensionCompat = "1.0.2" PolyesterForwardDiff = "0.1.1" diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl index 3cebee074..266d222a0 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl @@ -76,27 +76,14 @@ function DI.pushforward!(f, dy, backend::AutoGTPSA, x, dx, extras::GTPSAPushforw end yt = f(extras.xt) - if yt isa Number - if dx isa Number - return yt[1] - else - dy = 0 - for j in 1:length(dx) - dy += yt[j] - end - - return dy - end - else - dy .= 0 - for i in eachindex(yt) - for j in 1:length(dx) - dy[i] += yt[i][j] - end + dy .= 0 + for i in eachindex(yt) + for j in 1:length(dx) + dy[i] += yt[i][j] end - - return dy end + + return dy end function DI.value_and_pushforward( @@ -156,94 +143,15 @@ function DI.value_and_pushforward!( end yt = f(extras.xt) - if yt isa Number - if dx isa Number - return yt[0], yt[1] - - else - dy = 0 - for j in 1:length(dx) - dy += yt[j] - end - - return yt[0], dy - end - else - dy .= 0 - for i in eachindex(yt) - for j in 1:length(dx) - dy[i] += yt[i][j] - end - end - y = map(t -> t[0], yt) - - return y, dy - end -end - -## Derivative - -struct GTPSADerivativeExtras{X} <: DerivativeExtras - xt::X -end - -function DI.prepare_derivative(f, backend::AutoGTPSA{D}, x) where {D} - if D != Nothing - d = backend.descriptor - else - d = Descriptor(1, 1) - end - xt = TPS{promote_type(typeof(x), Float64)}(; use=d) - xt[1] = 1 - return GTPSADerivativeExtras(xt) -end - -function DI.derivative(f, ::AutoGTPSA, x, extras::GTPSADerivativeExtras) - extras.xt[0] = x - yt = f(extras.xt) - if yt isa Number - return yt[1] - else - der = similar(yt, eltype(eltype(yt))) - for i in eachindex(yt) - der[i] = yt[i][1] - end - return der - end -end - -function DI.derivative!(f, der, ::AutoGTPSA, x, extras::GTPSADerivativeExtras) - extras.xt[0] = x - yt = f(extras.xt) + dy .= 0 for i in eachindex(yt) - der[i] = yt[i][1] - end - return der -end - -function DI.value_and_derivative(f, ::AutoGTPSA, x, extras::GTPSADerivativeExtras) - extras.xt[0] = x - yt = f(extras.xt) - if yt isa Number - return yt[0], yt[1] - else - y = map(t -> t[0], yt) - der = similar(yt, eltype(eltype(yt))) - for i in eachindex(yt) - der[i] = yt[i][1] + for j in 1:length(dx) + dy[i] += yt[i][j] end - return y, der end -end - -function DI.value_and_derivative!(f, der, ::AutoGTPSA, x, extras::GTPSADerivativeExtras) - extras.xt[0] = x - yt = f(extras.xt) y = map(t -> t[0], yt) - for i in eachindex(yt) - der[i] = yt[i][1] - end - return y, der + + return y, dy end ## Gradient @@ -474,218 +382,83 @@ function DI.prepare_hessian(f, backend::AutoGTPSA{D}, x) where {D} return GTPSAHessianExtras(xt, m) end -function DI.hessian(f, ::AutoGTPSA, x, extras::GTPSAHessianExtras) +function DI.hessian(f, ::AutoGTPSA{D}, x, extras::GTPSAHessianExtras) where {D} foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part yt = f(extras.xt) hess = similar(x, eltype(eltype(yt)), (length(x), length(x))) - GTPSA.hessian!(hess, yt; include_params=true) + unsafe_fast = D == Nothing ? true : false + GTPSA.hessian!(hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=extras.m) return hess end -function DI.hessian!(f, hess, ::AutoGTPSA, x, extras::GTPSAHessianExtras) +function DI.hessian!(f, hess, ::AutoGTPSA{D}, x, extras::GTPSAHessianExtras) where {D} foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part yt = f(extras.xt) - GTPSA.hessian!(hess, yt; include_params=true) + unsafe_fast = D == Nothing ? true : false + GTPSA.hessian!(hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=extras.m) return hess end -function DI.value_gradient_and_hessian(f, ::AutoGTPSA, x, extras::GTPSAHessianExtras) +function DI.value_gradient_and_hessian(f, ::AutoGTPSA{D}, x, extras::GTPSAHessianExtras) where {D} foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part yt = f(extras.xt) y = map(t -> t[0], yt) grad = similar(x, eltype(eltype(yt))) GTPSA.gradient!(grad, yt; include_params=true) hess = similar(x, eltype(eltype(yt)), (length(x), length(x))) - GTPSA.hessian!(hess, yt; include_params=true) + unsafe_fast = D == Nothing ? true : false + GTPSA.hessian!(hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=extras.m) return y, grad, hess end function DI.value_gradient_and_hessian!( - f, grad, hess, ::AutoGTPSA, x, extras::GTPSAHessianExtras -) + f, grad, hess, ::AutoGTPSA{D}, x, extras::GTPSAHessianExtras +) where {D} foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part yt = f(extras.xt) y = map(t -> t[0], yt) GTPSA.gradient!(grad, yt; include_params=true) - GTPSA.hessian!(hess, yt; include_params=true) + unsafe_fast = D == Nothing ? true : false + GTPSA.hessian!(hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=extras.m) return y, grad, hess end # HVP -struct GTPSAHVPExtras{X,M} <: HVPExtras - xt::X - m::M +struct GTPSAHVPExtras{E,H} <: HVPExtras + hessextras::E + hess::H end function DI.prepare_hvp(f, backend::AutoGTPSA{D}, x, dx) where {D} - if D != Nothing - d = backend.descriptor - nn = GTPSA.numnn(d) - m = Vector{UInt8}(undef, nn) - 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(dx), 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(dx), eltype(x), Float64)}(; use=d) - xt[i][j] = 1 - j += 1 - end - - return GTPSAHVPExtras(xt, mono) + hessextras = DI.prepare_hessian(f, backend, x) + hess = similar(x, typeof(f(x)), (length(x), length(x))) + return GTPSAHVPExtras(hessextras, hess) end function DI.hvp(f, backend::AutoGTPSA{D}, x, dx, extras::GTPSAHVPExtras) where {D} - foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part - - yt = f(extras.xt) - - dg = similar(x, eltype(eltype(yt))) + DI.hessian!(f, extras.hess, backend, x, extras.hessextras) + dg = similar(x, eltype(extras.hess)) dg .= 0 - - d = GTPSA.getdesc(yt) - desc = unsafe_load(d.desc) - nn = desc.nn - - if D == Nothing - idx = desc.nv + desc.np - endidx = floor(nn * (nn + 1) / 2) + nn - curdiag = 1 - col = 1 - xt = Ref{eltype(yt)}() - idx = cycle!(yt, idx, 0, C_NULL, xt) - while idx <= endidx && idx > 0 - h_idx = idx - nn - while h_idx > curdiag - col += 1 - curdiag += col - end - row = col - (curdiag - h_idx) - #println("row = ", row, ", col = ", col) - if row == col - dg[row] += 2 * xt[] * dx[col] - else - dg[row] += xt[] * dx[col] - dg[col] += xt[] * dx[row] - end - idx = cycle!(yt, idx, 0, C_NULL, xt) - end - else - # If there are some variables/parameters with TO == 1, we have to do it "slow" - # because the indexing of TPSA index -> hessian index can be very complicated. - # I saw slow in quotes because it is likely still much faster than the calculation - # of the Hessian itself (this is just a getter) - idx = desc.nv + desc.np # start at 2nd order - xt = Ref{eltype(yt)}() - mono = extras.mono - idx = cycle!(yt, idx, nn, mono, xt) - while idx > 0 - if sum(mono) > 0x2 - return dg - end - i = findfirst(x -> x == 0x1, mono) - if isnothing(i) - i = findfirst(x -> x == 0x2, mono) - if isnothing(i) - return dg - end - if i <= nn - dg[i] += 2 * xt[] * dx[i] # Multiply by 2 because taylor coefficient on diagonal is 1/2!*d2f/dx2 - end - else - j = findlast(x -> x == 0x1, mono) - if isnothing(j) - return dg - end - if i <= nn && j <= nn - dg[i] += xt[] * dx[j] - dg[j] += xt[] * dx[i] - end - end - idx = cycle!(yt, idx, nn, mono, xt) + j = 1 + for dxi in dx + for i in 1:size(extras.hess,2) + dg[i] += extras.hess[i,j]*dxi end + j += 1 end return dg end function DI.hvp!(f, dg, backend::AutoGTPSA{D}, x, dx, extras::GTPSAHVPExtras) where {D} - foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part - - yt = f(extras.xt) - - #dg = similar(x, eltype(eltype(yt))) + DI.hessian!(f, extras.hess, backend, x, extras.hessextras) dg .= 0 - - d = GTPSA.getdesc(yt) - desc = unsafe_load(d.desc) - nn = desc.nn - - if D == Nothing - idx = desc.nv + desc.np - endidx = floor(nn * (nn + 1) / 2) + nn - curdiag = 1 - col = 1 - xt = Ref{eltype(yt)}() - idx = cycle!(yt, idx, 0, C_NULL, xt) - while idx <= endidx && idx > 0 - h_idx = idx - nn - while h_idx > curdiag - col += 1 - curdiag += col - end - row = col - (curdiag - h_idx) - #println("row = ", row, ", col = ", col) - if row == col - dg[row] += 2 * xt[] * dx[col] - else - dg[row] += xt[] * dx[col] - dg[col] += xt[] * dx[row] - end - idx = cycle!(yt, idx, 0, C_NULL, xt) - end - else - # If there are some variables/parameters with TO == 1, we have to do it "slow" - # because the indexing of TPSA index -> hessian index can be very complicated. - # I saw slow in quotes because it is likely still much faster than the calculation - # of the Hessian itself (this is just a getter) - idx = desc.nv + desc.np # start at 2nd order - xt = Ref{eltype(yt)}() - mono = extras.mono - idx = cycle!(yt, idx, nn, mono, xt) - while idx > 0 - if sum(mono) > 0x2 - return dg - end - i = findfirst(x -> x == 0x1, mono) - if isnothing(i) - i = findfirst(x -> x == 0x2, mono) - if isnothing(i) - return dg - end - if i <= nn - dg[i] += 2 * xt[] * dx[i] # Multiply by 2 because taylor coefficient on diagonal is 1/2!*d2f/dx2 - end - else - j = findlast(x -> x == 0x1, mono) - if isnothing(j) - return dg - end - if i <= nn && j <= nn - dg[i] += xt[] * dx[j] - dg[j] += xt[] * dx[i] - end - end - idx = cycle!(yt, idx, nn, mono, xt) + j = 1 + for dxi in dx + for i in 1:size(extras.hess,2) + dg[i] += extras.hess[i,j]*dxi end + j += 1 end return dg end + diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl index 43eed68dc..d216b1bee 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl @@ -22,16 +22,11 @@ function DI.prepare_pushforward(f!, y, backend::AutoGTPSA{D}, x, dx) where {D} end end - if y isa Number - yt = TPS{promote_type(typeof(y), Float64)}(; use=d) - else - yt = similar(y, TPS{promote_type(eltype(y), Float64)}) + 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 + for i in eachindex(yt) + yt[i] = TPS{promote_type(eltype(y), Float64)}(; use=d) end - return GTPSATwoArgPushforwardExtras(xt, yt) end @@ -51,29 +46,16 @@ function DI.pushforward( end f!(extras.yt, extras.xt) - if extras.yt isa Number - if dx isa Number - return extras.yt[1] - else - dy = 0 - for j in 1:length(dx) - dy += extras.yt[j] - end - - return dy - end - else - map!(t -> t[0], y, extras.yt) - dy = similar(extras.yt, eltype(eltype(extras.yt))) - dy .= 0 - for i in eachindex(extras.yt) - for j in 1:length(dx) - dy[i] += extras.yt[i][j] - end + map!(t -> t[0], y, extras.yt) + dy = similar(extras.yt, eltype(eltype(extras.yt))) + dy .= 0 + for i in eachindex(extras.yt) + for j in 1:length(dx) + dy[i] += extras.yt[i][j] end - - return dy end + + return dy end function DI.pushforward!( @@ -92,28 +74,15 @@ function DI.pushforward!( end f!(extras.yt, extras.xt) - if extras.yt isa Number - if dx isa Number - return extras.yt[1] - else - dy = 0 - for j in 1:length(dx) - dy += extras.yt[j] - end - - return extras.dy - end - else - map!(t -> t[0], y, extras.yt) - dy .= 0 - for i in eachindex(extras.yt) - for j in 1:length(dx) - dy[i] += extras.yt[i][j] - end + map!(t -> t[0], y, extras.yt) + dy .= 0 + for i in eachindex(extras.yt) + for j in 1:length(dx) + dy[i] += extras.yt[i][j] end - - return dy end + + return dy end function DI.value_and_pushforward( @@ -132,29 +101,16 @@ function DI.value_and_pushforward( end f!(extras.yt, extras.xt) - if extras.yt isa Number - if dx isa Number - return extras.yt[0], extras.yt[1] - else - dy = 0 - for j in 1:length(dx) - dy += extras.yt[j] - end - - return extras.yt[0], dy - end - else - dy = similar(extras.yt, eltype(eltype(extras.yt))) - dy .= 0 - for i in eachindex(extras.yt) - for j in 1:length(dx) - dy[i] += extras.yt[i][j] - end + dy = similar(extras.yt, eltype(eltype(extras.yt))) + dy .= 0 + for i in eachindex(extras.yt) + for j in 1:length(dx) + dy[i] += extras.yt[i][j] end - map!(t -> t[0], y, extras.yt) - - return y, dy end + map!(t -> t[0], y, extras.yt) + + return y, dy end function DI.value_and_pushforward!( @@ -173,110 +129,15 @@ function DI.value_and_pushforward!( end f!(extras.yt, extras.xt) - if extras.yt isa Number - if dx isa Number - return extras.yt[0], extras.yt[1] - - else - dy = 0 - for j in 1:length(dx) - dy += yt[j] - end - - return extras.yt[0], dy - end - else - map!(t -> t[0], y, extras.yt) - dy .= 0 - for i in eachindex(extras.yt) - for j in 1:length(dx) - dy[i] += extras.yt[i][j] - end - end - - return y, dy - end -end - -## Derivative - -struct GTPSATwoArgDerivativeExtras{X,Y} <: DerivativeExtras - xt::X - yt::Y -end - -function DI.prepare_derivative(f!, y, backend::AutoGTPSA{D}, x) where {D} - if D != Nothing - d = backend.descriptor - else - d = Descriptor(1, 1) - end - xt = TPS{promote_type(typeof(x), Float64)}(; use=d) - xt[1] = 1 - - if y isa Number - yt = TPS{promote_type(typeof(y), Float64)}(; use=d) - else - 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 - end - - return GTPSATwoArgDerivativeExtras(xt, yt) -end - -function DI.derivative(f!, y, ::AutoGTPSA, x, extras::GTPSATwoArgDerivativeExtras) - extras.xt[0] = x - f!(extras.yt, extras.xt) - if extras.yt isa Number - return extras.yt[1] - else - map!(t -> t[0], y, extras.yt) - der = similar(extras.yt, eltype(eltype(extras.yt))) - for i in eachindex(extras.yt) - der[i] = extras.yt[i][1] - end - return der - end -end - -function DI.derivative!(f!, y, der, ::AutoGTPSA, x, extras::GTPSATwoArgDerivativeExtras) - extras.xt[0] = x - f!(extras.yt, extras.xt) - for i in eachindex(extras.yt) - der[i] = extras.yt[i][1] - end map!(t -> t[0], y, extras.yt) - return der -end - -function DI.value_and_derivative(f!, y, ::AutoGTPSA, x, extras::GTPSATwoArgDerivativeExtras) - extras.xt[0] = x - f!(extras.yt, extras.xt) - if extras.yt isa Number - return extras.yt[0], extras.yt[1] - else - map!(t -> t[0], y, extras.yt) - der = similar(extras.yt, eltype(eltype(extras.yt))) - for i in eachindex(extras.yt) - der[i] = extras.yt[i][1] + dy .= 0 + for i in eachindex(extras.yt) + for j in 1:length(dx) + dy[i] += extras.yt[i][j] end - return y, der end -end -function DI.value_and_derivative!( - f!, y, der, ::AutoGTPSA, x, extras::GTPSATwoArgDerivativeExtras -) - extras.xt[0] = x - f!(extras.yt, extras.xt) - map!(t -> t[0], y, extras.yt) - for i in eachindex(extras.yt) - der[i] = extras.yt[i][1] - end - return y, der + return y, dy end ## Jacobian @@ -306,14 +167,10 @@ function DI.prepare_jacobian(f!, y, backend::AutoGTPSA{D}, x) where {D} j += 1 end - if y isa Number - yt = TPS{promote_type(typeof(y), Float64)}(; use=d) - else - yt = similar(y, TPS{promote_type(eltype(y), Float64)}) + 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 + for i in eachindex(yt) + yt[i] = TPS{promote_type(eltype(y), Float64)}(; use=d) end return GTPSATwoArgJacobianExtras(xt, yt) From 8a3a23978cdfa4d3bd2d10e72a9d917adf971fba Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Tue, 17 Sep 2024 13:24:22 +0200 Subject: [PATCH 26/44] Adapt to new Tangents --- .../onearg.jl | 306 ++++++++++-------- .../twoarg.jl | 187 ++++++----- .../test/Back/GTPSA/test.jl | 3 +- 3 files changed, 280 insertions(+), 216 deletions(-) diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl index 266d222a0..f19ce2e7d 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl @@ -3,13 +3,14 @@ struct GTPSAPushforwardExtras{X} <: PushforwardExtras xt::X end -function DI.prepare_pushforward(f, backend::AutoGTPSA{D}, x, dx) where {D} +function DI.prepare_pushforward(f, backend::AutoGTPSA{D}, x, tx::Tangents) where {D} if D != Nothing d = backend.descriptor else d = Descriptor(length(x), 1) end + dx = first(tx) if x isa Number xt = TPS{promote_type(typeof(dx), typeof(x), Float64)}(; use=d) return GTPSAPushforwardExtras(xt) @@ -24,134 +25,151 @@ function DI.prepare_pushforward(f, backend::AutoGTPSA{D}, x, dx) where {D} end end -function DI.pushforward(f, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras) - if x isa Number - extras.xt[0] = x - extras.xt[1] = dx - else - j = 1 - for i in eachindex(x) - extras.xt[i][0] = x[i] - extras.xt[i][j] = dx[i] - j += 1 +function DI.pushforward( + f, extras::GTPSAPushforwardExtras, backend::AutoGTPSA, x, tx::Tangents +) + ty = map(tx) do dx + if x isa Number + extras.xt[0] = x + extras.xt[1] = dx + else + j = 1 + for i in eachindex(x) + extras.xt[i][0] = x[i] + extras.xt[i][j] = dx[i] + j += 1 + end end - end - yt = f(extras.xt) - if yt isa Number - if dx isa Number - return yt[1] + yt = f(extras.xt) + if yt isa Number + if dx isa Number + return yt[1] + else + dy = 0 + for j in 1:length(dx) + dy += yt[j] + end + + return dy + end else - dy = 0 - for j in 1:length(dx) - dy += yt[j] + dy = similar(yt, eltype(eltype(yt))) + dy .= 0 + for i in eachindex(yt) + for j in 1:length(dx) + dy[i] += yt[i][j] + end end return dy end - else - dy = similar(yt, eltype(eltype(yt))) - dy .= 0 - for i in eachindex(yt) - for j in 1:length(dx) - dy[i] += yt[i][j] - end - end - - return dy end + return ty end -function DI.pushforward!(f, dy, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras) - if x isa Number - extras.xt[0] = x - extras.xt[1] = dx - else - j = 1 - for i in eachindex(x) - extras.xt[i][0] = x[i] - extras.xt[i][j] = dx[i] - j += 1 +function DI.pushforward!( + f, ty::Tangents, extras::GTPSAPushforwardExtras, backend::AutoGTPSA, x, tx::Tangents +) + for b in eachindex(tx.d, ty.d) + dx, dy = tx.d[b], ty.d[b] + if x isa Number + extras.xt[0] = x + extras.xt[1] = dx + else + j = 1 + for i in eachindex(x) + extras.xt[i][0] = x[i] + extras.xt[i][j] = dx[i] + j += 1 + end end - end - yt = f(extras.xt) - dy .= 0 - for i in eachindex(yt) - for j in 1:length(dx) - dy[i] += yt[i][j] + yt = f(extras.xt) + dy .= 0 + for i in eachindex(yt) + for j in 1:length(dx) + dy[i] += yt[i][j] + end end end - - return dy + return ty end function DI.value_and_pushforward( - f, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras + f, extras::GTPSAPushforwardExtras, backend::AutoGTPSA, x, dx ) - if x isa Number - extras.xt[0] = x - extras.xt[1] = dx - else - j = 1 - for i in eachindex(x) - extras.xt[i][0] = x[i] - extras.xt[i][j] = dx[i] - j += 1 + ys_and_dys = map(tx.d) do dx + if x isa Number + extras.xt[0] = x + extras.xt[1] = dx + else + j = 1 + for i in eachindex(x) + extras.xt[i][0] = x[i] + extras.xt[i][j] = dx[i] + j += 1 + end end - end - yt = f(extras.xt) - if yt isa Number - if dx isa Number - return yt[0], yt[1] + yt = f(extras.xt) + if yt isa Number + if dx isa Number + return yt[0], yt[1] + else + dy = 0 + for j in 1:length(dx) + dy += yt[j] + end + + return yt[0], dy + end else - dy = 0 - for j in 1:length(dx) - dy += yt[j] + dy = similar(yt, eltype(eltype(yt))) + dy .= 0 + for i in eachindex(yt) + for j in 1:length(dx) + dy[i] += yt[i][j] + end end + y = map(t -> t[0], yt) - return yt[0], dy - end - else - dy = similar(yt, eltype(eltype(yt))) - dy .= 0 - for i in eachindex(yt) - for j in 1:length(dx) - dy[i] += yt[i][j] - end + return y, dy end - y = map(t -> t[0], yt) - - return y, dy end + y = first(ys_and_dys[1]) + dys = last.(ys_and_dys) + ty = Tangents(dys...) + return y, ty end function DI.value_and_pushforward!( - f, dy, backend::AutoGTPSA, x, dx, extras::GTPSAPushforwardExtras + f, ty::Tangents, extras::GTPSAPushforwardExtras, backend::AutoGTPSA, x, tx::Tangents ) - if x isa Number - extras.xt[0] = x - extras.xt[1] = dx - else - j = 1 - for i in eachindex(x) - extras.xt[i][0] = x[i] - extras.xt[i][j] = dx[i] - j += 1 + for b in eachindex(tx.d, ty.d) + dx, dy = tx.d[b], ty.d[b] + if x isa Number + extras.xt[0] = x + extras.xt[1] = dx + else + j = 1 + for i in eachindex(x) + extras.xt[i][0] = x[i] + extras.xt[i][j] = dx[i] + j += 1 + end end - end - yt = f(extras.xt) - dy .= 0 - for i in eachindex(yt) - for j in 1:length(dx) - dy[i] += yt[i][j] + yt = f(extras.xt) + dy .= 0 + for i in eachindex(yt) + for j in 1:length(dx) + dy[i] += yt[i][j] + end end end - y = map(t -> t[0], yt) - - return y, dy + y = f(x) # TODO: optimize + return y, ty end ## Gradient @@ -183,7 +201,7 @@ function DI.prepare_gradient(f, backend::AutoGTPSA{D}, x) where {D} return GTPSAGradientExtras(xt) end -function DI.gradient(f, ::AutoGTPSA, x, extras::GTPSAGradientExtras) +function DI.gradient(f, extras::GTPSAGradientExtras, ::AutoGTPSA, x) foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part yt = f(extras.xt) grad = similar(x, eltype(eltype(yt))) @@ -191,14 +209,14 @@ function DI.gradient(f, ::AutoGTPSA, x, extras::GTPSAGradientExtras) return grad end -function DI.gradient!(f, grad, ::AutoGTPSA, x, extras::GTPSAGradientExtras) +function DI.gradient!(f, grad, extras::GTPSAGradientExtras, ::AutoGTPSA, x) foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part yt = f(extras.xt) GTPSA.gradient!(grad, yt; include_params=true) return grad end -function DI.value_and_gradient(f, ::AutoGTPSA, x, extras::GTPSAGradientExtras) +function DI.value_and_gradient(f, extras::GTPSAGradientExtras, ::AutoGTPSA, x) foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part yt = f(extras.xt) grad = similar(x, eltype(eltype(yt))) @@ -207,7 +225,7 @@ function DI.value_and_gradient(f, ::AutoGTPSA, x, extras::GTPSAGradientExtras) return y, grad end -function DI.value_and_gradient!(f, grad, ::AutoGTPSA, x, extras::GTPSAGradientExtras) +function DI.value_and_gradient!(f, grad, extras::GTPSAGradientExtras, ::AutoGTPSA, x) foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part yt = f(extras.xt) GTPSA.gradient!(grad, yt; include_params=true) @@ -244,7 +262,7 @@ function DI.prepare_jacobian(f, backend::AutoGTPSA{D}, x) where {D} return GTPSAJacobianExtras(xt) end -function DI.jacobian(f, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) +function DI.jacobian(f, extras::GTPSAJacobianExtras, ::AutoGTPSA, x) foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part yt = f(extras.xt) jac = similar(x, eltype(eltype(yt)), (length(yt), length(x))) @@ -252,14 +270,14 @@ function DI.jacobian(f, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) return jac end -function DI.jacobian!(f, jac, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) +function DI.jacobian!(f, jac, extras::GTPSAJacobianExtras, ::AutoGTPSA, x) foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part yt = f(extras.xt) GTPSA.jacobian!(jac, yt; include_params=true) return jac end -function DI.value_and_jacobian(f, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) +function DI.value_and_jacobian(f, extras::GTPSAJacobianExtras, ::AutoGTPSA, x) foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part yt = f(extras.xt) jac = similar(x, eltype(eltype(yt)), (length(yt), length(x))) @@ -268,7 +286,7 @@ function DI.value_and_jacobian(f, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) return y, jac end -function DI.value_and_jacobian!(f, jac, ::AutoGTPSA, x, extras::GTPSAJacobianExtras) +function DI.value_and_jacobian!(f, jac, extras::GTPSAJacobianExtras, ::AutoGTPSA, x) foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part yt = f(extras.xt) GTPSA.jacobian!(jac, yt; include_params=true) @@ -293,7 +311,7 @@ function DI.prepare_second_derivative(f, backend::AutoGTPSA{D}, x) where {D} return GTPSASecondDerivativeExtras(xt) end -function DI.second_derivative(f, ::AutoGTPSA, x, extras::GTPSASecondDerivativeExtras) +function DI.second_derivative(f, extras::GTPSASecondDerivativeExtras, ::AutoGTPSA, x) extras.xt[0] = x yt = f(extras.xt) if yt isa Number @@ -307,7 +325,7 @@ function DI.second_derivative(f, ::AutoGTPSA, x, extras::GTPSASecondDerivativeEx end end -function DI.second_derivative!(f, der2, ::AutoGTPSA, x, extras::GTPSASecondDerivativeExtras) +function DI.second_derivative!(f, der2, extras::GTPSASecondDerivativeExtras, ::AutoGTPSA, x) extras.xt[0] = x yt = f(extras.xt) for i in eachindex(yt) @@ -317,7 +335,7 @@ function DI.second_derivative!(f, der2, ::AutoGTPSA, x, extras::GTPSASecondDeriv end function DI.value_derivative_and_second_derivative( - f, ::AutoGTPSA, x, extras::GTPSASecondDerivativeExtras + f, extras::GTPSASecondDerivativeExtras, ::AutoGTPSA, x ) extras.xt[0] = x yt = f(extras.xt) @@ -336,7 +354,7 @@ function DI.value_derivative_and_second_derivative( end function DI.value_derivative_and_second_derivative!( - f, der, der2, ::AutoGTPSA, x, extras::GTPSASecondDerivativeExtras + f, der, der2, extras::GTPSASecondDerivativeExtras, ::AutoGTPSA, x ) extras.xt[0] = x yt = f(extras.xt) @@ -382,24 +400,30 @@ function DI.prepare_hessian(f, backend::AutoGTPSA{D}, x) where {D} return GTPSAHessianExtras(xt, m) end -function DI.hessian(f, ::AutoGTPSA{D}, x, extras::GTPSAHessianExtras) where {D} +function DI.hessian(f, extras::GTPSAHessianExtras, ::AutoGTPSA{D}, x) where {D} foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part yt = f(extras.xt) hess = similar(x, eltype(eltype(yt)), (length(x), length(x))) unsafe_fast = D == Nothing ? true : false - GTPSA.hessian!(hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=extras.m) + GTPSA.hessian!( + hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=extras.m + ) return hess end -function DI.hessian!(f, hess, ::AutoGTPSA{D}, x, extras::GTPSAHessianExtras) where {D} +function DI.hessian!(f, hess, extras::GTPSAHessianExtras, ::AutoGTPSA{D}, x) where {D} foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part yt = f(extras.xt) unsafe_fast = D == Nothing ? true : false - GTPSA.hessian!(hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=extras.m) + GTPSA.hessian!( + hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=extras.m + ) return hess end -function DI.value_gradient_and_hessian(f, ::AutoGTPSA{D}, x, extras::GTPSAHessianExtras) where {D} +function DI.value_gradient_and_hessian( + f, extras::GTPSAHessianExtras, ::AutoGTPSA{D}, x +) where {D} foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part yt = f(extras.xt) y = map(t -> t[0], yt) @@ -407,58 +431,70 @@ function DI.value_gradient_and_hessian(f, ::AutoGTPSA{D}, x, extras::GTPSAHessia GTPSA.gradient!(grad, yt; include_params=true) hess = similar(x, eltype(eltype(yt)), (length(x), length(x))) unsafe_fast = D == Nothing ? true : false - GTPSA.hessian!(hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=extras.m) + GTPSA.hessian!( + hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=extras.m + ) return y, grad, hess end function DI.value_gradient_and_hessian!( - f, grad, hess, ::AutoGTPSA{D}, x, extras::GTPSAHessianExtras + f, grad, hess, extras::GTPSAHessianExtras, ::AutoGTPSA{D}, x ) where {D} foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part yt = f(extras.xt) y = map(t -> t[0], yt) GTPSA.gradient!(grad, yt; include_params=true) unsafe_fast = D == Nothing ? true : false - GTPSA.hessian!(hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=extras.m) + GTPSA.hessian!( + hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=extras.m + ) return y, grad, hess end # HVP + struct GTPSAHVPExtras{E,H} <: HVPExtras hessextras::E hess::H end -function DI.prepare_hvp(f, backend::AutoGTPSA{D}, x, dx) where {D} +function DI.prepare_hvp(f, backend::AutoGTPSA{D}, x, tx::Tangents) where {D} hessextras = DI.prepare_hessian(f, backend, x) hess = similar(x, typeof(f(x)), (length(x), length(x))) return GTPSAHVPExtras(hessextras, hess) end -function DI.hvp(f, backend::AutoGTPSA{D}, x, dx, extras::GTPSAHVPExtras) where {D} +function DI.hvp(f, extras::GTPSAHVPExtras, backend::AutoGTPSA{D}, x, tx::Tangents) where {D} DI.hessian!(f, extras.hess, backend, x, extras.hessextras) - dg = similar(x, eltype(extras.hess)) - dg .= 0 - j = 1 - for dxi in dx - for i in 1:size(extras.hess,2) - dg[i] += extras.hess[i,j]*dxi + tg = map(tx) do dx + dg = similar(x, eltype(extras.hess)) + dg .= 0 + j = 1 + for dxi in dx + for i in 1:size(extras.hess, 2) + dg[i] += extras.hess[i, j] * dxi + end + j += 1 end - j += 1 + return dg end - return dg + return tg end -function DI.hvp!(f, dg, backend::AutoGTPSA{D}, x, dx, extras::GTPSAHVPExtras) where {D} +function DI.hvp!( + f, tg::Tangents, extras::GTPSAHVPExtras, backend::AutoGTPSA{D}, x, tx::Tangents +) where {D} DI.hessian!(f, extras.hess, backend, x, extras.hessextras) - dg .= 0 - j = 1 - for dxi in dx - for i in 1:size(extras.hess,2) - dg[i] += extras.hess[i,j]*dxi + for b in eachindex(tg.d) + dg = tg.d[b] + dg .= 0 + j = 1 + for dxi in dx + for i in 1:size(extras.hess, 2) + dg[i] += extras.hess[i, j] * dxi + end + j += 1 end - j += 1 end - return dg + return tg end - diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl index d216b1bee..17ea3893d 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl @@ -4,13 +4,14 @@ struct GTPSATwoArgPushforwardExtras{X,Y} <: PushforwardExtras yt::Y end -function DI.prepare_pushforward(f!, y, backend::AutoGTPSA{D}, x, dx) where {D} +function DI.prepare_pushforward(f!, y, backend::AutoGTPSA{D}, x, tx::Tangents) where {D} if D != Nothing d = backend.descriptor else d = Descriptor(length(x), 1) end + dx = first(tx) if x isa Number xt = TPS{promote_type(typeof(dx), typeof(x), Float64)}(; use=d) else @@ -31,113 +32,141 @@ function DI.prepare_pushforward(f!, y, backend::AutoGTPSA{D}, x, dx) where {D} end function DI.pushforward( - f!, y, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPushforwardExtras + f!, y, extras::GTPSATwoArgPushforwardExtras, backend::AutoGTPSA, x, tx::Tangents ) - if x isa Number - extras.xt[0] = x - extras.xt[1] = dx - else - j = 1 - for i in eachindex(x) - extras.xt[i][0] = x[i] - extras.xt[i][j] = dx[i] - j += 1 + ty = map(tx) do dx + if x isa Number + extras.xt[0] = x + extras.xt[1] = dx + else + j = 1 + for i in eachindex(x) + extras.xt[i][0] = x[i] + extras.xt[i][j] = dx[i] + j += 1 + end end - end - f!(extras.yt, extras.xt) - map!(t -> t[0], y, extras.yt) - dy = similar(extras.yt, eltype(eltype(extras.yt))) - dy .= 0 - for i in eachindex(extras.yt) - for j in 1:length(dx) - dy[i] += extras.yt[i][j] + f!(extras.yt, extras.xt) + map!(t -> t[0], y, extras.yt) + dy = similar(extras.yt, eltype(eltype(extras.yt))) + dy .= 0 + for i in eachindex(extras.yt) + for j in 1:length(dx) + dy[i] += extras.yt[i][j] + end end + return dy end - return dy + return ty end function DI.pushforward!( - f!, y, dy, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPushforwardExtras + f!, + y, + ty::Tangents, + extras::GTPSATwoArgPushforwardExtras, + backend::AutoGTPSA, + x, + tx::Tangents, ) - if x isa Number - extras.xt[0] = x - extras.xt[1] = dx - else - j = 1 - for i in eachindex(x) - extras.xt[i][0] = x[i] - extras.xt[i][j] = dx[i] - j += 1 + for b in eachindex(tx.d, ty.d) + dx, dy = tx.d[b], ty.d[b] + if x isa Number + extras.xt[0] = x + extras.xt[1] = dx + else + j = 1 + for i in eachindex(x) + extras.xt[i][0] = x[i] + extras.xt[i][j] = dx[i] + j += 1 + end end - end - f!(extras.yt, extras.xt) - map!(t -> t[0], y, extras.yt) - dy .= 0 - for i in eachindex(extras.yt) - for j in 1:length(dx) - dy[i] += extras.yt[i][j] + f!(extras.yt, extras.xt) + map!(t -> t[0], y, extras.yt) + dy .= 0 + for i in eachindex(extras.yt) + for j in 1:length(dx) + dy[i] += extras.yt[i][j] + end end end - return dy + return ty end function DI.value_and_pushforward( - f!, y, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPushforwardExtras + f!, y, extras::GTPSATwoArgPushforwardExtras, backend::AutoGTPSA, x, tx::Tangents ) - if x isa Number - extras.xt[0] = x - extras.xt[1] = dx - else - j = 1 - for i in eachindex(x) - extras.xt[i][0] = x[i] - extras.xt[i][j] = dx[i] - j += 1 + ys_and_dys = map(tx.d) do dx + if x isa Number + extras.xt[0] = x + extras.xt[1] = dx + else + j = 1 + for i in eachindex(x) + extras.xt[i][0] = x[i] + extras.xt[i][j] = dx[i] + j += 1 + end end - end - f!(extras.yt, extras.xt) - dy = similar(extras.yt, eltype(eltype(extras.yt))) - dy .= 0 - for i in eachindex(extras.yt) - for j in 1:length(dx) - dy[i] += extras.yt[i][j] + f!(extras.yt, extras.xt) + dy = similar(extras.yt, eltype(eltype(extras.yt))) + dy .= 0 + for i in eachindex(extras.yt) + for j in 1:length(dx) + dy[i] += extras.yt[i][j] + end end + map!(t -> t[0], y, extras.yt) + + return y, dy end - map!(t -> t[0], y, extras.yt) - return y, dy + y = first(ys_and_dys[1]) + dys = last.(ys_and_dys) + ty = Tangents(dys...) + return y, ty end function DI.value_and_pushforward!( - f!, y, dy, backend::AutoGTPSA, x, dx, extras::GTPSATwoArgPushforwardExtras + f!, + y, + ty::Tangents, + extras::GTPSATwoArgPushforwardExtras, + backend::AutoGTPSA, + x, + tx::Tangents, ) - if x isa Number - extras.xt[0] = x - extras.xt[1] = dx - else - j = 1 - for i in eachindex(x) - extras.xt[i][0] = x[i] - extras.xt[i][j] = dx[i] - j += 1 + for b in eachindex(tx.d, ty.d) + dx, dy = tx.d, ty.d + if x isa Number + extras.xt[0] = x + extras.xt[1] = dx + else + j = 1 + for i in eachindex(x) + extras.xt[i][0] = x[i] + extras.xt[i][j] = dx[i] + j += 1 + end end - end - f!(extras.yt, extras.xt) - map!(t -> t[0], y, extras.yt) - dy .= 0 - for i in eachindex(extras.yt) - for j in 1:length(dx) - dy[i] += extras.yt[i][j] + f!(extras.yt, extras.xt) + map!(t -> t[0], y, extras.yt) + dy .= 0 + for i in eachindex(extras.yt) + for j in 1:length(dx) + dy[i] += extras.yt[i][j] + end end end - return y, dy + return y, ty end ## Jacobian @@ -176,7 +205,7 @@ function DI.prepare_jacobian(f!, y, backend::AutoGTPSA{D}, x) where {D} return GTPSATwoArgJacobianExtras(xt, yt) end -function DI.jacobian(f!, y, ::AutoGTPSA, x, extras::GTPSATwoArgJacobianExtras) +function DI.jacobian(f!, y, extras::GTPSATwoArgJacobianExtras, ::AutoGTPSA, x) foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part f!(extras.yt, extras.xt) jac = similar(x, eltype(eltype(extras.yt)), (length(extras.yt), length(x))) @@ -185,7 +214,7 @@ function DI.jacobian(f!, y, ::AutoGTPSA, x, extras::GTPSATwoArgJacobianExtras) return jac end -function DI.jacobian!(f!, y, jac, ::AutoGTPSA, x, extras::GTPSATwoArgJacobianExtras) +function DI.jacobian!(f!, y, jac, extras::GTPSATwoArgJacobianExtras, ::AutoGTPSA, x) foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part f!(extras.yt, extras.xt) GTPSA.jacobian!(jac, extras.yt; include_params=true) @@ -193,7 +222,7 @@ function DI.jacobian!(f!, y, jac, ::AutoGTPSA, x, extras::GTPSATwoArgJacobianExt return jac end -function DI.value_and_jacobian(f!, y, ::AutoGTPSA, x, extras::GTPSATwoArgJacobianExtras) +function DI.value_and_jacobian(f!, y, extras::GTPSATwoArgJacobianExtras, ::AutoGTPSA, x) foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part f!(extras.yt, extras.xt) jac = similar(x, eltype(eltype(extras.yt)), (length(extras.yt), length(x))) @@ -203,7 +232,7 @@ function DI.value_and_jacobian(f!, y, ::AutoGTPSA, x, extras::GTPSATwoArgJacobia end function DI.value_and_jacobian!( - f!, y, jac, ::AutoGTPSA, x, extras::GTPSATwoArgJacobianExtras + f!, y, jac, extras::GTPSATwoArgJacobianExtras, ::AutoGTPSA, x ) foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part f!(extras.yt, extras.xt) diff --git a/DifferentiationInterface/test/Back/GTPSA/test.jl b/DifferentiationInterface/test/Back/GTPSA/test.jl index da98887ae..5895c08ca 100644 --- a/DifferentiationInterface/test/Back/GTPSA/test.jl +++ b/DifferentiationInterface/test/Back/GTPSA/test.jl @@ -9,8 +9,7 @@ LOGGING = get(ENV, "CI", "false") == "false" for backend in [AutoGTPSA()] @test check_available(backend) - @test check_twoarg(backend) - @test check_hessian(backend) + @test check_inplace(backend) end test_differentiation(AutoGTPSA(); type_stability=true, logging=LOGGING); From f663484cd282b9b9be4afb0fde4241b240ca4bb3 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Tue, 17 Sep 2024 13:26:26 +0200 Subject: [PATCH 27/44] Tangents import --- .../DifferentiationInterfaceGTPSAExt.jl | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl index 74432c3e7..982cada7a 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl @@ -9,7 +9,8 @@ using DifferentiationInterface: JacobianExtras, HessianExtras, PushforwardExtras, - HVPExtras + HVPExtras, + Tangents using GTPSA DI.check_available(::AutoGTPSA) = true From a388cfca3b5b07e719eb1bca32f52551d65aff8a Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Tue, 17 Sep 2024 13:31:26 +0200 Subject: [PATCH 28/44] Fix da fixes --- .../ext/DifferentiationInterfaceGTPSAExt/onearg.jl | 8 ++++---- .../ext/DifferentiationInterfaceGTPSAExt/twoarg.jl | 2 +- DifferentiationInterface/test/Back/GTPSA/test.jl | 2 +- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl index f19ce2e7d..2359f65f3 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl @@ -97,7 +97,7 @@ function DI.pushforward!( end function DI.value_and_pushforward( - f, extras::GTPSAPushforwardExtras, backend::AutoGTPSA, x, dx + f, extras::GTPSAPushforwardExtras, backend::AutoGTPSA, x, tx::Tangents ) ys_and_dys = map(tx.d) do dx if x isa Number @@ -465,7 +465,7 @@ function DI.prepare_hvp(f, backend::AutoGTPSA{D}, x, tx::Tangents) where {D} end function DI.hvp(f, extras::GTPSAHVPExtras, backend::AutoGTPSA{D}, x, tx::Tangents) where {D} - DI.hessian!(f, extras.hess, backend, x, extras.hessextras) + DI.hessian!(f, extras.hess, extras.hessextras, backend, x) tg = map(tx) do dx dg = similar(x, eltype(extras.hess)) dg .= 0 @@ -484,9 +484,9 @@ end function DI.hvp!( f, tg::Tangents, extras::GTPSAHVPExtras, backend::AutoGTPSA{D}, x, tx::Tangents ) where {D} - DI.hessian!(f, extras.hess, backend, x, extras.hessextras) + DI.hessian!(f, extras.hess, extras.hessextras, backend, x) for b in eachindex(tg.d) - dg = tg.d[b] + dx, dg = tx.d[b], tg.d[b] dg .= 0 j = 1 for dxi in dx diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl index 17ea3893d..9763135da 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl @@ -143,7 +143,7 @@ function DI.value_and_pushforward!( tx::Tangents, ) for b in eachindex(tx.d, ty.d) - dx, dy = tx.d, ty.d + dx, dy = tx.d[b], ty.d[b] if x isa Number extras.xt[0] = x extras.xt[1] = dx diff --git a/DifferentiationInterface/test/Back/GTPSA/test.jl b/DifferentiationInterface/test/Back/GTPSA/test.jl index 5895c08ca..c2d6cc715 100644 --- a/DifferentiationInterface/test/Back/GTPSA/test.jl +++ b/DifferentiationInterface/test/Back/GTPSA/test.jl @@ -12,4 +12,4 @@ for backend in [AutoGTPSA()] @test check_inplace(backend) end -test_differentiation(AutoGTPSA(); type_stability=true, logging=LOGGING); +test_differentiation(AutoGTPSA(); type_stability=false, logging=LOGGING); From bd70675254340960604f01747d044fcf722632ee Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Wed, 18 Sep 2024 10:27:28 +0200 Subject: [PATCH 29/44] Remove temporary backend --- DifferentiationInterface/Project.toml | 2 +- .../docs/src/explanation/backends.md | 2 ++ .../DifferentiationInterfaceGTPSAExt.jl | 2 +- .../src/DifferentiationInterface.jl | 3 --- .../src/misc/temp_backends.jl | 27 ------------------- 5 files changed, 4 insertions(+), 32 deletions(-) delete mode 100644 DifferentiationInterface/src/misc/temp_backends.jl diff --git a/DifferentiationInterface/Project.toml b/DifferentiationInterface/Project.toml index f0fd2641d..ce148f702 100644 --- a/DifferentiationInterface/Project.toml +++ b/DifferentiationInterface/Project.toml @@ -46,7 +46,7 @@ DifferentiationInterfaceTrackerExt = "Tracker" DifferentiationInterfaceZygoteExt = ["Zygote", "ForwardDiff"] [compat] -ADTypes = "1.7.0" +ADTypes = "1.8.0" ChainRulesCore = "1.23.0" Compat = "3.46,4.2" Diffractor = "=0.2.6" diff --git a/DifferentiationInterface/docs/src/explanation/backends.md b/DifferentiationInterface/docs/src/explanation/backends.md index 69ffe3cca..b5ad618e7 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) - [`AutoPolyesterForwardDiff`](@extref ADTypes.AutoPolyesterForwardDiff) - [`AutoReverseDiff`](@extref ADTypes.AutoReverseDiff) - [`AutoSymbolics`](@extref ADTypes.AutoSymbolics) @@ -55,6 +56,7 @@ In practice, many AD backends have custom implementations for high-level operato | `AutoFiniteDiff` | 🔀 | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | | `AutoFiniteDifferences` | 🔀 | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | | `AutoForwardDiff` | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | + | `AutoGTPSA` | ✅ | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | | `AutoPolyesterForwardDiff` | 🔀 | ❌ | 🔀 | ✅ | ✅ | 🔀 | 🔀 | 🔀 | | `AutoReverseDiff` | ❌ | 🔀 | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | | `AutoSymbolics` | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl index 982cada7a..7323d711c 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl @@ -1,7 +1,7 @@ module DifferentiationInterfaceGTPSAExt import DifferentiationInterface as DI -using DifferentiationInterface: AutoGTPSA # TODO: replace with ADTypes +using ADTypes: AutoGTPSA using DifferentiationInterface: DerivativeExtras, SecondDerivativeExtras, diff --git a/DifferentiationInterface/src/DifferentiationInterface.jl b/DifferentiationInterface/src/DifferentiationInterface.jl index 4668fdd04..3d95f8be9 100644 --- a/DifferentiationInterface/src/DifferentiationInterface.jl +++ b/DifferentiationInterface/src/DifferentiationInterface.jl @@ -62,8 +62,6 @@ include("misc/from_primitive.jl") include("misc/sparsity_detector.jl") include("misc/zero_backends.jl") -include("misc/temp_backends.jl") - function __init__() @require_extensions end @@ -121,7 +119,6 @@ export AutoSymbolics export AutoTapir export AutoTracker export AutoZygote -export AutoGTPSA # temporary export AutoSparse diff --git a/DifferentiationInterface/src/misc/temp_backends.jl b/DifferentiationInterface/src/misc/temp_backends.jl deleted file mode 100644 index 2763db91e..000000000 --- a/DifferentiationInterface/src/misc/temp_backends.jl +++ /dev/null @@ -1,27 +0,0 @@ -""" - AutoGTPSA{D} - -Struct used to select the [GTPSA.jl](https://github.com/bmad-sim/GTPSA.jl) backend for automatic differentiation. - -!!! warning - This type is not part of the public API, it is only used for testing before being moved to ADTypes.jl. - -# Constructors - - AutoGTPSA(; descriptor=nothing) - -# Fields - - - `descriptor::D`: can be either - - + a GTPSA `Descriptor` specifying the number of variables/parameters, parameter - order, individual variable/parameter truncation orders, and maximum order. See - the [GTPSA.jl documentation](https://bmad-sim.github.io/GTPSA.jl/stable/man/c_descriptor/) for more details. - + `nothing` to automatically use a `Descriptor` given the context. - -""" -Base.@kwdef struct AutoGTPSA{D} <: AbstractADType - descriptor::D = nothing -end - -ADTypes.mode(::AutoGTPSA) = ForwardMode() From 77e49d635db8a9d201f93033a10e4e159c39f641 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Wed, 18 Sep 2024 10:39:14 +0200 Subject: [PATCH 30/44] Fixes --- .../DifferentiationInterfaceGTPSAExt.jl | 1 + .../onearg.jl | 48 ++----------------- .../twoarg.jl | 48 ++----------------- .../DifferentiationInterfaceGTPSAExt/utils.jl | 12 +++++ .../src/DifferentiationInterface.jl | 1 + 5 files changed, 22 insertions(+), 88 deletions(-) create mode 100644 DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/utils.jl diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl index 7323d711c..437cda572 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl @@ -15,6 +15,7 @@ using GTPSA DI.check_available(::AutoGTPSA) = true +include("utils.jl") include("onearg.jl") include("twoarg.jl") diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl index 2359f65f3..4638f61ad 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl @@ -29,17 +29,7 @@ function DI.pushforward( f, extras::GTPSAPushforwardExtras, backend::AutoGTPSA, x, tx::Tangents ) ty = map(tx) do dx - if x isa Number - extras.xt[0] = x - extras.xt[1] = dx - else - j = 1 - for i in eachindex(x) - extras.xt[i][0] = x[i] - extras.xt[i][j] = dx[i] - j += 1 - end - end + initialize!(extras.xt, x, dx) yt = f(extras.xt) if yt isa Number @@ -73,17 +63,7 @@ function DI.pushforward!( ) for b in eachindex(tx.d, ty.d) dx, dy = tx.d[b], ty.d[b] - if x isa Number - extras.xt[0] = x - extras.xt[1] = dx - else - j = 1 - for i in eachindex(x) - extras.xt[i][0] = x[i] - extras.xt[i][j] = dx[i] - j += 1 - end - end + initialize!(extras.xt, x, dx) yt = f(extras.xt) dy .= 0 @@ -100,17 +80,7 @@ function DI.value_and_pushforward( f, extras::GTPSAPushforwardExtras, backend::AutoGTPSA, x, tx::Tangents ) ys_and_dys = map(tx.d) do dx - if x isa Number - extras.xt[0] = x - extras.xt[1] = dx - else - j = 1 - for i in eachindex(x) - extras.xt[i][0] = x[i] - extras.xt[i][j] = dx[i] - j += 1 - end - end + initialize!(extras.xt, x, dx) yt = f(extras.xt) if yt isa Number @@ -148,17 +118,7 @@ function DI.value_and_pushforward!( ) for b in eachindex(tx.d, ty.d) dx, dy = tx.d[b], ty.d[b] - if x isa Number - extras.xt[0] = x - extras.xt[1] = dx - else - j = 1 - for i in eachindex(x) - extras.xt[i][0] = x[i] - extras.xt[i][j] = dx[i] - j += 1 - end - end + initialize!(extras.xt, x, dx) yt = f(extras.xt) dy .= 0 diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl index 9763135da..849d0d1f3 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl @@ -35,17 +35,7 @@ function DI.pushforward( f!, y, extras::GTPSATwoArgPushforwardExtras, backend::AutoGTPSA, x, tx::Tangents ) ty = map(tx) do dx - if x isa Number - extras.xt[0] = x - extras.xt[1] = dx - else - j = 1 - for i in eachindex(x) - extras.xt[i][0] = x[i] - extras.xt[i][j] = dx[i] - j += 1 - end - end + initialize!(extras.xt, x, dx) f!(extras.yt, extras.xt) map!(t -> t[0], y, extras.yt) @@ -73,17 +63,7 @@ function DI.pushforward!( ) for b in eachindex(tx.d, ty.d) dx, dy = tx.d[b], ty.d[b] - if x isa Number - extras.xt[0] = x - extras.xt[1] = dx - else - j = 1 - for i in eachindex(x) - extras.xt[i][0] = x[i] - extras.xt[i][j] = dx[i] - j += 1 - end - end + initialize!(extras.xt, x, dx) f!(extras.yt, extras.xt) map!(t -> t[0], y, extras.yt) @@ -102,17 +82,7 @@ function DI.value_and_pushforward( f!, y, extras::GTPSATwoArgPushforwardExtras, backend::AutoGTPSA, x, tx::Tangents ) ys_and_dys = map(tx.d) do dx - if x isa Number - extras.xt[0] = x - extras.xt[1] = dx - else - j = 1 - for i in eachindex(x) - extras.xt[i][0] = x[i] - extras.xt[i][j] = dx[i] - j += 1 - end - end + initialize!(extras.xt, x, dx) f!(extras.yt, extras.xt) dy = similar(extras.yt, eltype(eltype(extras.yt))) @@ -144,17 +114,7 @@ function DI.value_and_pushforward!( ) for b in eachindex(tx.d, ty.d) dx, dy = tx.d[b], ty.d[b] - if x isa Number - extras.xt[0] = x - extras.xt[1] = dx - else - j = 1 - for i in eachindex(x) - extras.xt[i][0] = x[i] - extras.xt[i][j] = dx[i] - j += 1 - end - end + initialize!(extras.xt, x, dx) f!(extras.yt, extras.xt) map!(t -> t[0], y, extras.yt) diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/utils.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/utils.jl new file mode 100644 index 000000000..0e223f932 --- /dev/null +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/utils.jl @@ -0,0 +1,12 @@ +function initialize!(xt::TPS, x::Number, dx::Number) + xt[0] = x + xt[1] = dx + return xt +end + +function initialize!(xt::AbstractArray{<:TPS}, x::AbstractArray, dx::AbstractArray) + for i in eachindex(xt, x, dx) + initialize!(xt[i], x[i], dx[i]) + end + return xt +end diff --git a/DifferentiationInterface/src/DifferentiationInterface.jl b/DifferentiationInterface/src/DifferentiationInterface.jl index 3d95f8be9..b57f27343 100644 --- a/DifferentiationInterface/src/DifferentiationInterface.jl +++ b/DifferentiationInterface/src/DifferentiationInterface.jl @@ -23,6 +23,7 @@ using ADTypes: AutoFiniteDiff, AutoFiniteDifferences, AutoForwardDiff, + AutoGTPSA, AutoPolyesterForwardDiff, AutoReverseDiff, AutoSymbolics, From dacf84d6df92739c010d61971927c9d90b148232 Mon Sep 17 00:00:00 2001 From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> Date: Wed, 18 Sep 2024 11:12:56 +0200 Subject: [PATCH 31/44] Docs --- DifferentiationInterface/docs/src/api.md | 4 ---- 1 file changed, 4 deletions(-) diff --git a/DifferentiationInterface/docs/src/api.md b/DifferentiationInterface/docs/src/api.md index b9c235f27..0f46ea32f 100644 --- a/DifferentiationInterface/docs/src/api.md +++ b/DifferentiationInterface/docs/src/api.md @@ -136,8 +136,4 @@ The following is not part of the public API. Modules = [DifferentiationInterface] Public = false Filter = t -> !(Symbol(t) in [:outer, :inner]) -``` - -```@docs -AutoGTPSA ``` \ No newline at end of file From f67c9aa85568cf0b950a52f0345c5b3ddd3ff3b3 Mon Sep 17 00:00:00 2001 From: Matt Signorelli Date: Mon, 6 Jan 2025 15:31:33 -0500 Subject: [PATCH 32/44] GTPSAExt devel in progress --- .../DifferentiationInterfaceGTPSAExt.jl | 11 +- .../onearg.jl | 379 +++++++++--------- .../DifferentiationInterfaceGTPSAExt/utils.jl | 37 +- 3 files changed, 222 insertions(+), 205 deletions(-) diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl index 437cda572..a5017c6cb 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl @@ -2,21 +2,12 @@ module DifferentiationInterfaceGTPSAExt import DifferentiationInterface as DI using ADTypes: AutoGTPSA -using DifferentiationInterface: - DerivativeExtras, - SecondDerivativeExtras, - GradientExtras, - JacobianExtras, - HessianExtras, - PushforwardExtras, - HVPExtras, - Tangents using GTPSA DI.check_available(::AutoGTPSA) = true include("utils.jl") include("onearg.jl") -include("twoarg.jl") +#include("twoarg.jl") end diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl index 4638f61ad..03441af8e 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl @@ -1,57 +1,54 @@ ## Pushforward -struct GTPSAPushforwardExtras{X} <: PushforwardExtras + +# 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::Tangents) where {D} +function DI.prepare_pushforward( + ::F, + backend::AutoGTPSA{D}, + x, + tx::NTuple, + ::Vararg{DI.Context,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(length(x), 1) + d = Descriptor(1, 1) # 1 variable to first order end - - dx = first(tx) if x isa Number - xt = TPS{promote_type(typeof(dx), typeof(x), Float64)}(; use=d) - return GTPSAPushforwardExtras(xt) + xt = TPS{promote_type(typeof(first(tx)), typeof(x), Float64)}(; use=d) + return GTPSAOneArgPushforwardPrep(xt) else - xt = similar(x, TPS{promote_type(eltype(dx), eltype(x), Float64)}) - - # xt and x have same indexing because of similar + xt = similar(x, TPS{promote_type(eltype(first(tx)), eltype(x), Float64)}) for i in eachindex(xt) - xt[i] = TPS{promote_type(eltype(dx), eltype(x), Float64)}(; use=d) + xt[i] = TPS{promote_type(eltype(first(tx)), eltype(x), Float64)}(; use=d) end - return GTPSAPushforwardExtras(xt) + return GTPSAOneArgPushforwardPrep(xt) end end function DI.pushforward( - f, extras::GTPSAPushforwardExtras, backend::AutoGTPSA, x, tx::Tangents -) + f, + prep::GTPSAOneArgPushforwardPrep, + ::AutoGTPSA, + x, + tx::NTuple, + contexts::Vararg{DI.Context,C}, +) where {C} + fc = DI.with_contexts(f, contexts...) ty = map(tx) do dx - initialize!(extras.xt, x, dx) - - yt = f(extras.xt) + initialize!(prep.xt, x, dx, 1) + yt = fc(prep.xt) if yt isa Number - if dx isa Number - return yt[1] - else - dy = 0 - for j in 1:length(dx) - dy += yt[j] - end - - return dy - end + return yt[1] else - dy = similar(yt, eltype(eltype(yt))) - dy .= 0 - for i in eachindex(yt) - for j in 1:length(dx) - dy[i] += yt[i][j] - end - end - + dy = map(t -> t[1], yt) return dy end end @@ -59,135 +56,140 @@ function DI.pushforward( end function DI.pushforward!( - f, ty::Tangents, extras::GTPSAPushforwardExtras, backend::AutoGTPSA, x, tx::Tangents -) - for b in eachindex(tx.d, ty.d) - dx, dy = tx.d[b], ty.d[b] - initialize!(extras.xt, x, dx) - - yt = f(extras.xt) - dy .= 0 - for i in eachindex(yt) - for j in 1:length(dx) - dy[i] += yt[i][j] - end - end + f, + ty::NTuple, + prep::GTPSAOneArgPushforwardPrep, + ::AutoGTPSA, + x, + tx::NTuple, + contexts::Vararg{DI.Context,C}, +) where {C} + fc = DI.with_contexts(f, contexts...) + for b in eachindex(tx, ty) + dx, dy = tx[b], ty[b] + initialize!(prep.xt, x, dx, 1) + yt = fc(prep.xt) + map!(t -> t[1], dy, yt) end return ty end function DI.value_and_pushforward( - f, extras::GTPSAPushforwardExtras, backend::AutoGTPSA, x, tx::Tangents -) - ys_and_dys = map(tx.d) do dx - initialize!(extras.xt, x, dx) - - yt = f(extras.xt) - if yt isa Number - if dx isa Number - return yt[0], yt[1] - else - dy = 0 - for j in 1:length(dx) - dy += yt[j] - end - - return yt[0], dy - end - else - dy = similar(yt, eltype(eltype(yt))) - dy .= 0 - for i in eachindex(yt) - for j in 1:length(dx) - dy[i] += yt[i][j] - end - end - y = map(t -> t[0], yt) - - return y, dy - end - end - y = first(ys_and_dys[1]) - dys = last.(ys_and_dys) - ty = Tangents(dys...) + f, + prep::GTPSAOneArgPushforwardPrep, + backend::AutoGTPSA, + x, + tx::NTuple, + contexts::Vararg{DI.Context,C}, +) where {C} + + fc = DI.with_contexts(f, contexts...) + ty = DI.pushforward(fc, prep, backend, x, tx) + y = fc(x) # This actually seems faster than indexing the TPSs return y, ty end function DI.value_and_pushforward!( - f, ty::Tangents, extras::GTPSAPushforwardExtras, backend::AutoGTPSA, x, tx::Tangents -) - for b in eachindex(tx.d, ty.d) - dx, dy = tx.d[b], ty.d[b] - initialize!(extras.xt, x, dx) - - yt = f(extras.xt) - dy .= 0 - for i in eachindex(yt) - for j in 1:length(dx) - dy[i] += yt[i][j] - end - end - end - y = f(x) # TODO: optimize + f, + ty::NTuple, + prep::GTPSAOneArgPushforwardPrep, + ::AutoGTPSA, + x, + tx::NTuple, + contexts::Vararg{DI.Context,C}, +) where {C} + fc = DI.with_contexts(f, contexts...) + DI.pushforward!(fc, ty, prep, backend, x, tx) + y = fc(x) return y, ty end ## Gradient - -struct GTPSAGradientExtras{X} <: GradientExtras +# Contains either a single pre-allocated initial TPS +# or a vector of pre-allocated TPSs. +struct GTPSAOneArgGradientPrep{X} <: DI.GradientPrep xt::X end -function DI.prepare_gradient(f, backend::AutoGTPSA{D}, x) where {D} +# Unlike JVP, this requires us to use all variables +function DI.prepare_gradient( + f, + backend::AutoGTPSA{D}, + x, + contexts::Vararg{DI.Context,C}, +) where {D,C} + if D != Nothing d = backend.descriptor - nn = GTPSA.numnn(d) + length(x) == GTPSA.numnn(d) || error("Number of variables + parameters in Descriptor disagrees with function number of inputs!") else - nn = length(x) - d = Descriptor(nn, 1) + d = Descriptor(length(x), 1) # n variables to first order 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 + # We set the slopes of each variable to 1 here, this will always be the case for gradient + if x isa Number + xt = TPS{promote_type(typeof(x), Float64)}(; use=d) + xt[1] = 1 + return GTPSAOneArgGradientPrep(xt) + else + 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 + return GTPSAOneArgGradientPrep(xt) +end + +function DI.gradient( + f, + prep::GTPSAOneArgGradientPrep, + ::AutoGTPSA, + x, + contexts::Vararg{DI.Context,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) + if prep.xt isa Number + return yt[1] + else + grad = similar(x, GTPSA.numtype(eltype(yt))) + GTPSA.gradient!(grad, yt; include_params=true) end - - return GTPSAGradientExtras(xt) -end - -function DI.gradient(f, extras::GTPSAGradientExtras, ::AutoGTPSA, x) - foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part - yt = f(extras.xt) - grad = similar(x, eltype(eltype(yt))) - GTPSA.gradient!(grad, yt; include_params=true) return grad end -function DI.gradient!(f, grad, extras::GTPSAGradientExtras, ::AutoGTPSA, x) - foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part - yt = f(extras.xt) +function DI.gradient!( + f, + grad, + prep::GTPSAOneArgGradientPrep, + ::AutoGTPSA, + x, + contexts::Vararg{DI.Context,C}, +) where {C} + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part + yt = f(prep.xt) GTPSA.gradient!(grad, yt; include_params=true) return grad end -function DI.value_and_gradient(f, extras::GTPSAGradientExtras, ::AutoGTPSA, x) - foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part - yt = f(extras.xt) +#= +function DI.value_and_gradient(f, prep::GTPSAGradientPrep, ::AutoGTPSA, x) + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part + yt = f(prep.xt) grad = similar(x, eltype(eltype(yt))) GTPSA.gradient!(grad, yt; include_params=true) y = map(t -> t[0], yt) return y, grad end -function DI.value_and_gradient!(f, grad, extras::GTPSAGradientExtras, ::AutoGTPSA, x) - foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part - yt = f(extras.xt) +function DI.value_and_gradient!(f, grad, prep::GTPSAGradientPrep, ::AutoGTPSA, x) + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part + yt = f(prep.xt) GTPSA.gradient!(grad, yt; include_params=true) y = map(t -> t[0], yt) return y, grad @@ -195,7 +197,7 @@ end ## Jacobian -struct GTPSAJacobianExtras{X} <: JacobianExtras +struct GTPSAJacobianprep{X} <: Jacobianprep xt::X end @@ -219,36 +221,36 @@ function DI.prepare_jacobian(f, backend::AutoGTPSA{D}, x) where {D} j += 1 end - return GTPSAJacobianExtras(xt) + return GTPSAJacobianprep(xt) end -function DI.jacobian(f, extras::GTPSAJacobianExtras, ::AutoGTPSA, x) - foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part - yt = f(extras.xt) +function DI.jacobian(f, prep::GTPSAJacobianprep, ::AutoGTPSA, x) + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part + yt = f(prep.xt) jac = similar(x, eltype(eltype(yt)), (length(yt), length(x))) GTPSA.jacobian!(jac, yt; include_params=true) return jac end -function DI.jacobian!(f, jac, extras::GTPSAJacobianExtras, ::AutoGTPSA, x) - foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part - yt = f(extras.xt) +function DI.jacobian!(f, jac, prep::GTPSAJacobianprep, ::AutoGTPSA, x) + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part + yt = f(prep.xt) GTPSA.jacobian!(jac, yt; include_params=true) return jac end -function DI.value_and_jacobian(f, extras::GTPSAJacobianExtras, ::AutoGTPSA, x) - foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part - yt = f(extras.xt) +function DI.value_and_jacobian(f, prep::GTPSAJacobianprep, ::AutoGTPSA, x) + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part + yt = f(prep.xt) jac = similar(x, eltype(eltype(yt)), (length(yt), length(x))) GTPSA.jacobian!(jac, yt; include_params=true) y = map(t -> t[0], yt) return y, jac end -function DI.value_and_jacobian!(f, jac, extras::GTPSAJacobianExtras, ::AutoGTPSA, x) - foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part - yt = f(extras.xt) +function DI.value_and_jacobian!(f, jac, prep::GTPSAJacobianprep, ::AutoGTPSA, x) + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part + yt = f(prep.xt) GTPSA.jacobian!(jac, yt; include_params=true) y = map(t -> t[0], yt) return y, jac @@ -256,7 +258,7 @@ end ## Second derivative -struct GTPSASecondDerivativeExtras{X} <: SecondDerivativeExtras +struct GTPSASecondDerivativeprep{X} <: SecondDerivativeprep xt::X end @@ -268,12 +270,12 @@ function DI.prepare_second_derivative(f, backend::AutoGTPSA{D}, x) where {D} end xt = TPS{promote_type(typeof(x), Float64)}(; use=d) xt[1] = 1 - return GTPSASecondDerivativeExtras(xt) + return GTPSASecondDerivativeprep(xt) end -function DI.second_derivative(f, extras::GTPSASecondDerivativeExtras, ::AutoGTPSA, x) - extras.xt[0] = x - yt = f(extras.xt) +function DI.second_derivative(f, prep::GTPSASecondDerivativeprep, ::AutoGTPSA, x) + prep.xt[0] = x + yt = f(prep.xt) if yt isa Number return yt[2] * 2 else @@ -285,9 +287,9 @@ function DI.second_derivative(f, extras::GTPSASecondDerivativeExtras, ::AutoGTPS end end -function DI.second_derivative!(f, der2, extras::GTPSASecondDerivativeExtras, ::AutoGTPSA, x) - extras.xt[0] = x - yt = f(extras.xt) +function DI.second_derivative!(f, der2, prep::GTPSASecondDerivativeprep, ::AutoGTPSA, x) + prep.xt[0] = x + yt = f(prep.xt) for i in eachindex(yt) der2[i] = yt[i][2] * 2 end @@ -295,10 +297,10 @@ function DI.second_derivative!(f, der2, extras::GTPSASecondDerivativeExtras, ::A end function DI.value_derivative_and_second_derivative( - f, extras::GTPSASecondDerivativeExtras, ::AutoGTPSA, x + f, prep::GTPSASecondDerivativeprep, ::AutoGTPSA, x ) - extras.xt[0] = x - yt = f(extras.xt) + prep.xt[0] = x + yt = f(prep.xt) if yt isa Number return yt[0], yt[1], yt[2] * 2 else @@ -314,10 +316,10 @@ function DI.value_derivative_and_second_derivative( end function DI.value_derivative_and_second_derivative!( - f, der, der2, extras::GTPSASecondDerivativeExtras, ::AutoGTPSA, x + f, der, der2, prep::GTPSASecondDerivativeprep, ::AutoGTPSA, x ) - extras.xt[0] = x - yt = f(extras.xt) + prep.xt[0] = x + yt = f(prep.xt) y = map(t -> t[0], yt) for i in eachindex(yt) der[i] = yt[i][1] @@ -328,7 +330,7 @@ end ## Hessian -struct GTPSAHessianExtras{X,M} <: HessianExtras +struct GTPSAHessianprep{X,M} <: Hessianprep xt::X m::M end @@ -357,82 +359,82 @@ function DI.prepare_hessian(f, backend::AutoGTPSA{D}, x) where {D} j += 1 end - return GTPSAHessianExtras(xt, m) + return GTPSAHessianprep(xt, m) end -function DI.hessian(f, extras::GTPSAHessianExtras, ::AutoGTPSA{D}, x) where {D} - foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part - yt = f(extras.xt) +function DI.hessian(f, prep::GTPSAHessianprep, ::AutoGTPSA{D}, x) where {D} + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part + yt = f(prep.xt) hess = similar(x, eltype(eltype(yt)), (length(x), length(x))) unsafe_fast = D == Nothing ? true : false GTPSA.hessian!( - hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=extras.m + hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=prep.m ) return hess end -function DI.hessian!(f, hess, extras::GTPSAHessianExtras, ::AutoGTPSA{D}, x) where {D} - foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part - yt = f(extras.xt) +function DI.hessian!(f, hess, prep::GTPSAHessianprep, ::AutoGTPSA{D}, x) where {D} + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part + yt = f(prep.xt) unsafe_fast = D == Nothing ? true : false GTPSA.hessian!( - hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=extras.m + hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=prep.m ) return hess end function DI.value_gradient_and_hessian( - f, extras::GTPSAHessianExtras, ::AutoGTPSA{D}, x + f, prep::GTPSAHessianprep, ::AutoGTPSA{D}, x ) where {D} - foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part - yt = f(extras.xt) + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part + yt = f(prep.xt) y = map(t -> t[0], yt) grad = similar(x, eltype(eltype(yt))) GTPSA.gradient!(grad, yt; include_params=true) hess = similar(x, eltype(eltype(yt)), (length(x), length(x))) unsafe_fast = D == Nothing ? true : false GTPSA.hessian!( - hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=extras.m + hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=prep.m ) return y, grad, hess end function DI.value_gradient_and_hessian!( - f, grad, hess, extras::GTPSAHessianExtras, ::AutoGTPSA{D}, x + f, grad, hess, prep::GTPSAHessianprep, ::AutoGTPSA{D}, x ) where {D} - foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part - yt = f(extras.xt) + foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part + yt = f(prep.xt) y = map(t -> t[0], yt) GTPSA.gradient!(grad, yt; include_params=true) unsafe_fast = D == Nothing ? true : false GTPSA.hessian!( - hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=extras.m + hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=prep.m ) return y, grad, hess end # HVP -struct GTPSAHVPExtras{E,H} <: HVPExtras - hessextras::E +struct GTPSAHVPprep{E,H} <: HVPprep + hessprep::E hess::H end function DI.prepare_hvp(f, backend::AutoGTPSA{D}, x, tx::Tangents) where {D} - hessextras = DI.prepare_hessian(f, backend, x) + hessprep = DI.prepare_hessian(f, backend, x) hess = similar(x, typeof(f(x)), (length(x), length(x))) - return GTPSAHVPExtras(hessextras, hess) + return GTPSAHVPprep(hessprep, hess) end -function DI.hvp(f, extras::GTPSAHVPExtras, backend::AutoGTPSA{D}, x, tx::Tangents) where {D} - DI.hessian!(f, extras.hess, extras.hessextras, backend, x) +function DI.hvp(f, prep::GTPSAHVPprep, backend::AutoGTPSA{D}, x, tx::Tangents) where {D} + DI.hessian!(f, prep.hess, prep.hessprep, backend, x) tg = map(tx) do dx - dg = similar(x, eltype(extras.hess)) + dg = similar(x, eltype(prep.hess)) dg .= 0 j = 1 for dxi in dx - for i in 1:size(extras.hess, 2) - dg[i] += extras.hess[i, j] * dxi + for i in 1:size(prep.hess, 2) + dg[i] += prep.hess[i, j] * dxi end j += 1 end @@ -442,19 +444,20 @@ function DI.hvp(f, extras::GTPSAHVPExtras, backend::AutoGTPSA{D}, x, tx::Tangent end function DI.hvp!( - f, tg::Tangents, extras::GTPSAHVPExtras, backend::AutoGTPSA{D}, x, tx::Tangents + f, tg::Tangents, prep::GTPSAHVPprep, backend::AutoGTPSA{D}, x, tx::Tangents ) where {D} - DI.hessian!(f, extras.hess, extras.hessextras, backend, x) + DI.hessian!(f, prep.hess, prep.hessprep, backend, x) for b in eachindex(tg.d) dx, dg = tx.d[b], tg.d[b] dg .= 0 j = 1 for dxi in dx - for i in 1:size(extras.hess, 2) - dg[i] += extras.hess[i, j] * dxi + for i in 1:size(prep.hess, 2) + dg[i] += prep.hess[i, j] * dxi end j += 1 end end return tg end +=# \ No newline at end of file diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/utils.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/utils.jl index 0e223f932..6ba3d5b30 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/utils.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/utils.jl @@ -1,12 +1,35 @@ -function initialize!(xt::TPS, x::Number, dx::Number) - xt[0] = x - xt[1] = dx +function initialize!(xt::TPS, x::Union{Number,Nothing}, dx::Union{Number,Nothing}, varidx::Integer=1) + if !isnothing(x) + xt[0] = x + end + if !isnothing(dx) + xt[varidx] = dx + end return xt end -function initialize!(xt::AbstractArray{<:TPS}, x::AbstractArray, dx::AbstractArray) - for i in eachindex(xt, x, dx) - initialize!(xt[i], x[i], dx[i]) + +function initialize!(xt::AbstractArray{TPS{T}}, x::AbstractArray, dx::AbstractArray, varidx=1) where {T} + if varidx isa Number + for i in eachindex(xt, x, dx, dx) + initialize!(xt[i], x[i], dx[i], varidx) + end + else + for i in eachindex(xt, x, dx, varidxs) + initialize!(xt[i], x[i], dx[i], varidx[i]) + end + end +end +#= +function initialize!(xt::AbstractArray{TPS{T}}, x::AbstractArray, dx::AbstractArray, varidx=1) where {T} + if varidx isa Number + for i in eachindex(xt, x, dx, dx) + initialize!(xt[i], x[i], dx[i], varidx) + end + else + for i in eachindex(xt, x, dx, varidxs) + initialize!(xt[i], x[i], dx[i], varidx[i]) + end end - return xt end +=# From 97a2f5a4586012dd26bc4754fe7ccbd6b8d055f3 Mon Sep 17 00:00:00 2001 From: Matt Signorelli Date: Mon, 6 Jan 2025 21:40:05 -0500 Subject: [PATCH 33/44] one arg done --- .../onearg.jl | 312 +++++++++++++----- 1 file changed, 224 insertions(+), 88 deletions(-) diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl index 03441af8e..e32914eb6 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl @@ -85,7 +85,7 @@ function DI.value_and_pushforward( fc = DI.with_contexts(f, contexts...) ty = DI.pushforward(fc, prep, backend, x, tx) - y = fc(x) # This actually seems faster than indexing the TPSs + y = fc(x) # TO-DO: optimize return y, ty end @@ -100,7 +100,7 @@ function DI.value_and_pushforward!( ) where {C} fc = DI.with_contexts(f, contexts...) DI.pushforward!(fc, ty, prep, backend, x, tx) - y = fc(x) + y = fc(x) # TO-DO: optimize return y, ty end @@ -157,10 +157,10 @@ function DI.gradient( if prep.xt isa Number return yt[1] else - grad = similar(x, GTPSA.numtype(eltype(yt))) + grad = similar(x, GTPSA.numtype(yt)) GTPSA.gradient!(grad, yt; include_params=true) + return grad end - return grad end function DI.gradient!( @@ -172,114 +172,177 @@ function DI.gradient!( contexts::Vararg{DI.Context,C}, ) where {C} foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part - yt = f(prep.xt) + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) GTPSA.gradient!(grad, yt; include_params=true) return grad end -#= -function DI.value_and_gradient(f, prep::GTPSAGradientPrep, ::AutoGTPSA, x) - foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part - yt = f(prep.xt) - grad = similar(x, eltype(eltype(yt))) - GTPSA.gradient!(grad, yt; include_params=true) - y = map(t -> t[0], yt) - return y, grad +function DI.value_and_gradient( + f, + prep::GTPSAOneArgGradientPrep, + ::AutoGTPSA, + x, + contexts::Vararg{DI.Context,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) + if prep.xt isa Number + return yt[0], yt[1] + else + grad = similar(x, GTPSA.numtype(yt)) + GTPSA.gradient!(grad, yt; include_params=true) + return yt[0], grad + end end -function DI.value_and_gradient!(f, grad, prep::GTPSAGradientPrep, ::AutoGTPSA, x) - foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part - yt = f(prep.xt) +function DI.value_and_gradient!( + f, + grad, + prep::GTPSAOneArgGradientPrep, + ::AutoGTPSA, + x, + contexts::Vararg{DI.Context,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) - y = map(t -> t[0], yt) - return y, grad + return yt[0], grad end ## Jacobian - -struct GTPSAJacobianprep{X} <: Jacobianprep +# Contains a vector of pre-allocated TPSs +struct GTPSAOneArgJacobianPrep{X} <: DI.JacobianPrep xt::X end -function DI.prepare_jacobian(f, backend::AutoGTPSA{D}, x) where {D} +# To materialize the entire Jacobian we use all variables +function DI.prepare_jacobian( + f, + backend::AutoGTPSA{D}, + x, + contexts::Vararg{DI.Context,C}, +) where {D,C} if D != Nothing d = backend.descriptor - nn = GTPSA.numnn(d) + length(x) == GTPSA.numnn(d) || error("Number of variables + parameters in Descriptor disagrees with function number of inputs!") else - nn = length(x) - d = Descriptor(nn, 1) + d = Descriptor(length(x), 1) # n variables to first order 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. + # 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 GTPSAJacobianprep(xt) + return GTPSAOneArgJacobianPrep(xt) end -function DI.jacobian(f, prep::GTPSAJacobianprep, ::AutoGTPSA, x) +function DI.jacobian( + f, + prep::GTPSAOneArgJacobianPrep, + ::AutoGTPSA, + x, + contexts::Vararg{DI.Context,C}, +) where {C} foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part - yt = f(prep.xt) - jac = similar(x, eltype(eltype(yt)), (length(yt), length(x))) + 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) return jac end -function DI.jacobian!(f, jac, prep::GTPSAJacobianprep, ::AutoGTPSA, x) +function DI.jacobian!( + f, + jac, + prep::GTPSAOneArgJacobianPrep, + ::AutoGTPSA, + x, + contexts::Vararg{DI.Context,C}, +) where {C} foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part - yt = f(prep.xt) + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) GTPSA.jacobian!(jac, yt; include_params=true) return jac end -function DI.value_and_jacobian(f, prep::GTPSAJacobianprep, ::AutoGTPSA, x) + +function DI.value_and_jacobian( + f, + prep::GTPSAOneArgJacobianPrep, + ::AutoGTPSA, + x, + contexts::Vararg{DI.Context,C}, +) where {C} foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part - yt = f(prep.xt) - jac = similar(x, eltype(eltype(yt)), (length(yt), length(x))) + 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) y = map(t -> t[0], yt) return y, jac end -function DI.value_and_jacobian!(f, jac, prep::GTPSAJacobianprep, ::AutoGTPSA, x) +function DI.value_and_jacobian!( + f, + jac, + prep::GTPSAOneArgJacobianPrep, + ::AutoGTPSA, + x, + contexts::Vararg{DI.Context,C}, +) where {C} foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part - yt = f(prep.xt) + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) GTPSA.jacobian!(jac, yt; include_params=true) y = map(t -> t[0], yt) return y, jac end -## Second derivative -struct GTPSASecondDerivativeprep{X} <: SecondDerivativeprep +## 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) where {D} +function DI.prepare_second_derivative( + f, + backend::AutoGTPSA{D}, + x, + contexts::Vararg{DI.Context,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 - return GTPSASecondDerivativeprep(xt) + xt[1] = 1 # Set slope + return GTPSAOneArgSecondDerivativePrep(xt) end -function DI.second_derivative(f, prep::GTPSASecondDerivativeprep, ::AutoGTPSA, x) +function DI.second_derivative( + f, + prep::GTPSAOneArgSecondDerivativePrep, + ::AutoGTPSA, + x, + contexts::Vararg{DI.Context,C}, +) where {C} prep.xt[0] = x - yt = f(prep.xt) + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) if yt isa Number return yt[2] * 2 else - der2 = similar(yt, eltype(eltype(yt))) + der2 = similar(yt, GTPSA.numtype(eltype(yt))) for i in eachindex(yt) der2[i] = yt[i][2] * 2 # *2 because monomial coefficient is 1/2 end @@ -287,9 +350,17 @@ function DI.second_derivative(f, prep::GTPSASecondDerivativeprep, ::AutoGTPSA, x end end -function DI.second_derivative!(f, der2, prep::GTPSASecondDerivativeprep, ::AutoGTPSA, x) +function DI.second_derivative!( + f, + der2, + prep::GTPSAOneArgSecondDerivativePrep, + ::AutoGTPSA, + x, + contexts::Vararg{DI.Context,C}, +) where {C} prep.xt[0] = x - yt = f(prep.xt) + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) for i in eachindex(yt) der2[i] = yt[i][2] * 2 end @@ -297,16 +368,21 @@ function DI.second_derivative!(f, der2, prep::GTPSASecondDerivativeprep, ::AutoG end function DI.value_derivative_and_second_derivative( - f, prep::GTPSASecondDerivativeprep, ::AutoGTPSA, x -) + f, + prep::GTPSAOneArgSecondDerivativePrep, + ::AutoGTPSA, + x, + contexts::Vararg{DI.Context,C}, +) where {C} prep.xt[0] = x - yt = f(prep.xt) + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) if yt isa Number return yt[0], yt[1], yt[2] * 2 else y = map(t -> t[0], yt) - der = similar(yt, eltype(eltype(yt))) - der2 = similar(yt, eltype(eltype(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][2] * 2 @@ -316,10 +392,17 @@ function DI.value_derivative_and_second_derivative( end function DI.value_derivative_and_second_derivative!( - f, der, der2, prep::GTPSASecondDerivativeprep, ::AutoGTPSA, x -) + f, + der, + der2, + prep::GTPSAOneArgSecondDerivativePrep, + ::AutoGTPSA, + x, + contexts::Vararg{DI.Context,C}, +) where {C} prep.xt[0] = x - yt = f(prep.xt) + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) y = map(t -> t[0], yt) for i in eachindex(yt) der[i] = yt[i][1] @@ -329,16 +412,23 @@ function DI.value_derivative_and_second_derivative!( end ## Hessian - -struct GTPSAHessianprep{X,M} <: Hessianprep +# 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) where {D} +function DI.prepare_hessian( + f, + backend::AutoGTPSA{D}, + x, + contexts::Vararg{DI.Context,C}, +) where {D,C} if D != Nothing d = backend.descriptor nn = GTPSA.numnn(d) + length(x) == nn || error("Number of variables + parameters in Descriptor disagrees with function number of inputs!") m = Vector{UInt8}(undef, nn) else nn = length(x) @@ -359,13 +449,20 @@ function DI.prepare_hessian(f, backend::AutoGTPSA{D}, x) where {D} j += 1 end - return GTPSAHessianprep(xt, m) + return GTPSAOneArgHessianPrep(xt, m) end -function DI.hessian(f, prep::GTPSAHessianprep, ::AutoGTPSA{D}, x) where {D} +function DI.hessian( + f, + prep::GTPSAOneArgHessianPrep, + ::AutoGTPSA{D}, + x, + contexts::Vararg{DI.Context,C}, +) where {D,C} foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part - yt = f(prep.xt) - hess = similar(x, eltype(eltype(yt)), (length(x), length(x))) + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) + hess = similar(x, GTPSA.numtype(eltype(yt)), (length(x), length(x))) unsafe_fast = D == Nothing ? true : false GTPSA.hessian!( hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=prep.m @@ -373,9 +470,17 @@ function DI.hessian(f, prep::GTPSAHessianprep, ::AutoGTPSA{D}, x) where {D} return hess end -function DI.hessian!(f, hess, prep::GTPSAHessianprep, ::AutoGTPSA{D}, x) where {D} +function DI.hessian!( + f, + hess, + prep::GTPSAOneArgHessianPrep, + ::AutoGTPSA{D}, + x, + contexts::Vararg{DI.Context,C}, +) where {D,C} foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part - yt = f(prep.xt) + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) unsafe_fast = D == Nothing ? true : false GTPSA.hessian!( hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=prep.m @@ -383,15 +488,21 @@ function DI.hessian!(f, hess, prep::GTPSAHessianprep, ::AutoGTPSA{D}, x) where { return hess end + function DI.value_gradient_and_hessian( - f, prep::GTPSAHessianprep, ::AutoGTPSA{D}, x -) where {D} + f, + prep::GTPSAOneArgHessianPrep, + ::AutoGTPSA{D}, + x, + contexts::Vararg{DI.Context,C}, +) where {D,C} foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part - yt = f(prep.xt) + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) y = map(t -> t[0], yt) - grad = similar(x, eltype(eltype(yt))) + grad = similar(x, GTPSA.numtype(eltype(yt))) GTPSA.gradient!(grad, yt; include_params=true) - hess = similar(x, eltype(eltype(yt)), (length(x), length(x))) + hess = similar(x, GTPSA.numtype(eltype(yt)), (length(x), length(x))) unsafe_fast = D == Nothing ? true : false GTPSA.hessian!( hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=prep.m @@ -399,11 +510,19 @@ function DI.value_gradient_and_hessian( return y, grad, hess end + function DI.value_gradient_and_hessian!( - f, grad, hess, prep::GTPSAHessianprep, ::AutoGTPSA{D}, x -) where {D} + f, + grad, + hess, + prep::GTPSAOneArgHessianPrep, + ::AutoGTPSA{D}, + x, + contexts::Vararg{DI.Context,C}, +) where {D,C} foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part - yt = f(prep.xt) + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.xt) y = map(t -> t[0], yt) GTPSA.gradient!(grad, yt; include_params=true) unsafe_fast = D == Nothing ? true : false @@ -413,21 +532,33 @@ function DI.value_gradient_and_hessian!( return y, grad, hess end -# HVP - -struct GTPSAHVPprep{E,H} <: HVPprep +struct GTPSAOneArgHVPPrep{E,H} <: DI.HVPPrep hessprep::E hess::H end -function DI.prepare_hvp(f, backend::AutoGTPSA{D}, x, tx::Tangents) where {D} +function DI.prepare_hvp( + f, + backend::AutoGTPSA, + x, + tx::NTuple, + contexts::Vararg{DI.Context,C}, +) where {C} hessprep = DI.prepare_hessian(f, backend, x) - hess = similar(x, typeof(f(x)), (length(x), length(x))) - return GTPSAHVPprep(hessprep, hess) + 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::GTPSAHVPprep, backend::AutoGTPSA{D}, x, tx::Tangents) where {D} - DI.hessian!(f, prep.hess, prep.hessprep, backend, x) +function DI.hvp( + f, + prep::GTPSAOneArgHVPPrep, + backend::AutoGTPSA, + x, + tx::NTuple, + contexts::Vararg{DI.Context,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 @@ -444,11 +575,17 @@ function DI.hvp(f, prep::GTPSAHVPprep, backend::AutoGTPSA{D}, x, tx::Tangents) w end function DI.hvp!( - f, tg::Tangents, prep::GTPSAHVPprep, backend::AutoGTPSA{D}, x, tx::Tangents -) where {D} - DI.hessian!(f, prep.hess, prep.hessprep, backend, x) + f, + tg::NTuple, + prep::GTPSAOneArgHVPPrep, + backend::AutoGTPSA, + x, + tx::NTuple, + contexts::Vararg{DI.Context,C}, +) where {C} + DI.hessian!(f, prep.hess, prep.hessprep, backend, x, contexts...) for b in eachindex(tg.d) - dx, dg = tx.d[b], tg.d[b] + dx, dg = tx[b], tg[b] dg .= 0 j = 1 for dxi in dx @@ -460,4 +597,3 @@ function DI.hvp!( end return tg end -=# \ No newline at end of file From 9a3fddf99887ecd927094e87e051466b160000f6 Mon Sep 17 00:00:00 2001 From: Matt Signorelli Date: Tue, 7 Jan 2025 16:19:06 -0500 Subject: [PATCH 34/44] all tests except hvp passing, should have full code cov --- .../DifferentiationInterfaceGTPSAExt.jl | 4 +- .../onearg.jl | 169 ++++++++---- .../twoarg.jl | 241 +++++++++--------- .../DifferentiationInterfaceGTPSAExt/utils.jl | 35 --- .../test/Back/GTPSA/test.jl | 13 +- 5 files changed, 262 insertions(+), 200 deletions(-) delete mode 100644 DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/utils.jl diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl index a5017c6cb..4e5101772 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl @@ -3,11 +3,11 @@ module DifferentiationInterfaceGTPSAExt import DifferentiationInterface as DI using ADTypes: AutoGTPSA using GTPSA +using LinearAlgebra DI.check_available(::AutoGTPSA) = true -include("utils.jl") include("onearg.jl") -#include("twoarg.jl") +include("twoarg.jl") end diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl index e32914eb6..f844df6e5 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl @@ -40,10 +40,10 @@ function DI.pushforward( x, tx::NTuple, contexts::Vararg{DI.Context,C}, -) where {C} +) where {C} fc = DI.with_contexts(f, contexts...) ty = map(tx) do dx - initialize!(prep.xt, x, dx, 1) + 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] @@ -67,7 +67,7 @@ function DI.pushforward!( fc = DI.with_contexts(f, contexts...) for b in eachindex(tx, ty) dx, dy = tx[b], ty[b] - initialize!(prep.xt, x, dx, 1) + 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 @@ -93,7 +93,7 @@ function DI.value_and_pushforward!( f, ty::NTuple, prep::GTPSAOneArgPushforwardPrep, - ::AutoGTPSA, + backend::AutoGTPSA, x, tx::NTuple, contexts::Vararg{DI.Context,C}, @@ -121,7 +121,6 @@ function DI.prepare_gradient( if D != Nothing d = backend.descriptor - length(x) == GTPSA.numnn(d) || error("Number of variables + parameters in Descriptor disagrees with function number of inputs!") else d = Descriptor(length(x), 1) # n variables to first order end @@ -158,7 +157,7 @@ function DI.gradient( return yt[1] else grad = similar(x, GTPSA.numtype(yt)) - GTPSA.gradient!(grad, yt; include_params=true) + GTPSA.gradient!(grad, yt; include_params=true, unsafe_size=true) return grad end end @@ -174,7 +173,7 @@ function DI.gradient!( 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) + GTPSA.gradient!(grad, yt; include_params=true, unsafe_size=true) return grad end @@ -192,7 +191,7 @@ function DI.value_and_gradient( return yt[0], yt[1] else grad = similar(x, GTPSA.numtype(yt)) - GTPSA.gradient!(grad, yt; include_params=true) + GTPSA.gradient!(grad, yt; include_params=true, unsafe_size=true) return yt[0], grad end end @@ -208,7 +207,7 @@ function DI.value_and_gradient!( 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) + GTPSA.gradient!(grad, yt; include_params=true, unsafe_size=true) return yt[0], grad end @@ -227,7 +226,6 @@ function DI.prepare_jacobian( ) where {D,C} if D != Nothing d = backend.descriptor - length(x) == GTPSA.numnn(d) || error("Number of variables + parameters in Descriptor disagrees with function number of inputs!") else d = Descriptor(length(x), 1) # n variables to first order end @@ -254,7 +252,7 @@ function DI.jacobian( 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) + GTPSA.jacobian!(jac, yt; include_params=true, unsafe_size=true) return jac end @@ -269,7 +267,7 @@ function DI.jacobian!( 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) + GTPSA.jacobian!(jac, yt; include_params=true, unsafe_size=true) return jac end @@ -285,7 +283,7 @@ function DI.value_and_jacobian( 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) + GTPSA.jacobian!(jac, yt; include_params=true, unsafe_size=true) y = map(t -> t[0], yt) return y, jac end @@ -301,7 +299,7 @@ function DI.value_and_jacobian!( 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) + GTPSA.jacobian!(jac, yt; include_params=true, unsafe_size=true) y = map(t -> t[0], yt) return y, jac end @@ -332,19 +330,25 @@ end function DI.second_derivative( f, prep::GTPSAOneArgSecondDerivativePrep, - ::AutoGTPSA, + backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,C}, -) where {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[2] * 2 + return yt[idx2] * 2 else der2 = similar(yt, GTPSA.numtype(eltype(yt))) for i in eachindex(yt) - der2[i] = yt[i][2] * 2 # *2 because monomial coefficient is 1/2 + der2[i] = yt[i][idx2] * 2 # *2 because monomial coefficient is 1/2 end return der2 end @@ -354,15 +358,20 @@ function DI.second_derivative!( f, der2, prep::GTPSAOneArgSecondDerivativePrep, - ::AutoGTPSA, + backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,C}, -) where {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][2] * 2 + der2[i] = yt[i][idx2] * 2 end return der2 end @@ -370,22 +379,27 @@ end function DI.value_derivative_and_second_derivative( f, prep::GTPSAOneArgSecondDerivativePrep, - ::AutoGTPSA, + backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,C}, -) where {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[2] * 2 + 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][2] * 2 + der2[i] = yt[i][idx2] * 2 end return y, der, der2 end @@ -396,17 +410,22 @@ function DI.value_derivative_and_second_derivative!( der, der2, prep::GTPSAOneArgSecondDerivativePrep, - ::AutoGTPSA, + backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,C}, -) where {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][2] * 2 + der2[i] = yt[i][idx2] * 2 end return y, der, der2 end @@ -427,9 +446,7 @@ function DI.prepare_hessian( ) where {D,C} if D != Nothing d = backend.descriptor - nn = GTPSA.numnn(d) - length(x) == nn || error("Number of variables + parameters in Descriptor disagrees with function number of inputs!") - m = Vector{UInt8}(undef, nn) + m = Vector{UInt8}(undef, length(x)) else nn = length(x) d = Descriptor(nn, 2) @@ -462,10 +479,10 @@ function DI.hessian( 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(eltype(yt)), (length(x), length(x))) + hess = similar(x, GTPSA.numtype(yt), (length(x), length(x))) unsafe_fast = D == Nothing ? true : false GTPSA.hessian!( - hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=prep.m + hess, yt; include_params=true, unsafe_size=true, unsafe_fast=unsafe_fast, tmp_mono=prep.m ) return hess end @@ -483,7 +500,7 @@ function DI.hessian!( yt = fc(prep.xt) unsafe_fast = D == Nothing ? true : false GTPSA.hessian!( - hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=prep.m + hess, yt; include_params=true, unsafe_size=true, unsafe_fast=unsafe_fast, tmp_mono=prep.m ) return hess end @@ -499,15 +516,14 @@ function DI.value_gradient_and_hessian( foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part fc = DI.with_contexts(f, contexts...) yt = fc(prep.xt) - y = map(t -> t[0], yt) - grad = similar(x, GTPSA.numtype(eltype(yt))) - GTPSA.gradient!(grad, yt; include_params=true) - hess = similar(x, GTPSA.numtype(eltype(yt)), (length(x), length(x))) + grad = similar(x, GTPSA.numtype(yt)) + GTPSA.gradient!(grad, yt; include_params=true, unsafe_size=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_fast=unsafe_fast, tmp_mono=prep.m + hess, yt; include_params=true, unsafe_size=true, unsafe_fast=unsafe_fast, tmp_mono=prep.m ) - return y, grad, hess + return yt[0], grad, hess end @@ -523,13 +539,12 @@ function DI.value_gradient_and_hessian!( foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part fc = DI.with_contexts(f, contexts...) yt = fc(prep.xt) - y = map(t -> t[0], yt) - GTPSA.gradient!(grad, yt; include_params=true) + GTPSA.gradient!(grad, yt; include_params=true, unsafe_size=true) unsafe_fast = D == Nothing ? true : false GTPSA.hessian!( - hess, yt; include_params=true, unsafe_fast=unsafe_fast, tmp_mono=prep.m + hess, yt; include_params=true, unsafe_size=true, unsafe_fast=unsafe_fast, tmp_mono=prep.m ) - return y, grad, hess + return yt[0], grad, hess end struct GTPSAOneArgHVPPrep{E,H} <: DI.HVPPrep @@ -559,6 +574,7 @@ function DI.hvp( contexts::Vararg{DI.Context,C}, ) where {C} DI.hessian!(f, prep.hess, prep.hessprep, backend, x, contexts...) + #return prep.hess tg = map(tx) do dx dg = similar(x, eltype(prep.hess)) dg .= 0 @@ -584,7 +600,7 @@ function DI.hvp!( contexts::Vararg{DI.Context,C}, ) where {C} DI.hessian!(f, prep.hess, prep.hessprep, backend, x, contexts...) - for b in eachindex(tg.d) + for b in eachindex(tg) dx, dg = tx[b], tg[b] dg .= 0 j = 1 @@ -597,3 +613,66 @@ function DI.hvp!( end return tg end + +function DI.gradient_and_hvp( + f, + prep::GTPSAOneArgHVPPrep, + ::AutoGTPSA{D}, + x, + tx::NTuple, + contexts::Vararg{DI.Context,C}, +) where {D,C} + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.hessprep.xt) + grad = similar(x, GTPSA.numtype(yt)) + GTPSA.gradient!(grad, yt; include_params=true, unsafe_size=true) + unsafe_fast = D == Nothing ? true : false + GTPSA.hessian!( + prep.hess, yt; include_params=true, unsafe_size=true, unsafe_fast=unsafe_fast, tmp_mono=prep.hessprep.m + ) + 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, 2) + 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, + ::AutoGTPSA{D}, + x, + tx::NTuple, + contexts::Vararg{DI.Context,C}, +) where {D,C} + fc = DI.with_contexts(f, contexts...) + yt = fc(prep.hessprep.xt) + GTPSA.gradient!(grad, yt; include_params=true, unsafe_size=true) + unsafe_fast = D == Nothing ? true : false + GTPSA.hessian!( + prep.hess, yt; include_params=true, unsafe_size=true, unsafe_fast=unsafe_fast, tmp_mono=prep.hessprep.m + ) + 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, 2) + dg[i] += prep.hess[i, j] * dxi + end + j += 1 + end + return dg + end + return grad, tg +end \ No newline at end of file diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl index 849d0d1f3..8850f918b 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl @@ -1,25 +1,36 @@ ## Pushforward -struct GTPSATwoArgPushforwardExtras{X,Y} <: PushforwardExtras + +# 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::Tangents) where {D} +function DI.prepare_pushforward( + ::F, + y, + backend::AutoGTPSA{D}, + x, + tx::NTuple, + ::Vararg{DI.Context,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(length(x), 1) + d = Descriptor(1, 1) # 1 variable to first order end - - dx = first(tx) if x isa Number - xt = TPS{promote_type(typeof(dx), typeof(x), Float64)}(; use=d) + xt = TPS{promote_type(typeof(first(tx)), typeof(x), Float64)}(; use=d) else - xt = similar(x, TPS{promote_type(eltype(dx), eltype(x), Float64)}) - - # xt and x have same indexing because of similar + xt = similar(x, TPS{promote_type(eltype(first(tx)), eltype(x), Float64)}) for i in eachindex(xt) - xt[i] = TPS{promote_type(eltype(dx), eltype(x), Float64)}(; use=d) + xt[i] = TPS{promote_type(eltype(first(tx)), eltype(x), Float64)}(; use=d) end end @@ -28,127 +39,100 @@ function DI.prepare_pushforward(f!, y, backend::AutoGTPSA{D}, x, tx::Tangents) w for i in eachindex(yt) yt[i] = TPS{promote_type(eltype(y), Float64)}(; use=d) end - return GTPSATwoArgPushforwardExtras(xt, yt) + return GTPSATwoArgPushforwardPrep(xt, yt) end function DI.pushforward( - f!, y, extras::GTPSATwoArgPushforwardExtras, backend::AutoGTPSA, x, tx::Tangents -) + f!, + y, + prep::GTPSATwoArgPushforwardPrep, + ::AutoGTPSA, + x, + tx::NTuple, + contexts::Vararg{DI.Context,C}, +) where {C} + fc! = DI.with_contexts(f!, contexts...) ty = map(tx) do dx - initialize!(extras.xt, x, dx) - - f!(extras.yt, extras.xt) - map!(t -> t[0], y, extras.yt) - dy = similar(extras.yt, eltype(eltype(extras.yt))) - dy .= 0 - for i in eachindex(extras.yt) - for j in 1:length(dx) - dy[i] += extras.yt[i][j] - end - end + 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::Tangents, - extras::GTPSATwoArgPushforwardExtras, - backend::AutoGTPSA, + ty::NTuple, + prep::GTPSATwoArgPushforwardPrep, + ::AutoGTPSA, x, - tx::Tangents, -) - for b in eachindex(tx.d, ty.d) - dx, dy = tx.d[b], ty.d[b] - initialize!(extras.xt, x, dx) - - f!(extras.yt, extras.xt) - map!(t -> t[0], y, extras.yt) - dy .= 0 - for i in eachindex(extras.yt) - for j in 1:length(dx) - dy[i] += extras.yt[i][j] - end - end + tx::NTuple, + contexts::Vararg{DI.Context,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, extras::GTPSATwoArgPushforwardExtras, backend::AutoGTPSA, x, tx::Tangents -) - ys_and_dys = map(tx.d) do dx - initialize!(extras.xt, x, dx) - - f!(extras.yt, extras.xt) - dy = similar(extras.yt, eltype(eltype(extras.yt))) - dy .= 0 - for i in eachindex(extras.yt) - for j in 1:length(dx) - dy[i] += extras.yt[i][j] - end - end - map!(t -> t[0], y, extras.yt) - - return y, dy - end - - y = first(ys_and_dys[1]) - dys = last.(ys_and_dys) - ty = Tangents(dys...) + f!, + y, + prep::GTPSATwoArgPushforwardPrep, + backend::AutoGTPSA, + x, + tx::NTuple, + contexts::Vararg{DI.Context,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::Tangents, - extras::GTPSATwoArgPushforwardExtras, + ty::NTuple, + prep::GTPSATwoArgPushforwardPrep, backend::AutoGTPSA, x, - tx::Tangents, -) - for b in eachindex(tx.d, ty.d) - dx, dy = tx.d[b], ty.d[b] - initialize!(extras.xt, x, dx) - - f!(extras.yt, extras.xt) - map!(t -> t[0], y, extras.yt) - dy .= 0 - for i in eachindex(extras.yt) - for j in 1:length(dx) - dy[i] += extras.yt[i][j] - end - end - end - + tx::NTuple, + contexts::Vararg{DI.Context,C}, +) where {C} + DI.pushforward!(f!, y, ty, prep, backend, x, tx, contexts...) return y, ty end ## Jacobian - -struct GTPSATwoArgJacobianExtras{X,Y} <: JacobianExtras +# 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) where {D} +function DI.prepare_jacobian( + f, + y, + backend::AutoGTPSA{D}, + x, + contexts::Vararg{DI.Context,C}, +) where {D,C} if D != Nothing d = backend.descriptor - nn = GTPSA.numnn(d) else - nn = length(x) - d = Descriptor(nn, 1) + d = Descriptor(length(x), 1) # n variables to first order 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. + # 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) @@ -162,41 +146,64 @@ function DI.prepare_jacobian(f!, y, backend::AutoGTPSA{D}, x) where {D} yt[i] = TPS{promote_type(eltype(y), Float64)}(; use=d) end - return GTPSATwoArgJacobianExtras(xt, yt) + return GTPSATwoArgJacobianPrep(xt, yt) end -function DI.jacobian(f!, y, extras::GTPSATwoArgJacobianExtras, ::AutoGTPSA, x) - foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part - f!(extras.yt, extras.xt) - jac = similar(x, eltype(eltype(extras.yt)), (length(extras.yt), length(x))) - GTPSA.jacobian!(jac, extras.yt; include_params=true) - map!(t -> t[0], y, extras.yt) +function DI.jacobian( + f!, + y, + prep::GTPSATwoArgJacobianPrep, + ::AutoGTPSA, + x, + contexts::Vararg{DI.Context,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_size=true) + map!(t -> t[0], y, prep.yt) return jac end -function DI.jacobian!(f!, y, jac, extras::GTPSATwoArgJacobianExtras, ::AutoGTPSA, x) - foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part - f!(extras.yt, extras.xt) - GTPSA.jacobian!(jac, extras.yt; include_params=true) - map!(t -> t[0], y, extras.yt) +function DI.jacobian!( + f!, + y, + jac, + prep::GTPSATwoArgJacobianPrep, + ::AutoGTPSA, + x, + contexts::Vararg{DI.Context,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_size=true) + map!(t -> t[0], y, prep.yt) return jac end -function DI.value_and_jacobian(f!, y, extras::GTPSATwoArgJacobianExtras, ::AutoGTPSA, x) - foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part - f!(extras.yt, extras.xt) - jac = similar(x, eltype(eltype(extras.yt)), (length(extras.yt), length(x))) - GTPSA.jacobian!(jac, extras.yt; include_params=true) - map!(t -> t[0], y, extras.yt) +function DI.value_and_jacobian( + f!, + y, + prep::GTPSATwoArgJacobianPrep, + backend::AutoGTPSA, + x, + contexts::Vararg{DI.Context,C}, +) where {C} + jac = DI.jacobian(f!, y, prep, backend, x, contexts...) return y, jac end -function DI.value_and_jacobian!( - f!, y, jac, extras::GTPSATwoArgJacobianExtras, ::AutoGTPSA, x -) - foreach((t, xi) -> t[0] = xi, extras.xt, x) # Set the scalar part - f!(extras.yt, extras.xt) - GTPSA.jacobian!(jac, extras.yt; include_params=true) - map!(t -> t[0], y, extras.yt) +function DI.value_and_jacobian( + f!, + y, + jac, + prep::GTPSATwoArgJacobianPrep, + backend::AutoGTPSA, + x, + contexts::Vararg{DI.Context,C}, +) where {C} + DI.jacobian!(f!, y, jac, prep, backend, x, contexts...) return y, jac end diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/utils.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/utils.jl deleted file mode 100644 index 6ba3d5b30..000000000 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/utils.jl +++ /dev/null @@ -1,35 +0,0 @@ -function initialize!(xt::TPS, x::Union{Number,Nothing}, dx::Union{Number,Nothing}, varidx::Integer=1) - if !isnothing(x) - xt[0] = x - end - if !isnothing(dx) - xt[varidx] = dx - end - return xt -end - - -function initialize!(xt::AbstractArray{TPS{T}}, x::AbstractArray, dx::AbstractArray, varidx=1) where {T} - if varidx isa Number - for i in eachindex(xt, x, dx, dx) - initialize!(xt[i], x[i], dx[i], varidx) - end - else - for i in eachindex(xt, x, dx, varidxs) - initialize!(xt[i], x[i], dx[i], varidx[i]) - end - end -end -#= -function initialize!(xt::AbstractArray{TPS{T}}, x::AbstractArray, dx::AbstractArray, varidx=1) where {T} - if varidx isa Number - for i in eachindex(xt, x, dx, dx) - initialize!(xt[i], x[i], dx[i], varidx) - end - else - for i in eachindex(xt, x, dx, varidxs) - initialize!(xt[i], x[i], dx[i], varidx[i]) - end - end -end -=# diff --git a/DifferentiationInterface/test/Back/GTPSA/test.jl b/DifferentiationInterface/test/Back/GTPSA/test.jl index c2d6cc715..60db4ff92 100644 --- a/DifferentiationInterface/test/Back/GTPSA/test.jl +++ b/DifferentiationInterface/test/Back/GTPSA/test.jl @@ -12,4 +12,15 @@ for backend in [AutoGTPSA()] @test check_inplace(backend) end -test_differentiation(AutoGTPSA(); type_stability=false, logging=LOGGING); +# Test no Descriptor (use context) +test_differentiation(AutoGTPSA(); type_stability=:full, logging=LOGGING); + +# Test with Descriptor: +d1 = 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 = Descriptor(vos, 3) +test_differentiation(AutoGTPSA(d2); type_stability=:full, logging=LOGGING); From ea203ebe7da6e83ebc89920f2a647b87410a930b Mon Sep 17 00:00:00 2001 From: Matt Signorelli Date: Tue, 7 Jan 2025 19:00:13 -0500 Subject: [PATCH 35/44] bump gtpsa version --- DifferentiationInterface/Project.toml | 2 +- DifferentiationInterfaceTest/Project.toml | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/DifferentiationInterface/Project.toml b/DifferentiationInterface/Project.toml index 0a5e469f4..9836ae8a8 100644 --- a/DifferentiationInterface/Project.toml +++ b/DifferentiationInterface/Project.toml @@ -59,7 +59,7 @@ FastDifferentiation = "0.4.1" FiniteDiff = "2.23.1" FiniteDifferences = "0.12.31" ForwardDiff = "0.10.36" -GTPSA = "1.3.0" +GTPSA = "1.3.1" JuliaFormatter = "1" LinearAlgebra = "<0.0.1,1" Mooncake = "0.4.52" diff --git a/DifferentiationInterfaceTest/Project.toml b/DifferentiationInterfaceTest/Project.toml index cd632972d..f900204d1 100644 --- a/DifferentiationInterfaceTest/Project.toml +++ b/DifferentiationInterfaceTest/Project.toml @@ -50,6 +50,7 @@ FiniteDifferences = "0.12" Flux = "0.16" ForwardDiff = "0.10.36" Functors = "0.4, 0.5" +GTPSA = "1.3.1" JET = "0.4 - 0.8, 0.9" JLArrays = "0.1, 0.2" JuliaFormatter = "1" From 637f66e0933f32a8a097280db1ef05219bee52cc Mon Sep 17 00:00:00 2001 From: Matt Signorelli Date: Wed, 8 Jan 2025 09:48:05 -0500 Subject: [PATCH 36/44] fix hvp functions and docs update --- .../docs/src/explanation/backends.md | 7 ++- .../onearg.jl | 47 ++++++++++--------- .../twoarg.jl | 4 +- DifferentiationInterfaceTest/Project.toml | 1 + 4 files changed, 33 insertions(+), 26 deletions(-) diff --git a/DifferentiationInterface/docs/src/explanation/backends.md b/DifferentiationInterface/docs/src/explanation/backends.md index 080b3f062..2115d4d9b 100644 --- a/DifferentiationInterface/docs/src/explanation/backends.md +++ b/DifferentiationInterface/docs/src/explanation/backends.md @@ -143,7 +143,12 @@ For higher level operators, preparation creates a [config object](https://juliad ### GTPSA -TO-DO +For all operators, preparation preallocates the input [`TPS`s](https://bmad-sim.github.io/GTPSA.jl/stable/man/d_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/k_fastgtpsa/) macros are recommended. + +If a GTPSA [`Descriptor`](https://bmad-sim.github.io/GTPSA.jl/stable/man/c_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 variables + parameters is not less than the number of inputs of the provided function. Undefined behavior and crashes may occur if this is not the case. ### PolyesterForwardDiff diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl index f844df6e5..66a2630ca 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl @@ -157,7 +157,7 @@ function DI.gradient( return yt[1] else grad = similar(x, GTPSA.numtype(yt)) - GTPSA.gradient!(grad, yt; include_params=true, unsafe_size=true) + GTPSA.gradient!(grad, yt; include_params=true, unsafe_inbounds=true) return grad end end @@ -173,7 +173,7 @@ function DI.gradient!( 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_size=true) + GTPSA.gradient!(grad, yt; include_params=true, unsafe_inbounds=true) return grad end @@ -191,7 +191,7 @@ function DI.value_and_gradient( return yt[0], yt[1] else grad = similar(x, GTPSA.numtype(yt)) - GTPSA.gradient!(grad, yt; include_params=true, unsafe_size=true) + GTPSA.gradient!(grad, yt; include_params=true, unsafe_inbounds=true) return yt[0], grad end end @@ -207,7 +207,7 @@ function DI.value_and_gradient!( 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_size=true) + GTPSA.gradient!(grad, yt; include_params=true, unsafe_inbounds=true) return yt[0], grad end @@ -252,7 +252,7 @@ function DI.jacobian( 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_size=true) + GTPSA.jacobian!(jac, yt; include_params=true, unsafe_inbounds=true) return jac end @@ -267,7 +267,7 @@ function DI.jacobian!( 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_size=true) + GTPSA.jacobian!(jac, yt; include_params=true, unsafe_inbounds=true) return jac end @@ -283,7 +283,7 @@ function DI.value_and_jacobian( 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_size=true) + GTPSA.jacobian!(jac, yt; include_params=true, unsafe_inbounds=true) y = map(t -> t[0], yt) return y, jac end @@ -299,7 +299,7 @@ function DI.value_and_jacobian!( 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_size=true) + GTPSA.jacobian!(jac, yt; include_params=true, unsafe_inbounds=true) y = map(t -> t[0], yt) return y, jac end @@ -482,7 +482,7 @@ function DI.hessian( hess = similar(x, GTPSA.numtype(yt), (length(x), length(x))) unsafe_fast = D == Nothing ? true : false GTPSA.hessian!( - hess, yt; include_params=true, unsafe_size=true, unsafe_fast=unsafe_fast, tmp_mono=prep.m + hess, yt; include_params=true, unsafe_inbounds=true, unsafe_fast=unsafe_fast, tmp_mono=prep.m ) return hess end @@ -500,7 +500,7 @@ function DI.hessian!( yt = fc(prep.xt) unsafe_fast = D == Nothing ? true : false GTPSA.hessian!( - hess, yt; include_params=true, unsafe_size=true, unsafe_fast=unsafe_fast, tmp_mono=prep.m + hess, yt; include_params=true, unsafe_inbounds=true, unsafe_fast=unsafe_fast, tmp_mono=prep.m ) return hess end @@ -517,11 +517,11 @@ function DI.value_gradient_and_hessian( fc = DI.with_contexts(f, contexts...) yt = fc(prep.xt) grad = similar(x, GTPSA.numtype(yt)) - GTPSA.gradient!(grad, yt; include_params=true, unsafe_size=true) + 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_size=true, unsafe_fast=unsafe_fast, tmp_mono=prep.m + hess, yt; include_params=true, unsafe_inbounds=true, unsafe_fast=unsafe_fast, tmp_mono=prep.m ) return yt[0], grad, hess end @@ -539,10 +539,10 @@ function DI.value_gradient_and_hessian!( 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_size=true) + GTPSA.gradient!(grad, yt; include_params=true, unsafe_inbounds=true) unsafe_fast = D == Nothing ? true : false GTPSA.hessian!( - hess, yt; include_params=true, unsafe_size=true, unsafe_fast=unsafe_fast, tmp_mono=prep.m + hess, yt; include_params=true, unsafe_inbounds=true, unsafe_fast=unsafe_fast, tmp_mono=prep.m ) return yt[0], grad, hess end @@ -580,7 +580,7 @@ function DI.hvp( dg .= 0 j = 1 for dxi in dx - for i in 1:size(prep.hess, 2) + for i in 1:size(prep.hess, 1) dg[i] += prep.hess[i, j] * dxi end j += 1 @@ -605,7 +605,7 @@ function DI.hvp!( dg .= 0 j = 1 for dxi in dx - for i in 1:size(prep.hess, 2) + for i in 1:size(prep.hess, 1) dg[i] += prep.hess[i, j] * dxi end j += 1 @@ -622,20 +622,21 @@ function DI.gradient_and_hvp( tx::NTuple, contexts::Vararg{DI.Context,C}, ) where {D,C} + foreach((t, xi) -> t[0] = xi, prep.hessprep.xt, x) # Set the scalar part fc = DI.with_contexts(f, contexts...) yt = fc(prep.hessprep.xt) grad = similar(x, GTPSA.numtype(yt)) - GTPSA.gradient!(grad, yt; include_params=true, unsafe_size=true) + GTPSA.gradient!(grad, yt; include_params=true, unsafe_inbounds=true) unsafe_fast = D == Nothing ? true : false GTPSA.hessian!( - prep.hess, yt; include_params=true, unsafe_size=true, unsafe_fast=unsafe_fast, tmp_mono=prep.hessprep.m + prep.hess, yt; include_params=true, unsafe_inbounds=true, unsafe_fast=unsafe_fast, tmp_mono=prep.hessprep.m ) 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, 2) + for i in 1:size(prep.hess, 1) dg[i] += prep.hess[i, j] * dxi end j += 1 @@ -655,24 +656,24 @@ function DI.gradient_and_hvp!( tx::NTuple, contexts::Vararg{DI.Context,C}, ) where {D,C} + foreach((t, xi) -> t[0] = xi, prep.hessprep.xt, x) # Set the scalar part fc = DI.with_contexts(f, contexts...) yt = fc(prep.hessprep.xt) - GTPSA.gradient!(grad, yt; include_params=true, unsafe_size=true) + GTPSA.gradient!(grad, yt; include_params=true, unsafe_inbounds=true) unsafe_fast = D == Nothing ? true : false GTPSA.hessian!( - prep.hess, yt; include_params=true, unsafe_size=true, unsafe_fast=unsafe_fast, tmp_mono=prep.hessprep.m + prep.hess, yt; include_params=true, unsafe_inbounds=true, unsafe_fast=unsafe_fast, tmp_mono=prep.hessprep.m ) 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, 2) + 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 \ No newline at end of file diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl index 8850f918b..cd227204b 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl @@ -161,7 +161,7 @@ function DI.jacobian( 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_size=true) + GTPSA.jacobian!(jac, prep.yt; include_params=true, unsafe_inbounds=true) map!(t -> t[0], y, prep.yt) return jac end @@ -178,7 +178,7 @@ function DI.jacobian!( 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_size=true) + GTPSA.jacobian!(jac, prep.yt; include_params=true, unsafe_inbounds=true) map!(t -> t[0], y, prep.yt) return jac end diff --git a/DifferentiationInterfaceTest/Project.toml b/DifferentiationInterfaceTest/Project.toml index f900204d1..9ef2043f8 100644 --- a/DifferentiationInterfaceTest/Project.toml +++ b/DifferentiationInterfaceTest/Project.toml @@ -79,6 +79,7 @@ FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" 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" From cfe4c953a034acd4f3cf9e5b721dd35dca97840b Mon Sep 17 00:00:00 2001 From: Matt Signorelli Date: Wed, 8 Jan 2025 10:00:10 -0500 Subject: [PATCH 37/44] fix typo --- DifferentiationInterface/test/Back/GTPSA/test.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/DifferentiationInterface/test/Back/GTPSA/test.jl b/DifferentiationInterface/test/Back/GTPSA/test.jl index 60db4ff92..8a0d78f51 100644 --- a/DifferentiationInterface/test/Back/GTPSA/test.jl +++ b/DifferentiationInterface/test/Back/GTPSA/test.jl @@ -16,11 +16,11 @@ end test_differentiation(AutoGTPSA(); type_stability=:full, logging=LOGGING); # Test with Descriptor: -d1 = Descriptor(20, 2) # 20 variables to 2nd order +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 = Descriptor(vos, 3) +d2 = GTPSA.Descriptor(vos, 3) test_differentiation(AutoGTPSA(d2); type_stability=:full, logging=LOGGING); From 731e9da5aca04772afcfc9c5a3c118cd00884265 Mon Sep 17 00:00:00 2001 From: Matt Signorelli Date: Wed, 8 Jan 2025 10:25:15 -0500 Subject: [PATCH 38/44] one more typo (sorry) --- .../ext/DifferentiationInterfaceGTPSAExt/twoarg.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl index cd227204b..30df24213 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl @@ -195,7 +195,7 @@ function DI.value_and_jacobian( return y, jac end -function DI.value_and_jacobian( +function DI.value_and_jacobian!( f!, y, jac, From 6701e9576ca5509f2a9c5f4ef6d82790117ae2fd Mon Sep 17 00:00:00 2001 From: Matt Signorelli Date: Wed, 8 Jan 2025 11:18:53 -0500 Subject: [PATCH 39/44] JuliaFormatter --- .../onearg.jl | 269 ++++++++---------- .../twoarg.jl | 58 ++-- .../test/Back/GTPSA/test.jl | 2 +- 3 files changed, 135 insertions(+), 194 deletions(-) diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl index 66a2630ca..b193ba60b 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl @@ -7,13 +7,9 @@ struct GTPSAOneArgPushforwardPrep{X} <: DI.PushforwardPrep end function DI.prepare_pushforward( - ::F, - backend::AutoGTPSA{D}, - x, - tx::NTuple, - ::Vararg{DI.Context,C}, + ::F, backend::AutoGTPSA{D}, x, tx::NTuple, ::Vararg{DI.Context,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 @@ -34,13 +30,13 @@ function DI.prepare_pushforward( end function DI.pushforward( - f, - prep::GTPSAOneArgPushforwardPrep, - ::AutoGTPSA, - x, + f, + prep::GTPSAOneArgPushforwardPrep, + ::AutoGTPSA, + x, tx::NTuple, contexts::Vararg{DI.Context,C}, -) where {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) @@ -56,11 +52,11 @@ function DI.pushforward( end function DI.pushforward!( - f, - ty::NTuple, - prep::GTPSAOneArgPushforwardPrep, - ::AutoGTPSA, - x, + f, + ty::NTuple, + prep::GTPSAOneArgPushforwardPrep, + ::AutoGTPSA, + x, tx::NTuple, contexts::Vararg{DI.Context,C}, ) where {C} @@ -75,14 +71,13 @@ function DI.pushforward!( end function DI.value_and_pushforward( - f, - prep::GTPSAOneArgPushforwardPrep, - backend::AutoGTPSA, - x, + f, + prep::GTPSAOneArgPushforwardPrep, + backend::AutoGTPSA, + x, tx::NTuple, contexts::Vararg{DI.Context,C}, ) where {C} - fc = DI.with_contexts(f, contexts...) ty = DI.pushforward(fc, prep, backend, x, tx) y = fc(x) # TO-DO: optimize @@ -90,11 +85,11 @@ function DI.value_and_pushforward( end function DI.value_and_pushforward!( - f, - ty::NTuple, - prep::GTPSAOneArgPushforwardPrep, - backend::AutoGTPSA, - x, + f, + ty::NTuple, + prep::GTPSAOneArgPushforwardPrep, + backend::AutoGTPSA, + x, tx::NTuple, contexts::Vararg{DI.Context,C}, ) where {C} @@ -113,12 +108,8 @@ end # Unlike JVP, this requires us to use all variables function DI.prepare_gradient( - f, - backend::AutoGTPSA{D}, - x, - contexts::Vararg{DI.Context,C}, + f, backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,C} ) where {D,C} - if D != Nothing d = backend.descriptor else @@ -144,11 +135,7 @@ function DI.prepare_gradient( end function DI.gradient( - f, - prep::GTPSAOneArgGradientPrep, - ::AutoGTPSA, - x, - contexts::Vararg{DI.Context,C}, + f, prep::GTPSAOneArgGradientPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Context,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...) @@ -163,12 +150,7 @@ function DI.gradient( end function DI.gradient!( - f, - grad, - prep::GTPSAOneArgGradientPrep, - ::AutoGTPSA, - x, - contexts::Vararg{DI.Context,C}, + f, grad, prep::GTPSAOneArgGradientPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Context,C} ) where {C} foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part fc = DI.with_contexts(f, contexts...) @@ -178,11 +160,7 @@ function DI.gradient!( end function DI.value_and_gradient( - f, - prep::GTPSAOneArgGradientPrep, - ::AutoGTPSA, - x, - contexts::Vararg{DI.Context,C}, + f, prep::GTPSAOneArgGradientPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Context,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...) @@ -197,12 +175,7 @@ function DI.value_and_gradient( end function DI.value_and_gradient!( - f, - grad, - prep::GTPSAOneArgGradientPrep, - ::AutoGTPSA, - x, - contexts::Vararg{DI.Context,C}, + f, grad, prep::GTPSAOneArgGradientPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Context,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...) @@ -219,10 +192,7 @@ end # To materialize the entire Jacobian we use all variables function DI.prepare_jacobian( - f, - backend::AutoGTPSA{D}, - x, - contexts::Vararg{DI.Context,C}, + f, backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,C} ) where {D,C} if D != Nothing d = backend.descriptor @@ -242,11 +212,7 @@ function DI.prepare_jacobian( end function DI.jacobian( - f, - prep::GTPSAOneArgJacobianPrep, - ::AutoGTPSA, - x, - contexts::Vararg{DI.Context,C}, + f, prep::GTPSAOneArgJacobianPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Context,C} ) where {C} foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part fc = DI.with_contexts(f, contexts...) @@ -257,12 +223,7 @@ function DI.jacobian( end function DI.jacobian!( - f, - jac, - prep::GTPSAOneArgJacobianPrep, - ::AutoGTPSA, - x, - contexts::Vararg{DI.Context,C}, + f, jac, prep::GTPSAOneArgJacobianPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Context,C} ) where {C} foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part fc = DI.with_contexts(f, contexts...) @@ -271,13 +232,8 @@ function DI.jacobian!( return jac end - function DI.value_and_jacobian( - f, - prep::GTPSAOneArgJacobianPrep, - ::AutoGTPSA, - x, - contexts::Vararg{DI.Context,C}, + f, prep::GTPSAOneArgJacobianPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Context,C} ) where {C} foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part fc = DI.with_contexts(f, contexts...) @@ -289,12 +245,7 @@ function DI.value_and_jacobian( end function DI.value_and_jacobian!( - f, - jac, - prep::GTPSAOneArgJacobianPrep, - ::AutoGTPSA, - x, - contexts::Vararg{DI.Context,C}, + f, jac, prep::GTPSAOneArgJacobianPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Context,C} ) where {C} foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part fc = DI.with_contexts(f, contexts...) @@ -304,7 +255,6 @@ function DI.value_and_jacobian!( return y, jac end - ## Second derivative # Contains single pre-allocated TPS struct GTPSAOneArgSecondDerivativePrep{X} <: DI.SecondDerivativePrep @@ -312,10 +262,7 @@ struct GTPSAOneArgSecondDerivativePrep{X} <: DI.SecondDerivativePrep end function DI.prepare_second_derivative( - f, - backend::AutoGTPSA{D}, - x, - contexts::Vararg{DI.Context,C}, + f, backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,C} ) where {D,C} if D != Nothing d = backend.descriptor @@ -328,9 +275,9 @@ function DI.prepare_second_derivative( end function DI.second_derivative( - f, - prep::GTPSAOneArgSecondDerivativePrep, - backend::AutoGTPSA{D}, + f, + prep::GTPSAOneArgSecondDerivativePrep, + backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,C}, ) where {D,C} @@ -340,7 +287,7 @@ function DI.second_derivative( if D == Nothing idx2 = 2 else - idx2 = GTPSA.numnn(backend.descriptor)+1 # index of first second derivative + idx2 = GTPSA.numnn(backend.descriptor) + 1 # index of first second derivative end if yt isa Number @@ -355,10 +302,10 @@ function DI.second_derivative( end function DI.second_derivative!( - f, + f, der2, - prep::GTPSAOneArgSecondDerivativePrep, - backend::AutoGTPSA{D}, + prep::GTPSAOneArgSecondDerivativePrep, + backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,C}, ) where {D,C} @@ -368,7 +315,7 @@ function DI.second_derivative!( if D == Nothing idx2 = 2 else - idx2 = GTPSA.numnn(backend.descriptor)+1 # index of first second derivative + idx2 = GTPSA.numnn(backend.descriptor) + 1 # index of first second derivative end for i in eachindex(yt) der2[i] = yt[i][idx2] * 2 @@ -377,9 +324,9 @@ function DI.second_derivative!( end function DI.value_derivative_and_second_derivative( - f, - prep::GTPSAOneArgSecondDerivativePrep, - backend::AutoGTPSA{D}, + f, + prep::GTPSAOneArgSecondDerivativePrep, + backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,C}, ) where {D,C} @@ -389,7 +336,7 @@ function DI.value_derivative_and_second_derivative( if D == Nothing idx2 = 2 else - idx2 = GTPSA.numnn(backend.descriptor)+1 # index of first second derivative + idx2 = GTPSA.numnn(backend.descriptor) + 1 # index of first second derivative end if yt isa Number return yt[0], yt[1], yt[idx2] * 2 @@ -406,11 +353,11 @@ function DI.value_derivative_and_second_derivative( end function DI.value_derivative_and_second_derivative!( - f, + f, der, der2, - prep::GTPSAOneArgSecondDerivativePrep, - backend::AutoGTPSA{D}, + prep::GTPSAOneArgSecondDerivativePrep, + backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,C}, ) where {D,C} @@ -421,7 +368,7 @@ function DI.value_derivative_and_second_derivative!( if D == Nothing idx2 = 2 else - idx2 = GTPSA.numnn(backend.descriptor)+1 # index of first second derivative + idx2 = GTPSA.numnn(backend.descriptor) + 1 # index of first second derivative end for i in eachindex(yt) der[i] = yt[i][1] @@ -439,10 +386,7 @@ struct GTPSAOneArgHessianPrep{X,M} <: DI.HessianPrep end function DI.prepare_hessian( - f, - backend::AutoGTPSA{D}, - x, - contexts::Vararg{DI.Context,C}, + f, backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,C} ) where {D,C} if D != Nothing d = backend.descriptor @@ -470,11 +414,7 @@ function DI.prepare_hessian( end function DI.hessian( - f, - prep::GTPSAOneArgHessianPrep, - ::AutoGTPSA{D}, - x, - contexts::Vararg{DI.Context,C}, + f, prep::GTPSAOneArgHessianPrep, ::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,C} ) where {D,C} foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part fc = DI.with_contexts(f, contexts...) @@ -482,36 +422,36 @@ function DI.hessian( 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 + 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.Context,C}, + f, hess, prep::GTPSAOneArgHessianPrep, ::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,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 + 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.Context,C}, + f, prep::GTPSAOneArgHessianPrep, ::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,C} ) where {D,C} foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part fc = DI.with_contexts(f, contexts...) @@ -521,18 +461,22 @@ function DI.value_gradient_and_hessian( 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 + 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, + f, grad, hess, - prep::GTPSAOneArgHessianPrep, - ::AutoGTPSA{D}, + prep::GTPSAOneArgHessianPrep, + ::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,C}, ) where {D,C} @@ -542,7 +486,12 @@ function DI.value_gradient_and_hessian!( 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 + hess, + yt; + include_params=true, + unsafe_inbounds=true, + unsafe_fast=unsafe_fast, + tmp_mono=prep.m, ) return yt[0], grad, hess end @@ -553,11 +502,7 @@ struct GTPSAOneArgHVPPrep{E,H} <: DI.HVPPrep end function DI.prepare_hvp( - f, - backend::AutoGTPSA, - x, - tx::NTuple, - contexts::Vararg{DI.Context,C}, + f, backend::AutoGTPSA, x, tx::NTuple, contexts::Vararg{DI.Context,C} ) where {C} hessprep = DI.prepare_hessian(f, backend, x) fc = DI.with_contexts(f, contexts...) @@ -565,11 +510,11 @@ function DI.prepare_hvp( return GTPSAOneArgHVPPrep(hessprep, hess) end -function DI.hvp( - f, - prep::GTPSAOneArgHVPPrep, - backend::AutoGTPSA, - x, +function DI.hvp( + f, + prep::GTPSAOneArgHVPPrep, + backend::AutoGTPSA, + x, tx::NTuple, contexts::Vararg{DI.Context,C}, ) where {C} @@ -591,11 +536,11 @@ function DI.hvp( end function DI.hvp!( - f, + f, tg::NTuple, - prep::GTPSAOneArgHVPPrep, - backend::AutoGTPSA, - x, + prep::GTPSAOneArgHVPPrep, + backend::AutoGTPSA, + x, tx::NTuple, contexts::Vararg{DI.Context,C}, ) where {C} @@ -614,11 +559,11 @@ function DI.hvp!( return tg end -function DI.gradient_and_hvp( - f, - prep::GTPSAOneArgHVPPrep, - ::AutoGTPSA{D}, - x, +function DI.gradient_and_hvp( + f, + prep::GTPSAOneArgHVPPrep, + ::AutoGTPSA{D}, + x, tx::NTuple, contexts::Vararg{DI.Context,C}, ) where {D,C} @@ -629,7 +574,12 @@ function DI.gradient_and_hvp( GTPSA.gradient!(grad, yt; include_params=true, unsafe_inbounds=true) unsafe_fast = D == Nothing ? true : false GTPSA.hessian!( - prep.hess, yt; include_params=true, unsafe_inbounds=true, unsafe_fast=unsafe_fast, tmp_mono=prep.hessprep.m + prep.hess, + yt; + include_params=true, + unsafe_inbounds=true, + unsafe_fast=unsafe_fast, + tmp_mono=prep.hessprep.m, ) tg = map(tx) do dx dg = similar(x, eltype(prep.hess)) @@ -646,13 +596,13 @@ function DI.gradient_and_hvp( return grad, tg end -function DI.gradient_and_hvp!( - f, +function DI.gradient_and_hvp!( + f, grad, tg, - prep::GTPSAOneArgHVPPrep, - ::AutoGTPSA{D}, - x, + prep::GTPSAOneArgHVPPrep, + ::AutoGTPSA{D}, + x, tx::NTuple, contexts::Vararg{DI.Context,C}, ) where {D,C} @@ -662,7 +612,12 @@ function DI.gradient_and_hvp!( GTPSA.gradient!(grad, yt; include_params=true, unsafe_inbounds=true) unsafe_fast = D == Nothing ? true : false GTPSA.hessian!( - prep.hess, yt; include_params=true, unsafe_inbounds=true, unsafe_fast=unsafe_fast, tmp_mono=prep.hessprep.m + prep.hess, + yt; + include_params=true, + unsafe_inbounds=true, + unsafe_fast=unsafe_fast, + tmp_mono=prep.hessprep.m, ) for b in eachindex(tg, tx) dg, dx = tg[b], tx[b] @@ -676,4 +631,4 @@ function DI.gradient_and_hvp!( end end return grad, tg -end \ No newline at end of file +end diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl index 30df24213..7d0cde550 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl @@ -10,12 +10,7 @@ struct GTPSATwoArgPushforwardPrep{X,Y} <: DI.PushforwardPrep end function DI.prepare_pushforward( - ::F, - y, - backend::AutoGTPSA{D}, - x, - tx::NTuple, - ::Vararg{DI.Context,C}, + ::F, y, backend::AutoGTPSA{D}, x, tx::NTuple, ::Vararg{DI.Context,C} ) where {F,D,C} # For pushforward/JVP, we only actually need 1 single variable (in the GTPSA sense) @@ -43,11 +38,11 @@ function DI.prepare_pushforward( end function DI.pushforward( - f!, - y, - prep::GTPSATwoArgPushforwardPrep, - ::AutoGTPSA, - x, + f!, + y, + prep::GTPSATwoArgPushforwardPrep, + ::AutoGTPSA, + x, tx::NTuple, contexts::Vararg{DI.Context,C}, ) where {C} @@ -84,11 +79,11 @@ function DI.pushforward!( end function DI.value_and_pushforward( - f!, - y, - prep::GTPSATwoArgPushforwardPrep, - backend::AutoGTPSA, - x, + f!, + y, + prep::GTPSATwoArgPushforwardPrep, + backend::AutoGTPSA, + x, tx::NTuple, contexts::Vararg{DI.Context,C}, ) where {C} @@ -119,11 +114,7 @@ struct GTPSATwoArgJacobianPrep{X,Y} <: DI.JacobianPrep end function DI.prepare_jacobian( - f, - y, - backend::AutoGTPSA{D}, - x, - contexts::Vararg{DI.Context,C}, + f, y, backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,C} ) where {D,C} if D != Nothing d = backend.descriptor @@ -150,12 +141,7 @@ function DI.prepare_jacobian( end function DI.jacobian( - f!, - y, - prep::GTPSATwoArgJacobianPrep, - ::AutoGTPSA, - x, - contexts::Vararg{DI.Context,C}, + f!, y, prep::GTPSATwoArgJacobianPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Context,C} ) where {C} foreach((t, xi) -> t[0] = xi, prep.xt, x) # Set the scalar part fc! = DI.with_contexts(f!, contexts...) @@ -167,11 +153,11 @@ function DI.jacobian( end function DI.jacobian!( - f!, + f!, y, jac, - prep::GTPSATwoArgJacobianPrep, - ::AutoGTPSA, + prep::GTPSATwoArgJacobianPrep, + ::AutoGTPSA, x, contexts::Vararg{DI.Context,C}, ) where {C} @@ -184,10 +170,10 @@ function DI.jacobian!( end function DI.value_and_jacobian( - f!, + f!, y, - prep::GTPSATwoArgJacobianPrep, - backend::AutoGTPSA, + prep::GTPSATwoArgJacobianPrep, + backend::AutoGTPSA, x, contexts::Vararg{DI.Context,C}, ) where {C} @@ -196,11 +182,11 @@ function DI.value_and_jacobian( end function DI.value_and_jacobian!( - f!, + f!, y, jac, - prep::GTPSATwoArgJacobianPrep, - backend::AutoGTPSA, + prep::GTPSATwoArgJacobianPrep, + backend::AutoGTPSA, x, contexts::Vararg{DI.Context,C}, ) where {C} diff --git a/DifferentiationInterface/test/Back/GTPSA/test.jl b/DifferentiationInterface/test/Back/GTPSA/test.jl index 8a0d78f51..9d66aa562 100644 --- a/DifferentiationInterface/test/Back/GTPSA/test.jl +++ b/DifferentiationInterface/test/Back/GTPSA/test.jl @@ -20,7 +20,7 @@ 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 = 2 * ones(Int, 20) vos[1] = 3 d2 = GTPSA.Descriptor(vos, 3) test_differentiation(AutoGTPSA(d2); type_stability=:full, logging=LOGGING); From 52ecdf64b96a0e2269a40f1e99c7f96fe9228af0 Mon Sep 17 00:00:00 2001 From: mattsignorelli <113714234+mattsignorelli@users.noreply.github.com> Date: Thu, 23 Jan 2025 13:56:30 -0500 Subject: [PATCH 40/44] Update DifferentiationInterface/docs/Project.toml Co-authored-by: Guillaume Dalle <22795598+gdalle@users.noreply.github.com> --- DifferentiationInterface/docs/Project.toml | 1 - 1 file changed, 1 deletion(-) diff --git a/DifferentiationInterface/docs/Project.toml b/DifferentiationInterface/docs/Project.toml index a7a540626..18ba2902e 100644 --- a/DifferentiationInterface/docs/Project.toml +++ b/DifferentiationInterface/docs/Project.toml @@ -7,7 +7,6 @@ DocumenterInterLinks = "d12716ef-a0f6-4df4-a9f1-a5a34e75c656" DocumenterMermaid = "a078cd44-4d9c-4618-b545-3ab9d77f9177" FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" -GTPSA = "b27dd330-f138-47c5-815b-40db9dd9b6e8" Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a" PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d" SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5" From 840e136695bc63509e76d9aef510203c4e66ccc4 Mon Sep 17 00:00:00 2001 From: mattsignorelli <113714234+mattsignorelli@users.noreply.github.com> Date: Thu, 23 Jan 2025 14:03:29 -0500 Subject: [PATCH 41/44] Better clarification of "variables" in GTPSA --- DifferentiationInterface/docs/src/explanation/backends.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/DifferentiationInterface/docs/src/explanation/backends.md b/DifferentiationInterface/docs/src/explanation/backends.md index 2115d4d9b..f206f90d9 100644 --- a/DifferentiationInterface/docs/src/explanation/backends.md +++ b/DifferentiationInterface/docs/src/explanation/backends.md @@ -143,12 +143,12 @@ For higher level operators, preparation creates a [config object](https://juliad ### GTPSA -For all operators, preparation preallocates the input [`TPS`s](https://bmad-sim.github.io/GTPSA.jl/stable/man/d_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/k_fastgtpsa/) macros are recommended. +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/c_descriptor/) is not provided to `AutoGTPSA`, then a `Descriptor` will be generated in preparation based on the context. +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 variables + parameters is not less than the number of inputs of the provided function. Undefined behavior and crashes may occur if this is not the case. + 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 From c970426f4b786e849e91a446903a76e844ea4da3 Mon Sep 17 00:00:00 2001 From: Matt Signorelli Date: Thu, 23 Jan 2025 15:02:20 -0500 Subject: [PATCH 42/44] address comments --- DifferentiationInterface/Project.toml | 2 +- DifferentiationInterface/docs/src/api.md | 2 +- .../DifferentiationInterfaceGTPSAExt.jl | 3 +- .../onearg.jl | 35 +++---------------- .../twoarg.jl | 2 +- DifferentiationInterfaceTest/Project.toml | 2 -- 6 files changed, 9 insertions(+), 37 deletions(-) diff --git a/DifferentiationInterface/Project.toml b/DifferentiationInterface/Project.toml index 9836ae8a8..e850b45a8 100644 --- a/DifferentiationInterface/Project.toml +++ b/DifferentiationInterface/Project.toml @@ -59,7 +59,7 @@ FastDifferentiation = "0.4.1" FiniteDiff = "2.23.1" FiniteDifferences = "0.12.31" ForwardDiff = "0.10.36" -GTPSA = "1.3.1" +GTPSA = "1.4.0" JuliaFormatter = "1" LinearAlgebra = "<0.0.1,1" Mooncake = "0.4.52" diff --git a/DifferentiationInterface/docs/src/api.md b/DifferentiationInterface/docs/src/api.md index fee88f719..7a76d54c5 100644 --- a/DifferentiationInterface/docs/src/api.md +++ b/DifferentiationInterface/docs/src/api.md @@ -139,4 +139,4 @@ The following is not part of the public API. Modules = [DifferentiationInterface] Public = false Filter = t -> !(Symbol(t) in [:outer, :inner]) -``` \ No newline at end of file +``` diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl index 4e5101772..76b538109 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl @@ -2,8 +2,7 @@ module DifferentiationInterfaceGTPSAExt import DifferentiationInterface as DI using ADTypes: AutoGTPSA -using GTPSA -using LinearAlgebra +using GTPSA: GTPSA, TPS, Descriptor DI.check_available(::AutoGTPSA) = true diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl index b193ba60b..d6d5cc550 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl @@ -519,7 +519,6 @@ function DI.hvp( contexts::Vararg{DI.Context,C}, ) where {C} DI.hessian!(f, prep.hess, prep.hessprep, backend, x, contexts...) - #return prep.hess tg = map(tx) do dx dg = similar(x, eltype(prep.hess)) dg .= 0 @@ -562,25 +561,13 @@ end function DI.gradient_and_hvp( f, prep::GTPSAOneArgHVPPrep, - ::AutoGTPSA{D}, + backend::AutoGTPSA{D}, x, tx::NTuple, contexts::Vararg{DI.Context,C}, ) where {D,C} - foreach((t, xi) -> t[0] = xi, prep.hessprep.xt, x) # Set the scalar part - fc = DI.with_contexts(f, contexts...) - yt = fc(prep.hessprep.xt) - grad = similar(x, GTPSA.numtype(yt)) - GTPSA.gradient!(grad, yt; include_params=true, unsafe_inbounds=true) - unsafe_fast = D == Nothing ? true : false - GTPSA.hessian!( - prep.hess, - yt; - include_params=true, - unsafe_inbounds=true, - unsafe_fast=unsafe_fast, - tmp_mono=prep.hessprep.m, - ) + 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 @@ -601,24 +588,12 @@ function DI.gradient_and_hvp!( grad, tg, prep::GTPSAOneArgHVPPrep, - ::AutoGTPSA{D}, + backend::AutoGTPSA{D}, x, tx::NTuple, contexts::Vararg{DI.Context,C}, ) where {D,C} - foreach((t, xi) -> t[0] = xi, prep.hessprep.xt, x) # Set the scalar part - fc = DI.with_contexts(f, contexts...) - yt = fc(prep.hessprep.xt) - GTPSA.gradient!(grad, yt; include_params=true, unsafe_inbounds=true) - unsafe_fast = D == Nothing ? true : false - GTPSA.hessian!( - prep.hess, - yt; - include_params=true, - unsafe_inbounds=true, - unsafe_fast=unsafe_fast, - tmp_mono=prep.hessprep.m, - ) + 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 diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl index 7d0cde550..567f9b32d 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl @@ -177,7 +177,7 @@ function DI.value_and_jacobian( x, contexts::Vararg{DI.Context,C}, ) where {C} - jac = DI.jacobian(f!, y, prep, backend, x, contexts...) + jac = DI.jacobian(f!, y, prep, backend, x, contexts...) # y set on line 151 return y, jac end diff --git a/DifferentiationInterfaceTest/Project.toml b/DifferentiationInterfaceTest/Project.toml index 9ef2043f8..cd632972d 100644 --- a/DifferentiationInterfaceTest/Project.toml +++ b/DifferentiationInterfaceTest/Project.toml @@ -50,7 +50,6 @@ FiniteDifferences = "0.12" Flux = "0.16" ForwardDiff = "0.10.36" Functors = "0.4, 0.5" -GTPSA = "1.3.1" JET = "0.4 - 0.8, 0.9" JLArrays = "0.1, 0.2" JuliaFormatter = "1" @@ -79,7 +78,6 @@ FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" 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" From c1931c19ddc51e305c14daeb5d5ff9954821c302 Mon Sep 17 00:00:00 2001 From: Matt Signorelli Date: Thu, 23 Jan 2025 15:07:36 -0500 Subject: [PATCH 43/44] no implicit imports test --- DifferentiationInterface/test/Back/GTPSA/test.jl | 3 +++ 1 file changed, 3 insertions(+) diff --git a/DifferentiationInterface/test/Back/GTPSA/test.jl b/DifferentiationInterface/test/Back/GTPSA/test.jl index 9d66aa562..3dd9599d3 100644 --- a/DifferentiationInterface/test/Back/GTPSA/test.jl +++ b/DifferentiationInterface/test/Back/GTPSA/test.jl @@ -5,6 +5,9 @@ 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()] From 34734223cdc7091ef0aaac5a2eb63c64646b2735 Mon Sep 17 00:00:00 2001 From: Matt Signorelli Date: Thu, 23 Jan 2025 15:54:38 -0500 Subject: [PATCH 44/44] fixes and formatting --- .../DifferentiationInterfaceGTPSAExt.jl | 2 +- .../onearg.jl | 117 ++++++++---------- .../twoarg.jl | 20 +-- 3 files changed, 65 insertions(+), 74 deletions(-) diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl index 76b538109..b2d24dac3 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl @@ -2,7 +2,7 @@ module DifferentiationInterfaceGTPSAExt import DifferentiationInterface as DI using ADTypes: AutoGTPSA -using GTPSA: GTPSA, TPS, Descriptor +using GTPSA: GTPSA, TPS, Descriptor DI.check_available(::AutoGTPSA) = true diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl index d6d5cc550..ca7937d63 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/onearg.jl @@ -7,7 +7,7 @@ struct GTPSAOneArgPushforwardPrep{X} <: DI.PushforwardPrep end function DI.prepare_pushforward( - ::F, backend::AutoGTPSA{D}, x, tx::NTuple, ::Vararg{DI.Context,C} + ::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) @@ -35,7 +35,7 @@ function DI.pushforward( ::AutoGTPSA, x, tx::NTuple, - contexts::Vararg{DI.Context,C}, + contexts::Vararg{DI.Constant,C}, ) where {C} fc = DI.with_contexts(f, contexts...) ty = map(tx) do dx @@ -58,7 +58,7 @@ function DI.pushforward!( ::AutoGTPSA, x, tx::NTuple, - contexts::Vararg{DI.Context,C}, + contexts::Vararg{DI.Constant,C}, ) where {C} fc = DI.with_contexts(f, contexts...) for b in eachindex(tx, ty) @@ -76,7 +76,7 @@ function DI.value_and_pushforward( backend::AutoGTPSA, x, tx::NTuple, - contexts::Vararg{DI.Context,C}, + contexts::Vararg{DI.Constant,C}, ) where {C} fc = DI.with_contexts(f, contexts...) ty = DI.pushforward(fc, prep, backend, x, tx) @@ -91,7 +91,7 @@ function DI.value_and_pushforward!( backend::AutoGTPSA, x, tx::NTuple, - contexts::Vararg{DI.Context,C}, + contexts::Vararg{DI.Constant,C}, ) where {C} fc = DI.with_contexts(f, contexts...) DI.pushforward!(fc, ty, prep, backend, x, tx) @@ -100,57 +100,43 @@ function DI.value_and_pushforward!( end ## Gradient -# Contains either a single pre-allocated initial TPS -# or a vector of pre-allocated TPSs. +# 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.Context,C} + 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 gradient - if x isa Number - xt = TPS{promote_type(typeof(x), Float64)}(; use=d) - xt[1] = 1 - return GTPSAOneArgGradientPrep(xt) - else - 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) + 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.Context,C} + 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) - if prep.xt isa Number - return yt[1] - else - grad = similar(x, GTPSA.numtype(yt)) - GTPSA.gradient!(grad, yt; include_params=true, unsafe_inbounds=true) - return grad - end + 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.Context,C} + 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...) @@ -160,22 +146,18 @@ function DI.gradient!( end function DI.value_and_gradient( - f, prep::GTPSAOneArgGradientPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Context,C} + 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) - if prep.xt isa Number - return yt[0], yt[1] - else - grad = similar(x, GTPSA.numtype(yt)) - GTPSA.gradient!(grad, yt; include_params=true, unsafe_inbounds=true) - return yt[0], grad - end + 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.Context,C} + 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...) @@ -192,7 +174,7 @@ end # To materialize the entire Jacobian we use all variables function DI.prepare_jacobian( - f, backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,C} + f, backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Constant,C} ) where {D,C} if D != Nothing d = backend.descriptor @@ -212,7 +194,7 @@ function DI.prepare_jacobian( end function DI.jacobian( - f, prep::GTPSAOneArgJacobianPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Context,C} + 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...) @@ -223,7 +205,7 @@ function DI.jacobian( end function DI.jacobian!( - f, jac, prep::GTPSAOneArgJacobianPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Context,C} + 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...) @@ -233,7 +215,7 @@ function DI.jacobian!( end function DI.value_and_jacobian( - f, prep::GTPSAOneArgJacobianPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Context,C} + 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...) @@ -245,7 +227,7 @@ function DI.value_and_jacobian( end function DI.value_and_jacobian!( - f, jac, prep::GTPSAOneArgJacobianPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Context,C} + 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...) @@ -262,7 +244,7 @@ struct GTPSAOneArgSecondDerivativePrep{X} <: DI.SecondDerivativePrep end function DI.prepare_second_derivative( - f, backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,C} + f, backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Constant,C} ) where {D,C} if D != Nothing d = backend.descriptor @@ -279,7 +261,7 @@ function DI.second_derivative( prep::GTPSAOneArgSecondDerivativePrep, backend::AutoGTPSA{D}, x, - contexts::Vararg{DI.Context,C}, + contexts::Vararg{DI.Constant,C}, ) where {D,C} prep.xt[0] = x fc = DI.with_contexts(f, contexts...) @@ -307,7 +289,7 @@ function DI.second_derivative!( prep::GTPSAOneArgSecondDerivativePrep, backend::AutoGTPSA{D}, x, - contexts::Vararg{DI.Context,C}, + contexts::Vararg{DI.Constant,C}, ) where {D,C} prep.xt[0] = x fc = DI.with_contexts(f, contexts...) @@ -328,7 +310,7 @@ function DI.value_derivative_and_second_derivative( prep::GTPSAOneArgSecondDerivativePrep, backend::AutoGTPSA{D}, x, - contexts::Vararg{DI.Context,C}, + contexts::Vararg{DI.Constant,C}, ) where {D,C} prep.xt[0] = x fc = DI.with_contexts(f, contexts...) @@ -359,7 +341,7 @@ function DI.value_derivative_and_second_derivative!( prep::GTPSAOneArgSecondDerivativePrep, backend::AutoGTPSA{D}, x, - contexts::Vararg{DI.Context,C}, + contexts::Vararg{DI.Constant,C}, ) where {D,C} prep.xt[0] = x fc = DI.with_contexts(f, contexts...) @@ -386,7 +368,7 @@ struct GTPSAOneArgHessianPrep{X,M} <: DI.HessianPrep end function DI.prepare_hessian( - f, backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,C} + f, backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Constant,C} ) where {D,C} if D != Nothing d = backend.descriptor @@ -414,7 +396,7 @@ function DI.prepare_hessian( end function DI.hessian( - f, prep::GTPSAOneArgHessianPrep, ::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,C} + 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...) @@ -433,7 +415,12 @@ function DI.hessian( end function DI.hessian!( - f, hess, prep::GTPSAOneArgHessianPrep, ::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,C} + 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...) @@ -451,7 +438,7 @@ function DI.hessian!( end function DI.value_gradient_and_hessian( - f, prep::GTPSAOneArgHessianPrep, ::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,C} + 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...) @@ -478,7 +465,7 @@ function DI.value_gradient_and_hessian!( prep::GTPSAOneArgHessianPrep, ::AutoGTPSA{D}, x, - contexts::Vararg{DI.Context,C}, + 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...) @@ -502,7 +489,7 @@ struct GTPSAOneArgHVPPrep{E,H} <: DI.HVPPrep end function DI.prepare_hvp( - f, backend::AutoGTPSA, x, tx::NTuple, contexts::Vararg{DI.Context,C} + 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...) @@ -516,7 +503,7 @@ function DI.hvp( backend::AutoGTPSA, x, tx::NTuple, - contexts::Vararg{DI.Context,C}, + contexts::Vararg{DI.Constant,C}, ) where {C} DI.hessian!(f, prep.hess, prep.hessprep, backend, x, contexts...) tg = map(tx) do dx @@ -541,7 +528,7 @@ function DI.hvp!( backend::AutoGTPSA, x, tx::NTuple, - contexts::Vararg{DI.Context,C}, + contexts::Vararg{DI.Constant,C}, ) where {C} DI.hessian!(f, prep.hess, prep.hessprep, backend, x, contexts...) for b in eachindex(tg) @@ -564,10 +551,12 @@ function DI.gradient_and_hvp( backend::AutoGTPSA{D}, x, tx::NTuple, - contexts::Vararg{DI.Context,C}, + 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...) + 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 @@ -591,9 +580,11 @@ function DI.gradient_and_hvp!( backend::AutoGTPSA{D}, x, tx::NTuple, - contexts::Vararg{DI.Context,C}, + contexts::Vararg{DI.Constant,C}, ) where {D,C} - DI.value_gradient_and_hessian!(f, grad, prep.hess, prep.hessprep, backend, x, contexts...) + 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 diff --git a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl index 567f9b32d..3e2e88de5 100644 --- a/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl +++ b/DifferentiationInterface/ext/DifferentiationInterfaceGTPSAExt/twoarg.jl @@ -10,7 +10,7 @@ struct GTPSATwoArgPushforwardPrep{X,Y} <: DI.PushforwardPrep end function DI.prepare_pushforward( - ::F, y, backend::AutoGTPSA{D}, x, tx::NTuple, ::Vararg{DI.Context,C} + ::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) @@ -44,7 +44,7 @@ function DI.pushforward( ::AutoGTPSA, x, tx::NTuple, - contexts::Vararg{DI.Context,C}, + contexts::Vararg{DI.Constant,C}, ) where {C} fc! = DI.with_contexts(f!, contexts...) ty = map(tx) do dx @@ -65,7 +65,7 @@ function DI.pushforward!( ::AutoGTPSA, x, tx::NTuple, - contexts::Vararg{DI.Context,C}, + contexts::Vararg{DI.Constant,C}, ) where {C} fc! = DI.with_contexts(f!, contexts...) for b in eachindex(tx, ty) @@ -85,7 +85,7 @@ function DI.value_and_pushforward( backend::AutoGTPSA, x, tx::NTuple, - contexts::Vararg{DI.Context,C}, + contexts::Vararg{DI.Constant,C}, ) where {C} ty = DI.pushforward(f!, y, prep, backend, x, tx, contexts...) return y, ty @@ -99,7 +99,7 @@ function DI.value_and_pushforward!( backend::AutoGTPSA, x, tx::NTuple, - contexts::Vararg{DI.Context,C}, + contexts::Vararg{DI.Constant,C}, ) where {C} DI.pushforward!(f!, y, ty, prep, backend, x, tx, contexts...) return y, ty @@ -114,7 +114,7 @@ struct GTPSATwoArgJacobianPrep{X,Y} <: DI.JacobianPrep end function DI.prepare_jacobian( - f, y, backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Context,C} + f, y, backend::AutoGTPSA{D}, x, contexts::Vararg{DI.Constant,C} ) where {D,C} if D != Nothing d = backend.descriptor @@ -141,7 +141,7 @@ function DI.prepare_jacobian( end function DI.jacobian( - f!, y, prep::GTPSATwoArgJacobianPrep, ::AutoGTPSA, x, contexts::Vararg{DI.Context,C} + 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...) @@ -159,7 +159,7 @@ function DI.jacobian!( prep::GTPSATwoArgJacobianPrep, ::AutoGTPSA, x, - contexts::Vararg{DI.Context,C}, + 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...) @@ -175,7 +175,7 @@ function DI.value_and_jacobian( prep::GTPSATwoArgJacobianPrep, backend::AutoGTPSA, x, - contexts::Vararg{DI.Context,C}, + contexts::Vararg{DI.Constant,C}, ) where {C} jac = DI.jacobian(f!, y, prep, backend, x, contexts...) # y set on line 151 return y, jac @@ -188,7 +188,7 @@ function DI.value_and_jacobian!( prep::GTPSATwoArgJacobianPrep, backend::AutoGTPSA, x, - contexts::Vararg{DI.Context,C}, + contexts::Vararg{DI.Constant,C}, ) where {C} DI.jacobian!(f!, y, jac, prep, backend, x, contexts...) return y, jac