From 50ce141c0c223bc6ba702bc6048f756ff653a8b0 Mon Sep 17 00:00:00 2001
From: mattsignorelli <mgs255@cornell.edu>
Date: Thu, 8 Aug 2024 09:31:35 -0400
Subject: [PATCH] Add AutoGTPSA taylor mode

---
 docs/src/index.md |  6 ++++++
 src/ADTypes.jl    |  1 +
 src/dense.jl      | 36 ++++++++++++++++++++++++++++++++++++
 test/dense.jl     | 14 ++++++++++++++
 test/runtests.jl  |  3 +++
 5 files changed, 60 insertions(+)

diff --git a/docs/src/index.md b/docs/src/index.md
index 9640a33..2271540 100644
--- a/docs/src/index.md
+++ b/docs/src/index.md
@@ -30,6 +30,12 @@ AutoFiniteDiff
 AutoFiniteDifferences
 ```
 
+Taylor mode:
+
+```@docs
+AutoGTPSA
+```
+
 ### Reverse mode
 
 ```@docs
diff --git a/src/ADTypes.jl b/src/ADTypes.jl
index 0000a40..abdd2c4 100644
--- a/src/ADTypes.jl
+++ b/src/ADTypes.jl
@@ -39,6 +39,7 @@ export AutoChainRules,
        AutoFiniteDiff,
        AutoFiniteDifferences,
        AutoForwardDiff,
+       AutoGTPSA,
        AutoModelingToolkit,
        AutoPolyesterForwardDiff,
        AutoReverseDiff,
diff --git a/src/dense.jl b/src/dense.jl
index 8e6d960..c69d476 100644
--- a/src/dense.jl
+++ b/src/dense.jl
@@ -185,6 +185,42 @@ function Base.show(io::IO, backend::AutoForwardDiff{chunksize}) where {chunksize
     print(io, ")")
 end
 
+
+"""
+    AutoGTPSA{D}
+
+Struct used to select the [GTPSA.jl](https://github.com/bmad-sim/GTPSA.jl) backend for automatic differentiation.
+
+Defined by [ADTypes.jl](https://github.com/SciML/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
+
+mode(::AutoGTPSA) = ForwardMode()
+
+function Base.show(io::IO, backend::AutoGTPSA{D}) where {D}
+    print(io, AutoGTPSA, "(")
+    D != Nothing && print(io, "descriptor=\n", repr(backend.descriptor; context = io))
+    print(io, ")")
+end
+
+
+
 """
     AutoPolyesterForwardDiff{chunksize,T}
 
diff --git a/test/dense.jl b/test/dense.jl
index 7554565..267f152 100644
--- a/test/dense.jl
+++ b/test/dense.jl
@@ -104,6 +104,20 @@ end
     @test ad.tag == CustomTag()
 end
 
+@testset "AutoGTPSA" begin
+    ad = AutoGTPSA(; descriptor = nothing)
+    @test ad isa AbstractADType
+    @test ad isa AutoGTPSA{Nothing}
+    @test mode(ad) isa ForwardMode
+    @test ad.descriptor === nothing
+
+    ad = AutoGTPSA(; descriptor = Val(:descriptor))
+    @test ad isa AbstractADType
+    @test ad isa AutoGTPSA{Val{:descriptor}}
+    @test mode(ad) isa ForwardMode
+    @test ad.descriptor == Val(:descriptor)
+end
+
 @testset "AutoPolyesterForwardDiff" begin
     ad = AutoPolyesterForwardDiff()
     @test ad isa AbstractADType
diff --git a/test/runtests.jl b/test/runtests.jl
index 76fccba..6b82466 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -43,6 +43,7 @@ function every_ad()
         AutoFiniteDiff(),
         AutoFiniteDifferences(; fdm = :fdm),
         AutoForwardDiff(),
+        AutoGTPSA(),
         AutoPolyesterForwardDiff(),
         AutoReverseDiff(),
         AutoSymbolics(),
@@ -64,6 +65,8 @@ function every_ad_with_options()
         AutoFiniteDifferences(; fdm = :fdm),
         AutoForwardDiff(),
         AutoForwardDiff(chunksize = 3, tag = :tag),
+        AutoGTPSA(),
+        AutoGTPSA(descriptor=:descriptor),
         AutoPolyesterForwardDiff(),
         AutoPolyesterForwardDiff(chunksize = 3, tag = :tag),
         AutoReverseDiff(),