From 5c53d904a58c54e2386268b8bc385ae48c9af9c0 Mon Sep 17 00:00:00 2001 From: Olivier Cots Date: Thu, 25 Jun 2026 17:03:34 +0200 Subject: [PATCH 1/4] Add CTBase.Differentiation submodule + DI extension (moved from CTFlows) AD backend strategy layer built on CTBase.Strategies: AbstractADBackend + DifferentiationInterface concrete strategy, with the differentiation methods (gradient, derivative, differentiate, pushforward, hamiltonian_gradient, variable_gradient) provided by the CTBaseDifferentiationInterface extension. - ADTypes becomes a hard dependency (default backend AutoForwardDiff). - DifferentiationInterface added as weakdep + extension. - __ad_backend default moved here from CTFlows.Common. - Version bumped to 0.24.0-beta; weakdeps mirrored into [extras]. Co-Authored-By: Claude Opus 4.8 --- Project.toml | 17 +- ext/CTBaseDifferentiationInterface.jl | 228 +++++++++++++++ src/CTBase.jl | 4 + src/Differentiation/Differentiation.jl | 85 ++++++ src/Differentiation/abstract_ad_backend.jl | 272 ++++++++++++++++++ src/Differentiation/building.jl | 24 ++ src/Differentiation/default.jl | 18 ++ .../differentiation_interface.jl | 104 +++++++ 8 files changed, 751 insertions(+), 1 deletion(-) create mode 100644 ext/CTBaseDifferentiationInterface.jl create mode 100644 src/Differentiation/Differentiation.jl create mode 100644 src/Differentiation/abstract_ad_backend.jl create mode 100644 src/Differentiation/building.jl create mode 100644 src/Differentiation/default.jl create mode 100644 src/Differentiation/differentiation_interface.jl diff --git a/Project.toml b/Project.toml index 6ce69e9c..db0f3896 100644 --- a/Project.toml +++ b/Project.toml @@ -1,42 +1,57 @@ name = "CTBase" uuid = "54762871-cc72-4466-b8e8-f6c8b58076cd" -version = "0.23.0-beta" +version = "0.24.0-beta" authors = ["Olivier Cots ", "Jean-Baptiste Caillau "] [deps] +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" [weakdeps] Coverage = "a2441757-f6aa-5fb2-8edb-039e3f45d037" +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a" MarkdownAST = "d0879d2d-cac2-40c8-9cee-1863dc0c7391" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [extensions] +CTBaseDifferentiationInterface = ["DifferentiationInterface"] CoveragePostprocessing = ["Coverage"] DocumenterReference = ["Documenter", "Markdown", "MarkdownAST"] TestRunner = ["Test"] [extras] +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +Coverage = "a2441757-f6aa-5fb2-8edb-039e3f45d037" +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" +Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a" +MarkdownAST = "d0879d2d-cac2-40c8-9cee-1863dc0c7391" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] test = [ "Aqua", "Coverage", + "DifferentiationInterface", "Documenter", + "ForwardDiff", "Markdown", "MarkdownAST", "Test" ] [compat] +ADTypes = "1" Aqua = "0.8" Coverage = "1" +DifferentiationInterface = "0.7" DocStringExtensions = "0.9" Documenter = "1" +ForwardDiff = "0.10, 1" Markdown = "1" MarkdownAST = "0.1" Test = "1" diff --git a/ext/CTBaseDifferentiationInterface.jl b/ext/CTBaseDifferentiationInterface.jl new file mode 100644 index 00000000..b72b170e --- /dev/null +++ b/ext/CTBaseDifferentiationInterface.jl @@ -0,0 +1,228 @@ +""" + CTBaseDifferentiationInterface + +Package extension providing DifferentiationInterface.jl backend implementations +for automatic differentiation in `CTBase.Differentiation`. + +Activated automatically when `DifferentiationInterface` is loaded together with `CTBase`. + +This extension provides: +- `Differentiation.hamiltonian_gradient` — Hamiltonian gradient (∂H/∂x, ∂H/∂p) via DI +- `Differentiation.variable_gradient` — variable gradient ∂H/∂v via DI +- `Differentiation.gradient` / `Differentiation.derivative` — general AD primitives +- `Differentiation.differentiate` / `Differentiation.pushforward` — partial derivative / JVP primitives +""" +module CTBaseDifferentiationInterface + +import DocStringExtensions: TYPEDEF, TYPEDSIGNATURES +using CTBase.Data: Data +using CTBase.Differentiation: Differentiation +using DifferentiationInterface: DifferentiationInterface as DI + +# ============================================================================== +# Differentiation.hamiltonian_gradient / variable_gradient +# ============================================================================== + +""" +$(TYPEDSIGNATURES) + +Return the scalar DI differentiation primitive for a `Number` active argument. + +Dispatches to `DI.derivative`, which computes `df/dx` for scalar `x`. + +See also: [`_derivator(::Type{<:AbstractArray})`](@ref) +""" +function _derivator(::Type{<:Number}) + return DI.derivative +end + +""" +$(TYPEDSIGNATURES) + +Return the array DI differentiation primitive for an `AbstractArray` active argument. + +Dispatches to `DI.gradient`, which computes `∇f` for array `x`. + +See also: [`_derivator(::Type{<:Number})`](@ref) +""" +function _derivator(::Type{<:AbstractArray}) + return DI.gradient +end + +""" +$(TYPEDSIGNATURES) + +Compute Hamiltonian gradients (∂H/∂x, ∂H/∂p) via DifferentiationInterface.jl. + +Anonymous closures are used deliberately so that ForwardDiff `tagcount` values +are assigned at runtime in the correct left-to-right order inside `ForwardDiff.≺`, +avoiding silent zero-gradient bugs in nested-AD contexts (e.g. inside NonlinearSolve). + +# Returns +- Tuple `(grad_x, grad_p)` where `grad_x` = ∂H/∂x, `grad_p` = ∂H/∂p. +""" +function Differentiation.hamiltonian_gradient( + backend::Differentiation.DifferentiationInterface, + h::Data.AbstractHamiltonian, + t, x, p, v, +) + di_backend = Differentiation.ad_backend(backend) + h_x(x_) = h(t, x_, p, v) + h_p(p_) = h(t, x, p_, v) + grad_x = _derivator(typeof(x))(h_x, di_backend, x) + grad_p = _derivator(typeof(p))(h_p, di_backend, p) + return (grad_x, grad_p) +end + +""" +$(TYPEDSIGNATURES) + +Compute variable gradient ∂H/∂v via DifferentiationInterface.jl. + +See the note in [`hamiltonian_gradient`](@ref) on why anonymous closures are used. + +# Returns +- `grad_v` = ∂H/∂v. +""" +function Differentiation.variable_gradient( + backend::Differentiation.DifferentiationInterface, + h::Data.AbstractHamiltonian, + t, x, p, v, +) + di_backend = Differentiation.ad_backend(backend) + h_v(v_) = h(t, x, p, v_) + return _derivator(typeof(v))(h_v, di_backend, v) +end + +# ============================================================================= +# Differentiation.gradient — extension contract methods +# ============================================================================= + +""" +$(TYPEDSIGNATURES) + +Compute the gradient of a scalar function using DifferentiationInterface.jl. + +# Arguments +- `backend::Differentiation.DifferentiationInterface`: The AD backend. +- `f::Function`: The scalar function to differentiate. +- `x::AbstractArray`: The input vector. + +# Returns +- `∇f`: The gradient of `f` at `x`. + +# See also +- [`CTBase.Differentiation.derivative`](@ref) +""" +function Differentiation.gradient( + backend::Differentiation.DifferentiationInterface, + f::Function, + x::AbstractArray, +) + ad = Differentiation.ad_backend(backend) + return DI.gradient(f, ad, x) +end + +""" +$(TYPEDSIGNATURES) + +Compute the gradient of a scalar function using DifferentiationInterface.jl (scalar case). + +# Arguments +- `backend::Differentiation.DifferentiationInterface`: The AD backend. +- `f::Function`: The scalar function to differentiate. +- `x::Real`: The input scalar. + +# Returns +- `df/dx`: The derivative of `f` at `x`. + +# See also +- [`CTBase.Differentiation.derivative`](@ref) +""" +function Differentiation.gradient( + backend::Differentiation.DifferentiationInterface, + f::Function, + x::Real, +) + ad = Differentiation.ad_backend(backend) + return DI.derivative(f, ad, x) +end + +""" +$(TYPEDSIGNATURES) + +Compute the derivative of a scalar function using DifferentiationInterface.jl. + +# Arguments +- `backend::Differentiation.DifferentiationInterface`: The AD backend. +- `g::Function`: The scalar function to differentiate. +- `t::Real`: The input scalar. + +# Returns +- `dg/dt`: The derivative of `g` at `t`. + +# See also +- [`CTBase.Differentiation.gradient`](@ref) +""" +function Differentiation.derivative( + backend::Differentiation.DifferentiationInterface, + g::Function, + t::Real, +) + ad = Differentiation.ad_backend(backend) + return DI.derivative(g, ad, t) +end + +# ============================================================================= +# Differentiation.differentiate / pushforward — new primitives +# ============================================================================= + +""" +$(TYPEDSIGNATURES) + +Compute the partial derivative or gradient of `f` with respect to the argument at +slot `Slot`, using DifferentiationInterface.jl. + +An anonymous closure captures the constant arguments and places `active_` at `Slot` +via `ntuple` — same rationale as [`hamiltonian_gradient`](@ref) (ForwardDiff tag ordering). +`_derivator` dispatches to `DI.gradient` for array `active` and `DI.derivative` for scalar. + +See also: [`CTBase.Differentiation.pushforward`](@ref). +""" +function Differentiation.differentiate( + backend::Differentiation.DifferentiationInterface, + f, + ::Val{Slot}, + active, + consts::Vararg{Any, N}, +) where {Slot, N} + di = Differentiation.ad_backend(backend) + f_active(active_) = f(ntuple(i -> i == Slot ? active_ : consts[i < Slot ? i : i - 1], Val(N + 1))...) + return _derivator(typeof(active))(f_active, di, active) +end + +""" +$(TYPEDSIGNATURES) + +Compute the pushforward (Jacobian-vector product) of `f` at `x` in direction `dx`, +using DifferentiationInterface.jl. + +An anonymous closure captures `consts` and reconstructs the full argument tuple via +`ntuple`, placing `x_` at slot `Slot`. The single tangent is extracted with `only`. + +See also: [`CTBase.Differentiation.differentiate`](@ref). +""" +function Differentiation.pushforward( + backend::Differentiation.DifferentiationInterface, + f, + ::Val{Slot}, + x, + dx, + consts::Vararg{Any, N}, +) where {Slot, N} + di = Differentiation.ad_backend(backend) + f_slot(x_) = f(ntuple(i -> i == Slot ? x_ : consts[i < Slot ? i : i - 1], Val(N + 1))...) + return only(DI.pushforward(f_slot, di, x, (dx,))) +end + +end # module diff --git a/src/CTBase.jl b/src/CTBase.jl index 8a31d0cd..4c99ab53 100644 --- a/src/CTBase.jl +++ b/src/CTBase.jl @@ -40,6 +40,10 @@ using .Traits include(joinpath(@__DIR__, "Data", "Data.jl")) using .Data +# Differentiation module - AD backend strategies (moved from CTFlows) +include(joinpath(@__DIR__, "Differentiation", "Differentiation.jl")) +using .Differentiation + # Unicode module - Unicode character utilities include(joinpath(@__DIR__, "Unicode", "Unicode.jl")) using .Unicode diff --git a/src/Differentiation/Differentiation.jl b/src/Differentiation/Differentiation.jl new file mode 100644 index 00000000..35883b71 --- /dev/null +++ b/src/Differentiation/Differentiation.jl @@ -0,0 +1,85 @@ +# ============================================================================== +# Differentiation — AD Backend Strategies for Hamiltonian Gradients +# ============================================================================== + +""" +Module `CTBase.Differentiation` provides automatic differentiation backend strategies +for computing gradients of scalar Hamiltonian functions. + +## Architecture + +The module defines an abstract contract `AbstractADBackend` with the following methods: +- `hamiltonian_gradient(backend, h, t, x, p, v)` → (∂H/∂x, ∂H/∂p) +- `variable_gradient(backend, h, t, x, p, v)` → ∂H/∂v +- `gradient(backend, f, x)` → ∇f (extension contract) +- `derivative(backend, g, t)` → dg/dt (extension contract) +- `differentiate(backend, f, ::Val{Slot}, active, consts...)` → partial derivative at slot +- `pushforward(backend, f, ::Val{Slot}, x, dx, consts...)` → JVP along direction dx + +The concrete strategy `DifferentiationInterface` wraps DifferentiationInterface.jl backends +(e.g., `AutoForwardDiff()`) and stores them in its `:ad_backend` option. + +## Dependencies + +- `ADTypes.jl` (hard dependency) — provides `AutoForwardDiff` type +- `CTBase.Data` — `AbstractHamiltonian` in the gradient contract signatures +- `CTBase.Strategies` — strategy contract +- `CTBase.Exceptions` — `NotImplemented` for stub methods + +## Extension + +Gradient computation requires the `CTBaseDifferentiationInterface` extension, +which implements the contract methods using `DifferentiationInterface.gradient` +and friends. It is loaded automatically when `DifferentiationInterface` is loaded +together with `CTBase`. + +## Exports + +- `AbstractADBackend` +- `DifferentiationInterface` +- `build_ad_backend` +- `ad_backend` +- `hamiltonian_gradient` +- `variable_gradient` +- `gradient` +- `derivative` +- `differentiate` +- `pushforward` +""" +module Differentiation + +# ============================================================================== +# External Imports +# ============================================================================== + +import ADTypes +import DocStringExtensions: TYPEDEF, TYPEDSIGNATURES +import CTBase.Data +import CTBase.Exceptions +import CTBase.Strategies + +# ============================================================================== +# Includes (in dependency order) +# ============================================================================== + +include(joinpath(@__DIR__, "default.jl")) +include(joinpath(@__DIR__, "abstract_ad_backend.jl")) +include(joinpath(@__DIR__, "differentiation_interface.jl")) +include(joinpath(@__DIR__, "building.jl")) + +# ============================================================================== +# Exports +# ============================================================================== + +export AbstractADBackend +export DifferentiationInterface +export build_ad_backend +export ad_backend +export hamiltonian_gradient +export variable_gradient +export gradient +export derivative +export differentiate +export pushforward + +end # module diff --git a/src/Differentiation/abstract_ad_backend.jl b/src/Differentiation/abstract_ad_backend.jl new file mode 100644 index 00000000..f19d6b5c --- /dev/null +++ b/src/Differentiation/abstract_ad_backend.jl @@ -0,0 +1,272 @@ +# ============================================================================== +# AbstractADBackend — AD Backend Strategy Contract +# ============================================================================== + +""" +$(TYPEDEF) + +Abstract base type for automatic differentiation backends. + +An `AbstractADBackend` is a strategy that defines how to compute gradients of a +scalar Hamiltonian function. Concrete backends (e.g., `DifferentiationInterface`) +implement the contract methods to provide actual gradient computation. + +# Notes + - `AbstractADBackend` subtypes `CTBase.Strategies.AbstractStrategy` — they are + first-class strategies in the CTBase.Strategies ecosystem. + - The contract consists of two methods: `hamiltonian_gradient` and `variable_gradient`. + - Gradient methods return **non-negated** partial derivatives; the RHS closures + apply the signs (ṗ = -∂H/∂x, ṽ = -∂H/∂v). + +See also: [`CTBase.Differentiation.DifferentiationInterface`](@ref), +[`CTBase.Differentiation.hamiltonian_gradient`](@ref), +[`CTBase.Differentiation.variable_gradient`](@ref). +""" +abstract type AbstractADBackend <: Strategies.AbstractStrategy end + +# ============================================================================== +# Contract Methods +# ============================================================================== + +""" +$(TYPEDSIGNATURES) + +Compute the Hamiltonian gradient (∂H/∂x, ∂H/∂p) using the backend. + +# Arguments +- `backend::AbstractADBackend`: The AD backend. +- `h`: The Hamiltonian function or type. +- `t`: Time (scalar). +- `x`: State vector. +- `p`: Costate vector. +- `v`: Variable (scalar or `nothing` for Fixed problems). + +# Returns +- `(∂H_∂x, ∂H_∂p)`: Tuple of partial derivatives, **non-negated**. The RHS closure + is responsible for applying the signs (ṗ = -∂H/∂x). + +# Throws +- `CTBase.Exceptions.NotImplemented`: If the concrete backend does not implement this method. + +See also: [`CTBase.Differentiation.variable_gradient`](@ref). +""" +function hamiltonian_gradient( + backend::AbstractADBackend, + h::Data.AbstractHamiltonian, + t, x, p, v, +) + throw(Exceptions.NotImplemented( + "hamiltonian_gradient not implemented for $(typeof(backend))", + required_method = "hamiltonian_gradient(backend::$(typeof(backend)), h, t, x, p, v)", + suggestion = "Implement hamiltonian_gradient for $(typeof(backend)) or load an extension that provides gradient computation (e.g., CTBaseDifferentiationInterface)", + context = "AD backend contract" + )) +end + +""" +$(TYPEDSIGNATURES) + +Compute the variable gradient ∂H/∂v using the backend. + +# Arguments +- `backend::AbstractADBackend`: The AD backend. +- `h`: The Hamiltonian function or type. +- `t`: Time (scalar). +- `x`: State vector. +- `p`: Costate vector. +- `v`: Variable (scalar or `nothing` for Fixed problems). + +# Returns +- `∂H_∂v`: Partial derivative with respect to the variable, **non-negated**. The RHS + closure is responsible for applying the sign (ṽ = -∂H/∂v). + +# Throws +- `CTBase.Exceptions.NotImplemented`: If the concrete backend does not implement this method. + +See also: [`CTBase.Differentiation.hamiltonian_gradient`](@ref). +""" +function variable_gradient( + backend::AbstractADBackend, + h::Data.AbstractHamiltonian, + t, x, p, v, +) + throw(Exceptions.NotImplemented( + "variable_gradient not implemented for $(typeof(backend))", + required_method = "variable_gradient(backend::$(typeof(backend)), h, t, x, p, v)", + suggestion = "Implement variable_gradient for $(typeof(backend)) or load an extension that provides gradient computation (e.g., CTBaseDifferentiationInterface)", + context = "AD backend contract" + )) +end + +""" +$(TYPEDSIGNATURES) + +Extract the AD backend from a backend strategy. + +# Arguments +- `backend::AbstractADBackend`: The AD backend. + +# Returns +- `ADTypes.AbstractADType`: The concrete AD backend (e.g., `AutoForwardDiff()`). + +# Throws +- `CTBase.Exceptions.NotImplemented`: If the concrete backend does not implement + this method. + +# Notes + - This method is used by Hamiltonian vector-field getters to extract the AD backend + from a Hamiltonian system's backend. + - For `DifferentiationInterface`, this extracts the `:ad_backend` option. + +See also: [`CTBase.Differentiation.DifferentiationInterface`](@ref). +""" +function ad_backend(backend::AbstractADBackend) + throw(Exceptions.NotImplemented( + "ad_backend not implemented for $(typeof(backend))"; + required_method = "ad_backend(backend::$(typeof(backend)))", + suggestion = "Implement ad_backend for $(typeof(backend))", + context = "AD backend contract", + )) +end + +# ============================================================================= +# Extension contract — gradient and derivative methods +# ============================================================================= + +""" +$(TYPEDSIGNATURES) + +Compute the gradient of a scalar function using the backend. + +# Arguments +- `backend::AbstractADBackend`: The AD backend. +- `f`: The scalar function to differentiate. +- `x`: The input vector. + +# Returns +- `∇f`: The gradient of `f` at `x`. + +# Throws +- `CTBase.Exceptions.NotImplemented`: If the concrete backend does not implement + this method. + +# Notes + - This method is provided for extensions that implement gradient computation + via DifferentiationInterface.jl. + +See also: [`CTBase.Differentiation.derivative`](@ref), +[`CTBase.Differentiation.hamiltonian_gradient`](@ref). +""" +function gradient(backend::AbstractADBackend, f::Function, x) + throw(Exceptions.NotImplemented( + "gradient not implemented for $(typeof(backend))", + required_method = "gradient(backend::$(typeof(backend)), f::Function, x)", + suggestion = "Load CTBaseDifferentiationInterface (load DifferentiationInterface)", + context = "AD backend contract", + )) +end + +""" +$(TYPEDSIGNATURES) + +Compute the derivative of a scalar function using the backend. + +# Arguments +- `backend::AbstractADBackend`: The AD backend. +- `g`: The scalar function to differentiate. +- `t::Real`: The input scalar. + +# Returns +- `dg/dt`: The derivative of `g` at `t`. + +# Throws +- `CTBase.Exceptions.NotImplemented`: If the concrete backend does not implement + this method. + +# Notes + - This method is provided for extensions that implement derivative computation + via DifferentiationInterface.jl. + +See also: [`CTBase.Differentiation.gradient`](@ref), +[`CTBase.Differentiation.variable_gradient`](@ref). +""" +function derivative(backend::AbstractADBackend, g::Function, t::Real) + throw(Exceptions.NotImplemented( + "derivative not implemented for $(typeof(backend))", + required_method = "derivative(backend::$(typeof(backend)), g::Function, t::Real)", + suggestion = "Load CTBaseDifferentiationInterface (load DifferentiationInterface)", + context = "AD backend contract", + )) +end + +# ============================================================================= +# Partial differentiation and JVP primitives +# ============================================================================= + +""" +$(TYPEDSIGNATURES) + +Compute the partial derivative or gradient of `f` with respect to the argument at +slot `Slot`, holding all other arguments fixed. + +`f` is called as `f(arg₁, …, argₙ)`. The active argument is `active` (placed at +position `Slot`); the remaining `n-1` arguments are supplied as `consts...` in +slot order (skipping `Slot`). + +# Arguments +- `backend::AbstractADBackend`: The AD backend. +- `f`: The function to differentiate (any callable). +- `::Val{Slot}`: Compile-time slot index of the active argument. +- `active`: The point at which to differentiate. +- `consts...`: The fixed arguments, in order of their slot positions (excluding `Slot`). + +# Returns +- Gradient vector if `active isa AbstractArray`, scalar derivative if `active isa Real`. + +# Throws +- `CTBase.Exceptions.NotImplemented`: If the concrete backend does not implement this method. + +See also: [`CTBase.Differentiation.pushforward`](@ref). +""" +function differentiate(backend::AbstractADBackend, f, ::Val{Slot}, active, consts...) where {Slot} + throw(Exceptions.NotImplemented( + "differentiate not implemented for $(typeof(backend))"; + required_method = "differentiate(backend::$(typeof(backend)), f, ::Val{Slot}, active, consts...)", + suggestion = "Load CTBaseDifferentiationInterface (load DifferentiationInterface)", + context = "AD backend contract", + )) +end + +""" +$(TYPEDSIGNATURES) + +Compute the pushforward (Jacobian-vector product) of `f` at `x` in direction `dx`, +holding fixed arguments `consts...` at the slots other than `Slot`. + +Returns `d/ds f(x + s·dx, consts…)|_{s=0}` — the directional derivative of `f` +at `x` along `dx`, with all other arguments frozen. + +# Arguments +- `backend::AbstractADBackend`: The AD backend. +- `f`: The function to differentiate (any callable). +- `::Val{Slot}`: Compile-time slot index of the active (differentiated) argument. +- `x`: The point at which to differentiate. +- `dx`: The direction (tangent vector, same shape as `x`). +- `consts...`: The fixed arguments, in slot order (excluding `Slot`). + +# Returns +- The directional derivative `J_f(x) · dx`, same shape as `f(x, …)`. + +# Throws +- `CTBase.Exceptions.NotImplemented`: If the concrete backend does not implement this method. + +See also: [`CTBase.Differentiation.differentiate`](@ref). +""" +function pushforward(backend::AbstractADBackend, f, ::Val{Slot}, x, dx, consts...) where {Slot} + throw(Exceptions.NotImplemented( + "pushforward not implemented for $(typeof(backend))"; + required_method = "pushforward(backend::$(typeof(backend)), f, ::Val{Slot}, x, dx, consts...)", + suggestion = "Load CTBaseDifferentiationInterface (load DifferentiationInterface)", + context = "AD backend contract", + )) +end diff --git a/src/Differentiation/building.jl b/src/Differentiation/building.jl new file mode 100644 index 00000000..4e910e7c --- /dev/null +++ b/src/Differentiation/building.jl @@ -0,0 +1,24 @@ +# ============================================================================== +# build_ad_backend — Factory for AD Backends +# ============================================================================== + +""" +$(TYPEDSIGNATURES) + +Factory function to build an AD backend with default options. + +# Arguments +- `kwargs...`: Options passed to `DifferentiationInterface` constructor. + +# Returns +- `DifferentiationInterface`: A new AD backend strategy instance. + +# Notes + - This is a convenience factory that always returns a `DifferentiationInterface` + instance with the provided options. + +See also: [`CTBase.Differentiation.DifferentiationInterface`](@ref). +""" +function build_ad_backend(; kwargs...) + return DifferentiationInterface(; kwargs...) +end diff --git a/src/Differentiation/default.jl b/src/Differentiation/default.jl new file mode 100644 index 00000000..41e18b3f --- /dev/null +++ b/src/Differentiation/default.jl @@ -0,0 +1,18 @@ +# ============================================================================= +# Default values for the Differentiation module +# ============================================================================= + +""" +$(TYPEDSIGNATURES) + +Default AD backend type for gradient computation. + +Returns `AutoForwardDiff()` by default, meaning gradient getters use ForwardDiff +for differentiation unless specified otherwise. + +# Returns +- `ADTypes.AutoForwardDiff`: the default automatic differentiation backend. + +See also: [`CTBase.Differentiation.DifferentiationInterface`](@ref). +""" +__ad_backend() = ADTypes.AutoForwardDiff() diff --git a/src/Differentiation/differentiation_interface.jl b/src/Differentiation/differentiation_interface.jl new file mode 100644 index 00000000..a63efc35 --- /dev/null +++ b/src/Differentiation/differentiation_interface.jl @@ -0,0 +1,104 @@ +# ============================================================================== +# DifferentiationInterface — Concrete AD Backend Strategy +# ============================================================================== + +""" +$(TYPEDEF) + +Concrete AD backend strategy wrapping DifferentiationInterface.jl backends. + +`DifferentiationInterface` stores a `:ad_backend` option (e.g., `AutoForwardDiff()`, +`AutoZygote()`) and uses it to compute gradients via the DifferentiationInterface.jl +ecosystem. + +# Arguments +- `backend=AutoForwardDiff()`: The DifferentiationInterface.jl backend to use. Defaults + to `AutoForwardDiff()` from `ADTypes.jl` (a hard dependency). +- `kwargs...`: Additional options passed to `StrategyOptions`. + +# Notes + - `ADTypes.jl` is a hard dependency, so `AutoForwardDiff()` is always available in core. + - Gradient computation requires the `CTBaseDifferentiationInterface` extension. + - Without the extension, the gradient methods throw `NotImplemented` with a helpful message. + +See also: [`CTBase.Differentiation.AbstractADBackend`](@ref), +[`CTBase.Differentiation.hamiltonian_gradient`](@ref), +[`CTBase.Differentiation.variable_gradient`](@ref). +""" +struct DifferentiationInterface{O<:Strategies.StrategyOptions} <: + AbstractADBackend + options::O +end + +""" +$(TYPEDSIGNATURES) + +Constructor for `DifferentiationInterface` with a specific backend. + +# Arguments +- `backend=AutoForwardDiff()`: The DifferentiationInterface.jl backend. +- `kwargs...`: Additional options passed to `StrategyOptions`. + +# Returns +- `DifferentiationInterface`: A new backend strategy instance. +""" +function DifferentiationInterface(; mode::Symbol=:strict, kwargs...) + opts = Strategies.build_strategy_options(DifferentiationInterface; mode=mode, kwargs...) + return DifferentiationInterface{typeof(opts)}(opts) +end + +# ============================================================================== +# CTBase.Strategies Contract +# ============================================================================== + +""" +$(TYPEDSIGNATURES) + +Return the strategy identifier for `DifferentiationInterface`. +""" +Strategies.id(::Type{<:DifferentiationInterface}) = :di + +""" +$(TYPEDSIGNATURES) + +Return a human-readable description of the `DifferentiationInterface` strategy. +""" +Strategies.description(::Type{<:DifferentiationInterface}) = + "AD backend wrapping DifferentiationInterface.jl backends (e.g., AutoForwardDiff)." + +""" +$(TYPEDSIGNATURES) + +Return metadata defining `DifferentiationInterface` options and their specifications. +""" +function Strategies.metadata(::Type{<:DifferentiationInterface}) + return Strategies.StrategyMetadata( + Strategies.OptionDefinition(; + name = :ad_backend, + type = ADTypes.AbstractADType, + default = __ad_backend(), + description = "DifferentiationInterface.jl backend (e.g. AutoForwardDiff()).", + aliases=(:backend, :ad), + ), + ) +end + +""" +$(TYPEDSIGNATURES) + +Extract the AD backend from a `DifferentiationInterface` strategy. + +# Arguments +- `backend::DifferentiationInterface`: The AD backend. + +# Returns +- `ADTypes.AbstractADType`: The concrete AD backend from the `:ad_backend` option. + +# Notes + - This extracts the `:ad_backend` option from the strategy's options. + - Used by Hamiltonian vector-field getters to delegate to the AD-backed getter. + +See also: [`CTBase.Differentiation.ad_backend`](@ref). +""" +ad_backend(backend::DifferentiationInterface) = + Base.get(Strategies.options(backend), Val(:ad_backend)) From 13ad0a5069170470227665f0fb502a4e8ca668b1 Mon Sep 17 00:00:00 2001 From: Olivier Cots Date: Thu, 25 Jun 2026 17:03:35 +0200 Subject: [PATCH 2/4] test: add CTBase.Differentiation suite + DI extension test Co-Authored-By: Claude Opus 4.8 --- test/suite/differentiation/test_ad_backend.jl | 136 ++++++++++++++++ .../differentiation/test_arg_placement.jl | 148 ++++++++++++++++++ .../test_differentiation_module.jl | 147 +++++++++++++++++ ...est_differentiation_interface_extension.jl | 71 +++++++++ 4 files changed, 502 insertions(+) create mode 100644 test/suite/differentiation/test_ad_backend.jl create mode 100644 test/suite/differentiation/test_arg_placement.jl create mode 100644 test/suite/differentiation/test_differentiation_module.jl create mode 100644 test/suite/extensions/test_differentiation_interface_extension.jl diff --git a/test/suite/differentiation/test_ad_backend.jl b/test/suite/differentiation/test_ad_backend.jl new file mode 100644 index 00000000..70987ffc --- /dev/null +++ b/test/suite/differentiation/test_ad_backend.jl @@ -0,0 +1,136 @@ +""" +Unit and error tests for the AD backend contract and DifferentiationInterface strategy. +""" + +module TestADBackend + +import Test +import CTBase.Differentiation +import CTBase.Data +import CTBase.Traits +import CTBase.Exceptions +import CTBase.Strategies +import ADTypes + +const VERBOSE = isdefined(Main, :TestData) ? Main.TestData.VERBOSE : true +const SHOWTIMING = isdefined(Main, :TestData) ? Main.TestData.SHOWTIMING : true + +# ============================================================================== +# Fake AD Backend for Testing (at module top-level) +# ============================================================================== + +struct FakeADBackend <: Differentiation.AbstractADBackend + options::Strategies.StrategyOptions +end + +FakeADBackend() = FakeADBackend(Strategies.StrategyOptions()) + +# ============================================================================== +# Fake Hamiltonian for Testing (at module top-level) +# ============================================================================== + +struct FakeHamiltonian <: Data.AbstractHamiltonian{Traits.Autonomous, Traits.Fixed} + f::Function +end + +FakeHamiltonian() = FakeHamiltonian((x, p) -> 0.0) + +function test_ad_backend() + Test.@testset "AD Backend Tests" verbose=VERBOSE showtiming=SHOWTIMING begin + + # ============================================================================== + # Unit Tests + # ============================================================================== + + Test.@testset "Unit: AbstractADBackend abstract type" begin + backend = FakeADBackend() + Test.@test backend isa Differentiation.AbstractADBackend + Test.@test backend isa Strategies.AbstractStrategy + end + + Test.@testset "Unit: DifferentiationInterface construction" begin + # Default construction + di = Differentiation.DifferentiationInterface() + Test.@test di isa Differentiation.DifferentiationInterface + Test.@test di isa Differentiation.AbstractADBackend + + # Default backend is AutoForwardDiff + metadata = Strategies.metadata(Differentiation.DifferentiationInterface) + Test.@test metadata[:ad_backend].default === ADTypes.AutoForwardDiff() + + # Custom backend + di_custom = Differentiation.DifferentiationInterface(backend=ADTypes.AutoZygote()) + Test.@test di_custom isa Differentiation.DifferentiationInterface + end + + Test.@testset "Unit: CTBase.Strategies contract" begin + # id + Test.@test Strategies.id(Differentiation.DifferentiationInterface) === + :di + + # description + desc = Strategies.description(Differentiation.DifferentiationInterface) + Test.@test desc isa String + Test.@test !isempty(desc) + + # metadata + metadata = Strategies.metadata(Differentiation.DifferentiationInterface) + Test.@test metadata isa Strategies.StrategyMetadata + Test.@test length(metadata) > 0 + Test.@test :ad_backend in keys(metadata) + end + + Test.@testset "Unit: build_ad_backend" begin + backend = Differentiation.build_ad_backend() + Test.@test backend isa Differentiation.DifferentiationInterface + end + + Test.@testset "Unit: __ad_backend default" begin + Test.@test Differentiation.__ad_backend() === ADTypes.AutoForwardDiff() + end + + # ============================================================================== + # Error Tests + # ============================================================================== + + Test.@testset "Error: hamiltonian_gradient stub throws NotImplemented" begin + backend = FakeADBackend() + h = FakeHamiltonian() + t = 0.0 + x = [1.0, 2.0] + p = [3.0, 4.0] + v = 5.0 + + try + Differentiation.hamiltonian_gradient(backend, h, t, x, p, v) + Test.@test false # Should not reach here + catch err + Test.@test err isa Exceptions.NotImplemented + Test.@test occursin("hamiltonian_gradient", err.required_method) + end + end + + Test.@testset "Error: variable_gradient stub throws NotImplemented" begin + backend = FakeADBackend() + h = FakeHamiltonian() + t = 0.0 + x = [1.0, 2.0] + p = [3.0, 4.0] + v = 5.0 + + try + Differentiation.variable_gradient(backend, h, t, x, p, v) + Test.@test false # Should not reach here + catch err + Test.@test err isa Exceptions.NotImplemented + Test.@test occursin("variable_gradient", err.required_method) + end + end + + end +end + +end # module + +# CRITICAL: Redefine in outer scope for TestRunner +test_ad_backend() = TestADBackend.test_ad_backend() diff --git a/test/suite/differentiation/test_arg_placement.jl b/test/suite/differentiation/test_arg_placement.jl new file mode 100644 index 00000000..b8ca4079 --- /dev/null +++ b/test/suite/differentiation/test_arg_placement.jl @@ -0,0 +1,148 @@ +""" +Tests for Differentiation.differentiate and Differentiation.pushforward. +""" + +module TestArgPlacement + +import Test +import ForwardDiff # ensure DI ForwardDiff extension is loaded (AutoForwardDiff backend) +import DifferentiationInterface # activates CTBaseDifferentiationInterface extension +import CTBase.Differentiation +import CTBase.Exceptions +import CTBase.Strategies + +const VERBOSE = isdefined(Main, :TestData) ? Main.TestData.VERBOSE : true +const SHOWTIMING = isdefined(Main, :TestData) ? Main.TestData.SHOWTIMING : true + +# ============================================================================== +# Fake backend for stub error tests (at module top-level — world-age rule) +# ============================================================================== + +struct FakeBackendAP <: Differentiation.AbstractADBackend + options::Strategies.StrategyOptions +end +FakeBackendAP() = FakeBackendAP(Strategies.StrategyOptions()) + +# ============================================================================== +# Helper — default backend (AutoForwardDiff via DI) +# ============================================================================== + +_default_backend() = Differentiation.build_ad_backend() + +function test_arg_placement() + Test.@testset "differentiate/pushforward" verbose=VERBOSE showtiming=SHOWTIMING begin + + # ====================================================================== + # Stub error tests + # ====================================================================== + + Test.@testset "Error: differentiate stub throws NotImplemented" begin + b = FakeBackendAP() + f(x) = x^2 + try + Differentiation.differentiate(b, f, Val(1), 2.0) + Test.@test false + catch err + Test.@test err isa Exceptions.NotImplemented + Test.@test occursin("differentiate", err.required_method) + end + end + + Test.@testset "Error: pushforward stub throws NotImplemented" begin + b = FakeBackendAP() + f(x) = x .^ 2 + try + Differentiation.pushforward(b, f, Val(1), [1.0, 2.0], [1.0, 0.0]) + Test.@test false + catch err + Test.@test err isa Exceptions.NotImplemented + Test.@test occursin("pushforward", err.required_method) + end + end + + # ====================================================================== + # Differentiation.differentiate — with DI backend + # ====================================================================== + + Test.@testset "differentiate: gradient wrt slot (array active)" begin + b = _default_backend() + # H(x, p) = ½‖p‖² + ‖x‖² → ∂H/∂x = 2x, ∂H/∂p = p + H(x, p) = 0.5 * sum(p .^ 2) + sum(x .^ 2) + x = [1.0, 2.0] + p = [3.0, 4.0] + gx = Differentiation.differentiate(b, H, Val(1), x, p) + gp = Differentiation.differentiate(b, H, Val(2), p, x) + Test.@test gx ≈ 2x + Test.@test gp ≈ p + end + + Test.@testset "differentiate: derivative wrt slot (scalar active)" begin + b = _default_backend() + # f(t, x) = t * x[1] → ∂f/∂t = x[1] + f(t, x) = t * x[1] + x = [3.0, 4.0] + t = 2.0 + dt = Differentiation.differentiate(b, f, Val(1), t, x) + Test.@test dt ≈ x[1] + end + + Test.@testset "differentiate: 4-arg Hamiltonian (uniform signature)" begin + b = _default_backend() + # H(t, x, p, v) = t*v + ½‖p‖² + ‖x‖² + H(t, x, p, v) = t * v + 0.5 * sum(p .^ 2) + sum(x .^ 2) + t = 1.0; x = [1.0, 2.0]; p = [3.0, 4.0]; v = 5.0 + # ∂H/∂x = 2x (slot 2), consts in order t,p,v + gx = Differentiation.differentiate(b, H, Val(2), x, t, p, v) + # ∂H/∂p = p (slot 3), consts in order t,x,v + gp = Differentiation.differentiate(b, H, Val(3), p, t, x, v) + # ∂H/∂v = t (slot 4, scalar), consts in order t,x,p + dv = Differentiation.differentiate(b, H, Val(4), v, t, x, p) + Test.@test gx ≈ 2x + Test.@test gp ≈ p + Test.@test dv ≈ t + end + + # ====================================================================== + # Differentiation.pushforward — with DI backend + # ====================================================================== + + Test.@testset "pushforward: linear vector field (array x)" begin + b = _default_backend() + # X(x) = A*x, J_X = A → pushforward(X, x, d) = A*d + A = [0.0 1.0; -1.0 0.0] + X(x) = A * x + x = [1.0, 2.0] + d = [5.0, 6.0] + jvp = Differentiation.pushforward(b, X, Val(1), x, d) + Test.@test jvp ≈ A * d + end + + Test.@testset "pushforward: scalar function (directional derivative)" begin + b = _default_backend() + # f(x) = ‖x‖², ∇f = 2x → JVP = 2x⋅d + f(x) = sum(x .^ 2) + x = [1.0, 2.0] + d = [3.0, 4.0] + jvp = Differentiation.pushforward(b, f, Val(1), x, d) + Test.@test jvp ≈ 2 * sum(x .* d) + end + + Test.@testset "pushforward: with frozen constant arguments" begin + b = _default_backend() + # foo(x, v) = v * x, active = x (slot 1), const = v + # JVP: d(v*x)/dx · d = v * d + foo(x, v) = v .* x + x = [1.0, 2.0] + d = [1.0, 0.0] + v = 3.0 + jvp = Differentiation.pushforward(b, foo, Val(1), x, d, v) + Test.@test jvp ≈ v .* d + end + + end +end + +end # module + +# CRITICAL: Redefine in outer scope for TestRunner +test_arg_placement() = TestArgPlacement.test_arg_placement() diff --git a/test/suite/differentiation/test_differentiation_module.jl b/test/suite/differentiation/test_differentiation_module.jl new file mode 100644 index 00000000..df46fcd9 --- /dev/null +++ b/test/suite/differentiation/test_differentiation_module.jl @@ -0,0 +1,147 @@ +""" +# ============================================================================ +# Differentiation Module Exports Tests +# ============================================================================ +# This file tests the exports from the `Differentiation` module. It verifies that +# the expected types and functions are properly exported by +# `CTBase.Differentiation` and readily accessible to the end user. +# +# Functionality tests are in separate files: +# - test_ad_backend.jl for AD backend functionality +# - test_arg_placement.jl for differentiate/pushforward primitives +""" + +module TestDifferentiationModule + +import Test +import CTBase +import CTBase.Differentiation +using CTBase.Differentiation # For testing exported symbols + +const VERBOSE = isdefined(Main, :TestData) ? Main.TestData.VERBOSE : true +const SHOWTIMING = isdefined(Main, :TestData) ? Main.TestData.SHOWTIMING : true + +const CurrentModule = TestDifferentiationModule + +# ============================================================================ +# Hardcoded export lists +# ============================================================================ +# These lists define the expected public API of the Differentiation module. + +const EXPORTED_ABSTRACT_TYPES = ( + :AbstractADBackend, +) + +const EXPORTED_CONCRETE_TYPES = ( + :DifferentiationInterface, +) + +const EXPORTED_FUNCTIONS = ( + :build_ad_backend, + :ad_backend, + :hamiltonian_gradient, + :variable_gradient, + :gradient, + :derivative, + :differentiate, + :pushforward, +) + +# Note: Differentiation module has no private symbols (after filtering Julia internals) +# All symbols are exported + +# ============================================================================ +# Helper functions (generic for reuse in other modules) +# ============================================================================ + +""" + test_exported_symbols(module_ref::Module, symbols::Tuple, test_module::Module) + +Test that symbols are exported from a module and available via `using`. +""" +function test_exported_symbols(module_ref::Module, symbols::Tuple, test_module::Module) + for sym in symbols + Test.@testset "$(sym)" begin + Test.@test isdefined(module_ref, sym) + Test.@test isdefined(test_module, sym) + end + end +end + +""" + test_internal_symbols(module_ref::Module, symbols::Tuple, test_module::Module) + +Test that symbols are defined in a module but NOT exported (not available via `using`). +Generic helper for modules with private symbols. +""" +function test_internal_symbols(module_ref::Module, symbols::Tuple, test_module::Module) + for sym in symbols + Test.@testset "$(sym)" begin + Test.@test isdefined(module_ref, sym) + Test.@test !isdefined(test_module, sym) + end + end +end + +# ============================================================================ +# Test function +# ============================================================================ + +function test_differentiation_module() + Test.@testset "Differentiation Module Exports" verbose=VERBOSE showtiming=SHOWTIMING begin + + # ==================================================================== + # Module availability + # ==================================================================== + + Test.@testset "Module availability" begin + Test.@testset "Differentiation module exists" begin + Test.@test isdefined(CTBase, :Differentiation) + Test.@test CTBase.Differentiation isa Module + end + end + + # ==================================================================== + # Exported abstract types verification + # ==================================================================== + + Test.@testset "Exported abstract types" begin + test_exported_symbols(Differentiation, EXPORTED_ABSTRACT_TYPES, CurrentModule) + end + + # ==================================================================== + # Exported concrete types verification + # ==================================================================== + + Test.@testset "Exported concrete types" begin + test_exported_symbols(Differentiation, EXPORTED_CONCRETE_TYPES, CurrentModule) + end + + # ==================================================================== + # Exported functions verification + # ==================================================================== + + Test.@testset "Exported functions" begin + test_exported_symbols(Differentiation, EXPORTED_FUNCTIONS, CurrentModule) + end + + # ==================================================================== + # Type hierarchy tests + # ==================================================================== + + Test.@testset "Type hierarchy" begin + Test.@testset "Abstract types are abstract" begin + Test.@test isabstracttype(Differentiation.AbstractADBackend) + end + + Test.@testset "Concrete types inherit from abstract types" begin + Test.@test Differentiation.DifferentiationInterface <: Differentiation.AbstractADBackend + end + end + end +end + +end # module + +# CRITICAL: Redefine in outer scope for TestRunner +test_differentiation_module() = TestDifferentiationModule.test_differentiation_module() diff --git a/test/suite/extensions/test_differentiation_interface_extension.jl b/test/suite/extensions/test_differentiation_interface_extension.jl new file mode 100644 index 00000000..34621fa0 --- /dev/null +++ b/test/suite/extensions/test_differentiation_interface_extension.jl @@ -0,0 +1,71 @@ +""" +Unit and integration tests for the CTBaseDifferentiationInterface extension. +""" + +module TestDifferentiationInterfaceExtension + +import Test +import ForwardDiff # ensure DI ForwardDiff extension is loaded (AutoForwardDiff backend) +import CTBase: CTBase +import CTBase.Data: Data +import CTBase.Differentiation: Differentiation +import ADTypes +import DifferentiationInterface + +const VERBOSE = isdefined(Main, :TestData) ? Main.TestData.VERBOSE : true +const SHOWTIMING = isdefined(Main, :TestData) ? Main.TestData.SHOWTIMING : true + +# ============================================================================== +# Fake Hamiltonians for testing (module top-level) +# ============================================================================== + +# H(t, x, p) = 0.5*(‖x‖² + ‖p‖²) → ∂H/∂x = x, ∂H/∂p = p +const FAKE_HAMILTONIAN_FIXED = Data.Hamiltonian( + (x, p) -> 0.5 * (sum(abs2, x) + sum(abs2, p)); + is_autonomous=true, is_variable=false) + +# H(t, x, p, v) = 0.5*(‖x‖² + ‖p‖² + ‖v‖²) → ∂H/∂v = v +const FAKE_HAMILTONIAN_NONFIXED = Data.Hamiltonian( + (x, p, v) -> 0.5 * (sum(abs2, x) + sum(abs2, p) + sum(abs2, v)); + is_autonomous=true, is_variable=true) + +# ============================================================================== +# Test function +# ============================================================================== + +function test_differentiation_interface_extension() + Test.@testset "DifferentiationInterface Extension Tests" verbose=VERBOSE showtiming=SHOWTIMING begin + + Test.@testset "Integration: hamiltonian_gradient (Fixed)" begin + backend = Differentiation.DifferentiationInterface(; ad_backend=ADTypes.AutoForwardDiff()) + t = 0.0 + x = [1.0, 2.0] + p = [3.0, 4.0] + v = nothing + + grad_x, grad_p = Differentiation.hamiltonian_gradient(backend, FAKE_HAMILTONIAN_FIXED, t, x, p, v) + Test.@test grad_x isa AbstractVector + Test.@test grad_x ≈ x atol=1e-8 + Test.@test grad_p isa AbstractVector + Test.@test grad_p ≈ p atol=1e-8 + end + + Test.@testset "Integration: variable_gradient (NonFixed)" begin + backend = Differentiation.DifferentiationInterface(; ad_backend=ADTypes.AutoForwardDiff()) + t = 0.0 + x = [1.0, 2.0] + p = [3.0, 4.0] + v = [5.0, 6.0] + + grad_v = Differentiation.variable_gradient(backend, FAKE_HAMILTONIAN_NONFIXED, t, x, p, v) + Test.@test grad_v isa AbstractVector + Test.@test grad_v ≈ v atol=1e-8 + end + + end +end + +end # module + +# CRITICAL: Redefine in outer scope for TestRunner +test_differentiation_interface_extension() = TestDifferentiationInterfaceExtension.test_differentiation_interface_extension() From 9c8bcbd3e4614a967f528385dd2d39507edc4ab6 Mon Sep 17 00:00:00 2001 From: Olivier Cots Date: Thu, 25 Jun 2026 17:03:35 +0200 Subject: [PATCH 3/4] docs: add Differentiation guide + API reference Strategy-style guide (based on the Options/Strategies guides), api_reference entry for the module + the CTBaseDifferentiationInterface extension, and the make.jl page. Docs env gains ADTypes/DifferentiationInterface/ForwardDiff for runnable examples. Co-Authored-By: Claude Opus 4.8 --- docs/Project.toml | 7 ++ docs/api_reference.jl | 9 ++ docs/make.jl | 1 + docs/src/guide/differentiation.md | 185 ++++++++++++++++++++++++++++++ 4 files changed, 202 insertions(+) create mode 100644 docs/src/guide/differentiation.md diff --git a/docs/Project.toml b/docs/Project.toml index eb28eb76..683a108c 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,14 +1,21 @@ [deps] +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +CTBase = "54762871-cc72-4466-b8e8-f6c8b58076cd" Coverage = "a2441757-f6aa-5fb2-8edb-039e3f45d037" +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" DocumenterVitepress = "4710194d-e776-4893-9690-8d956a29c365" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" LiveServer = "16fef848-5104-11e9-1b77-fb7a48bbb589" MarkdownAST = "d0879d2d-cac2-40c8-9cee-1863dc0c7391" [compat] +ADTypes = "1" Coverage = "1" +DifferentiationInterface = "0.7" Documenter = "1" DocumenterVitepress = "0.3" +ForwardDiff = "0.10, 1" LiveServer = "1" MarkdownAST = "0.1" julia = "1.10" diff --git a/docs/api_reference.jl b/docs/api_reference.jl index b87cea63..d66e8d5e 100644 --- a/docs/api_reference.jl +++ b/docs/api_reference.jl @@ -32,6 +32,12 @@ function generate_api_reference(src_dir::String) joinpath("Data", "abstract_hamiltonian_vector_field.jl"), joinpath("Data", "hamiltonian_vector_field.jl"), )), + (mod=CTBase.Differentiation, title="Differentiation", filename="differentiation", files=src( + joinpath("Differentiation", "Differentiation.jl"), joinpath("Differentiation", "default.jl"), + joinpath("Differentiation", "abstract_ad_backend.jl"), + joinpath("Differentiation", "differentiation_interface.jl"), + joinpath("Differentiation", "building.jl"), + )), (mod=CTBase.Descriptions, title="Descriptions", filename="descriptions", files=src( joinpath("Descriptions", "Descriptions.jl"), joinpath("Descriptions", "types.jl"), joinpath("Descriptions", "similarity.jl"), joinpath("Descriptions", "display.jl"), @@ -106,6 +112,9 @@ function generate_api_reference(src_dir::String) internals_modules = Any[cfg.mod => cfg.files for cfg in modules_config] for (sym, files) in [ + (:CTBaseDifferentiationInterface, ext( + joinpath("CTBaseDifferentiationInterface.jl"), + )), (:DocumenterReference, ext( joinpath("DocumenterReference", "DocumenterReference.jl"), joinpath("DocumenterReference", "config_helpers.jl"), diff --git a/docs/make.jl b/docs/make.jl index 3817c9b0..7d30c6a1 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -80,6 +80,7 @@ with_api_reference(src_dir) do api_pages "Implementing a Strategy" => joinpath("guide", "implementing-a-strategy.md"), "Strategy Parameters" => joinpath("guide", "strategy-parameters.md"), "Orchestration & Routing" => joinpath("guide", "orchestration-and-routing.md"), + "Differentiation" => joinpath("guide", "differentiation.md"), ], "Developer Tools" => [ "Test Runner" => joinpath("guide", "test-runner.md"), diff --git a/docs/src/guide/differentiation.md b/docs/src/guide/differentiation.md new file mode 100644 index 00000000..53622bee --- /dev/null +++ b/docs/src/guide/differentiation.md @@ -0,0 +1,185 @@ +# Differentiation: AD backend strategies + +```@meta +CurrentModule = CTBase +``` + +The [`CTBase.Differentiation`](@ref CTBase.Differentiation) submodule provides an +**automatic differentiation (AD) backend** abstraction built on the +[Strategies](implementing-a-strategy.md) contract. It exposes a small set of +differentiation primitives — gradients, derivatives, partial derivatives and +Jacobian–vector products — behind a single strategy type, so that consumers +(Hamiltonian systems, flows, differential geometry) never depend on a concrete +AD package directly. + +!!! tip "Prerequisites" + Read the [Implementing a Strategy](@ref) and [Options System](@ref) guides + first: an AD backend is a strategy with a single `:ad_backend` option. + +```@setup diff +using CTBase +using CTBase.Differentiation +using CTBase.Data +using CTBase.Strategies +# Loading DifferentiationInterface (+ an AD package) activates the extension +# that provides the actual differentiation methods. +using DifferentiationInterface +using ForwardDiff +import ADTypes +``` + +## Overview + +The module is organised around a two-level strategy contract: + +| Level | Symbol | Role | +|---|---|---| +| Family | [`Differentiation.AbstractADBackend`](@ref CTBase.Differentiation.AbstractADBackend) | abstract strategy; defines the AD contract | +| Concrete | [`Differentiation.DifferentiationInterface`](@ref CTBase.Differentiation.DifferentiationInterface) | wraps a [DifferentiationInterface.jl](https://github.com/JuliaDiff/DifferentiationInterface.jl) backend | + +The concrete strategy stores a single option, `:ad_backend`, holding an +[ADTypes.jl](https://github.com/SciML/ADTypes.jl) backend (default +`AutoForwardDiff()`). `ADTypes` is a hard dependency of CTBase, so the default is +always available. + +!!! note "The differentiation methods live in an extension" + The contract methods (`gradient`, `derivative`, `differentiate`, + `pushforward`, `hamiltonian_gradient`, `variable_gradient`) are implemented in + the `CTBaseDifferentiationInterface` **package extension**. They become + available only once `DifferentiationInterface` *and* a concrete AD package + (e.g. `ForwardDiff`) are loaded. Without the extension, every contract method + throws a [`CTBase.Exceptions.NotImplemented`](@ref) with a helpful message. + +## Building a backend + +Use [`Differentiation.build_ad_backend`](@ref CTBase.Differentiation.build_ad_backend) +to construct a backend with the default AD type: + +```@example diff +backend = Differentiation.build_ad_backend() +``` + +Pass an explicit ADTypes backend through the `ad_backend` option (the aliases +`backend` and `ad` resolve to the same option): + +```@example diff +Differentiation.build_ad_backend(; ad_backend = ADTypes.AutoForwardDiff()) +``` + +The wrapped AD type is recovered with +[`Differentiation.ad_backend`](@ref CTBase.Differentiation.ad_backend): + +```@example diff +Differentiation.ad_backend(backend) +``` + +Because it is a regular strategy, its metadata is introspectable at the type level: + +```@example diff +Strategies.metadata(Differentiation.DifferentiationInterface) +``` + +## Differentiation primitives + +### Gradient and derivative + +[`Differentiation.gradient`](@ref CTBase.Differentiation.gradient) differentiates a +scalar function of a vector, and +[`Differentiation.derivative`](@ref CTBase.Differentiation.derivative) a scalar +function of a scalar: + +```@example diff +f(x) = sum(x .^ 2) # ∇f = 2x +Differentiation.gradient(backend, f, [1.0, 2.0]) +``` + +```@example diff +g(t) = t^3 # g'(t) = 3t² +Differentiation.derivative(backend, g, 2.0) +``` + +### Partial derivative at a slot + +[`Differentiation.differentiate`](@ref CTBase.Differentiation.differentiate) +computes the partial derivative of `f` with respect to the argument at a +**compile-time slot** `Val(Slot)`, holding the other arguments fixed. The active +argument comes first, the frozen ones follow in slot order: + +```@example diff +# H(x, p) = ½‖p‖² + ‖x‖² → ∂H/∂x = 2x, ∂H/∂p = p +H(x, p) = 0.5 * sum(p .^ 2) + sum(x .^ 2) +x = [1.0, 2.0]; p = [3.0, 4.0] + +∂x = Differentiation.differentiate(backend, H, Val(1), x, p) # active = x (slot 1), const = p +∂p = Differentiation.differentiate(backend, H, Val(2), p, x) # active = p (slot 2), const = x +(∂x, ∂p) +``` + +A scalar slot yields a scalar derivative: + +```@example diff +# f(t, x) = t * x[1] → ∂f/∂t = x[1] +ft(t, x) = t * x[1] +Differentiation.differentiate(backend, ft, Val(1), 2.0, [3.0, 4.0]) +``` + +### Pushforward (Jacobian–vector product) + +[`Differentiation.pushforward`](@ref CTBase.Differentiation.pushforward) computes the +directional derivative of `f` at `x` along `dx`, freezing the other arguments: + +```@example diff +A = [0.0 1.0; -1.0 0.0] +X(x) = A * x # J_X = A → pushforward = A·dx +Differentiation.pushforward(backend, X, Val(1), [1.0, 2.0], [5.0, 6.0]) +``` + +## Hamiltonian gradients + +The two domain-specific methods operate directly on a +[`Data.Hamiltonian`](@ref CTBase.Data.Hamiltonian) (see the [Data](data.md) guide). +[`Differentiation.hamiltonian_gradient`](@ref CTBase.Differentiation.hamiltonian_gradient) +returns `(∂H/∂x, ∂H/∂p)`, **non-negated** — the caller applies the signs of +Hamilton's equations (`ẋ = ∂H/∂p`, `ṗ = -∂H/∂x`): + +```@example diff +# H(x, p) = ½(‖x‖² + ‖p‖²) → ∂H/∂x = x, ∂H/∂p = p +h = Data.Hamiltonian((x, p) -> 0.5 * (sum(abs2, x) + sum(abs2, p))) +∂x, ∂p = Differentiation.hamiltonian_gradient(backend, h, 0.0, [1.0, 2.0], [3.0, 4.0], nothing) +(∂x, ∂p) +``` + +For a non-fixed Hamiltonian, +[`Differentiation.variable_gradient`](@ref CTBase.Differentiation.variable_gradient) +returns `∂H/∂v`: + +```@example diff +# H(x, p, v) = ½(‖x‖² + ‖p‖² + ‖v‖²) → ∂H/∂v = v +hv = Data.Hamiltonian((x, p, v) -> 0.5 * (sum(abs2, x) + sum(abs2, p) + sum(abs2, v)); is_variable=true) +Differentiation.variable_gradient(backend, hv, 0.0, [1.0, 2.0], [3.0, 4.0], [5.0, 6.0]) +``` + +## The contract without the extension + +Every contract method has a fallback on `AbstractADBackend` that throws +[`CTBase.Exceptions.NotImplemented`](@ref) (see the [Exceptions](exceptions.md) +guide). This is what users hit when the AD package is not loaded, and what custom +backends must override: + +```@repl diff +struct MyBackend <: Differentiation.AbstractADBackend + options::Strategies.StrategyOptions +end +try # hide +Differentiation.gradient(MyBackend(Strategies.StrategyOptions()), x -> sum(x), [1.0]) +catch e # hide +showerror(IOContext(stdout, :color => false), e) # hide +end # hide +``` + +## See also + +- [`Differentiation.AbstractADBackend`](@ref CTBase.Differentiation.AbstractADBackend), [`Differentiation.DifferentiationInterface`](@ref CTBase.Differentiation.DifferentiationInterface) — the strategy types. +- [Implementing a Strategy](@ref), [Options System](@ref) — the contract these build on. +- [Data](data.md) — `Hamiltonian` wrappers consumed by the gradient methods. +- [Exceptions](exceptions.md) — `NotImplemented`, raised by the contract fallbacks. From 109acb123b068f6f11fb563434fa9caa6eecde42 Mon Sep 17 00:00:00 2001 From: Olivier Cots Date: Thu, 25 Jun 2026 20:57:16 +0200 Subject: [PATCH 4/4] docs+test: refine Differentiation docstrings, guide sync, gradient/derivative tests Docstrings: - AbstractADBackend: replace stale "two methods" note with the full contract - Differentiation module: list ad_backend among the contract methods - DifferentiationInterface ctor: document the mode kwarg and ad_backend aliases Docs: - differentiation.md: clarify ad_backend (core) vs the extension primitives - index.md / getting-started.md: list Options, Strategies, Orchestration, Differentiation, Interpolation submodules and the DI extension Tests: - cover Differentiation.gradient/derivative (real computation via DI extension) - add NotImplemented stub-error tests for gradient/derivative Co-Authored-By: Claude Opus 4.8 --- docs/src/getting-started.md | 9 +++++++- docs/src/guide/differentiation.md | 5 +++++ docs/src/index.md | 10 +++++++++ src/Differentiation/Differentiation.jl | 1 + src/Differentiation/abstract_ad_backend.jl | 5 ++++- .../differentiation_interface.jl | 6 +++-- test/suite/differentiation/test_ad_backend.jl | 22 +++++++++++++++++++ ...est_differentiation_interface_extension.jl | 20 +++++++++++++++++ 8 files changed, 74 insertions(+), 4 deletions(-) diff --git a/docs/src/getting-started.md b/docs/src/getting-started.md index 7199f1ba..e515714a 100644 --- a/docs/src/getting-started.md +++ b/docs/src/getting-started.md @@ -31,7 +31,8 @@ Three things to keep in mind: CTBase.Exceptions.NotImplemented ``` 2. **Submodule-first API.** The public API lives in named submodules - (`Core`, `Descriptions`, `Exceptions`, `Traits`, `Data`, `DevTools`, `Unicode`, …). + (`Core`, `Exceptions`, `Traits`, `Data`, `Descriptions`, `Options`, `Strategies`, + `Orchestration`, `Differentiation`, `Interpolation`, `DevTools`, `Unicode`, …). You can bring a submodule's exports into scope explicitly: ```julia using CTBase.Exceptions # brings IncorrectArgument, NotImplemented, … into scope @@ -40,6 +41,8 @@ Three things to keep in mind: 3. **Extension-backed features.** `run_tests`, `postprocess_coverage`, and `automatic_reference_documentation` require loading the matching weak dependency (`Test`, `Coverage`, `Documenter` respectively) before they become active. + Likewise, the differentiation primitives of `CTBase.Differentiation` become active + only once `DifferentiationInterface` and an AD package (e.g. `ForwardDiff`) are loaded. ## 5-Minute Walkthrough @@ -134,6 +137,10 @@ For more, see the **[Data guide](guide/data.md)**. | Compile-time traits and dispatch | [Traits](guide/traits.md) | | Trait-carrying vector fields and Hamiltonians | [Data](guide/data.md) | | Descriptions catalogue and completion | [Descriptions](guide/descriptions.md) | +| Option schema, validation, and aliases | [Options System](guide/options-system.md) | +| Strategy contract and registration | [Implementing a Strategy](guide/implementing-a-strategy.md) | +| Routing options to strategies | [Orchestration & Routing](guide/orchestration-and-routing.md) | +| AD backends and differentiation primitives | [Differentiation](guide/differentiation.md) | | Modular test runner setup | [Test Runner](guide/test-runner.md) | | Coverage report generation | [Coverage](guide/coverage.md) | | Auto-generated API reference | [API Documentation](guide/api-documentation.md) | diff --git a/docs/src/guide/differentiation.md b/docs/src/guide/differentiation.md index 53622bee..a5dc2088 100644 --- a/docs/src/guide/differentiation.md +++ b/docs/src/guide/differentiation.md @@ -42,6 +42,11 @@ The concrete strategy stores a single option, `:ad_backend`, holding an `AutoForwardDiff()`). `ADTypes` is a hard dependency of CTBase, so the default is always available. +The contract has seven methods. +[`Differentiation.ad_backend`](@ref CTBase.Differentiation.ad_backend) — the accessor +returning the wrapped ADTypes backend — is resolved in core and always available; the +six differentiation primitives below live in an extension. + !!! note "The differentiation methods live in an extension" The contract methods (`gradient`, `derivative`, `differentiate`, `pushforward`, `hamiltonian_gradient`, `variable_gradient`) are implemented in diff --git a/docs/src/index.md b/docs/src/index.md index 858ef2ef..c1377288 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -29,6 +29,11 @@ It provides the **base layer** shared by all packages: common types, structured | [`CTBase.Traits`](@ref) | Compile-time trait types for time dependence, variable dependence, mutability, and dynamics dispatch | | [`CTBase.Data`](@ref) | Trait-carrying function wrappers: `VectorField`, `Hamiltonian`, `HamiltonianVectorField` | | [`CTBase.Descriptions`](@ref) | Symbolic description tuples: catalogue management, pattern completion, similarity search | +| [`CTBase.Options`](@ref) | Generic option handling: provenance tracking, schema definition, validation, and aliases | +| [`CTBase.Strategies`](@ref) | Strategy contract, registry, building/validation, and metadata for pluggable algorithmic components | +| [`CTBase.Orchestration`](@ref) | Option routing and disambiguation between problem-level actions and strategies | +| [`CTBase.Differentiation`](@ref) | AD-backend strategies for gradients, derivatives, partial derivatives, and Jacobian–vector products | +| [`CTBase.Interpolation`](@ref) | Linear and piecewise-constant interpolation with flat extrapolation | | [`CTBase.DevTools`](@ref) | Developer tools with tag-based dispatch for `run_tests`, `postprocess_coverage`, and `automatic_reference_documentation` | | [`CTBase.Unicode`](@ref) | Unicode subscript/superscript helpers for display | @@ -39,6 +44,11 @@ It provides the **base layer** shared by all packages: common types, structured - **[Traits](guide/traits.md)** — compile-time trait types, the opt-in contract, and predicate functions. - **[Data](guide/data.md)** — trait-carrying wrappers for vector fields and Hamiltonians. - **[Descriptions](guide/descriptions.md)** — catalogue API, pattern matching, error handling. +- **[Options System](guide/options-system.md)** — option schema, validation, aliases, and provenance. +- **[Implementing a Strategy](guide/implementing-a-strategy.md)** — the strategy contract and how to add one. +- **[Strategy Parameters](guide/strategy-parameters.md)** — declaring and resolving strategy options. +- **[Orchestration & Routing](guide/orchestration-and-routing.md)** — routing options to strategies with disambiguation. +- **[Differentiation](guide/differentiation.md)** — AD-backend strategies and differentiation primitives (extension-backed). - **[Test Runner](guide/test-runner.md)** — modular test infrastructure with `CTBase.DevTools.run_tests`. - **[Coverage](guide/coverage.md)** — post-processing coverage artifacts with `CTBase.postprocess_coverage`. - **[API Documentation](guide/api-documentation.md)** — auto-generating per-module API pages. diff --git a/src/Differentiation/Differentiation.jl b/src/Differentiation/Differentiation.jl index 35883b71..7dec2e70 100644 --- a/src/Differentiation/Differentiation.jl +++ b/src/Differentiation/Differentiation.jl @@ -9,6 +9,7 @@ for computing gradients of scalar Hamiltonian functions. ## Architecture The module defines an abstract contract `AbstractADBackend` with the following methods: +- `ad_backend(backend)` → the wrapped ADTypes backend (resolved in core) - `hamiltonian_gradient(backend, h, t, x, p, v)` → (∂H/∂x, ∂H/∂p) - `variable_gradient(backend, h, t, x, p, v)` → ∂H/∂v - `gradient(backend, f, x)` → ∇f (extension contract) diff --git a/src/Differentiation/abstract_ad_backend.jl b/src/Differentiation/abstract_ad_backend.jl index f19d6b5c..78a88daa 100644 --- a/src/Differentiation/abstract_ad_backend.jl +++ b/src/Differentiation/abstract_ad_backend.jl @@ -14,7 +14,10 @@ implement the contract methods to provide actual gradient computation. # Notes - `AbstractADBackend` subtypes `CTBase.Strategies.AbstractStrategy` — they are first-class strategies in the CTBase.Strategies ecosystem. - - The contract consists of two methods: `hamiltonian_gradient` and `variable_gradient`. + - The contract comprises `ad_backend` (the wrapped ADTypes backend), the domain + methods `hamiltonian_gradient` and `variable_gradient`, and the generic primitives + `gradient`, `derivative`, `differentiate` and `pushforward`. All but `ad_backend` + are supplied by the `CTBaseDifferentiationInterface` extension. - Gradient methods return **non-negated** partial derivatives; the RHS closures apply the signs (ṗ = -∂H/∂x, ṽ = -∂H/∂v). diff --git a/src/Differentiation/differentiation_interface.jl b/src/Differentiation/differentiation_interface.jl index a63efc35..6ad554ce 100644 --- a/src/Differentiation/differentiation_interface.jl +++ b/src/Differentiation/differentiation_interface.jl @@ -36,8 +36,10 @@ $(TYPEDSIGNATURES) Constructor for `DifferentiationInterface` with a specific backend. # Arguments -- `backend=AutoForwardDiff()`: The DifferentiationInterface.jl backend. -- `kwargs...`: Additional options passed to `StrategyOptions`. +- `mode::Symbol=:strict`: Validation mode forwarded to `build_strategy_options`. +- `kwargs...`: Options passed to `StrategyOptions`. The AD backend is set through the + `:ad_backend` option (aliases `backend` and `ad`), e.g. `backend=AutoForwardDiff()`; + it defaults to `AutoForwardDiff()`. # Returns - `DifferentiationInterface`: A new backend strategy instance. diff --git a/test/suite/differentiation/test_ad_backend.jl b/test/suite/differentiation/test_ad_backend.jl index 70987ffc..d0d8007c 100644 --- a/test/suite/differentiation/test_ad_backend.jl +++ b/test/suite/differentiation/test_ad_backend.jl @@ -127,6 +127,28 @@ function test_ad_backend() end end + Test.@testset "Error: gradient stub throws NotImplemented" begin + backend = FakeADBackend() + try + Differentiation.gradient(backend, x -> sum(x), [1.0, 2.0]) + Test.@test false # Should not reach here + catch err + Test.@test err isa Exceptions.NotImplemented + Test.@test occursin("gradient", err.required_method) + end + end + + Test.@testset "Error: derivative stub throws NotImplemented" begin + backend = FakeADBackend() + try + Differentiation.derivative(backend, t -> t^2, 1.0) + Test.@test false # Should not reach here + catch err + Test.@test err isa Exceptions.NotImplemented + Test.@test occursin("derivative", err.required_method) + end + end + end end diff --git a/test/suite/extensions/test_differentiation_interface_extension.jl b/test/suite/extensions/test_differentiation_interface_extension.jl index 34621fa0..d48a415e 100644 --- a/test/suite/extensions/test_differentiation_interface_extension.jl +++ b/test/suite/extensions/test_differentiation_interface_extension.jl @@ -62,6 +62,26 @@ function test_differentiation_interface_extension() Test.@test grad_v ≈ v atol=1e-8 end + Test.@testset "Integration: gradient (vector input)" begin + backend = Differentiation.DifferentiationInterface(; ad_backend=ADTypes.AutoForwardDiff()) + f(x) = sum(x .^ 2) # ∇f = 2x + g = Differentiation.gradient(backend, f, [1.0, 2.0]) + Test.@test g isa AbstractVector + Test.@test g ≈ [2.0, 4.0] atol=1e-8 + end + + Test.@testset "Integration: gradient (scalar overload)" begin + backend = Differentiation.DifferentiationInterface(; ad_backend=ADTypes.AutoForwardDiff()) + f(x) = x^3 # f'(x) = 3x² + Test.@test Differentiation.gradient(backend, f, 2.0) ≈ 12.0 atol=1e-8 + end + + Test.@testset "Integration: derivative (scalar input)" begin + backend = Differentiation.DifferentiationInterface(; ad_backend=ADTypes.AutoForwardDiff()) + g(t) = t^3 # g'(t) = 3t² + Test.@test Differentiation.derivative(backend, g, 2.0) ≈ 12.0 atol=1e-8 + end + end end