Density matrix traces not remaining at 1 after certain operations

[mvchalupnik]

Description of the issue
When looping over certain gate operations and measurements, the trace of the density matrix decreases and eventually approaches 0.

How to reproduce the issue
The following script (@kinetic-cipher) reproduces the issue:

import cirq
import numpy as np

q0 = cirq.NamedQubit('q0')
q1 = cirq.NamedQubit('q1')
qc = cirq.Circuit()
qc.append( cirq.CNOT.on(q1,q0) )
qc.append( cirq.H.on(q1) )
qc.append( cirq.measure(q1) )
sim = cirq.DensityMatrixSimulator()
initial_state = None

for _ in range(100):
    output = sim.simulate(qc, initial_state=initial_state)
    initial_state = output.final_density_matrix
    print( np.trace(initial_state) )

The printed output is:

(0.9999999+0j)
(0.99999976+0j)
(0.9999994+0j)
(0.9999988+0j)
(0.9999976+0j)
(0.99999523+0j)
(0.99999034+0j)
(0.99998057+0j)
(0.999961+0j)
(0.99992204+0j)
(0.99984396+0j)
(0.9996879+0j)
(0.99937594+0j)
(0.9987523+0j)
(0.997506+0j)
(0.9950181+0j)
(0.99006104+0j)
(0.98022074+0j)
(0.9608326+0j)
(0.9231992+0j)
(0.8522966+0j)
(0.7264095+0j)
(0.5276707+0j)
(0.27843636+0j)
(0.07752679+0j)
(0.0060104034+0j)
(3.6124948e-05+0j)
(1.3050119e-09+0j)
(1.703056e-18+0j)
(2.9003994e-36+0j)
0j

with error (cirq v0.14.0):

ValueError: probabilities contain NaN

or with error (cirq v1.0.0):

When modifying the same script to use np.complex128 as the dtype of the DensityMatrixSimulator

sim = cirq.DensityMatrixSimulator(dtype=np.complex128) 

the density matrix traces get very large.

(1.0000000000000004+0j) (1.0000000000000013+0j) (1.0000000000000027+0j) (1.0000000000000058+0j) (1.000000000000012+0j) (1.000000000000024+0j) (1.0000000000000484+0j) (1.0000000000000968+0j) (1.000000000000194+0j) (1.0000000000003886+0j) (1.0000000000007776+0j) (1.0000000000015556+0j) (1.0000000000031117+0j) (1.000000000006224+0j) (1.0000000000124478+0j) (1.000000000024896+0j) (1.0000000000497922+0j) (1.0000000000995848+0j) (1.00000000019917+0j) (1.0000000003983405+0j) (1.0000000007966814+0j) (1.0000000015933632+0j) (1.0000000031867269+0j) (1.0000000063734538+0j) (1.0000000127469075+0j) (1.0000000254938157+0j) (1.0000000509876326+0j) (1.0000001019752678+0j) (1.0000002039505465+0j) (1.0000004079011349+0j) (1.000000815802436+0j) (1.000001631605538+0j) (1.0000032632137386+0j) (1.0000065264381257+0j) (1.0000130529188458+0j) (1.0000261060080706+0j) (1.000052212697665+0j) (1.0001044281214957+0j) (1.0002088671482243+0j) (1.0004177779219348+0j) (1.000835730382262+0j) (1.0016721592097961+0j) (1.003347114536015+0j) (1.0067054322477473+0j) (1.0134558273171237+0j) (1.0270927139230361+0j) (1.0549194429937883+0j) (1.112855031206325+0j) (1.238446320481231+0j) (1.5337492887135007+0j) (2.35238688062917+0j) (5.533724036156239+0j) (30.622101708333307+0j) (937.7131130355101+0j) (879305.8823587475+0j) (773178834750.6958+0j) (5.9780551050644395e+23+0j) (3.5737142839187013e+47+0j) (1.2771433783084559e+95+0j) (1.6310952087571362e+190+0j)

with the same error:

ValueError: probabilities contain NaN

Cirq version
We tested this in cirq version 1.0.0 as well as 0.14.0 and found the error in both cases.

[daxfohl]

This seems to be happening in the factor when the measurement gate is encountered. https://github.com/quantumlib/cirq/blob/8953eebe05e2510d798356dfd9fc2d5ba922cfe9/cirq-core/cirq/linalg/transformations.py#L636 likely needs to renormalize after getting the partial traces.

For now, OP can create the simulator as DensityMatrixSimulator(split_untangled_qubits=False), which avoids the factoring after measurement.

[daxfohl]

Note reverting #4300 fixes this, though that implementation is still incorrectly (and worsely) normalized in other cases. I think the pre-#4300 implementation should be faster since it calculates both factors at the same time rather than two calls to partial_trace (I haven't tested this). However it needs renormalized as a density matrix rather than a state vector (extracted /= np.trace(extracted) and same for remainder).

[kinetic-cipher]

This seems to be happening in the factor when the measurement gate is encountered. https://github.com/quantumlib/cirq/blob/8953eebe05e2510d798356dfd9fc2d5ba922cfe9/cirq-core/cirq/linalg/transformations.py#L636 likely needs to renormalize after getting the partial traces.

For now, OP can create the simulator as DensityMatrixSimulator(split_untangled_qubits=False), which avoids the factoring after measurement.

split_untangled_states = False seems to work in the provided test script

[daxfohl]

Ran some tests and the newer implementation using partial_trace is faster than the old implementation. Will push a PR that restabilizes the DM by trace after factoring.

[daxfohl]

I closed the linked PR with the fix because I think the right thing to do here is to do the renormalization in partial_trace itself. The einsum there is losing precision and ending up with a not-quite-one trace. Doing the renormalization factor_density_matrix (which calls partial_trace), as I did in the PR, solves the issue for density matrix simulations, but leaves the problem open for any algorithms or users using partial_trace directly.

Of course fixing this in partial_trace will have larger perf impact and the extent of that impact needs to be understood before making the change. This investigation and the determination of whether that is acceptable or not is a something that needs to be managed by a core maintainer.

[pavoljuhas]

NaN seems to be a problem. Otherwise this appears to be accumulation of numerical errors.
We may consider if the input state should be renormalized to have trace equal 1.