Improve inferrability

There is no way to make the `S` keyword argument of `LinearOperator{T}`
inferrable.

This moves `S` to the second position, so anyone who wishes to write
inferrable code can use `LinearOperator{T, Vector{T}}(args...)`.

More inferrability & performance work

Author: timholy

ext/LinearOperatorsLDLFactorizationsExt.jl

@@ -13,7 +13,7 @@ function LinearOperators.opLDL(M::AbstractMatrix; check::Bool = false)
   tprod! = @closure (res, u, α, β) -> LinearOperators.tmulFact!(res, LDL, u, α, β)  # M.' = conj(M)
   ctprod! = @closure (res, w, α, β) -> LinearOperators.mulFact!(res, LDL, w, α, β)
   S = eltype(LDL)
-  return LinearOperator{S}(m, m, isreal(M), true, prod!, tprod!, ctprod!)
+  return LinearOperator{S, Vector{S}}(m, m, isreal(M), true, prod!, tprod!, ctprod!)
   #TODO: use iterative refinement.
 end
 
@@ -31,7 +31,7 @@ function LinearOperators.opLDL(
   tprod! = @closure (res, u) -> ldiv!(res, LDL, u)  # M.' = conj(M)
   ctprod! = @closure (res, w) -> ldiv!(res, LDL, w)
   S = eltype(LDL)
-  return LinearOperator{S}(m, m, isreal(M), true, prod!, tprod!, ctprod!)
+  return LinearOperator{S, Vector{S}}(m, m, isreal(M), true, prod!, tprod!, ctprod!)
 end
 
 end # module

src/abstract.jl

@@ -43,7 +43,7 @@ to combine or otherwise alter them. They can be combined with
 other operators, with matrices and with scalars. Operators may
 be transposed and conjugate-transposed using the usual Julia syntax.
 """
-mutable struct LinearOperator{T, I <: Integer, F, Ft, Fct, S} <: AbstractLinearOperator{T}
+mutable struct LinearOperator{T, S, I <: Integer, F, Ft, Fct} <: AbstractLinearOperator{T}
   nrow::I
   ncol::I
   symmetric::Bool
@@ -61,7 +61,7 @@ mutable struct LinearOperator{T, I <: Integer, F, Ft, Fct, S} <: AbstractLinearO
   allocated5::Bool # true for 5-args mul!, false for 3-args mul! until the vectors are allocated
 end
 
-function LinearOperator{T}(
+function LinearOperator{T, S}(
   nrow::I,
   ncol::I,
   symmetric::Bool,
@@ -71,16 +71,15 @@ function LinearOperator{T}(
   ctprod!::Fct,
   nprod::I,
   ntprod::I,
-  nctprod::I;
-  S::Type = Vector{T},
-) where {T, I <: Integer, F, Ft, Fct}
+  nctprod::I
+) where {T, S, I <: Integer, F, Ft, Fct}
   Mv5, Mtu5 = S(undef, 0), S(undef, 0)
   nargs = get_nargs(prod!)
   args5 = (nargs == 4)
   (args5 == false) || (nargs != 2) || throw(LinearOperatorException("Invalid number of arguments"))
   allocated5 = args5 ? true : false
   use_prod5! = args5 ? true : false
-  return LinearOperator{T, I, F, Ft, Fct, S}(
+  return LinearOperator{T, S, I, F, Ft, Fct}(
     nrow,
     ncol,
     symmetric,
@@ -99,19 +98,27 @@ function LinearOperator{T}(
   )
 end
 
-LinearOperator{T}(
+# backward compatibility (not inferrable; use LinearOperator{T, S} if you want something inferrable)
+LinearOperator{T}(nrow, ncol, symmetric, hermitian, prod!, tprod!, ctprod!, nprod, ntprod, nctprod; S::Type = Vector{T}) where {T} =
+  LinearOperator{T, S}(nrow, ncol, symmetric, hermitian, prod!, tprod!, ctprod!, nprod, ntprod, nctprod)
+
+LinearOperator{T, S}(
   nrow::I,
   ncol::I,
   symmetric::Bool,
   hermitian::Bool,
   prod!,
   tprod!,
-  ctprod!;
-  S::Type = Vector{T},
-) where {T, I <: Integer} =
-  LinearOperator{T}(nrow, ncol, symmetric, hermitian, prod!, tprod!, ctprod!, 0, 0, 0, S = S)
+  ctprod!
+) where {T, S, I <: Integer} =
+  LinearOperator{T, S}(nrow, ncol, symmetric, hermitian, prod!, tprod!, ctprod!, 0, 0, 0)
+
+# backward compatibility (not inferrable)
+LinearOperator{T}(nrow, ncol, symmetric, hermitian, prod!, tprod!, ctprod!; S::Type = Vector{T}) where {T} =
+  LinearOperator{T, S}(nrow, ncol, symmetric, hermitian, prod!, tprod!, ctprod!)
 
 # create operator from other operators with +, *, vcat,...
+# TODO: this is not a type, so it should not be uppercase
 function CompositeLinearOperator(
   T::Type,
   nrow::I,
@@ -121,13 +128,13 @@ function CompositeLinearOperator(
   prod!::F,
   tprod!::Ft,
   ctprod!::Fct,
-  args5::Bool;
-  S::Type = Vector{T},
-) where {I <: Integer, F, Ft, Fct}
+  args5::Bool,
+  ::Type{S},
+) where {S, I <: Integer, F, Ft, Fct}
   Mv5, Mtu5 = S(undef, 0), S(undef, 0)
   allocated5 = true
   use_prod5! = true
-  return LinearOperator{T, I, F, Ft, Fct, S}(
+  return LinearOperator{T, S, I, F, Ft, Fct}(
     nrow,
     ncol,
     symmetric,
@@ -146,6 +153,20 @@ function CompositeLinearOperator(
   )
 end
 
+# backward compatibility (not inferrable)
+CompositeLinearOperator(
+  T::Type,
+  nrow::I,
+  ncol::I,
+  symmetric::Bool,
+  hermitian::Bool,
+  prod!::F,
+  tprod!::Ft,
+  ctprod!::Fct,
+  args5::Bool;
+  S::Type = Vector{T}) where {I <: Integer, F, Ft, Fct} =
+  CompositeLinearOperator(T, nrow, ncol, symmetric, hermitian, prod!, tprod!, ctprod!, args5, S)
+
 nprod(op::AbstractLinearOperator) = op.nprod
 ntprod(op::AbstractLinearOperator) = op.ntprod
 nctprod(op::AbstractLinearOperator) = op.nctprod

src/cat.jl

@@ -48,7 +48,7 @@ function hcat(A::AbstractLinearOperator, B::AbstractLinearOperator)
   S = promote_type(storage_type(A), storage_type(B))
   isconcretetype(S) ||
     throw(LinearOperatorException("storage types cannot be promoted to a concrete type"))
-  CompositeLinearOperator(T, nrow, ncol, false, false, prod!, tprod!, ctprod!, args5, S = S)
+  CompositeLinearOperator(T, nrow, ncol, false, false, prod!, tprod!, ctprod!, args5, S)
 end
 
 function hcat(ops::AbstractLinearOperator...)
@@ -107,7 +107,7 @@ function vcat(A::AbstractLinearOperator, B::AbstractLinearOperator)
   S = promote_type(storage_type(A), storage_type(B))
   isconcretetype(S) ||
     throw(LinearOperatorException("storage types cannot be promoted to a concrete type"))
-  CompositeLinearOperator(T, nrow, ncol, false, false, prod!, tprod!, ctprod!, args5, S = S)
+  CompositeLinearOperator(T, nrow, ncol, false, false, prod!, tprod!, ctprod!, args5, S)
 end
 
 function vcat(ops::AbstractLinearOperator...)

src/constructors.jl

@@ -16,7 +16,7 @@ function LinearOperator(
   prod! = @closure (res, v, α, β) -> mul!(res, M, v, α, β)
   tprod! = @closure (res, u, α, β) -> mul!(res, transpose(M), u, α, β)
   ctprod! = @closure (res, w, α, β) -> mul!(res, adjoint(M), w, α, β)
-  LinearOperator{T}(nrow, ncol, symmetric, hermitian, prod!, tprod!, ctprod!, S = S)
+  LinearOperator{T, S}(nrow, ncol, symmetric, hermitian, prod!, tprod!, ctprod!)
 end
 
 """
@@ -43,7 +43,7 @@ end
 
 """
     LinearOperator(M::Hermitian{T}, S = Vector{T}) where {T}
-    
+
 Constructs a linear operator from a Hermitian matrix. If
 its elements are real, it is also symmetric.
 Change `S` to use LinearOperators on GPU.
@@ -57,7 +57,7 @@ end
     LinearOperator(type::Type{T}, nrow, ncol, symmetric, hermitian, prod!,
                     [tprod!=nothing, ctprod!=nothing],
                     S = Vector{T}) where {T}
-                    
+
 Construct a linear operator from functions where the type is specified as the first argument.
 Change `S` to use LinearOperators on GPU.
 Notice that the linear operator does not enforce the type, so using a wrong type can
@@ -67,29 +67,29 @@ A = [im 1.0; 0.0 1.0] # Complex matrix
 function mulOp!(res, M, v, α, β)
   mul!(res, M, v, α, β)
 end
-op = LinearOperator(Float64, 2, 2, false, false, 
-                    (res, v, α, β) -> mulOp!(res, A, v, α, β), 
-                    (res, u, α, β) -> mulOp!(res, transpose(A), u, α, β), 
+op = LinearOperator(Float64, 2, 2, false, false,
+                    (res, v, α, β) -> mulOp!(res, A, v, α, β),
+                    (res, u, α, β) -> mulOp!(res, transpose(A), u, α, β),
                     (res, w, α, β) -> mulOp!(res, A', w, α, β))
 Matrix(op) # InexactError
 ```
 The error is caused because `Matrix(op)` tries to create a Float64 matrix with the
 contents of the complex matrix `A`.
 
 Using `*` may generate a vector that contains `NaN` values.
-This can also happen if you use the 3-args `mul!` function with a preallocated vector such as 
+This can also happen if you use the 3-args `mul!` function with a preallocated vector such as
 `Vector{Float64}(undef, n)`.
 To fix this issue you will have to deal with the cases `β == 0` and `β != 0` separately:
 ```
 d1 = [2.0; 3.0]
 function mulSquareOpDiagonal!(res, d, v, α, β::T) where T
   if β == zero(T)
     res .= α .* d .* v
-  else 
+  else
     res .= α .* d .* v .+ β .* res
   end
 end
-op = LinearOperator(Float64, 2, 2, true, true, 
+op = LinearOperator(Float64, 2, 2, true, true,
                     (res, v, α, β) -> mulSquareOpDiagonal!(res, d, v, α, β))
 ```
 
@@ -98,7 +98,7 @@ In this case, using the 5-args `mul!` will generate storage vectors.
 
 ```
 A = rand(2, 2)
-op = LinearOperator(Float64, 2, 2, false, false, 
+op = LinearOperator(Float64, 2, 2, false, false,
                     (res, v) -> mul!(res, A, v),
                     (res, w) -> mul!(res, A', w))
 ```
@@ -114,7 +114,7 @@ function LinearOperator(
   prod!,
   tprod! = nothing,
   ctprod! = nothing;
-  S = Vector{T},
+  S::Type = Vector{T},
 ) where {T, I <: Integer}
-  return LinearOperator{T}(nrow, ncol, symmetric, hermitian, prod!, tprod!, ctprod!, S = S)
+  return LinearOperator{T, S}(nrow, ncol, symmetric, hermitian, prod!, tprod!, ctprod!)
 end

src/kron.jl

@@ -41,7 +41,7 @@ function kron(A::AbstractLinearOperator, B::AbstractLinearOperator)
   symm = issymmetric(A) && issymmetric(B)
   herm = ishermitian(A) && ishermitian(B)
   nrow, ncol = m * p, n * q
-  return LinearOperator{T}(nrow, ncol, symm, herm, prod!, tprod!, ctprod!)
+  return LinearOperator{T, Vector{T}}(nrow, ncol, symm, herm, prod!, tprod!, ctprod!)
 end
 
 kron(A::AbstractMatrix, B::AbstractLinearOperator) = kron(LinearOperator(A), B)

src/linalg.jl

@@ -28,14 +28,14 @@ function opInverse(M::AbstractMatrix{T}; symm = false, herm = false) where {T}
   prod! = @closure (res, v, α, β) -> mulFact!(res, M, v, α, β)
   tprod! = @closure (res, u, α, β) -> mulFact!(res, transpose(M), u, α, β)
   ctprod! = @closure (res, w, α, β) -> mulFact!(res, adjoint(M), w, α, β)
-  LinearOperator{T}(size(M, 2), size(M, 1), symm, herm, prod!, tprod!, ctprod!)
+  LinearOperator{T, Vector{T}}(size(M, 2), size(M, 1), symm, herm, prod!, tprod!, ctprod!)
 end
 
 """
     opCholesky(M, [check=false])
 
 Inverse of a Hermitian and positive definite matrix as a linear operator
-using its Cholesky factorization. 
+using its Cholesky factorization.
 The factorization is computed only once.
 The optional `check` argument will perform cheap hermicity and definiteness
 checks.
@@ -53,7 +53,7 @@ function opCholesky(M::AbstractMatrix; check::Bool = false)
   tprod! = @closure (res, u, α, β) -> tmulFact!(res, LL, u, α, β)  # M.' = conj(M)
   ctprod! = @closure (res, w, α, β) -> mulFact!(res, LL, w, α, β)
   S = eltype(LL)
-  LinearOperator{S}(m, m, isreal(M), true, prod!, tprod!, ctprod!)
+  LinearOperator{S, Vector{S}}(m, m, isreal(M), true, prod!, tprod!, ctprod!)
   #TODO: use iterative refinement.
 end
 
@@ -64,7 +64,7 @@ Inverse of a symmetric matrix as a linear operator using its LDLᵀ factorizatio
 if it exists. The factorization is computed only once. The optional `check`
 argument will perform a cheap hermicity check.
 
-If M is sparse and real, then only the upper triangle should be stored in order to use 
+If M is sparse and real, then only the upper triangle should be stored in order to use
 [`LDLFactorizations.jl`](https://github.com/JuliaSmoothOptimizers/LDLFactorizations.jl):
 
     using LDLFactorizations
@@ -91,7 +91,7 @@ The result is `x -> (I - 2 h hᵀ) x`.
 function opHouseholder(h::AbstractVector{T}) where {T}
   n = length(h)
   prod! = @closure (res, v, α, β) -> mulHouseholder!(res, h, v, α, β)  # tprod will be inferred
-  LinearOperator{T}(n, n, isreal(h), true, prod!, nothing, prod!)
+  LinearOperator{T, Vector{T}}(n, n, isreal(h), true, prod!, nothing, prod!)
 end
 
 function mulHermitian!(res, d, L, v, α, β::T) where {T}
@@ -113,7 +113,7 @@ function opHermitian(d::AbstractVector{S}, A::AbstractMatrix{T}) where {S, T}
   L = tril(A, -1)
   U = promote_type(S, T)
   prod! = @closure (res, v, α, β) -> mulHermitian!(res, d, L, v, α, β)
-  LinearOperator{U}(m, m, isreal(A), true, prod!, nothing, nothing)
+  LinearOperator{U, Vector{U}}(m, m, isreal(A), true, prod!, nothing, nothing)
 end
 
 """

src/operations.jl

@@ -150,7 +150,7 @@ function *(op1::AbstractLinearOperator, op2::AbstractLinearOperator)
   tprod! = @closure (res, u, α, β) -> prod_op!(res, transpose(op2), transpose(op1), utmp, u, α, β)
   ctprod! = @closure (res, w, α, β) -> prod_op!(res, adjoint(op2), adjoint(op1), wtmp, w, α, β)
   args5 = (has_args5(op1) && has_args5(op2))
-  CompositeLinearOperator(T, m1, n2, false, false, prod!, tprod!, ctprod!, args5, S = S)
+  CompositeLinearOperator(T, m1, n2, false, false, prod!, tprod!, ctprod!, args5, S)
 end
 
 ## Matrix times operator.
@@ -213,7 +213,7 @@ function +(op1::AbstractLinearOperator, op2::AbstractLinearOperator)
   S = promote_type(storage_type(op1), storage_type(op2))
   isconcretetype(S) ||
     throw(LinearOperatorException("storage types cannot be promoted to a concrete type"))
-  return CompositeLinearOperator(T, m1, n1, symm, herm, prod!, tprod!, ctprod!, args5, S = S)
+  return CompositeLinearOperator(T, m1, n1, symm, herm, prod!, tprod!, ctprod!, args5, S)
 end
 
 # Operator + matrix.

src/special-operators.jl

@@ -52,7 +52,7 @@ Change `S` to use LinearOperators on GPU.
 """
 function opEye(T::Type, n::Int; S = Vector{T})
   prod! = @closure (res, v, α, β) -> mulOpEye!(res, v, α, β, n)
-  LinearOperator{T}(n, n, true, true, prod!, prod!, prod!, S = S)
+  LinearOperator{T, S}(n, n, true, true, prod!, prod!, prod!)
 end
 
 opEye(n::Int) = opEye(Float64, n)
@@ -71,7 +71,7 @@ function opEye(T::Type, nrow::I, ncol::I; S = Vector{T}) where {I <: Integer}
     return opEye(T, nrow; S = S)
   end
   prod! = @closure (res, v, α, β) -> mulOpEye!(res, v, α, β, min(nrow, ncol))
-  return LinearOperator{T}(nrow, ncol, false, false, prod!, prod!, prod!, S = S)
+  return LinearOperator{T, S}(nrow, ncol, false, false, prod!, prod!, prod!)
 end
 
 opEye(nrow::I, ncol::I) where {I <: Integer} = opEye(Float64, nrow, ncol)
@@ -94,7 +94,7 @@ Change `S` to use LinearOperators on GPU.
 """
 function opOnes(T::Type, nrow::I, ncol::I; S = Vector{T}) where {I <: Integer}
   prod! = @closure (res, v, α, β) -> mulOpOnes!(res, v, α, β)
-  LinearOperator{T}(nrow, ncol, nrow == ncol, nrow == ncol, prod!, prod!, prod!, S = S)
+  LinearOperator{T, S}(nrow, ncol, nrow == ncol, nrow == ncol, prod!, prod!, prod!)
 end
 
 opOnes(nrow::I, ncol::I) where {I <: Integer} = opOnes(Float64, nrow, ncol)
@@ -117,7 +117,7 @@ Change `S` to use LinearOperators on GPU.
 """
 function opZeros(T::Type, nrow::I, ncol::I; S = Vector{T}) where {I <: Integer}
   prod! = @closure (res, v, α, β) -> mulOpZeros!(res, v, α, β)
-  LinearOperator{T}(nrow, ncol, nrow == ncol, nrow == ncol, prod!, prod!, prod!, S = S)
+  LinearOperator{T, S}(nrow, ncol, nrow == ncol, nrow == ncol, prod!, prod!, prod!)
 end
 
 opZeros(nrow::I, ncol::I) where {I <: Integer} = opZeros(Float64, nrow, ncol)
@@ -138,7 +138,7 @@ Diagonal operator with the vector `d` on its main diagonal.
 function opDiagonal(d::AbstractVector{T}) where {T}
   prod! = @closure (res, v, α, β) -> mulSquareOpDiagonal!(res, d, v, α, β)
   ctprod! = @closure (res, w, α, β) -> mulSquareOpDiagonal!(res, conj.(d), w, α, β)
-  LinearOperator{T}(length(d), length(d), true, isreal(d), prod!, prod!, ctprod!, S = typeof(d))
+  LinearOperator{T, typeof(d)}(length(d), length(d), true, isreal(d), prod!, prod!, ctprod!)
 end
 
 function mulOpDiagonal!(res, d, v, α, β::T, n_min) where {T}
@@ -161,11 +161,11 @@ function opDiagonal(nrow::I, ncol::I, d::AbstractVector{T}) where {T, I <: Integ
   prod! = @closure (res, v, α, β) -> mulOpDiagonal!(res, d, v, α, β, n_min)
   tprod! = @closure (res, u, α, β) -> mulOpDiagonal!(res, d, u, α, β, n_min)
   ctprod! = @closure (res, w, α, β) -> mulOpDiagonal!(res, conj.(d), w, α, β, n_min)
-  LinearOperator{T}(nrow, ncol, false, false, prod!, tprod!, ctprod!, S = typeof(d))
+  LinearOperator{T, typeof(d)}(nrow, ncol, false, false, prod!, tprod!, ctprod!)
 end
 
 function mulRestrict!(res, I, v, α, β)
-  res .= v[I]
+  res .= view(v, I)
 end
 
 function multRestrict!(res, I, u, α, β)
@@ -190,9 +190,9 @@ function opRestriction(Idx::LinearOperatorIndexType{I}, ncol::I; S = nothing) wh
   prod! = @closure (res, v, α, β) -> mulRestrict!(res, Idx, v, α, β)
   tprod! = @closure (res, u, α, β) -> multRestrict!(res, Idx, u, α, β)
   if isnothing(S)
-    return LinearOperator{I}(nrow, ncol, false, false, prod!, tprod!, tprod!)
+    return LinearOperator{I, Vector{I}}(nrow, ncol, false, false, prod!, tprod!, tprod!)
   else
-    return LinearOperator{I}(nrow, ncol, false, false, prod!, tprod!, tprod!; S = S)
+    return LinearOperator{I, S}(nrow, ncol, false, false, prod!, tprod!, tprod!)
   end
 end
 
@@ -291,5 +291,5 @@ function BlockDiagonalOperator(ops...; S = promote_type(storage_type.(ops)...))
   symm = all((issymmetric(op) for op ∈ ops))
   herm = all((ishermitian(op) for op ∈ ops))
   args5 = all((has_args5(op) for op ∈ ops))
-  CompositeLinearOperator(T, nrow, ncol, symm, herm, prod!, tprod!, ctprod!, args5, S = S)
+  CompositeLinearOperator(T, nrow, ncol, symm, herm, prod!, tprod!, ctprod!, args5, S)
 end