Issue with dropping terminal measurements for simulation

[glanzz]

Description of the issue
The drop terminal measurements (https://quantumai.google/reference/python/cirq/drop_terminal_measurements) allows user to drop the terminal measures before simulation to get the proper simulation results.
I am trying to optimize a circuit and drop the terminal measurements, but the fidelity and final vector of the optimized circuit is different.
Seems like the measurement operation is absorbing the Z operation even though the Z gate is not separate (part of PhasedXZ gate).
Expected it to replace back the terminal measures like it does for X gate.
Eg:

Optimized:
1: ───PhX(-0.5)^-0.5───!M───

After dropping terminal measures:
1: ───PhX(-0.5)^-0.5───X───

How to reproduce the issue

import cirq as c
import numpy as np



def optimize_circuit(circuit):
    optimized_circuit = c.eject_z(circuit)
    optimized_circuit = c.eject_phased_paulis(optimized_circuit)
    c.testing.assert_circuits_with_terminal_measurements_are_equivalent(
        circuit, optimized_circuit
    )

    return optimized_circuit

qubit = c.LineQubit(1)
circuit = c.Circuit()
circuit.append(c.ZPowGate(exponent=0.71)(qubit))
circuit.append(c.Rz(rads=np.pi*0.125)(qubit))
circuit.append(c.PhasedXPowGate(**{"phase_exponent": 0.409, "exponent": 0.248})(qubit))
circuit.append(c.PhasedXZGate(**{"axis_phase_exponent": 0.555, "x_exponent": 0.803, "z_exponent": 0.0358})(qubit))
circuit.append(c.XPowGate(exponent=0.012)(qubit))
circuit.append(c.S(qubit))
circuit.append(c.Y(qubit))
circuit.append(c.PhasedXZGate(**{"x_exponent": 0.452, "z_exponent": 0.0275, "axis_phase_exponent": 0.489})(qubit))
circuit.append(c.Rx(rads=np.pi*0.333)(qubit))
circuit.append(c.Rz(rads=np.pi*0.25)(qubit))
circuit.append(c.Rx(rads=np.pi)(qubit))
circuit += c.inverse(circuit)
circuit.insert(0, [c.H(qubit)], c.InsertStrategy.NEW)

# Copy circuits for with and without measurements
circuit_without_measurement = circuit.copy()

circuit_with_measurement = circuit.copy()
circuit_with_measurement.append(c.measure_each(qubit))


# Transpile circuit which as measurement appended
transpiled_circuit_with_measurement = c.optimize_for_target_gateset(circuit_with_measurement, gateset=c.SqrtIswapTargetGateset())
print("Transpiled circuit with measurements after gateset SqrtIswap & then optimization:")
print(transpiled_circuit_with_measurement)
transpiled_circuit_with_measurement = optimize_circuit(transpiled_circuit_with_measurement)
print(transpiled_circuit_with_measurement)

# Transpile circuit
transpiled_circuit_without_measurement = c.optimize_for_target_gateset(circuit_without_measurement, gateset=c.SqrtIswapTargetGateset())
transpiled_circuit_without_measurement = optimize_circuit(transpiled_circuit_without_measurement)
print("Transpiled circuit without measurements:")
print(transpiled_circuit_without_measurement)


hqubit = c.LineQubit(0)
h_circuit = c.Circuit()
h_circuit.append(c.H(hqubit))


s = c.Simulator()

transpiled_circuit_with_measurement_but_dropped_measures = c.drop_terminal_measurements(transpiled_circuit_with_measurement)
print("Transpiled circuit with measurement but dropped:")
print(transpiled_circuit_with_measurement_but_dropped_measures)

h_circuit_STATE = s.simulate(h_circuit, qubit_order=[hqubit]).final_state_vector
transpiled_circuit_without_measurement_STATE = s.simulate(transpiled_circuit_without_measurement, qubit_order=[q for q in transpiled_circuit_without_measurement.all_qubits()]).final_state_vector
transpiled_circuit_with_measurement_but_dropped_measures_STATE = s.simulate(transpiled_circuit_with_measurement_but_dropped_measures, qubit_order=[q for q in transpiled_circuit_with_measurement_but_dropped_measures.all_qubits()]).final_state_vector

actual_without_measurement_circuit_fidelity = c.fidelity(h_circuit_STATE, transpiled_circuit_without_measurement_STATE)
dropped_measurement_circuit_fidelity = c.fidelity(h_circuit_STATE, transpiled_circuit_with_measurement_but_dropped_measures_STATE)


print("Fidelity of the circuits should be equal and close to 1 with H circuit (w/o measure, dropped):")
print(actual_without_measurement_circuit_fidelity)
print(dropped_measurement_circuit_fidelity)

print("Simulated States of the circuits (w/o measure, dropped):")
print(transpiled_circuit_without_measurement_STATE)
print(transpiled_circuit_with_measurement_but_dropped_measures_STATE)

Output:

Transpiled circuit with measurements after gateset SqrtIswap & then optimization:
1: ───PhXZ(a=0.5,x=-0.5,z=1.0)───M───
1: ───PhXZ(a=0.5,x=-0.5,z=0)───M───
Transpiled circuit without measurements:
1: ───PhXZ(a=0.5,x=-0.5,z=1.0)───
Transpiled circuit with measurement but dropped:
1: ───PhXZ(a=0.5,x=-0.5,z=0)───I───
Fidelity of the circuits should be equal and close to 1 with H circuit (w/o measure, dropped):
0.999999880790714
0.0

Cirq version: 1.4.0

[github-actions]

This issue is stale because it has been open 30 days with no activity. Remove stale label or comment or this will be closed in 30 days