Merge #101 from branch operator-matrix-interface

Author: goerz

src/generators.jl

@@ -180,6 +180,33 @@ Base.size(O::Operator) = size(O.ops[1])
 Base.size(O::Operator, dim::Integer) = size(O.ops[1], dim)
 Base.eltype(::Type{Operator{OT,CT}}) where {OT,CT} = promote_type(eltype(OT), CT)
 
+function Base.getindex(O::Operator, i::Int, j::Int)
+    drift_offset = length(O.ops) - length(O.coeffs)
+    result = (drift_offset == 0 ? O.coeffs[1] : one(eltype(O))) * O.ops[1][i, j]
+    for k = 2:length(O.ops)
+        c = k <= drift_offset ? one(eltype(O)) : O.coeffs[k-drift_offset]
+        result += c * O.ops[k][i, j]
+    end
+    return result
+end
+
+Base.length(O::Operator) = prod(size(O))
+
+function Base.iterate(O::Operator, k = 1)
+    n = length(O)
+    k > n && return nothing
+    n_rows = size(O, 1)
+    i = (k - 1) % n_rows + 1
+    j = (k - 1) ÷ n_rows + 1
+    return (O[i, j], k + 1)
+end
+
+Base.similar(O::Operator) = Array{eltype(O)}(undef, size(O))
+Base.similar(O::Operator, ::Type{S}) where {S} = Array{S}(undef, size(O))
+Base.similar(O::Operator, dims::Tuple{Vararg{Int}}) = Array{eltype(O)}(undef, dims)
+Base.similar(O::Operator, ::Type{S}, dims::Tuple{Vararg{Int}}) where {S} =
+    Array{S}(undef, dims)
+
 
 function LinearAlgebra.ishermitian(O::Operator{OT,CT}) where {OT,CT}
     return all(ishermitian(op) for op in O.ops) && all(isreal(c) for c in O.coeffs)
@@ -254,8 +281,28 @@ Base.Array(O::ScaledOperator) = Array{eltype(O)}(O)
 
 Base.size(O::ScaledOperator) = size(O.operator)
 Base.size(O::ScaledOperator, dim::Integer) = size(O.operator, dim)
-Base.eltype(::Type{ScaledOperator{CT,Operator{OOT,OCT}}}) where {CT,OOT,OCT} =
-    promote_type(CT, eltype(OOT), OCT)
+Base.eltype(::Type{ScaledOperator{CT,OT}}) where {CT,OT} = promote_type(CT, eltype(OT))
+
+function Base.getindex(O::ScaledOperator, i::Int, j::Int)
+    return O.coeff * O.operator[i, j]
+end
+
+Base.length(O::ScaledOperator) = prod(size(O))
+
+function Base.iterate(O::ScaledOperator, k = 1)
+    n = length(O)
+    k > n && return nothing
+    n_rows = size(O, 1)
+    i = (k - 1) % n_rows + 1
+    j = (k - 1) ÷ n_rows + 1
+    return (O[i, j], k + 1)
+end
+
+Base.similar(O::ScaledOperator) = Array{eltype(O)}(undef, size(O))
+Base.similar(O::ScaledOperator, ::Type{S}) where {S} = Array{S}(undef, size(O))
+Base.similar(O::ScaledOperator, dims::Tuple{Vararg{Int}}) = Array{eltype(O)}(undef, dims)
+Base.similar(O::ScaledOperator, ::Type{S}, dims::Tuple{Vararg{Int}}) where {S} =
+    Array{S}(undef, dims)
 
 LinearAlgebra.ishermitian(O::ScaledOperator) = (isreal(O.coeff) && ishermitian(O.operator))
 

src/interfaces/supports_matrix_interface.jl

@@ -13,9 +13,6 @@ of `AbstractMatrix`. Calling `supports_matrix_interface` on an instance
 `x` also works via a convenience fallback that forwards to
 `supports_matrix_interface(typeof(x))`.
 
-Depending the value of [`supports_inplace`](@ref), `T` may implement a
-read-write matrix (`setindex!` etc.) or a read-only matrix.
-
 The matrix interface is encouraged for [operators](@ref Operators), and the
 specific conditions of the required interface in that context are checked via
 [`QuantumPropagators.Interfaces.check_operator`](@ref).
