Merge branch 'master' into os/use-LinearSolve-precs

SciML · Sep 28, 2024 · 631c4d6 · 631c4d6
2 parents c4e6def + d6d741b
commit 631c4d6
Show file tree

Hide file tree

Showing 14 changed files with 252 additions and 166 deletions.
diff --git a/Project.toml b/Project.toml
@@ -8,6 +8,7 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
 FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
@@ -58,20 +59,21 @@ NonlinearSolveSymbolicsExt = "Symbolics"
 NonlinearSolveZygoteExt = "Zygote"
 
 [compat]
-ADTypes = "1.1.0"
+ADTypes = "1.9"
 Aqua = "0.8"
-ArrayInterface = "7.9"
+ArrayInterface = "7.16"
 BandedMatrices = "1.5"
 BenchmarkTools = "1.4"
-CUDA = "5.2"
+CUDA = "5.5"
 ConcreteStructs = "0.2.3"
-DiffEqBase = "6.149.0"
-Enzyme = "0.12"
+DiffEqBase = "6.155.3"
+DifferentiationInterface = "0.6.1"
+Enzyme = "0.13.2"
 ExplicitImports = "1.5"
-FastBroadcast = "0.2.8, 0.3"
+FastBroadcast = "0.3.5"
 FastClosures = "0.3.2"
 FastLevenbergMarquardt = "0.1"
-FiniteDiff = "2.22"
+FiniteDiff = "2.24"
 FixedPointAcceleration = "0.3"
 ForwardDiff = "0.10.36"
 Hwloc = "3"
@@ -80,36 +82,36 @@ LazyArrays = "1.8.2, 2"
 LeastSquaresOptim = "0.8.5"
 LineSearches = "7.2"
 LinearAlgebra = "1.10"
-LinearSolve = "2.30"
+LinearSolve = "2.35"
 MINPACK = "1.2"
 MaybeInplace = "0.1.3"
-ModelingToolkit = "9.15.0"
+ModelingToolkit = "9.41.0"
 NLSolvers = "0.5"
 NLsolve = "4.5"
 NaNMath = "1"
 NonlinearProblemLibrary = "0.1.2"
-OrdinaryDiffEq = "6.75"
+OrdinaryDiffEqTsit5 = "1.1.0"
 Pkg = "1.10"
 PrecompileTools = "1.2"
 Preferences = "1.4"
 Printf = "1.10"
 Random = "1.91"
 ReTestItems = "1.24"
-RecursiveArrayTools = "3.8"
+RecursiveArrayTools = "3.27"
 Reexport = "1.2"
 SIAMFANLEquations = "1.0.1"
-SciMLBase = "2.34.0"
+SciMLBase = "2.54.0"
 SciMLJacobianOperators = "0.1"
-SimpleNonlinearSolve = "1.8"
+SimpleNonlinearSolve = "1.12.3"
 SparseArrays = "1.10"
-SparseDiffTools = "2.19"
+SparseDiffTools = "2.22"
 SpeedMapping = "0.3"
 StableRNGs = "1"
 StaticArrays = "1.9"
 StaticArraysCore = "1.4"
 Sundials = "4.23.1"
-SymbolicIndexingInterface = "0.3.15"
-Symbolics = "5.26, 6"
+SymbolicIndexingInterface = "0.3.31"
+Symbolics = "6.12"
 Test = "1.10"
 TimerOutputs = "0.5.23"
 Zygote = "0.6.69"
@@ -133,7 +135,7 @@ NLSolvers = "337daf1e-9722-11e9-073e-8b9effe078ba"
 NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 NonlinearProblemLibrary = "b7050fa9-e91f-4b37-bcee-a89a063da141"
-OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
@@ -147,4 +149,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Aqua", "BandedMatrices", "BenchmarkTools", "CUDA", "Enzyme", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "Hwloc", "InteractiveUtils", "LeastSquaresOptim", "MINPACK", "ModelingToolkit", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEq", "Pkg", "Random", "ReTestItems", "SIAMFANLEquations", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Symbolics", "Test", "Zygote"]
+test = ["Aqua", "BandedMatrices", "BenchmarkTools", "CUDA", "Enzyme", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "Hwloc", "InteractiveUtils", "LeastSquaresOptim", "MINPACK", "ModelingToolkit", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEqTsit5", "Pkg", "Random", "ReTestItems", "SIAMFANLEquations", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Symbolics", "Test", "Zygote"]
diff --git a/docs/Project.toml b/docs/Project.toml
@@ -10,7 +10,7 @@ InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
-OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
 Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
