def test_density_matrix_diagonal():
nn_state = DensityMatrix(5, gpu=False)
v = nn_state.generate_hilbert_space(5)
rho = nn_state.rho(v, expand=True)
diag = nn_state.rho(v, expand=False)
msg = "Diagonal of density matrix is wrong!"
assertAlmostEqual(torch.einsum("cii...->ci...", rho), diag, TOL, msg=msg)
def test_density_matrix_tr1():
nn_state = DensityMatrix(5, gpu=False)
space = nn_state.generate_hilbert_space(5)
matrix = nn_state.rho(space, space) / nn_state.normalization(space)
msg = f"Trace of density matrix is not within {TOL} of 1!"
assertAlmostEqual(torch.trace(matrix[0]), torch.Tensor([1]), TOL, msg=msg)
def test_density_matrix_sizes():
nn_state = DensityMatrix(5, gpu=False)
v = nn_state.generate_hilbert_space(5)
vp = v[:4, :]
rho = nn_state.rho(v, vp)
assert rho.shape == (2, v.shape[0], vp.shape[0])
def test_density_matrix_hermiticity():
nn_state = DensityMatrix(5, 5, 5, gpu=False)
space = nn_state.generate_hilbert_space(5)
Z = nn_state.normalization(space)
rho = nn_state.rho(space, space) / Z
assert torch.equal(rho, cplx.conjugate(rho)), "DensityMatrix should be Hermitian!"