Noise modeling updates for Stim backend (#2258)

* Noise modeling updates for Stim backend * DCO Remediation Commit for Ben Howe <[email protected]> I, Ben Howe <[email protected]>, hereby add my Signed-off-by to this commit: 86681ef Signed-off-by: Ben Howe <[email protected]> --------- Signed-off-by: Ben Howe <[email protected]> Co-authored-by: Bettina Heim <[email protected]>
NVIDIA · Oct 15, 2024 · 7bdab51 · 7bdab51
2 parents 126084d + 6e035fe
commit 7bdab51
Show file tree

Hide file tree

Showing 3 changed files with 106 additions and 18 deletions.
diff --git a/python/tests/builder/test_NoiseModel.py b/python/tests/builder/test_NoiseModel.py
@@ -15,10 +15,11 @@
 import random
 
 
-def test_depolarization_channel():
+@pytest.mark.parametrize('target', ['density-matrix-cpu', 'stim'])
+def test_depolarization_channel(target: str):
     """Tests the depolarization channel in the case of a non-zero probability."""
+    cudaq.set_target(target)
     cudaq.set_random_seed(13)
-    cudaq.set_target('density-matrix-cpu')
     circuit = cudaq.make_kernel()
     q = circuit.qalloc()
     circuit.x(q)
@@ -43,10 +44,11 @@ def test_depolarization_channel():
     assert ('1' in counts)
 
 
-def test_depolarization_channel_simple():
+@pytest.mark.parametrize('target', ['density-matrix-cpu', 'stim'])
+def test_depolarization_channel_simple(target: str):
     """Tests the depolarization channel in the case of `probability = 1.0`"""
+    cudaq.set_target(target)
     cudaq.set_random_seed(13)
-    cudaq.set_target('density-matrix-cpu')
     kernel = cudaq.make_kernel()
     qubit = kernel.qalloc()
     noise = cudaq.NoiseModel()
@@ -81,8 +83,8 @@ def test_depolarization_channel_simple():
 
 def test_amplitude_damping_simple():
     """Tests the amplitude damping channel in the case of `probability = 1.0`"""
-    cudaq.set_random_seed(13)
     cudaq.set_target('density-matrix-cpu')
+    cudaq.set_random_seed(13)
     noise = cudaq.NoiseModel()
     # Amplitude damping channel with `1.0` probability of the qubit
     # decaying to the ground state.
@@ -114,10 +116,11 @@ def test_amplitude_damping_simple():
     cudaq.reset_target()
 
 
-def test_phase_flip_simple():
+@pytest.mark.parametrize('target', ['density-matrix-cpu', 'stim'])
+def test_phase_flip_simple(target: str):
     """Tests the phase flip channel in the case of `probability = 1.0`"""
+    cudaq.set_target(target)
     cudaq.set_random_seed(13)
-    cudaq.set_target('density-matrix-cpu')
     noise = cudaq.NoiseModel()
     # Phase flip channel with `1.0` probability of the qubit
     # undergoing a phase rotation of 180 degrees (π).
@@ -150,13 +153,14 @@ def test_phase_flip_simple():
     cudaq.reset_target()
 
 
-def test_bit_flip_simple():
+@pytest.mark.parametrize('target', ['density-matrix-cpu', 'stim'])
+def test_bit_flip_simple(target: str):
     """
     Tests the bit flip channel with the probability at `0.0` on qubit 0, 
     and `1.0` on qubit 1.
     """
+    cudaq.set_target(target)
     cudaq.set_random_seed(13)
-    cudaq.set_target('density-matrix-cpu')
     noise = cudaq.NoiseModel()
     # Bit flip channel with `0.0` probability of the qubit flipping 180 degrees.
     bit_flip_zero = cudaq.BitFlipChannel(0.0)
@@ -194,8 +198,8 @@ def test_bit_flip_simple():
 
 def test_kraus_channel():
     """Tests the Kraus Channel with a series of custom Kraus Operators."""
-    cudaq.set_random_seed(13)
     cudaq.set_target('density-matrix-cpu')
+    cudaq.set_random_seed(13)
     k0 = np.array([[0.05773502691896258, 0.0], [0., -0.05773502691896258]],
                   dtype=np.complex128)
     k1 = np.array([[0., 0.05773502691896258], [0.05773502691896258, 0.]],
@@ -309,8 +313,9 @@ def test_noise_u3():
     cudaq.reset_target()
 
 
-def test_all_qubit_channel():
-    cudaq.set_target('density-matrix-cpu')
+@pytest.mark.parametrize('target', ['density-matrix-cpu', 'stim'])
+def test_all_qubit_channel(target: str):
+    cudaq.set_target(target)
     cudaq.set_random_seed(13)
     noise = cudaq.NoiseModel()
     bf = cudaq.BitFlipChannel(1.0)
@@ -484,8 +489,8 @@ def noise_cb(qubits, params):
 
 
 def check_custom_op_noise(noise_model):
-    cudaq.set_random_seed(13)
     cudaq.set_target('density-matrix-cpu')
+    cudaq.set_random_seed(13)
 
     @cudaq.kernel
     def basic():

diff --git a/runtime/nvqir/stim/StimCircuitSimulator.cpp b/runtime/nvqir/stim/StimCircuitSimulator.cpp
@@ -69,14 +69,28 @@ class StimCircuitSimulator : public nvqir::CircuitSimulatorBase<double> {
 
     // Get the Kraus channels specified for this gate and qubits
     auto krausChannels = executionContext->noiseModel->get_channels(
-        gName, controls, targets, params);
+        gName, targets, controls, params);
 
     // If none, do nothing
     if (krausChannels.empty())
       return;
 
-    // TODO
-    return;
+    cudaq::info("Applying {} kraus channels to qubits {}", krausChannels.size(),
+                stimTargets);
+
+    for (auto &channel : krausChannels) {
+      if (channel.noise_type == cudaq::noise_model_type::bit_flip_channel)
+        stimCircuit.safe_append_ua("X_ERROR", stimTargets,
+                                   channel.parameters[0]);
+      else if (channel.noise_type ==
+               cudaq::noise_model_type::phase_flip_channel)
+        stimCircuit.safe_append_ua("Z_ERROR", stimTargets,
+                                   channel.parameters[0]);
+      else if (channel.noise_type ==
+               cudaq::noise_model_type::depolarization_channel)
+        stimCircuit.safe_append_ua("DEPOLARIZE1", stimTargets,
+                                   channel.parameters[0]);
+    }
   }
 
   void applyGate(const GateApplicationTask &task) override {
@@ -169,7 +183,6 @@ class StimCircuitSimulator : public nvqir::CircuitSimulatorBase<double> {
                                         randomEngine, false);
     size_t bits_per_sample = stimCircuit.count_measurements();
     std::vector<std::string> sequentialData;
-    sequentialData.reserve(shots);
     // Only retain the final "qubits.size()" measurements. All other
     // measurements were mid-circuit measurements that have been previously
     // accounted for and saved.

diff --git a/unittests/integration/noise_tester.cpp b/unittests/integration/noise_tester.cpp
@@ -11,7 +11,7 @@
 #include <set>
 #include <stdio.h>
 
-#ifdef CUDAQ_BACKEND_DM
+#if defined(CUDAQ_BACKEND_DM) || defined(CUDAQ_BACKEND_STIM)
 struct xOp {
   void operator()() __qpu__ {
     cudaq::qubit q;
@@ -27,6 +27,10 @@ struct bell {
   }
 };
 
+#endif
+#if defined(CUDAQ_BACKEND_DM)
+// Stim does not support arbitrary cudaq::kraus_channel specification.
+
 CUDAQ_TEST(NoiseTest, checkSimple) {
   cudaq::set_random_seed(13);
   cudaq::kraus_channel depol({cudaq::complex{0.99498743710662, 0.0},
@@ -68,6 +72,10 @@ CUDAQ_TEST(NoiseTest, checkSimple) {
   EXPECT_EQ(1, counts.size());
 }
 
+#endif
+#if defined(CUDAQ_BACKEND_DM)
+// Stim does not support arbitrary cudaq::kraus_channel specification.
+
 CUDAQ_TEST(NoiseTest, checkAmplitudeDamping) {
   cudaq::set_random_seed(13);
   cudaq::kraus_channel amplitudeDamping{{1., 0., 0., .8660254037844386},
@@ -84,6 +92,10 @@ CUDAQ_TEST(NoiseTest, checkAmplitudeDamping) {
   cudaq::unset_noise(); // clear for subsequent tests
 }
 
+#endif
+#if defined(CUDAQ_BACKEND_DM)
+// Stim does not support arbitrary cudaq::kraus_op specification.
+
 CUDAQ_TEST(NoiseTest, checkCNOT) {
   cudaq::set_random_seed(13);
   cudaq::kraus_op op0{cudaq::complex{0.99498743710662, 0.0},
@@ -161,6 +173,10 @@ CUDAQ_TEST(NoiseTest, checkCNOT) {
   cudaq::unset_noise(); // clear for subsequent tests
 }
 
+#endif
+#if defined(CUDAQ_BACKEND_DM)
+// Stim does not support arbitrary cudaq::kraus_channel specification.
+
 CUDAQ_TEST(NoiseTest, checkExceptions) {
   cudaq::set_random_seed(13);
   cudaq::kraus_channel amplitudeDamping{{1., 0., 0., .8660254037844386},
@@ -171,6 +187,9 @@ CUDAQ_TEST(NoiseTest, checkExceptions) {
   });
 }
 
+#endif
+#if defined(CUDAQ_BACKEND_DM) || defined(CUDAQ_BACKEND_STIM)
+
 CUDAQ_TEST(NoiseTest, checkDepolType) {
   cudaq::set_random_seed(13);
   cudaq::depolarization_channel depol(.1);
@@ -183,6 +202,9 @@ CUDAQ_TEST(NoiseTest, checkDepolType) {
   cudaq::unset_noise(); // clear for subsequent tests
 }
 
+#endif
+#if defined(CUDAQ_BACKEND_DM) || defined(CUDAQ_BACKEND_STIM)
+
 CUDAQ_TEST(NoiseTest, checkDepolTypeSimple) {
   cudaq::set_random_seed(13);
   cudaq::depolarization_channel depol(1.);
@@ -197,6 +219,10 @@ CUDAQ_TEST(NoiseTest, checkDepolTypeSimple) {
   cudaq::unset_noise(); // clear for subsequent tests
 }
 
+#endif
+#if defined(CUDAQ_BACKEND_DM)
+// Stim does not support cudaq::amplitude_damping_channel.
+
 CUDAQ_TEST(NoiseTest, checkAmpDampType) {
   cudaq::set_random_seed(13);
   cudaq::amplitude_damping_channel ad(.25);
@@ -211,6 +237,10 @@ CUDAQ_TEST(NoiseTest, checkAmpDampType) {
   cudaq::unset_noise(); // clear for subsequent tests
 }
 
+#endif
+#if defined(CUDAQ_BACKEND_DM)
+// Stim does not support cudaq::amplitude_damping_channel.
+
 CUDAQ_TEST(NoiseTest, checkAmpDampTypeSimple) {
   cudaq::set_random_seed(13);
   cudaq::amplitude_damping_channel ad(1.);
@@ -224,6 +254,9 @@ CUDAQ_TEST(NoiseTest, checkAmpDampTypeSimple) {
   cudaq::unset_noise(); // clear for subsequent tests
 }
 
+#endif
+#if defined(CUDAQ_BACKEND_DM) || defined(CUDAQ_BACKEND_STIM)
+
 CUDAQ_TEST(NoiseTest, checkBitFlipType) {
   cudaq::set_random_seed(13);
   cudaq::bit_flip_channel bf(.1);
@@ -238,6 +271,9 @@ CUDAQ_TEST(NoiseTest, checkBitFlipType) {
   cudaq::unset_noise(); // clear for subsequent tests
 }
 
+#endif
+#if defined(CUDAQ_BACKEND_DM) || defined(CUDAQ_BACKEND_STIM)
+
 CUDAQ_TEST(NoiseTest, checkBitFlipTypeSimple) {
   cudaq::set_random_seed(13);
   cudaq::bit_flip_channel bf(1.);
@@ -251,6 +287,8 @@ CUDAQ_TEST(NoiseTest, checkBitFlipTypeSimple) {
   cudaq::unset_noise(); // clear for subsequent tests
 }
 
+#endif
+#if defined(CUDAQ_BACKEND_DM) || defined(CUDAQ_BACKEND_STIM)
 // Same as above but use alternate sample interface that specifies the number of
 // shots and the noise model to use.
 CUDAQ_TEST(NoiseTest, checkBitFlipTypeSimpleOptions) {
@@ -269,6 +307,9 @@ CUDAQ_TEST(NoiseTest, checkBitFlipTypeSimpleOptions) {
   EXPECT_EQ(totalShots, shots);
 }
 
+#endif
+#if defined(CUDAQ_BACKEND_DM) || defined(CUDAQ_BACKEND_STIM)
+
 CUDAQ_TEST(NoiseTest, checkPhaseFlipType) {
   cudaq::set_random_seed(13);
 
@@ -291,6 +332,9 @@ CUDAQ_TEST(NoiseTest, checkPhaseFlipType) {
   cudaq::unset_noise(); // clear for subsequent tests
 }
 
+#endif
+#if defined(CUDAQ_BACKEND_DM) || defined(CUDAQ_BACKEND_STIM)
+
 template <std::size_t N>
 struct xOpAll {
   void operator()() __qpu__ {
@@ -299,6 +343,9 @@ struct xOpAll {
   }
 };
 
+#endif
+#if defined(CUDAQ_BACKEND_DM) || defined(CUDAQ_BACKEND_STIM)
+
 CUDAQ_TEST(NoiseTest, checkAllQubitChannel) {
   cudaq::set_random_seed(13);
   cudaq::bit_flip_channel bf(1.);
@@ -316,6 +363,10 @@ CUDAQ_TEST(NoiseTest, checkAllQubitChannel) {
   EXPECT_EQ(totalShots, shots);
 }
 
+#endif
+#if defined(CUDAQ_BACKEND_DM)
+// Stim does not support arbitrary cudaq::kraus_op specification.
+
 static cudaq::kraus_channel create2pNoiseChannel() {
   cudaq::kraus_op op0{cudaq::complex{0.99498743710662, 0.0},
                       {0.0, 0.0},
@@ -395,6 +446,10 @@ struct bellRandom {
   }
 };
 
+#endif
+#if defined(CUDAQ_BACKEND_DM)
+// Stim does not support arbitrary cudaq::kraus_op specification.
+
 CUDAQ_TEST(NoiseTest, checkAllQubitChannelWithControl) {
   cudaq::set_random_seed(13);
   cudaq::noise_model noise;
@@ -420,6 +475,10 @@ CUDAQ_TEST(NoiseTest, checkAllQubitChannelWithControl) {
   } while (std::next_permutation(qubitIds.begin(), qubitIds.end()));
 }
 
+#endif
+#if defined(CUDAQ_BACKEND_DM)
+// Stim does not support arbitrary cudaq::kraus_op specification.
+
 CUDAQ_TEST(NoiseTest, checkAllQubitChannelWithControlPrefix) {
   cudaq::set_random_seed(13);
   cudaq::noise_model noise;
@@ -444,6 +503,9 @@ CUDAQ_TEST(NoiseTest, checkAllQubitChannelWithControlPrefix) {
   } while (std::next_permutation(qubitIds.begin(), qubitIds.end()));
 }
 
+#endif
+#if defined(CUDAQ_BACKEND_DM) || defined(CUDAQ_BACKEND_STIM)
+
 CUDAQ_TEST(NoiseTest, checkCallbackChannel) {
   cudaq::set_random_seed(13);
   cudaq::noise_model noise;
@@ -474,6 +536,10 @@ struct rxOp {
   }
 };
 
+#endif
+#if defined(CUDAQ_BACKEND_DM)
+// Stim does not support rx gate.
+
 CUDAQ_TEST(NoiseTest, checkCallbackChannelWithParams) {
   cudaq::set_random_seed(13);
   cudaq::noise_model noise;
@@ -507,6 +573,10 @@ CUDAQ_TEST(NoiseTest, checkCallbackChannelWithParams) {
   }
 }
 
+#endif
+#if defined(CUDAQ_BACKEND_DM)
+// Stim does not support custom operations.
+
 CUDAQ_REGISTER_OPERATION(CustomX, 1, 0, {0, 1, 1, 0});
 CUDAQ_TEST(NoiseTest, checkCustomOperation) {
   auto kernel = []() {