Guard linalg for objects not declaring matrix/vector interface

Author: goerz

src/linalg.jl

@@ -80,48 +80,45 @@ function Base.copy(Ψ::GradVector{num_controls,T}) where {num_controls,T}
 end
 
 
-function Base.length(Ψ::GradVector)
+# === Vector interface for GradVector ===
+#
+# The following methods are part of the vector interface and are only
+# meaningful when `supports_vector_interface` is true for the state type T.
+# Each method delegates to a private `_name(::Val{supports}, ...)` function:
+# the Val{true} method contains the implementation, and the Val{false} method
+# throws an error.
+
+function _length(::Val{true}, Ψ::GradVector)
     return length(Ψ.state) * (1 + length(Ψ.grad_states))
 end
 
-
-function Base.size(Ψ::GradVector{num_controls,T}) where {num_controls,T}
-    return ((num_controls + 1) * length(Ψ.state),)
+function _length(::Val{false}, Ψ::GradVector)
+    error("$(typeof(Ψ)) does not support the vector interface")
 end
 
-
-function Base.size(O::GradgenOperator{num_controls,GT,CGT}) where {num_controls,GT,CGT}
-    return (num_controls + 1) .* size(O.G)
+function Base.length(Ψ::T) where {T<:GradVector}
+    return _length(Val(supports_vector_interface(T)), Ψ)
 end
 
 
-function Base.size(
-    O::GradgenOperator{num_controls,GT,CGT},
-    dim::Integer
-) where {num_controls,GT,CGT}
-    return (num_controls + 1) * size(O.G, dim)
+function _size(::Val{true}, Ψ::GradVector{num_controls,T}) where {num_controls,T}
+    return ((num_controls + 1) * length(Ψ.state),)
 end
 
-
-function Base.similar(Ψ::GradVector{num_controls,T}) where {num_controls,T}
-    state_sim = similar(Ψ.state)
-    grad_states_sim = [similar(ϕ) for ϕ ∈ Ψ.grad_states]
-    return GradVector{num_controls,typeof(state_sim)}(state_sim, grad_states_sim)
+function _size(::Val{false}, Ψ::GradVector)
+    error("$(typeof(Ψ)) does not support the vector interface")
 end
 
-Base.similar(Ψ::GradVector, ::Type{S}) where {S} = Vector{S}(undef, length(Ψ))
-
-Base.similar(Ψ::GradVector, dims::Tuple{Vararg{Int}}) = Array{eltype(Ψ)}(undef, dims)
+function Base.size(Ψ::T) where {T<:GradVector}
+    return _size(Val(supports_vector_interface(T)), Ψ)
+end
 
-# These definitions of `similar` exist to make ExponentialUtilities happy, but
-# it's not clear at all that `similar` with a custom shape really makes sense
-Base.similar(::GradVector, ::Type{T}, dims::Tuple{Int,Int}) where {T} =
-    Matrix{T}(undef, dims...)
 
-Base.similar(::GradVector, ::Type{T}, dims::Tuple{Int}) where {T} =
-    Vector{T}(undef, dims[1])
-
-function Base.getindex(Ψ::GradVector{num_controls,T}, k::Int) where {num_controls,T}
+function _getindex(
+    ::Val{true},
+    Ψ::GradVector{num_controls,T},
+    k::Int
+) where {num_controls,T}
     N = length(Ψ.state)
     L = num_controls
     block = (k - 1) ÷ N + 1
