support * for Compoundimensions

Author: ytdHuang

src/negativity.jl

@@ -78,7 +78,7 @@ end
 # for dense matrices
 function _partial_transpose(ρ::QuantumObject{<:AbstractArray,OperatorQuantumObject}, mask::Vector{Bool})
     isa(ρ.dims, CompoundDimensions) &&
-        (ρ.dims.to != ρ.dims.from) &&
+        (ρ.to != ρ.from) &&
         throw(ArgumentError("Invalid partial transpose for dims = $(ρ.dims)"))
 
     mask2 = [1 + Int(i) for i in mask]
@@ -87,7 +87,7 @@ function _partial_transpose(ρ::QuantumObject{<:AbstractArray,OperatorQuantumObj
     #   2 - the subsystem need be transposed
 
     nsys = length(mask2)
-    dimslist = dims_to_list(ρ.dims)
+    dimslist = dims_to_list(ρ.dims.to) # need `dims.to` here if QO is CompoundDimensions
     pt_dims = reshape(Vector(1:(2*nsys)), (nsys, 2))
     pt_idx = [
         [pt_dims[n, mask2[n]] for n in 1:nsys] # origin   value in mask2
@@ -103,11 +103,11 @@ end
 # for sparse matrices
 function _partial_transpose(ρ::QuantumObject{<:AbstractSparseArray,OperatorQuantumObject}, mask::Vector{Bool})
     isa(ρ.dims, CompoundDimensions) &&
-        (ρ.dims.to != ρ.dims.from) &&
+        (ρ.to != ρ.from) &&
         throw(ArgumentError("Invalid partial transpose for dims = $(ρ.dims)"))
 
     M, N = size(ρ)
-    dimsTuple = Tuple(dims_to_list(ρ.dims))
+    dimsTuple = Tuple(dims_to_list(ρ.dims.to)) # need `dims.to` here if QO is CompoundDimensions
     colptr = ρ.data.colptr
     rowval = ρ.data.rowval
     nzval = ρ.data.nzval

src/qobj/arithmetic_and_attributes.jl

@@ -50,19 +50,54 @@ for op in (:(+), :(-), :(*))
     end
 end
 
+function check_mul_dims(from::SVector{NA,AbstractSpace}, to::SVector{NB,AbstractSpace}) where {NA,NB}
+    (from != to) && throw(
+        DimensionMismatch(
+            "The quantum object with dims = $from can not multiply a quantum object with dims = $to on the right-hand side.",
+        ),
+    )
+    return nothing
+end
+
+for ADimType in (:Dimensions, :CompoundDimensions)
+    for BDimType in (:Dimensions, :CompoundDimensions)
+        if ADimType == BDimType == :Dimensions
+            @eval begin
+                function LinearAlgebra.:(*)(
+                    A::AbstractQuantumObject{DT1,OperatorQuantumObject,$ADimType{NA}},
+                    B::AbstractQuantumObject{DT2,OperatorQuantumObject,$BDimType{NB}},
+                ) where {DT1,DT2,NA,NB}
+                    check_dims(A, B)
+                    return QuantumObject(A.data * B.data, Operator, A.dims)
+                end
+            end
+        else
+            @eval begin
+                function LinearAlgebra.:(*)(
+                    A::AbstractQuantumObject{DT1,OperatorQuantumObject,$ADimType{NA}},
+                    B::AbstractQuantumObject{DT2,OperatorQuantumObject,$BDimType{NB}},
+                ) where {DT1,DT2,NA,NB}
+                    check_mul_dims(A.from, B.to)
+                    return QuantumObject(A.data * B.data, Operator, CompoundDimensions(A.to, B.from))
+                end
+            end
+        end
+    end
+end
+
 function LinearAlgebra.:(*)(
     A::AbstractQuantumObject{DT1,OperatorQuantumObject},
     B::QuantumObject{DT2,KetQuantumObject},
 ) where {DT1,DT2}
-    check_dims(A, B)
-    return QuantumObject(A.data * B.data, Ket, A.dims)
+    check_mul_dims(A.from, B.to)
+    return QuantumObject(A.data * B.data, Ket, Dimensions(A.to))
 end
 function LinearAlgebra.:(*)(
     A::QuantumObject{DT1,BraQuantumObject},
     B::AbstractQuantumObject{DT2,OperatorQuantumObject},
 ) where {DT1,DT2}
-    check_dims(A, B)
-    return QuantumObject(A.data * B.data, Bra, A.dims)
+    check_mul_dims(A.from, B.to)
+    return QuantumObject(A.data * B.data, Bra, Dimensions(B.from))
 end
 function LinearAlgebra.:(*)(
     A::QuantumObject{DT1,KetQuantumObject},
@@ -569,7 +604,7 @@ ptrace(QO::QuantumObject{<:AbstractArray,BraQuantumObject}, sel::Union{AbstractV
 
 function ptrace(QO::QuantumObject{<:AbstractArray,OperatorQuantumObject}, sel::Union{AbstractVector{Int},Tuple})
     isa(QO.dims, CompoundDimensions) &&
-        (QO.dims.to != QO.dims.from) &&
+        (QO.to != QO.from) &&
         throw(ArgumentError("Invalid partial trace for dims = $(QO.dims)"))
 
     _non_static_array_warning("sel", sel)
@@ -587,7 +622,7 @@ function ptrace(QO::QuantumObject{<:AbstractArray,OperatorQuantumObject}, sel::U
         (n_d == 1) && return QO
     end
 
-    dimslist = dims_to_list(QO.dims)
+    dimslist = dims_to_list(QO.dims.to) # need `dims.to` here if QO is CompoundDimensions
     _sort_sel = sort(SVector{length(sel),Int}(sel))
     ρtr, dkeep = _ptrace_oper(QO.data, dimslist, _sort_sel)
     return QuantumObject(ρtr, type = Operator, dims = Dimensions(dkeep))
@@ -757,7 +792,7 @@ function permute(
     order::Union{AbstractVector{Int},Tuple},
 ) where {T,ObjType<:Union{KetQuantumObject,BraQuantumObject,OperatorQuantumObject}}
     isa(A.dims, CompoundDimensions) &&
-        (A.dims.to != A.dims.from) &&
+        (A.to != A.from) &&
         throw(ArgumentError("Invalid permutation for dims = $(A.dims)"))
 
     (length(order) != length(A.dims)) &&
@@ -770,7 +805,7 @@ function permute(
     order_svector = SVector{length(order),Int}(order) # convert it to SVector for performance
 
     # obtain the arguments: dims for reshape; perm for PermutedDimsArray
-    dimslist = dims_to_list(A.dims)
+    dimslist = dims_to_list(A.dims.to) # need `dims.to` here if QO is CompoundDimensions
     dims, perm = _dims_and_perm(A.type, dimslist, order_svector, length(order_svector))
 
     return QuantumObject(

src/qobj/functions.jl

@@ -178,27 +178,36 @@ julia> a.dims, O.dims
 ([20], [20, 20])
 ```
 """
-function LinearAlgebra.kron( # for same OpType. Note that `<:DimType` means same DimType but can have different length N : DimType{N}
-    A::AbstractQuantumObject{DT1,OpType,<:DimType},
-    B::AbstractQuantumObject{DT2,OpType,<:DimType},
-) where {DT1,DT2,OpType<:Union{KetQuantumObject,BraQuantumObject,OperatorQuantumObject},DimType<:AbstractDimensions}
+function LinearAlgebra.kron(
+    A::AbstractQuantumObject{DT1,OpType,Dimensions{NA}},
+    B::AbstractQuantumObject{DT2,OpType,Dimensions{NB}},
+) where {DT1,DT2,OpType<:Union{KetQuantumObject,BraQuantumObject,OperatorQuantumObject},NA,NB}
     QType = promote_op_type(A, B)
-    return QType(kron(A.data, B.data), A.type, kron(A.dims, B.dims))
+    return QType(kron(A.data, B.data), A.type, Dimensions{NA + NB}(vcat(A.dims.to, B.dims.to)))
 end
-function LinearAlgebra.kron( # if A and B are both Operator but different Dimensions type
-    A::AbstractQuantumObject{DT1,OperatorQuantumObject,<:Dimensions},
-    B::AbstractQuantumObject{DT2,OperatorQuantumObject,<:CompoundDimensions},
-) where {DT1,DT2}
-    QType = promote_op_type(A, B)
-    return QType(kron(A.data, B.data), Operator, kron(CompoundDimensions(A.type, A.dims), B.dims))
-end
-function LinearAlgebra.kron( # if A and B are both Operator but different Dimensions type
-    A::AbstractQuantumObject{DT1,OperatorQuantumObject,<:CompoundDimensions},
-    B::AbstractQuantumObject{DT2,OperatorQuantumObject,<:Dimensions},
-) where {DT1,DT2}
-    QType = promote_op_type(A, B)
-    return QType(kron(A.data, B.data), Operator, kron(A.dims, CompoundDimensions(B.type, B.dims)))
+
+# if A and B are both Operator but either one of them has CompoundDimensions
+for ADimType in (:Dimensions, :CompoundDimensions)
+    for BDimType in (:Dimensions, :CompoundDimensions)
+        if !(ADimType == BDimType == :Dimensions) # not for this case because it's already implemented
+            @eval begin
+                function LinearAlgebra.kron(
+                    A::AbstractQuantumObject{DT1,OperatorQuantumObject,$ADimType{NA}},
+                    B::AbstractQuantumObject{DT2,OperatorQuantumObject,$BDimType{NB}},
+                ) where {DT1,DT2,NA,NB}
+                    QType = promote_op_type(A, B)
+                    return QType(
+                        kron(A.data, B.data),
+                        Operator,
+                        CompoundDimensions{NA + NB}(vcat(A.to, B.to), vcat(A.from, B.from)),
+                    )
+                end
+            end
+        end
+    end
 end
+
+# if A and B are differet type (must return Operator with CompoundDimensions)
 for AOpType in (:KetQuantumObject, :BraQuantumObject, :OperatorQuantumObject)
     for BOpType in (:KetQuantumObject, :BraQuantumObject, :OperatorQuantumObject)
         if (AOpType != BOpType)
@@ -208,17 +217,17 @@ for AOpType in (:KetQuantumObject, :BraQuantumObject, :OperatorQuantumObject)
                     B::AbstractQuantumObject{DT2,$BOpType},
                 ) where {DT1,DT2}
                     QType = promote_op_type(A, B)
-
-                    # transfer to CompoundDimensions
-                    _Adims = CompoundDimensions(A.type, A.dims)
-                    _Bdims = CompoundDimensions(B.type, B.dims)
-
-                    return QType(kron(A.data, B.data), Operator, kron(_Adims, _Bdims))
+                    return QType(
+                        kron(A.data, B.data),
+                        Operator,
+                        CompoundDimensions(vcat(A.to, B.to), vcat(A.from, B.from)),
+                    )
                 end
             end
         end
     end
 end
+
 LinearAlgebra.kron(A::AbstractQuantumObject) = A
 function LinearAlgebra.kron(A::Vector{<:AbstractQuantumObject})
     @warn "`tensor(A)` or `kron(A)` with `A` is a `Vector` can hurt performance. Try to use `tensor(A...)` or `kron(A...)` instead."

src/qobj/operators.jl

@@ -435,7 +435,7 @@ function eye(
     if dims isa Nothing
         dims = isa(type, OperatorQuantumObject) ? N : isqrt(N)
     end
-    QuantumObject(Diagonal(ones(ComplexF64, N)); type = type, dims = dims)
+    return QuantumObject(Diagonal(ones(ComplexF64, N)); type = type, dims = dims)
 end
 
 @doc raw"""
@@ -497,7 +497,8 @@ end
 
 Generates the projection operator ``\hat{O} = |i \rangle\langle j|`` with Hilbert space dimension `N`.
 """
-projection(N::Int, i::Int, j::Int) = QuantumObject(sparse([i + 1], [j + 1], [1.0 + 0.0im], N, N), type = Operator, dims = N)
+projection(N::Int, i::Int, j::Int) =
+    QuantumObject(sparse([i + 1], [j + 1], [1.0 + 0.0im], N, N), type = Operator, dims = N)
 
 @doc raw"""
     tunneling(N::Int, m::Int=1; sparse::Union{Bool,Val{<:Bool}}=Val(false))

src/qobj/quantum_object_base.jl

@@ -209,23 +209,26 @@ function _check_QuantumObject(type::OperatorBraQuantumObject, dims::Dimensions,
     return nothing
 end
 
-CompoundDimensions(::KetQuantumObject, dims::Dimensions{N}) where {N} =
-    CompoundDimensions{N}(dims.to, Field_list(N))
-CompoundDimensions(::BraQuantumObject, dims::Dimensions{N}) where {N} =
-    CompoundDimensions{N}(Field_list(N), dims.to)
-CompoundDimensions(::OperatorQuantumObject, dims::Dimensions{N}) where {N} = CompoundDimensions{N}(dims.to, dims.to)
-CompoundDimensions(::OperatorQuantumObject, dims::CompoundDimensions) = dims
-
-# support Qobj.to and Qobj.from (but maybe this is not a good idea)
-# Base.getproperty(A::AbstractQuantumObject, key::Symbol) = getproperty(A, Val{key}())
-# Base.getproperty(A::AbstractQuantumObject{DT,KetQuantumObject,<:Dimensions}, ::Val{:to}) where {DT} = A.dims.to
-# Base.getproperty(A::AbstractQuantumObject{DT,KetQuantumObject,<:Dimensions}, ::Val{:from}) where {DT} = Dimensions(Field_list(length(A.dims))
-# Base.getproperty(A::AbstractQuantumObject{DT,BraQuantumObject,<:Dimensions}, ::Val{:to}) where {DT} = Dimensions(Field_list(length(A.dims))
-# Base.getproperty(A::AbstractQuantumObject{DT,BraQuantumObject,<:Dimensions}, ::Val{:from}) where {DT} = A.dims.to
-# Base.getproperty(A::AbstractQuantumObject{DT,OperatorQuantumObject}, ::Val{:to}) where {DT} = A.dims.to
-# Base.getproperty(A::AbstractQuantumObject{DT,OperatorQuantumObject,<:Dimensions}, ::Val{:from}) where {DT} = A.dims.to
-# Base.getproperty(A::AbstractQuantumObject{DT,OperatorQuantumObject,<:CompoundDimensions}, ::Val{:from}) where {DT} = A.dims.from
-# Base.getproperty(A::AbstractQuantumObject{DT,ObjType}, ::Val{T}) where {DT,ObjType<:QuantumObjectType,T} = throw(ArgumentError("Invalid property `$(T)` for $(A.type)"))
+Base.getproperty(A::AbstractQuantumObject, key::Symbol) = getproperty(A, Val{key}())
+Base.getproperty(A::AbstractQuantumObject, ::Val{K}) where {K} = getfield(A, K)
+
+# support `AbstractQuantumObject.to` and `AbstractQuantumObject.from`
+Base.getproperty(A::AbstractQuantumObject{DT,KetQuantumObject,<:Dimensions}, ::Val{:to}) where {DT} = A.dims.to
+Base.getproperty(A::AbstractQuantumObject{DT,KetQuantumObject,Dimensions{N}}, ::Val{:from}) where {DT,N} = Field_list(N)
+Base.getproperty(A::AbstractQuantumObject{DT,BraQuantumObject,Dimensions{N}}, ::Val{:to}) where {DT,N} = Field_list(N)
+Base.getproperty(A::AbstractQuantumObject{DT,BraQuantumObject,<:Dimensions}, ::Val{:from}) where {DT} = A.dims.to
+Base.getproperty(A::AbstractQuantumObject{DT,OperatorQuantumObject}, ::Val{:to}) where {DT} = A.dims.to
+Base.getproperty(A::AbstractQuantumObject{DT,OperatorQuantumObject,<:Dimensions}, ::Val{:from}) where {DT} = A.dims.to
+Base.getproperty(A::AbstractQuantumObject{DT,OperatorQuantumObject,<:CompoundDimensions}, ::Val{:from}) where {DT} =
+    A.dims.from
+Base.getproperty(
+    A::AbstractQuantumObject{DT,ObjType,<:Dimensions},
+    ::KeyType,
+) where {
+    DT,
+    ObjType<:Union{SuperOperatorQuantumObject,OperatorBraQuantumObject,OperatorKetQuantumObject},
+    KeyType<:Union{Val{:to},Val{:from}},
+} = A.dims.to
 
 # functions for getting Float or Complex element type
 _FType(A::AbstractQuantumObject) = _FType(eltype(A))

src/qobj/space.jl

@@ -14,15 +14,14 @@ Field_list(N::Int) = SVector{N,AbstractSpace}(ntuple(i -> Field(), Val(N)))
 struct Space <: AbstractSpace
     size::Int
 
-    function Space(size::Int)
+    Space(size::Int) =
         if size == 1
             return Field()
         elseif size > 1
             return new(size)
         else
             throw(DomainError(size, "The size of Space must be positive integer (≥ 1)."))
         end
-    end
 end
 Base.string(s::Space) = string(s.size) # this is only used when printing AbstractDimensions
 

test/core-test/quantum_objects.jl

@@ -652,8 +652,26 @@
         ρ = kron(ρ1, ρ2)
         ρ1_ptr = ptrace(ρ, 1)
         ρ2_ptr = ptrace(ρ, 2)
-        @test ρ1.data ≈ ρ1_ptr.data atol = 1e-10
-        @test ρ2.data ≈ ρ2_ptr.data atol = 1e-10
+
+        # use CompoundDimensions to do partial trace
+        ρ1_compound = Qobj(zeros(ComplexF64, 2, 2), dims = CompoundDimensions((2, 1), (2, 1)))
+        basis2 = [tensor(eye(2), basis(2, i)) for i in 0:1]
+        for b in basis2
+            ρ1_compound += b' * ρ * b
+        end
+        ρ2_compound = Qobj(zeros(ComplexF64, 2, 2), dims = CompoundDimensions((1, 2), (1, 2)))
+        basis1 = [tensor(basis(2, i), eye(2)) for i in 0:1]
+        for b in basis1
+            ρ2_compound += b' * ρ * b
+        end
+        @test ρ1.data ≈ ρ1_ptr.data ≈ ρ1_compound.data
+        @test ρ2.data ≈ ρ2_ptr.data ≈ ρ2_compound.data
+        @test ρ1.dims != ρ1_compound.dims
+        @test ρ2.dims != ρ2_compound.dims
+        ρ1_compound = ptrace(ρ1_compound, 1)
+        ρ2_compound = ptrace(ρ2_compound, 2)
+        @test ρ1.dims == ρ1_compound.dims
+        @test ρ2.dims == ρ2_compound.dims
 
         ψlist = [rand_ket(2), rand_ket(3), rand_ket(4), rand_ket(5)]
         ρlist = [rand_dm(2), rand_dm(3), rand_dm(4), rand_dm(5)]