change into adding a new functionality

qsimcirq/qsim_simulator.py

@@ -435,99 +435,13 @@ def compute_amplitudes_sweep_iter(
             options["s"] = self.get_seed()
             yield simulator_fn(options)
 
-    def simulate(
-        self,
-        program: cirq.AbstractCircuit,
-        param_resolver: cirq.ParamResolverOrSimilarType = None,
-        qubit_order: cirq.QubitOrderOrList = cirq.ops.QubitOrder.DEFAULT,
-        initial_state: Any = None,
-        as_1d_state_vector: bool = False,
-    ) -> cirq.sim.simulator.TSimulationTrialResult:
-        """Simulates the supplied Circuit.
-
-        This method returns a result which allows access to the entire
-        simulator's final state.
-
-        Args:
-            program: The circuit to simulate.
-            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
-                ordering of the computational basis states.
-            initial_state: The initial state for the simulation. The form of
-                this state depends on the simulation implementation. See
-                documentation of the implementing class for details.
-            as_1d_state_vector: Whether the returned state vector is 1D or
-                has number of dimensions equal number of qubits.
-                Operations on 1D representation are significantly slower than
-                the other representation and might not even work.
-                The 1D representation should only be used when the number of qubits
-                is larger than the dimension limit on numpy arrays (numpy/numpy#5744).
-
-        Returns:
-            SimulationTrialResults for the simulation. Includes the final state.
-        """
-        return self.simulate_sweep(
-            program,
-            cirq.study.ParamResolver(param_resolver),
-            qubit_order,
-            initial_state,
-            as_1d_state_vector,
-        )[0]
-
-    def simulate_sweep(
-        self,
-        program: cirq.AbstractCircuit,
-        params: cirq.Sweepable,
-        qubit_order: cirq.QubitOrderOrList = cirq.ops.QubitOrder.DEFAULT,
-        initial_state: Any = None,
-        as_1d_state_vector: bool = False,
-    ) -> List[cirq.sim.simulator.TSimulationTrialResult]:
-        return list(
-            self.simulate_sweep_iter(
-                program, params, qubit_order, initial_state, as_1d_state_vector
-            )
-        )
-
-    def simulate_sweep_iter(
+    def _simulate_impl(
         self,
         program: cirq.Circuit,
         params: cirq.Sweepable,
         qubit_order: cirq.QubitOrderOrList = cirq.QubitOrder.DEFAULT,
         initial_state: Optional[Union[int, np.ndarray]] = None,
-        as_1d_state_vector: bool = False,
-    ) -> Iterator[cirq.StateVectorTrialResult]:
-        """Simulates the supplied Circuit.
-
-        This method returns a result which allows access to the entire
-        wave function. In contrast to simulate, this allows for sweeping
-        over different parameter values.
-
-        Avoid using this method with `use_gpu=True` in the simulator options;
-        when used with GPU this method must copy state from device to host memory
-        multiple times, which can be very slow. This issue is not present in
-        `simulate_expectation_values_sweep`.
-
-        Args:
-            program: The circuit to simulate.
-            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 ordering of the
-              computational basis states.
-            initial_state: The initial state for the simulation. This can either
-              be an integer representing a pure state (e.g. 11010) or a numpy
-              array containing the full state vector. If none is provided, this
-              is assumed to be the all-zeros state.
-            as_1d_state_vector: Whether the returned state vector is 1D or
-                has number of dimensions equal number of qubits.
-
-        Returns:
-            List of SimulationTrialResults for this run, one for each
-            possible parameter resolver.
-
-        Raises:
-            TypeError: if an invalid initial_state is provided.
-        """
+    ) -> Iterator[Tuple[cirq.ParamResolver, np.ndarray, Sequence[int]]]:
         if initial_state is None:
             initial_state = 0
         if not isinstance(initial_state, (int, np.ndarray)):
@@ -583,9 +497,64 @@ def simulate_sweep_iter(
             assert qsim_state.dtype == np.float32
             assert qsim_state.ndim == 1
 
+            yield prs, qsim_state.view(np.complex64), cirq_order
+
+    def simulate_into_1d_array(
+        self,
+        program: cirq.AbstractCircuit,
+        param_resolver: cirq.ParamResolverOrSimilarType = None,
+        qubit_order: cirq.QubitOrderOrList = cirq.ops.QubitOrder.DEFAULT,
+        initial_state: Any = None,
+    ) -> np.ndarray:
+        """Same as simulate() but returns the state vector as 1D np.ndarray."""
+        params = cirq.study.ParamResolver(param_resolver)
+        _, state_vector, _ = next(
+            self._simulate_impl(program, params, qubit_order, initial_state)
+        )
+        return state_vector
+
+    def simulate_sweep_iter(
+        self,
+        program: cirq.Circuit,
+        params: cirq.Sweepable,
+        qubit_order: cirq.QubitOrderOrList = cirq.QubitOrder.DEFAULT,
+        initial_state: Optional[Union[int, np.ndarray]] = None,
+    ) -> Iterator[cirq.StateVectorTrialResult]:
+        """Simulates the supplied Circuit.
+
+        This method returns a result which allows access to the entire
+        wave function. In contrast to simulate, this allows for sweeping
+        over different parameter values.
+
+        Avoid using this method with `use_gpu=True` in the simulator options;
+        when used with GPU this method must copy state from device to host memory
+        multiple times, which can be very slow. This issue is not present in
+        `simulate_expectation_values_sweep`.
+
+        Args:
+            program: The circuit to simulate.
+            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 ordering of the
+              computational basis states.
+            initial_state: The initial state for the simulation. This can either
+              be an integer representing a pure state (e.g. 11010) or a numpy
+              array containing the full state vector. If none is provided, this
+              is assumed to be the all-zeros state.
+
+        Returns:
+            List of SimulationTrialResults for this run, one for each
+            possible parameter resolver.
+
+        Raises:
+            TypeError: if an invalid initial_state is provided.
+        """
+
+        for prs, state_vector, cirq_order in self._simulate_impl(
+            program, params, qubit_order, initial_state
+        ):
             final_state = cirq.StateVectorSimulationState(
-                initial_state=qsim_state.view(np.complex64),
-                qubits=cirq_order if not as_1d_state_vector else None,
+                initial_state=state_vector, qubits=cirq_order
             )
             # create result for this parameter
             # TODO: We need to support measurements.

qsimcirq_tests/qsimcirq_test.py

@@ -2066,7 +2066,5 @@ def test_1d_representation():
     c = cirq.Circuit(cirq.H.on_each(qs), cirq.X(qs[0]), cirq.Y(qs[1]))
 
     want = np.array([0.0 - 0.5j, 0.0 + 0.5j, 0.0 - 0.5j, 0.0 + 0.5j])
-    res = qsim_sim.simulate(c, as_1d_state_vector=True)
-    np.testing.assert_allclose(
-        res.final_state_vector, np.array(want, dtype=np.complex64)
-    )
+    res = qsim_sim.simulate_into_1d_array(c)
+    np.testing.assert_allclose(res, np.array(want, dtype=np.complex64))