Rename and add default behaviour for auto_unsqueeze_arg decorator

Author: emerali

qucumber/nn_states/complex_wavefunction.py

@@ -176,11 +176,7 @@ def rotated_gradient(self, basis, sample):
         )
         inv_Upsi = cplx.inverse(Upsi)
 
-        vr = v.reshape(-1, v.shape[-1])
-        raw_grads = [
-            self.am_grads(vr).reshape(2, *v.shape[:-1], -1),
-            self.ph_grads(vr).reshape(2, *v.shape[:-1], -1),
-        ]
+        raw_grads = [self.am_grads(v), self.ph_grads(v)]
 
         rotated_grad = [cplx.einsum("s...,s...g->...g", Upsi_v, g) for g in raw_grads]
         grad = [

qucumber/nn_states/neural_state.py

@@ -177,6 +177,9 @@ def generate_hilbert_space(self, size=None, device=None):
 
     def normalization(self, space):
         r"""Compute the normalization constant of the state.
+        In the case of a pure state, this is the norm of the unnormalized wavefunction.
+        In the case of a mixed state, this is the trace of the unnormalized density
+        matrix.
 
         .. math::
 
@@ -190,7 +193,7 @@ def normalization(self, space):
         return self.rbm_am.partition(space)
 
     def compute_normalization(self, space):
-        """Alias for `normalization`"""
+        """Alias for :func:`normalization<qucumber.nn_states.NeuralStateBase.normalization>`"""
         return self.normalization(space)
 
     def save(self, location, metadata=None):

qucumber/nn_states/positive_wavefunction.py

@@ -17,7 +17,7 @@
 
 from qucumber import _warn_on_missing_gpu
 from qucumber.rbm import BinaryRBM
-from qucumber.utils import cplx, auto_unsqueeze_arg
+from qucumber.utils import cplx, auto_unsqueeze_args
 from .wavefunction import WaveFunctionBase
 
 
@@ -91,7 +91,7 @@ def amplitude(self, v):
         """
         return super().amplitude(v)
 
-    @auto_unsqueeze_arg(1)
+    @auto_unsqueeze_args()
     def phase(self, v):
         r"""Compute the phase of a given vector/matrix of visible states.
 

qucumber/nn_states/wavefunction.py

@@ -86,20 +86,6 @@ def importance_sampling_numerator(self, v, vp):
     def importance_sampling_denominator(self, vp):
         return self.psi(vp)
 
-    def normalization(self, space):
-        r"""Compute the norm of the wavefunction.
-
-        .. math::
-
-            Z_{\bm{\lambda}}=
-            \sum_{\bm{\sigma}} p_{\bm{\lambda}}(\bm{\sigma})
-
-        :param space: A rank 2 tensor of the entire visible space.
-        :type space: torch.Tensor
-
-        """
-        return super().normalization(space)
-
 
 # make module path show up properly in sphinx docs
 WaveFunctionBase.__module__ = "qucumber.nn_states"

qucumber/rbm/binary_rbm.py

@@ -20,7 +20,7 @@
 from torch.nn.utils import parameters_to_vector
 
 from qucumber import _warn_on_missing_gpu
-from qucumber.utils import auto_unsqueeze_arg
+from qucumber.utils import auto_unsqueeze_args
 
 
 class BinaryRBM(nn.Module):
@@ -71,6 +71,7 @@ def initialize_parameters(self, zero_weights=False):
             requires_grad=False,
         )
 
+    @auto_unsqueeze_args()
     def effective_energy(self, v):
         r"""The effective energies of the given visible states.
 
@@ -88,7 +89,7 @@ def effective_energy(self, v):
         :returns: The effective energies of the given visible states.
         :rtype: torch.Tensor
         """
-        v = (v.unsqueeze(0) if v.dim() < 2 else v).to(self.weights)
+        v = v.to(self.weights)
         visible_bias_term = torch.matmul(v, self.visible_bias)
         hid_bias_term = F.softplus(F.linear(v, self.weights, self.hidden_bias)).sum(-1)
 
@@ -124,7 +125,7 @@ def effective_energy_gradient(self, v, reduce=True):
             vec = [W_grad.view(*v.shape[:-1], -1), vb_grad, hb_grad]
             return torch.cat(vec, dim=-1)
 
-    @auto_unsqueeze_arg(1)
+    @auto_unsqueeze_args()
     def prob_v_given_h(self, h, out=None):
         """Given a hidden unit configuration, compute the probability
         vector of the visible units being on.
@@ -145,7 +146,7 @@ def prob_v_given_h(self, h, out=None):
             .clamp_(min=0, max=1)
         )
 
-    @auto_unsqueeze_arg(1)
+    @auto_unsqueeze_args()
     def prob_h_given_v(self, v, out=None):
         """Given a visible unit configuration, compute the probability
         vector of the hidden units being on.

qucumber/rbm/purification_rbm.py

@@ -19,7 +19,7 @@
 from torch.nn import functional as F
 from torch.nn.utils import parameters_to_vector
 
-from qucumber.utils import cplx, auto_unsqueeze_arg
+from qucumber.utils import cplx, auto_unsqueeze_args
 from qucumber import _warn_on_missing_gpu
 
 
@@ -127,6 +127,7 @@ def initialize_parameters(self, zero_weights=False):
             requires_grad=False,
         )
 
+    @auto_unsqueeze_args()
     def effective_energy(self, v, a=None):
         r"""Computes the equivalent of the "effective energy" for the RBM. If
         `a` is `None`, will analytically trace out the auxiliary units.
