diff --git a/build/ImplicitDifferentiation.jl b/build/ImplicitDifferentiation.jl deleted file mode 100644 index a8aadeb0..00000000 --- a/build/ImplicitDifferentiation.jl +++ /dev/null @@ -1,35 +0,0 @@ -""" - ImplicitDifferentiation - -A Julia package for automatic differentiation of implicit functions. - -Its main export is the type [`ImplicitFunction`](@ref). -""" -module ImplicitDifferentiation - -using ADTypes: AbstractADType -using DifferentiationInterface: - Constant, - jacobian, - prepare_jacobian, - prepare_pullback, - prepare_pullback_same_point, - prepare_pushforward, - prepare_pushforward_same_point, - pullback!, - pushforward! -using Krylov: gmres -using LinearOperators: LinearOperator -using LinearAlgebra: factorize - -include("settings.jl") -include("preparation.jl") -include("implicit_function.jl") -include("execution.jl") - -export KrylovLinearSolver -export MatrixRepresentation, OperatorRepresentation -export NoPreparation, ForwardPreparation, ReversePreparation, BothPreparation -export ImplicitFunction - -end diff --git a/build/execution.jl b/build/execution.jl deleted file mode 100644 index b34f64a4..00000000 --- a/build/execution.jl +++ /dev/null @@ -1,182 +0,0 @@ -const SYMMETRIC = false -const HERMITIAN = false - -struct JVP!{F,P,B,X,C} - f::F - prep::P - backend::B - x::X - contexts::C -end - -struct VJP!{F,P,B,X,C} - f::F - prep::P - backend::B - x::X - contexts::C -end - -function (po::JVP!)(res::AbstractVector, v::AbstractVector) - (; f, backend, x, contexts, prep) = po - pushforward!(f, (res,), prep, backend, x, (v,), contexts...) - return res -end - -function (po::VJP!)(res::AbstractVector, v::AbstractVector) - (; f, backend, x, contexts, prep) = po - pullback!(f, (res,), prep, backend, x, (v,), contexts...) - return res -end - -## A - -function build_A( - implicit::ImplicitFunction, - x::AbstractVector, - y::AbstractVector, - z, - args...; - suggested_backend::AbstractADType, -) - return build_A_aux( - implicit.representation, implicit, x, y, z, args...; suggested_backend - ) -end - -function build_A_aux(::MatrixRepresentation, implicit, x, y, z, args...; suggested_backend) - (; conditions, backend, prep_A) = implicit - actual_backend = isnothing(backend) ? suggested_backend : backend - contexts = (Constant(x), Constant(z), map(Constant, args)...) - A = jacobian(Switch12(conditions), prep_A..., actual_backend, y, contexts...) - return factorize(A) -end - -function build_A_aux( - ::OperatorRepresentation, implicit, x, y, z, args...; suggested_backend -) - (; conditions, backend, prep_A) = implicit - actual_backend = isnothing(backend) ? suggested_backend : backend - contexts = (Constant(x), Constant(z), map(Constant, args)...) - prep_A_same = prepare_pushforward_same_point( - Switch12(conditions), prep_A..., actual_backend, y, (zero(y),), contexts... - ) - prod! = JVP!(Switch12(conditions), prep_A_same, actual_backend, y, contexts) - return LinearOperator( - eltype(y), length(y), length(y), SYMMETRIC, HERMITIAN, prod!, typeof(y) - ) -end - -## Aᵀ - -function build_Aᵀ( - implicit::ImplicitFunction, - x::AbstractVector, - y::AbstractVector, - z, - args...; - suggested_backend::AbstractADType, -) - return build_Aᵀ_aux( - implicit.representation, implicit, x, y, z, args...; suggested_backend - ) -end - -function build_Aᵀ_aux(::MatrixRepresentation, implicit, x, y, z, args...; suggested_backend) - (; conditions, backend, prep_Aᵀ) = implicit - actual_backend = isnothing(backend) ? suggested_backend : backend - contexts = (Constant(x), Constant(z), map(Constant, args)...) - Aᵀ = transpose( - jacobian(Switch12(conditions), prep_Aᵀ..., actual_backend, y, contexts...) - ) - return factorize(Aᵀ) -end - -function build_Aᵀ_aux( - ::OperatorRepresentation, implicit, x, y, z, args...; suggested_backend -) - (; conditions, backend, prep_Aᵀ) = implicit - actual_backend = isnothing(backend) ? suggested_backend : backend - contexts = (Constant(x), Constant(z), map(Constant, args)...) - prep_Aᵀ_same = prepare_pullback_same_point( - Switch12(conditions), prep_Aᵀ..., actual_backend, y, (zero(y),), contexts... - ) - prod! = VJP!(Switch12(conditions), prep_Aᵀ_same, actual_backend, y, contexts) - return LinearOperator( - eltype(y), length(y), length(y), SYMMETRIC, HERMITIAN, prod!, typeof(y) - ) -end - -## B - -function build_B( - implicit::ImplicitFunction, - x::AbstractVector, - y::AbstractVector, - z, - args...; - suggested_backend::AbstractADType, -) - return build_B_aux( - implicit.representation, implicit, x, y, z, args...; suggested_backend - ) -end - -function build_B_aux(::MatrixRepresentation, implicit, x, y, z, args...; suggested_backend) - (; conditions, backend, prep_B) = implicit - actual_backend = isnothing(backend) ? suggested_backend : backend - contexts = (Constant(y), Constant(z), map(Constant, args)...) - return jacobian(conditions, prep_B..., actual_backend, x, contexts...) -end - -function build_B_aux( - ::OperatorRepresentation, implicit, x, y, z, args...; suggested_backend -) - (; conditions, backend, prep_B) = implicit - actual_backend = isnothing(backend) ? suggested_backend : backend - contexts = (Constant(y), Constant(z), map(Constant, args)...) - prep_B_same = prepare_pushforward_same_point( - conditions, prep_B..., actual_backend, x, (zero(x),), contexts... - ) - prod! = JVP!(conditions, prep_B_same, actual_backend, x, contexts) - return LinearOperator( - eltype(y), length(y), length(x), SYMMETRIC, HERMITIAN, prod!, typeof(x) - ) -end - -## Bᵀ - -function build_Bᵀ( - implicit::ImplicitFunction, - x::AbstractVector, - y::AbstractVector, - z, - args...; - suggested_backend::AbstractADType, -) - return build_Bᵀ_aux( - implicit.representation, implicit, x, y, z, args...; suggested_backend - ) -end - -function build_Bᵀ_aux(::MatrixRepresentation, implicit, x, y, z, args...; suggested_backend) - (; conditions, backend, prep_Bᵀ) = implicit - actual_backend = isnothing(backend) ? suggested_backend : backend - contexts = (Constant(y), Constant(z), map(Constant, args)...) - return transpose(jacobian(conditions, prep_Bᵀ..., actual_backend, x, contexts...)) -end - -function build_Bᵀ_aux( - ::OperatorRepresentation, implicit, x, y, z, args...; suggested_backend -) - (; conditions, backend, prep_Bᵀ) = implicit - actual_backend = isnothing(backend) ? suggested_backend : backend - contexts = (Constant(y), Constant(z), map(Constant, args)...) - prep_Bᵀ_same = prepare_pullback_same_point( - conditions, prep_Bᵀ..., actual_backend, x, (zero(y),), contexts... - ) - prod! = VJP!(conditions, prep_Bᵀ_same, actual_backend, x, contexts) - return LinearOperator( - eltype(y), length(x), length(y), SYMMETRIC, HERMITIAN, prod!, typeof(x) - ) -end diff --git a/build/implicit_function.jl b/build/implicit_function.jl deleted file mode 100644 index d0c41cba..00000000 --- a/build/implicit_function.jl +++ /dev/null @@ -1,132 +0,0 @@ -""" - ImplicitFunction - -Wrapper for an implicit function defined by a _solver_ and a set of _conditions_ which the solution satisfies. - -An `ImplicitFunction` object behaves like a function, with the following signature: - - y, z = (implicit::ImplicitFunction)(x, args...) - -The first output `y` is differentiable with respect to the first argument `x`, while the second output `z` (a byproduct of the solve) and the following positional arguments `args` are considered constant. - -When a derivative is queried, the Jacobian of `y(x)` is computed using the implicit function theorem applied to the conditions `c(x, y)` (we ignore `z` for concision): - - ∂₂c(x, y(x)) * ∂y(x) = -∂₁c(x, y(x)) - -This requires solving a linear system `A * J = -B` where `A = ∂₂c`, `B = ∂₁c` and `J = ∂y`. - -# Constructor - - ImplicitFunction( - solver, - conditions, - linear_solver=KrylovLinearSolver(), - backend=nothing, - representation=OperatorRepresentation(), - preparation=NoPreparation(), - input_example=nothing, - ) - -## Positional arguments - -- `solver`: a callable returning `(x, args...) -> (y, z)` where `z` is an arbitrary byproduct of the solve. Both `x` and `y` must be subtypes of `AbstractVector`, while `z` and `args` can be anything. -- `conditions`: a callable returning a vector of optimality conditions `(x, y, z, args...) -> c`, must be compatible with automatic differentiation - -## Keyword arguments - -- `linear_solver`: a callable to solve linear systems with two methods, one for `(A, b)` (single solve) and one for `(A, B)` (batched solve) (defaults to [`KrylovLinearSolver`](@ref)) -- `backend::AbstractADType`: either `nothing` or an object from [ADTypes.jl](https://github.com/SciML/ADTypes.jl) dictating how how the conditions will be differentiated -- `representation`: either [`MatrixRepresentation`](@ref) or [`OperatorRepresentation`](@ref) -- `preparation`: one of [`NoPreparation`](@ref), [`ForwardPreparation`](@ref), [`ReversePreparation`](@ref), or [`BothPreparation`](@ref) -- `input_example`: either `nothing` or a tuple `(x, args...)` used to prepare differentiation -""" -struct ImplicitFunction{ - F, - C, - L, - B<:Union{Nothing,AbstractADType}, - R<:AbstractRepresentation, - P<:AbstractPreparation, - PA, - PAT, - PB, - PBT, -} - solver::F - conditions::C - linear_solver::L - backend::B - representation::R - preparation::P - prep_A::PA - prep_Aᵀ::PAT - prep_B::PB - prep_Bᵀ::PBT -end - -function ImplicitFunction( - solver, - conditions; - linear_solver=KrylovLinearSolver(), - backend=nothing, - representation=OperatorRepresentation(), - preparation=NoPreparation(), - input_example=nothing, -) - if preparation isa NoPreparation || isnothing(backend) || isnothing(input_example) - prep_A = () - prep_Aᵀ = () - prep_B = () - prep_Bᵀ = () - else - x, args = first(input_example), Base.tail(input_example) - y, z = solver(x, args...) - if preparation isa Union{ForwardPreparation,BothPreparation} - prep_A = (prepare_A(representation, x, y, z, args...; conditions, backend),) - prep_B = (prepare_B(representation, x, y, z, args...; conditions, backend),) - else - prep_A = () - prep_B = () - end - if preparation isa Union{ReversePreparation,BothPreparation} - prep_Aᵀ = (prepare_Aᵀ(representation, x, y, z, args...; conditions, backend),) - prep_Bᵀ = (prepare_Bᵀ(representation, x, y, z, args...; conditions, backend),) - else - prep_Aᵀ = () - prep_Bᵀ = () - end - end - return ImplicitFunction( - solver, - conditions, - linear_solver, - backend, - representation, - preparation, - prep_A, - prep_Aᵀ, - prep_B, - prep_Bᵀ, - ) -end - -function Base.show(io::IO, implicit::ImplicitFunction) - (; solver, conditions, backend, linear_solver, representation, preparation) = implicit - return print( - io, - """ - ImplicitFunction( - $solver, - $conditions; - linear_solver=$linear_solver, - backend=$backend, - representation=$representation, - preparation=$preparation, - ) - """, - ) -end - -function (implicit::ImplicitFunction)(x::AbstractVector, args::Vararg{Any,N}) where {N} - return implicit.solver(x, args...) -end diff --git a/build/preparation.jl b/build/preparation.jl deleted file mode 100644 index 7ab1efe8..00000000 --- a/build/preparation.jl +++ /dev/null @@ -1,111 +0,0 @@ -struct Switch12{F} - f::F -end - -function (s12::Switch12)(arg1, arg2, other_args::Vararg{Any,N}) where {N} - return s12.f(arg2, arg1, other_args...) -end - -function prepare_A( - ::MatrixRepresentation, - x::AbstractVector, - y::AbstractVector, - z, - args...; - conditions, - backend::AbstractADType, -) - contexts = (Constant(x), Constant(z), map(Constant, args)...) - return prepare_jacobian(Switch12(conditions), backend, y, contexts...) -end - -function prepare_A( - ::OperatorRepresentation, - x::AbstractVector, - y::AbstractVector, - z, - args...; - conditions, - backend::AbstractADType, -) - contexts = (Constant(x), Constant(z), map(Constant, args)...) - return prepare_pushforward(Switch12(conditions), backend, y, (zero(y),), contexts...) -end - -function prepare_Aᵀ( - ::MatrixRepresentation, - x::AbstractVector, - y::AbstractVector, - z, - args...; - conditions, - backend::AbstractADType, -) - contexts = (Constant(x), Constant(z), map(Constant, args)...) - return prepare_jacobian(Switch12(conditions), backend, y, contexts...) -end - -function prepare_Aᵀ( - ::OperatorRepresentation, - x::AbstractVector, - y::AbstractVector, - z, - args...; - conditions, - backend::AbstractADType, -) - contexts = (Constant(x), Constant(z), map(Constant, args)...) - return prepare_pullback(Switch12(conditions), backend, y, (zero(y),), contexts...) -end - -function prepare_B( - ::MatrixRepresentation, - x::AbstractVector, - y::AbstractVector, - z, - args...; - conditions, - backend::AbstractADType, -) - contexts = (Constant(y), Constant(z), map(Constant, args)...) - return prepare_jacobian(conditions, backend, x, contexts...) -end - -function prepare_B( - ::OperatorRepresentation, - x::AbstractVector, - y::AbstractVector, - z, - args...; - conditions, - backend::AbstractADType, -) - contexts = (Constant(y), Constant(z), map(Constant, args)...) - return prepare_pushforward(conditions, backend, x, (zero(x),), contexts...) -end - -function prepare_Bᵀ( - ::MatrixRepresentation, - x::AbstractVector, - y::AbstractVector, - z, - args...; - conditions, - backend::AbstractADType, -) - contexts = (Constant(y), Constant(z), map(Constant, args)...) - return prepare_jacobian(conditions, backend, x, contexts...) -end - -function prepare_Bᵀ( - ::OperatorRepresentation, - x::AbstractVector, - y::AbstractVector, - z, - args...; - conditions, - backend::AbstractADType, -) - contexts = (Constant(y), Constant(z), map(Constant, args)...) - return prepare_pullback(conditions, backend, x, (zero(y),), contexts...) -end diff --git a/build/settings.jl b/build/settings.jl deleted file mode 100644 index 1ac91ce8..00000000 --- a/build/settings.jl +++ /dev/null @@ -1,105 +0,0 @@ -## Linear solver - -""" - KrylovLinearSolver - -Callable object that can solve linear systems `Ax = b` and `AX = B` in the same way as the built-in `\\`. -Uses an iterative solver from [Krylov.jl](https://github.com/JuliaSmoothOptimizers/Krylov.jl) under the hood. - -# Constructor - - KrylovLinearSolver(; verbose=true) - -If `verbose` is `true`, the solver logs a warning in case of failure. -Otherwise it will fail silently, and may return solutions that do not exactly satisfy the linear system. - -# Callable behavior - - (::KylovLinearSolver)(A, b::AbstractVector) - -Solve a linear system with a single right-hand side. - - (::KrylovLinearSolver)(A, B::AbstractMatrix) - -Solve a linear system with multiple right-hand sides. -""" -Base.@kwdef struct KrylovLinearSolver - verbose::Bool = true -end - -function (solver::KrylovLinearSolver)(A, b::AbstractVector) - x, stats = gmres(A, b) - if !stats.solved || stats.inconsistent - solver.verbose && - @warn "Failed to solve the linear system in the implicit function theorem with `Krylov.gmres`" stats - end - return x -end - -function (solver::KrylovLinearSolver)(A, B::AbstractMatrix) - # X, stats = block_gmres(A, B) # https://github.com/JuliaSmoothOptimizers/Krylov.jl/issues/854 - X = mapreduce(hcat, eachcol(B)) do b - solver(A, b) - end - return X -end - -## Representation - -abstract type AbstractRepresentation end - -""" - MatrixRepresentation - -Specify that the matrices involved in the implicit function theorem should be represented explicitly, with all their coefficients. - -# See also - -- [`ImplicitFunction`](@ref) -- [`OperatorRepresentation`](@ref) -""" -struct MatrixRepresentation <: AbstractRepresentation end - -""" - OperatorRepresentation - -Specify that the matrices involved in the implicit function theorem should be represented lazily, as linear operators from [LinearOperators.jl](https://github.com/JuliaSmoothOptimizers/LinearOperators.jl). - -# See also - -- [`ImplicitFunction`](@ref) -- [`MatrixRepresentation`](@ref) -""" -struct OperatorRepresentation <: AbstractRepresentation end - -## Preparation - -abstract type AbstractPreparation end - -""" - ForwardPreparation - -Specify that the derivatives of the conditions should be prepared for subsequent forward-mode differentiation of the implicit function. -""" -struct ForwardPreparation <: AbstractPreparation end - -""" - ReversePreparation - -Specify that the derivatives of the conditions should be prepared for subsequent reverse-mode differentiation of the implicit function. -""" -struct ReversePreparation <: AbstractPreparation end - -""" - BothPreparation - -Specify that the derivatives of the conditions should be prepared for subsequent forward- or reverse-mode differentiation of the implicit function. -""" -struct BothPreparation <: AbstractPreparation end - -""" - NoPreparation - -Specify that the derivatives of the conditions should not be prepared for subsequent differentiation of the implicit function. -""" -struct NoPreparation <: AbstractPreparation end