Repository: quantumlib/Cirq

Issue:
Density matrix simulator doesn't accept array with shape=(2,)*2**len(qubits) as an initial state
See https://quantumcomputing.stackexchange.com/questions/16636/


Discussion:
Seems pretty self-explanatory: `cirq.qis.to_valid_density_matrix` assumes inputs are either rank-2 (and therefore density matrices) or rank-1 (and therefore state vectors, which must be converted to density matrices).
https://github.com/quantumlib/Cirq/blob/c27dcffcab6a62540bd939fb05055fabd474a1a1/cirq/qis/states.py#L788-L795

Since either `num_qubits` or `qid_shape` is required, it should always be possible to tell whether `density_matrix_rep` is a state vector or density matrix, but we'll need to break the existing assumptions about input tensor rank.
Also doesn't allow specifying initial state as a state vector, e.g. a vector of shape `(2**len(qubits),)`:
```python
>>> sim = cirq.DensityMatrixSimulator()
>>> circuit = cirq.Circuit(cirq.measure(cirq.LineQubit(0)))
>>> sim.simulate(circuit, initial_state=np.array([1, 0]))
...
ValueError: cannot reshape array of size 2 into shape (2,2)
```

<details>

```python
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [20], in <cell line: 1>()
----> 1 sim.simulate(circuit, initial_state=np.array([1, 0]))

File ~/.virtualenvs/pyle/lib/python3.8/site-packages/cirq/sim/simulator.py:518, in SimulatesFinalState.simulate(self, program, param_resolver, qubit_order, initial_state)
    493 def simulate(
    494     self,
    495     program: 'cirq.AbstractCircuit',
   (...)
    498     initial_state: Any = None,
    499 ) -> TSimulationTrialResult:
    500     """Simulates the supplied Circuit.
    501 
    502     This method returns a result which allows access to the entire
   (...)
    516         SimulationTrialResults for the simulation. Includes the final state.
    517     """
--> 518     return self.simulate_sweep(
    519         program, study.ParamResolver(param_resolver), qubit_order, initial_state
    520     )[0]

File ~/.virtualenvs/pyle/lib/python3.8/site-packages/cirq/sim/simulator.py:533, in SimulatesFinalState.simulate_sweep(self, program, params, qubit_order, initial_state)
    522 def simulate_sweep(
    523     self,
    524     program: 'cirq.AbstractCircuit',
   (...)
    527     initial_state: Any = None,
    528 ) -> List[TSimulationTrialResult]:
    529     """Wraps computed states in a list.
    530 
    531     Prefer overriding `simulate_sweep_iter`.
    532     """
--> 533     return list(self.simulate_sweep_iter(program, params, qubit_order, initial_state))

File ~/.virtualenvs/pyle/lib/python3.8/site-packages/cirq/sim/simulator_base.py:345, in SimulatorBase.simulate_sweep_iter(self, program, params, qubit_order, initial_state)
    343 qubits = ops.QubitOrder.as_qubit_order(qubit_order).order_for(program.all_qubits())
    344 initial_state = 0 if initial_state is None else initial_state
--> 345 sim_state = self._create_act_on_args(initial_state, qubits)
    346 prefix, suffix = (
    347     split_into_matching_protocol_then_general(program, sweep_prefixable)
    348     if self._can_be_in_run_prefix(self.noise)
    349     else (program[0:0], program)
    350 )
    351 step_result = None

File ~/.virtualenvs/pyle/lib/python3.8/site-packages/cirq/sim/simulator_base.py:380, in SimulatorBase._create_act_on_args(self, initial_state, qubits)
    378         initial_state = int(initial_state / q.dimension)
    379 else:
--> 380     args = self._create_partial_act_on_args(
    381         initial_state=initial_state,
    382         qubits=qubits,
    383         classical_data=classical_data,
    384     )
    385     for q in qubits:
    386         args_map[q] = args

File ~/.virtualenvs/pyle/lib/python3.8/site-packages/cirq/sim/density_matrix_simulator.py:221, in DensityMatrixSimulator._create_partial_act_on_args(self, initial_state, qubits, classical_data)
    212 if self._ignore_measurement_results:
    213     return act_on_density_matrix_args.ActOnDensityMatrixArgs(
    214         qubits=qubits,
    215         prng=self._prng,
   (...)
    219         dtype=self._dtype,
    220     )
--> 221 return act_on_density_matrix_args.ActOnDensityMatrixArgs(
    222     qubits=qubits,
    223     prng=self._prng,
    224     classical_data=classical_data,
    225     initial_state=initial_state,
    226     dtype=self._dtype,
    227 )

File ~/.virtualenvs/pyle/lib/python3.8/site-packages/cirq/_compat.py:330, in deprecated_parameter.<locals>.decorator.<locals>.decorated_func(*args, **kwargs)
    322     qualname = func.__qualname__ if func_name is None else func_name
    323     _warn_or_error(
    324         f'The {parameter_desc} parameter of {qualname} was '
    325         f'used but is deprecated.\n'
    326         f'It will be removed in cirq {deadline}.\n'
    327         f'{fix}\n',
    328     )
--> 330 return func(*args, **kwargs)

File ~/.virtualenvs/pyle/lib/python3.8/site-packages/cirq/_compat.py:330, in deprecated_parameter.<locals>.decorator.<locals>.decorated_func(*args, **kwargs)
    322     qualname = func.__qualname__ if func_name is None else func_name
    323     _warn_or_error(
    324         f'The {parameter_desc} parameter of {qualname} was '
    325         f'used but is deprecated.\n'
    326         f'It will be removed in cirq {deadline}.\n'
    327         f'{fix}\n',
    328     )
--> 330 return func(*args, **kwargs)

File ~/.virtualenvs/pyle/lib/python3.8/site-packages/cirq/_compat.py:330, in deprecated_parameter.<locals>.decorator.<locals>.decorated_func(*args, **kwargs)
    322     qualname = func.__qualname__ if func_name is None else func_name
    323     _warn_or_error(
    324         f'The {parameter_desc} parameter of {qualname} was '
    325         f'used but is deprecated.\n'
    326         f'It will be removed in cirq {deadline}.\n'
    327         f'{fix}\n',
    328     )
--> 330 return func(*args, **kwargs)

File ~/.virtualenvs/pyle/lib/python3.8/site-packages/cirq/sim/act_on_density_matrix_args.py:307, in ActOnDensityMatrixArgs.__init__(self, target_tensor, available_buffer, qid_shape, prng, log_of_measurement_results, qubits, ignore_measurement_results, initial_state, dtype, classical_data)
    242 @_compat.deprecated_parameter(
    243     deadline='v0.15',
    244     fix='Use classical_data.',
   (...)
    271     classical_data: Optional['cirq.ClassicalDataStore'] = None,
    272 ):
    273     """Inits ActOnDensityMatrixArgs.
    274 
    275     Args:
   (...)
    305             and `qid_shape` is not provided.
    306     """
--> 307     state = _BufferedDensityMatrix.create(
    308         initial_state=target_tensor if target_tensor is not None else initial_state,
    309         qid_shape=tuple(q.dimension for q in qubits) if qubits is not None else None,
    310         dtype=dtype,
    311         buffer=available_buffer,
    312     )
    313     if ignore_measurement_results:
    314         super().__init__(
    315             state=state,
    316             prng=prng,
   (...)
    320             classical_data=classical_data,
    321         )

File ~/.virtualenvs/pyle/lib/python3.8/site-packages/cirq/sim/act_on_density_matrix_args.py:82, in _BufferedDensityMatrix.create(cls, initial_state, qid_shape, dtype, buffer)
     80 else:
     81     if qid_shape is not None:
---> 82         density_matrix = initial_state.reshape(qid_shape * 2)
     83     else:
     84         density_matrix = initial_state

ValueError: cannot reshape array of size 2 into shape (2,2)
```
</details>