@@ -30,10 +30,11 @@ DiffEqBase = "6.136"
 Documenter = "1"
 DocumenterCitations = "1"
 IncompleteLU = "0.2"
+InteractiveUtils = "<0.0.1, 1"
 LinearSolve = "2"
 ModelingToolkit = "8, 9"
 NonlinearSolve = "3"
-OrdinaryDiffEq = "6"
+OrdinaryDiffEqTsit5 = "1.1.0"
 Plots = "1"
 Random = "<0.0.1, 1"
 SciMLBase = "2.4"

diff --git a/docs/src/tutorials/code_optimization.md b/docs/src/tutorials/code_optimization.md
@@ -33,9 +33,7 @@ Take for example a prototypical small nonlinear solver code in its out-of-place
 ```@example small_opt
 using NonlinearSolve
 
-function f(u, p)
-    u .* u .- p
-end
+f(u, p) = u .* u .- p
 u0 = [1.0, 1.0]
 p = 2.0
 prob = NonlinearProblem(f, u0, p)
@@ -53,9 +51,7 @@ using BenchmarkTools
 Note that this way of writing the function is a shorthand for:
 
 ```@example small_opt
-function f(u, p)
-    [u[1] * u[1] - p, u[2] * u[2] - p]
-end
+f(u, p) = [u[1] * u[1] - p, u[2] * u[2] - p]
 ```
 
 where the function `f` returns an array. This is a common pattern from things like MATLAB's
@@ -71,7 +67,7 @@ by hand, this looks like:
 function f(du, u, p)
     du[1] = u[1] * u[1] - p
     du[2] = u[2] * u[2] - p
-    nothing
+    return nothing
 end
 
 prob = NonlinearProblem(f, u0, p)
@@ -84,6 +80,7 @@ the `.=` in-place broadcasting.
 ```@example small_opt
 function f(du, u, p)
     du .= u .* u .- p
+    return nothing
 end
 
 @benchmark sol = solve(prob, NewtonRaphson())
@@ -114,6 +111,7 @@ to normal array expressions, for example:
 
 ```@example small_opt
 using StaticArrays
+
 A = SA[2.0, 3.0, 5.0]
 typeof(A)
 ```
@@ -135,22 +133,20 @@ want to use the out-of-place allocating form, but this time we want to output a
 array. Doing it with broadcasting looks like:
 
 ```@example small_opt
-function f_SA(u, p)
-    u .* u .- p
-end
+f_SA(u, p) = u .* u .- p
+
 u0 = SA[1.0, 1.0]
 p = 2.0
 prob = NonlinearProblem(f_SA, u0, p)
+
 @benchmark solve(prob, NewtonRaphson())
 ```
 
 Note that only change here is that `u0` is made into a StaticArray! If we needed to write
 `f` out for a more complex nonlinear case, then we'd simply do the following:
 
 ```@example small_opt
-function f_SA(u, p)
-    SA[u[1] * u[1] - p, u[2] * u[2] - p]
-end
+f_SA(u, p) = SA[u[1] * u[1] - p, u[2] * u[2] - p]
 
 @benchmark solve(prob, NewtonRaphson())
 ```

diff --git a/docs/src/tutorials/optimizing_parameterized_ode.md b/docs/src/tutorials/optimizing_parameterized_ode.md
@@ -4,7 +4,7 @@ Let us fit a parameterized ODE to some data. We will use the Lotka-Volterra mode
 example. We will use Single Shooting to fit the parameters.
 
 ```@example parameterized_ode
-using OrdinaryDiffEq, NonlinearSolve, Plots
+using OrdinaryDiffEqTsit5, NonlinearSolve, Plots
 ```
 
 Let us simulate some real data from the Lotka-Volterra model.

diff --git a/lib/SciMLJacobianOperators/Project.toml b/lib/SciMLJacobianOperators/Project.toml
@@ -8,7 +8,6 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
-EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
@@ -19,9 +18,8 @@ ADTypes = "1.8.1"
 Aqua = "0.8.7"
 ConcreteStructs = "0.2.3"
 ConstructionBase = "1.5"
-DifferentiationInterface = "0.5"
+DifferentiationInterface = "0.6.1"
 Enzyme = "0.12, 0.13"
-EnzymeCore = "0.7, 0.8"
 ExplicitImports = "1.9.0"
 FastClosures = "0.3.2"
 FiniteDiff = "2.24.0"

