Rebase master, rename & add some multiplication methods

src/QuantumOpticsBase.jl

@@ -28,7 +28,7 @@ export Basis, GenericBasis, CompositeBasis, basis,
         #operators_lazytensor
                 LazyTensor, lazytensor_use_cache, lazytensor_clear_cache,
                 lazytensor_cachesize, lazytensor_disable_cache, lazytensor_enable_cache,
-        #operators_lazyket
+        #states_lazyket
                 LazyKet,
         #time_dependent_operators
                 AbstractTimeDependentOperator, TimeDependentSum, set_time!,
@@ -77,11 +77,8 @@ include("operators_sparse.jl")
 include("operators_lazysum.jl")
 include("operators_lazyproduct.jl")
 include("operators_lazytensor.jl")
-<<<<<<< HEAD
 include("time_dependent_operator.jl")
-=======
-include("operators_lazyket.jl")
->>>>>>> 6c13438 (lazy ket implementation)
+include("states_lazyket.jl")
 include("superoperators.jl")
 include("spin.jl")
 include("fock.jl")

src/operators.jl

@@ -1,7 +1,7 @@
 import Base: ==, +, -, *, /, ^, length, one, exp, conj, conj!, transpose
 import LinearAlgebra: tr, ishermitian
 import SparseArrays: sparse
-import QuantumInterface: AbstractOperator
+import QuantumInterface: AbstractOperator, AbstractKet
 
 """
 Abstract type for operators with a data field.
@@ -111,6 +111,9 @@ Expectation value of the given operator `op` for the specified `state`.
 """
 expect(op::AbstractOperator{B,B}, state::Ket{B}) where B = dot(state.data, (op * state).data)
 
+# TODO upstream this one
+# expect(op::AbstractOperator{B,B}, state::AbstractKet{B}) where B = norm(op * state) ^ 2
+
 function expect(indices, op::AbstractOperator{B,B}, state::Ket{B2}) where {B,B2<:CompositeBasis}
     N = length(state.basis.shape)
     indices_ = complement(N, indices)

src/states_lazyket.jl

