Fix str/repr explosion in separated states (#4518)

* Fix str/repr explosion in separated states * format * Fix str/repr explosion in state vector * format * nondirac_str * TrialResultBase * cleanup * Add get_state_containing_qubit * Ensure eval(repr(obj)) holds * substates not a property * fix test * coverage * test * fix windows Co-authored-by: Cirq Bot <[email protected]>
quantumlib · Dec 21, 2021 · 1d27a00 · 1d27a00
1 parent 4ebfb1c
commit 1d27a00
Show file tree

Hide file tree

Showing 17 changed files with 348 additions and 84 deletions.
diff --git a/cirq-core/cirq/__init__.py b/cirq-core/cirq/__init__.py
@@ -436,6 +436,7 @@
     SimulatesIntermediateStateVector,
     SimulatesSamples,
     SimulationTrialResult,
+    SimulationTrialResultBase,
     Simulator,
     SimulatorBase,
     SparseSimulatorStep,

diff --git a/cirq-core/cirq/_compat.py b/cirq-core/cirq/_compat.py
@@ -87,6 +87,9 @@ def _print(self, expr, **kwargs):
             f'\n)'
         )
 
+    if isinstance(value, Dict):
+        return '{' + ','.join(f"{proper_repr(k)}: {proper_repr(v)}" for k, v in value.items()) + '}'
+
     return repr(value)
 
 

diff --git a/cirq-core/cirq/contrib/quimb/mps_simulator.py b/cirq-core/cirq/contrib/quimb/mps_simulator.py
@@ -25,7 +25,7 @@
 import quimb.tensor as qtn
 
 from cirq import devices, study, ops, protocols, value
-from cirq.sim import simulator, simulator_base
+from cirq.sim import simulator_base
 from cirq.sim.act_on_args import ActOnArgs
 
 if TYPE_CHECKING:
@@ -146,7 +146,7 @@ def _create_simulator_trial_result(
         )
 
 
