Merge pull request #406 from AayushSabharwal/as/nlintegrator-indexing

feat: support indexing AbstractNonlinearSolveCache

Author: ChrisRackauckas

.github/workflows/CI.yml

@@ -19,6 +19,7 @@ jobs:
       matrix:
         group:
           - Core
+          - Downstream
         version:
           - "1"
         os:

Project.toml

@@ -27,6 +27,7 @@ SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
+SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [weakdeps]
@@ -35,8 +36,8 @@ FastLevenbergMarquardt = "7a0df574-e128-4d35-8cbd-3d84502bf7ce"
 FixedPointAcceleration = "817d07cb-a79a-5c30-9a31-890123675176"
 LeastSquaresOptim = "0fc2ff8b-aaa3-5acd-a817-1944a5e08891"
 MINPACK = "4854310b-de5a-5eb6-a2a5-c1dee2bd17f9"
-NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
 NLSolvers = "337daf1e-9722-11e9-073e-8b9effe078ba"
+NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
 SIAMFANLEquations = "084e46ad-d928-497d-ad5e-07fa361a48c4"
 SpeedMapping = "f1835b91-879b-4a3f-a438-e4baacf14412"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
@@ -48,8 +49,8 @@ NonlinearSolveFastLevenbergMarquardtExt = "FastLevenbergMarquardt"
 NonlinearSolveFixedPointAccelerationExt = "FixedPointAcceleration"
 NonlinearSolveLeastSquaresOptimExt = "LeastSquaresOptim"
 NonlinearSolveMINPACKExt = "MINPACK"
-NonlinearSolveNLsolveExt = "NLsolve"
 NonlinearSolveNLSolversExt = "NLSolvers"
+NonlinearSolveNLsolveExt = "NLsolve"
 NonlinearSolveSIAMFANLEquationsExt = "SIAMFANLEquations"
 NonlinearSolveSpeedMappingExt = "SpeedMapping"
 NonlinearSolveSymbolicsExt = "Symbolics"
@@ -61,8 +62,8 @@ Aqua = "0.8"
 ArrayInterface = "7.9"
 BandedMatrices = "1.4"
 BenchmarkTools = "1.4"
-ConcreteStructs = "0.2.3"
 CUDA = "5.2"
+ConcreteStructs = "0.2.3"
 DiffEqBase = "6.149.0"
 Enzyme = "0.11.15"
 FastBroadcast = "0.2.8"
@@ -78,8 +79,8 @@ LinearAlgebra = "1.10"
 LinearSolve = "2.21"
 MINPACK = "1.2"
 MaybeInplace = "0.1.1"
-NLsolve = "4.5"
 NLSolvers = "0.5"
+NLsolve = "4.5"
 NaNMath = "1"
 NonlinearProblemLibrary = "0.1.2"
 OrdinaryDiffEq = "6.74"
@@ -103,6 +104,7 @@ StaticArrays = "1.7"
 StaticArraysCore = "1.4"
 Sundials = "4.23.1"
 Symbolics = "5.13"
+SymbolicIndexingInterface = "0.3.3"
 Test = "1.10"
 TimerOutputs = "0.5.23"
 Zygote = "0.6.69"
@@ -122,8 +124,8 @@ LeastSquaresOptim = "0fc2ff8b-aaa3-5acd-a817-1944a5e08891"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 MINPACK = "4854310b-de5a-5eb6-a2a5-c1dee2bd17f9"
-NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56"
 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"

docs/make.jl

