Implement offset in FockBasis (qojulia#14)

* Implement offset in FockBasis

* Update transform to handle offset

src/fock.jl

@@ -1,30 +1,32 @@
 """
-    FockBasis(N)
+    FockBasis(N,offset=0)
 
 Basis for a Fock space where `N` specifies a cutoff, i.e. what the highest
-included fock state is. Note that the dimension of this basis then is N+1.
+included fock state is. Similarly, the `offset` defines the lowest included
+fock state (default is 0). Note that the dimension of this basis is `N+1-offset`.
 """
 struct FockBasis{T} <: Basis
     shape::Vector{T}
     N::T
-    function FockBasis(N::T) where T<:Integer
-        if N < 0
+    offset::T
+    function FockBasis(N::T,offset::T=0) where T<:Integer
+        if N < 0 || offset < 0 || N <= offset
             throw(DimensionMismatch())
         end
-        new{T}([N+1], N)
+        new{T}([N-offset+1], N, offset)
     end
 end
 
 
-==(b1::FockBasis, b2::FockBasis) = b1.N==b2.N
+==(b1::FockBasis, b2::FockBasis) = (b1.N==b2.N && b1.offset==b2.offset)
 
 """
     number(b::FockBasis)
 
 Number operator for the specified Fock space.
 """
 function number(b::FockBasis)
-    diag = complex.(0.:b.N)
+    diag = complex.(b.offset:b.N)
     data = spdiagm(0 => diag)
     SparseOperator(b, data)
 end
@@ -35,7 +37,7 @@ end
 Annihilation operator for the specified Fock space.
 """
 function destroy(b::FockBasis)
-    diag = complex.(sqrt.(1.:b.N))
+    diag = complex.(sqrt.(b.offset+1.:b.N))
     data = spdiagm(1 => diag)
     SparseOperator(b, data)
 end
@@ -46,7 +48,7 @@ end
 Creation operator for the specified Fock space.
 """
 function create(b::FockBasis)
-    diag = complex.(sqrt.(1.:b.N))
+    diag = complex.(sqrt.(b.offset+1.:b.N))
     data = spdiagm(-1 => diag)
     SparseOperator(b, data)
 end
@@ -64,8 +66,8 @@ displace(b::FockBasis, alpha::Number) = exp(dense(alpha*create(b) - conj(alpha)*
 Fock state ``|n⟩`` for the specified Fock space.
 """
 function fockstate(b::FockBasis, n::Int)
-    @assert n <= b.N
-    basisstate(b, n+1)
+    @assert b.offset <= n <= b.N
+    basisstate(b, n+1-b.offset)
 end
 
 """
@@ -88,7 +90,17 @@ function coherentstate!(ket::Ket, b::FockBasis, alpha::Number)
     alpha = complex(alpha)
     data = ket.data
     data[1] = exp(-abs2(alpha)/2)
-    @inbounds for n=1:b.N
-        data[n+1] = data[n]*alpha/sqrt(n)
+
+    # Compute coefficient up to offset
+    offset = b.offset
+    @inbounds for n=1:offset
+        data[1] *= alpha/sqrt(n)
+    end
+
+    # Write coefficients to state
+    @inbounds for n=1:b.N-offset
+        data[n+1] = data[n]*alpha/sqrt(n+offset)
     end
+
+    return ket
 end

src/printing.jl

@@ -50,7 +50,11 @@ function show(stream::IO, x::SpinBasis)
 end
 
 function show(stream::IO, x::FockBasis)
-    write(stream, "Fock(cutoff=$(x.N))")
+    if iszero(x.offset)
+        write(stream, "Fock(cutoff=$(x.N))")
+    else
+        write(stream, "Fock(cutoff=$(x.N), offset=$(x.offset))")
+    end
 end
 
 function show(stream::IO, x::NLevelBasis)

src/transformations.jl

@@ -24,8 +24,9 @@ function transform(b1::PositionBasis, b2::FockBasis; x0::Real=1)
     for i in 1:length(b1)
         u = xvec[i]/x0
         C = c*exp(-u^2/2)
-        for n=0:b2.N
-            T[i,n+1] = C*horner(A[n+1], u)
+        for n=b2.offset:b2.N
+            idx = n-b2.offset+1
+            T[i,idx] = C*horner(A[n+1], u)
         end
     end
     DenseOperator(b1, b2, T)
@@ -40,8 +41,9 @@ function transform(b1::FockBasis, b2::PositionBasis; x0::Real=1)
     for i in 1:length(b2)
         u = xvec[i]/x0
         C = c*exp(-u^2/2)
-        for n=0:b1.N
-            T[n+1,i] = C*horner(A[n+1], u)
+        for n=b1.offset:b1.N
+            idx = n-b1.offset+1
+            T[idx,i] = C*horner(A[n+1], u)
         end
     end
     DenseOperator(b1, b2, T)

test/test_fock.jl

@@ -69,4 +69,17 @@ rho = dm(psi)
 @test 1e-13 > abs(variance(n, psi) - abs(alpha)^2)
 @test 1e-13 > abs(variance(n, rho) - abs(alpha)^2)
 
+# Test FockBasis with offset
+b_off = FockBasis(100,4)
+@test_throws AssertionError fockstate(b_off, 0)
+n = 55
+psi = fockstate(b_off, n)
+@test expect(number(b_off), psi)==n==expect(create(b_off)*destroy(b_off), psi)
+
+alpha = 5
+psi = coherentstate(b, alpha)
+psi_off = coherentstate(b_off, alpha)
+@test psi.data[b_off.offset+1:end] == psi_off.data
+@test isapprox(norm(psi_off), 1, atol=1e-7)
+
 end # testset

test/test_transformations.jl

@@ -38,6 +38,22 @@ psi_n = coherentstate(b_fock, α0)
 psi_x = gaussianstate(b_position, x0/σ0, p0/σ0, σ0)
 @test 1e-10 > D(psi_x, Txn*psi_n)
 
+# Test with offset in FockBasis
+b_fock = FockBasis(50,1)
+b_position = PositionBasis(-10, 10, 300)
 
+x0 = 6.1
+p0 = 0.3
+α0 = (x0 + 1im*p0)/sqrt(2)
+
+psi_n = coherentstate(b_fock, α0)
+psi_x = gaussianstate(b_position, x0, p0, 1)
+@test isapprox(norm(psi_n), 1, atol=1e-8)
+
+Txn = transform(b_position, b_fock)
+Tnx = transform(b_fock, b_position)
+@test 1e-10 > D(dagger(Txn), Tnx)
+@test 1e-4 > D(psi_x, Txn*psi_n)
+@test 1e-4 > D(psi_n, Tnx*psi_x)
 
 end # testset