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)