Resolving patch:
diff --git a/cirq-core/cirq/qis/states.py b/cirq-core/cirq/qis/states.py
--- a/cirq-core/cirq/qis/states.py
+++ b/cirq-core/cirq/qis/states.py
@@ -968,9 +968,13 @@ def to_valid_density_matrix(
         ValueError if the density_matrix_rep is not valid.
     """
     qid_shape = _qid_shape_from_args(num_qubits, qid_shape)
-    if isinstance(density_matrix_rep, np.ndarray) and density_matrix_rep.ndim == 2:
-        validate_density_matrix(density_matrix_rep, qid_shape=qid_shape, dtype=dtype, atol=atol)
-        return density_matrix_rep
+    if isinstance(density_matrix_rep, np.ndarray):
+        N = np.prod(qid_shape, dtype=np.int64)
+        if len(qid_shape) > 1 and density_matrix_rep.shape == qid_shape * 2:
+            density_matrix_rep = density_matrix_rep.reshape((N, N))
+        if density_matrix_rep.shape == (N, N):
+            validate_density_matrix(density_matrix_rep, qid_shape=qid_shape, dtype=dtype, atol=atol)
+            return density_matrix_rep
 
     state_vector = to_valid_state_vector(
         density_matrix_rep, len(qid_shape), qid_shape=qid_shape, dtype=dtype
diff --git a/cirq-core/cirq/sim/act_on_density_matrix_args.py b/cirq-core/cirq/sim/act_on_density_matrix_args.py
--- a/cirq-core/cirq/sim/act_on_density_matrix_args.py
+++ b/cirq-core/cirq/sim/act_on_density_matrix_args.py
@@ -79,7 +79,11 @@ def create(
             ).reshape(qid_shape * 2)
         else:
             if qid_shape is not None:
-                density_matrix = initial_state.reshape(qid_shape * 2)
+                if dtype and initial_state.dtype != dtype:
+                    initial_state = initial_state.astype(dtype)
+                density_matrix = qis.to_valid_density_matrix(
+                    initial_state, len(qid_shape), qid_shape=qid_shape, dtype=dtype
+                ).reshape(qid_shape * 2)
             else:
                 density_matrix = initial_state
             if np.may_share_memory(density_matrix, initial_state):


Test patch:
diff --git a/cirq-core/cirq/qis/states_test.py b/cirq-core/cirq/qis/states_test.py
--- a/cirq-core/cirq/qis/states_test.py
+++ b/cirq-core/cirq/qis/states_test.py
@@ -607,6 +607,21 @@ def test_to_valid_density_matrix_from_density_matrix():
     assert_valid_density_matrix(np.diag([0.2, 0.8, 0, 0]), qid_shape=(4,))
 
 
+def test_to_valid_density_matrix_from_density_matrix_tensor():
+    np.testing.assert_almost_equal(
+        cirq.to_valid_density_matrix(
+            cirq.one_hot(shape=(2, 2, 2, 2, 2, 2), dtype=np.complex64), num_qubits=3
+        ),
+        cirq.one_hot(shape=(8, 8), dtype=np.complex64),
+    )
+    np.testing.assert_almost_equal(
+        cirq.to_valid_density_matrix(
+            cirq.one_hot(shape=(2, 3, 4, 2, 3, 4), dtype=np.complex64), qid_shape=(2, 3, 4)
+        ),
+        cirq.one_hot(shape=(24, 24), dtype=np.complex64),
+    )
+
+
 def test_to_valid_density_matrix_not_square():
     with pytest.raises(ValueError, match='shape'):
         cirq.to_valid_density_matrix(np.array([[1], [0]]), num_qubits=1)
@@ -614,7 +629,7 @@ def test_to_valid_density_matrix_not_square():
 
 def test_to_valid_density_matrix_size_mismatch_num_qubits():
     with pytest.raises(ValueError, match='shape'):
-        cirq.to_valid_density_matrix(np.array([[1, 0], [0, 0]]), num_qubits=2)
+        cirq.to_valid_density_matrix(np.array([[[1, 0], [0, 0]], [[0, 0], [0, 0]]]), num_qubits=2)
     with pytest.raises(ValueError, match='shape'):
         cirq.to_valid_density_matrix(np.eye(4) / 4.0, num_qubits=1)
 
@@ -690,6 +705,16 @@ def test_to_valid_density_matrix_from_state_vector():
     )
 
 
+def test_to_valid_density_matrix_from_state_vector_tensor():
+    np.testing.assert_almost_equal(
+        cirq.to_valid_density_matrix(
+            density_matrix_rep=np.array(np.full((2, 2), 0.5), dtype=np.complex64),
+            num_qubits=2,
+        ),
+        0.25 * np.ones((4, 4)),
+    )
+
+
 def test_to_valid_density_matrix_from_state_invalid_state():
     with pytest.raises(ValueError, match="Invalid quantum state"):
         cirq.to_valid_density_matrix(np.array([1, 0, 0]), num_qubits=2)
diff --git a/cirq-core/cirq/sim/act_on_density_matrix_args_test.py b/cirq-core/cirq/sim/act_on_density_matrix_args_test.py
--- a/cirq-core/cirq/sim/act_on_density_matrix_args_test.py
+++ b/cirq-core/cirq/sim/act_on_density_matrix_args_test.py
@@ -98,3 +98,50 @@ def test_with_qubits():
 def test_qid_shape_error():
     with pytest.raises(ValueError, match="qid_shape must be provided"):
         cirq.sim.act_on_density_matrix_args._BufferedDensityMatrix.create(initial_state=0)
+
+
+def test_initial_state_vector():
+    qubits = cirq.LineQubit.range(3)
+    args = cirq.ActOnDensityMatrixArgs(
+        qubits=qubits, initial_state=np.full((8,), 1 / np.sqrt(8)), dtype=np.complex64
+    )
+    assert args.target_tensor.shape == (2, 2, 2, 2, 2, 2)
+
+    args2 = cirq.ActOnDensityMatrixArgs(
+        qubits=qubits, initial_state=np.full((2, 2, 2), 1 / np.sqrt(8)), dtype=np.complex64
+    )
+    assert args2.target_tensor.shape == (2, 2, 2, 2, 2, 2)
+
+
+def test_initial_state_matrix():
+    qubits = cirq.LineQubit.range(3)
+    args = cirq.ActOnDensityMatrixArgs(
+        qubits=qubits, initial_state=np.full((8, 8), 1 / 8), dtype=np.complex64
+    )
+    assert args.target_tensor.shape == (2, 2, 2, 2, 2, 2)
+
+    args2 = cirq.ActOnDensityMatrixArgs(
+        qubits=qubits, initial_state=np.full((2, 2, 2, 2, 2, 2), 1 / 8), dtype=np.complex64
+    )
+    assert args2.target_tensor.shape == (2, 2, 2, 2, 2, 2)
+
+
+def test_initial_state_bad_shape():
+    qubits = cirq.LineQubit.range(3)
+    with pytest.raises(ValueError, match="Invalid quantum state"):
+        cirq.ActOnDensityMatrixArgs(
+            qubits=qubits, initial_state=np.full((4,), 1 / 2), dtype=np.complex64
+        )
+    with pytest.raises(ValueError, match="Invalid quantum state"):
+        cirq.ActOnDensityMatrixArgs(
+            qubits=qubits, initial_state=np.full((2, 2), 1 / 2), dtype=np.complex64
+        )
+
+    with pytest.raises(ValueError, match="Invalid quantum state"):
+        cirq.ActOnDensityMatrixArgs(
+            qubits=qubits, initial_state=np.full((4, 4), 1 / 4), dtype=np.complex64
+        )
+    with pytest.raises(ValueError, match="Invalid quantum state"):
+        cirq.ActOnDensityMatrixArgs(
+            qubits=qubits, initial_state=np.full((2, 2, 2, 2), 1 / 4), dtype=np.complex64
+        )