Fix get_probability after register updates

ProjectQ-Framework · Jun 9, 2021 · d1a65ab · d1a65ab
1 parent 7a7ec3c
commit d1a65ab
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/projectq/backends/_ionq/_ionq.py b/projectq/backends/_ionq/_ionq.py
@@ -261,7 +261,10 @@ def get_probability(self, state, qureg):
         Returns:
             float: The probability for the provided state.
         """
-        probs = self.get_probabilities(qureg[: len(state)])
+        if len(state) != len(qureg):
+            raise ValueError('Desired state and register must be the same length!')
+
+        probs = self.get_probabilities(qureg)
         return probs[state]
 
     def get_probabilities(self, qureg):

diff --git a/projectq/backends/_ionq/_ionq_test.py b/projectq/backends/_ionq/_ionq_test.py
@@ -189,6 +189,10 @@ def mock_retrieve(*args, **kwargs):
     assert eng.backend.get_probability('11', qureg) == pytest.approx(0.4)
     assert eng.backend.get_probability('00', qureg) == pytest.approx(0.6)
 
+    with pytest.raises(ValueError) as excinfo:
+        eng.backend.get_probability('111', qureg)
+    assert str(excinfo.value) == 'Desired state and register must be the same length!'
+
 
 def test_ionq_get_probabilities(monkeypatch, mapper_factory):
     """Test a shortcut for getting a specific state's probability"""