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_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!"