Generalize broadcasting to StaticMatrixLike (#1220)

* Generalize broadcasting to `StaticMatrixLike`

* Bump version to 1.7.0

Author: devmotion

Project.toml

@@ -1,6 +1,6 @@
 name = "StaticArrays"
 uuid = "90137ffa-7385-5640-81b9-e52037218182"
-version = "1.6.5"
+version = "1.7.0"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/abstractarray.jl

@@ -8,13 +8,10 @@ length(a::Type{SA}) where {SA <: StaticArrayLike} = prod(Size(SA))::Int
 end
 @inline size(a::StaticArrayLike) = Tuple(Size(a))
 
-Base.axes(s::StaticArray) = _axes(Size(s))
+Base.axes(s::StaticArrayLike) = _axes(Size(s))
 @pure function _axes(::Size{sizes}) where {sizes}
     map(SOneTo, sizes)
 end
-Base.axes(rv::Adjoint{<:Any,<:StaticVector})   = (SOneTo(1), axes(rv.parent)...)
-Base.axes(rv::Transpose{<:Any,<:StaticVector}) = (SOneTo(1), axes(rv.parent)...)
-Base.axes(d::Diagonal{<:Any,<:StaticVector}) = (ax = axes(d.diag, 1); (ax, ax))
 
 Base.eachindex(::IndexLinear, a::StaticArray) = SOneTo(length(a))
 

src/broadcast.jl

@@ -9,9 +9,7 @@ import Base.Broadcast: _bcs1  # for SOneTo axis information
 using Base.Broadcast: _bcsm
 
 BroadcastStyle(::Type{<:StaticArray{<:Tuple, <:Any, N}}) where {N} = StaticArrayStyle{N}()
-BroadcastStyle(::Type{<:Transpose{<:Any, <:StaticArray}}) = StaticArrayStyle{2}()
-BroadcastStyle(::Type{<:Adjoint{<:Any, <:StaticArray}}) = StaticArrayStyle{2}()
-BroadcastStyle(::Type{<:Diagonal{<:Any, <:StaticArray{<:Tuple, <:Any, 1}}}) = StaticArrayStyle{2}()
+BroadcastStyle(::Type{<:StaticMatrixLike}) = StaticArrayStyle{2}()
 # Precedence rules
 BroadcastStyle(::StaticArrayStyle{M}, ::DefaultArrayStyle{N}) where {M,N} =
     DefaultArrayStyle(Val(max(M, N)))
@@ -104,10 +102,7 @@ function broadcast_getindex(oldsize::Tuple, i::Int, newindex::CartesianIndex)
     return :(a[$i][$ind])
 end
 
-isstatic(::StaticArray) = true
-isstatic(::Transpose{<:Any, <:StaticArray}) = true
-isstatic(::Adjoint{<:Any, <:StaticArray}) = true
-isstatic(::Diagonal{<:Any, <:StaticArray}) = true
+isstatic(::StaticArrayLike) = true
 isstatic(_) = false
 
 @inline first_statictype(x, y...) = isstatic(x) ? typeof(x) : first_statictype(y...)

test/abstractarray.jl

@@ -10,6 +10,20 @@ using StaticArrays, Test, LinearAlgebra
         @test eachindex(m) isa SOneTo
     end
 
+    @testset "axes" begin
+        v = @SVector [1, 2, 3]
+        @test @inferred(axes(v)) == (SOneTo(3),)
+        for T in (Adjoint, Transpose)
+            @test @inferred(axes(T(v))) == (SOneTo(1), SOneTo(3))
+        end
+
+        m = @SMatrix [1 2; 3 4]
+        @test @inferred(axes(m)) == (SOneTo(2), SOneTo(2))
+        for T in (Adjoint, Transpose, Diagonal, Symmetric, Hermitian, UpperTriangular, LowerTriangular, UnitUpperTriangular, UnitLowerTriangular)
+            @test @inferred(axes(T(m))) == (SOneTo(2), SOneTo(2))
+        end
+    end
+
     @testset "strides" begin
         m1 = MArray{Tuple{3, 4, 5}}(rand(Int, 3, 4, 5))
         m2 = SizedArray{Tuple{3,4,5}}(rand(Int, 3, 4, 5))

test/broadcast.jl

@@ -48,37 +48,44 @@ end
     @testset "2x2 StaticMatrix with StaticVector" begin
         m = @SMatrix [1 2; 3 4]
         v = SVector(1, 4)
-        @test @inferred(broadcast(+, m, v)) === @SMatrix [2 3; 7 8]
-        @test @inferred(m .+ v) === @SMatrix [2 3; 7 8]
-        @test @inferred(v .+ m) === @SMatrix [2 3; 7 8]
-        @test @inferred(m .* v) === @SMatrix [1 2; 12 16]
-        @test @inferred(v .* m) === @SMatrix [1 2; 12 16]
-        @test @inferred(m ./ v) === @SMatrix [1 2; 3/4 1]
-        @test @inferred(v ./ m) === @SMatrix [1 1/2; 4/3 1]
-        @test @inferred(m .- v) === @SMatrix [0 1; -1 0]
-        @test @inferred(v .- m) === @SMatrix [0 -1; 1 0]
-        @test @inferred(m .^ v) === @SMatrix [1 2; 81 256]
-        @test @inferred(v .^ m) === @SMatrix [1 1; 64 256]
-        # Issue #546
-        @test @inferred(m ./ (v .* v')) === @SMatrix [1.0 0.5; 0.75 0.25]
-        testinf(m, v) = m ./ (v .* v')
-        @test @inferred(testinf(m, v)) === @SMatrix [1.0 0.5; 0.75 0.25]
+        vrep = @SMatrix [1 1; 4 4]
+        for m in (m, Transpose(m), Adjoint(m), Diagonal(m), Symmetric(m, :U), Symmetric(m, :L), Hermitian(m, :U), Hermitian(m, :L), UpperTriangular(m), LowerTriangular(m), UnitUpperTriangular(m), UnitLowerTriangular(m)) 
+            @test @inferred(broadcast(+, m, v)) === map(+, m, vrep)::SMatrix
+            @test @inferred(m .+ v) === map(+, m, vrep)::SMatrix
+            @test @inferred(v .+ m) === map(+, vrep, m)::SMatrix
+            @test @inferred(m .* v) === map(*, m, vrep)::SMatrix
+            @test @inferred(v .* m) === map(*, vrep, m)::SMatrix
+            @test @inferred(m ./ v) === map(/, m, vrep)::SMatrix
+            @test @inferred(v ./ m) === map(/, vrep, m)::SMatrix
+            @test @inferred(m .- v) === map(-, m, vrep)::SMatrix
+            @test @inferred(v .- m) === map(-, vrep, m)::SMatrix
+            @test @inferred(m .^ v) === map(^, m, vrep)::SMatrix
+            @test @inferred(v .^ m) === map(^, vrep, m)::SMatrix
+            # Issue #546
+            @test @inferred(m ./ (v .* v')) === map(/, m, v .* v')::SMatrix
+            testinf(m, v) = m ./ (v .* v')
+            @test @inferred(testinf(m, v)) === map(/, m, v .* v')::SMatrix
+        end
     end
 
     @testset "2x2 StaticMatrix with 1x2 StaticMatrix" begin
         # Issues #197, #242: broadcast between SArray and row-like SMatrix
         m1 = @SMatrix [1 2; 3 4]
         m2 = @SMatrix [1 4]
-        @test @inferred(broadcast(+, m1, m2)) === @SMatrix [2 6; 4 8]
-        @test @inferred(m1 .+ m2) === @SMatrix [2 6; 4 8]
-        @test @inferred(m2 .+ m1) === @SMatrix [2 6; 4 8]
-        @test @inferred(m1 .* m2) === @SMatrix [1 8; 3 16]
-        @test @inferred(m2 .* m1) === @SMatrix [1 8; 3 16]
-        @test @inferred(m1 ./ m2) === @SMatrix [1 1/2; 3 1]
-        @test @inferred(m2 ./ m1) === @SMatrix [1 2; 1/3 1]
-        @test @inferred(m1 .- m2) === @SMatrix [0 -2; 2 0]
-        @test @inferred(m2 .- m1) === @SMatrix [0 2; -2 0]
-        @test @inferred(m1 .^ m2) === @SMatrix [1 16; 3 256]
+        m2rep = @SMatrix [1 4; 1 4]
+        m1s = (m1, Transpose(m1), Adjoint(m1), Diagonal(m1), Symmetric(m1, :U), Symmetric(m1, :L), Hermitian(m1, :U), Hermitian(m1, :L), UpperTriangular(m1), LowerTriangular(m1), UnitUpperTriangular(m1), UnitLowerTriangular(m1))
+        for m1 in m1s
+            @test @inferred(broadcast(+, m1, m2)) === map(+, m1, m2rep)::SMatrix
+            @test @inferred(m1 .+ m2) === map(+, m1, m2rep)::SMatrix
+            @test @inferred(m2 .+ m1) === map(+, m2rep, m1)::SMatrix
+            @test @inferred(m1 .* m2) === map(*, m1, m2rep)::SMatrix
+            @test @inferred(m2 .* m1) === map(*, m2rep, m1)::SMatrix
+            @test @inferred(m1 ./ m2) === map(/, m1, m2rep)::SMatrix
+            @test @inferred(m2 ./ m1) === map(/, m2rep, m1)::SMatrix
+            @test @inferred(m1 .- m2) === map(-, m1, m2rep)::SMatrix
+            @test @inferred(m2 .- m1) === map(-, m2rep, m1)::SMatrix
+            @test @inferred(m1 .^ m2) === map(^, m1, m2rep)::SMatrix
+        end
     end
 
     @testset "1x2 StaticMatrix with StaticVector" begin