fix: handle non-standard indices in linear_expansion

Author: AayushSabharwal

Project.toml

@@ -25,6 +25,7 @@ LogExpFunctions = "2ab3a3ac-af41-5b50-aa03-7779005ae688"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
+OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Primes = "27ebfcd6-29c5-5fa9-bf4b-fb8fc14df3ae"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
@@ -79,6 +80,7 @@ Lux = "1"
 MacroTools = "0.5"
 NaNMath = "1"
 Nemo = "0.46, 0.47, 0.48"
+OffsetArrays = "1.15.0"
 PreallocationTools = "0.4"
 PrecompileTools = "1"
 Primes = "0.5"

src/Symbolics.jl

@@ -47,6 +47,8 @@ import SymbolicLimits
 
 using ADTypes: ADTypes
 
+import OffsetArrays
+
 @reexport using SymbolicUtils
 RuntimeGeneratedFunctions.init(@__MODULE__)
 

src/linear_algebra.jl

@@ -329,7 +329,7 @@ function _linear_expansion(t, x)
         arrt, idxst... = arguments(t)
         isequal(arrt, arrx) && return (0, t, true)
 
-        indexed_t = Symbolics.scalarize(arrt)[idxst...]
+        indexed_t = OffsetArrays.Origin(map(first, axes(arrt)))(Symbolics.scalarize(arrt))[idxst...]
         # when indexing a registered function/callable symbolic
         # scalarizing and indexing leads to the same symbolic variable
         # which causes a StackOverflowError without this

test/linear_solver.jl

@@ -73,4 +73,8 @@ a, b, islinear = Symbolics.linear_expansion(D(x) - x, x)
     @test !Symbolics.linear_expansion(z([x...]), x[1])[3]
     @test !Symbolics.linear_expansion(z(collect(Symbolics.unwrap(x))), x[1])[3]
     @test !Symbolics.linear_expansion(z([x, 2x]), x[1])[3]
+
+    @variables x[0:2]
+    a, b, islin = Symbolics.linear_expansion(x[0] - z(x[1]), z(x[1]))
+    @test islin && isequal(a, -1) && isequal(b, x[0])
 end