|
| 1 | +# Differentiation: AD backend strategies |
| 2 | + |
| 3 | +```@meta |
| 4 | +CurrentModule = CTBase |
| 5 | +``` |
| 6 | + |
| 7 | +The [`CTBase.Differentiation`](@ref CTBase.Differentiation) submodule provides an |
| 8 | +**automatic differentiation (AD) backend** abstraction built on the |
| 9 | +[Strategies](implementing-a-strategy.md) contract. It exposes a small set of |
| 10 | +differentiation primitives — gradients, derivatives, partial derivatives and |
| 11 | +Jacobian–vector products — behind a single strategy type, so that consumers |
| 12 | +(Hamiltonian systems, flows, differential geometry) never depend on a concrete |
| 13 | +AD package directly. |
| 14 | + |
| 15 | +!!! tip "Prerequisites" |
| 16 | + Read the [Implementing a Strategy](@ref) and [Options System](@ref) guides |
| 17 | + first: an AD backend is a strategy with a single `:ad_backend` option. |
| 18 | + |
| 19 | +```@setup diff |
| 20 | +using CTBase |
| 21 | +using CTBase.Differentiation |
| 22 | +using CTBase.Data |
| 23 | +using CTBase.Strategies |
| 24 | +# Loading DifferentiationInterface (+ an AD package) activates the extension |
| 25 | +# that provides the actual differentiation methods. |
| 26 | +using DifferentiationInterface |
| 27 | +using ForwardDiff |
| 28 | +import ADTypes |
| 29 | +``` |
| 30 | + |
| 31 | +## Overview |
| 32 | + |
| 33 | +The module is organised around a two-level strategy contract: |
| 34 | + |
| 35 | +| Level | Symbol | Role | |
| 36 | +|---|---|---| |
| 37 | +| Family | [`Differentiation.AbstractADBackend`](@ref CTBase.Differentiation.AbstractADBackend) | abstract strategy; defines the AD contract | |
| 38 | +| Concrete | [`Differentiation.DifferentiationInterface`](@ref CTBase.Differentiation.DifferentiationInterface) | wraps a [DifferentiationInterface.jl](https://github.com/JuliaDiff/DifferentiationInterface.jl) backend | |
| 39 | + |
| 40 | +The concrete strategy stores a single option, `:ad_backend`, holding an |
| 41 | +[ADTypes.jl](https://github.com/SciML/ADTypes.jl) backend (default |
| 42 | +`AutoForwardDiff()`). `ADTypes` is a hard dependency of CTBase, so the default is |
| 43 | +always available. |
| 44 | + |
| 45 | +The contract has seven methods. |
| 46 | +[`Differentiation.ad_backend`](@ref CTBase.Differentiation.ad_backend) — the accessor |
| 47 | +returning the wrapped ADTypes backend — is resolved in core and always available; the |
| 48 | +six differentiation primitives below live in an extension. |
| 49 | + |
| 50 | +!!! note "The differentiation methods live in an extension" |
| 51 | + The contract methods (`gradient`, `derivative`, `differentiate`, |
| 52 | + `pushforward`, `hamiltonian_gradient`, `variable_gradient`) are implemented in |
| 53 | + the `CTBaseDifferentiationInterface` **package extension**. They become |
| 54 | + available only once `DifferentiationInterface` *and* a concrete AD package |
| 55 | + (e.g. `ForwardDiff`) are loaded. Without the extension, every contract method |
| 56 | + throws a [`CTBase.Exceptions.NotImplemented`](@ref) with a helpful message. |
| 57 | + |
| 58 | +## Building a backend |
| 59 | + |
| 60 | +Use [`Differentiation.build_ad_backend`](@ref CTBase.Differentiation.build_ad_backend) |
| 61 | +to construct a backend with the default AD type: |
| 62 | + |
| 63 | +```@example diff |
| 64 | +backend = Differentiation.build_ad_backend() |
| 65 | +``` |
| 66 | + |
| 67 | +Pass an explicit ADTypes backend through the `ad_backend` option (the aliases |
| 68 | +`backend` and `ad` resolve to the same option): |
| 69 | + |
| 70 | +```@example diff |
| 71 | +Differentiation.build_ad_backend(; ad_backend = ADTypes.AutoForwardDiff()) |
| 72 | +``` |
| 73 | + |
| 74 | +The wrapped AD type is recovered with |
| 75 | +[`Differentiation.ad_backend`](@ref CTBase.Differentiation.ad_backend): |
| 76 | + |
| 77 | +```@example diff |
| 78 | +Differentiation.ad_backend(backend) |
| 79 | +``` |
| 80 | + |
| 81 | +Because it is a regular strategy, its metadata is introspectable at the type level: |
| 82 | + |
| 83 | +```@example diff |
| 84 | +Strategies.metadata(Differentiation.DifferentiationInterface) |
| 85 | +``` |
| 86 | + |
| 87 | +## Differentiation primitives |
| 88 | + |
| 89 | +### Gradient and derivative |
| 90 | + |
| 91 | +[`Differentiation.gradient`](@ref CTBase.Differentiation.gradient) differentiates a |
| 92 | +scalar function of a vector, and |
| 93 | +[`Differentiation.derivative`](@ref CTBase.Differentiation.derivative) a scalar |
| 94 | +function of a scalar: |
| 95 | + |
| 96 | +```@example diff |
| 97 | +f(x) = sum(x .^ 2) # ∇f = 2x |
| 98 | +Differentiation.gradient(backend, f, [1.0, 2.0]) |
| 99 | +``` |
| 100 | + |
| 101 | +```@example diff |
| 102 | +g(t) = t^3 # g'(t) = 3t² |
| 103 | +Differentiation.derivative(backend, g, 2.0) |
| 104 | +``` |
| 105 | + |
| 106 | +### Partial derivative at a slot |
| 107 | + |
| 108 | +[`Differentiation.differentiate`](@ref CTBase.Differentiation.differentiate) |
| 109 | +computes the partial derivative of `f` with respect to the argument at a |
| 110 | +**compile-time slot** `Val(Slot)`, holding the other arguments fixed. The active |
| 111 | +argument comes first, the frozen ones follow in slot order: |
| 112 | + |
| 113 | +```@example diff |
| 114 | +# H(x, p) = ½‖p‖² + ‖x‖² → ∂H/∂x = 2x, ∂H/∂p = p |
| 115 | +H(x, p) = 0.5 * sum(p .^ 2) + sum(x .^ 2) |
| 116 | +x = [1.0, 2.0]; p = [3.0, 4.0] |
| 117 | +
|
| 118 | +∂x = Differentiation.differentiate(backend, H, Val(1), x, p) # active = x (slot 1), const = p |
| 119 | +∂p = Differentiation.differentiate(backend, H, Val(2), p, x) # active = p (slot 2), const = x |
| 120 | +(∂x, ∂p) |
| 121 | +``` |
| 122 | + |
| 123 | +A scalar slot yields a scalar derivative: |
| 124 | + |
| 125 | +```@example diff |
| 126 | +# f(t, x) = t * x[1] → ∂f/∂t = x[1] |
| 127 | +ft(t, x) = t * x[1] |
| 128 | +Differentiation.differentiate(backend, ft, Val(1), 2.0, [3.0, 4.0]) |
| 129 | +``` |
| 130 | + |
| 131 | +### Pushforward (Jacobian–vector product) |
| 132 | + |
| 133 | +[`Differentiation.pushforward`](@ref CTBase.Differentiation.pushforward) computes the |
| 134 | +directional derivative of `f` at `x` along `dx`, freezing the other arguments: |
| 135 | + |
| 136 | +```@example diff |
| 137 | +A = [0.0 1.0; -1.0 0.0] |
| 138 | +X(x) = A * x # J_X = A → pushforward = A·dx |
| 139 | +Differentiation.pushforward(backend, X, Val(1), [1.0, 2.0], [5.0, 6.0]) |
| 140 | +``` |
| 141 | + |
| 142 | +## Hamiltonian gradients |
| 143 | + |
| 144 | +The two domain-specific methods operate directly on a |
| 145 | +[`Data.Hamiltonian`](@ref CTBase.Data.Hamiltonian) (see the [Data](data.md) guide). |
| 146 | +[`Differentiation.hamiltonian_gradient`](@ref CTBase.Differentiation.hamiltonian_gradient) |
| 147 | +returns `(∂H/∂x, ∂H/∂p)`, **non-negated** — the caller applies the signs of |
| 148 | +Hamilton's equations (`ẋ = ∂H/∂p`, `ṗ = -∂H/∂x`): |
| 149 | + |
| 150 | +```@example diff |
| 151 | +# H(x, p) = ½(‖x‖² + ‖p‖²) → ∂H/∂x = x, ∂H/∂p = p |
| 152 | +h = Data.Hamiltonian((x, p) -> 0.5 * (sum(abs2, x) + sum(abs2, p))) |
| 153 | +∂x, ∂p = Differentiation.hamiltonian_gradient(backend, h, 0.0, [1.0, 2.0], [3.0, 4.0], nothing) |
| 154 | +(∂x, ∂p) |
| 155 | +``` |
| 156 | + |
| 157 | +For a non-fixed Hamiltonian, |
| 158 | +[`Differentiation.variable_gradient`](@ref CTBase.Differentiation.variable_gradient) |
| 159 | +returns `∂H/∂v`: |
| 160 | + |
| 161 | +```@example diff |
| 162 | +# H(x, p, v) = ½(‖x‖² + ‖p‖² + ‖v‖²) → ∂H/∂v = v |
| 163 | +hv = Data.Hamiltonian((x, p, v) -> 0.5 * (sum(abs2, x) + sum(abs2, p) + sum(abs2, v)); is_variable=true) |
| 164 | +Differentiation.variable_gradient(backend, hv, 0.0, [1.0, 2.0], [3.0, 4.0], [5.0, 6.0]) |
| 165 | +``` |
| 166 | + |
| 167 | +## The contract without the extension |
| 168 | + |
| 169 | +Every contract method has a fallback on `AbstractADBackend` that throws |
| 170 | +[`CTBase.Exceptions.NotImplemented`](@ref) (see the [Exceptions](exceptions.md) |
| 171 | +guide). This is what users hit when the AD package is not loaded, and what custom |
| 172 | +backends must override: |
| 173 | + |
| 174 | +```@repl diff |
| 175 | +struct MyBackend <: Differentiation.AbstractADBackend |
| 176 | + options::Strategies.StrategyOptions |
| 177 | +end |
| 178 | +try # hide |
| 179 | +Differentiation.gradient(MyBackend(Strategies.StrategyOptions()), x -> sum(x), [1.0]) |
| 180 | +catch e # hide |
| 181 | +showerror(IOContext(stdout, :color => false), e) # hide |
| 182 | +end # hide |
| 183 | +``` |
| 184 | + |
| 185 | +## See also |
| 186 | + |
| 187 | +- [`Differentiation.AbstractADBackend`](@ref CTBase.Differentiation.AbstractADBackend), [`Differentiation.DifferentiationInterface`](@ref CTBase.Differentiation.DifferentiationInterface) — the strategy types. |
| 188 | +- [Implementing a Strategy](@ref), [Options System](@ref) — the contract these build on. |
| 189 | +- [Data](data.md) — `Hamiltonian` wrappers consumed by the gradient methods. |
| 190 | +- [Exceptions](exceptions.md) — `NotImplemented`, raised by the contract fallbacks. |
0 commit comments