def rotated_gradient(self, basis, sites, sample):
Upsi, vp, Us, rotated_grad = self.init_gradient(basis, sites)
int_sample = sample[sites].round().int().cpu().numpy()
vp = sample.round().clone()
grad_size = (self.num_visible * self.num_hidden + self.num_hidden +
self.num_visible)
Upsi_v = torch.zeros_like(Upsi, device=self.device)
Z = torch.zeros(grad_size, dtype=torch.double, device=self.device)
Z2 = torch.zeros((2, grad_size),
dtype=torch.double,
device=self.device)
U = torch.tensor([1.0, 1.0], dtype=torch.double, device=self.device)
Ut = np.zeros_like(Us[:, 0], dtype=complex)
ints_size = np.arange(sites.size)
for x in range(2**sites.size):
# overwrite rotated elements
vp = sample.round().clone()
vp[sites] = self.subspace_vector(x, size=sites.size)
int_vp = vp[sites].int().cpu().numpy()
all_Us = Us[ints_size, :, int_sample, int_vp]
# Gradient from the rotation
Ut = np.prod(all_Us[:, 0] + (1j * all_Us[:, 1]))
U[0] = Ut.real
U[1] = Ut.imag
cplx.scalar_mult(U, self.psi(vp), out=Upsi_v)
Upsi += Upsi_v
# Gradient on the current configuration
grad_vp0 = self.rbm_am.effective_energy_gradient(vp)
grad_vp1 = self.rbm_ph.effective_energy_gradient(vp)
rotated_grad[0] += cplx.scalar_mult(Upsi_v,
cplx.make_complex(grad_vp0, Z),
out=Z2)
rotated_grad[1] += cplx.scalar_mult(Upsi_v,
cplx.make_complex(grad_vp1, Z),
out=Z2)
grad = [
cplx.scalar_divide(rotated_grad[0], Upsi)[0, :], # Real
-cplx.scalar_divide(rotated_grad[1], Upsi)[1, :], # Imaginary
]
return grad
def test_vector_scalar_divide(self):
scalar = torch.tensor([1, 2], dtype=torch.double)
vector = torch.tensor([[1, 2], [3, 4]], dtype=torch.double)
expect = torch.tensor([[1.4, 2.0], [0.2, 0.0]], dtype=torch.double)
self.assertTensorsAlmostEqual(
cplx.scalar_divide(vector, scalar),
expect,
msg="Vector / Scalar divide failed!",
)
def test_matrix_scalar_divide(self):
scalar = torch.tensor([1, 2], dtype=torch.double)
matrix = torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]], dtype=torch.double)
expect = torch.tensor(
[[[2.2, 2.8], [3.4, 4.0]], [[0.6, 0.4], [0.2, 0.0]]], dtype=torch.double
)
self.assertTensorsAlmostEqual(
cplx.scalar_divide(matrix, scalar),
expect,
msg="Matrix / Scalar divide failed!",
)
def rotated_gradient(self, basis, sites, sample):
Upsi, vp, Us, rotated_grad = self.init_gradient(basis, sites)
int_sample = sample[sites].round().int().cpu().numpy()
Upsi_v = torch.zeros_like(Upsi, device=self.device)
ints_size = np.arange(sites.size)
# if the number of rotated sites is too large, fallback to loop
# since memory may be unable to store the entire expanded set of
# visible states
if sites.size > self.max_size or (
hasattr(self, "debug_gradient_rotation") and self.debug_gradient_rotation
):
grad_size = (
self.num_visible * self.num_hidden + self.num_hidden + self.num_visible
)
vp = sample.round().clone()
Z = torch.zeros(grad_size, dtype=torch.double, device=self.device)
Z2 = torch.zeros((2, grad_size), dtype=torch.double, device=self.device)
U = torch.tensor([1.0, 1.0], dtype=torch.double, device=self.device)
Ut = np.zeros_like(Us[:, 0], dtype=complex)
for x in range(2 ** sites.size):
# overwrite rotated elements
vp = sample.round().clone()
vp[sites] = self.subspace_vector(x, size=sites.size)
int_vp = vp[sites].int().cpu().numpy()
all_Us = Us[ints_size, :, int_sample, int_vp]
# Gradient from the rotation
Ut = np.prod(all_Us[:, 0] + (1j * all_Us[:, 1]))
U[0] = Ut.real
U[1] = Ut.imag
cplx.scalar_mult(U, self.psi(vp), out=Upsi_v)
Upsi += Upsi_v
# Gradient on the current configuration
grad_vp0 = self.rbm_am.effective_energy_gradient(vp)
grad_vp1 = self.rbm_ph.effective_energy_gradient(vp)
rotated_grad[0] += cplx.scalar_mult(
Upsi_v, cplx.make_complex(grad_vp0, Z), out=Z2
)
rotated_grad[1] += cplx.scalar_mult(
Upsi_v, cplx.make_complex(grad_vp1, Z), out=Z2
)
else:
vp = sample.round().clone().unsqueeze(0).repeat(2 ** sites.size, 1)
vp[:, sites] = self.generate_hilbert_space(size=sites.size)
vp = vp.contiguous()
# overwrite rotated elements
int_vp = vp[:, sites].long().cpu().numpy()
all_Us = Us[ints_size, :, int_sample, int_vp]
Ut = np.prod(all_Us[..., 0] + (1j * all_Us[..., 1]), axis=1)
U = (
cplx.make_complex(torch.tensor(Ut.real), torch.tensor(Ut.imag))
.to(vp)
.contiguous()
)
Upsi_v = cplx.scalar_mult(U, self.psi(vp).detach())
Upsi = torch.sum(Upsi_v, dim=1)
grad_vp0 = self.rbm_am.effective_energy_gradient(vp, reduce=False)
grad_vp1 = self.rbm_ph.effective_energy_gradient(vp, reduce=False)
# since grad_vp0/1 are real, can just treat the scalar multiplication
# and addition as a matrix multiplication
torch.matmul(Upsi_v, grad_vp0, out=rotated_grad[0])
torch.matmul(Upsi_v, grad_vp1, out=rotated_grad[1])
grad = [
cplx.scalar_divide(rotated_grad[0], Upsi)[0, :], # Real
-cplx.scalar_divide(rotated_grad[1], Upsi)[1, :], # Imaginary
]
return grad
def rotated_gradient(self, basis, sites, sample):
r"""Computes the gradients rotated into the measurement basis
:param basis: The bases in which the measurement is made
:type basis: numpy.ndarray
:param sites: The sites where the measurements are not made
in the computational basis
:type sites: numpy.ndarray
:param sample: The measurement (either 0 or 1)
:type sample: torch.Tensor
:returns: A list of two tensors, representing the rotated gradients
of the amplitude and phase RBMS
:rtype: list[torch.Tensor, torch.Tensor]
"""
UrhoU, v, vp, Us, Us_dag, rotated_grad = self.init_gradient(
basis, sites)
int_sample = sample[sites].round().int().cpu().numpy()
ints_size = np.arange(sites.size)
U_ = torch.tensor([1.0, 1.0], dtype=torch.double, device=self.device)
UrhoU = torch.zeros(2, dtype=torch.double, device=self.device)
UrhoU_ = torch.zeros_like(UrhoU)
grad_size = (self.num_visible * self.num_hidden +
self.num_visible * self.num_aux + self.num_visible +
self.num_hidden + self.num_aux)
Z2 = torch.zeros((2, grad_size),
dtype=torch.double,
device=self.device)
v = sample.round().clone()
vp = sample.round().clone()
for x in range(2**sites.size):
v = sample.round().clone()
v[sites] = self.subspace_vector(x, sites.size)
int_v = v[sites].int().cpu().numpy()
all_Us = Us[ints_size, :, int_sample, int_v]
for y in range(2**sites.size):
vp = sample.round().clone()
vp[sites] = self.subspace_vector(y, sites.size)
int_vp = vp[sites].int().cpu().numpy()
all_Us_dag = Us[ints_size, :, int_sample, int_vp]
Ut = np.prod(all_Us[:, 0] + (1j * all_Us[:, 1]))
Ut *= np.prod(
np.conj(all_Us_dag[:, 0] + (1j * all_Us_dag[:, 1])))
U_[0] = Ut.real
U_[1] = Ut.imag
cplx.scalar_mult(U_, self.rhoRBM_tilde(v, vp), out=UrhoU_)
UrhoU += UrhoU_
grad0 = self.am_grads(v, vp)
grad1 = self.ph_grads(v, vp)
rotated_grad[0] += cplx.scalar_mult(UrhoU_, grad0, out=Z2)
rotated_grad[1] += cplx.scalar_mult(UrhoU_, grad1, out=Z2)
grad = [
cplx.scalar_divide(rotated_grad[0], UrhoU),
cplx.scalar_divide(rotated_grad[1], UrhoU),
]
return grad
def gradient(self, basis, sample):
r"""Compute the gradient of a set (v_state) of samples, measured
in different bases
:param basis: A set of basis, (i.e.vector of strings)
:type basis: np.array
"""
num_U = 0 # Number of 1-local unitary rotations
rotated_sites = [] # List of site where the rotations are applied
grad = [] # Gradient
# Read where the unitary rotations are applied
for j in range(self.num_visible):
if (basis[j] != 'Z'):
num_U += 1
rotated_sites.append(j)
# If the basis is the reference one ('ZZZ..Z')
if (num_U == 0):
grad.append(self.rbm_am.effective_energy_gradient(sample)) # Real
grad.append(0.0) # Imaginary
else:
# Initialize
vp = torch.zeros(self.num_visible,
dtype=torch.double,
device=self.device)
rotated_grad = [
torch.zeros(2,
self.rbm_am.num_pars,
dtype=torch.double,
device=self.device),
torch.zeros(2,
self.rbm_ph.num_pars,
dtype=torch.double,
device=self.device)
]
Upsi = torch.zeros(2, dtype=torch.double, device=self.device)
# Sum over the full subspace where the rotation are applied
#sub_state = self.generate_visible_space(num_U)
sub_space = self.generate_Hilbert_space(num_U)
for x in range(1 << num_U):
# Create the correct state for the full system (given the data)
cnt = 0
for j in range(self.num_visible):
if (basis[j] != 'Z'):
#vp[j]=sub_state[x][cnt] # This site sums (it is rotated)
vp[j] = sub_space[x][cnt]
cnt += 1
else:
vp[j] = sample[j] # This site is left unchanged
U = torch.tensor(
[1., 0.], dtype=torch.double, device=self.device
) #Product of the matrix elements of the unitaries
for ii in range(num_U):
tmp = self.unitary_dict[basis[
rotated_sites[ii]]][:,
int(sample[rotated_sites[ii]]),
int(vp[rotated_sites[ii]])]
U = cplx.scalar_mult(U, tmp.to(self.device))
# Gradient on the current configuration
grad_vp = [
self.rbm_am.effective_energy_gradient(vp),
self.rbm_ph.effective_energy_gradient(vp)
]
# NN state rotated in this bases
Upsi_v = cplx.scalar_mult(U, self.psi(vp))
Upsi += Upsi_v
rotated_grad[0] += cplx.scalar_mult(
Upsi_v,
cplx.make_complex(grad_vp[0],
torch.zeros_like(grad_vp[0])))
rotated_grad[1] += cplx.scalar_mult(
Upsi_v,
cplx.make_complex(grad_vp[1],
torch.zeros_like(grad_vp[1])))
grad.append(cplx.scalar_divide(rotated_grad[0], Upsi)[0, :])
grad.append(-cplx.scalar_divide(rotated_grad[1], Upsi)[1, :])
return grad