def rho(self, v, vp=None, expand=True):
r"""Computes the matrix elements of the (unnormalized) density matrix.
If `expand` is `True`, will return a complex matrix
:math:`A_{ij} = \langle\sigma_i|\widetilde{\rho}|\sigma'_j\rangle`.
Otherwise will return a complex vector
:math:`A_{i} = \langle\sigma_i|\widetilde{\rho}|\sigma'_i\rangle`.
:param v: One of the visible states, :math:`\sigma`.
:type v: torch.Tensor
:param vp: The other visible state, :math:`\sigma'`.
If `None`, will be set to `v`.
:type vp: torch.Tensor
:param expand: Whether to return a matrix (`True`) or a vector (`False`).
:type expand: bool
:returns: The elements of the current density matrix
:math:`\langle\sigma|\widetilde{\rho}|\sigma'\rangle`
:rtype: torch.Tensor
"""
if expand is False and vp is None:
return cplx.make_complex(self.probability(v))
elif vp is None:
vp = v
pi_ = self.pi(v, vp, expand=expand)
amp = (self.rbm_am.gamma(v, vp, eta=+1, expand=expand) +
cplx.real(pi_)).exp()
phase = self.rbm_ph.gamma(v, vp, eta=-1,
expand=expand) + cplx.imag(pi_)
return cplx.make_complex(amp * phase.cos(), amp * phase.sin())
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 importance_sampling_numerator(self, iter_sample, drawn_sample):
out = torch.zeros(iter_sample.shape[0],
dtype=torch.double,
device=self.device)
idx = torch.all(iter_sample == drawn_sample, dim=-1)
out[idx] = self.probability(iter_sample[idx, :])
return cplx.make_complex(out)
def apply(self, nn_state, samples):
r"""Computes the magnetization along Y of each sample in the given batch of samples.
:param nn_state: The WaveFunction that drew the samples.
:type nn_state: qucumber.nn_states.WaveFunction
:param samples: A batch of samples to calculate the observable on.
Must be using the :math:`\sigma_i = 0, 1` convention.
:type samples: torch.Tensor
"""
samples = samples.to(device=nn_state.device)
# vectors of shape: (2, num_samples,)
psis = nn_state.psi(samples)
psi_ratio_sum = torch.zeros_like(psis)
for i in range(samples.shape[-1]): # sum over spin sites
coeff = -to_pm1(samples[:, i])
coeff = cplx.make_complex(torch.zeros_like(coeff), coeff)
flip_spin(i, samples) # flip the spin at site i
# compute ratio of psi_(-i) / psi, multiply it by the appropriate
# eigenvalue, and add it to the running sum
psi_ratio = nn_state.psi(samples)
psi_ratio = cplx.elementwise_division(psi_ratio, psis)
psi_ratio = cplx.elementwise_mult(psi_ratio, coeff)
psi_ratio_sum.add_(psi_ratio)
flip_spin(i, samples) # flip it back
# take real part (imaginary part should be approximately zero)
# and divide by number of spins
return psi_ratio_sum[0].div_(samples.shape[-1])
def load_data_DM(
tr_samples_path,
tr_mtx_real_path=None,
tr_mtx_imag_path=None,
tr_bases_path=None,
bases_path=None,
):
r"""Load the data required for training.
:param tr_samples_path: The path to the training data.
:type tr_samples_path: str
:param tr_mtx_real_path: The path to the real part of the density matrix
:type tr_mtx_real_path: str
:param tr_mtx_imag_path: The path to the imaginary part of the density matrix
:type tr_mtx_imag_path: str
:param tr_bases_path: The path to the basis data.
:type tr_bases_path: str
:param bases_path: The path to a file containing all possible bases used in
the tr_bases_path file.
:type bases_path: str
:returns: A list of all input parameters, with the real and imaginary parts
combined into one (PyTorch-hack) complex matrix.
:rtype: list
"""
data = []
data.append(
torch.tensor(np.loadtxt(tr_samples_path, dtype="float32"),
dtype=torch.double))
if tr_mtx_real_path is not None:
mtx_real = torch.tensor(np.loadtxt(tr_mtx_real_path, dtype="float32"),
dtype=torch.double)
if tr_mtx_imag_path is not None:
mtx_imag = torch.tensor(np.loadtxt(tr_mtx_imag_path, dtype="float32"),
dtype=torch.double)
if tr_mtx_real_path is not None or tr_mtx_imag_path is not None:
if tr_mtx_real_path is None or tr_mtx_imag_path is None:
raise ValueError(
"Must provide a real and imaginary part of target matrix!")
else:
data.append(cplx.make_complex(mtx_real, mtx_imag))
if tr_bases_path is not None:
data.append(np.loadtxt(tr_bases_path, dtype=str))
if bases_path is not None:
bases_data = np.loadtxt(bases_path, dtype=str, ndmin=1)
bases = []
for i in range(len(bases_data)):
tmp = ""
for j in range(len(bases_data[i])):
if bases_data[i][j] != " ":
tmp += bases_data[i][j]
bases.append(tmp)
data.append(bases)
return data
def test_imag_part_of_tensor(self):
x = torch.randn(3, 3, 3)
y = torch.randn(3, 3, 3)
z = cplx.make_complex(x, y)
self.assertTensorsEqual(
y, cplx.imag(z), msg="Imaginary part of rank-3 tensor failed!"
)
def test_real_part_of_tensor(self):
x = torch.randn(3, 3, 3)
y = torch.randn(3, 3, 3)
z = cplx.make_complex(x, y)
self.assertTensorsEqual(
x, cplx.real(z), msg="Real part of rank-3 tensor failed!"
)
def test_real_part_of_vector(self):
x = torch.tensor([1, 2])
y = torch.tensor([5, 6])
z = cplx.make_complex(x, y)
self.assertTensorsEqual(x,
cplx.real(z),
msg="Real part of vector failed!")
def test_make_complex_vector(self):
x = torch.tensor([1, 2, 3, 4])
y = torch.tensor([5, 6, 7, 8])
z = cplx.make_complex(x, y)
expect = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]])
self.assertTensorsEqual(expect, z, msg="Make Complex Vector failed!")
def test_make_complex_matrix(self):
x = torch.tensor([[1, 2], [3, 4]])
y = torch.tensor([[5, 6], [7, 8]])
z = cplx.make_complex(x, y)
expect = torch.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
self.assertTensorsEqual(expect, z, msg="Make Complex Matrix failed!")
def test_imag_part_of_matrix(self):
x = torch.tensor([[1, 2], [3, 4]])
y = torch.tensor([[5, 6], [7, 8]])
z = cplx.make_complex(x, y)
self.assertTensorsEqual(y,
cplx.imag(z),
msg="Imaginary part of matrix failed!")
def test_real_part_of_matrix(self):
x = torch.tensor([[1, 2], [3, 4]])
y = torch.tensor([[5, 6], [7, 8]])
z = cplx.make_complex(x, y)
self.assertTensorsEqual(x,
cplx.real(z),
msg="Real part of matrix failed!")
def test_imag_part_of_vector(self):
x = torch.tensor([1, 2])
y = torch.tensor([5, 6])
z = cplx.make_complex(x, y)
self.assertTensorsEqual(y,
cplx.imag(z),
msg="Imaginary part of vector failed!")
def test_make_complex_vector_with_zero_imaginary_part(self):
x = torch.tensor([1, 2, 3, 4])
z = cplx.make_complex(x)
expect = torch.tensor([[1, 2, 3, 4], [0, 0, 0, 0]])
self.assertTensorsEqual(
expect, z, msg="Making a complex vector with zero imaginary part failed!"
)
def am_grads(self, v):
r"""Computes the gradients of the amplitude RBM for given input states
:param v: The input state, :math:`\sigma`
:type v: torch.Tensor
:returns: The gradients of all amplitude RBM parameters
:rtype: torch.Tensor
"""
return cplx.make_complex(self.rbm_am.effective_energy_gradient(v, reduce=False))
def gamma_grad(self, v, vp, eta=1, expand=False):
r"""Calculates elements of the gradient of
the :math:`\Gamma^{(\eta)}` matrix, where :math:`\eta = \pm`.
:param v: A batch of visible states, :math:`\sigma`
:type v: torch.Tensor
:param vp: The other batch of visible states, :math:`\sigma'`
:type vp: torch.Tensor
:param eta: Determines which gamma matrix elements to compute.
:type eta: int
:param expand: Whether to return a rank-3 tensor (`True`) or a matrix (`False`).
:type expand: bool
:returns: The matrix element given by
:math:`\langle\sigma|\nabla_\lambda\Gamma^{(\eta)}|\sigma'\rangle`
:rtype: torch.Tensor
"""
sign = np.sign(eta)
unsqueezed = v.dim() < 2 or vp.dim() < 2
v = (v.unsqueeze(0) if v.dim() < 2 else v).to(self.weights_W)
vp = (vp.unsqueeze(0) if vp.dim() < 2 else vp).to(self.weights_W)
prob_h = self.prob_h_given_v(v)
prob_hp = self.prob_h_given_v(vp)
W_grad_ = torch.einsum("...j,...k->...jk", prob_h, v)
W_grad_p = torch.einsum("...j,...k->...jk", prob_hp, vp)
if expand:
W_grad = 0.5 * (W_grad_.unsqueeze_(1) +
sign * W_grad_p.unsqueeze_(0))
vb_grad = 0.5 * (v.unsqueeze(1) + sign * vp.unsqueeze(0))
hb_grad = 0.5 * (prob_h.unsqueeze_(1) +
sign * prob_hp.unsqueeze_(0))
else:
W_grad = 0.5 * (W_grad_ + sign * W_grad_p)
vb_grad = 0.5 * (v + sign * vp)
hb_grad = 0.5 * (prob_h + sign * prob_hp)
batch_sizes = ((v.shape[0], vp.shape[0], *v.shape[1:-1]) if expand else
(*v.shape[:-1], ))
U_grad = torch.zeros_like(self.weights_U).expand(*batch_sizes, -1, -1)
ab_grad = torch.zeros_like(self.aux_bias).expand(*batch_sizes, -1)
vec = [
W_grad.view(*batch_sizes, -1),
U_grad.view(*batch_sizes, -1),
vb_grad,
hb_grad,
ab_grad,
]
if unsqueezed and not expand:
vec = [grad.squeeze_(0) for grad in vec]
return cplx.make_complex(torch.cat(vec, dim=-1))
def ph_grads(self, v):
r"""Computes the gradients of the phase RBM for given input states
:param v: The input state, :math:`\sigma`
:type v: torch.Tensor
:returns: The gradients of all phase RBM parameters
:rtype: torch.Tensor
"""
return cplx.scalar_mult(
cplx.make_complex(self.rbm_ph.effective_energy_gradient(v, reduce=False)),
cplx.I, # need to multiply phase gradient by i
)
def rotate_rho_probs(
nn_state, basis, states, unitaries=None, rho=None, include_extras=False
):
r"""A function that rotates the wavefunction to a different
basis and then computes the resulting Born rule probability of a basis
element.
:param nn_state: The density matrix neural network state.
:type nn_state: qucumber.nn_states.DensityMatrix
:param basis: The basis to rotate the density matrix to.
:type basis: str
:param states: The batch of basis elements to compute the probabilities of.
:type states: torch.Tensor
:param unitaries: A dictionary of (2x2) unitary operators.
:type unitaries: dict(str, torch.Tensor)
:param rho: A density matrix that the user can input to override the neural
network state's density matrix.
:type rho: torch.Tensor
:param include_extras: Whether to include all the terms of the summation as
well as the expanded basis states in the output.
:type include_extras: bool
:returns: Probability of the states wrt the rotated density matrix, and
possibly some extras.
:rtype: torch.Tensor or tuple(torch.Tensor)
"""
Ut, v = _rotate_basis_state(nn_state, basis, states, unitaries=unitaries)
Ut = np.einsum("ib,jb->ijb", Ut, np.conj(Ut))
if rho is None:
rho = nn_state.rho(v).detach()
else:
# pick out the entries of rho that we actually need
idx = _convert_basis_element_to_index(v).long()
rho = rho[:, idx.unsqueeze(0), idx.unsqueeze(1)]
rho = cplx.numpy(rho.cpu())
Ut *= rho
UrhoU_v = cplx.make_complex(Ut).to(dtype=torch.double, device=nn_state.device)
UrhoU = torch.sum(
cplx.real(UrhoU_v), dim=(0, 1)
) # imaginary parts will cancel out anyway
if include_extras:
return UrhoU, UrhoU_v, v
else:
return UrhoU
def psi(self, v):
r"""Compute the (unnormalized) wavefunction of a given vector/matrix of visible states.
.. math::
\psi_{\bm{\lambda}}(\bm{\sigma})
= e^{-\mathcal{E}_{\bm{\lambda}}(\bm{\sigma})/2}
:param v: visible states :math:`\bm{\sigma}`
:type v: torch.Tensor
:returns: Complex object containing the value of the wavefunction for
each visible state
:rtype: torch.Tensor
"""
# vector/tensor of shape (2, len(v))
return cplx.make_complex(self.amplitude(v))
def Pi(self, v, vp):
r"""Calculates an element of the :math:`\Pi` matrix
:param v: One of the visible states, :math:`\sigma`
:type v: torch.Tensor
:param vp: The other visible state, :math`\sigma'`
:type vp: torch.Tensor
:returns: The matrix element given by :math:`\langle\sigma|\Pi|\sigma'\rangle`
:rtype: torch.Tensor
"""
if len(v.shape) < 2 and len(vp.shape) < 2:
exp_arg = self.rbm_am.mixing_term(v + vp)
phase = self.rbm_ph.mixing_term(v - vp)
log_term = ((1 + 2 * exp_arg.exp() * phase.cos() +
(2 * exp_arg).exp()).sqrt().log())
phase_term = ((exp_arg.exp() * phase.sin()) /
(1 + exp_arg.exp() * phase.cos())).atan()
return cplx.make_complex(log_term.sum(), phase_term.sum())
else:
out = torch.zeros(
2,
2**self.num_visible,
2**self.num_visible,
dtype=torch.double,
device=self.device,
)
for i in range(2**self.num_visible):
for j in range(2**self.num_visible):
exp_arg = self.rbm_am.mixing_term(v[i] + vp[j])
phase = self.rbm_ph.mixing_term(v[i] - vp[j])
log_term = ((1 + 2 * exp_arg.exp() * phase.cos() +
(2 * exp_arg).exp()).sqrt().log())
phase_term = ((exp_arg.exp() * phase.sin()) /
(1 + exp_arg.exp() * phase.cos())).atan()
out[0][i][j] = log_term.sum()
out[1][i][j] = phase_term.sum()
return out
def rotate_psi_inner_prod(
nn_state, basis, states, unitaries=None, psi=None, include_extras=False
):
r"""A function that rotates the wavefunction to a different
basis and then computes the resulting amplitude of a batch of basis elements.
:param nn_state: The neural network state (i.e. complex wavefunction or
positive wavefunction).
:type nn_state: qucumber.nn_states.WaveFunctionBase
:param basis: The basis to rotate the wavefunction to.
:type basis: str
:param states: The batch of basis elements to compute the amplitudes of.
:type states: torch.Tensor
:param unitaries: A dictionary of (2x2) unitary operators.
:type unitaries: dict(str, torch.Tensor)
:param psi: A wavefunction that the user can input to override the neural
network state's wavefunction.
:type psi: torch.Tensor
:param include_extras: Whether to include all the terms of the summation as
well as the expanded basis states in the output.
:type include_extras: bool
:returns: Amplitude of the states wrt the rotated wavefunction, and
possibly some extras.
:rtype: torch.Tensor or tuple(torch.Tensor)
"""
Ut, v = _rotate_basis_state(nn_state, basis, states, unitaries=unitaries)
if psi is None:
psi = nn_state.psi(v).detach()
else:
# pick out the entries of psi that we actually need
idx = _convert_basis_element_to_index(v).long()
psi = psi[:, idx]
psi = cplx.numpy(psi.cpu())
Ut *= psi
Upsi_v = cplx.make_complex(Ut).to(dtype=torch.double, device=nn_state.device)
Upsi = torch.sum(Upsi_v, dim=1)
if include_extras:
return Upsi, Upsi_v, v
else:
return Upsi
def psi(self, v):
r"""Compute the (unnormalized) wavefunction of a given vector/matrix of
visible states.
.. math::
\psi(\bm{\sigma})
:param v: visible states :math:`\bm{\sigma}`
:type v: torch.Tensor
:returns: Complex object containing the value of the wavefunction for
each visible state
:rtype: torch.Tensor
"""
# vectors/tensors of shape (len(v),)
amplitude, phase = self.amplitude(v), self.phase(v)
return cplx.make_complex(
amplitude * phase.cos(), # real part
amplitude * phase.sin(), # imaginary part
)
def pi(self, v, vp, expand=True):
r"""Calculates elements of the :math:`\Pi` matrix.
If `expand` is `True`, will return a complex matrix
:math:`A_{ij} = \langle\sigma_i|\Pi|\sigma'_j\rangle`.
Otherwise will return a complex vector
:math:`A_{i} = \langle\sigma_i|\Pi|\sigma'_i\rangle`.
:param v: A batch of visible states, :math:`\sigma`.
:type v: torch.Tensor
:param vp: The other batch of visible state, :math:`\sigma'`.
:type vp: torch.Tensor
:param expand: Whether to return a matrix (`True`) or a vector (`False`).
:type expand: bool
:returns: The matrix elements given by :math:`\langle\sigma|\Pi|\sigma'\rangle`
:rtype: torch.Tensor
"""
m_am = F.linear(v, self.rbm_am.weights_U, self.rbm_am.aux_bias)
mp_am = F.linear(vp, self.rbm_am.weights_U, self.rbm_am.aux_bias)
m_ph = F.linear(v, self.rbm_ph.weights_U)
mp_ph = F.linear(vp, self.rbm_ph.weights_U)
if expand and v.dim() >= 2:
m_am = m_am.unsqueeze_(1)
m_ph = m_ph.unsqueeze_(1)
if expand and vp.dim() >= 2:
mp_am = mp_am.unsqueeze_(0)
mp_ph = mp_ph.unsqueeze_(0)
exp_arg = (m_am + mp_am) / 2
phase = (m_ph - mp_ph) / 2
real = ((1 + 2 * exp_arg.exp() * phase.cos() +
(2 * exp_arg).exp()).sqrt().log().sum(-1))
imag = torch.atan2((exp_arg.exp() * phase.sin()),
(1 + exp_arg.exp() * phase.cos())).sum(-1)
return cplx.make_complex(real, imag)
def apply(self, nn_state, samples):
r"""Computes the average magnetization along Y of each sample in the given batch of samples.
Assumes that the computational basis that the NeuralState was trained
on was the Z basis.
:param nn_state: The NeuralState that drew the samples.
:type nn_state: qucumber.nn_states.NeuralStateBase
:param samples: A batch of samples to calculate the observable on.
Must be using the :math:`\sigma_i = 0, 1` convention.
:type samples: torch.Tensor
"""
samples = samples.to(device=nn_state.device)
# vectors of shape: (2, num_samples,)
denom = nn_state.importance_sampling_denominator(samples)
numer_sum = torch.zeros_like(denom)
for i in range(samples.shape[-1]): # sum over spin sites
coeff = to_pm1(samples[..., i])
coeff = cplx.make_complex(torch.zeros_like(coeff), coeff)
samples_ = flip_spin(i, samples.clone()) # flip the spin at site i
# compute the numerator of the importance and add it to the running sum
numer = nn_state.importance_sampling_numerator(samples_, samples)
numer = cplx.elementwise_mult(numer, coeff)
numer_sum.add_(numer)
numer_sum = cplx.elementwise_division(numer_sum, denom)
# take real part (imaginary part should be approximately zero)
# and divide by number of spins
res = cplx.real(numer_sum).div_(samples.shape[-1])
if self.absolute:
return res.abs_()
else:
return res
def psi(self, v):
return cplx.make_complex(self.amplitude(v))
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 pi_grad(self, v, vp, phase=False, expand=False):
r"""Calculates the gradient of the :math:`\Pi` matrix with
respect to the amplitude RBM parameters for two input states
:param v: One of the visible states, :math:`\sigma`
:type v: torch.Tensor
:param vp: The other visible state, :math`\sigma'`
:type vp: torch.Tensor
:param phase: Whether to compute the gradients for the phase RBM (`True`)
or the amplitude RBM (`False`)
:type phase: bool
:returns: The matrix element of the gradient given by
:math:`\langle\sigma|\nabla_\lambda\Pi|\sigma'\rangle`
:rtype: torch.Tensor
"""
unsqueezed = v.dim() < 2 or vp.dim() < 2
v = (v.unsqueeze(0) if v.dim() < 2 else v).to(self.rbm_am.weights_W)
vp = (vp.unsqueeze(0) if vp.dim() < 2 else vp).to(
self.rbm_am.weights_W)
if expand:
arg_real = 0.5 * (F.linear(v, self.rbm_am.weights_U,
self.rbm_am.aux_bias).unsqueeze_(1) +
F.linear(vp, self.rbm_am.weights_U,
self.rbm_am.aux_bias).unsqueeze_(0))
arg_imag = 0.5 * (
F.linear(v, self.rbm_ph.weights_U).unsqueeze_(1) -
F.linear(vp, self.rbm_ph.weights_U).unsqueeze_(0))
else:
arg_real = self.rbm_am.mixing_term(v + vp)
arg_imag = self.rbm_ph.mixing_term(v - vp)
sig = cplx.sigmoid(arg_real, arg_imag)
batch_sizes = ((v.shape[0], vp.shape[0], *v.shape[1:-1]) if expand else
(*v.shape[:-1], ))
W_grad = torch.zeros_like(self.rbm_am.weights_W).expand(
*batch_sizes, -1, -1)
vb_grad = torch.zeros_like(self.rbm_am.visible_bias).expand(
*batch_sizes, -1)
hb_grad = torch.zeros_like(self.rbm_am.hidden_bias).expand(
*batch_sizes, -1)
if phase:
temp = (v.unsqueeze(1) - vp.unsqueeze(0)) if expand else (v - vp)
sig = cplx.scalar_mult(sig, cplx.I)
ab_grad_real = torch.zeros_like(self.rbm_ph.aux_bias).expand(
*batch_sizes, -1)
ab_grad_imag = ab_grad_real.clone()
else:
temp = (v.unsqueeze(1) + vp.unsqueeze(0)) if expand else (v + vp)
ab_grad_real = cplx.real(sig)
ab_grad_imag = cplx.imag(sig)
U_grad = 0.5 * torch.einsum("c...j,...k->c...jk", sig, temp)
U_grad_real = cplx.real(U_grad)
U_grad_imag = cplx.imag(U_grad)
vec_real = [
W_grad.view(*batch_sizes, -1),
U_grad_real.view(*batch_sizes, -1),
vb_grad,
hb_grad,
ab_grad_real,
]
vec_imag = [
W_grad.view(*batch_sizes, -1).clone(),
U_grad_imag.view(*batch_sizes, -1),
vb_grad.clone(),
hb_grad.clone(),
ab_grad_imag,
]
if unsqueezed and not expand:
vec_real = [grad.squeeze_(0) for grad in vec_real]
vec_imag = [grad.squeeze_(0) for grad in vec_imag]
return cplx.make_complex(torch.cat(vec_real, dim=-1),
torch.cat(vec_imag, dim=-1))
def GammaM_grad(self, v, vp, reduce=False):
r"""Calculates an element of the gradient of
the :math:`\Gamma^{(-)}` matrix
:param v: One of the visible states, :math:`\sigma`
:type v: torch.Tensor
:param vp: The other visible state, :math`\sigma'`
:type vp: torch.Tensor
:returns: The matrix element given by
:math:`\langle\sigma|\nabla_\mu\Gamma^{(-)}|\sigma'\rangle`
:rtype: torch.Tensor
"""
prob_h = self.prob_h_given_v(v)
prob_hp = self.prob_h_given_v(vp)
if v.dim() < 2:
W_grad = 0.5 * (torch.ger(prob_h, v) - torch.ger(prob_hp, vp))
U_grad = torch.zeros_like(self.weights_U)
vb_grad = 0.5 * (v - vp)
hb_grad = 0.5 * (prob_h - prob_hp)
ab_grad = torch.zeros_like(self.aux_bias)
else:
# Don't think this works, but the 'else' option does. Never use reduce
if reduce:
W_grad = 0.5 * (torch.matmul(prob_h.t(), v) -
torch.matmul(prob_hp.t(), vp))
U_grad = torch.zeros(
(v.shape[0], self.weights_U.shape[0],
self.weights_U.shape[1]),
dtype=torch.double,
device=self.device,
)
vb_grad = 0.5 * torch.sum(v - vp, 0)
hb_grad = 0.5 * torch.sum(prob_h - prob_hp, 0)
ab_grad = torch.zeros((v.shape[0], self.num_aux),
dtype=torch.double,
device=self.device)
else:
W_grad = 0.5 * (torch.einsum("ij,ik->ijk", prob_h, v) -
torch.einsum("ij,ik->ijk", prob_hp, vp))
U_grad = torch.zeros(
(v.shape[0], self.weights_U.shape[0],
self.weights_U.shape[1]),
dtype=torch.double,
device=self.device,
)
vb_grad = 0.5 * (v - vp)
hb_grad = 0.5 * (prob_h - prob_hp)
ab_grad = torch.zeros((v.shape[0], self.num_aux),
dtype=torch.double,
device=self.device)
vec = [
W_grad.view(v.size()[0], -1),
U_grad.view(v.size()[0], -1),
vb_grad,
hb_grad,
ab_grad,
]
return cplx.make_complex(torch.cat(vec, dim=1))
return cplx.make_complex(
parameters_to_vector([W_grad, U_grad, vb_grad, hb_grad, ab_grad]))
def importance_sampling_denominator(self, v):
return cplx.make_complex(self.probability(v))
def test_bad_complex_matrix(self):
with self.assertRaises(RuntimeError):
x = torch.tensor([[1, 2, 3]])
y = torch.tensor([[4, 5, 6, 7]])
return cplx.make_complex(x, y)
