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:

Where the theoretical value was calculted in this paper: https://iopscience.iop.org/article/10.1088/0031-8949/90/7/074029/meta

Does anyone know what I am doing wrong?

Thanks in advance!