Specialize Base.axes for [Composite/Adjoint/Transpose]Maps (#164)

Author: dkarrasch

Project.toml

@@ -1,6 +1,6 @@
 name = "LinearMaps"
 uuid = "7a12625a-238d-50fd-b39a-03d52299707e"
-version = "3.5.0"
+version = "3.5.1"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

docs/src/history.md

@@ -2,12 +2,13 @@
 
 ## What's new in v3.5
 
-* `WrappedMap`, `ScaledMap`, and `LinearCombination`, instead of using the default `axes(A) 
-  = map(oneto, size(A))`, now forward calls to `axes` to the underlying wrapped linear map. 
-  This allows allocating operations such as `*` to determine the appropriate storage and axes 
-  type of their outputs. For example, linear maps that wrap `BlockArrays` will, upon
-  multiplicative action, produce a `BlockArrays.PseudoBlockVector` with block structure
-  inherited from the operator's *output* axes `axes(A,1)`.
+* `WrappedMap`, `ScaledMap`, `LinearCombination`, `AdjointMap`, `TransposeMap` and
+  `CompositeMap`, instead of using the default `axes(A) = map(oneto, size(A))`, now forward
+  calls to `axes` to the underlying wrapped linear map. This allows allocating operations
+  such as `*` to determine the appropriate storage and axes type of their outputs.
+  For example, linear maps that wrap `BlockArrays` will, upon multiplicative action,
+  produce a `BlockArrays.PseudoBlockVector` with block structure inherited from the
+  operator's *output* axes `axes(A,1)`.
 
 ## What's new in v3.4
 

src/LinearMaps.jl

@@ -49,7 +49,9 @@ LinearAlgebra.isposdef(::LinearMap) = false # default assumptions
 Base.ndims(::LinearMap) = 2
 Base.size(A::LinearMap, n) =
     (n == 1 || n == 2 ? size(A)[n] : error("LinearMap objects have only 2 dimensions"))
-Base.length(A::LinearMap) = size(A)[1] * size(A)[2]
+Base.axes(A::LinearMap, n::Integer) =
+    (n == 1 || n == 2 ? axes(A)[n] : error("LinearMap objects have only 2 dimensions"))
+Base.length(A::LinearMap) = prod(size(A))
 
 # check dimension consistency for multiplication A*B
 _iscompatible((A, B)) = size(A, 2) == size(B, 1)

src/composition.jl

@@ -17,6 +17,7 @@ CompositeMap{T}(maps::As) where {T, As<:LinearMapTuple} = CompositeMap{T, As}(ma
 
 # basic methods
 Base.size(A::CompositeMap) = (size(A.maps[end], 1), size(A.maps[1], 2))
+Base.axes(A::CompositeMap) = (axes(A.maps[end])[1], axes(A.maps[1])[2])
 Base.isreal(A::CompositeMap) = all(isreal, A.maps) # sufficient but not necessary
 
 # the following rules are sufficient but not necessary

src/transpose.jl

@@ -34,6 +34,7 @@ LinearAlgebra.adjoint(A::LinearMap{<:Real}) = transpose(A)
 
 # properties
 Base.size(A::Union{TransposeMap, AdjointMap}) = reverse(size(A.lmap))
+Base.axes(A::Union{TransposeMap, AdjointMap}) = reverse(axes(A.lmap))
 LinearAlgebra.issymmetric(A::Union{TransposeMap, AdjointMap}) = issymmetric(A.lmap)
 LinearAlgebra.ishermitian(A::Union{TransposeMap, AdjointMap}) = ishermitian(A.lmap)
 LinearAlgebra.isposdef(A::Union{TransposeMap, AdjointMap}) = isposdef(A.lmap)

test/nontradaxes.jl

@@ -12,14 +12,15 @@ using Test, LinearMaps, LinearAlgebra, BlockArrays
 
     N = @inferred LinearMap(B)
     @test axes(N) == (ax1,ax2)
+    @test axes(N, 1) == ax1
+    @test axes(N, 2) == ax2
+    @test_throws ErrorException axes(N, 3)
 
     @test eltype(N) == eltype(B)
 
     u = similar(Array{ComplexF64}, ax2)
     v = PseudoBlockVector{ComplexF64}(undef, [3,5])
-    w = PseudoBlockVector{ComplexF64}(undef, [4,3])
-    # v = similar(Array{ComplexF64}, blockedrange([3,5]))
-    # w = similar(Array{ComplexF64}, blockedrange([4,3]))
+    w = similar(Array{ComplexF64}, ax1)
 
     for i in eachindex(u) u[i] = rand(ComplexF64) end
     for i in eachindex(v) v[i] = rand(ComplexF64) end
@@ -34,11 +35,39 @@ using Test, LinearMaps, LinearAlgebra, BlockArrays
     @test axes(Nu)[1] == axes(N)[1] == axes(B)[1]
     @test blocklengths(axes(Nu)[1]) == blocklengths(axes(N)[1]) == blocklengths(axes(B)[1]) == [2,3]
 
+    for trans in (adjoint, transpose)
+        Nt = trans(LinearMap(N))
+        Bt = trans(B)
+        Ntw = Nt*w
+        Btw = Bt*w
+        @test Ntw ≈ Btw
+        @test axes(Ntw)[1] == axes(Nt)[1] == axes(Bt)[1]
+        @test blocklengths(axes(Ntw)[1]) == blocklengths(axes(Nt)[1]) == blocklengths(axes(Bt)[1]) == [3,4]
+    end
+
     C = B + 2N
     @test axes(C) === axes(B) === axes(N)
     @test C*u ≈ 3*Nu
 
     Cu = C*u
     @test axes(C)[1] == ax1
     @test blocklengths(axes(C)[1]) == blocklengths(ax1)
+
+    A = rand(ComplexF64,2,2)
+    B = PseudoBlockMatrix{ComplexF64}(undef, [2,2], [2,2])
+    ax1 = axes(B)[1]
+    ax2 = axes(B)[2]
+    fill!(B,0)
+    B[Block(1),Block(2)] .= A
+    L = LinearMap(B)
+    L2 = L*L
+    B2 = B*B
+    @test axes(L2) == axes(B2)
+    B2 = B*Matrix(B)
+    L2 = L*LinearMap(Matrix(B))
+    @test axes(L2) == axes(B2)
+    u = similar(Array{ComplexF64}, ax2)
+    B2u = B2*u; L2u = L2*u
+    @test axes(B2u)[1] == axes(L2u)[1] == axes(B2)[1] == axes(L2)[1]
+    @test blocklengths(axes(B2u)[1]) == blocklengths(axes(L2u)[1]) == [2,2]
 end