# Advanced tutorial

We present contexts and sparsity handling with DifferentiationInterface.jl.

```@example tuto_advanced
using ADTypes
using BenchmarkTools
using DifferentiationInterface
using ForwardDiff: ForwardDiff
using Zygote: Zygote
using Random
using SparseConnectivityTracer
using SparseMatrixColorings
```

## Contexts

Assume you want differentiate a multi-argument function with respect to the first argument.

```@example tuto_advanced
f_multiarg(x, c) = c * sum(abs2, x)
nothing  # hide
```

The first way, which works with every backend, is to create a closure:

```@example tuto_advanced
f_singlearg(c) = x -> f_multiarg(x, c)
nothing  # hide
```

Let's see it in action:

```@example tuto_advanced
backend = AutoForwardDiff()
x = float.(1:3)

gradient(f_singlearg(10), backend, x)
```

However, for performance reasons, it is sometimes preferrable to avoid closures and pass all arguments to the original function.
We can do this by wrapping `c` into a [`Constant`](@ref) and giving this constant to the `gradient` operator.

```@example tuto_advanced
gradient(f_multiarg, backend, x, Constant(10))
```

Preparation also works in this case, even if the constant changes before execution:

```@example tuto_advanced
prep_other_constant = prepare_gradient(f_multiarg, backend, x, Constant(-1))
gradient(f_multiarg, prep_other_constant, backend, x, Constant(10))
```

For additional arguments which act as mutated buffers, the [`Cache`](@ref) wrapper is the appropriate choice instead of [`Constant`](@ref).

## Sparsity