diff --git a/lib/SciMLJacobianOperators/README.md b/lib/SciMLJacobianOperators/README.md
@@ -0,0 +1,59 @@
+# SciMLJacobianOperators.jl
+
+SciMLJacobianOperators provides a convenient way to compute Jacobian-Vector Product (JVP)
+and Vector-Jacobian Product (VJP) using
+[SciMLOperators.jl](https://github.com/SciML/SciMLOperators.jl) and
+[DifferentiationInterface.jl](https://github.com/gdalle/DifferentiationInterface.jl).
+
+Currently we have interfaces for:
+
+  - `NonlinearProblem`
+  - `NonlinearLeastSquaresProblem`
+
+and all autodiff backends supported by DifferentiationInterface.jl are supported.
+
+## Example
+
+```julia
+using SciMLJacobianOperators, NonlinearSolve, Enzyme, ForwardDiff
+
+# Define the problem
+f(u, p) = u .* u .- p
+u0 = ones(4)
+p = 2.0
+prob = NonlinearProblem(f, u0, p)
+fu0 = f(u0, p)
+v = ones(4) .* 2
+
+# Construct the operator
+jac_op = JacobianOperator(
+    prob, fu0, u0;
+    jvp_autodiff = AutoForwardDiff(),
+    vjp_autodiff = AutoEnzyme(; mode = Enzyme.Reverse)
+)
+sjac_op = StatefulJacobianOperator(jac_op, u0, p)
+
+sjac_op * v  # Computes the JVP
+# 4-element Vector{Float64}:
+#  4.0
+#  4.0
+#  4.0
+#  4.0
+
+sjac_op' * v  # Computes the VJP
+# 4-element Vector{Float64}:
+#  4.0
+#  4.0
+#  4.0
+#  4.0
+
+# What if we multiply the VJP and JVP?
+snormal_form = sjac_op' * sjac_op
+
+snormal_form * v  # Computes JᵀJ * v
+# 4-element Vector{Float64}:
+#  8.0
+#  8.0
+#  8.0
+#  8.0
+```
diff --git a/lib/SciMLJacobianOperators/src/SciMLJacobianOperators.jl b/lib/SciMLJacobianOperators/src/SciMLJacobianOperators.jl
@@ -1,10 +1,9 @@
 module SciMLJacobianOperators
 
-using ADTypes: ADTypes, AutoSparse, AutoEnzyme
+using ADTypes: ADTypes, AutoSparse
 using ConcreteStructs: @concrete
 using ConstructionBase: ConstructionBase
-using DifferentiationInterface: DifferentiationInterface
-using EnzymeCore: EnzymeCore
+using DifferentiationInterface: DifferentiationInterface, Constant
 using FastClosures: @closure
 using LinearAlgebra: LinearAlgebra
 using SciMLBase: SciMLBase, AbstractNonlinearProblem, AbstractNonlinearFunction
@@ -112,10 +111,10 @@ function JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff =
     iip = SciMLBase.isinplace(prob)
     T = promote_type(eltype(u), eltype(fu))
 
-    vjp_autodiff = set_function_as_const(get_dense_ad(vjp_autodiff))
+    vjp_autodiff = get_dense_ad(vjp_autodiff)
     vjp_op = prepare_vjp(skip_vjp, prob, f, u, fu; autodiff = vjp_autodiff)
 
-    jvp_autodiff = set_function_as_const(get_dense_ad(jvp_autodiff))
+    jvp_autodiff = get_dense_ad(jvp_autodiff)
     jvp_op = prepare_jvp(skip_jvp, prob, f, u, fu; autodiff = jvp_autodiff)
 
     output_cache = fu isa Number ? T(fu) : similar(fu, T)
@@ -295,23 +294,21 @@ function prepare_vjp(::Val{false}, prob::AbstractNonlinearProblem,
 
     @assert autodiff!==nothing "`vjp_autodiff` must be provided if `f` doesn't have \
                                 analytic `vjp` or `jac`."
-    # TODO: Once DI supports const params we can use `p`
-    fₚ = SciMLBase.JacobianWrapper{SciMLBase.isinplace(f)}(f, prob.p)
     if SciMLBase.isinplace(f)
-        @assert DI.check_twoarg(autodiff) "Backend: $(autodiff) doesn't support in-place \
-                                           problems."
+        @assert DI.check_inplace(autodiff) "Backend: $(autodiff) doesn't support in-place \
+                                            problems."
         fu_cache = copy(fu)
-        v_fake = copy(fu)
-        di_extras = DI.prepare_pullback(fₚ, fu_cache, autodiff, u, v_fake)
+        di_extras = DI.prepare_pullback(f, fu_cache, autodiff, u, (fu,), Constant(prob.p))
         return @closure (vJ, v, u, p) -> begin
-            DI.pullback!(fₚ, fu_cache, reshape(vJ, size(u)), autodiff,
-                u, reshape(v, size(fu_cache)), di_extras)
+            DI.pullback!(f, fu_cache, (reshape(vJ, size(u)),), di_extras, autodiff,
+                u, (reshape(v, size(fu_cache)),), Constant(p))
             return
         end
     else
-        di_extras = DI.prepare_pullback(fₚ, autodiff, u, fu)
+        di_extras = DI.prepare_pullback(f, autodiff, u, (fu,), Constant(prob.p))
         return @closure (v, u, p) -> begin
-            return DI.pullback(fₚ, autodiff, u, reshape(v, size(fu)), di_extras)
+            return only(DI.pullback(
+                f, di_extras, autodiff, u, (reshape(v, size(fu)),), Constant(p)))
         end
     end
 end
@@ -342,23 +339,21 @@ function prepare_jvp(::Val{false}, prob::AbstractNonlinearProblem,
 
     @assert autodiff!==nothing "`jvp_autodiff` must be provided if `f` doesn't have \
                                 analytic `vjp` or `jac`."
-    # TODO: Once DI supports const params we can use `p`
-    fₚ = SciMLBase.JacobianWrapper{SciMLBase.isinplace(f)}(f, prob.p)
     if SciMLBase.isinplace(f)
-        @assert DI.check_twoarg(autodiff) "Backend: $(autodiff) doesn't support in-place \
-                                           problems."
+        @assert DI.check_inplace(autodiff) "Backend: $(autodiff) doesn't support in-place \
+                                            problems."
         fu_cache = copy(fu)
-        di_extras = DI.prepare_pushforward(fₚ, fu_cache, autodiff, u, u)
+        di_extras = DI.prepare_pushforward(f, fu_cache, autodiff, u, (u,), Constant(prob.p))
         return @closure (Jv, v, u, p) -> begin
-            DI.pushforward!(
-                fₚ, fu_cache, reshape(Jv, size(fu_cache)),
-                autodiff, u, reshape(v, size(u)), di_extras)
+            DI.pushforward!(f, fu_cache, (reshape(Jv, size(fu_cache)),), di_extras,
+                autodiff, u, (reshape(v, size(u)),), Constant(p))
             return
         end
     else
-        di_extras = DI.prepare_pushforward(fₚ, autodiff, u, u)
+        di_extras = DI.prepare_pushforward(f, autodiff, u, (u,), Constant(prob.p))
         return @closure (v, u, p) -> begin
-            return DI.pushforward(fₚ, autodiff, u, reshape(v, size(u)), di_extras)
+            return only(DI.pushforward(
+                f, di_extras, autodiff, u, (reshape(v, size(u)),), Constant(p)))
         end
     end
 end
@@ -371,10 +366,8 @@ function prepare_scalar_op(::Val{false}, prob::AbstractNonlinearProblem,
 
     @assert autodiff!==nothing "`autodiff` must be provided if `f` doesn't have \
                                 analytic `vjp` or `jvp` or `jac`."
-    # TODO: Once DI supports const params we can use `p`
-    fₚ = Base.Fix2(f, prob.p)
-    di_extras = DI.prepare_derivative(fₚ, autodiff, u)
-    return @closure (v, u, p) -> DI.derivative(fₚ, autodiff, u, di_extras) * v
+    di_extras = DI.prepare_derivative(f, autodiff, u, Constant(prob.p))
+    return @closure (v, u, p) -> DI.derivative(f, di_extras, autodiff, u, Constant(p)) * v
 end
 
 get_dense_ad(::Nothing) = nothing
@@ -386,12 +379,6 @@ function get_dense_ad(ad::AutoSparse)
     return dense_ad
 end
 
-# In our case we know that it is safe to mark the function as const
-set_function_as_const(ad) = ad
-function set_function_as_const(ad::AutoEnzyme{M, Nothing}) where {M}
-    return AutoEnzyme(; ad.mode, function_annotation = EnzymeCore.Const)
-end
-
 export JacobianOperator, VecJacOperator, JacVecOperator
 export StatefulJacobianOperator
 export StatefulJacobianNormalFormOperator