Add the GTPSA.jl backend (#329)

* add extension directory * derivative,second derivative, gradient * jacobian, hessian * add GTPSA to Project.toml, move AutoGTPSA to ADTypes * Fix tests * Lighter workflow * pushforward * trigger workflow * fix tests? * fix typo * hvp done and two args done * fix typo * Update Test.yml * no code coverage * add GTPSA test * almost all tests passing, bugs in pushforward * all tests passing * latest GTPSA version specified, docs updated * no coverage test * Fix formatting and docs * Exclude GTPSA from LTS testing * Reactivate coverage * Reactivate all tests * Test group * simplified code and removed derivative * Adapt to new Tangents * Tangents import * Fix da fixes * Remove temporary backend * Fixes * Docs * GTPSAExt devel in progress * one arg done * all tests except hvp passing, should have full code cov * bump gtpsa version * fix hvp functions and docs update * fix typo * one more typo (sorry) * JuliaFormatter * Update DifferentiationInterface/docs/Project.toml Co-authored-by: Guillaume Dalle <[email protected]> * Better clarification of "variables" in GTPSA * address comments * no implicit imports test * fixes and formatting --------- Co-authored-by: Guillaume Dalle <[email protected]>
JuliaDiff · Jan 27, 2025 · 8dcd924 · 8dcd924
1 parent 59dc865
commit 8dcd924
Show file tree

Hide file tree

Showing 9 changed files with 856 additions and 0 deletions.
diff --git a/.github/workflows/Test.yml b/.github/workflows/Test.yml
@@ -42,6 +42,7 @@ jobs:
           - Back/FiniteDiff
           - Back/FiniteDifferences
           - Back/ForwardDiff
+          - Back/GTPSA
           - Back/Mooncake
           - Back/PolyesterForwardDiff
           - Back/ReverseDiff

diff --git a/DifferentiationInterface/Project.toml b/DifferentiationInterface/Project.toml
@@ -17,6 +17,7 @@ FastDifferentiation = "eb9bf01b-bf85-4b60-bf87-ee5de06c00be"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+GTPSA = "b27dd330-f138-47c5-815b-40db9dd9b6e8"
 Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
@@ -35,6 +36,7 @@ DifferentiationInterfaceFastDifferentiationExt = "FastDifferentiation"
 DifferentiationInterfaceFiniteDiffExt = "FiniteDiff"
 DifferentiationInterfaceFiniteDifferencesExt = "FiniteDifferences"
 DifferentiationInterfaceForwardDiffExt = ["ForwardDiff", "DiffResults"]
+DifferentiationInterfaceGTPSAExt = "GTPSA"
 DifferentiationInterfaceMooncakeExt = "Mooncake"
 DifferentiationInterfacePolyesterForwardDiffExt = "PolyesterForwardDiff"
 DifferentiationInterfaceReverseDiffExt = ["ReverseDiff", "DiffResults"]
@@ -57,6 +59,7 @@ FastDifferentiation = "0.4.1"
 FiniteDiff = "2.23.1"
 FiniteDifferences = "0.12.31"
 ForwardDiff = "0.10.36"
+GTPSA = "1.4.0"
 JuliaFormatter = "1"
 LinearAlgebra = "<0.0.1,1"
 Mooncake = "0.4.52"
@@ -84,6 +87,7 @@ FastDifferentiation = "eb9bf01b-bf85-4b60-bf87-ee5de06c00be"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+GTPSA = "b27dd330-f138-47c5-815b-40db9dd9b6e8"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb"
 JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"

diff --git a/DifferentiationInterface/README.md b/DifferentiationInterface/README.md
@@ -37,6 +37,7 @@ We support the following backends defined by [ADTypes.jl](https://github.com/Sci
 - [FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl)
 - [FiniteDifferences.jl](https://github.com/JuliaDiff/FiniteDifferences.jl)
 - [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl)
+- [GTPSA.jl](https://github.com/bmad-sim/GTPSA.jl)
 - [Mooncake.jl](https://github.com/compintell/Mooncake.jl)
 - [PolyesterForwardDiff.jl](https://github.com/JuliaDiff/PolyesterForwardDiff.jl)
 - [ReverseDiff.jl](https://github.com/JuliaDiff/ReverseDiff.jl)

diff --git a/DifferentiationInterface/docs/src/explanation/backends.md b/DifferentiationInterface/docs/src/explanation/backends.md
@@ -11,6 +11,7 @@ We support the following dense backend choices from [ADTypes.jl](https://github.
 - [`AutoFiniteDiff`](@extref ADTypes.AutoFiniteDiff)
 - [`AutoFiniteDifferences`](@extref ADTypes.AutoFiniteDifferences)
 - [`AutoForwardDiff`](@extref ADTypes.AutoForwardDiff)
+- [`AutoGTPSA`](@extref ADTypes.AutoGTPSA)
 - [`AutoMooncake`](@extref ADTypes.AutoMooncake)
 - [`AutoPolyesterForwardDiff`](@extref ADTypes.AutoPolyesterForwardDiff)
 - [`AutoReverseDiff`](@extref ADTypes.AutoReverseDiff)
@@ -45,6 +46,7 @@ In practice, many AD backends have custom implementations for high-level operato
     | `AutoFiniteDiff`           | 🔀    | ❌    | ✅     | ✅      | ✅     | ✅      | ❌     | ❌      |
     | `AutoFiniteDifferences`    | 🔀    | ❌    | ❌     | ✅      | ✅     | ❌      | ❌     | ❌      |
     | `AutoForwardDiff`          | ✅    | ❌    | ✅     | ✅      | ✅     | ✅      | ✅     | ✅      |
+    | `AutoGTPSA`                | ✅    | ❌    | ❌     | ✅      | ✅     | ✅      | ✅     | ✅      |
     | `AutoMooncake`             | ❌    | ✅    | ❌     | ❌      | ❌     | ❌      | ❌     | ❌      |
     | `AutoPolyesterForwardDiff` | 🔀    | ❌    | 🔀     | ✅      | ✅     | 🔀      | 🔀     | 🔀      |
     | `AutoReverseDiff`          | ❌    | 🔀    | ❌     | ✅      | ✅     | ✅      | ❌     | ❌      |
@@ -64,6 +66,7 @@ Moreover, each context type is supported by a specific subset of backends:
 | `AutoFiniteDiff`           | ✅                  |
 | `AutoFiniteDifferences`    | ✅                  |
 | `AutoForwardDiff`          | ✅                  |
+| `AutoGTPSA`                | ✅                  |
 | `AutoMooncake`             | ✅                  |
 | `AutoPolyesterForwardDiff` | ✅                  |
 | `AutoReverseDiff`          | ✅                  |
@@ -138,6 +141,15 @@ Nothing specific to mention.
 We implement [`pushforward`](@ref) directly using [`Dual` numbers](https://juliadiff.org/ForwardDiff.jl/stable/dev/how_it_works/), and preparation allocates the necessary space.
 For higher level operators, preparation creates a [config object](https://juliadiff.org/ForwardDiff.jl/stable/user/api/#Preallocating/Configuring-Work-Buffers), which can be type-unstable.
 
+### GTPSA
+
+For all operators, preparation preallocates the input [`TPS`s](https://bmad-sim.github.io/GTPSA.jl/stable/man/c_tps/), and for in-place functions the output `TPS`s as well. For minimal allocations of `TPS` temporaries inside of a function, the [`@FastGTPSA`/`@FastGTPSA!`](https://bmad-sim.github.io/GTPSA.jl/stable/man/j_fastgtpsa/) macros are recommended.
+
+If a GTPSA [`Descriptor`](https://bmad-sim.github.io/GTPSA.jl/stable/man/b_descriptor/) is not provided to `AutoGTPSA`, then a `Descriptor` will be generated in preparation based on the context.
+
+!!! danger
+    When providing a custom GTPSA `Descriptor` to `AutoGTPSA`, it is the responsibility of the user to ensure that the number of [GTPSA "variables"](https://bmad-sim.github.io/GTPSA.jl/stable/quickstart/#Calculating-a-Truncated-Power-Series) specified in the `Descriptor` is consistent with the number of inputs of the provided function. Undefined behavior and crashes may occur if this is not the case.
+
 ### PolyesterForwardDiff
 
 Most operators fall back on `AutoForwardDiff`.

diff --git a/...tiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl b/...tiationInterface/ext/DifferentiationInterfaceGTPSAExt/DifferentiationInterfaceGTPSAExt.jl
@@ -0,0 +1,12 @@
+module DifferentiationInterfaceGTPSAExt
+
+import DifferentiationInterface as DI
+using ADTypes: AutoGTPSA
+using GTPSA: GTPSA, TPS, Descriptor
+
+DI.check_available(::AutoGTPSA) = true
+
+include("onearg.jl")
+include("twoarg.jl")
+
+end