Merge pull request #1686 from JuliaSymbolics/as/more-stuff

fix: fix some method ambiguties, improve testing

Author: AayushSabharwal

src/array-lib.jl

@@ -11,6 +11,7 @@ function wrapper_fn_from_idxs(x::Arr{T, N}, idxs...) where {T, N}
     nd = _indexed_ndims(idxs...)
     return nd == 0 ? is_wrapper_type(T) ? T : identity : Arr{T, nd}
 end
+wrapper_fn_from_idxs(x::Arr{T, N}, idx::SymbolicUtils.StableIndex{Int}) where {T, N} = T
 # Wrapped array should wrap the elements too
 function Base.getindex(x::Arr{T, N}, idx::CartesianIndex{N}) where {T, N}
     if is_wrapper_type(T)
@@ -80,10 +81,10 @@ end
 function *(a::PolyadicT, b::Arr, bs::PolyadicT...)
     return *(a, unwrap(b), bs...)
 end
-function *(a::LinearAlgebra.Adjoint{T, <: AbstractVector}, b::Arr, bs::PolyadicT...) where {T}
+function *(a::LinearAlgebra.Adjoint{T, <: AbstractVector}, b::Arr, bs::PolyadicT...) where {T <: Number}
     return *(a, unwrap(b), bs...)
 end
-function *(a::LinearAlgebra.Adjoint{T, <: AbstractVector}, b::Arr, c::AbstractVector, bs::PolyadicT...) where {T}
+function *(a::LinearAlgebra.Adjoint{T, <: AbstractVector}, b::Arr, c::AbstractVector, bs::PolyadicT...) where {T <: Number}
     return *(a, unwrap(b), unwrap(c), bs...)
 end
 function *(a::Number, b::Arr, bs::PolyadicT...)
@@ -121,6 +122,20 @@ function +(x1::Arr, x2::Arr, args::AbstractArray...)
     return +(unwrap(x1), unwrap(x2), args...)
 end
 
+for T1 in [Arr, AbstractArray], T2 in [Arr, AbstractArray]
+    T1 == T2 == AbstractArray && continue
+    @eval Base.:(\)(x1::$T1{Num, 1}, x2::$T2{Num, 1}) = Num(unwrap(x1) \ unwrap(x2))
+    @eval Base.:(\)(x1::$T1{Num, 1}, x2::$T2{Num, 2}) = Arr{Num, 2}(unwrap(x1) \ unwrap(x2))
+    @eval Base.:(\)(x1::$T1{Num, 2}, x2::$T2{Num, 1}) = Arr{Num, 1}(unwrap(x1) \ unwrap(x2))
+    @eval Base.:(\)(x1::$T1{Num, 2}, x2::$T2{Num, 2}) = Arr{Num, 2}(unwrap(x1) \ unwrap(x2))
+
+    @eval Base.:(/)(x1::$T1{Num, 1}, x2::$T2{Num, 1}) = Arr{Num, 2}(unwrap(x1) / unwrap(x2))
+    @eval Base.:(/)(x1::$T1{Num, 1}, x2::$T2{Num, 2}) = Arr{Num, 2}(unwrap(x1) / unwrap(x2))
+    @eval Base.:(/)(x1::$T1{Num, 2}, x2::$T2{Num, 2}) = Arr{Num, 2}(unwrap(x1) / unwrap(x2))
+end
+
+Base.:(/)(x1::Num, x2::Arr{Num, 1}) = Arr{Num, 2}(unwrap(x1) / unwrap(x2))
+
 #################### MAP-REDUCE ################
 
 SymbolicUtils.@map_methods Arr unwrap wrap

src/num.jl

@@ -65,7 +65,7 @@ for (T1, T2) in Iterators.product([Num, Integer], [Num, Integer])
     end
 end
 
-for f in [/, \, ^]
+for f in [\, ^]
     @eval function (::$(typeof(f)))(x1::AbstractArray{<:Real}, x2::Num)
         $f(x1, unwrap(x2))
     end
@@ -75,6 +75,14 @@ for f in [/, \, ^]
     end
 end
 
+function Base.:(/)(x1::AbstractArray{<:Real}, x2::Num)
+    /(unwrap(x1), unwrap(x2))
+end
+
+function Base.:(/)(x1::Num, x2::AbstractVector{<:Real})
+    /(unwrap(x1), unwrap(x2))
+end
+
 Base.conj(x::Num) = x
 Base.transpose(x::Num) = x
 

test/arrays.jl

@@ -3,7 +3,8 @@ using SymbolicUtils, Test
 using Symbolics: symtype, shape, wrap, unwrap, Arr, jacobian, @variables, value, get_variables, @arrayop, getname, metadata, scalarize
 using Base: Slice
 using SymbolicUtils: Sym, term, operation, search_variables
-import LinearAlgebra: dot, Adjoint
+import SymbolicUtils.Code: toexpr
+import LinearAlgebra: dot, Adjoint, cross, diagm
 import ..limit2
 
 struct TestMetaT end
@@ -446,7 +447,78 @@ end
     @test operation(sym) === +
 end
 
-@testset "matmul ambiguity" begin
-    @variables x[1:3] y
-    @test_nowarn x * y
+@testset "Basic functions test" begin
+    @variables a b c
+    @variables x[1:3] y[1:3, 1:3]
+    arrays = [
+        x,
+        y,
+        unwrap(x),
+        unwrap(y),
+        collect(y),
+        [a, b, c],
+        diagm([a, b, c]) * x,
+        diagm([a, b, c]) * [a, b, c],
+        y * [a, b, c],
+        diagm([a, b, c]) * y,
+    ]
+
+    xval = rand(3)
+    yval = rand(3, 3)
+    setup(ex) = quote
+        let a = 1.0, b = 2.3, c = 4.5, x = $xval, y = $yval
+            $(ex)
+        end
+    end
+
+    @testset "$arr1 AND $arr2" for arr1 in arrays, arr2 in arrays
+        val1 = eval(setup(toexpr(arr1)))
+        val2 = eval(setup(toexpr(arr2)))
+        if ndims(arr1) == ndims(arr2) == 1 && length(arr1) == length(arr2) == 3
+            t1 = setup(toexpr(cross(arr1, arr2)))
+            @test eval(t1) ≈ cross(val1, val2)
+        end
+        if size(arr1) == size(arr2)
+            t1 = setup(toexpr(dot(arr1, arr2)))
+            @test eval(t1) ≈ dot(val1, val2)
+
+            t1 = setup(toexpr(arr1 .* arr2))
+            @test eval(t1) ≈ val1 .* val2
+
+            t1 = setup(toexpr(arr1 + arr2))
+            @test eval(t1) ≈ val1 + val2
+        end
+        t1 = try
+            setup(toexpr(arr1'arr2))
+        catch
+            nothing
+        end
+        if t1 !== nothing
+            @test eval(t1) ≈ val1'val2
+        end
+
+        truth = try
+            val1 / val2
+        catch
+            nothing
+        end
+        if truth !== nothing
+            t1 = setup(toexpr(arr1 / arr2))
+            @test eval(t1) ≈ truth
+        end
+
+        truth = try
+            val1 \ val2
+        catch
+            nothing
+        end
+        t1 = try
+            setup(toexpr(arr1 \ arr2))
+        catch
+            nothing
+        end
+        if truth !== nothing && t1 !== nothing
+            @test eval(t1) ≈ truth
+        end
+    end
 end