Update phase-space methods for offsets in FockBasis (#290)

* Update qfunc for offsets in FockBasis

* Handle offset in wigner

* Bump versions

* Fix indexing error

Project.toml

@@ -1,6 +1,6 @@
 name = "QuantumOptics"
 uuid = "6e0679c1-51ea-5a7c-ac74-d61b76210b0c"
-version = "v0.8.3"
+version = "v0.8.4"
 
 [deps]
 Arpack = "7d9fca2a-8960-54d3-9f78-7d1dccf2cb97"
@@ -21,7 +21,7 @@ Arpack = "0.5.1"
 DiffEqCallbacks = "2"
 FFTW = "1"
 OrdinaryDiffEq = "5"
-QuantumOpticsBase = "0.2"
+QuantumOpticsBase = "0.2.5"
 RecursiveArrayTools = "2"
 Reexport = "0.2, 1.0"
 StochasticDiffEq = "6"

src/phasespace.jl

@@ -30,38 +30,52 @@ end
 
 function qfunc(psi::Ket{B}, alpha::Number) where B<:FockBasis
  b = basis(psi)
- _alpha = convert(ComplexF64, alpha)
- _conj_alpha = conj(_alpha)
- N = length(psi.basis)
- s = psi.data[N]/sqrt(N-1)
- @inbounds for i=1:N-2
- s = (psi.data[N-i] + s*_conj_alpha)/sqrt(N-i-1)
+ _conj_alpha = conj(alpha)
+ c = one(alpha)
+ @inbounds for n=1:b.offset
+ c *= _conj_alpha/sqrt(n)
  end
- s = psi.data[1] + s*_conj_alpha
- return abs2(s)*exp(-abs2(_alpha))/pi
+ s = c*psi.data[1]
+ @inbounds for n=b.offset+1:b.N
+ c *= _conj_alpha/sqrt(n)
+ s += c*psi.data[n+1-b.offset]
+ end
+ return abs2(s)*exp(-abs2(_conj_alpha))/pi
 end
 
 function qfunc(psi::Ket{B}, xvec::AbstractVector, yvec::AbstractVector) where B<:FockBasis
  b = basis(psi)
- Nx = length(xvec)
- Ny = length(yvec)
- points = Nx*Ny
+ points = length(xvec)*length(yvec)
  N = length(b)::Int
+ N0 = b.offset
  _conj_alpha = [complex(x, -y)/sqrt(2) for x=xvec, y=yvec]
- q = fill(psi.data[N]/sqrt(N-1), size(_conj_alpha))
- @inbounds for n=1:N-2
- f0_ = 1/sqrt(N-n-1)
- x = psi.data[N-n]
+
+ # Compute overlap <α|ψ> as reversed sum
+ q = psi.data[N]/sqrt(b.N) .* _conj_alpha
+ @inbounds for n=b.N:-1:N0+2
+ x = psi.data[n-N0]
+ f0_ = 1/sqrt(n-1)
  for i=1:points
- q[i] = (x + q[i]*_conj_alpha[i])*f0_
+ q[i] = (x + q[i])*_conj_alpha[i]*f0_
  end
  end
- result = similar(q, Float64)
+
+ # 1/sqrt(n!) up to offset for first term in sum
+ nfac = 1.0
+ @inbounds for n=1:N0
+ nfac /= sqrt(n)
+ end
  x = psi.data[1]
  @inbounds for i=1:points
- result[i] = abs2(x + q[i]*_conj_alpha[i])*exp(-abs2(_conj_alpha[i]))/pi
+ q[i] = (x + q[i])*nfac*_conj_alpha[i]^N0
  end
- result
+
+ # Return e^(-|α|^2)*|q|^2
+ result = similar(q, float(real(eltype(psi))))
+ @inbounds for i=1:points
+ result[i] = abs2(q[i])*exp(-abs2(_conj_alpha[i]))/pi
+ end
+ return result
 end
 
 function qfunc(state::Union{Ket{B}, AbstractOperator{B,B}}, x::Number, y::Number) where B<:FockBasis
@@ -70,7 +84,7 @@ end
 
 function _qfunc_operator(rho::AbstractOperator{B,B}, alpha::Complex, tmp1::Ket, tmp2::Ket) where B<:FockBasis
  coherentstate!(tmp1, basis(rho), alpha)
- QuantumOpticsBase.mul!(tmp2,rho,tmp1,complex(1.),complex(0.))
+ QuantumOpticsBase.mul!(tmp2,rho,tmp1,true,false)
  a = dot(tmp1.data, tmp2.data)
  return a/pi
 end
@@ -89,32 +103,34 @@ done via the relation ``α = \\frac{1}{\\sqrt{2}}(x + i y)``.
 function wigner(rho::Operator{B,B}, x::Number, y::Number) where B<:FockBasis
  b = basis(rho)
  N = b.N::Int
+ N0 = b.offset
  _2α = complex(convert(Float64, x), convert(Float64, y))*sqrt(2)
  abs2_2α = abs2(_2α)
  w = complex(0.)
  coefficient = complex(0.)
  @inbounds for L=N:-1:1
- coefficient = 2*_clenshaw(L, abs2_2α, rho.data)
+ coefficient = 2*_clenshaw(L, abs2_2α, rho.data, N, N0)
  w = coefficient + w*_2α/sqrt(L+1)
  end
- coefficient = _clenshaw(0, abs2_2α, rho.data)
+ coefficient = _clenshaw(0, abs2_2α, rho.data, N, N0)
  w = coefficient + w*_2α
  exp(-abs2_2α/2)/pi*real(w)
 end
 
 function wigner(rho::Operator{B,B}, xvec::AbstractVector, yvec::AbstractVector) where B<:FockBasis
  b = basis(rho)
  N = b.N::Int
+ N0 = b.offset
  _2α = [complex(x, y)*sqrt(2) for x=xvec, y=yvec]
  abs2_2α = abs2.(_2α)
  w = zero(_2α)
  b0 = similar(_2α)
  b1 = similar(_2α)
  b2 = similar(_2α)
  @inbounds for L=N:-1:1
- _clenshaw_grid(L, rho.data, abs2_2α, _2α, w, b0, b1, b2, 2)
+ _clenshaw_grid!(w, L, rho.data, abs2_2α, _2α, b0, b1, b2, 2, N, N0)
  end
- _clenshaw_grid(0, rho.data, abs2_2α, _2α, w, b0, b1, b2, 1)
+ _clenshaw_grid!(w, 0, rho.data, abs2_2α, _2α, b0, b1, b2, 1, N, N0)
  @inbounds for i=eachindex(w)
  abs2_2α[i] = exp(-abs2_2α[i]/2)/pi.*real(w[i])
  end
@@ -125,33 +141,62 @@ wigner(psi::Ket, x, y) = wigner(dm(psi), x, y)
 wigner(state, alpha::Number) = wigner(state, real(alpha)*sqrt(2), imag(alpha)*sqrt(2))
 
 
-function _clenshaw_grid(L::Int, ρ::Matrix{ComplexF64},
- abs2_2α::Matrix{Float64}, _2α::Matrix{ComplexF64}, w::Matrix{ComplexF64},
- b0::Matrix{ComplexF64}, b1::Matrix{ComplexF64}, b2::Matrix{ComplexF64}, scale::Int)
- n = size(ρ, 1)-L-1
+function _clenshaw_grid!(w, L::Integer, ρ, abs2_2α, _2α, b0, b1, b2,
+ scale::Integer, N::Integer, offset::Integer)
+ n = N-L
  points = length(w)
  if n==0
- f = scale*ρ[1, L+1]
+ if iszero(offset)
+ f = scale*ρ[1, L+1]
+ else
+ f = zero(eltype(ρ))
+ end
  @inbounds for i=1:points
  w[i] = f + w[i]*_2α[i]/sqrt(L+1)
  end
  elseif n==1
  f1 = 1/sqrt(L+1)
- @inbounds for i=1:points
- w[i] = scale*(ρ[1, L+1] - ρ[2, L+2]*(L+1-abs2_2α[i])*f1) + w[i]*_2α[i]*f1
+ if iszero(offset)
+ @inbounds for i=1:points
+ w[i] = scale*(ρ[1, L+1] - ρ[2, L+2]*(L+1-abs2_2α[i])*f1) + w[i]*_2α[i]*f1
+ end
+ elseif isone(offset)
+ @inbounds for i=1:points
+ w[i] = -scale*ρ[1, L+1]*(L+1-abs2_2α[i])*f1 + w[i]*_2α[i]*f1
+ end
+ else # offset > 1
+ @inbounds for i=1:points
+ w[i] *= _2α[i]*f1
+ end
  end
  else
  f0 = sqrt(float((n+L-1)*(n-1)))
  f1 = sqrt(float((n+L)*n))
  f0_ = 1/f0
  f1_ = 1/f1
- fill!(b1, ρ[n+1, L+n+1])
- @inbounds for i=1:points
- b0[i] = ρ[n, L+n] - (2*n-1+L-abs2_2α[i])*f1_*b1[i]
+ if n < offset
+ fill!(b1, zero(eltype(ρ)))
+ fill!(b0, zero(eltype(ρ)))
+ else
+ fill!(b1, ρ[n+1-offset, L+n+1-offset])
+ if n <= offset
+ @inbounds for i=1:points
+ b0[i] = -(2*n-1+L-abs2_2α[i])*f1_*b1[i]
+ end
+ else
+ @inbounds for i=1:points
+ b0[i] = ρ[n-offset, L+n-offset] - (2*n-1+L-abs2_2α[i])*f1_*b1[i]
+ end
+ end
  end
+
  @inbounds for k=n-2:-1:1
  b1, b2, b0 = b0, b1, b2
- x = ρ[k+1, L+k+1]
+ if k < offset
+ x = zero(eltype(ρ))
+ else
+ x = ρ[k+1-offset, L+k+1-offset]
+ end
  a1 = -(2*k+1+L)
  a2 = -f0*f1_
  @inbounds for i=1:points
@@ -161,35 +206,69 @@ function _clenshaw_grid(L::Int, ρ::Matrix{ComplexF64},
  f0 = sqrt((k+L)*k)
  f0_ = 1/f0
  end
- @inbounds for i=1:points
- w[i] = scale*(ρ[1, L+1] - (L+1-abs2_2α[i])*f0_*b0[i] - f0*f1_*b1[i]) + w[i]*_2α[i]*f0_
+ if iszero(offset)
+ @inbounds for i=1:points
+ w[i] = scale*(ρ[1, L+1] - (L+1-abs2_2α[i])*f0_*b0[i] - f0*f1_*b1[i]) + w[i]*_2α[i]*f0_
+ end
+ else
+ @inbounds for i=1:points
+ w[i] = scale*(-(L+1-abs2_2α[i])*f0_*b0[i] - f0*f1_*b1[i]) + w[i]*_2α[i]*f0_
+ end
  end
  end
+ return w
 end
 
-function _clenshaw(L::Int, abs2_2α::Float64, ρ::Matrix{ComplexF64})
- n = size(ρ, 1)-L-1
+function _clenshaw(L::Integer, abs2_2α::Real, ρ, N::Integer, offset::Integer)
+ n = N-L
  if n==0
- return ρ[1, L+1]
+ if iszero(offset)
+ return ρ[1, L+1]
+ else
+ return zero(eltype(ρ))
+ end
  elseif n==1
- ϕ1 = -(L+1-abs2_2α)/sqrt(L+1)
- return ρ[1, L+1] + ρ[2, L+2]*ϕ1
+ if iszero(offset)
+ ϕ1 = -(L+1-abs2_2α)/sqrt(L+1)
+ return ρ[1, L+1] + ρ[2, L+2]*ϕ1
+ elseif isone(offset)
+ ϕ1 = -(L+1-abs2_2α)/sqrt(L+1)
+ return ρ[1,L+1]*ϕ1
+ else
+ return zero(eltype(ρ))
+ end
  else
  f0 = sqrt(float((n+L-1)*(n-1)))
  f1 = sqrt(float((n+L)*n))
  f0_ = 1/f0
  f1_ = 1/f1
  b2 = complex(0.)
- b1 = ρ[n+1, L+n+1]
- b0 = ρ[n, L+n] - (2*n-1+L-abs2_2α)*f1_*b1
+ if n < offset
+ b1 = zero(eltype(ρ))
+ else
+ b1 = ρ[n+1-offset, L+n+1-offset]
+ end
+ if n <= offset
+ b0 = - (2*n-1+L-abs2_2α)*f1_*b1
+ else
+ b0 = ρ[n-offset, L+n-offset] - (2*n-1+L-abs2_2α)*f1_*b1
+ end
  @inbounds for k=n-2:-1:1
  b1, b2 = b0, b1
- b0 = ρ[k+1, L+k+1] - (2*k+1+L-abs2_2α)*f0_*b1 - f0*f1_*b2
+ if k < offset
+ b0 = - (2*k+1+L-abs2_2α)*f0_*b1 - f0*f1_*b2
+ else
+ b0 = ρ[k+1-offset, L+k+1-offset] - (2*k+1+L-abs2_2α)*f0_*b1 - f0*f1_*b2
+ end
  f1, f1_ = f0, f0_
  f0 = sqrt((k+L)*k)
  f0_ = 1/f0
  end
- return ρ[1, L+1] - (L+1-abs2_2α)*f0_*b0 - f0*f1_*b1
+ if iszero(offset)
+ return ρ[1, L+1] - (L+1-abs2_2α)*f0_*b0 - f0*f1_*b1
+ else
+ return - (L+1-abs2_2α)*f0_*b0 - f0*f1_*b1
+ end
  end
 end
 

test/test_phasespace.jl

@@ -66,6 +66,49 @@ for (i,x)=enumerate(X), (j,y)=enumerate(Y)
  @test wigner(rho_fock, beta) ≈ w_fock
 end
 
+# Test with offset
+alpha = complex(4.3, 2.7)
+nfock = 3
+
+X = [0.1:0.1:4;] .+ real(alpha)
+Y = [0.13:0.1:3;] .+ imag(alpha)
+psi_coherent = coherentstate(b, alpha)
+rho_coherent = dm(psi_coherent)
+psi_fock = fockstate(b, nfock)
+rho_fock = dm(psi_fock)
+
+
+Qpsi_coherent = qfunc(psi_coherent, X, Y)
+Qrho_coherent = qfunc(rho_coherent, X, Y)
+Qpsi_fock = qfunc(psi_fock, X, Y)
+Qrho_fock = qfunc(rho_fock, X, Y)
+
+Wpsi_coherent = wigner(psi_coherent, X, Y)
+Wrho_coherent = wigner(rho_coherent, X, Y)
+Wpsi_fock = wigner(psi_fock, X, Y)
+Wrho_fock = wigner(rho_fock, X, Y)
+
+b_off = FockBasis(100,3)
+psi_coherent_off = coherentstate(b_off, alpha)
+rho_coherent_off = dm(psi_coherent_off)
+psi_fock_off = fockstate(b_off, nfock)
+rho_fock_off = dm(psi_fock)
+
+@test isapprox(qfunc(psi_coherent, alpha), qfunc(psi_coherent_off, alpha))
+@test isapprox(Qpsi_coherent, qfunc(psi_coherent_off, X, Y), atol=1e-6)
+@test isapprox(Qrho_coherent, qfunc(rho_coherent_off, X, Y), atol=1e-6)
+@test isapprox(qfunc(psi_fock, alpha), qfunc(psi_fock_off, alpha))
+@test isapprox(Qpsi_fock, qfunc(psi_fock_off, X, Y), atol=1e-6)
+@test isapprox(Qrho_fock, qfunc(rho_fock_off, X, Y), atol=1e-6)
+
+@test isapprox(wigner(psi_coherent, alpha), wigner(psi_coherent_off, alpha), atol=1e-6)
+@test isapprox(Wpsi_coherent, wigner(psi_coherent_off, X, Y), atol=1e-5, rtol=1e-5)
+@test isapprox(Wrho_coherent, wigner(rho_coherent_off, X, Y), atol=1e-5, rtol=1e-5)
+@test isapprox(wigner(psi_fock, alpha), wigner(psi_fock_off, alpha), atol=1e-5, rtol=1e-5)
+@test isapprox(Wpsi_fock, wigner(psi_fock_off, X, Y), atol=1e-5, rtol=1e-5)
+@test isapprox(Wrho_fock, wigner(rho_fock_off, X, Y), atol=1e-5, rtol=1e-5)
+
+
 # Test qfunc with rand operators
 b = FockBasis(50)
 psi = randstate(b)