GTPSA.jl: A Julia interface to the new GTPSA library for fast, high-order automatic differentiation

[mattsignorelli]

GTPSA.jl is a new Julia package that provides a full-featured interface to the Generalised Truncated Power Series Algebra library for computing Taylor expansions, or Truncated Power Series (TPSs) of real and complex multivariable functions to arbitrary orders.

Truncated Power Series Algebra (TPSA) performs forward-mode automatic differentation (AD) similar to the dual-number implementation as in ForwardDiff.jl. However, instead of nesting derivatives for higher orders, TPSA naturally extends to arbitrary orders by directly using the power series expansions. This, paired with a highly optimized monomial indexing function and storage for propagating the partial derivatives, makes GTPSA.jl significantly faster for 2nd-order calculations and above, and have similar performance at 1st-order. In our example, GTPSA was x3.3 faster than ForwardDiff to 2nd order, and x18.5 faster to 3rd order.

GTPSA provides several advantages over current Julia AD packages:

1. Speed: GTPSA.jl is significantly faster than ForwardDiff.jl for 2nd-order calculations and above, and has very similar performance at 1st-order
2. Custom Orders in Individual Variables: Other packages use a single maximum order for all variables. With GTPSA, the maximum order can be set differently for individual variables, as well as for a separate part of the monomial. For example, computing the Taylor expansion of f(x_1,x_2) to 1st order in x_1 and 3rd order in x_2 is possible
3. Complex Numbers: GTPSA.jl natively supports complex numbers and allows for mixing of complex and real truncated power series
4. Distinction Between State Variables and Parameters: Distinguishing between dependent variables and parameters in the solution of a differential equation expressed as a power series in the dependent variables/parameters can be advantageous in analysis

To read more about the inner-workings of the GTPSA C library, advanced users are referred to this paper, written by the creators of GTPSA.

As we’d like to prepare for a v1.0 release, we welcome any comments or criticisms of the package interface and documentation!

Because of the power of the GTPSA library, our long-term goal is to rewrite the library fully in Julia. If you are interested in this effort, feel free to reach out or add to the current issue on Github.

[lrnv]

How does it compare to TaylorDiff.jl ? It seems like the same kind of ideas (leverage Faa di Bruno formula), but TaylorDiff does that at Julia’s compile time.

[mattsignorelli]

For multivariable functions, TaylorDiff.jl only seems to give you the directional derivatives, whereas GTPSA.jl will give you all of the higher-order partial derivatives. Because of this, I haven’t been able to come up with an apples-to-apples benchmark to compare the two.

[mattsignorelli]

Also, GTPSA does much more than just use the power series directly for calculating higher orders; the library plays a lot of tricks in both storing and propagating efficiently the many monomial coefficients, making it especially powerful for calculating derivatives w.r.t. to many variables to high orders.

[nsajko]

I opened two issues on your Github repo page. Some questions:

How does GTPSA.jl compare to Diffractor.jl?

How does GTPSA compare to the ADOL-C library? It also seems to have an in-development Julia wrapper package: GitHub - TimSiebert1/ADOLC.jl

[gdalle]

Congrats on the release! Do you think it might make sense to add GPTSA.jl bindings to DifferentiationInterface.jl?
I didn’t add any for TaylorDiff.jl because it does not support things like gradient or Hessian.

[mattsignorelli]

@nsajko

I opened two issues on your Github repo page

Thank you! These are very helpful

@nsajko

How does GTPSA.jl compare to Diffractor.jl?

I just tried to add this as a benchmark in our example, however Diffractor.jl is returning an incorrect Jacobian. The benchmark uses a function m: \mathbb{R} ^{58} \rightarrow \mathbb{R}^6:

julia> j = AD.jacobian(AD.ForwardDiffBackend(), m, zeros(58)) |> only
6×58 Matrix{Float64}:
 -0.699483      -3.08238     …   7.18128      4.7683      2.17832
  0.0406188     -1.25064         1.49797      1.36552     1.18237
  0.0            0.0             0.0          0.0         0.0
  0.0            0.0             0.0          0.0         0.0
  1.33161      -14.3074         -0.789128    -0.266338    0.0
  0.000133161   -0.00143074  …  -7.89128e-5  -2.66338e-5  0.0

julia> j = AD.jacobian(DiffractorForwardBackend(), m, zeros(58)) |> only
6×58 Matrix{Float64}:
 -1.13036e-19  -2.52907e-21  …   0.170851    1.36948    2.17832
  3.21068e-22  -1.14018e-19     -0.133861    0.180479   1.18237
  0.0           0.0              0.0         0.0        0.0
  0.0           0.0              0.0         0.0        0.0
 -0.0115275    -0.113072        -0.17197    -0.12261    0.0
 -1.15275e-6   -1.13072e-5   …  -1.7197e-5  -1.2261e-5  0.0

The Jacobian calculated with GTPSA.jl agrees exactly with ForwardDiff.jl. Therefore since Diffractor.jl is not calculating the correct Jacobian, it is probably not appropriate to add it as a benchmark yet. I will submit this as an issue to Diffractor.jl, though.

@nsajko

How does GTPSA compare to the ADOL-C library? It also seems to have an in-development Julia wrapper package: GitHub - TimSiebert1/ADOLC.jl

I haven’t heard of this library, but I would be really interested to see how GTPSA compares. I’ll keep an eye on GitHub - TimSiebert1/ADOLC.jl, and once it is partially usable I can look into creating a benchmark.

[mattsignorelli]

@gdalle

Congrats on the release! Do you think it might make sense to add GPTSA.jl bindings to DifferentiationInterface.jl?

Absolutely! We definitely want to use GTPSA.jl for AD with Julia’s many already-available optimizers, and this seems like an important first step. @nsajko already opened an issue about this on the GTPSA Github, and I see in the documentation some steps on how to do this. I’ll take a look and try to make sense of how to implement this interface, and will probably reach out on the issues with questions.

[gdalle]

@mattsignorelli

I just tried to add this as a benchmark in our example,

Note that for ForwardDiff.jl to be optimally efficient, you need to provide it with a precomputed config object. DifferentiationInterface.jl allows that easily with the backend-agnostic “preparation” mechanism.

[lrnv]

@gdalle

I didn’t add any for TaylorDiff.jl because it does not support things like gradient or Hessian.

This is weird, it should. Maybe not as “interfaced” as other packages, but it should work.

[gdalle]

The only interface it has is TaylorDiff.derivative, see the API reference. Of course I could reconstruct everything else from there, I just didn’t take the time, and there might be smarter ways to do it from within TaylorDiff

[oxinabox]

One interesting case for a taylor mode AD is how it handles nesting.
The goal is to avoid nested forward-mode and use taylor mode instead.
But you can’t away do so easily at top level, because of encapsulation – some function is called in the middle of code you are ADing and that function happens to internally call AD.

In Diffractor, we have a flag (era-mode) that transforms nested forwards mode into taylor forward mode where it is encountered.
I think it is possible to do this with operator overloading AD.

[gdalle]

@oxinabox

In Diffractor, we have a flag (era-mode) that transforms nested forwards mode into taylor forward mode where it is encountered.

In DifferentiationInterface.jl, we have the internal DifferentiationInterface.nested which lets the autodiff backend know it is inside another autodiff operator.

In Lux.jl there is something similar with Nested Automatic Differentiation | Lux.jl Documentation

[oxinabox]

Code in diffractor for reference:

JuliaDiff/Diffractor.jl/blob/main/src/stage1/forward.jl#L261-L263