Merge pull request #447 from AayushSabharwal/as/integrator-setindex

feat: support `setindex!` for nonlinear integrators

Author: ChrisRackauckas

src/NonlinearSolve.jl

@@ -53,7 +53,7 @@ using PrecompileTools: @recompile_invalidations, @compile_workload, @setup_workl
     using StaticArraysCore: StaticArray, SVector, SArray, MArray, Size, SMatrix
     using SymbolicIndexingInterface: SymbolicIndexingInterface, ParameterIndexingProxy,
                                      symbolic_container, parameter_values, state_values,
-                                     getu
+                                     getu, setu
 end
 
 @reexport using SciMLBase, SimpleNonlinearSolve

src/abstract_types.jl

@@ -220,6 +220,10 @@ function Base.getindex(cache::AbstractNonlinearSolveCache, sym)
     return getu(cache, sym)(cache)
 end
 
+function Base.setindex!(cache::AbstractNonlinearSolveCache, val, sym)
+    return setu(cache, sym)(cache, val)
+end
+
 function Base.show(io::IO, cache::AbstractNonlinearSolveCache)
     __show_cache(io, cache, 0)
 end

test/downstream/mtk_cache_indexing_tests.jl

@@ -10,39 +10,24 @@
     # 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
+    @testset "$integtype" for (alg, integtype) in [
+            (NewtonRaphson(), NonlinearSolve.GeneralizedFirstOrderAlgorithmCache),
+            (FastShortcutNonlinearPolyalg(), NonlinearSolve.NonlinearSolvePolyAlgorithmCache),
+            (SimpleNewtonRaphson(), NonlinearSolve.NonlinearSolveNoInitCache),
+        ]
+        nint = init(nlprob, alg)
+        @test nint isa integtype
+
+        for (i, sym) in enumerate([X, nlsys.X, :X])
+            # test both getindex and setindex!
+            nint[sym] = 1.5i
+            @test nint[sym] == 1.5i
+        end
+
+        for (i, sym) in enumerate([p, nlsys.p, :p])
+            # test both getindex and setindex!
+            nint.ps[sym] = 2.5i
+            @test nint.ps[sym] == 2.5i
+        end
     end
 end