SimulatorBase independent qubits optimization (#4100)

Add optimization that ensures independent qubit sets are simulated independently. This is done by adding join, extract, and reorder methods to ActOnArgs, and updating SimulatorBase with the logic to merge qubit sets when necessary and split them when possible.

This optimization is enabled or disabled via a new parameter in the simulator constructors: `split_entangled_qubits`. Currently the PR has this set to True by default, though perhaps it should be disabled by default lest it breaks anything? The MPS simulator does not yet have `extract` defined and thus there's no option to enable this feature in MPS simulator's constructor yet, though nothing prevents this from being added later.

The perf boost of this implementation is limited because each StepResult still requires the full product state. It's still a speedup because full product state calculations will only have to occur once per moment rather than once per operation, but not as nice as avoiding full product state calculations entirely. *That* optimization will be available in a subsequent PR that never creates the full product state if possible: StepResults will join the product state only on demand, and sampling will sample each substate independently and zip up the results, avoiding the full state join: The WIP is here https://github.com/daxfohl/Cirq/compare/split...daxfohl:sample?expand=1.

I ramped up the number of qubits in the benchmarks to 25 for sparse and 12 for DM:

From master:
```
(cirq-py3) dax@DESKTOP-Q5MLJ3J:~/cirq$ time pytest dev_tools/profiling/benchmark_simulators_test.py
platform linux -- Python 3.8.5, pytest-5.4.3, py-1.10.0, pluggy-0.13.1
benchmark: 3.2.3 (defaults: timer=time.perf_counter disable_gc=False min_rounds=5 min_time=0.000005 max_time=1.0 calibration_precision=10 warmup=False warmup_iterations=100000)
rootdir: /home/dax/cirq
plugins: cov-2.5.1, asyncio-0.12.0, benchmark-3.2.3
collected 5 items

dev_tools/profiling/benchmark_simulators_test.py .....                                                                                                                                                    [100%]

real    0m16.973s
user    0m15.754s
sys     0m3.862s
```

From split branch (the current PR):
```
(cirq-py3) dax@DESKTOP-Q5MLJ3J:~/cirq$ time pytest dev_tools/profiling/benchmark_simulators_test.py
platform linux -- Python 3.8.5, pytest-5.4.3, py-1.10.0, pluggy-0.13.1
benchmark: 3.2.3 (defaults: timer=time.perf_counter disable_gc=False min_rounds=5 min_time=0.000005 max_time=1.0 calibration_precision=10 warmup=False warmup_iterations=100000)
rootdir: /home/dax/cirq
plugins: cov-2.5.1, asyncio-0.12.0, benchmark-3.2.3
collected 5 items

dev_tools/profiling/benchmark_simulators_test.py .....                                                                                                                                                    [100%]

real    0m10.073s
user    0m9.082s
sys     0m3.805s
```

From sample branch (future iteration mentioned above):
```
(cirq-py3) dax@DESKTOP-Q5MLJ3J:~/cirq$ time pytest dev_tools/profiling/benchmark_simulators_test.py
platform linux -- Python 3.8.5, pytest-5.4.3, py-1.10.0, pluggy-0.13.1
benchmark: 3.2.3 (defaults: timer=time.perf_counter disable_gc=False min_rounds=5 min_time=0.000005 max_time=1.0 calibration_precision=10 warmup=False warmup_iterations=100000)
rootdir: /home/dax/cirq
plugins: cov-2.5.1, asyncio-0.12.0, benchmark-3.2.3
collected 5 items

dev_tools/profiling/benchmark_simulators_test.py .....                                                                                                                                                    [100%]

real    0m2.885s
user    0m3.523s
sys     0m2.597s
```

Initial PR for #3240
Closes #882

Author: daxfohl

cirq-core/cirq/__init__.py

@@ -358,6 +358,7 @@
 
 from cirq.sim import (
     ActOnArgs,
+    ActOnArgsContainer,
     ActOnCliffordTableauArgs,
     ActOnDensityMatrixArgs,
     ActOnStabilizerCHFormArgs,
@@ -376,6 +377,7 @@
     measure_state_vector,
     final_density_matrix,
     final_state_vector,
+    OperationTarget,
     sample,
     sample_density_matrix,
     sample_state_vector,

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

@@ -85,10 +85,11 @@ def __init__(
             seed=seed,
         )
 
-    def _create_act_on_args(
+    def _create_partial_act_on_args(
         self,
         initial_state: Union[int, 'MPSState'],
         qubits: Sequence['cirq.Qid'],
+        logs: Dict[str, Any],
     ) -> 'MPSState':
         """Creates MPSState args for simulating the Circuit.
 
@@ -111,6 +112,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(
@@ -317,6 +319,7 @@ def copy(self) -> 'MPSState':
             prng=self.prng,
             simulation_options=self.simulation_options,
             grouping=self.grouping,
+            log_of_measurement_results=self.log_of_measurement_results.copy(),
         )
         state.M = [x.copy() for x in self.M]
         state.estimated_gate_error_list = self.estimated_gate_error_list

cirq-core/cirq/linalg/transformations.py

@@ -500,3 +500,160 @@ def to_special(u: np.ndarray) -> np.ndarray:
         the special unitary matrix
     """
     return u * (np.linalg.det(u) ** (-1 / len(u)))
+
+
+def state_vector_kronecker_product(
+    t1: np.ndarray,
+    t2: np.ndarray,
+) -> np.ndarray:
+    """Merges two state vectors into a single unified state vector.
+
+    The resulting vector's shape will be `t1.shape + t2.shape`.
+
+    Args:
+        t1: The first state vector.
+        t2: The second state vector.
+    Returns:
+        A new state vector representing the unified state.
+    """
+    return np.outer(t1, t2).reshape(t1.shape + t2.shape)
+
+
+def density_matrix_kronecker_product(
+    t1: np.ndarray,
+    t2: np.ndarray,
+) -> np.ndarray:
+    """Merges two density matrices into a single unified density matrix.
+
+    The resulting matrix's shape will be `(t1.shape/2 + t2.shape/2) * 2`. In
+    other words, if t1 has shape [A,B,C,A,B,C] and t2 has shape [X,Y,Z,X,Y,Z],
+    the resulting matrix will have shape [A,B,C,X,Y,Z,A,B,C,X,Y,Z].
+
+    Args:
+        t1: The first density matrix.
+        t2: The second density matrix.
+    Returns:
+        A density matrix representing the unified state.
+    """
+    t = state_vector_kronecker_product(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 factor_state_vector(
+    t: np.ndarray,
+    axes: Sequence[int],
+    *,
+    validate=True,
+    atol=1e-07,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Factors a state vector into two independent state vectors.
+
+    This function should only be called on state vectors that are known to be
+    separable, such as immediately after a measurement or reset operation. It
+    does not verify that the provided state vector is indeed separable, and
+    will return nonsense results for vectors representing entangled states.
+
+    Args:
+        t: The state vector to factor.
+        axes: The axes to factor out.
+        validate: Perform a validation that the density matrix factors cleanly.
+        atol: The absolute tolerance for the validation.
+    Returns:
+        A tuple with the `(extracted, remainder)` state vectors, where
+        `extracted` means the sub-state vector which corresponds to the axes
+        requested, and with the axes in the requested order, and where
+        `remainder` means the sub-state vector on the remaining axes, in the
+        same order as the original state vector.
+    """
+    n_axes = len(axes)
+    t1 = np.moveaxis(t, axes, range(n_axes))
+    pivot = np.unravel_index(np.abs(t1).argmax(), t1.shape)
+    slices1 = (slice(None),) * n_axes + pivot[n_axes:]
+    slices2 = pivot[:n_axes] + (slice(None),) * (t1.ndim - n_axes)
+    extracted = t1[slices1]
+    extracted = extracted / np.sum(abs(extracted) ** 2) ** 0.5
+    remainder = t1[slices2]
+    remainder = remainder / np.sum(abs(remainder) ** 2) ** 0.5
+    if validate:
+        t2 = state_vector_kronecker_product(extracted, remainder)
+        axes2 = list(axes) + [i for i in range(t1.ndim) if i not in axes]
+        t3 = transpose_state_vector_to_axis_order(t2, axes2)
+        if not np.allclose(t3, t, atol=atol):
+            raise ValueError('The tensor cannot be factored by the requested axes')
+    return extracted, remainder
+
+
+def factor_density_matrix(
+    t: np.ndarray,
+    axes: Sequence[int],
+    *,
+    validate=True,
+    atol=1e-07,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Factors a density matrix into two independent density matrices.
+
+    This function should only be called on density matrices that are known to
+    be separable, such as immediately after a measurement or reset operation.
+    It does not verify that the provided density matrix is indeed separable,
+    and will return nonsense results for matrices representing entangled
+    states.
+
+    Args:
+        t: The density matrix to factor.
+        axes: The axes to factor out. Only the left axes should be provided.
+            For example, to extract [C,A] from density matrix of shape
+            [A,B,C,D,A,B,C,D], `axes` should be [2,0], and the return value
+            will be two density matrices ([C,A,C,A], [B,D,B,D]).
+        validate: Perform a validation that the density matrix factors cleanly.
+        atol: The absolute tolerance for the validation.
+    Returns:
+        A tuple with the `(extracted, remainder)` density matrices, where
+        `extracted` means the sub-matrix which corresponds to the axes
+        requested, and with the axes in the requested order, and where
+        `remainder` means the sub-matrix on the remaining axes, in the same
+        order as the original density matrix.
+    """
+    axes1 = list(axes) + [i + int(t.ndim / 2) for i in axes]
+    extracted, remainder = factor_state_vector(t, axes1, validate=False)
+    if validate:
+        t1 = density_matrix_kronecker_product(extracted, remainder)
+        axes2 = list(axes) + [i for i in range(int(t.ndim / 2)) if i not in axes]
+        t2 = transpose_density_matrix_to_axis_order(t1, axes2)
+        if not np.allclose(t2, t, atol=atol):
+            raise ValueError('The tensor cannot be factored by the requested axes')
+    return extracted, remainder
+
+
+def transpose_state_vector_to_axis_order(t: np.ndarray, axes: Sequence[int]):
+    """Transposes the axes of a state vector to a specified order.
+
+    Args:
+        t: The state vector to transpose.
+        axes: The desired axis order.
+    Returns:
+        The transposed state vector.
+    """
+    assert set(axes) == set(range(int(t.ndim))), "All axes must be provided."
+    return np.moveaxis(t, axes, range(len(axes)))
+
+
+def transpose_density_matrix_to_axis_order(t: np.ndarray, axes: Sequence[int]):
+    """Transposes the axes of a density matrix to a specified order.
+
+    Args:
+        t: The density matrix to transpose.
+        axes: The desired axis order. Only the left axes should be provided.
+            For example, to transpose [A,B,C,A,B,C] to [C,B,A,C,B,A], `axes`
+            should be [2,1,0].
+    Returns:
+        The transposed density matrix.
+    """
+    axes = list(axes) + [i + len(axes) for i in axes]
+    return transpose_state_vector_to_axis_order(t, axes)

cirq-core/cirq/protocols/act_on_protocol_test.py

@@ -82,7 +82,7 @@ def _act_on_fallback_(self, action, qubits, allow_decompose):
             return True
 
     args = Args()
-    args.qubits = tuple(cirq.LineQubit.range(3))
+    args._qubits = tuple(cirq.LineQubit.range(3))
     with cirq.testing.assert_deprecated(
         "ActOnArgs.axes", "Use `protocols.act_on` instead.", deadline="v0.13"
     ):

cirq-core/cirq/protocols/json_test_data/spec.py

@@ -83,13 +83,15 @@
         'TwoQubitInteractionHeatmap',
         # Intermediate states with work buffers and unknown external prng guts.
         'ActOnArgs',
+        'ActOnArgsContainer',
         'ActOnCliffordTableauArgs',
         'ActOnDensityMatrixArgs',
         'ActOnStabilizerCHFormArgs',
         'ActOnStateVectorArgs',
         'ApplyChannelArgs',
         'ApplyMixtureArgs',
         'ApplyUnitaryArgs',
+        'OperationTarget',
         # Circuit optimizers are function-like. Only attributes
         # are ignore_failures, tolerance, and other feature flags
         'AlignLeft',

cirq-core/cirq/sim/__init__.py

@@ -19,6 +19,10 @@
     ActOnArgs,
 )
 
+from cirq.sim.act_on_args_container import (
+    ActOnArgsContainer,
+)
+
 from cirq.sim.act_on_density_matrix_args import (
     ActOnDensityMatrixArgs,
 )
@@ -39,6 +43,8 @@
     DensityMatrixTrialResult,
 )
 
+from cirq.sim.operation_target import OperationTarget
+
 from cirq.sim.mux import (
     CIRCUIT_LIKE,
     final_density_matrix,

cirq-core/cirq/sim/act_on_args.py

@@ -13,26 +13,39 @@
 # limitations under the License.
 """Objects and methods for acting efficiently on a state tensor."""
 import abc
-from typing import Any, Dict, List, TypeVar, TYPE_CHECKING, Sequence, Tuple, Iterable
+from typing import (
+    Any,
+    Iterable,
+    Dict,
+    List,
+    TypeVar,
+    TYPE_CHECKING,
+    Sequence,
+    Tuple,
+    cast,
+    Optional,
+    Iterator,
+)
 
 import numpy as np
 
 from cirq import protocols
 from cirq._compat import deprecated
 from cirq.protocols.decompose_protocol import _try_decompose_into_operations_and_qubits
+from cirq.sim.operation_target import OperationTarget
 
 TSelf = TypeVar('TSelf', bound='ActOnArgs')
 
 if TYPE_CHECKING:
     import cirq
 
 
-class ActOnArgs:
+class ActOnArgs(OperationTarget[TSelf]):
     """State and context for an operation acting on a state tensor."""
 
     def __init__(
         self,
-        prng: np.random.RandomState,
+        prng: np.random.RandomState = None,
         qubits: Sequence['cirq.Qid'] = None,
         axes: Iterable[int] = None,
         log_of_measurement_results: Dict[str, Any] = None,
@@ -50,17 +63,19 @@ def __init__(
                 being recorded into. Edit it easily by calling
                 `ActOnStateVectorArgs.record_measurement_result`.
         """
+        if prng is None:
+            prng = cast(np.random.RandomState, np.random)
         if qubits is None:
             qubits = ()
         if axes is None:
             axes = ()
         if log_of_measurement_results is None:
             log_of_measurement_results = {}
-        self.qubits = tuple(qubits)
+        self._qubits = tuple(qubits)
         self.qubit_map = {q: i for i, q in enumerate(self.qubits)}
         self._axes = tuple(axes)
         self.prng = prng
-        self.log_of_measurement_results = log_of_measurement_results
+        self._log_of_measurement_results = log_of_measurement_results
 
     def measure(self, qubits: Sequence['cirq.Qid'], key: str, invert_mask: Sequence[bool]):
         """Adds a measurement result to the log.
@@ -90,6 +105,51 @@ def _perform_measurement(self, qubits: Sequence['cirq.Qid']) -> List[int]:
     def copy(self: TSelf) -> TSelf:
         """Creates a copy of the object."""
 
+    def create_merged_state(self: TSelf) -> TSelf:
+        """Creates a final merged state."""
+        return self
+
+    def apply_operation(self, op: 'cirq.Operation'):
+        """Applies the operation to the state."""
+        protocols.act_on(op, self)
+
+    def kronecker_product(self: TSelf, other: TSelf) -> TSelf:
+        """Joins two state spaces together."""
+        raise NotImplementedError()
+
+    def factor(
+        self: TSelf,
+        qubits: Sequence['cirq.Qid'],
+        *,
+        validate=True,
+        atol=1e-07,
+    ) -> Tuple[TSelf, TSelf]:
+        """Splits two state spaces after a measurement or reset."""
+        raise NotImplementedError()
+
+    def transpose_to_qubit_order(self: TSelf, qubits: Sequence['cirq.Qid']) -> TSelf:
+        """Physically reindexes the state by the new basis."""
+        raise NotImplementedError()
+
+    @property
+    def log_of_measurement_results(self) -> Dict[str, Any]:
+        return self._log_of_measurement_results
+
+    @property
+    def qubits(self) -> Tuple['cirq.Qid', ...]:
+        return self._qubits
+
+    def __getitem__(self: TSelf, item: Optional['cirq.Qid']) -> TSelf:
+        if item not in self.qubit_map:
+            raise IndexError(f'{item} not in {self.qubits}')
+        return self
+
+    def __len__(self) -> int:
+        return len(self.qubits)
+
+    def __iter__(self) -> Iterator[Optional['cirq.Qid']]:
+        return iter(self.qubits)
+
     @abc.abstractmethod
     def _act_on_fallback_(self, action: Any, qubits: Sequence['cirq.Qid'], allow_decompose: bool):
         """Handles the act_on protocol fallback implementation."""