feat: add symbolic interface for LinearProblem

Author: AayushSabharwal

src/problems/linear_problems.jl

@@ -1,3 +1,36 @@
+"""
+    $(TYPEDEF)
+
+A utility struct stored inside `LinearProblem` to enable a symbolic interface.
+
+# Fields
+
+$(TYPEDFIELDS)
+"""
+struct SymbolicLinearInterface{F1, F2, S, M}
+    """
+    A function which takes `A` and the parameter object `p` and updates `A` in-place.
+    """
+    update_A!::F1
+    """
+    A function which takes `b` and the parameter object `p` and updates `b` in-place.
+    """
+    update_b!::F2
+    """
+    The symbolic backend for the `LinearProblem`.
+    """
+    sys::S
+    """
+    Arbitrary metadata useful for the symbolic backend.
+    """
+    metadata::M
+end
+
+__has_sys(::SymbolicLinearInterface) = true
+has_sys(::SymbolicLinearInterface) = true
+
+SymbolicIndexingInterface.symbolic_container(sli::SymbolicLinearInterface) = sli.sys
+
 @doc doc"""
 
 Defines a linear system problem.
@@ -50,20 +83,23 @@ parameters. Any extra keyword arguments are passed on to the solvers.
 * `b`: The right-hand side of the linear system.
 * `p`: The parameters for the problem. Defaults to `NullParameters`. Currently unused.
 * `u0`: The initial condition used by iterative solvers.
+* `symbolic_interface`: An instance of `SymbolicLinearInterface` if the problem was
+  generated by a symbolic backend.
 * `kwargs`: The keyword arguments passed on to the solvers.
 """
-struct LinearProblem{uType, isinplace, F, bType, P, K} <:
+struct LinearProblem{uType, isinplace, F, bType, P, I <: Union{SymbolicLinearInterface, Nothing}, K} <:
        AbstractLinearProblem{bType, isinplace}
     A::F
     b::bType
     u0::uType
     p::P
+    f::I
     kwargs::K
     @add_kwonly function LinearProblem{iip}(A, b, p = NullParameters(); u0 = nothing,
-            kwargs...) where {iip}
+            f = nothing, kwargs...) where {iip}
         warn_paramtype(p)
-        new{typeof(u0), iip, typeof(A), typeof(b), typeof(p), typeof(kwargs)}(A, b, u0, p,
-            kwargs)
+        new{typeof(u0), iip, typeof(A), typeof(b), typeof(p), typeof(f), typeof(kwargs)}(A, b, u0, p,
+            f, kwargs)
     end
 end
 
@@ -77,6 +113,16 @@ function LinearProblem(A, b, args...; kwargs...)
     end
 end
 
+SymbolicIndexingInterface.symbolic_container(prob::LinearProblem) = prob.f
+SymbolicIndexingInterface.state_values(prob::LinearProblem) = prob.u0
+SymbolicIndexingInterface.parameter_values(prob::LinearProblem) = prob.p
+SymbolicIndexingInterface.is_time_dependent(::LinearProblem) = false
+function SymbolicIndexingInterface.set_parameter!(valp::LinearProblem{A, B, C, D, E, <:SymbolicLinearInterface}, val, idx) where {A, B, C, D, E}
+    set_parameter!(parameter_values(valp), val, idx)
+    valp.f.update_A!(valp.A, valp.p)
+    valp.f.update_b!(valp.b, valp.p)
+end
+
 @doc doc"""
 Holds information on what variables to alias
 when solving a LinearProblem. Conforms to the AbstractAliasSpecifier interface. 

test/downstream/problem_interface.jl

@@ -332,3 +332,48 @@ prob = SteadyStateProblem(osys, [u0; ps])
         @test scc.ps[p] ≈ 2.5
     end
 end
+
+@testset "LinearProblem" begin
+    # TODO update when MTK codegen exists
+    sys = SymbolCache([:x, :y, :z], [:p, :q, :r])
+    update_A! = function (A, p)
+        A[1, 1] = p[1]
+        A[2, 2] = p[2]
+        A[3, 3] = p[3]
+    end
+    update_b! = function (b, p)
+        b[1] = p[3]
+        b[2] = -8p[2] - p[1]
+    end
+    f = SciMLBase.SymbolicLinearInterface(update_A!, update_b!, sys, nothing)
+    A = Float64[1 1 1; 6 -4 5; 5 2 2]
+    b = Float64[2, 31, 13]
+    p = Float64[1, -4, 2]
+    u0 = Float64[1, 2, 3]
+    prob = LinearProblem(A, b, p; u0, f)
+    @test prob[:x] ≈ 1.0
+    @test prob[:y] ≈ 2.0
+    @test prob[:z] ≈ 3.0
+    @test prob.ps[:p] ≈ 1.0
+    @test prob.ps[:q] ≈ -4.0
+    @test prob.ps[:r] ≈ 2.0
+    prob.ps[:p] = 2.0
+    @test prob.ps[:p] ≈ 2.0
+    @test prob.A[1, 1] ≈ 2.0
+    @test prob.b[2] ≈ 30.0
+
+    prob2 = remake(prob; u0 = 2u0)
+    @test prob2.u0 ≈ 2u0
+    prob2 = remake(prob; p = 2p)
+    @test prob2.p ≈ 2p
+    prob2 = remake(prob; u0 = [:x => 3.0], p = [:q => 1.5])
+    @test prob2.u0[1] ≈ 3.0
+    @test prob2.p[2] ≈ 1.5
+
+    # no u0
+    prob = LinearProblem(A, b, p; f)
+    prob2 = remake(prob; p = 2p)
+    @test prob2.p ≈ 2p
+    prob2 = remake(prob; p = [:q => 1.5])
+    @test prob2.p[2] ≈ 1.5
+end