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_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!"
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 density_matrix_data(request, gpu, num_hidden):
with open(
os.path.join(request.fspath.dirname, "data", "test_grad_data.pkl"),
"rb") as f:
test_data = pickle.load(f)
qucumber.set_random_seed(SEED, cpu=True, gpu=gpu, quiet=True)
data_bases = test_data["density_matrix"]["train_bases"]
data_samples = torch.tensor(test_data["density_matrix"]["train_samples"],
dtype=torch.double)
all_bases = test_data["density_matrix"]["bases"]
target = torch.tensor(test_data["density_matrix"]["density_matrix"],
dtype=torch.double)
num_visible = data_samples.shape[-1]
num_aux = num_visible + 1
nn_state = DensityMatrix(num_visible, num_hidden, num_aux, gpu=gpu)
unitary_dict = nn_state.unitary_dict
DGU = DensityGradsUtils(nn_state)
all_bases = DGU.transform_bases(all_bases)
space = nn_state.generate_hilbert_space()
data_samples = data_samples.to(device=nn_state.device)
target = target.to(device=nn_state.device)
DensityMatrixFixture = namedtuple(
"DensityMatrixFixture",
[
"data_samples",
"data_bases",
"grad_utils",
"all_bases",
"target",
"space",
"nn_state",
"unitary_dict",
],
)
return DensityMatrixFixture(
data_samples=data_samples,
data_bases=data_bases,
grad_utils=DGU,
all_bases=all_bases,
target=target,
space=space,
nn_state=nn_state,
unitary_dict=unitary_dict,
)
def test_density_matrix_expansion(prop):
qucumber.set_random_seed(INIT_SEED, cpu=True, gpu=False, quiet=True)
nn_state = DensityMatrix(5, gpu=False)
v = nn_state.generate_hilbert_space(5)
vp = v[torch.randperm(v.shape[0]), :]
prop_name = prop[0]
is_complex = prop[1]
args = prop[2:]
fn = attrgetter(prop_name)(nn_state)
matrix = fn(v, vp, *args, expand=True)
diag = fn(v, vp, *args, expand=False)
msg = f"Diagonal of matrix {prop_name} is wrong!"
equation = "cii...->ci..." if is_complex else "ii->i"
assertAlmostEqual(torch.einsum(equation, matrix), diag, 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)
def density_matrix_data(request, gpu, num_hidden):
with open(
os.path.join(request.fspath.dirname, "data", "test_grad_data.pkl"),
"rb") as f:
test_data = pickle.load(f)
qucumber.set_random_seed(SEED, cpu=True, gpu=gpu, quiet=True)
data_bases = test_data["density_matrix"]["train_bases"]
data_samples = torch.tensor(test_data["density_matrix"]["train_samples"],
dtype=torch.double)
bases_data = test_data["density_matrix"]["bases"]
target_matrix = torch.tensor(test_data["density_matrix"]["density_matrix"],
dtype=torch.double)
num_visible = data_samples.shape[-1]
num_aux = num_hidden + 1 # this is not a rule, will change with data
unitary_dict = unitaries.create_dict()
nn_state = DensityMatrix(num_visible,
num_hidden,
num_aux,
unitary_dict=unitary_dict,
gpu=gpu)
DGU = DensityGradsUtils(nn_state)
bases = DGU.transform_bases(bases_data)
v_space = nn_state.generate_hilbert_space(num_visible)
a_space = nn_state.generate_hilbert_space(num_aux)
data_samples = data_samples.to(device=nn_state.device)
unitary_dict = {
b: v.to(device=nn_state.device)
for b, v in unitary_dict.items()
}
DensityMatrixFixture = namedtuple(
"DensityMatrixFixture",
[
"data_samples",
"data_bases",
"grad_utils",
"bases",
"target",
"v_space",
"a_space",
"nn_state",
"unitary_dict",
],
)
return DensityMatrixFixture(
data_samples=data_samples,
data_bases=data_bases,
grad_utils=DGU,
bases=bases,
target=target_matrix,
v_space=v_space,
a_space=a_space,
nn_state=nn_state,
unitary_dict=unitary_dict,
)