allow dispatching on axes for dot (#78)

* allow dispatching on axes for dot

* Update Project.toml

* add adtrans mul tests

Author: dlfivefifty

Project.toml

@@ -15,7 +15,7 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 ArrayLayouts = "0.7"
 DomainSets = "0.5"
 FillArrays = "0.11, 0.12"
-LazyArrays = "0.21.5"
+LazyArrays = "0.21.5, 0.22"
 StaticArrays = "1"
 julia = "1.6"
 

src/QuasiArrays.jl

@@ -38,7 +38,7 @@ import Base.Broadcast: materialize, materialize!, BroadcastStyle, AbstractArrayS
 
 import LinearAlgebra: transpose, adjoint, checkeltype_adjoint, checkeltype_transpose, Diagonal,
                         AbstractTriangular, pinv, inv, promote_leaf_eltypes, power_by_squaring,
-                        integerpow, schurpow, tr, factorize, copy_oftype, rank
+                        integerpow, schurpow, tr, factorize, copy_oftype, rank, dot
 
 import ArrayLayouts: indextype, concretize, fillzeros, OnesLayout, AbstractFillLayout, FillLayout, ZerosLayout, diagonallayout, diagonaldata, diagonal
 import LazyArrays: MemoryLayout, UnknownLayout, Mul, ApplyLayout, BroadcastLayout,

src/quasiadjtrans.jl

@@ -167,12 +167,21 @@ map(f, tvs::QuasiTransposeAbsVec...) = transpose(map((xs...) -> transpose(f(tran
 ## multiplication *
 
 # QuasiAdjoint/QuasiTranspose-vector * vector
+
+# allow re-expansion using Legendre
+_dot(ax, a, b) = Base.invoke(dot, NTuple{2,Any}, a, b)
+
+function dot(a::AbstractQuasiArray, b::AbstractQuasiArray)
+    ax = axes(a,1)
+    ax == axes(b,1) || throw(DimensionMismatch())
+    _dot(ax, a, b)
+end
+
 *(u::QuasiAdjointAbsVec, v::AbstractQuasiVector) = dot(u.parent, v)
 *(u::QuasiTransposeAbsVec{T}, v::AbstractQuasiVector{T}) where {T<:Real} = dot(u.parent, v)
 function *(u::QuasiTransposeAbsVec, v::AbstractQuasiVector)
-    @assert !has_offset_axes(u, v)
-    @boundscheck length(u) == length(v) || throw(DimensionMismatch())
-    return sum(@inbounds(u[k]*v[k]) for k in 1:length(u))
+    @boundscheck axes(u,2) == axes(v,1) || throw(DimensionMismatch())
+    return sum(@inbounds(u.parent[k]*v[k]) for k in axes(v,1))
 end
 # vector * QuasiAdjoint/QuasiTranspose-vector
 *(u::AbstractQuasiVector, v::AdjOrTransAbsVec) = broadcast(*, u, v)
@@ -229,4 +238,18 @@ arguments(LAY::ApplyLayout{typeof(*)}, V::QuasiTranspose) = reverse(transpose.(a
 
 # This is used in ContinuumArrays.jl to ensure x' is lazy
 BroadcastStyle(::Type{<:QuasiAdjoint{<:Any,<:Inclusion}}) = LazyQuasiArrayStyle{2}()
-BroadcastStyle(::Type{<:QuasiTranspose{<:Any,<:Inclusion}}) = LazyQuasiArrayStyle{2}()
+BroadcastStyle(::Type{<:QuasiTranspose{<:Any,<:Inclusion}}) = LazyQuasiArrayStyle{2}()
+
+
+
+###
+# adjoint concat support
+###
+
+arguments(::ApplyLayout{typeof(vcat)}, A::QuasiAdjoint) = map(adjoint, arguments(ApplyLayout{typeof(hcat)}(), parent(A)))
+arguments(::ApplyLayout{typeof(hcat)}, A::QuasiAdjoint) = map(adjoint, arguments(ApplyLayout{typeof(vcat)}(), parent(A)))
+arguments(::ApplyLayout{typeof(vcat)}, A::QuasiTranspose) = map(transpose, arguments(ApplyLayout{typeof(hcat)}(), parent(A)))
+arguments(::ApplyLayout{typeof(hcat)}, A::QuasiTranspose) = map(transpose, arguments(ApplyLayout{typeof(vcat)}(), parent(A)))
+
+
+copy(M::Mul{ApplyLayout{typeof(vcat)},QuasiArrayLayout}) = vcat((arguments(vcat, M.A) .* Ref(M.B))...)

src/quasiconcat.jl

@@ -71,4 +71,4 @@ function _getindex(::Type{IND}, A::UnionVcat{T,2}, (x,j)::IND) where {IND,T}
         x in axes(a,1) && return convert(T,a[x,j])::T
     end
     throw(BoundsError(A, I))
-end
+end

test/test_matmul.jl

@@ -1,4 +1,4 @@
-using QuasiArrays, ArrayLayouts, Test
+using QuasiArrays, ArrayLayouts, LinearAlgebra, Test
 import QuasiArrays: apply
 
 @testset "Multiplication" begin
@@ -35,4 +35,12 @@ import QuasiArrays: apply
         @test Array(A*A) ≈ Array(A)^2
         @test Array(A*A*A) ≈ Array(A)^3
     end
+
+    @testset "absvec" begin
+        a = QuasiArray(rand(3),(0:0.5:1,))
+        c = QuasiArray(rand(3) .+ im .* randn(3),(0:0.5:1,))
+        @test a'a ≈ transpose(a)a ≈ dot(a,a)
+        @test c'c ≈ dot(c,c)
+        @test transpose(c)c ≈ transpose(c.parent)*c.parent
+    end
 end