@@ -25,3 +22,10 @@ supports_matrix_interface(::Type{<:AbstractMatrix}) = true
 supports_matrix_interface(::Type{T}) where {T} = false
 
 supports_matrix_interface(x) = supports_matrix_interface(typeof(x))
+
+import ..Operator
+import ..ScaledOperator
+
+supports_matrix_interface(::Type{<:Operator{OT}}) where {OT} = supports_matrix_interface(OT)
+supports_matrix_interface(::Type{<:ScaledOperator{<:Any,OT}}) where {OT} =
+    supports_matrix_interface(OT)

test/test_operator_linalg.jl

@@ -1,11 +1,29 @@
 using Test
 using LinearAlgebra
 using QuantumControlTestUtils.RandomObjects: random_matrix, random_state_vector
-using QuantumPropagators.Interfaces: check_operator
+using QuantumPropagators.Interfaces: check_operator, supports_matrix_interface
+import QuantumPropagators.Interfaces: supports_inplace
+import QuantumPropagators.Controls: get_controls, evaluate
 using StaticArrays: SMatrix, SVector
 
 using QuantumPropagators: Generator, Operator, ScaledOperator
 
+# A minimal operator type that does not subtype AbstractMatrix, to test that
+# supports_matrix_interface propagates correctly through Operator/ScaledOperator.
+struct MatrixFreeOp
+    mat::Matrix{ComplexF64}
+end
+
+Base.size(op::MatrixFreeOp) = size(op.mat)
+Base.size(op::MatrixFreeOp, dim::Integer) = size(op.mat, dim)
+Base.:*(α::Number, op::MatrixFreeOp) = MatrixFreeOp(α * op.mat)
+Base.:*(op::MatrixFreeOp, Ψ::AbstractVector) = op.mat * Ψ
+LinearAlgebra.mul!(ϕ, op::MatrixFreeOp, Ψ, α::Number, β::Number) = mul!(ϕ, op.mat, Ψ, α, β)
+LinearAlgebra.dot(ϕ, op::MatrixFreeOp, Ψ) = dot(ϕ, op.mat, Ψ)
+supports_inplace(::Type{MatrixFreeOp}) = true
+get_controls(::MatrixFreeOp) = ()
+evaluate(op::MatrixFreeOp, args...; kwargs...) = op
+
 
 @testset "Operator mul!" begin
 
@@ -71,6 +89,7 @@ end
     H₂ = random_matrix(5; hermitian = true)
     Op = Operator([H₀, H₁, H₂], [2.1, 1.1])
     Ψ = random_state_vector(5)
+    @test supports_matrix_interface(typeof(Op))
     @test check_operator(Op; state = Ψ)
 end
 
@@ -277,5 +296,39 @@ end
     Op = Operator([H₀, H₁, H₂], [2.1, 1.1]) * 0.5
     @test Op isa ScaledOperator
     Ψ = random_state_vector(5)
+    @test supports_matrix_interface(typeof(Op))
     @test check_operator(Op; state = Ψ)
 end
+
+
+@testset "supports_matrix_interface" begin
+
+    H₀ = random_matrix(5; hermitian = true)
+    H₁ = random_matrix(5; hermitian = true)
+    H₂ = random_matrix(5; hermitian = true)
+    Ψ = random_state_vector(5)
+
+    # Operator wrapping standard Matrix → true
+    Op = Operator([H₀, H₁, H₂], [2.1, 1.1])
+    @test supports_matrix_interface(typeof(Op))
+
+    # ScaledOperator wrapping Operator{Matrix} → true
+    SOp = Op * 0.5
+    @test SOp isa ScaledOperator
+    @test supports_matrix_interface(typeof(SOp))
+
+    # Operator wrapping MatrixFreeOp (not AbstractMatrix) → false, but valid operator
+    h₀ = MatrixFreeOp(H₀)
+    h₁ = MatrixFreeOp(H₁)
+    h₂ = MatrixFreeOp(H₂)
+    FreeOp = Operator([h₀, h₁, h₂], [2.1, 1.1])
+    @test !supports_matrix_interface(typeof(FreeOp))
+    @test check_operator(FreeOp; state = Ψ)
+
+    # ScaledOperator wrapping Operator{MatrixFreeOp} → false, but valid operator
+    SFreeOp = FreeOp * 0.5
+    @test SFreeOp isa ScaledOperator
+    @test !supports_matrix_interface(typeof(SFreeOp))
+    @test check_operator(SFreeOp; state = Ψ)
+
+end