Computing the logarithmic negativity in the Fock basis

I am trying to calculate the logarithmic negativity of a density matrix consisting of 2 modes. I wrote the following functions, but they do not give me the expected result:

def basis_states(cutoff_dim):
    """
    Gives a list of all possible Fock basis states (for a system consisting of 2 modes)
    """
    res = []
    for i in range(cutoff_dim):
        for j in range(cutoff_dim):
            res.append([i,j])
    return res
def dm_4d_to_2d(rho):
    """
    Converts the StrawberryFields representiation of a density matrix to 
    a 2 dimensional density matrix 
    (for a system consisting of 2 modes)
    rho [cutoff_dim,cutoff_dim,cutoff_dim,cutoff_dim]
    res [cutoff_dim**2,cutoff_dim**2]
    """
    assert isinstance(rho,np.ndarray)
    basis = basis_states(cutoff_dim)
    res = np.zeros([len(basis),len(basis)],dtype=rho.dtype)
    for idx1, psi1 in enumerate(basis):
        for idx2, psi2 in enumerate(basis):
            res[idx1,idx2] = rho[psi1[0],psi2[0],psi1[1],psi2[1]]
    return res

def partial_transpose(rho):
    """
    Partial transpose (for a system consisting of 2 modes)
    https://en.wikipedia.org/wiki/Peres%E2%80%93Horodecki_criterion

    """
    assert isinstance(rho, np.ndarray)
    assert len(rho.shape) == 4 # only works for two modes + not batched
    res = rho.copy()
    for i in range(rho.shape[0]):
        for j in range(rho.shape[1]):
            res[i,j] = rho[i,j].transpose()
    return res

def negativity_via_2d(np_4d):
    """
    Negativity (for a system consisting of 2 modes)
    https://en.wikipedia.org/wiki/Negativity_(quantum_mechanics)
    """
    assert len(np_4d.shape) == 4
    rho_4d = np_4d.copy()
    
    rho_4d = partial_transpose(rho_4d)
    rho_2d = dm_4d_to_2d(rho_4d)

    rho_2d = rho_2d.conj().T @ rho_2d
    rho_2d = scipy.linalg.sqrtm(rho_2d)

    return (np.trace(rho_2d) - 1.) * 0.5

def log_negativity_via_2d(rho_4d):
    neg = negativity_via_2d(rho_4d)
    return np.log2(1+2*neg)

I now use these functions to plot the log negativity of 2 mode squeezed vacuum states (as a function of the squeezing parameter), using the following simulation:

log_neg_list = []
range_r_params = np.arange(0,3,0.1)
for r_param in range_r_params:
    eng = sf.Engine("tf", backend_options={"cutoff_dim": 6})
    prog = sf.Program(m)
    with prog as q:
        Fock(0) | q[0]
        Fock(0) | q[1]
        S2gate(r_param) | q

    state = eng.run(prog).state
    rho = state.dm().numpy()
    log_neg_list.append(log_negativity_via_2d(rho))

This gives the following result:
image image

Where the theoretical value was calculted in this paper: ShieldSquare Captcha

Does anyone know what I am doing wrong?
Thanks in advance!

Moreover, this alternative computation (using eigenvalues) gives me the same plots:

def log_negativity_via_eigenvalues(np_4d):
    assert len(np_4d.shape) == 4
    rho_4d = np_4d.copy()

    rho_4d = partial_transpose(rho_4d)
    rho_2d = dm_4d_to_2d(rho_4d)

    eigvals = np.linalg.eigvals(rho_2d)
    N = np.abs(sum([i for i in eigvals if i<0.]))
    return np.log2(1+2*N)

I found the problem. Apparently, the trace of my simulation output does not always equal 1, so my cutoff dimension quickly becomes too small when squeezing:

image

I’m glad you found the solution @Robbe! Thanks for posting it here too so that others can benefit from it.

Enjoy using StrawberryFields!