def test_density_matrix_tr1():
nn_state = DensityMatrix(5, 5, 5, gpu=False)
v_space = nn_state.generate_hilbert_space(5)
matrix = nn_state.rhoRBM(v_space, v_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)
v_space = nn_state.generate_hilbert_space(5)
matrix = nn_state.rhoRBM(v_space, v_space)
# Pick 10 random elements to sample, row and column index
elements = torch.randint(0, 2**5, (2, 10))
real_reg_elements = torch.zeros(10)
real_dag_elements = torch.zeros(10)
imag_reg_elements = torch.zeros(10)
imag_dag_elements = torch.zeros(10)
for i in range(10):
real_reg_elements[i] = matrix[0, elements[0][i], elements[1][i]]
real_dag_elements[i] = matrix[0, elements[1][i], elements[0][i]]
imag_reg_elements[i] = matrix[1, elements[0][i], elements[1][i]]
imag_dag_elements[i] = -matrix[1, elements[1][i], elements[0][i]]
assert torch.equal(real_reg_elements, real_dag_elements)
assert torch.equal(imag_reg_elements, imag_dag_elements)