@@ -133,7 +130,21 @@ function Base.getindex(Ψ::GradVector{num_controls,T}, k::Int) where {num_contro
     end
 end
 
-function Base.setindex!(Ψ::GradVector{num_controls,T}, v, k::Int) where {num_controls,T}
+function _getindex(::Val{false}, Ψ::GradVector, k::Int)
+    error("$(typeof(Ψ)) does not support the vector interface")
+end
+
+function Base.getindex(Ψ::T, k::Int) where {T<:GradVector}
+    return _getindex(Val(supports_vector_interface(T)), Ψ, k)
+end
+
+
+function _setindex!(
+    ::Val{true},
+    Ψ::GradVector{num_controls,T},
+    v,
+    k::Int
+) where {num_controls,T}
     N = length(Ψ.state)
     L = num_controls
     block = (k - 1) ÷ N + 1
@@ -146,28 +157,122 @@ function Base.setindex!(Ψ::GradVector{num_controls,T}, v, k::Int) where {num_co
     return Ψ
 end
 
-function Base.iterate(Ψ::GradVector, k = 1)
+function _setindex!(::Val{false}, Ψ::GradVector, v, k::Int)
+    error("$(typeof(Ψ)) does not support the vector interface")
+end
+
+function Base.setindex!(Ψ::T, v, k::Int) where {T<:GradVector}
+    return _setindex!(Val(supports_vector_interface(T)), Ψ, v, k)
+end
+
+
+function _iterate(::Val{true}, Ψ::GradVector, k)
     k > length(Ψ) && return nothing
     return (Ψ[k], k + 1)
 end
 
+function _iterate(::Val{false}, Ψ::GradVector, k)
+    error("$(typeof(Ψ)) does not support the vector interface")
+end
+
+function Base.iterate(Ψ::T, k = 1) where {T<:GradVector}
+    return _iterate(Val(supports_vector_interface(T)), Ψ, k)
+end
+
+
+function Base.similar(Ψ::GradVector{num_controls,T}) where {num_controls,T}
+    state_sim = similar(Ψ.state)
+    grad_states_sim = [similar(ϕ) for ϕ ∈ Ψ.grad_states]
+    return GradVector{num_controls,typeof(state_sim)}(state_sim, grad_states_sim)
+end
+
+# similar(Ψ, S) calls length(Ψ), which will error if !supports_vector_interface
+Base.similar(Ψ::GradVector, ::Type{S}) where {S} = Vector{S}(undef, length(Ψ))
+
+# similar(Ψ, dims) calls eltype(Ψ) but not length/size, so no vector interface needed
+Base.similar(Ψ::GradVector, dims::Tuple{Vararg{Int}}) = Array{eltype(Ψ)}(undef, dims)
+
+# These definitions of `similar` exist to make ExponentialUtilities happy, but
+# it's not clear at all that `similar` with a custom shape really makes sense
+Base.similar(::GradVector, ::Type{T}, dims::Tuple{Int,Int}) where {T} =
+    Matrix{T}(undef, dims...)
+
+Base.similar(::GradVector, ::Type{T}, dims::Tuple{Int}) where {T} =
+    Vector{T}(undef, dims[1])
+
+
+function Base.fill!(Ψ::GradVector, v)
+    Base.fill!(Ψ.state, v)
+    for i = 1:length(Ψ.grad_states)
+        Base.fill!(Ψ.grad_states[i], v)
+    end
+    return Ψ
+end
+
+
+# === Matrix interface for GradgenOperator ===
+#
+# The following methods are part of the matrix interface and are only
+# meaningful when `supports_matrix_interface` is true for both component types.
+# Each method delegates to a private `_name(::Val{supports}, ...)` function:
+# the Val{true} method contains the implementation, and the Val{false} method
+# throws an error.
+
+function _size(
+    ::Val{true},
+    O::GradgenOperator{num_controls,GT,CGT}
+) where {num_controls,GT,CGT}
+    return (num_controls + 1) .* size(O.G)
+end
+
+function _size(::Val{false}, O::GradgenOperator)
+    error("$(typeof(O)) does not support the matrix interface")
+end
+
+function Base.size(O::T) where {T<:GradgenOperator}
+    return _size(Val(supports_matrix_interface(T)), O)
+end
+
+
+function _size(
+    ::Val{true},
+    O::GradgenOperator{num_controls,GT,CGT},
+    dim::Integer
+) where {num_controls,GT,CGT}
+    return (num_controls + 1) * size(O.G, dim)
+end
+
+function _size(::Val{false}, O::GradgenOperator, dim::Integer)
+    error("$(typeof(O)) does not support the matrix interface")
+end
+
+function Base.size(O::T, dim::Integer) where {T<:GradgenOperator}
+    return _size(Val(supports_matrix_interface(T)), O, dim)
+end
+
+
 # As for an `Operator`, we implement `similar` to return a standard `Array`
-# because `GradgenOperator` does not `setindex!`, so it's arguable not a
-# "mutable array"even if its components are mutable.
+# because `GradgenOperator` does not `setindex!`, so it's arguably not a
+# "mutable array" even if its components are mutable.
+# similar(O) and similar(O, S) call size(O), which will error if
+# !supports_matrix_interface. The dims-based variants need no guard.
 Base.similar(G::GradgenOperator) = Array{eltype(G)}(undef, size(G))
 
 Base.similar(O::GradgenOperator, ::Type{S}) where {S} = Array{S}(undef, size(O))
 Base.similar(O::GradgenOperator, dims::Tuple{Vararg{Int}}) = Array{eltype(O)}(undef, dims)
 Base.similar(O::GradgenOperator, ::Type{S}, dims::Tuple{Vararg{Int}}) where {S} =
     Array{S}(undef, dims)
 
+
 function Base.eltype(
     ::Type{GradgenOperator{num_controls,GT,CGT}}
 ) where {num_controls,GT,CGT}
     return promote_type(eltype(GT), eltype(CGT))
 end
 
-function Base.getindex(
+
+function _getindex(
+    ::Val{true},
     O::GradgenOperator{num_controls,GT,CGT},
     row::Int,
     col::Int
@@ -188,9 +293,29 @@ function Base.getindex(
     end
 end
 
-Base.length(O::GradgenOperator) = prod(size(O))
+function _getindex(::Val{false}, O::GradgenOperator, row::Int, col::Int)
+    error("$(typeof(O)) does not support the matrix interface")
+end
+
+function Base.getindex(O::T, row::Int, col::Int) where {T<:GradgenOperator}
+    return _getindex(Val(supports_matrix_interface(T)), O, row, col)
+end
+
+
+function _length(::Val{true}, O::GradgenOperator)
+    return prod(size(O))
+end
+
+function _length(::Val{false}, O::GradgenOperator)
+    error("$(typeof(O)) does not support the matrix interface")
+end
+
+function Base.length(O::T) where {T<:GradgenOperator}
+    return _length(Val(supports_matrix_interface(T)), O)
+end
+
 
-function Base.iterate(O::GradgenOperator, k = 1)
+function _iterate(::Val{true}, O::GradgenOperator, k)
     n = length(O)
     k > n && return nothing
     n_rows = size(O, 1)
@@ -199,6 +324,15 @@ function Base.iterate(O::GradgenOperator, k = 1)
     return (O[i, j], k + 1)
 end
 
+function _iterate(::Val{false}, O::GradgenOperator, k)
+    error("$(typeof(O)) does not support the matrix interface")
+end
+
+function Base.iterate(O::T, k = 1) where {T<:GradgenOperator}
+    return _iterate(Val(supports_matrix_interface(T)), O, k)
+end
+
+
 function Base.eltype(::Type{GradVector{num_controls,T}}) where {num_controls,T}
     return eltype(T)
 end
@@ -209,15 +343,6 @@ function Base.copyto!(dest::GradgenOperator, src::GradgenOperator)
 end
 
 
-function Base.fill!(Ψ::GradVector, v)
-    Base.fill!(Ψ.state, v)
-    for i = 1:length(Ψ.grad_states)
-        Base.fill!(Ψ.grad_states[i], v)
-    end
-    return Ψ
-end
-
-
 function Base.zero(Ψ::GradVector{num_controls,T}) where {num_controls,T}
     return GradVector{num_controls,T}(zero(Ψ.state), [zero(ϕ) for ϕ ∈ Ψ.grad_states])
 end

test/test_interface.jl

@@ -103,7 +103,7 @@ end
     # iterate visits elements in column-major order, consistent with vec(Array(op))
     @test all(collect(op) .≈ vec(dense))
 
-    # 3-arg mul! agrees with 5-arg mul!(Phi, G, Psi, 1, 0)
+    # 3-arg mul! agrees with 5-arg mul!(Phi, G, Psi, true, false)
     Psi = GradVector(rand(ComplexF64, N), L)
     Phi1 = GradVector(zeros(ComplexF64, N), L)
     Phi2 = GradVector(zeros(ComplexF64, N), L)
@@ -211,4 +211,37 @@ end
     # check_state still passes via the basic (non-vector) state interface
     @test check_state(gradvec)
 
+    # Vector interface methods must throw an error when not supported
+    @test_throws "does not support the vector interface" gradvec[1]
+    @test_throws "does not support the vector interface" (gradvec[1] = 0.0)
+    @test_throws "does not support the vector interface" size(gradvec)
+    @test_throws "does not support the vector interface" length(gradvec)
+    @test_throws "does not support the vector interface" iterate(gradvec)
+
+end
+
+
+
+# A wrapper type with no supports_matrix_interface declaration (defaults to false)
+struct NonMatrixOp
+    data::Matrix{ComplexF64}
+end
+
+@testset "GradgenOperator without Matrix Interface" begin
+
+    N = 5
+    L = 2
+    G = NonMatrixOp(rand(ComplexF64, N, N))
+    mu = [NonMatrixOp(rand(ComplexF64, N, N)) for _ = 1:L]
+    op = GradgenOperator{L,NonMatrixOp,NonMatrixOp}(G, mu)
+
+    @test !supports_matrix_interface(typeof(op))
+
+    # Matrix interface methods must throw an error when not supported
+    @test_throws "does not support the matrix interface" op[1, 1]
+    @test_throws "does not support the matrix interface" size(op)
+    @test_throws "does not support the matrix interface" size(op, 1)
+    @test_throws "does not support the matrix interface" length(op)
+    @test_throws "does not support the matrix interface" iterate(op)
+
 end