Create a new method to return the final state vector array instead of wrapping it

.github/workflows/bazeltest.yml

@@ -63,7 +63,7 @@ jobs:
       - name: Upgrade libc 
       # An LLVM update broke this test, fix per is https://bugs.llvm.org/show_bug.cgi?id=27310.
         run: |
-          sudo apt-get upgrade libc-bin=2.31-0ubuntu9.9
+          sudo apt-get upgrade libc-bin
       - name: Run C++ tests
         run: |
           bazel test --config=avx --config=openmp \

qsimcirq/qsim_simulator.py

@@ -435,42 +435,13 @@ def compute_amplitudes_sweep_iter(
             options["s"] = self.get_seed()
             yield simulator_fn(options)
 
-    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,
-    ) -> 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.
-        """
+    ) -> Iterator[Tuple[cirq.ParamResolver, np.ndarray, Sequence[int]]]:
         if initial_state is None:
             initial_state = 0
         if not isinstance(initial_state, (int, np.ndarray)):
@@ -526,8 +497,68 @@ 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,
+    ) -> Tuple[cirq.ParamResolver, np.ndarray, Sequence[int]]:
+        """Same as simulate() but returns raw simulation result without wrapping it.
+
+            The returned result is not wrapped in a StateVectorTrialResult but can be used
+            to create a StateVectorTrialResult.
+
+        Returns:
+            Tuple of (param resolver, final state, qubit order)
+        """
+        params = cirq.study.ParamResolver(param_resolver)
+        return next(self._simulate_impl(program, params, qubit_order, initial_state))
+
+    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:
+            Iterator over 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
+                initial_state=state_vector, qubits=cirq_order
             )
             # create result for this parameter
             # TODO: We need to support measurements.

qsimcirq_tests/qsimcirq_test.py

@@ -2058,3 +2058,13 @@ def test_cirq_global_phase_gate():
     assert cirq.approx_eq(
         qsim_result.state_vector(), cirq_result.state_vector(), atol=1e-6
     )
+
+
+def test_1d_representation():
+    qsim_sim = qsimcirq.QSimSimulator()
+    qs = cirq.LineQubit.range(2)
+    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_into_1d_array(c)
+    np.testing.assert_allclose(res, np.array(want, dtype=np.complex64))