Normalizes state after every W gate in simulator

Author: dabacon

cirq/google/sim/xmon_stepper.py

@@ -313,6 +313,13 @@ def simulate_w(self,
             # W gate is within a shard.
             self._pool.map(_w_within_shard, args)
 
+        # Normalize after every w.
+        norm = np.sum(self._pool.map(_norm, args))
+        args = self._shard_num_args({
+            'norm': norm
+        })
+        self._pool.map(_renorm, args)
+
     def simulate_measurement(self, index: int) -> bool:
         """Simulates a single qubit measurement in the computational basis.
 
@@ -507,6 +514,19 @@ def _one_prob_per_shard(args: Dict[str, Any]) -> float:
     return norm * norm
 
 
+def _norm(args: Dict[str, Any]) -> float:
+    """Returns the norm for each state shard."""
+    state = _state_shard(args)
+    return np.dot(state, np.conjugate(state))
+
+
+def _renorm(args: Dict[str, Any]):
+    """Renormalizes the state using the norm arg."""
+    state = _state_shard(args)
+    # If our gate is so bad that we have norm of zero, we have bigger problems.
+    state /= args['norm']
+
+
 def _collapse_state(args: Dict[str, Any]):
     """Projects state shards onto the appropriate post measurement state.
 

cirq/google/sim/xmon_stepper_test.py

@@ -594,3 +594,33 @@ def test_shard_for_more_prefix_qubits_than_qubits():
         expected = np.zeros(2 ** 2, dtype=np.complex64)
         expected[0] = 1.0
         np.testing.assert_almost_equal(expected, s.current_state)
+
+
+@pytest.mark.parametrize('num_prefix_qubits', (0, 2))
+def test_precision(num_prefix_qubits):
+    # 16 random W's followed by their inverses on five qubits.
+    # The floating point epsilon for np.float32 is about 1e-7.
+    # Floating point epsilon is about 1e-7.
+    # Each qubits error will add across gates on that qubit, but it is like a
+    # random walk, so error should be about sqrt(16) per qubit.
+    # The total error should then be about 2e-6.
+    with xmon_stepper.Stepper(num_qubits=5,
+                              num_prefix_qubits=num_prefix_qubits,
+                              min_qubits_before_shard=0) as s:
+        half_turns_list = [np.random.rand() for _ in range(16)]
+        axis_half_turns_list = [np.random.rand() for _ in range(16)]
+
+        for half_turns, axis_half_turns in zip(half_turns_list,
+                                               axis_half_turns_list):
+            for index in range(5):
+                s.simulate_w(index=index, axis_half_turns=axis_half_turns,
+                             half_turns=half_turns)
+        for half_turns, axis_half_turns in zip(half_turns_list[::-1],
+                                               axis_half_turns_list[::-1]):
+            for index in range(5):
+                s.simulate_w(index=index, axis_half_turns=axis_half_turns,
+                             half_turns=-half_turns)
+        expected = np.zeros(2 ** 5, dtype=np.complex64)
+        expected[0] = 1.0
+        # asserts that abs value of arrays is < 1.5 * 10^(-decimal)
+        np.testing.assert_almost_equal(expected, s.current_state, decimal=6)