-class MPSTrialResult(simulator.SimulationTrialResult):
+class MPSTrialResult(simulator_base.SimulationTrialResultBase['MPSState', 'MPSState']):
     """A single trial reult"""
 
     def __init__(

diff --git a/cirq-core/cirq/protocols/json_test_data/spec.py b/cirq-core/cirq/protocols/json_test_data/spec.py
@@ -63,6 +63,7 @@
         'QuilFormatter',
         'QuilOutput',
         'SimulationTrialResult',
+        'SimulationTrialResultBase',
         'SparseSimulatorStep',
         'StateVectorMixin',
         'TextDiagramDrawer',

diff --git a/cirq-core/cirq/sim/__init__.py b/cirq-core/cirq/sim/__init__.py
@@ -64,8 +64,9 @@
 )
 
 from cirq.sim.simulator_base import (
-    StepResultBase,
+    SimulationTrialResultBase,
     SimulatorBase,
+    StepResultBase,
 )
 
 from cirq.sim.sparse_simulator import (

diff --git a/cirq-core/cirq/sim/act_on_density_matrix_args.py b/cirq-core/cirq/sim/act_on_density_matrix_args.py
@@ -18,6 +18,7 @@
 import numpy as np
 
 from cirq import protocols, sim
+from cirq._compat import proper_repr
 from cirq.sim.act_on_args import ActOnArgs, strat_act_on_from_apply_decompose
 from cirq.linalg import transformations
 
@@ -37,8 +38,8 @@ def __init__(
         target_tensor: np.ndarray,
         available_buffer: List[np.ndarray],
         qid_shape: Tuple[int, ...],
-        prng: np.random.RandomState,
-        log_of_measurement_results: Dict[str, Any],
+        prng: np.random.RandomState = None,
+        log_of_measurement_results: Dict[str, Any] = None,
         qubits: Sequence['cirq.Qid'] = None,
     ):
         """Inits ActOnDensityMatrixArgs.
@@ -168,6 +169,16 @@ def sample(
             seed=seed,
         )
 
+    def __repr__(self) -> str:
+        return (
+            'cirq.ActOnDensityMatrixArgs('
+            f'target_tensor={proper_repr(self.target_tensor)},'
+            f' available_buffer={proper_repr(self.available_buffer)},'
+            f' qid_shape={self.qid_shape!r},'
+            f' qubits={self.qubits!r},'
+            f' log_of_measurement_results={proper_repr(self.log_of_measurement_results)})'
+        )
+
 
 def _strat_apply_channel_to_state(
     action: Any, args: ActOnDensityMatrixArgs, qubits: Sequence['cirq.Qid']

diff --git a/cirq-core/cirq/sim/act_on_state_vector_args.py b/cirq-core/cirq/sim/act_on_state_vector_args.py
@@ -18,6 +18,7 @@
 import numpy as np
 
 from cirq import linalg, protocols, sim
+from cirq._compat import proper_repr
 from cirq.sim.act_on_args import ActOnArgs, strat_act_on_from_apply_decompose
 from cirq.linalg import transformations
 
@@ -40,8 +41,8 @@ def __init__(
         self,
         target_tensor: np.ndarray,
         available_buffer: np.ndarray,
-        prng: np.random.RandomState,
-        log_of_measurement_results: Dict[str, Any],
+        prng: np.random.RandomState = None,
+        log_of_measurement_results: Dict[str, Any] = None,
         qubits: Sequence['cirq.Qid'] = None,
     ):
         """Inits ActOnStateVectorArgs.
@@ -224,6 +225,15 @@ def sample(
             seed=seed,
         )
 
+    def __repr__(self) -> str:
+        return (
+            'cirq.ActOnStateVectorArgs('
+            f'target_tensor={proper_repr(self.target_tensor)},'
+            f' available_buffer={proper_repr(self.available_buffer)},'
+            f' qubits={self.qubits!r},'
+            f' log_of_measurement_results={proper_repr(self.log_of_measurement_results)})'
+        )
+
 
 def _strat_act_on_state_vector_from_apply_unitary(
     unitary_value: Any,

diff --git a/cirq-core/cirq/sim/clifford/clifford_simulator.py b/cirq-core/cirq/sim/clifford/clifford_simulator.py
@@ -36,7 +36,7 @@
 import cirq
 from cirq import study, protocols, value
 from cirq.protocols import act_on
-from cirq.sim import clifford, simulator, simulator_base
+from cirq.sim import clifford, simulator_base
 
 
 class CliffordSimulator(
@@ -114,7 +114,11 @@ def _create_simulator_trial_result(
         )
 
 
-class CliffordTrialResult(simulator.SimulationTrialResult):
+class CliffordTrialResult(
+    simulator_base.SimulationTrialResultBase[
+        'clifford.CliffordState', 'clifford.ActOnStabilizerCHFormArgs'
+    ]
+):
     def __init__(
         self,
         params: study.ParamResolver,

diff --git a/cirq-core/cirq/sim/density_matrix_simulator.py b/cirq-core/cirq/sim/density_matrix_simulator.py
@@ -17,6 +17,7 @@
 import numpy as np
 
 from cirq import ops, protocols, qis, study, value
+from cirq._compat import proper_repr
 from cirq.sim import (
     simulator,
     act_on_density_matrix_args,
@@ -283,7 +284,7 @@ class DensityMatrixStepResult(
     def __init__(
         self,
         sim_state: 'cirq.OperationTarget[cirq.ActOnDensityMatrixArgs]',
-        simulator: DensityMatrixSimulator,
+        simulator: DensityMatrixSimulator = None,
         dtype: 'DTypeLike' = np.complex64,
     ):
         """DensityMatrixStepResult.
@@ -315,7 +316,8 @@ def set_density_matrix(self, density_matrix_repr: Union[int, np.ndarray]):
             mixed state it must be correctly sized and positive semidefinite
             with trace one.
         """
-        self._sim_state = self._simulator._create_act_on_args(density_matrix_repr, self._qubits)
+        if self._simulator:
+            self._sim_state = self._simulator._create_act_on_args(density_matrix_repr, self._qubits)
 
     def density_matrix(self, copy=True):
         """Returns the density matrix at this step in the simulation.
@@ -361,6 +363,12 @@ def density_matrix(self, copy=True):
                 self._density_matrix = np.reshape(matrix, (size, size))
         return self._density_matrix.copy() if copy else self._density_matrix
 
+    def __repr__(self) -> str:
+        return (
+            f'cirq.DensityMatrixStepResult(sim_state={self._sim_state!r},'
+            f' dtype=np.{self._dtype.__name__})'
+        )
+
 
 @value.value_equality(unhashable=True)
 class DensityMatrixSimulatorState:
@@ -381,7 +389,7 @@ def _qid_shape_(self) -> Tuple[int, ...]:
         return self._qid_shape
 
     def _value_equality_values_(self) -> Any:
-        return (self.density_matrix.tolist(), self.qubit_map)
+        return self.density_matrix.tolist(), self.qubit_map
 
     def __repr__(self) -> str:
         return (
@@ -392,7 +400,11 @@ def __repr__(self) -> str:
 
 
 @value.value_equality(unhashable=True)
-class DensityMatrixTrialResult(simulator.SimulationTrialResult):
+class DensityMatrixTrialResult(
+    simulator_base.SimulationTrialResultBase[
+        'DensityMatrixSimulatorState', act_on_density_matrix_args.ActOnDensityMatrixArgs
+    ]
+):
     """A `SimulationTrialResult` for `DensityMatrixSimulator` runs.
 
     The density matrix that is stored in this result is returned in the
@@ -452,17 +464,24 @@ def final_density_matrix(self):
 
     def _value_equality_values_(self) -> Any:
         measurements = {k: v.tolist() for k, v in sorted(self.measurements.items())}
-        return (self.params, measurements, self._final_simulator_state)
+        return self.params, measurements, self._final_simulator_state
 
     def __str__(self) -> str:
         samples = super().__str__()
-        return f'measurements: {samples}\nfinal density matrix:\n{self.final_density_matrix}'
+        ret = f'measurements: {samples}'
+        for substate in self._get_substates():
+            tensor = substate.target_tensor
+            size = np.prod([tensor.shape[i] for i in range(tensor.ndim // 2)], dtype=np.int64)
+            dm = tensor.reshape((size, size))
+            label = f'qubits: {substate.qubits}' if substate.qubits else 'phase:'
+            ret += f'\n\n{label}\nfinal density matrix:\n{dm}'
+        return ret
 
     def __repr__(self) -> str:
         return (
             'cirq.DensityMatrixTrialResult('
-            f'params={self.params!r}, measurements={self.measurements!r}, '
-            f'final_simulator_state={self._final_simulator_state!r})'
+            f'params={self.params!r}, measurements={proper_repr(self.measurements)}, '
+            f'final_step_result={self._final_step_result!r})'
         )
 
     def _repr_pretty_(self, p: Any, cycle: bool):
-Original file line number
+Diff line change
@@ Expand Up / @@ -87,6 +87,9 @@ def _print(self, expr, **kwargs): @@
                 f'\n)'
             )
+        if isinstance(value, Dict):
+            return '{' + ','.join(f"{proper_repr(k)}: {proper_repr(v)}" for k, v in value.items()) + '}'
         return repr(value)
@@ Expand Down @@