Set a wavefunction's amplitudes to be a complex number instead of an array

referenceqvm/qvm_wavefunction.py

@@ -115,7 +115,7 @@ def measurement(self, qubit_index, psi=None):
             proj_psi = measure_0.dot(self.wf)
         else:
             proj_psi = measure_0.dot(psi)
-        prob_zero = np.dot(np.conj(proj_psi).T, proj_psi)[0, 0]
+        prob_zero = np.dot(np.conj(proj_psi).T, proj_psi)
 
         # generate random number to 'roll' for measurement
         if np.random.random() < prob_zero:
@@ -372,8 +372,8 @@ def wavefunction(self, pyquil_program, classical_addresses=None):
             mask = None
 
         # setup wavefunction
-        self.wf = np.zeros((2 ** self.num_qubits, 1), dtype=np.complex128)
-        self.wf[0, 0] = 1.0
+        self.wf = np.zeros(2 ** self.num_qubits, dtype=np.complex128)
+        self.wf[0] = complex(1.0, 0)
 
         # evolve wf with program, via kernel
         self.kernel()

referenceqvm/tests/test_arbitrary_state_prep.py

@@ -13,7 +13,7 @@ def test_generate_arbitrary_states(qvm):
         wf, _ = qvm.wavefunction(p)
 
         # check actual part of wavefunction
-        assert np.allclose(v.reshape((-1, 1)), wf.amplitudes[:len(v), :] * norm)
+        assert np.allclose(v.reshape(-1), wf.amplitudes[:len(v)] * norm)
 
         # check remaining zeros part of wavefunction
         assert np.allclose(np.zeros((wf.amplitudes.shape[0] - len(v), 1)),

referenceqvm/tests/test_unitary.py

@@ -37,7 +37,7 @@ def test_identity(qvm_unitary):
 def test_qaoa_unitary(qvm_unitary):
     wf_true = [0.00167784 + 1.00210180e-05*1j, 0.50000000 - 4.99997185e-01*1j,
                0.50000000 - 4.99997185e-01*1j, 0.00167784 + 1.00210180e-05*1j]
-    wf_true = np.reshape(np.array(wf_true), (4, 1))
+
     prog = Program()
     prog.inst([RY(np.pi/2)(0), RX(np.pi)(0),
                RY(np.pi/2)(1), RX(np.pi)(1),
@@ -46,8 +46,8 @@ def test_qaoa_unitary(qvm_unitary):
                RX(-2*2.74973750579)(0), RX(-2*2.74973750579)(1)])
 
     test_unitary = qvm_unitary.unitary(prog)
-    wf_test = np.zeros((4, 1))
-    wf_test[0, 0] = 1.0
+    wf_test = np.zeros(4)
+    wf_test[0] = 1.0
     wf_test = test_unitary.dot(wf_test)
     assert np.allclose(wf_test, wf_true)
 

referenceqvm/tests/test_wavefunction.py

@@ -30,28 +30,28 @@ def test_belltest(qvm):
     """
     prog = Program().inst([Hgate(0), CNOTgate(0, 1)])
     bellout, _ = qvm.wavefunction(prog)
-    bell = np.zeros((4, 1))
-    bell[0, 0] = bell[-1, 0] = 1.0 / np.sqrt(2)
+    bell = np.zeros(4)
+    bell[0] = bell[-1] = 1.0 / np.sqrt(2)
     assert np.allclose(bellout.amplitudes, bell)
 
 
 def test_occupation_basis(qvm):
     prog = Program().inst([Xgate(0), Xgate(1), Igate(2), Igate(3)])
-    state = np.zeros((2 ** 4, 1))
-    state[3, 0] = 1.0
+    state = np.zeros(2 ** 4)
+    state[3] = 1.0
     meanfield_state, _ = qvm.wavefunction(prog)
     assert np.allclose(meanfield_state.amplitudes, state)
 
 
 def test_exp_circuit(qvm):
-    true_wf = np.array([[ 0.54030231-0.84147098j], 
-                        [ 0.00000000+0.j],
-                        [ 0.00000000+0.j],
-                        [ 0.00000000+0.j],
-                        [ 0.00000000+0.j],
-                        [ 0.00000000+0.j],
-                        [ 0.00000000+0.j],
-                        [ 0.00000000+0.j]])
+    true_wf = np.array([0.54030231-0.84147098j,
+                        0.00000000+0.j,
+                        0.00000000+0.j,
+                        0.00000000+0.j,
+                        0.00000000+0.j,
+                        0.00000000+0.j,
+                        0.00000000+0.j,
+                        0.00000000+0.j])
 
     create2kill1 = PauliTerm("X", 1, -0.25)*PauliTerm("Y", 2)
     create2kill1 += PauliTerm("Y", 1, 0.25)*PauliTerm("Y", 2)
@@ -64,14 +64,13 @@ def test_exp_circuit(qvm):
         prog += single_exp_prog
 
     wf, _ = qvm.wavefunction(prog)
-    wf = np.reshape(wf.amplitudes, (-1, 1))
+    wf = np.reshape(wf.amplitudes, -1)
     assert np.allclose(wf, true_wf)
 
 
 def test_qaoa_circuit(qvm):
     wf_true = [0.00167784 + 1.00210180e-05*1j, 0.50000000 - 4.99997185e-01*1j,
                0.50000000 - 4.99997185e-01*1j, 0.00167784 + 1.00210180e-05*1j]
-    wf_true = np.reshape(np.array(wf_true), (4, 1))
     prog = Program()
     prog.inst([RYgate(np.pi/2)(0), RXgate(np.pi)(0),
                RYgate(np.pi/2)(1), RXgate(np.pi)(1),
@@ -139,5 +138,4 @@ def test_larger_qaoa_circuit(qvm):
                         -1.24927731e-01+0.00329533*1j,
                         8.43771693e-05-0.1233845*1j])
 
-    wf_true = np.reshape(wf_true, (2 ** 4, 1))
     assert np.allclose(wf_test.amplitudes, wf_true)