SimulatorBase independent qubits optimization

cirq-core/cirq/contrib/quimb/mps_simulator.py

@@ -19,7 +19,7 @@
 
 import dataclasses
 import math
-from typing import Any, Dict, List, Optional, Sequence, Set, TYPE_CHECKING, Iterable, Union
+from typing import Any, Dict, List, Optional, Sequence, Set, TYPE_CHECKING, Iterable, Union, Tuple
 
 import numpy as np
 import quimb.tensor as qtn
@@ -84,10 +84,11 @@ def __init__(
             seed=seed,
         )
 
-    def _create_act_on_args(
+    def _create_act_on_arg(
         self,
         initial_state: Union[int, 'MPSState'],
         qubits: Sequence['cirq.Qid'],
+        logs: Dict[str, Any],
     ) -> 'MPSState':
         """Creates MPSState args for simulating the Circuit.
 
@@ -110,6 +111,7 @@ def _create_act_on_args(
             simulation_options=self.simulation_options,
             grouping=self.grouping,
             initial_state=initial_state,
+            log_of_measurement_results=logs,
         )
 
     def _create_step_result(
@@ -220,7 +222,6 @@ def sample(
         return np.array(measurements, dtype=int)
 
 
-@value.value_equality
 class MPSState(ActOnArgs):
     """A state of the MPS simulation."""
 
@@ -310,11 +311,33 @@ def copy(self) -> 'MPSState':
             prng=self.prng,
             simulation_options=self.simulation_options,
             grouping=self.grouping,
+            log_of_measurement_results=self.log_of_measurement_results,
         )
         state.M = [x.copy() for x in self.M]
         state.estimated_gate_error_list = self.estimated_gate_error_list
         return state
 
+    def join(self, other: 'MPSState') -> 'MPSState':
+        # TODO MPS simulator currently does not enable split_untangled_states
+        # so this will never be called during simulation, and MPS gains nothing
+        # from running in split_untangled_states mode, so this is not necessary,
+        # however it may be useful if other use cases arise.
+        raise NotImplementedError()
+
+    def extract(self, qubits: Sequence['cirq.Qid']) -> Tuple['MPSState', 'MPSState']:
+        # TODO MPS simulator currently does not enable split_untangled_states
+        # so this will never be called during simulation, and MPS gains nothing
+        # from running in split_untangled_states mode, so this is not necessary,
+        # however it may be useful if other use cases arise.
+        raise NotImplementedError()
+
+    def reorder(self, qubits: Sequence['cirq.Qid']) -> 'MPSState':
+        # TODO MPS simulator currently does not enable split_untangled_states
+        # so this will never be called during simulation, and MPS gains nothing
+        # from running in split_untangled_states mode, so this is not necessary,
+        # however may be useful if other use cases arise.
+        raise NotImplementedError()
+
     def state_vector(self) -> np.ndarray:
         """Returns the full state vector.
 

cirq-core/cirq/linalg/transformations.py

@@ -500,3 +500,49 @@ def to_special(u: np.ndarray) -> np.ndarray:
         the special unitary matrix
     """
     return u * (np.linalg.det(u) ** (-1 / len(u)))
+
+
+def merge_state_vectors(
+    t1: np.ndarray,
+    t2: np.ndarray,
+) -> np.ndarray:
+    return np.outer(t1, t2).reshape(t1.shape + t2.shape)
+
+
+def merge_density_matrices(
+    t1: np.ndarray,
+    t2: np.ndarray,
+) -> np.ndarray:
+    t = merge_state_vectors(t1, t2)
+    t1_len = len(t1.shape)
+    t1_dim = int(t1_len / 2)
+    t2_len = len(t2.shape)
+    t2_dim = int(t2_len / 2)
+    shape = t1.shape[:t1_dim] + t2.shape[:t2_dim]
+    return np.moveaxis(t, range(t1_len, t1_len + t2_dim), range(t1_dim, t1_dim + t2_dim)).reshape(
+        shape * 2
+    )
+
+
+def split_state_vectors(
+    t: np.ndarray,
+    axes: Sequence[int],
+) -> Tuple[np.ndarray, np.ndarray]:
+    n_axes = len(axes)
+    t = np.moveaxis(t, axes, range(n_axes))
+    pivot = np.unravel_index(np.abs(t).argmax(), t.shape)
+    slices1 = (slice(None),) * n_axes + pivot[n_axes:]
+    slices2 = pivot[:n_axes] + (slice(None),) * (t.ndim - n_axes)
+    extracted = t[slices1]
+    extracted = extracted / np.sum(abs(extracted) ** 2) ** 0.5
+    remainder = t[slices2]
+    remainder = remainder / np.sum(abs(remainder) ** 2) ** 0.5
+    return extracted, remainder
+
+
+def split_density_matrices(
+    t: np.ndarray,
+    axes: Sequence[int],
+) -> Tuple[np.ndarray, np.ndarray]:
+    axes = list(axes) + [i + int(t.ndim / 2) for i in axes]
+    return split_state_vectors(t, axes)

cirq-core/cirq/sim/act_on_args.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 """Objects and methods for acting efficiently on a state tensor."""
 import abc
-from typing import Any, Iterable, Dict, List, TypeVar, TYPE_CHECKING, Sequence
+from typing import Any, Iterable, Dict, List, TypeVar, TYPE_CHECKING, Sequence, Tuple
 
 import numpy as np
 
@@ -85,6 +85,18 @@ def _perform_measurement(self) -> List[int]:
     def copy(self: TSelf) -> TSelf:
         """Creates a copy of the object."""
 
+    @abc.abstractmethod
+    def join(self: TSelf, other: TSelf) -> TSelf:
+        """Joins two state spaces together."""
+
+    @abc.abstractmethod
+    def extract(self: TSelf, qubits: Sequence['cirq.Qid']) -> Tuple[TSelf, TSelf]:
+        """Splits two state spaces after a measurement or reset."""
+
+    @abc.abstractmethod
+    def reorder(self: TSelf, qubits: Sequence['cirq.Qid']) -> TSelf:
+        """Physically reindexes the state by the new basis."""
+
 
 def strat_act_on_from_apply_decompose(
     val: Any,

cirq-core/cirq/sim/act_on_density_matrix_args.py

@@ -19,6 +19,7 @@
 
 from cirq import protocols, sim
 from cirq.sim.act_on_args import ActOnArgs, strat_act_on_from_apply_decompose
+from cirq.linalg import transformations as tf
 
 if TYPE_CHECKING:
     import cirq
@@ -107,7 +108,63 @@ def copy(self) -> 'cirq.ActOnDensityMatrixArgs':
             axes=self.axes,
             qid_shape=self.qid_shape,
             prng=self.prng,
-            log_of_measurement_results=self.log_of_measurement_results.copy(),
+            log_of_measurement_results=self.log_of_measurement_results,
+        )
+
+    def join(self, other: 'cirq.ActOnDensityMatrixArgs') -> 'cirq.ActOnDensityMatrixArgs':
+        target_tensor = tf.merge_density_matrices(self.target_tensor, other.target_tensor)
+        buffer = [np.empty_like(target_tensor) for _ in self.available_buffer]
+        return ActOnDensityMatrixArgs(
+            target_tensor=target_tensor,
+            available_buffer=buffer,
+            qubits=self.qubits + other.qubits,
+            axes=(),
+            qid_shape=target_tensor.shape[: int(target_tensor.ndim / 2)],
+            prng=self.prng,
+            log_of_measurement_results=self.log_of_measurement_results,
+        )
+
+    def extract(
+        self, qubits: Sequence['cirq.Qid']
+    ) -> Tuple['cirq.ActOnDensityMatrixArgs', 'cirq.ActOnDensityMatrixArgs']:
+        axes = [self.qubit_map[q] for q in qubits]
+        extracted_tensor, remainder_tensor = tf.split_density_matrices(self.target_tensor, axes)
+        buffer = [np.empty_like(extracted_tensor) for _ in self.available_buffer]
+        extracted_args = ActOnDensityMatrixArgs(
+            target_tensor=extracted_tensor,
+            available_buffer=buffer,
+            qubits=qubits,
+            qid_shape=extracted_tensor.shape[: int(extracted_tensor.ndim / 2)],
+            axes=(),
+            prng=self.prng,
+            log_of_measurement_results=self.log_of_measurement_results,
+        )
+        buffer = [np.empty_like(remainder_tensor) for _ in self.available_buffer]
+        remainder_args = ActOnDensityMatrixArgs(
+            target_tensor=remainder_tensor,
+            available_buffer=buffer,
+            qubits=tuple(q for q in self.qubits if q not in qubits),
+            qid_shape=remainder_tensor.shape[: int(remainder_tensor.ndim / 2)],
+            axes=(),
+            prng=self.prng,
+            log_of_measurement_results=self.log_of_measurement_results,
+        )
+        return extracted_args, remainder_args
+
+    def reorder(self, qubits: Sequence['cirq.Qid']) -> 'cirq.ActOnDensityMatrixArgs':
+        assert len(qubits) == len(self.qubits)
+        axes = [self.qubit_map[q] for q in qubits]
+        axes = axes + [i + len(qubits) for i in axes]
+        new_tensor = np.moveaxis(self.target_tensor, axes, range(len(qubits) * 2))
+        buffer = [np.empty_like(new_tensor) for _ in self.available_buffer]
+        return ActOnDensityMatrixArgs(
+            target_tensor=new_tensor,
+            available_buffer=buffer,
+            qubits=qubits,
+            qid_shape=new_tensor.shape[: int(new_tensor.ndim / 2)],
+            axes=(),
+            prng=self.prng,
+            log_of_measurement_results=self.log_of_measurement_results,
         )
 
 

cirq-core/cirq/sim/act_on_state_vector_args.py

@@ -19,6 +19,7 @@
 
 from cirq import linalg, protocols, sim
 from cirq.sim.act_on_args import ActOnArgs, strat_act_on_from_apply_decompose
+from cirq.linalg import transformations as tf
 
 if TYPE_CHECKING:
     import cirq
@@ -178,9 +179,60 @@ def copy(self) -> 'cirq.ActOnStateVectorArgs':
             qubits=self.qubits,
             axes=self.axes,
             prng=self.prng,
-            log_of_measurement_results=self.log_of_measurement_results.copy(),
+            log_of_measurement_results=self.log_of_measurement_results,
         )
 
+    def join(self, other: 'cirq.ActOnStateVectorArgs') -> 'cirq.ActOnStateVectorArgs':
+        target_tensor = tf.merge_state_vectors(self.target_tensor, other.target_tensor)
+        buffer = np.empty_like(target_tensor)
+        offset = len(self.target_tensor.shape)
+        axes = self.axes + tuple(a + offset for a in other.axes)
+        return ActOnStateVectorArgs(
+            target_tensor=target_tensor,
+            available_buffer=buffer,
+            qubits=self.qubits + other.qubits,
+            axes=axes,
+            prng=self.prng,
+            log_of_measurement_results=self.log_of_measurement_results,
+        )
+
+    def extract(
+        self, qubits: Sequence['cirq.Qid']
+    ) -> Tuple['cirq.ActOnStateVectorArgs', 'cirq.ActOnStateVectorArgs']:
+        axes = [self.qubit_map[q] for q in qubits]
+        extracted_tensor, remainder_tensor = tf.split_state_vectors(self.target_tensor, axes)
+        extracted_args = ActOnStateVectorArgs(
+            target_tensor=extracted_tensor,
+            available_buffer=np.empty_like(extracted_tensor),
+            qubits=qubits,
+            axes=(),
+            prng=self.prng,
+            log_of_measurement_results=self.log_of_measurement_results,
+        )
+        remainder_args = ActOnStateVectorArgs(
+            target_tensor=remainder_tensor,
+            available_buffer=np.empty_like(remainder_tensor),
+            qubits=tuple(q for q in self.qubits if q not in qubits),
+            axes=(),
+            prng=self.prng,
+            log_of_measurement_results=self.log_of_measurement_results,
+        )
+        return extracted_args, remainder_args
+
+    def reorder(self, qubits: Sequence['cirq.Qid']) -> 'cirq.ActOnStateVectorArgs':
+        assert len(qubits) == len(self.qubits)
+        axes = [self.qubit_map[q] for q in qubits]
+        new_tensor = np.moveaxis(self.target_tensor, axes, range(len(qubits)))
+        new_args = ActOnStateVectorArgs(
+            target_tensor=new_tensor,
+            available_buffer=np.empty_like(new_tensor),
+            qubits=qubits,
+            axes=(),
+            prng=self.prng,
+            log_of_measurement_results=self.log_of_measurement_results,
+        )
+        return new_args
+
 
 def _strat_act_on_state_vector_from_apply_unitary(
     unitary_value: Any,

cirq-core/cirq/sim/clifford/act_on_clifford_tableau_args.py

@@ -14,7 +14,7 @@
 """A protocol for implementing high performance clifford tableau evolutions
  for Clifford Simulator."""
 
-from typing import Any, Dict, Iterable, TYPE_CHECKING, List, Sequence
+from typing import Any, Dict, Iterable, TYPE_CHECKING, List, Sequence, Tuple
 
 import numpy as np
 
@@ -87,9 +87,23 @@ def copy(self) -> 'cirq.ActOnCliffordTableauArgs':
             qubits=self.qubits,
             axes=self.axes,
             prng=self.prng,
-            log_of_measurement_results=self.log_of_measurement_results.copy(),
+            log_of_measurement_results=self.log_of_measurement_results,
         )
 
+    def join(self, other: 'ActOnCliffordTableauArgs') -> 'ActOnCliffordTableauArgs':
+        # Unnecessary for now but can be added later if there is a use case.
+        raise NotImplementedError()
+
+    def extract(
+        self, qubits: Sequence['cirq.Qid']
+    ) -> Tuple['ActOnCliffordTableauArgs', 'ActOnCliffordTableauArgs']:
+        # Unnecessary for now but can be added later if there is a use case.
+        raise NotImplementedError()
+
+    def reorder(self, qubits: Sequence['cirq.Qid']) -> 'ActOnCliffordTableauArgs':
+        # Unnecessary for now but can be added later if there is a use case.
+        raise NotImplementedError()
+
 
 def _strat_act_on_clifford_tableau_from_single_qubit_decompose(
     val: Any, args: 'cirq.ActOnCliffordTableauArgs'

cirq-core/cirq/sim/clifford/act_on_clifford_tableau_args_test.py

@@ -98,5 +98,4 @@ def test_copy():
     assert args.tableau == args1.tableau
     assert args.axes == args1.axes
     assert args.prng is args1.prng
-    assert args.log_of_measurement_results is not args1.log_of_measurement_results
-    assert args.log_of_measurement_results == args.log_of_measurement_results
+    assert args.log_of_measurement_results is args1.log_of_measurement_results

cirq-core/cirq/sim/clifford/act_on_stabilizer_ch_form_args.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, Iterable, TYPE_CHECKING, List, Sequence
+from typing import Any, Dict, Iterable, TYPE_CHECKING, List, Sequence, Tuple
 
 import numpy as np
 
@@ -82,9 +82,23 @@ def copy(self) -> 'cirq.ActOnStabilizerCHFormArgs':
             qubits=self.qubits,
             axes=self.axes,
             prng=self.prng,
-            log_of_measurement_results=self.log_of_measurement_results.copy(),
+            log_of_measurement_results=self.log_of_measurement_results,
         )
 
+    def join(self, other: 'cirq.ActOnStabilizerCHFormArgs') -> 'cirq.ActOnStabilizerCHFormArgs':
+        # Unnecessary for now but can be added later if there is a use case.
+        raise NotImplementedError()
+
+    def extract(
+        self, qubits: Sequence['cirq.Qid']
+    ) -> Tuple['cirq.ActOnStabilizerCHFormArgs', 'cirq.ActOnStabilizerCHFormArgs']:
+        # Unnecessary for now but can be added later if there is a use case.
+        raise NotImplementedError()
+
+    def reorder(self, qubits: Sequence['cirq.Qid']) -> 'cirq.ActOnStabilizerCHFormArgs':
+        # Unnecessary for now but can be added later if there is a use case.
+        raise NotImplementedError()
+
 
 def _strat_act_on_stabilizer_ch_form_from_single_qubit_decompose(
     val: Any, args: 'cirq.ActOnStabilizerCHFormArgs'

cirq-core/cirq/sim/clifford/act_on_stabilizer_ch_form_args_test.py

@@ -122,5 +122,4 @@ def test_copy():
     np.testing.assert_equal(args.state.state_vector(), args1.state.state_vector())
     assert args.axes == args1.axes
     assert args.prng is args1.prng
-    assert args.log_of_measurement_results is not args1.log_of_measurement_results
-    assert args.log_of_measurement_results == args.log_of_measurement_results
+    assert args.log_of_measurement_results is args1.log_of_measurement_results