Merge pull request #1570 from AayushSabharwal/as/fast-sub-sparse

feat: implement `fast_substitute` and `fixpoint_sub` for `SparseMatrixCSC`

Author: ChrisRackauckas

src/variable.jl

@@ -589,6 +589,12 @@ function fixpoint_sub(x, dict; operator = Nothing, maxiters = 1000)
 
     return x
 end
+function fixpoint_sub(x::SparseMatrixCSC, dict; operator = Nothing, maxiters = 1000)
+    I, J, V = findnz(x)
+    V = fixpoint_sub(V, dict; operator, maxiters)
+    m, n = size(x)
+    return sparse(I, J, V, m, n)
+end
 
 const Eq = Union{Equation, Inequality}
 """
@@ -620,6 +626,12 @@ end
 function fast_substitute(eqs::AbstractArray, subs::Pair; operator = Nothing)
     fast_substitute.(eqs, (subs,); operator)
 end
+function fast_substitute(eqs::SparseMatrixCSC, subs; operator = Nothing)
+    I, J, V = findnz(eqs)
+    V = fast_substitute(V, subs; operator)
+    m, n = size(eqs)
+    return sparse(I, J, V, m, n)
+end
 for (exprType, subsType) in Iterators.product((Num, Symbolics.Arr), (Any, Pair))
     @eval function fast_substitute(expr::$exprType, subs::$subsType; operator = Nothing)
         fast_substitute(value(expr), subs; operator)

test/utils.jl

@@ -3,7 +3,7 @@ import Symbolics: symbolic_to_float, var_from_nested_derivative, unwrap,
                   isblock, flatten_expr!, build_expr, get_variables, 
                   is_singleton, diff2term, tosymbol, lower_varname, 
                   makesubscripts, degree, coeff
-
+using SparseArrays
 using Test
 
 @testset "get_variables" begin
@@ -188,6 +188,17 @@ end
     @test isequal(Symbolics.fast_substitute(x[1], unwrap(x) => collect(unwrap(x))), x[1])
 end
 
+@testset "`fixpoint_sub` and `fast_substitute` on sparse arrays" begin
+    @variables x y z
+    mat = Num[x 0 0; 0 y 0; 0 0 z]
+    mat = sparse(mat)
+    mat = unwrap.(mat)
+    rules = Dict(x => y, y => z, z => 1)
+    res = Symbolics.fixpoint_sub(mat, rules)
+    @test res isa SparseMatrixCSC
+    @test res[1, 1] == res[2, 2] == res[3, 3] == 1
+end
+
 @testset "numerator and denominator" begin
     @variables x y
     num_den(x) = (numerator(x), denominator(x))