From f18da045bb46e9894f240666191bb89eebe96785 Mon Sep 17 00:00:00 2001
From: Guillaume Dalle <22795598+gdalle@users.noreply.github.com>
Date: Wed, 29 May 2024 09:51:25 +0200
Subject: [PATCH] Start using DifferentiationInterface

---
 Project.toml  |  4 +++-
 src/CTBase.jl |  3 ++-
 src/utils.jl  | 24 ++++++++++++++++++------
 3 files changed, 23 insertions(+), 8 deletions(-)

diff --git a/Project.toml b/Project.toml
index 421e937e..0a76d1a7 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,10 +1,11 @@
 name = "CTBase"
 uuid = "54762871-cc72-4466-b8e8-f6c8b58076cd"
 authors = ["Olivier Cots <olivier.cots@enseeiht.fr>"]
-version = "0.9.0"
+version = "0.9.1"
 
 [deps]
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
@@ -21,6 +22,7 @@ Unicode = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
 
 [compat]
 DataStructures = "0.18"
+DifferentiationInterface = "0.5"
 DocStringExtensions = "0.9"
 ForwardDiff = "0.10"
 Interpolations = "0.15"
diff --git a/src/CTBase.jl b/src/CTBase.jl
index 44d30887..c4fd3602 100644
--- a/src/CTBase.jl
+++ b/src/CTBase.jl
@@ -14,7 +14,8 @@ module CTBase
 # using
 import Base
 using DocStringExtensions
-using ForwardDiff: jacobian, gradient, ForwardDiff # automatic differentiation
+using DifferentiationInterface: AutoForwardDiff, derivative, gradient, jacobian, prepare_derivative, prepare_gradient, prepare_jacobian
+using ForwardDiff: ForwardDiff # automatic differentiation
 using Interpolations: linear_interpolation, Line, Interpolations # for default interpolation
 using MLStyle # pattern matching
 using Parameters # @with_kw: to have default values in struct
diff --git a/src/utils.jl b/src/utils.jl
index 2a597f1e..476b0dee 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -74,7 +74,9 @@ $(TYPEDSIGNATURES)
 Return the gradient of `f` at `x`.
 """
 function ctgradient(f::Function, x::ctNumber)
-    return ForwardDiff.derivative(x -> f(x), x)
+    backend = AutoForwardDiff()
+    extras = prepare_derivative(f, backend, x)
+    return derivative(f, backend, x, extras)
 end
 
 """
@@ -83,7 +85,9 @@ $(TYPEDSIGNATURES)
 Return the gradient of `f` at `x`.
 """
 function ctgradient(f::Function, x)
-    return ForwardDiff.gradient(f, x)
+    backend = AutoForwardDiff()
+    extras = prepare_gradient(f, backend, x)
+    return gradient(f, backend, x, extras)
 end
 
 """
@@ -98,8 +102,12 @@ $(TYPEDSIGNATURES)
 
 Return the Jacobian of `f` at `x`.
 """
-function ctjacobian(f::Function, x::ctNumber) 
-    return ForwardDiff.jacobian(x -> f(x[1]), [x])
+function ctjacobian(f::Function, x::ctNumber)
+    f_number_to_number = only ∘ f ∘ only
+    backend = AutoForwardDiff()
+    extras = prepare_derivative(f_number_to_number, backend, x)
+    der = derivative(f_number_to_number, backend, x, extras)
+    return [der;;]
 end
 
 """
@@ -107,7 +115,11 @@ $(TYPEDSIGNATURES)
 
 Return the Jacobian of `f` at `x`.
 """
-ctjacobian(f::Function, x) = ForwardDiff.jacobian(f, x)
+function ctjacobian(f::Function, x)
+    backend = AutoForwardDiff()
+    extras = prepare_jacobian(f, backend, x)
+    return jacobian(f, backend, x, extras)
+end
 
 """
 $(TYPEDSIGNATURES)
@@ -195,4 +207,4 @@ function matrix2vec(x::Matrix{<:ctNumber}, dim::Integer=__matrix_dimension_stock
         end
     end
     return y
-end
\ No newline at end of file
+end