Specialisation of mul! for left multiplication with scalars (#173)

Co-authored-by: Daniel Karrasch <[email protected]>

Author: krcools

docs/Project.toml

@@ -8,4 +8,4 @@ Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 [compat]
 BenchmarkTools = "1"
 Documenter = "0.25, 0.26, 0.27"
-Literate = "2"
+Literate = "2"

docs/src/custom.jl

@@ -29,10 +29,13 @@ end
 Base.size(A::MyFillMap) = A.size
 
 # By a couple of defaults provided for all subtypes of `LinearMap`, we only need to define
-# a `LinearAlgebra.mul!` method to have minimal, operational type.
+# a `LinearMaps._unsafe_mul!` method to have a minimal, operational type. The (internal)
+# function `_unsafe_mul!` is called by `LinearAlgebra.mul!`, constructors, and conversions
+# and only needs to be concerned with the bare computing kernel. Dimension checking is done
+# on the level of `mul!` etc. Factoring out dimension checking is done to minimise overhead
+# caused by repetitive checking.
 
-function LinearAlgebra.mul!(y::AbstractVecOrMat, A::MyFillMap, x::AbstractVector)
-    LinearMaps.check_dim_mul(y, A, x)
+function LinearMaps._unsafe_mul!(y::AbstractVecOrMat, A::MyFillMap, x::AbstractVector)
     return fill!(y, iszero(A.λ) ? zero(eltype(y)) : A.λ*sum(x))
 end
 
@@ -84,7 +87,7 @@ using BenchmarkTools
 
 LinearMaps.MulStyle(A::MyFillMap) = FiveArg()
 
-function LinearAlgebra.mul!(
+function LinearMaps._unsafe_mul!(
     y::AbstractVecOrMat,
     A::MyFillMap,
     x::AbstractVector,
@@ -126,7 +129,7 @@ typeof(A')
 try A'x catch e println(e) end
 
 # If the operator is symmetric or Hermitian, the transpose and the adjoint, respectively,
-# of the linear map `A` is given by `A` itself. So let's define corresponding checks.
+# of the linear map `A` is given by `A` itself. So let us define corresponding checks.
 
 LinearAlgebra.issymmetric(A::MyFillMap) = A.size[1] == A.size[2]
 LinearAlgebra.ishermitian(A::MyFillMap) = isreal(A.λ) && A.size[1] == A.size[2]
@@ -154,22 +157,35 @@ try MyFillMap(5.0, (3, 4))' * ones(3) catch e println(e) end
 # The first option is to write `LinearAlgebra.mul!` methods for the corresponding wrapped
 # map types; for instance,
 
-function LinearAlgebra.mul!(
+function LinearMaps._unsafe_mul!(
     y::AbstractVecOrMat,
     transA::LinearMaps.TransposeMap{<:Any,<:MyFillMap},
     x::AbstractVector
 )
-    LinearMaps.check_dim_mul(y, transA, x)
     λ = transA.lmap.λ
     return fill!(y, iszero(λ) ? zero(eltype(y)) : transpose(λ)*sum(x))
 end
 
+# Now, the adjoint multiplication works.
+
+MyFillMap(5.0, (3, 4))' * ones(3)
+
 # If you have set the `MulStyle` trait to `FiveArg()`, you should provide a corresponding
 # 5-arg `mul!` method for `LinearMaps.TransposeMap{<:Any,<:MyFillMap}` and
 # `LinearMaps.AdjointMap{<:Any,<:MyFillMap}`.
 
 # ### Path 2: Invariant `LinearMap` subtypes
 
+# Before we start, let us delete the previously defined method to make sure we use the
+# following definitions.
+
+Base.delete_method(
+    first(methods(
+        LinearMaps._unsafe_mul!,
+        (AbstractVecOrMat, LinearMaps.TransposeMap{<:Any,<:MyFillMap}, AbstractVector))
+    )
+)
+
 # The seconnd option is when your class of linear maps that are modelled by your custom
 # `LinearMap` subtype are invariant under taking adjoints and transposes.
 

docs/src/history.md

@@ -1,6 +1,14 @@
 # Version history
 
 ## What's new in v3.7
+* `mul!(M::AbstractMatrix, A::LinearMap, s::Number, a, b)` methods are provided, mimicking
+  similar methods in `Base.LinearAlgebra`. This version allows for the memory efficient
+  implementation of in-place addition and conversion of a `LinearMap` to `Matrix`.
+  Efficient specialisations for `WrappedMap`, `ScaledMap`, and `LinearCombination` are
+  provided. If users supply the corresponding `_unsafe_mul!` method for their custom maps,
+  conversion, construction, and inplace addition will benefit from this supplied efficient
+  implementation. If no specialisation is supplied, a generic fallback is used that is based
+  on feeding the canonical basis of unit vectors to the linear map.
 
 * A new map type called `EmbeddedMap` is introduced. It is a wrapper of a "small" `LinearMap`
   (or a suitably converted `AbstractVecOrMat`) embedded into a "larger" zero map. Hence,

docs/src/types.md

@@ -103,6 +103,8 @@ Base.:*(::AbstractMatrix,::LinearMap)
 LinearAlgebra.mul!(::AbstractVecOrMat,::LinearMap,::AbstractVector)
 LinearAlgebra.mul!(::AbstractVecOrMat,::LinearMap,::AbstractVector,::Number,::Number)
 LinearAlgebra.mul!(::AbstractMatrix,::AbstractMatrix,::LinearMap)
+LinearAlgebra.mul!(::AbstractVecOrMat,::LinearMap,::Number)
+LinearAlgebra.mul!(::AbstractVecOrMat,::LinearMap,::Number,::Number,::Number)
 *(::LinearAlgebra.AdjointAbsVec,::LinearMap)
 *(::LinearAlgebra.TransposeAbsVec,::LinearMap)
 ```

src/LinearMaps.jl

@@ -143,14 +143,14 @@ with either `A` or `B`.
 
 ## Examples
 ```jldoctest; setup=(using LinearAlgebra, LinearMaps)
-julia> A=LinearMap([1.0 2.0; 3.0 4.0]); B=[1.0, 1.0]; Y = similar(B); mul!(Y, A, B);
+julia> A=LinearMap([1.0 2.0; 3.0 4.0]); B=ones(2); Y = similar(B); mul!(Y, A, B);
 
 julia> Y
 2-element Array{Float64,1}:
  3.0
  7.0
 
-julia> A=LinearMap([1.0 2.0; 3.0 4.0]); B=[1.0 1.0; 1.0 1.0]; Y = similar(B); mul!(Y, A, B);
+julia> A=LinearMap([1.0 2.0; 3.0 4.0]); B=ones(4,4); Y = similar(B); mul!(Y, A, B);
 
 julia> Y
 2×2 Array{Float64,2}:
@@ -162,6 +162,35 @@ function mul!(y::AbstractVecOrMat, A::LinearMap, x::AbstractVector)
     check_dim_mul(y, A, x)
     return _unsafe_mul!(y, A, x)
 end
+# the following is of interest in, e.g., subspace-iteration methods
+function mul!(Y::AbstractMatrix, A::LinearMap, X::AbstractMatrix)
+    check_dim_mul(Y, A, X)
+    return _unsafe_mul!(Y, A, X)
+end
+
+"""
+    mul!(Y::AbstractMatrix, A::LinearMap, b::Number) -> Y
+
+Scales the matrix representation of the linear map `A` by `b` and stores the result in `Y`,
+overwriting the existing value of `Y`.
+
+## Examples
+```jldoctest; setup=(using LinearAlgebra, LinearMaps)
+julia> A = LinearMap{Int}(cumsum, 3); b = 2; Y = Matrix{Int}(undef, (3,3));
+
+julia> mul!(Y, A, b)
+3×3 Matrix{Int64}:
+ 2  0  0
+ 2  2  0
+ 2  2  2
+```
+"""
+function mul!(y::AbstractVecOrMat, A::LinearMap, s::Number)
+    size(y) == size(A) ||     
+        throw(
+            DimensionMismatch("y has size $(size(y)), A has size $(size(A))."))
+    return _unsafe_mul!(y, A, s)
+end
 
 """
     mul!(C::AbstractVecOrMat, A::LinearMap, B::AbstractVector, α, β) -> C
@@ -197,6 +226,46 @@ function mul!(y::AbstractVecOrMat, A::LinearMap, x::AbstractVector, α::Number,
     check_dim_mul(y, A, x)
     return _unsafe_mul!(y, A, x, α, β)
 end
+function mul!(Y::AbstractMatrix, A::LinearMap, X::AbstractMatrix, α::Number, β::Number)
+    check_dim_mul(Y, A, X)
+    return _unsafe_mul!(Y, A, X, α, β)
+end
+
+"""
+    mul!(Y::AbstractMatrix, A::LinearMap, b::Number, α::Number, β::Number) -> Y
+
+Scales the matrix representation of the linear map `A` by `b*α`, adds the result to `Y*β`
+and stores the final result in `Y`, overwriting the existing value of `Y`.
+
+## Examples
+```jldoctest; setup=(using LinearAlgebra, LinearMaps)
+julia> A = LinearMap{Int}(cumsum, 3); b = 2; Y = ones(Int, (3,3));
+
+julia> mul!(Y, A, b, 2, 1)
+3×3 Matrix{Int64}:
+ 5  1  1
+ 5  5  1
+ 5  5  5
+```
+"""
+function mul!(y::AbstractMatrix, A::LinearMap, s::Number, α::Number, β::Number)
+    size(y) == size(A) ||     
+        throw(
+            DimensionMismatch("y has size $(size(y)), A has size $(size(A))."))
+    return _unsafe_mul!(y, A, s, α, β)
+end
+
+_unsafe_mul!(y, A::MapOrVecOrMat, x) = mul!(y, A, x)
+_unsafe_mul!(y, A::AbstractVecOrMat, x, α, β) = mul!(y, A, x, α, β)
+_unsafe_mul!(y::AbstractVecOrMat, A::LinearMap, x::AbstractVector, α, β) =
+    _generic_mapvec_mul!(y, A, x, α, β)
+_unsafe_mul!(y::AbstractMatrix, A::LinearMap, x::AbstractMatrix) =
+    _generic_mapmat_mul!(y, A, x)
+_unsafe_mul!(y::AbstractMatrix, A::LinearMap, x::AbstractMatrix, α::Number, β::Number) =
+    _generic_mapmat_mul!(y, A, x, α, β)
+_unsafe_mul!(Y::AbstractMatrix, A::LinearMap, s::Number) = _generic_mapnum_mul!(Y, A, s)
+_unsafe_mul!(Y::AbstractMatrix, A::LinearMap, s::Number, α::Number, β::Number) =
+    _generic_mapnum_mul!(Y, A, s, α, β)
 
 function _generic_mapvec_mul!(y, A, x, α, β)
     # this function needs to call mul! for, e.g.,  AdjointMap{...,<:CustomMap}
@@ -226,33 +295,40 @@ function _generic_mapvec_mul!(y, A, x, α, β)
     end
 end
 
-# the following is of interest in, e.g., subspace-iteration methods
-function mul!(Y::AbstractMatrix, A::LinearMap, X::AbstractMatrix)
-    check_dim_mul(Y, A, X)
-    return _unsafe_mul!(Y, A, X)
-end
-function mul!(Y::AbstractMatrix, A::LinearMap, X::AbstractMatrix, α::Number, β::Number)
-    check_dim_mul(Y, A, X)
-    return _unsafe_mul!(Y, A, X, α, β)
+function _generic_mapmat_mul!(Y, A, X)
+    for (Xi, Yi) in zip(eachcol(X), eachcol(Y))
+        mul!(Yi, A, Xi)
+    end
+    return Y
 end
-
-function _generic_mapmat_mul!(Y, A, X, α=true, β=false)
+function _generic_mapmat_mul!(Y, A, X, α, β)
     for (Xi, Yi) in zip(eachcol(X), eachcol(Y))
         mul!(Yi, A, Xi, α, β)
     end
     return Y
 end
 
-_unsafe_mul!(y, A::MapOrVecOrMat, x) = mul!(y, A, x)
-_unsafe_mul!(y, A::AbstractVecOrMat, x, α, β) = mul!(y, A, x, α, β)
-function _unsafe_mul!(y::AbstractVecOrMat, A::LinearMap, x::AbstractVector, α, β)
-    return _generic_mapvec_mul!(y, A, x, α, β)
-end
-function _unsafe_mul!(y::AbstractMatrix, A::LinearMap, x::AbstractMatrix)
-    return _generic_mapmat_mul!(y, A, x)
+function _generic_mapnum_mul!(Y, A, s)
+    T = promote_type(eltype(A), typeof(s))
+    ax2 = axes(A)[2]
+    xi = zeros(T, ax2)
+    @inbounds for (i, Yi) in zip(ax2, eachcol(Y))
+        xi[i] = s
+        mul!(Yi, A, xi)
+        xi[i] = zero(T)
+    end
+    return Y
 end
-function _unsafe_mul!(y::AbstractMatrix, A::LinearMap, x::AbstractMatrix, α, β)
-    return _generic_mapmat_mul!(y, A, x, α, β)
+function _generic_mapnum_mul!(Y, A, s, α, β)
+    T = promote_type(eltype(A), typeof(s))
+    ax2 = axes(A)[2]
+    xi = zeros(T, ax2)
+    @inbounds for (i, Yi) in zip(ax2, eachcol(Y))
+        xi[i] = s
+        mul!(Yi, A, xi, α, β)
+        xi[i] = zero(T)
+    end
+    return Y
 end
 
 include("left.jl") # left multiplication by a transpose or adjoint vector