Release v0.7.0: optional LinearSolve.jl backend, rename Mem to F1Cache

.gitignore

@@ -1,3 +1,5 @@
 Manifest.toml
 docs/build/
 *html
+*plan*md
+.vscode/

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`.

Project.toml

@@ -1,7 +1,7 @@
 name = "F1Method"
 uuid = "d5605bae-6d76-11e9-2fd7-71bcf42edbaa"
 authors = ["Benoit Pasquier <briochemc@gmail.com>"]
-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"

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.

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