@@ -139,7 +140,7 @@ def effective_energy(self, v, a=None):
         :returns: The "effective energy" of the RBM. Shape (b,) or (1,).
         :rtype: torch.Tensor
         """
-        v = (v.unsqueeze(0) if v.dim() < 2 else v).to(self.weights_W)
+        v = v.to(self.weights_W)
 
         vis_term = torch.matmul(v, self.visible_bias) + F.softplus(
             F.linear(v, self.weights_W, self.hidden_bias)
@@ -191,7 +192,7 @@ def effective_energy_gradient(self, v, reduce=True):
             vec = [W_grad, U_grad, vb_grad, hb_grad, ab_grad]
             return torch.cat(vec, dim=-1)
 
-    @auto_unsqueeze_arg(1)
+    @auto_unsqueeze_args()
     def prob_h_given_v(self, v, out=None):
         r"""Given a visible unit configuration, compute the probability
         vector of the hidden units being on
@@ -212,7 +213,7 @@ def prob_h_given_v(self, v, out=None):
             .clamp_(min=0, max=1)
         )
 
-    @auto_unsqueeze_arg(1)
+    @auto_unsqueeze_args()
     def prob_a_given_v(self, v, out=None):
         r"""Given a visible unit configuration, compute the probability
         vector of the auxiliary units being on
@@ -233,7 +234,7 @@ def prob_a_given_v(self, v, out=None):
             .clamp_(min=0, max=1)
         )
 
-    @auto_unsqueeze_arg(1, 2)
+    @auto_unsqueeze_args(1, 2)
     def prob_v_given_ha(self, h, a, out=None):
         r"""Given a hidden and auxiliary unit configuration, compute
         the probability vector of the hidden units being on
@@ -335,7 +336,7 @@ def gibbs_steps(self, k, initial_state, overwrite=False):
 
         return v
 
-    @auto_unsqueeze_arg(1)
+    @auto_unsqueeze_args()
     def mixing_term(self, v):
         r"""Describes the extent of mixing in the system,
             :math:`V_\theta = \frac{1}{2}U_\theta \bm{\sigma} + d_\theta`

qucumber/utils/__init__.py

@@ -17,9 +17,12 @@
 from functools import wraps
 
 
-class auto_unsqueeze_arg:
+class auto_unsqueeze_args:
     def __init__(self, *arg_indices):
-        self.arg_indices = arg_indices
+        self.arg_indices = list(arg_indices)
+
+        if len(self.arg_indices) == 0:
+            self.arg_indices.append(1)
 
     def __call__(self, f):
         @wraps(f)

qucumber/utils/data.py

@@ -77,7 +77,7 @@ def load_data_DM(
     :type bases_path: str
 
     :returns: A list of all input parameters, with the real and imaginary parts
-              combined into one (PyTorch-hack) complex matrix.
+              of the target density matrix (if provided) combined into one complex matrix.
     :rtype: list
     """
     data = []

tests/test_models_misc.py

@@ -150,11 +150,11 @@ def test_positive_wavefunction_psi():
     nn_state = PositiveWaveFunction(10, gpu=False)
 
     vis_state = torch.ones(10).to(dtype=torch.double)
-    actual_psi = nn_state.psi(vis_state)[1].to(vis_state)
-    expected_psi = torch.zeros(1).to(vis_state)
+    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, expected_psi), msg
+    assert torch.equal(actual_psi_im, expected_psi_im), msg
 
 
 def test_density_matrix_hermiticity():

tests/test_observables.py

@@ -18,8 +18,9 @@
 import pytest
 import torch
 
-import qucumber.observables as observables
-from qucumber.nn_states import WaveFunctionBase
+from qucumber import observables
+from qucumber.nn_states import WaveFunctionBase, NeuralStateBase
+from qucumber.utils import cplx, auto_unsqueeze_args
 
 
 class MockWaveFunction(WaveFunctionBase):
@@ -66,35 +67,16 @@ def device(self, new_val):
     def networks(self):
         return ["rbm_am"]
 
+    @auto_unsqueeze_args()
     def phase(self, v):
-        if v.dim() == 1:
-            v = v.unsqueeze(0)
-            unsqueezed = True
-        else:
-            unsqueezed = False
-
-        phase = torch.zeros(v.shape[0])
-
-        if unsqueezed:
-            return phase.squeeze_(0)
-        else:
-            return phase
+        return torch.zeros(v.shape[0])
 
     def amplitude(self, v):
         return torch.ones(v.size(0)) / torch.sqrt(torch.tensor(float(self.nqubits)))
 
     def psi(self, v):
         # vector/tensor of shape (len(v),)
-        amplitude = self.amplitude(v)
-
-        # complex vector; shape: (2, len(v))
-        psi = torch.zeros((2,) + amplitude.shape).to(
-            dtype=torch.double, device=self.device
-        )
-        psi[0] = amplitude
-
-        # squeeze down to complex scalar if there was only one visible state
-        return psi.squeeze()
+        return cplx.make_complex(self.amplitude(v))
 
 
 @pytest.fixture(scope="module")