Replace dense matrix expansion with a tensordot implementation (#17)

* replaced expand with mat_vec_product

* suggested changes

* suggested changes

pennylane_forest/device.py

@@ -246,3 +246,77 @@ def reset(self):
     @property
     def operations(self):
         return set(self._operation_map.keys())
+
+    def mat_vec_product(self, mat, vec, wires):
+        r"""Apply multiplication of a matrix to subsystems of the quantum state.
+
+        Args:
+            mat (array): matrix to multiply
+            vec (array): state vector to multiply
+            wires (Sequence[int]): target subsystems
+
+        Returns:
+            array: output vector after applying ``mat`` to input ``vec`` on specified subsystems
+        """
+        num_wires = len(wires)
+
+        if mat.shape != (2 ** num_wires, 2 ** num_wires):
+            raise ValueError(
+                f"Please specify a {2**num_wires} x {2**num_wires} matrix for {num_wires} wires."
+            )
+
+        # first, we need to reshape both the matrix and vector
+        # into blocks of 2x2 matrices, in order to do the higher
+        # order matrix multiplication
+
+        # Reshape the matrix to ``size=[2, 2, 2, ..., 2]``,
+        # where ``len(size) == 2*len(wires)``
+        #
+        # The first half of the dimensions correspond to a
+        # 'ket' acting on each wire/qubit, while the second
+        # half of the dimensions correspond to a 'bra' acting
+        # on each wire/qubit.
+        #
+        # E.g., if mat = \sum_{ijkl} (c_{ijkl} |ij><kl|),
+        # and wires=[0, 1], then
+        # the reshaped dimensions of mat are such that
+        # mat[i, j, k, l] == c_{ijkl}.
+        mat = np.reshape(mat, [2] * len(wires) * 2)
+
+        # Reshape the state vector to ``size=[2, 2, ..., 2]``,
+        # where ``len(size) == num_wires``.
+        # Each wire corresponds to a subsystem.
+        #
+        # E.g., if vec = \sum_{ijk}c_{ijk}|ijk>,
+        # the reshaped dimensions of vec are such that
+        # vec[i, j, k] == c_{ijk}.
+        vec = np.reshape(vec, [2] * self.num_wires)
+
+        # Calculate the axes on which the matrix multiplication
+        # takes place. For the state vector, this simply
+        # corresponds to the requested wires. For the matrix,
+        # it is the latter half of the dimensions (the 'bra' dimensions).
+        #
+        # For example, if num_wires=3 and wires=[2, 0], then
+        # axes=((2, 3), (2, 0)). This is equivalent to doing
+        # np.einsum("ijkl,lnk", mat, vec).
+        axes = (np.arange(len(wires), 2 * len(wires)), wires)
+
+        # After the tensor dot operation, the resulting array
+        # will have shape ``size=[2, 2, ..., 2]``,
+        # where ``len(size) == num_wires``, corresponding
+        # to a valid state of the system.
+        tdot = np.tensordot(mat, vec, axes=axes)
+
+        # Tensordot causes the axes given in `wires` to end up in the first positions
+        # of the resulting tensor. This corresponds to a (partial) transpose of
+        # the correct output state
+        # We'll need to invert this permutation to put the indices in the correct place
+        unused_idxs = [idx for idx in range(self.num_wires) if idx not in wires]
+        perm = wires + unused_idxs
+
+        # argsort gives the inverse permutation
+        inv_perm = np.argsort(perm)
+        state_multi_index = np.transpose(tdot, inv_perm)
+
+        return np.reshape(state_multi_index, 2 ** self.num_wires)

pennylane_forest/numpy_wavefunction.py

@@ -100,50 +100,5 @@ def ev(self, A, wires):
           float: expectation value :math:`\left\langle{A}\right\rangle = \left\langle{\psi}\mid A\mid{\psi}\right\rangle`
         """
         # Expand the Hermitian observable over the entire subsystem
-        A = self.expand(A, wires)
-        return np.vdot(self.state, A @ self.state).real
-
-    def expand(self, U, wires):
-        r"""Expand a multi-qubit operator into a full system operator.
-
-        Args:
-            U (array): :math:`2^n \times 2^n` matrix where n = len(wires).
-            wires (Sequence[int]): Target subsystems (order matters! the
-                left-most Hilbert space is at index 0).
-
-        Returns:
-            array: :math:`2^N\times 2^N` matrix. The full system operator.
-        """
-        if self.num_wires == 1:
-            # total number of wires is 1, simply return the matrix
-            return U
-
-        N = self.num_wires
-        wires = np.asarray(wires)
-
-        if np.any(wires < 0) or np.any(wires >= N) or len(set(wires)) != len(wires):
-            raise ValueError("Invalid target subsystems provided in 'wires' argument.")
-
-        if U.shape != (2 ** len(wires), 2 ** len(wires)):
-            raise ValueError("Matrix parameter must be of size (2**len(wires), 2**len(wires))")
-
-        # generate N qubit basis states via the cartesian product
-        tuples = np.array(list(itertools.product([0, 1], repeat=N)))
-
-        # wires not acted on by the operator
-        inactive_wires = list(set(range(N)) - set(wires))
-
-        # expand U to act on the entire system
-        U = np.kron(U, np.identity(2 ** len(inactive_wires)))
-
-        # move active wires to beginning of the list of wires
-        rearranged_wires = np.array(list(wires) + inactive_wires)
-
-        # convert to computational basis
-        # i.e., converting the list of basis state bit strings into
-        # a list of decimal numbers that correspond to the computational
-        # basis state. For example, [0, 1, 0, 1, 1] = 2^3+2^1+2^0 = 11.
-        perm = np.ravel_multi_index(tuples[:, rearranged_wires].T, [2] * N)
-
-        # permute U to take into account rearranged wires
-        return U[:, perm][perm]
+        As = self.mat_vec_product(A, self.state, wires)
+        return np.vdot(self.state, As).real

pennylane_forest/wavefunction.py

@@ -164,50 +164,5 @@ def ev(self, A, wires):
           float: expectation value :math:`\left\langle{A}\right\rangle = \left\langle{\psi}\mid A\mid{\psi}\right\rangle`
         """
         # Expand the Hermitian observable over the entire subsystem
-        A = self.expand(A, wires)
-        return np.vdot(self.state, A @ self.state).real
-
-    def expand(self, U, wires):
-        r"""Expand a multi-qubit operator into a full system operator.
-
-        Args:
-            U (array): :math:`2^n \times 2^n` matrix where n = len(wires).
-            wires (Sequence[int]): Target subsystems (order matters! the
-                left-most Hilbert space is at index 0).
-
-        Returns:
-            array: :math:`2^N\times 2^N` matrix. The full system operator.
-        """
-        if self.num_wires == 1:
-            # total number of wires is 1, simply return the matrix
-            return U
-
-        N = self.num_wires
-        wires = np.asarray(wires)
-
-        if np.any(wires < 0) or np.any(wires >= N) or len(set(wires)) != len(wires):
-            raise ValueError("Invalid target subsystems provided in 'wires' argument.")
-
-        if U.shape != (2 ** len(wires), 2 ** len(wires)):
-            raise ValueError("Matrix parameter must be of size (2**len(wires), 2**len(wires))")
-
-        # generate N qubit basis states via the cartesian product
-        tuples = np.array(list(itertools.product([0, 1], repeat=N)))
-
-        # wires not acted on by the operator
-        inactive_wires = list(set(range(N)) - set(wires))
-
-        # expand U to act on the entire system
-        U = np.kron(U, np.identity(2 ** len(inactive_wires)))
-
-        # move active wires to beginning of the list of wires
-        rearranged_wires = np.array(list(wires) + inactive_wires)
-
-        # convert to computational basis
-        # i.e., converting the list of basis state bit strings into
-        # a list of decimal numbers that correspond to the computational
-        # basis state. For example, [0, 1, 0, 1, 1] = 2^3+2^1+2^0 = 11.
-        perm = np.ravel_multi_index(tuples[:, rearranged_wires].T, [2] * N)
-
-        # permute U to take into account rearranged wires
-        return U[:, perm][perm]
+        As = self.mat_vec_product(A, self.state, wires)
+        return np.vdot(self.state, As).real

tests/conftest.py

@@ -13,6 +13,7 @@
 from pyquil.api import QVMConnection, QVMCompiler, local_qvm
 from pyquil.api._config import PyquilConfig
 from pyquil.api._errors import UnknownApiError
+from pyquil.api._qvm import QVMNotRunning
 
 
 # defaults
@@ -144,7 +145,7 @@ def qvm():
         qvm = QVMConnection(random_seed=52)
         qvm.run(Program(Id(0)), [])
         return qvm
-    except (RequestException, UnknownApiError, TypeError) as e:
+    except (RequestException, UnknownApiError, QVMNotRunning, TypeError) as e:
         return pytest.skip("This test requires QVM connection: {}".format(e))
 
 
@@ -156,7 +157,7 @@ def compiler():
         compiler = QVMCompiler(endpoint=config.quilc_url, device=device)
         compiler.quil_to_native_quil(Program(Id(0)))
         return compiler
-    except (RequestException, UnknownApiError, TypeError) as e:
+    except (RequestException, UnknownApiError, QVMNotRunning, TypeError) as e:
         return pytest.skip("This test requires compiler connection: {}".format(e))
 
 

tests/test_device.py

@@ -0,0 +1,111 @@
+"""
+Unit tests for the abstract ForestDevice class
+"""
+from unittest.mock import patch
+import pytest
+
+import pennylane as qml
+from pennylane import numpy as np
+
+from conftest import I, U, U2, H
+
+from pennylane_forest.device import ForestDevice
+
+
+# make test deterministic
+np.random.seed(42)
+
+
+@patch.multiple(ForestDevice, __abstractmethods__=set())
+class TestMatVecProduct:
+    """Unit tests matrix-vector product method of the ForestDevice class"""
+
+    def test_incorrect_matrix_size(self, tol):
+        """Test that an exception is raised if the input matrix is
+        applied to the incorrect number of wires"""
+        wires = 3
+        dev = ForestDevice(wires=wires, shots=1)
+
+        # create a random length 2**wires vector
+        vec = np.random.random([2 ** wires])
+
+        # apply 2 qubit unitary to the full system
+        with pytest.raises(ValueError, match="specify a 8 x 8 matrix for 3 wires"):
+            res = dev.mat_vec_product(U2, vec, wires=[0, 1, 2])
+
+    def test_full_system(self, tol):
+        """Test that matrix-vector multiplication
+        over the entire subsystem agrees with standard
+        dense matrix multiplication"""
+        wires = 3
+        dev = ForestDevice(wires=wires, shots=1)
+
+        # create a random length 2**wires vector
+        vec = np.random.random([2 ** wires])
+
+        # make a 3 qubit unitary
+        mat = np.kron(U2, H)
+
+        # apply to the system
+        res = dev.mat_vec_product(mat, vec, wires=[0, 1, 2])
+
+        # perform the same operation using dense matrix multiplication
+        expected = mat @ vec
+
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_permutation_full_system(self, tol):
+        """Test that matrix-vector multiplication
+        over a permutation of the system agrees with standard
+        dense matrix multiplication"""
+        wires = 2
+        dev = ForestDevice(wires=wires, shots=1)
+
+        # create a random length 2**wires vector
+        vec = np.random.random([2 ** wires])
+
+        # apply to the system
+        res = dev.mat_vec_product(U2, vec, wires=[1, 0])
+
+        # perform the same operation using dense matrix multiplication
+        perm = np.array([0, 2, 1, 3])
+        expected = U2[:, perm][perm] @ vec
+
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_consecutive_subsystem(self, tol):
+        """Test that matrix-vector multiplication
+        over a consecutive subset of the system agrees with standard
+        dense matrix multiplication"""
+        wires = 3
+        dev = ForestDevice(wires=wires, shots=1)
+
+        # create a random length 2**wires vector
+        vec = np.random.random([2 ** wires])
+
+        # apply a 2 qubit unitary to wires 1, 2
+        res = dev.mat_vec_product(U2, vec, wires=[1, 2])
+
+        # perform the same operation using dense matrix multiplication
+        expected = np.kron(I, U2) @ vec
+
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_non_consecutive_subsystem(self, tol):
+        """Test that matrix-vector multiplication
+        over a non-consecutive subset of the system agrees with standard
+        dense matrix multiplication"""
+        wires = 3
+        dev = ForestDevice(wires=wires, shots=1)
+
+        # create a random length 2**wires vector
+        vec = np.random.random([2 ** wires])
+
+        # apply a 2 qubit unitary to wires 1, 2
+        res = dev.mat_vec_product(U2, vec, wires=[2, 0])
+
+        # perform the same operation using dense matrix multiplication
+        perm = np.array([0, 4, 1, 5, 2, 6, 3, 7])
+        expected = np.kron(U2, I)[:, perm][perm] @ vec
+
+        assert np.allclose(res, expected, atol=tol, rtol=0)