diff --git a/.gitignore b/.gitignore index 6a1febc..2f782e7 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,5 @@ Manifest.toml docs/build/ *html +*plan*md +.vscode/ diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..3442352 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,40 @@ +# Changelog + +## 0.7.0 + +Breaking renames (deprecation wrappers retained for 0.7.x; remove in 0.8): + +- `Mem` → `F1Cache` +- `initialize_mem(...)` → `F1Cache(...)` constructor +- `update_mem!` → `update_cache!` (internal) +- `Mem.A` field → `F1Cache.linear_cache` + +Migration: + +```julia +# 0.6.x +mem = F1Method.initialize_mem(F, ∇ₓf, x₀, p₀, alg) +F1Method.objective(f, F, mem, p, alg) +F1Method.gradient(f, F, ∇ₓf, mem, p, alg) +F1Method.hessian(f, F, ∇ₓf, mem, p, alg) + +# 0.7.0 +cache = F1Method.F1Cache(F, ∇ₓf, x₀, p₀, alg) +F1Method.objective(f, F, cache, p, alg) +F1Method.gradient(f, F, ∇ₓf, cache, p, alg) +F1Method.hessian(f, F, ∇ₓf, cache, p, alg) +``` + +`initialize_mem(...)` still works in 0.7.x with a deprecation warning; +it will be removed in 0.8. + +New: + +- Optional LinearSolve.jl backend, exposed via the new `linsolve` + kwarg on the `F1Cache` constructor. With it, + `F1Cache.linear_cache` holds a `LinearSolve.LinearCache`, enabling + symbolic factorization reuse for sparse direct solvers and buffer + reuse for dense. Default behaviour (`linsolve === nothing`) is + unchanged from 0.6.x. +- Tested with `LUFactorization`, `GenericLUFactorization`, + `UMFPACKFactorization`, `KLUFactorization`. diff --git a/Project.toml b/Project.toml index d5cbd28..b2de951 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "F1Method" uuid = "d5605bae-6d76-11e9-2fd7-71bcf42edbaa" authors = ["Benoit Pasquier "] -version = "0.6.1" +version = "0.7.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" @@ -11,9 +11,11 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" [weakdeps] +LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba" [extensions] +F1MethodLinearSolveExt = "LinearSolve" F1MethodOptimizationExt = "Optimization" [compat] @@ -25,6 +27,7 @@ DifferentiationInterface = "0.7" Distributions = "0.25" FiniteDiff = "2" ForwardDiff = "1" +LinearSolve = "3" Optimization = "4, 5.5" OptimizationOptimJL = "0.4" SciMLBase = "2, 3.9" diff --git a/README.md b/README.md index 7d5a496..ae480f4 100644 --- a/README.md +++ b/README.md @@ -135,12 +135,12 @@ Pick initial values for the state `x` and parameters `p`, initialize the cache, x₀, p₀ = [1.0, 2.0], [3.0, 4.0] # Initialize the cache for storing reusable objects -mem = F1Method.initialize_mem(F, ∇ₓf, x₀, p₀, MyAlg()) +cache = F1Method.F1Cache(F, ∇ₓf, x₀, p₀, MyAlg()) # Define the functions via the F1 method -objective(p) = F1Method.objective(f, F, mem, p, MyAlg()) -gradient(p) = F1Method.gradient(f, F, ∇ₓf, mem, p, MyAlg()) -hessian(p) = F1Method.hessian(f, F, ∇ₓf, mem, p, MyAlg()) +objective(p) = F1Method.objective(f, F, cache, p, MyAlg()) +gradient(p) = F1Method.gradient(f, F, ∇ₓf, cache, p, MyAlg()) +hessian(p) = F1Method.hessian(f, F, ∇ₓf, cache, p, MyAlg()) ``` You can now call `objective(p₀)`, `gradient(p₀)`, and `hessian(p₀)` directly: @@ -164,6 +164,25 @@ That's it. You were told it was simple, weren't you? Now you can test how fast and accurate it is! +## Custom linear solver + +By default F1Method uses Julia's `factorize` / `\` for the inner linear solves. +You can opt into [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) by passing a `linsolve` algorithm to the `F1Cache` constructor: + +```julia +using LinearSolve +cache = F1Method.F1Cache(F, ∇ₓf, x₀, p₀, MyAlg(); linsolve = UMFPACKFactorization()) +``` + +With a `linsolve` set, the `F1Cache.linear_cache` field holds a `LinearSolve.LinearCache` (it holds a plain `Factorization` otherwise). +Symbolic factorizations are reused across parameter updates (for sparse direct solvers), and the same factors are reused for the adjoint solve inside `hessian`. + +F1Method expects `F.jac(s, p)` to return either a `Matrix{Float64}` or `SparseMatrixCSC{Float64}`. +CI exercises `LUFactorization`, `GenericLUFactorization`, `UMFPACKFactorization`, and `KLUFactorization`. +Other LU-family algorithms whose `cacheval` is a Julia `Factorization` should also work (`MKLLUFactorization`, `AppleAccelerateLUFactorization`, `RFLUFactorization`, `SparspakFactorization`, `ParUFactorization`). + +Pardiso, MUMPS, CUDSS, STRUMPACK, and cusolverRF require per-solver adjoint shims and are not supported in 0.7. + ## Citing the software If you use this package, or implement your own package based on the F-1 algorithm please cite us. @@ -173,6 +192,7 @@ If you also use this package directly, please cite it! (Use [the Zenodo link](ht # Future This package is developed mainly for use with [AIBECS.jl](https://github.com/JuliaOcean/AIBECS.jl) and is likely not in its final form. -The API was last changed in v0.6: `gradient` now returns a `Vector` (previously a `1 × m` row matrix), parameter-side AD is routed through [DifferentiationInterface.jl](https://github.com/JuliaDiff/DifferentiationInterface.jl) (configurable via the `ad` kwarg on `initialize_mem`), and an `F1MethodOptimizationExt` weak-dep extension exposes `optimization_function(...)` for use with [Optimization.jl](https://github.com/SciML/Optimization.jl). +The API was last changed in v0.7: the `Mem` type was renamed to `F1Cache` and `initialize_mem(...)` to the `F1Cache(...)` constructor (deprecation wrappers retained for 0.7.x, removed in 0.8), and a new `linsolve` keyword on `F1Cache` opts into a [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl)-backed linear-solve path via the `F1MethodLinearSolveExt` weak-dep extension. +v0.6 changed `gradient` to return a `Vector` (previously a `1 × m` row matrix), routed parameter-side AD through [DifferentiationInterface.jl](https://github.com/JuliaDiff/DifferentiationInterface.jl) (configurable via the `ad` kwarg on the `F1Cache` constructor), and added an `F1MethodOptimizationExt` weak-dep extension exposing `optimization_function(...)` for use with [Optimization.jl](https://github.com/SciML/Optimization.jl). That being said, ultimately, it would make sense for the shortcuts used here to be integrated into a package like ChainRules.jl. For the time being, AIBECS users can use F1Method.jl to speed up their optimizations. diff --git a/ext/F1MethodLinearSolveExt.jl b/ext/F1MethodLinearSolveExt.jl new file mode 100644 index 0000000..2c12e34 --- /dev/null +++ b/ext/F1MethodLinearSolveExt.jl @@ -0,0 +1,55 @@ +module F1MethodLinearSolveExt + +using F1Method +using F1Method: F1Cache +using LinearSolve +using SciMLBase +using DifferentiationInterface +const DI = DifferentiationInterface + +function F1Method._init_F1Cache(linsolve::SciMLBase.AbstractLinearAlgorithm, + F, ∇ₓf, x, p, alg, ad; options...) + x = copy(x) + p = copy(p) + psol = copy(p) + s = solve(SteadyStateProblem(F, x, p), alg; options...).u + J = F.jac(s, p) + ∇ₚF = DI.jacobian(p -> F(s, p), ad, p) + # `u0 = similar(∇ₚF)` forces LinearCache.u to match the matrix-RHS + # shape; without it LinearSolve.init allocates `u` as a Vector even + # when `b` is a Matrix, and the subsequent solve! errors on the + # shape mismatch. + lc = init(LinearProblem(J, -∇ₚF; u0 = similar(∇ₚF)), linsolve) + sol = solve!(lc) + ∇s = copy(sol.u) # F1Cache owns its own ∇s buffer + return F1Cache(lc, s, ∇s, ∇ₓf(s, p), p, psol, ad) +end + +# cache.linear_cache is a LinearSolve.LinearCache; .A is the cache's +# stored Jacobian. Assigning to it triggers LinearSolve's setproperty! +# → isfresh = true, so the next solve!(linear_cache) refactors. +# UMFPACK/KLU reuse the symbolic factorization automatically when +# colptr/rowval match — that's the main win beyond plain `\`. +F1Method.set_A!(cache::F1Cache{<:LinearSolve.LinearCache}, A) = + cache.linear_cache.A = A + +# solve!(lc).u aliases lc.u. The only call site is +# `cache.∇s .= solveA!(...)`, which copies — don't "optimize" the +# broadcast away or you'll alias cache.∇s to the LinearCache buffer. +function F1Method.solveA!(cache::F1Cache{<:LinearSolve.LinearCache}, b) + cache.linear_cache.b = b + return solve!(cache.linear_cache).u +end + +# `cacheval` is the algorithm's stored Factorization (or, for +# GenericLUFactorization, a `(LU, ipiv)` tuple whose first element is +# the factorization). Group A algorithms make `factor' \ B` legal. +# Upstream-recommended idiom for concrete matrices until +# SciML/LinearSolve.jl#92 ships a first-class adjoint toggle on +# LinearCache. +_factor(cv) = cv +_factor(cv::Tuple) = first(cv) +F1Method.solveAᵀ!(cache::F1Cache{<:LinearSolve.LinearCache}, B) = + _factor(cache.linear_cache.cacheval)' \ B + +end # module diff --git a/src/F1Method.jl b/src/F1Method.jl index fb0b085..446a44e 100644 --- a/src/F1Method.jl +++ b/src/F1Method.jl @@ -14,106 +14,130 @@ using DifferentiationInterface const DI = DifferentiationInterface """ - Mem + F1Cache Memory cache holding the steady-state solution and the linear-algebra artifacts that the F-1 method reuses across `objective`, `gradient`, and `hessian` calls at the same parameter point. # Fields -- `s` : the steady-state solution, `𝒔(𝒑)` -- `A` : the factors of `𝐀 = ∇ₓ𝑭(𝒔, 𝒑)` -- `∇s` : the parameter-Jacobian of the steady state, `∇𝒔(𝒑)` -- `∇ₓf` : the state-Jacobian of the objective at the steady state, - `∇ₓ𝑓(𝒔, 𝒑)` -- `p` : the parameter vector that matches `A`, `∇s`, `∇ₓf` (and `s`) -- `psol` : the parameter vector that matches `s` (may be ahead of `p` +- `linear_cache`: a linear-solve cache holding the factors of + `𝐀 = ∇ₓ𝑭(𝒔, 𝒑)` — either a `Factorization` (default) + or a `LinearSolve.LinearCache` when a `linsolve` + algorithm is configured +- `s`: the steady-state solution, `𝒔(𝒑)` +- `∇s`: the parameter-Jacobian of the steady state, `∇𝒔(𝒑)` +- `∇ₓf`: the state-Jacobian of the objective at the steady state, `∇ₓ𝑓(𝒔, 𝒑)` +- `p`: the parameter vector that matches `linear_cache`, `∇s`, `∇ₓf` (and `s`) +- `psol`: the parameter vector that matches `s` (may be ahead of `p` when only the steady-state solve has been refreshed) -- `ad` : the AD backend used for parameter-side derivatives +- `ad`: the AD backend used for parameter-side derivatives -Initialise with [`initialize_mem`](@ref). +Construct with `F1Cache`. """ -mutable struct Mem{Ts, TA, T∇s, T∇ₓf, Tp, AD <: AbstractADType} - s::Ts - A::TA - ∇s::T∇s - ∇ₓf::T∇ₓf - p::Tp - psol::Tp +mutable struct F1Cache{C, S, JACS, JACF, P, AD <: AbstractADType} + linear_cache::C + s::S + ∇s::JACS + ∇ₓf::JACF + p::P + psol::P ad::AD end -function update_mem!(F, ∇ₓf, mem, p, alg; options...) - if p ≠ mem.p # only update if 𝒑 has changed - update_solution!(F, mem, p, alg; options...) - ∇ₚF = DI.jacobian(p -> F(mem.s, p), mem.ad, p) - mem.A = factorize(F.jac(mem.s, p)) # update factors of ∇ₓ𝑭(𝒔, 𝒑) - mem.∇s .= mem.A \ -∇ₚF # update ∇𝒔 - mem.∇ₓf .= ∇ₓf(mem.s, p) # update ∇ₓ𝑓(𝒔, 𝒑) - mem.p .= p # update 𝒑 for the cached values +# Solve-operation helpers on the mathematical operator 𝐀 = ∇ₓ𝑭. The +# `Factorization` branch lives here; the `LinearCache` branch lives in +# F1MethodLinearSolveExt. +set_A!(cache::F1Cache, J) = cache.linear_cache = factorize(J) +solveA!(cache::F1Cache, b) = cache.linear_cache \ b +solveAᵀ!(cache::F1Cache, B) = cache.linear_cache' \ B + + +function update_cache!(F, ∇ₓf, cache, p, alg; options...) + if p ≠ cache.p # only update if 𝒑 has changed + update_solution!(F, cache, p, alg; options...) + ∇ₚF = DI.jacobian(p -> F(cache.s, p), cache.ad, p) + set_A!(cache, F.jac(cache.s, p)) # update factors of ∇ₓ𝑭(𝒔, 𝒑) + cache.∇s .= solveA!(cache, -∇ₚF) # update ∇𝒔 + cache.∇ₓf .= ∇ₓf(cache.s, p) # update ∇ₓ𝑓(𝒔, 𝒑) + cache.p .= p # update 𝒑 for the cached values end + return cache end -function update_solution!(F, mem, p, alg; options...) - if p ≠ mem.psol - prob = SteadyStateProblem(F, mem.s, p) - mem.s .= solve(prob, alg; options...).u - mem.psol .= p +function update_solution!(F, cache, p, alg; options...) + if p ≠ cache.psol + prob = SteadyStateProblem(F, cache.s, p) + cache.s .= solve(prob, alg; options...).u + cache.psol .= p end + return cache end """ - objective(f, F, mem, p, alg; options...) + objective(f, F, cache, p, alg; options...) Evaluate `𝑓̂(𝒑) = 𝑓(𝒔(𝒑), 𝒑)`, where `𝒔(𝒑)` is the steady-state solution computed by the iterative solver `alg` (so that `𝑭(𝒔, 𝒑) = 0`). The -memory cache `mem` (built with [`initialize_mem`](@ref)) is updated -in place if `𝒑` has changed since the last call. +`cache` (built with `F1Cache`) is updated in place if `𝒑` has +changed since the last call. """ -function objective(f, F, mem, p, alg; options...) - update_solution!(F, mem, p, alg; options...) - return f(mem.s, p) +function objective(f, F, cache, p, alg; options...) + update_solution!(F, cache, p, alg; options...) + return f(cache.s, p) end """ - gradient(f, F, ∇ₓf, mem, p, alg; options...) + gradient(f, F, ∇ₓf, cache, p, alg; options...) Return the gradient `∇𝑓̂(𝒑)` as a `Vector` using the F-1 method. (Prior to F1Method 0.6 this returned a `1 × m` row matrix; the new shape is `Vector{T}` of length `m = length(p)`, which is what Optim.jl / Optimization.jl expect.) """ -function gradient(f, F, ∇ₓf, mem, p, alg; options...) - update_mem!(F, ∇ₓf, mem, p, alg; options...) - s, ∇s = mem.s, mem.∇s - ∇ₚf = DI.gradient(p -> f(s, p), mem.ad, p) - return vec(mem.∇ₓf * ∇s) + ∇ₚf +function gradient(f, F, ∇ₓf, cache, p, alg; options...) + update_cache!(F, ∇ₓf, cache, p, alg; options...) + s, ∇s = cache.s, cache.∇s + ∇ₚf = DI.gradient(p -> f(s, p), cache.ad, p) + return vec(cache.∇ₓf * ∇s) + ∇ₚf end """ - hessian(f, F, ∇ₓf, mem, p, alg; options...) + hessian(f, F, ∇ₓf, cache, p, alg; options...) Return the Hessian `∇²𝑓̂(𝒑)` as an `m × m` `Matrix` using the F-1 method. """ -function hessian(f, F, ∇ₓf, mem, p, alg; options...) - update_mem!(F, ∇ₓf, mem, p, alg; options...) - s, A, ∇s, m = mem.s, mem.A, mem.∇s, length(p) - A⁻ᵀ∇ₓfᵀ = vec(A' \ mem.∇ₓf') # independent of (𝑗, 𝑘) +function hessian(f, F, ∇ₓf, cache, p, alg; options...) + update_cache!(F, ∇ₓf, cache, p, alg; options...) + s, ∇s, m = cache.s, cache.∇s, length(p) + A⁻ᵀ∇ₓfᵀ = vec(solveAᵀ!(cache, cache.∇ₓf')) # independent of (𝑗, 𝑘) H(λ) = f(s + ∇s * λ, p + λ) - F(s + ∇s * λ, p + λ)' * A⁻ᵀ∇ₓfᵀ - return DI.hessian(H, mem.ad, zeros(m)) + return DI.hessian(H, cache.ad, zeros(m)) end """ - initialize_mem(F, ∇ₓf, x, p, alg; ad=AutoForwardDiff(), options...) + F1Cache(F, ∇ₓf, x, p, alg; ad=AutoForwardDiff(), linsolve=nothing, options...) -Initialise the F-1 memory cache. `ad` selects the AD backend used for +Construct the F-1 memory cache. `ad` selects the AD backend used for parameter-side derivatives (any `ADTypes.AbstractADType`); it defaults -to `AutoForwardDiff()`, matching pre-0.7 behaviour. Remaining keyword -arguments are forwarded to `solve(::SteadyStateProblem, alg; ...)`. +to `AutoForwardDiff()`, matching pre-0.7 behaviour. Passing a +`linsolve <: SciMLBase.AbstractLinearAlgorithm` (from +[LinearSolve.jl](https://github.com/SciML/LinearSolve.jl)) routes the +internal linear solves through a `LinearSolve.LinearCache`, enabling +symbolic-factorization reuse for sparse direct solvers and buffer +reuse for dense; the default `linsolve = nothing` keeps the legacy +`factorize` / `\\` path. Remaining keyword arguments are forwarded to +`solve(::SteadyStateProblem, alg; ...)`. """ -function initialize_mem(F, ∇ₓf, x, p, alg; - ad::AbstractADType = AutoForwardDiff(), options...) +function F1Cache(F, ∇ₓf, x, p, alg; + ad::AbstractADType = AutoForwardDiff(), + linsolve = nothing, + options...) + return _init_F1Cache(linsolve, F, ∇ₓf, x, p, alg, ad; options...) +end + +function _init_F1Cache(::Nothing, F, ∇ₓf, x, p, alg, ad; options...) x = copy(x) p = copy(p) psol = copy(p) @@ -121,14 +145,28 @@ function initialize_mem(F, ∇ₓf, x, p, alg; s = solve(prob, alg; options...).u A = factorize(F.jac(s, p)) ∇ₚF = DI.jacobian(p -> F(s, p), ad, p) - return Mem(s, A, A \ -∇ₚF, ∇ₓf(s, p), p, psol, ad) + return F1Cache(A, s, A \ -∇ₚF, ∇ₓf(s, p), p, psol, ad) +end + +""" + initialize_mem(args...; kwargs...) + +Deprecated; use `F1Cache` instead. Retained for 0.7.x; will be +removed in 0.8. +""" +function initialize_mem(args...; kwargs...) + Base.depwarn( + "initialize_mem is deprecated; use F1Cache(...) instead", + :initialize_mem, + ) + return F1Cache(args...; kwargs...) end """ - optimization_function(f, F, ∇ₓf, mem, alg; options...) + optimization_function(f, F, ∇ₓf, cache, alg; options...) -Wrap the F-1-method `objective` / `gradient` / `hessian` triple at memory -cache `mem` into a `SciMLBase.OptimizationFunction`, ready to feed into +Wrap the F-1-method `objective` / `gradient` / `hessian` triple at the +F-1 `cache` into a `SciMLBase.OptimizationFunction`, ready to feed into `OptimizationProblem` and any solver from the [Optimization.jl][1] ecosystem. @@ -139,4 +177,11 @@ when `Optimization.jl` has been loaded. """ function optimization_function end +# TODO(F1Method 0.8?): consider whether `optimization_function` should +# accept its own `linsolve` and rebuild the F1Cache internally, vs. +# always inheriting from the F1Cache the caller passes in. Today the +# latter is simpler — F1Cache is the carrier of every cached artifact +# including the linear solver — but if F1Cache construction grows more +# args this could become awkward. + end # module diff --git a/test/dense_rosenbrock.jl b/test/dense_rosenbrock.jl new file mode 100644 index 0000000..d1d20e1 --- /dev/null +++ b/test/dense_rosenbrock.jl @@ -0,0 +1,78 @@ +module SubModuleLinearSolveDense + +using Test +using F1Method +using LinearAlgebra +using SciMLBase +using ForwardDiff +using LinearSolve + +# Reuse the quasi-Rosenbrock setup from test/rosenbrock.jl: same statefun, +# same ODEFunction, same MyAlg Newton solver, same mismatch f / ∇ₓf. The +# only thing we vary is the `linsolve` kwarg on F1Cache. + +function newton_solve(F, ∇ₓF, x; Ftol = 1e-10) + while norm(F(x)) ≥ Ftol + x .-= ∇ₓF(x) \ F(x) + end + return x +end + +struct MyAlg <: SciMLBase.AbstractSteadyStateAlgorithm end + +function SciMLBase.solve(prob::SciMLBase.AbstractSteadyStateProblem, + alg::MyAlg; + Ftol = 1e-10) + p = prob.p + x0 = copy(prob.u0) + F(x) = prob.f(x, p) + ∇ₓF(x) = prob.f.jac(x, p) + x_steady = newton_solve(F, ∇ₓF, x0; Ftol = Ftol) + resid = F(x_steady) + return SciMLBase.build_solution(prob, alg, x_steady, resid; retcode = ReturnCode.Success) +end + +statefun(x, p, t = 0) = [ + -2 * (p[1] - x[1]) - 4 * p[2] * (x[2] - x[1]^2) * x[1] + p[2] * (x[2] - x[1]^2) +] +F = ODEFunction{false}(statefun; jac = (x, p, t = 0) -> ForwardDiff.jacobian(x -> statefun(x, p), x)) + +function state_mismatch(x) + δ(x) = x .- 1 + return 0.5δ(x)'δ(x) +end +function parameter_mismatch(p) + δ(p) = log.(p) + return 0.5δ(p)'δ(p) +end +f(x, p) = state_mismatch(x) + parameter_mismatch(p) +∇ₓf(x, p) = ForwardDiff.jacobian(x -> [f(x, p)], x) + +x₀ = rand(2) +p₀ = rand(2) + +exact_solution(p) = [p[1], p[1]^2] +exact_objective(p) = f(exact_solution(p), p) + +# `linsolve = nothing` skips LinearSolve entirely and uses Julia's +# `factorize`, which on a dense `Matrix{Float64}` returns a pivoted +# LAPACK `LU` (same family as LUFactorization, but invoked directly +# through LinearAlgebra rather than via a LinearSolve.LinearCache). +for linsolve in (nothing, LUFactorization(), GenericLUFactorization()) + @testset "linsolve = $(linsolve)" begin + cache = F1Method.F1Cache(F, ∇ₓf, x₀, p₀, MyAlg(); linsolve = linsolve) + F1_objective(p) = F1Method.objective(f, F, cache, p, MyAlg()) + F1_gradient(p) = F1Method.gradient(f, F, ∇ₓf, cache, p, MyAlg()) + F1_hessian(p) = F1Method.hessian(f, F, ∇ₓf, cache, p, MyAlg()) + + @test exact_objective(p₀) ≈ F1_objective(p₀) + @test exact_objective(2p₀) ≈ F1_objective(2p₀) + @test ForwardDiff.gradient(exact_objective, p₀) ≈ F1_gradient(p₀) + @test ForwardDiff.gradient(exact_objective, 2p₀) ≈ F1_gradient(2p₀) + @test ForwardDiff.hessian(exact_objective, p₀) ≈ F1_hessian(p₀) + @test ForwardDiff.hessian(exact_objective, 2p₀) ≈ F1_hessian(2p₀) + end +end + +end # submodule diff --git a/test/optimization_ext.jl b/test/optimization_ext.jl index 35e9f18..7ededd5 100644 --- a/test/optimization_ext.jl +++ b/test/optimization_ext.jl @@ -45,10 +45,10 @@ f(x, p) = state_mismatch(x) + parameter_mismatch(p) x₀ = ones(2) .+ 0.1 p₀ = ones(2) .* 1.5 -mem = F1Method.initialize_mem(F, ∇ₓf, x₀, p₀, MyAlg()) +cache = F1Method.F1Cache(F, ∇ₓf, x₀, p₀, MyAlg()) # Build the OptimizationFunction via the extension method -optfn = F1Method.optimization_function(f, F, ∇ₓf, mem, MyAlg()) +optfn = F1Method.optimization_function(f, F, ∇ₓf, cache, MyAlg()) @testset "Optimization.jl extension" begin @test optfn isa SciMLBase.OptimizationFunction diff --git a/test/rosenbrock.jl b/test/rosenbrock.jl deleted file mode 100644 index 435c027..0000000 --- a/test/rosenbrock.jl +++ /dev/null @@ -1,90 +0,0 @@ -module SubModuleRosenBrock - -using Test -using F1Method -using LinearAlgebra -using SciMLBase -using ForwardDiff - -# Set up: -# - overload `SteadyStateProblem` solve via a tiny Newton method -# - define solver algorithm (basic Newton here) -# - define type for that algorithm (here `MyAlg`) - -function newton_solve(F, ∇ₓF, x; Ftol=1e-10) - while norm(F(x)) ≥ Ftol - x .-= ∇ₓF(x) \ F(x) - end - return x -end - -# Create a type for the solver's algorithm -struct MyAlg <: SciMLBase.AbstractSteadyStateAlgorithm end - -# Overload SciMLBase's solve function -function SciMLBase.solve(prob::SciMLBase.AbstractSteadyStateProblem, - alg::MyAlg; - Ftol=1e-10) - p = prob.p - x0 = copy(prob.u0) - F(x) = prob.f(x, p) - ∇ₓF(x) = prob.f.jac(x, p) - x_steady = newton_solve(F, ∇ₓF, x0, Ftol=Ftol) - resid = F(x_steady) - return SciMLBase.build_solution(prob, alg, x_steady, resid; retcode = ReturnCode.Success) -end - - -# Quasi-Rosenbrock: F(x,p) is almost the gradient of the Rosenbrock function, -# except the Hessian matrix is made non-symmetric by slightly altering F(x,p). -# With p = [a, b] the parameters defining the constants of Rosenbrock's -# function, and the slight alteration to make the Hessian non-symmetric. -# Specifically, if r(x,y) denotes Rosenbrock, -# F₁([x, y], p) = ∂r/∂x -# F₂([x, y], p) = 0.5 * ∂r/∂y -statefun(x, p, t = 0) = [ - -2 * (p[1] - x[1]) - 4 * p[2] * (x[2] - x[1]^2) * x[1] - p[2] * (x[2] - x[1]^2) -] -F = ODEFunction{false}(statefun; jac = (x, p, t = 0) -> ForwardDiff.jacobian(x -> statefun(x, p), x)) - -# Define mismatch f(x,p) and its state-jacobian ∇ₓf(x,p) -function state_mismatch(x) - δ(x) = x .- 1 - return 0.5δ(x)'δ(x) -end -function parameter_mismatch(p) - δ(p) = log.(p) - return 0.5δ(p)'δ(p) -end -f(x, p) = state_mismatch(x) + parameter_mismatch(p) -∇ₓf(x, p) = ForwardDiff.jacobian(x -> [f(x, p)], x) # 1×n row matrix - -# Initialize the starting state and parameters -x₀ = rand(2) -p₀ = rand(2) -mem = F1Method.initialize_mem(F, ∇ₓf, x₀, p₀, MyAlg()) - -# Define the F-1 wrappers -F1_objective(p) = F1Method.objective(f, F, mem, p, MyAlg()) -F1_gradient(p) = F1Method.gradient(f, F, ∇ₓf, mem, p, MyAlg()) -F1_hessian(p) = F1Method.hessian(f, F, ∇ₓf, mem, p, MyAlg()) - -# Exact solution and objective for comparison -exact_solution(p) = [p[1], p[1]^2] -exact_objective(p) = f(exact_solution(p), p) - -@testset "objective" begin - @test exact_objective(p₀) ≈ F1_objective(p₀) - @test exact_objective(2p₀) ≈ F1_objective(2p₀) -end -@testset "gradient" begin - @test ForwardDiff.gradient(exact_objective, p₀) ≈ F1_gradient(p₀) - @test ForwardDiff.gradient(exact_objective, 2p₀) ≈ F1_gradient(2p₀) -end -@testset "Hessian" begin - @test ForwardDiff.hessian(exact_objective, p₀) ≈ F1_hessian(p₀) - @test ForwardDiff.hessian(exact_objective, 2p₀) ≈ F1_hessian(2p₀) -end - -end # submodule diff --git a/test/runtests.jl b/test/runtests.jl index 15d92e1..cc30ddb 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -3,14 +3,14 @@ ENV["DATADEPS_ALWAYS_ACCEPT"] = true using Test -@testset "quasi-Rosenbrock derivative" begin - include("rosenbrock.jl") +@testset "dense (quasi-Rosenbrock)" begin + include("dense_rosenbrock.jl") end @testset "Optimization.jl extension" begin include("optimization_ext.jl") end -@testset "AIBECS test" begin - include("AIBECS_test.jl") +@testset "sparse (AIBECS P-model)" begin + include("sparse_aibecs.jl") end diff --git a/test/AIBECS_test.jl b/test/sparse_aibecs.jl similarity index 73% rename from test/AIBECS_test.jl rename to test/sparse_aibecs.jl index b803dff..26cccda 100644 --- a/test/AIBECS_test.jl +++ b/test/sparse_aibecs.jl @@ -154,12 +154,8 @@ f, ∇ₓf = f_and_∇ₓf(ωs, ωp, grd, modify, obs, PmodelParameters) # steady state. solver_kwargs = (; abstol = 0.0, maxItNewton = 20) # CTKAlg uses maxItNewton — push to machine floor nl_solver_kwargs = (; abstol = 0.0, maxiters = 20) # NewtonRaphson uses maxiters — push to machine floor -mem = F1Method.initialize_mem(F, ∇ₓf, x, λ, CTKAlg(); solver_kwargs...) -gradient(λ) = F1Method.gradient(f, F, ∇ₓf, mem, λ, CTKAlg(); solver_kwargs...) -hessian(λ) = F1Method.hessian(f, F, ∇ₓf, mem, λ, CTKAlg(); solver_kwargs...) - # Note: we deliberately do not test `F1Method.objective` here. It just -# returns `f(mem.s, p)` after re-solving the steady state for `p` — there +# returns `f(cache.s, p)` after re-solving the steady state for `p` — there # is no analytical adjoint involved, so the only thing it would check is # whether the steady-state solve + `f` evaluation agree with themselves. # The interesting validation surface is `gradient` / `hessian`, which is @@ -168,7 +164,7 @@ hessian(λ) = F1Method.hessian(f, F, ∇ₓf, mem, λ, CTKAlg(); solver_kwargs. # Cold-started reference objective for verification. # # F1Method's gradient/hessian use the analytical adjoint formula at the -# cached `mem.s`. To validate them, we need an independent ground truth +# cached `cache.s`. To validate them, we need an independent ground truth # that does *not* depend on the cached state. We re-solve from scratch # (cold start) for each evaluation and let ForwardDiff propagate Duals # through Newton + the analytical ∂Gs Jacobian. With analytical ∂Gs in @@ -187,20 +183,45 @@ function objective_cold_nl(λ_eval) return f(s_eval, λ_eval) end -# F1 vs ForwardDiff + CTKAlg (existing checks) -@test ForwardDiff.gradient(objective_cold, 2λ) ≈ gradient(2λ) rtol = 1e-6 -@test ForwardDiff.hessian(objective_cold, 2λ) ≈ hessian(2λ) rtol = 1e-6 -@test ForwardDiff.gradient(objective_cold, λ) ≈ gradient(λ) rtol = 1e-6 -@test ForwardDiff.hessian(objective_cold, λ) ≈ hessian(λ) rtol = 1e-6 - -# F1 vs ForwardDiff + NonlinearSolve -@test ForwardDiff.gradient(objective_cold_nl, 2λ) ≈ gradient(2λ) rtol = 1e-6 -@test ForwardDiff.hessian(objective_cold_nl, 2λ) ≈ hessian(2λ) rtol = 1e-6 -@test ForwardDiff.gradient(objective_cold_nl, λ) ≈ gradient(λ) rtol = 1e-6 -@test ForwardDiff.hessian(objective_cold_nl, λ) ≈ hessian(λ) rtol = 1e-6 - -# FD reference (sanity check at looser rtol — second-order FD is noisy). -@test FiniteDiff.finite_difference_gradient(objective_cold, λ) ≈ gradient(λ) rtol = 1e-4 -@test FiniteDiff.finite_difference_hessian(objective_cold, λ) ≈ hessian(λ) rtol = 1e-2 +# Pre-compute ground-truth oracles once — they are invariant under +# `linsolve` (they don't touch F1Method at all), so the LinearSolve loop +# below reuses them across every variant. Each ForwardDiff hessian here +# costs ~1.6s warm; computing them per-variant would add ~10s for no +# information gain. +fd_g = (λ = ForwardDiff.gradient(objective_cold, λ), twoλ = ForwardDiff.gradient(objective_cold, 2λ)) +fd_h = (λ = ForwardDiff.hessian(objective_cold, λ), twoλ = ForwardDiff.hessian(objective_cold, 2λ)) +fdnl_g = (λ = ForwardDiff.gradient(objective_cold_nl, λ), twoλ = ForwardDiff.gradient(objective_cold_nl, 2λ)) +fdnl_h = (λ = ForwardDiff.hessian(objective_cold_nl, λ), twoλ = ForwardDiff.hessian(objective_cold_nl, 2λ)) +finitediff_g = FiniteDiff.finite_difference_gradient(objective_cold, λ) +finitediff_h = FiniteDiff.finite_difference_hessian(objective_cold, λ) + +# `linsolve = nothing` skips LinearSolve entirely and uses Julia's +# `factorize`, which on a `SparseMatrixCSC{Float64}` returns a UMFPACK +# `LU` from SparseArrays (same family as UMFPACKFactorization, but +# invoked directly through LinearAlgebra rather than via a +# LinearSolve.LinearCache). +for linsolve in (nothing, UMFPACKFactorization(), KLUFactorization()) + @testset "linsolve = $(linsolve)" begin + cache = F1Method.F1Cache(F, ∇ₓf, x, λ, CTKAlg(); linsolve = linsolve, solver_kwargs...) + gradient(λ) = F1Method.gradient(f, F, ∇ₓf, cache, λ, CTKAlg(); solver_kwargs...) + hessian(λ) = F1Method.hessian(f, F, ∇ₓf, cache, λ, CTKAlg(); solver_kwargs...) + + # F1 vs ForwardDiff + CTKAlg + @test fd_g.λ ≈ gradient(λ) rtol = 1e-6 + @test fd_h.λ ≈ hessian(λ) rtol = 1e-6 + @test fd_g.twoλ ≈ gradient(2λ) rtol = 1e-6 + @test fd_h.twoλ ≈ hessian(2λ) rtol = 1e-6 + + # F1 vs ForwardDiff + NonlinearSolve + @test fdnl_g.λ ≈ gradient(λ) rtol = 1e-6 + @test fdnl_h.λ ≈ hessian(λ) rtol = 1e-6 + @test fdnl_g.twoλ ≈ gradient(2λ) rtol = 1e-6 + @test fdnl_h.twoλ ≈ hessian(2λ) rtol = 1e-6 + + # FD reference (sanity check at looser rtol — second-order FD is noisy). + @test finitediff_g ≈ gradient(λ) rtol = 1e-4 + @test finitediff_h ≈ hessian(λ) rtol = 1e-2 + end +end end # submodule