Overload axes for ScaledMap and LinearCombination (#162)

Author: krcools

Project.toml

@@ -1,6 +1,6 @@
 name = "LinearMaps"
 uuid = "7a12625a-238d-50fd-b39a-03d52299707e"
-version = "3.4.1"
+version = "3.5.0"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -12,11 +12,12 @@ julia = "1"
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Quaternions = "94ee1d12-ae83-5a48-8b1c-48b8ff168ae0"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "BenchmarkTools", "InteractiveUtils", "LinearAlgebra", "Quaternions", "SparseArrays", "Test"]
+test = ["Aqua", "BenchmarkTools", "BlockArrays", "InteractiveUtils", "LinearAlgebra", "Quaternions", "SparseArrays", "Test"]

docs/src/history.md

@@ -1,5 +1,14 @@
 # Version history
 
+## 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)`.
+
 ## What's new in v3.4
 
 * In `WrappedMap` constructors, as implicitly called in addition and mutliplication

src/LinearMaps.jl

@@ -111,7 +111,9 @@ julia> A(x)
 """
 function Base.:(*)(A::LinearMap, x::AbstractVector)
     check_dim_mul(A, x)
-    return mul!(similar(x, promote_type(eltype(A), eltype(x)), size(A, 1)), A, x)
+    T = promote_type(eltype(A), eltype(x))
+    y = similar(x, T, axes(A)[1])
+    return mul!(y, A, x)
 end
 if VERSION ≥ v"1.3"
     (L::LinearMap)(x::AbstractVector) = L*x

src/linearcombination.jl

@@ -2,10 +2,10 @@ struct LinearCombination{T, As<:LinearMapTuple} <: LinearMap{T}
     maps::As
     function LinearCombination{T, As}(maps::As) where {T, As}
         N = length(maps)
-        sz = size(maps[1])
+        ax = axes(maps[1])
         for n in eachindex(maps)
             A = maps[n]
-            size(A) == sz || throw(DimensionMismatch("LinearCombination"))
+            axes(A) == ax || throw(DimensionMismatch("LinearCombination"))
             @assert promote_type(T, eltype(A)) == T  "eltype $(eltype(A)) cannot be promoted to $T in LinearCombination constructor"
         end
         new{T, As}(maps)
@@ -18,6 +18,7 @@ MulStyle(A::LinearCombination) = MulStyle(A.maps...)
 
 # basic methods
 Base.size(A::LinearCombination) = size(A.maps[1])
+Base.axes(A::LinearMaps.LinearCombination) = axes(A.maps[1])
 # following conditions are sufficient but not necessary
 LinearAlgebra.issymmetric(A::LinearCombination) = all(issymmetric, A.maps)
 LinearAlgebra.ishermitian(A::LinearCombination) = all(ishermitian, A.maps)

src/scaledmap.jl

@@ -18,6 +18,7 @@ ScaledMap(λ::RealOrComplex, lmap::LinearMap{<:RealOrComplex}) =
 
 # basic methods
 Base.size(A::ScaledMap) = size(A.lmap)
+Base.axes(A::ScaledMap) = axes(A.lmap)
 Base.isreal(A::ScaledMap) = isreal(A.λ) && isreal(A.lmap)
 LinearAlgebra.issymmetric(A::ScaledMap) = issymmetric(A.lmap)
 LinearAlgebra.ishermitian(A::ScaledMap) = isreal(A.λ) && ishermitian(A.lmap)

src/wrappedmap.jl

@@ -59,6 +59,7 @@ Base.:(==)(A::VecOrMatMap, B::VecOrMatMap) =
 # properties
 Base.size(A::WrappedMap) = size(A.lmap)
 Base.size(A::WrappedMap{<:Any,<:AbstractVector}) = (Int(length(A.lmap))::Int, 1)
+Base.axes(A::WrappedMap) = axes(A.lmap)
 LinearAlgebra.issymmetric(A::WrappedMap) = A._issymmetric
 LinearAlgebra.ishermitian(A::WrappedMap) = A._ishermitian
 LinearAlgebra.isposdef(A::WrappedMap) = A._isposdef

test/nontradaxes.jl

@@ -0,0 +1,44 @@
+using Test, LinearMaps, LinearAlgebra, BlockArrays
+
+@testset "Non-traditional axes" begin
+
+    A = rand(ComplexF64,2,4)
+    B = PseudoBlockMatrix{ComplexF64}(undef, [2,3], [3,4])
+
+    ax1 = axes(B)[1]
+    ax2 = axes(B)[2]
+    fill!(B,0)
+    B[Block(1),Block(2)] .= A
+
+    N = @inferred LinearMap(B)
+    @test axes(N) == (ax1,ax2)
+
+    @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]))
+
+    for i in eachindex(u) u[i] = rand(ComplexF64) end
+    for i in eachindex(v) v[i] = rand(ComplexF64) end
+    for i in eachindex(w) w[i] = rand(ComplexF64) end
+
+    @test B*u == N*u
+    @test_throws DimensionMismatch N*v
+
+    # Lu = L*u
+    Nu = N*u
+
+    @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]
+
+    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)
+end

test/runtests.jl

@@ -36,3 +36,7 @@ include("conversion.jl")
 include("left.jl")
 
 include("fillmap.jl")
+
+if VERSION ≥ v"1.1"
+    include("nontradaxes.jl")
+end