def density_matrix_fidelity(nn_state, target, v_space, **kwargs):
r"""Calculate the fidelity of the reconstructed density matrix
given the exact target density matrix
:param nn_state: The neural network state (i.e. current density matrix)
:type nn_state: qucumber.nn_states.DensityMatrix
:param target: The true density matrix of the system
:type target: torch.Tensor
:param v_space: The basis elements of the visible space
:type v_space: torch.Tensor
:param \**kwargs: Extra keyword arguments that may be passed.
Will be ignored.
:returns: The fidelity
:rtype: float
"""
rhoRBM_ = nn_state.rhoRBM(v_space, v_space)
argReal = cplx.real(rhoRBM_).numpy()
argIm = cplx.imag(rhoRBM_).numpy()
rho_rbm_ = argReal + 1j * argIm
argReal = cplx.real(target).numpy()
argIm = cplx.imag(target).numpy()
target_ = argReal + 1j * argIm
sqrt_rho_rbm = sqrtm(rho_rbm_)
arg = sqrtm(np.matmul(sqrt_rho_rbm, np.matmul(target_, sqrt_rho_rbm)))
return np.trace(arg).real
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 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_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_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_positive_wavefunction_psi():
nn_state = PositiveWaveFunction(10, gpu=False)
vis_state = torch.ones(10).to(dtype=torch.double)
actual_psi_im = cplx.imag(nn_state.psi(vis_state)).to(vis_state)
expected_psi_im = torch.zeros(1).squeeze().to(vis_state)
msg = "PositiveWaveFunction is giving a non-zero imaginary part!"
assert torch.equal(actual_psi_im, expected_psi_im), msg
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))