!!! tip

    If you use DifferentiationInterface's Sparse AD functionality in your research,
    please cite our preprint [*Sparser, Better, Faster, Stronger: Efficient Automatic Differentiation for Sparse Jacobians and Hessians*](https://arxiv.org/abs/2501.17737).

Sparse AD is very useful when Jacobian or Hessian matrices have a lot of zeros.
So let us write functions that satisfy this property.

```@example tuto_advanced
f_sparse_vector(x::AbstractVector) = diff(x .^ 2) + diff(reverse(x .^ 2))
f_sparse_scalar(x::AbstractVector) = sum(f_sparse_vector(x) .^ 2)
nothing  # hide
```

### Dense backends

When we use the [`jacobian`](@ref) or [`hessian`](@ref) operator with a dense backend, we get a dense matrix with plenty of zeros.

```@example tuto_advanced
x = float.(1:8);
```

```@example tuto_advanced
dense_forward_backend = AutoForwardDiff()
J_dense = jacobian(f_sparse_vector, dense_forward_backend, x)
```

```@example tuto_advanced
dense_second_order_backend = SecondOrder(AutoForwardDiff(), AutoZygote())
H_dense = hessian(f_sparse_scalar, dense_second_order_backend, x)
```

The results are correct but the procedure is very slow.
By using a sparse backend, we can get the runtime to increase with the number of nonzero elements, instead of the total number of elements.

### Sparse backends

Recipe to create a sparse backend: combine a dense backend, a sparsity detector and a compatible coloring algorithm inside [`AutoSparse`](@extref ADTypes.AutoSparse).
The following are reasonable defaults:

```@example tuto_advanced
sparse_forward_backend = AutoSparse(
    dense_forward_backend;  # any object from ADTypes
    sparsity_detector=TracerSparsityDetector(),
    coloring_algorithm=GreedyColoringAlgorithm(),
)

sparse_second_order_backend = AutoSparse(
    dense_second_order_backend;  # any object from ADTypes or a SecondOrder from DI
    sparsity_detector=TracerSparsityDetector(),
    coloring_algorithm=GreedyColoringAlgorithm(),
)
nothing  # hide
```

Now the resulting matrices are sparse:

```@example tuto_advanced
jacobian(f_sparse_vector, sparse_forward_backend, x)
```

```@example tuto_advanced
hessian(f_sparse_scalar, sparse_second_order_backend, x)
```

### Sparse preparation

In the examples above, we didn't use preparation.
Sparse preparation is more costly than dense preparation, but it is even more essential.
Indeed, once preparation is done, sparse differentiation is much faster than dense differentiation, because it makes fewer calls to the underlying function.

Some result analysis functions from [SparseMatrixColorings.jl](https://github.com/JuliaDiff/SparseMatrixColorings.jl) can help you figure out what the preparation contains.
First, it records the sparsity pattern itself (the one returned by the detector).

```@example tuto_advanced
jac_prep = prepare_jacobian(f_sparse_vector, sparse_forward_backend, x)
sparsity_pattern(jac_prep)
```

In forward mode, each column of the sparsity pattern gets a color.

```@example tuto_advanced
column_colors(jac_prep)
```

And the colors in turn define non-overlapping groups (for Jacobians at least, Hessians are a bit more complicated).

```@example tuto_advanced
column_groups(jac_prep)
```

### Sparsity speedup

When preparation is used, the speedup due to sparsity becomes very visible in large dimensions.

```@example tuto_advanced
xbig = rand(1000)
nothing  # hide
```

```@example tuto_advanced
jac_prep_dense = prepare_jacobian(f_sparse_vector, dense_forward_backend, zero(xbig))
@benchmark jacobian($f_sparse_vector, $jac_prep_dense, $dense_forward_backend, $xbig)
```

```@example tuto_advanced
jac_prep_sparse = prepare_jacobian(f_sparse_vector, sparse_forward_backend, zero(xbig))
@benchmark jacobian($f_sparse_vector, $jac_prep_sparse, $sparse_forward_backend, $xbig)
```

Better memory use can be achieved by pre-allocating the matrix from the preparation result (so that it has the correct structure).

```@example tuto_advanced
jac_buffer = similar(sparsity_pattern(jac_prep_sparse), eltype(xbig))
@benchmark jacobian!(
    $f_sparse_vector, $jac_buffer, $jac_prep_sparse, $sparse_forward_backend, $xbig
)
```

And for optimal speed, one should write non-allocating and type-stable functions.

```@example tuto_advanced
function f_sparse_vector!(y::AbstractVector, x::AbstractVector)
    n = length(x)
    for i in eachindex(y)
        y[i] = abs2(x[i + 1]) - abs2(x[i]) + abs2(x[n - i]) - abs2(x[n - i + 1])
    end
    return nothing
end

ybig = zeros(length(xbig) - 1)
f_sparse_vector!(ybig, xbig)
ybig ≈ f_sparse_vector(xbig)
```

In this case, the sparse Jacobian should also become non-allocating (for our specific choice of backend).

```@example tuto_advanced
jac_prep_sparse_nonallocating = prepare_jacobian(
    f_sparse_vector!, zero(ybig), sparse_forward_backend, zero(xbig)
)
jac_buffer = similar(sparsity_pattern(jac_prep_sparse_nonallocating), eltype(xbig))
@benchmark jacobian!(
    $f_sparse_vector!,
    $ybig,
    $jac_buffer,
    $jac_prep_sparse_nonallocating,
    $sparse_forward_backend,
    $xbig,
)
```

### Mixed mode

Some Jacobians have a structure which includes dense rows and dense columns, like this one:

```@example tuto_advanced
arrowhead(x) = x .+ x[1] .+ vcat(sum(x), zeros(eltype(x), length(x)-1))

jacobian_sparsity(arrowhead, x, TracerSparsityDetector())
```

In such cases, sparse AD is only beneficial in "mixed mode", where we combine a forward and a reverse backend.
This is achieved using the [`MixedMode`](@ref) wrapper, for which we recommend a random coloring order (see [`RandomOrder`](@extref SparseMatrixColorings.RandomOrder)):

```@example tuto_advanced
sparse_mixed_backend = AutoSparse(
    MixedMode(AutoForwardDiff(), AutoZygote());
    sparsity_detector=TracerSparsityDetector(),
    coloring_algorithm=GreedyColoringAlgorithm(RandomOrder(MersenneTwister(), 0)),
)
```

It unlocks a large speedup compared to pure forward mode, and the same would be true compared to reverse mode:

```@example tuto_advanced
@benchmark jacobian($arrowhead, prep, $sparse_forward_backend, $xbig) setup=(
    prep=prepare_jacobian(arrowhead, sparse_forward_backend, xbig)
)
```

```@example tuto_advanced
@benchmark jacobian($arrowhead, prep, $sparse_mixed_backend, $xbig) setup=(
    prep=prepare_jacobian(arrowhead, sparse_mixed_backend, xbig)
)
```
