Repository: quantumlib/Cirq

Issue:
Change input to SimulatesAmplitudes.compute_amplitudes from 2-d numpy array to 1-d int sequence
@viathor are you doing this in #2114 ? If not I can do it.


Discussion:
I'd rather keep the PR focused on one change, so nope. Feel free to do this. If you don't have time, I can also take it. Your call.
Ok I'll do it.

Resolving patch:
diff --git a/cirq/sim/simulator.py b/cirq/sim/simulator.py
--- a/cirq/sim/simulator.py
+++ b/cirq/sim/simulator.py
@@ -33,6 +33,7 @@
     Hashable,
     Iterator,
     List,
+    Sequence,
     Tuple,
     Union,
     Optional,
@@ -191,10 +192,10 @@ class SimulatesAmplitudes(metaclass=abc.ABCMeta):
     def compute_amplitudes(
             self,
             program: Union[circuits.Circuit, schedules.Schedule],
-            bitstrings: np.ndarray,
+            bitstrings: Sequence[int],
             param_resolver: 'study.ParamResolverOrSimilarType' = None,
             qubit_order: ops.QubitOrderOrList = ops.QubitOrder.DEFAULT,
-    ) -> List[complex]:
+    ) -> Sequence[complex]:
         """Computes the desired amplitudes.
 
         The initial state is assumed to be the all zeros state.
@@ -202,9 +203,9 @@ def compute_amplitudes(
         Args:
             program: The circuit or schedule to simulate.
             bitstrings: The bitstrings whose amplitudes are desired, input
-                as a two-dimensional array of bools. The first dimension
-                indexes the bitstrings and the second dimension indexes
-                the bits within a bitstring.
+                as an integer array where each integer is formed from measured
+                qubit values according to `qubit_order` from most to least
+                significant qubit, i.e. in big-endian ordering.
             param_resolver: Parameters to run with the program.
             qubit_order: Determines the canonical ordering of the qubits. This
                 is often used in specifying the initial state, i.e. the
@@ -221,10 +222,10 @@ def compute_amplitudes(
     def compute_amplitudes_sweep(
             self,
             program: Union[circuits.Circuit, schedules.Schedule],
-            bitstrings: np.ndarray,
+            bitstrings: Sequence[int],
             params: study.Sweepable,
             qubit_order: ops.QubitOrderOrList = ops.QubitOrder.DEFAULT,
-    ) -> List[List[complex]]:
+    ) -> Sequence[Sequence[complex]]:
         """Computes the desired amplitudes.
 
         The initial state is assumed to be the all zeros state.
@@ -232,9 +233,9 @@ def compute_amplitudes_sweep(
         Args:
             program: The circuit or schedule to simulate.
             bitstrings: The bitstrings whose amplitudes are desired, input
-                as a two-dimensional array of bools. The first dimension
-                indexes the bitstrings and the second dimension indexes
-                the bits within a bitstring.
+                as an integer array where each integer is formed from measured
+                qubit values according to `qubit_order` from most to least
+                significant qubit, i.e. in big-endian ordering.
             params: Parameters to run with the program.
             qubit_order: Determines the canonical ordering of the qubits. This
                 is often used in specifying the initial state, i.e. the
diff --git a/cirq/sim/wave_function_simulator.py b/cirq/sim/wave_function_simulator.py
--- a/cirq/sim/wave_function_simulator.py
+++ b/cirq/sim/wave_function_simulator.py
@@ -16,7 +16,8 @@
 
 import abc
 
-from typing import Any, Dict, Iterator, Hashable, List, Optional, Union
+from typing import (Any, Dict, Iterator, Hashable, List, Optional, Sequence,
+                    Union, cast)
 
 import numpy as np
 
@@ -164,27 +165,28 @@ def compute_displays_sweep(
     def compute_amplitudes_sweep(
             self,
             program: Union[circuits.Circuit, schedules.Schedule],
-            bitstrings: np.ndarray,
+            bitstrings: Sequence[int],
             params: study.Sweepable,
             qubit_order: ops.QubitOrderOrList = ops.QubitOrder.DEFAULT,
-    ) -> List[List[complex]]:
+    ) -> Sequence[Sequence[complex]]:
+        if isinstance(bitstrings, np.ndarray) and len(bitstrings.shape) > 1:
+            raise ValueError('The list of bitstrings must be input as a '
+                             '1-dimensional array of ints. Got an array with '
+                             f'shape {bitstrings.shape}.')
+
         circuit = (program if isinstance(program, circuits.Circuit) else
                    program.to_circuit())
-
         trial_results = self.simulate_sweep(circuit, params, qubit_order)
 
-        qid_shape = circuit.qid_shape(qubit_order=qubit_order)
-        amplitude_indices = [
-            value.big_endian_digits_to_int(bitstring, base=qid_shape)
-            for bitstring in bitstrings
-        ]
+        # 1-dimensional tuples don't trigger advanced Numpy array indexing
+        # https://docs.scipy.org/doc/numpy/reference/arrays.indexing.html
+        if isinstance(bitstrings, tuple):
+            bitstrings = list(bitstrings)
 
         all_amplitudes = []
         for trial_result in trial_results:
-            # mypy doesn't know that this trial result has a final_state
-            # attribute
-            final_state = trial_result.final_state  # type: ignore
-            amplitudes = [final_state[index] for index in amplitude_indices]
+            trial_result = cast(WaveFunctionTrialResult, trial_result)
+            amplitudes = trial_result.final_state[bitstrings]
             all_amplitudes.append(amplitudes)
 
         return all_amplitudes


Test patch:
diff --git a/cirq/sim/sparse_simulator_test.py b/cirq/sim/sparse_simulator_test.py
--- a/cirq/sim/sparse_simulator_test.py
+++ b/cirq/sim/sparse_simulator_test.py
@@ -714,18 +714,25 @@ def test_compute_amplitudes():
     c = cirq.Circuit.from_ops(cirq.X(a), cirq.H(a), cirq.H(b))
     sim = cirq.Simulator()
 
-    result = sim.compute_amplitudes(c, np.array([[0, 0]]))
+    result = sim.compute_amplitudes(c, [0])
     np.testing.assert_allclose(np.array(result), np.array([0.5]))
 
-    result = sim.compute_amplitudes(c, np.array([[0, 1], [1, 0], [1, 1]]))
+    result = sim.compute_amplitudes(c, [1, 2, 3])
     np.testing.assert_allclose(np.array(result), np.array([0.5, -0.5, -0.5]))
 
-    result = sim.compute_amplitudes(c,
-                                    np.array([[0, 1], [1, 0], [1, 1]]),
-                                    qubit_order=(b, a))
+    result = sim.compute_amplitudes(c, (1, 2, 3), qubit_order=(b, a))
     np.testing.assert_allclose(np.array(result), np.array([-0.5, 0.5, -0.5]))
 
 
+def test_compute_amplitudes_bad_input():
+    a, b = cirq.LineQubit.range(2)
+    c = cirq.Circuit.from_ops(cirq.X(a), cirq.H(a), cirq.H(b))
+    sim = cirq.Simulator()
+
+    with pytest.raises(ValueError, match='1-dimensional'):
+        _ = sim.compute_amplitudes(c, np.array([[0, 0]]))
+
+
 def test_run_sweep_parameters_not_resolved():
     a = cirq.LineQubit(0)
     simulator = cirq.Simulator()