Two mode Wigner function in python

Question

I have attempted to adjust the qutip Wigner function, for it to process two mode states, specifically for the iterative method.

However the size of the array my output gives out is too big and I am unsure of why? That is when I try and calculate the Wigner logarithmic negativity using it, the integrals come out as arrays rather than singular values.

The code and description of what it is meant to do is below:

`import numpy as np
from scipy import (zeros, array, arange, exp, real, conj, pi,
               copy, sqrt, meshgrid, size, polyval, fliplr, conjugate,
               cos, sin)
import scipy.sparse as sp
import scipy.fftpack as ft
import scipy.linalg as la
from scipy.special import genlaguerre
from scipy.special import binom
from scipy.special import sph_harm

from qutip.qobj import Qobj, isket, isoper
from qutip.states import ket2dm
from qutip.parallel import parfor
from qutip.utilities import clebsch
from scipy.special import factorial
from qutip.cy.sparse_utils import _csr_get_diag
from qutip import *

def wigner2(psi, xvec1, yvec1, xvec2, yvec2, method='iterative', g=np.sqrt(2)):
    """Wigner function for a state vector or density matrix at points
     `xvec1 + i * yvec1` `xvec2 + i * yvec2`

Parameters

state : qobj A state vector or density matrix.

xvec1 : array_like x-coordinates at which to calculate the Wigner function.

yvec1 : array_like y-coordinates at which to calculate the Wigner function.

xvec2 : array_like x-coordinates at which to calculate the Wigner function.

yvec2 : array_like y-coordinates at which to calculate the Wigner function.

g : float Scaling factor for a = 0.5 * g * (x + iy), default g = sqrt(2).

method : string {'iterative'} Select method 'iterative', where 'iterative' uses an iterative method to evaluate the Wigner functions for density matrices :math:|m><n|. The 'iterative' method is default, and in general recommended, but the 'laguerre' method is more efficient for very sparse density matrices (e.g., superpositions of Fock states in a large Hilbert space). The 'fft' method is the preferred method for dealing with density matrices that have a large number of excitations (>~50).

Returns

W : array Values representing the Wigner function calculated over the specified range [xvec1,yvec1] and [xvec2,yvec2] """

    if not (psi.type == 'ket' or psi.type == 'oper' or psi.type == 'bra'):
        raise TypeError('Input state is not a valid operator.')
    if psi.type == 'ket' or psi.type == 'bra':
        rho = ket2dm(psi)   #always use  density matrix
    else:
        rho = psi
    if method == 'iterative':
        return _wigner2_iterative(rho, xvec1, yvec1, xvec2, yvec2, g)
    else:
        raise TypeError(
            "method must be 'iterative'")


def _wigner2_iterative(rho, xvec1, yvec1, xvec2, yvec2, g=np.sqrt(2)):

"""Using an iterative method to evaluate the wigner functions for the Fock state :math:|mp><nq| The Wigner function is calculated as :math:W = \sum_{mpnq} \\rho_{mpnq} W_{mpnq} where :math:W_{mpnq} is the Wigner function for the density matrix :math:|mp><nq|. In this implementation, for each row m*p, Wlist contains the Wigner functions Wlist = [0, ..., W_mpmp, ..., W_mpnq]. As soon as one W_mpnq Wigner function is calculated, the corresponding contribution is added to the total Wigner function, weighted by the corresponding element in the density matrix :math:rho_{mpnq}."""

    M1 = np.prod(ptrace(rho, 0).shape[0])
    M2 = np.prod(ptrace(rho, 1).shape[0])
    M = np.prod(rho.shape[0])  
    X1, Y1, X2, Y2 = np.meshgrid(xvec1, yvec1, xvec2, yvec2)
    A1 = 0.5 * g * (X1 + 1.0j * Y1 + 0 * X2 + 0 * Y2)
    A2 = 0.5 * g * (0 * X1 + 0 * Y1 + X2 + 1.0j * Y2)


    Wlist1 = array([zeros(np.shape(A1), dtype=complex) for k in range(M)])

    Wlist2 = array([zeros(np.shape(A2), dtype=complex) for k in range(M)])


    W = real(rho[0, 0]) * real(Wlist1[0] * Wlist2[0])


    for m in range(0,M1):
        if m==0:
            Wlist1[0] = exp(-2.0 * abs(A1) ** 2) / (pi) 
        else:
            Wlist1[m] = ((2.0 * A1 * Wlist1[m - 1]) / sqrt(m))
        for n in range(0, M2):
            if n==0:
                Wlist2[0] = exp(-2.0 * abs(A2) ** 2) / (pi) 
            else:
                Wlist2[n] = ((2.0 * A2 * Wlist2[n - 1]) / sqrt(n))
            if m != 0 and n != 0:
                W += 2 * real(rho[0, m * M2 + n] * Wlist1[m] * Wlist2[n])

    for p in range(0, M1):
        temp1 = copy(Wlist1[m])
        temp2 = copy(Wlist2[n])
        if p==0:
            Wlist1[p] = exp(-2.0 * abs(A1) ** 2) / (pi) 
        else:
            Wlist1[p] = ((2.0 * conj(A1) * temp1 -sqrt(p) * Wlist1[p-1]) / sqrt(p))
            for q in range(0, M2):
                if q==0:
                    Wlist2[q] = exp(-2.0 * abs(A2) ** 2) / (pi) 
            else:
                Wlist2[q] = ((2.0 * conj(A2) * temp2 - sqrt(q) * Wlist1[q - 1]) / sqrt(q))                
                W += 2 * real(rho[p * M2 + q, p * M2 + q] * Wlist1[p] * Wlist2[q])

            if p != 0 and q !=0:
                for k in range(p + 1, M1):
                    temp3 = (2 * A1 * Wlist1[k-1] - sqrt(k) * temp1) / sqrt(k)
                    temp1 = copy(Wlist1[k])
                    Wlist1[k] = temp3
                    for l in range(q +1, M2):
                        temp4 = (2 * A2 * Wlist2[l-1] - sqrt(l) * temp2) / sqrt(l)
                        temp2 = copy(Wlist2[l])
                        Wlist2[l] = temp4
                        W += 2 * real(rho[p * M2 + q, k *M2 +l] * Wlist1[k] * Wlist2[l])

    return 0.5 * W * g **2'

Answer 1

Same problem for me, you can try this code.

def wigner2(rho,x1,p1,x2,p2,d1,d2):
    # calculates the wigner function at point (alpha1,alpha2) for two modes
    b1=tensor(destroy(d1),identity(d2)) #mechanical oscillator 1
    b2=tensor(identity(d1),destroy(d2)) #mechanical oscillator 2
    
    alpha1=(x1+1j*p1)/np.sqrt(2)
    alpha2=(x2+1j*p2)/np.sqrt(2)
    wig2=(rho*tensor(displace(d1,2*alpha1),displace(d2,2*alpha2))*(1j*np.pi*(b1.dag()*b1+b2.dag()*b2)).expm()).tr()/np.pi**2
    return np.real(wig2)