Merge pull request #1095 from AayushSabharwal/as/similarterm-arrayop

feat: support similarterm for arrayop

src/arrays.jl

@@ -62,6 +62,34 @@ end
 
 ConstructionBase.constructorof(s::Type{<:ArrayOp{T}}) where {T} = ArrayOp{T}
 
+function SymbolicUtils.similarterm(t::ArrayOp, f, args, _symtype = nothing; metadata = nothing)
+    oldargs = arguments(t)
+    if _symtype === nothing
+        _symtype = symtype(t)
+    end
+
+    if !all(isequal.(args, oldargs)) || !isequal(f, operation(t))
+        term = similarterm(t.term, f, args)
+        subs = Dict()
+        for (orig, new) in zip(oldargs, args)
+            isequal(orig, new) && continue
+            subs[orig] = new
+        end
+        if !isequal(f, operation(t))
+            subs[operation(t)] = f
+        end
+        expr = substitute(t.expr, subs)
+        expr = SymbolicUtils.term(operation(expr), arguments(expr)...)
+    else
+        term = t.term
+        expr = t.expr
+    end
+    if _symtype === nothing
+        _symtype = symtype(t)
+    end
+    return ArrayOp{_symtype}(t.output_idx, expr, t.reduce, term, t.shape, t.ranges, metadata)
+end
+
 shape(aop::ArrayOp) = aop.shape
 
 const show_arrayop = Ref{Bool}(false)

test/arrays.jl

@@ -73,6 +73,15 @@ end
     @test getmetadata(unwrap(v[1]), TestMetaT) == 4
 end
 
+@testset "similarterm" begin
+    @variables A[1:5, 1:5] B[1:5, 1:5]
+
+    T = unwrap(3A)
+    @test isequal(T, similarterm(T, operation(T), arguments(T)))
+    T2 = unwrap(3B)
+    @test isequal(T2, similarterm(T, operation(T), [*, 3, unwrap(B)]))
+end
+
 getdef(v) = getmetadata(v, Symbolics.VariableDefaultValue)
 @testset "broadcast & scalarize" begin
     @variables A[1:5,1:3]=42 b[1:3]=[2, 3, 5] t x(t)[1:4] u[1:1]