@@ -18,7 +18,8 @@ makedocs(; sitename = "NonlinearSolve.jl",
     clean = true,
     doctest = false,
     linkcheck = true,
-    linkcheck_ignore = ["https://twitter.com/ChrisRackauckas/status/1544743542094020615"],
+    linkcheck_ignore = ["https://twitter.com/ChrisRackauckas/status/1544743542094020615",
+                        "https://link.springer.com/article/10.1007/s40096-020-00339-4"],
     checkdocs = :exports,
     warnonly = [:missing_docs],
     plugins = [bib],

src/NonlinearSolve.jl

@@ -29,6 +29,9 @@ import PrecompileTools: @recompile_invalidations, @compile_workload, @setup_work
                       AbstractSciMLOperator, NLStats, _unwrap_val, has_jac, isinplace
     import SparseDiffTools: AbstractSparsityDetection, AutoSparseEnzyme
     import StaticArraysCore: StaticArray, SVector, SArray, MArray, Size, SMatrix, MMatrix
+    import SymbolicIndexingInterface: SymbolicIndexingInterface, ParameterIndexingProxy,
+                                      symbolic_container, parameter_values, state_values,
+                                      getu
 end
 
 @reexport using ADTypes, SciMLBase, SimpleNonlinearSolve

src/abstract_types.jl

@@ -207,6 +207,27 @@ Abstract Type for all NonlinearSolve.jl Caches.
 """
 abstract type AbstractNonlinearSolveCache{iip, timeit} end
 
+function SymbolicIndexingInterface.symbolic_container(cache::AbstractNonlinearSolveCache)
+    cache.prob
+end
+function SymbolicIndexingInterface.parameter_values(cache::AbstractNonlinearSolveCache)
+    parameter_values(symbolic_container(cache))
+end
+function SymbolicIndexingInterface.state_values(cache::AbstractNonlinearSolveCache)
+    state_values(symbolic_container(cache))
+end
+
+function Base.getproperty(cache::AbstractNonlinearSolveCache, sym::Symbol)
+    if sym == :ps
+        return ParameterIndexingProxy(cache)
+    end
+    return getfield(cache, sym)
+end
+
+function Base.getindex(cache::AbstractNonlinearSolveCache, sym)
+    return getu(cache, sym)(cache)
+end
+
 function Base.show(io::IO, cache::AbstractNonlinearSolveCache)
     __show_cache(io, cache, 0)
 end

src/core/generalized_first_order.jl

@@ -113,6 +113,11 @@ concrete_jac(::GeneralizedFirstOrderAlgorithm{CJ}) where {CJ} = CJ
     force_stop::Bool
 end
 
+SymbolicIndexingInterface.state_values(cache::GeneralizedFirstOrderAlgorithmCache) = cache.u
+function SymbolicIndexingInterface.parameter_values(cache::GeneralizedFirstOrderAlgorithmCache)
+    cache.p
+end
+
 function __reinit_internal!(
         cache::GeneralizedFirstOrderAlgorithmCache{iip}, args...; p = cache.p, u0 = cache.u,
         alias_u0::Bool = false, maxiters = 1000, maxtime = nothing, kwargs...) where {iip}

src/default.jl

@@ -70,6 +70,11 @@ end
     alias_u0::Bool
 end
 
+function SymbolicIndexingInterface.symbolic_container(cache::NonlinearSolvePolyAlgorithmCache)
+    cache.caches[cache.current]
+end
+SymbolicIndexingInterface.state_values(cache::NonlinearSolvePolyAlgorithmCache) = cache.u0
+
 function Base.show(
         io::IO, cache::NonlinearSolvePolyAlgorithmCache{pType, N}) where {pType, N}
     problem_kind = ifelse(pType == :NLS, "NonlinearProblem", "NonlinearLeastSquaresProblem")

test/downstream/Project.toml

@@ -0,0 +1,4 @@
+[deps]
+ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
+NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
+SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"

test/downstream/cache_indexing.jl

@@ -0,0 +1,46 @@
+using ModelingToolkit, NonlinearSolve
+using ModelingToolkit: t_nounits as t
+
+@parameters p d
+@variables X(t)
+eqs = [0 ~ sin(X + p) - d * sqrt(X + 1)]
+@mtkbuild nlsys = NonlinearSystem(eqs, [X], [p, d])
+
+# Creates an integrator.
+nlprob = NonlinearProblem(nlsys, [X => 1.0], [p => 2.0, d => 3.0])
+
+@testset "GeneralizedFirstOrderAlgorithmCache" begin
+    nint = init(nlprob, NewtonRaphson())
+    @test nint isa NonlinearSolve.GeneralizedFirstOrderAlgorithmCache
+
+    @test nint[X] == 1.0
+    @test nint[nlsys.X] == 1.0
+    @test nint[:X] == 1.0
+    @test nint.ps[p] == 2.0
+    @test nint.ps[nlsys.p] == 2.0
+    @test nint.ps[:p] == 2.0
+end
+
+@testset "NonlinearSolvePolyAlgorithmCache" begin
+    nint = init(nlprob, FastShortcutNonlinearPolyalg())
+    @test nint isa NonlinearSolve.NonlinearSolvePolyAlgorithmCache
+
+    @test nint[X] == 1.0
+    @test nint[nlsys.X] == 1.0
+    @test nint[:X] == 1.0
+    @test nint.ps[p] == 2.0
+    @test nint.ps[nlsys.p] == 2.0
+    @test nint.ps[:p] == 2.0
+end
+
+@testset "NonlinearSolveNoInitCache" begin
+    nint = init(nlprob, SimpleNewtonRaphson())
+    @test nint isa NonlinearSolve.NonlinearSolveNoInitCache
+
+    @test nint[X] == 1.0
+    @test nint[nlsys.X] == 1.0
+    @test nint[:X] == 1.0
+    @test nint.ps[p] == 2.0
+    @test nint.ps[nlsys.p] == 2.0
+    @test nint.ps[:p] == 2.0
+end

test/downstream/downstream_tests.jl

@@ -0,0 +1,11 @@
+using Pkg
+using SafeTestsets
+
+function activate_downstream_env()
+    Pkg.activate("downstream")
+    Pkg.develop(PackageSpec(path = dirname(dirname(@__DIR__))))
+    Pkg.instantiate()
+end
+
+activate_downstream_env()
+@safetestset "Cache indexing test" include("cache_indexing.jl")

test/runtests.jl

@@ -7,6 +7,10 @@ if GROUP == "All" || GROUP == "Core"
         joinpath(@__DIR__, "wrappers/"))
 end
 
+if GROUP == "Downstream"
+    include("downstream/downstream_tests.jl")
+end
+
 if GROUP == "GPU"
     ReTestItems.runtests(joinpath(@__DIR__, "gpu/"))
 end