feat: support array symbolics in BatchedInterface

Author: AayushSabharwal

Project.toml

@@ -15,6 +15,7 @@ Accessors = "0.1.36"
 Aqua = "0.8"
 ArrayInterface = "7.9"
 MacroTools = "0.5.13"
+Pkg = "1"
 RuntimeGeneratedFunctions = "0.5"
 SafeTestsets = "0.0.1"
 StaticArrays = "1.9"
@@ -24,9 +25,10 @@ julia = "1.10"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "Test", "SafeTestsets", "StaticArrays"]
+test = ["Aqua", "Pkg", "Test", "SafeTestsets", "StaticArrays"]

src/batched_interface.jl

@@ -50,10 +50,18 @@ function BatchedInterface(syssyms::Tuple...)
         symbol_subset = Int[]
         symbol_indexes = []
         system_isstate = BitVector()
+        allsyms = []
         for sym in syms
             if symbolic_type(sym) === NotSymbolic()
                 error("Only symbolic variables allowed in BatchedInterface.")
             end
+            if symbolic_type(sym) === ArraySymbolic()
+                append!(allsyms, collect(sym))
+            else
+                push!(allsyms, sym)
+            end
+        end
+        for sym in allsyms
             if !is_variable(sys, sym) && !is_parameter(sys, sym)
                 error("Only variables and parameters allowed in BatchedInterface.")
             end

test/downstream/Project.toml

@@ -0,0 +1,3 @@
+[deps]
+SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
+Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"

test/downstream/batchedinterface_arrayvars.jl

@@ -0,0 +1,42 @@
+using Symbolics
+using SymbolicIndexingInterface
+
+@variables x[1:2] y z
+
+syss = [
+    SymbolCache([x..., y]),
+    SymbolCache([x[1], y, z])
+]
+syms = [
+    [x, y],
+    [x[1], y]
+]
+probs = [
+    ProblemState(; u = [1.0, 2.0, 3.0]),
+    ProblemState(; u = [4.0, 5.0, 6.0])
+]
+
+bi = BatchedInterface(zip(syss, syms)...)
+
+@test all(isequal.(variable_symbols(bi), [x..., y]))
+@test variable_index.((bi,), [x..., y, z]) == [1, 2, 3, nothing]
+@test is_variable.((bi,), [x..., y, z]) == Bool[1, 1, 1, 0]
+@test associated_systems(bi) == [1, 1, 1]
+
+getter = getu(bi)
+@test (@inferred getter(probs...)) == [1.0, 2.0, 3.0]
+buf = zeros(3)
+@inferred getter(buf, probs...)
+@test buf == [1.0, 2.0, 3.0]
+
+setter! = setu(bi)
+buf .*= 10
+setter!(probs..., buf)
+
+@test state_values(probs[1]) == [10.0, 20.0, 30.0]
+@test state_values(probs[2]) == [10.0, 30.0, 6.0]
+
+buf ./= 10
+
+setter!(probs[1], 1, buf)
+@test state_values(probs[1]) == [1.0, 2.0, 3.0]

test/runtests.jl

@@ -1,34 +1,52 @@
 using SymbolicIndexingInterface
 using SafeTestsets
 using Test
+using Pkg
 
-@safetestset "Quality Assurance" begin
-    @time include("qa.jl")
-end
-@safetestset "Interface test" begin
-    @time include("example_test.jl")
-end
-@safetestset "Trait test" begin
-    @time include("trait_test.jl")
-end
-@safetestset "SymbolCache test" begin
-    @time include("symbol_cache_test.jl")
-end
-@safetestset "Fallback test" begin
-    @time include("fallback_test.jl")
-end
-@safetestset "Parameter indexing test" begin
-    @time include("parameter_indexing_test.jl")
-end
-@safetestset "State indexing test" begin
-    @time include("state_indexing_test.jl")
-end
-@safetestset "Remake test" begin
-    @time include("remake_test.jl")
+const GROUP = get(ENV, "GROUP", "All")
+
+function activate_downstream_env()
+    Pkg.activate("downstream")
+    Pkg.develop(PackageSpec(path = dirname(@__DIR__)))
+    Pkg.instantiate()
 end
-@safetestset "ProblemState test" begin
-    @time include("problem_state_test.jl")
+
+if GROUP == "All" || GROUP == "Core"
+    @safetestset "Quality Assurance" begin
+        @time include("qa.jl")
+    end
+    @safetestset "Interface test" begin
+        @time include("example_test.jl")
+    end
+    @safetestset "Trait test" begin
+        @time include("trait_test.jl")
+    end
+    @safetestset "SymbolCache test" begin
+        @time include("symbol_cache_test.jl")
+    end
+    @safetestset "Fallback test" begin
+        @time include("fallback_test.jl")
+    end
+    @safetestset "Parameter indexing test" begin
+        @time include("parameter_indexing_test.jl")
+    end
+    @safetestset "State indexing test" begin
+        @time include("state_indexing_test.jl")
+    end
+    @safetestset "Remake test" begin
+        @time include("remake_test.jl")
+    end
+    @safetestset "ProblemState test" begin
+        @time include("problem_state_test.jl")
+    end
+    @safetestset "BatchedInterface test" begin
+        @time include("batched_interface_test.jl")
+    end
 end
-@safetestset "BatchedInterface test" begin
-    @time include("batched_interface_test.jl")
+
+if GROUP == "All" || GROUP == "Downstream"
+    activate_downstream_env()
+    @safetestset "BatchedInterface with array symbolics test" begin
+        @time include("downstream/batchedinterface_arrayvars.jl")
+    end
 end