@@ -0,0 +1,143 @@
+"""
+    LazyKet(b, kets)
+
+Lazy implementation of a tensor product of kets.
+
+The subkets are stored in the `kets` field.
+The main purpose of such a ket are simple computations for large product states, such as expectation values.
+It's used to compute numeric initial states in QuantumCumulants.jl (see QuantumCumulants.initial_values).
+"""
+mutable struct LazyKet{B,T} <: AbstractKet{B,T}
+    basis::B
+    kets::T
+    function LazyKet(b::B, kets::T) where {B<:CompositeBasis,T<:Tuple}
+        N = length(b.bases)
+        for n=1:N
+            @assert isa(kets[n], Ket)
+            @assert kets[n].basis == b.bases[n] #
+        end
+        new{B,T}(b, kets)
+    end
+end
+function LazyKet(b::CompositeBasis, kets::Vector)
+    return LazyKet(b,Tuple(kets))
+end
+
+Base.eltype(ket::LazyKet) = Base.promote_type(eltype.(ket.kets)...)
+
+Base.isequal(x::LazyKet, y::LazyKet) = isequal(x.basis, y.basis) && isequal(x.kets, y.kets)
+Base.:(==)(x::LazyKet, y::LazyKet) = (x.basis == y.basis) && (x.kets == y.kets)
+
+# conversion to dense
+Ket(ket::LazyKet) = ⊗(ket.kets...)
+
+# no lazy bras for now
+dagger(x::LazyKet) = throw(MethodError("dagger not implemented for LazyKet: LazyBra is currently not implemented at all!"))
+
+# tensor with other kets
+function tensor(x::LazyKet, y::Ket)
+    return LazyKet(x.basis ⊗ y.basis, (x.kets..., y))
+end
+function tensor(x::Ket, y::LazyKet)
+    return LazyKet(x.basis ⊗ y.basis, (x, y.kets...))
+end
+function tensor(x::LazyKet, y::LazyKet)
+    return LazyKet(x.basis ⊗ y.basis, (x.kets..., y.kets...))
+end
+
+# norms
+norm(state::LazyKet) = prod(norm.(state.kets))
+function normalize!(state::LazyKet)
+    for ket in state.kets
+        normalize!(ket)
+    end
+    return state
+end
+function normalize(state::LazyKet)
+    y = deepcopy(state)
+    normalize!(y)
+    return y
+end
+
+# expect
+function expect(op::LazyTensor{B, B}, state::LazyKet{B}) where B <: Basis
+    check_samebases(op); check_samebases(op.basis_l, state.basis)
+    ops = op.operators
+    inds = op.indices
+    kets = state.kets
+
+    T = promote_type(eltype(op), eltype(state))
+    exp_val = convert(T, op.factor)
+    for i in 1:length(kets)
+        if i ∈ inds
+            exp_val *= expect(ops[i], kets[i])
+        else
+            exp_val *= dot(kets[i], kets[i])
+        end
+    end
+
+    return exp_val
+end
+
+function expect(op::LazyProduct{B,B}, state::LazyKet{B}) where B <: Basis
+    check_samebases(op); check_samebases(op.basis_l, state.basis)
+
+    tmp_state1 = deepcopy(state)
+    tmp_state2 = deepcopy(state)
+    for i = length(op.operators):-1:1
+        mul!(tmp_state2, op.operators[i], tmp_state1)
+        for j = 1:length(state.kets)
+            copyto!(tmp_state1.kets[j].data, tmp_state2.kets[j].data)
+        end
+    end
+
+    T = promote_type(eltype(op), eltype(state))
+    exp_val = convert(T, op.factor)
+    for i = 1:length(state.kets)
+        exp_val *= dot(state.kets[i].data, tmp_state2.kets[i].data)
+    end
+
+    return exp_val
+end
+
+function expect(op::LazySum{B,B}, state::LazyKet{B}) where B <: Basis
+    check_samebases(op); check_samebases(op.basis_l, state.basis)
+
+    T = promote_type(eltype(op), eltype(state))
+    exp_val = zero(T)
+    for (factor, sub_op) in zip(op.factors, op.operators)
+        exp_val += factor * expect(sub_op, state)
+    end
+
+    return exp_val
+end
+
+
+# mul! with lazytensor -- needed to compute lazyproduct averages (since ⟨op1 * op2⟩ doesn't factorize)
+# this mul! is the only one that really makes sense
+function mul!(y::LazyKet{BL}, op::LazyOperator{BL,BR}, x::LazyKet{BR}) where {BL, BR}
+    T = promote_type(eltype(y), eltype(op), eltype(x))
+    mul!(y, op, x, one(T), zero(T))
+end
+function mul!(y::LazyKet{BL}, op::LazyTensor{BL, BR}, x::LazyKet{BR}, alpha, beta) where {BL, BR}
+    iszero(beta) || throw("Error: cannot perform muladd operation on LazyKets since addition is not implemented. Convert them to dense kets using Ket(x) in order to perform muladd operations.")
+
+    iszero(alpha) && (_zero_op_mul!(y.kets[1].data, beta); return result)
+
+    missing_index_allowed = samebases(op)
+    (length(y.basis.bases) == length(x.basis.bases)) || throw(IncompatibleBases())
+
+    for i in 1:length(y.kets)
+        if i ∈ op.indices
+            mul!(y.kets[i], op.operators[i], x.kets[i])
+        else
+            missing_index_allowed || throw("Can't multiply a LazyOperator with a Ket when there's missing indices and the bases are different.
+                A missing index is equivalent to applying an identity operator, but that's not possible when mapping from one basis to another!")
+
+            copyto!(y.kets[i].data, x.kets[i].data)
+        end
+    end
+
+    rmul!(y.kets[1].data, op.factor * alpha)
+    return y
+end

test/test_operators_lazytensor.jl

@@ -440,19 +440,6 @@ Lop1 = LazyTensor(b1^2, b2^2, 2, sparse(randoperator(b1, b2)))
 @test_throws DimensionMismatch Lop1*dense(Lop1)
 @test_throws DimensionMismatch Lop1*sparse(Lop1)
 
-# LazyKet
-b1 = SpinBasis(1//2)
-b2 = SpinBasis(1)
-b = b1⊗b2
-ψ1 = spindown(b1)
-ψ2 = spinup(b2)
-ψ = LazyKet(b, (ψ1,ψ2))
-sz1 = LazyTensor(b,1,(sigmaz(b1),))
-sz2 = LazyTensor(b,2,(sigmaz(b2),))
-sz = sz1*sz2
-@test expect(sz, ψ) == expect(sigmaz(b1)⊗sigmaz(b2), ψ1⊗ψ2)
-
-
 end # testset
 
 @testset "LazyTensor: explicit isometries" begin

test/test_states.jl

@@ -170,3 +170,54 @@ bra_ .= 3*bra123
 @test_throws ErrorException cos.(bra_)
 
 end # testset
+
+
+@testset "LazyKet" begin
+
+Random.seed!(1)
+
+# LazyKet
+b1 = SpinBasis(1//2)
+b2 = SpinBasis(1)
+b = b1⊗b2
+ψ1 = spindown(b1)
+ψ2 = spinup(b2)
+ψ = LazyKet(b, (ψ1,ψ2))
+sz = LazyTensor(b,(1, 2),(sigmaz(b1), sigmaz(b2)))
+@test expect(sz, ψ) == expect(sigmaz(b1)⊗sigmaz(b2), ψ1⊗ψ2)
+
+@test ψ ⊗ ψ == LazyKet(b ⊗ b, (ψ1, ψ2, ψ1, ψ2))
+
+ψ2 = deepcopy(ψ)
+mul!(ψ2, sz, ψ)
+@test Ket(ψ2) == dense(sz) * Ket(ψ)
+
+# randomize data
+b1 = GenericBasis(4)
+b2 = FockBasis(5)
+b3 = SpinBasis(1//2)
+ψ1 = randstate(b1)
+ψ2 = randstate(b2)
+ψ3 = randstate(b3)
+
+b = b1⊗b2⊗b3
+ψ = LazyKet(b1⊗b2⊗b3, [ψ1, ψ2, ψ3])
+
+op1 = randoperator(b1)
+op2 = randoperator(b2)
+op3 = randoperator(b3)
+
+op = rand(ComplexF64) * LazyTensor(b, b, (1, 2, 3), (op1, op2, op3))
+
+@test expect(op, ψ) ≈ expect(dense(op), Ket(ψ))
+
+op_sum = rand(ComplexF64) * LazySum(op, op)
+@test expect(op_sum, ψ) ≈ expect(op, ψ) * sum(op_sum.factors)
+
+op_prod = rand(ComplexF64) * LazyProduct(op, op)
+@test expect(op_prod, ψ) ≈ expect(dense(op)^2, Ket(ψ)) * op_prod.factor
+
+op_nested = rand(ComplexF64) * LazySum(op_prod, op)
+@test expect(op_nested, ψ) ≈ expect(dense(op_prod), Ket(ψ)) * op_nested.factors[1] + expect(dense(op), Ket(ψ)) * op_nested.factors[2]
+
+end