Fix snapshot probabilities extension (#380)

* fix probabilities snapshot extension

* Add snapshot probabilities tests

CHANGELOG.md

@@ -24,10 +24,12 @@ Added
 - Added tests for the cu1 gate (#360)
 - Added tests for statevector and stabilizer snapshots (\#355)
 - Added tests for density matrix snapshot (\#374)
+- Added tests for probabilities snapshot (\#380)
 
 Changed
 -------
-- Change signature of density matrix snapshot extension (\#374)
+- Change signature of SnapshotProbabilities extension (\#380)
+- Change signature of SnapshotDensityMatrix extension (\#374)
 - Stabilizer snapshot returns stabilizer instead of full Clifford table (\#355)
 - Signature of SnapshotStatevector and SnapshotStabilizer (\#355)
 - Changed all names from tensor_network_state to matrix_product_state (\#356)

qiskit/providers/aer/extensions/snapshot_probabilities.py

@@ -18,47 +18,46 @@
 from qiskit.providers.aer.extensions import Snapshot
 
 
-class SnapshotProbabilites(Snapshot):
+class SnapshotProbabilities(Snapshot):
     """Snapshot instruction for all methods of Qasm simulator."""
 
-    def __init__(self,
-                 label,
-                 num_qubits=0,
-                 num_clbits=0,
-                 variance=False,
-                 params=None,):
+    def __init__(self, label, num_qubits, variance=False):
+        """Create a probability snapshot instruction.
 
-        if variance:
-            super().__init__(label, 'probabilities_with_variance', num_qubits, num_clbits, params)
-        else:
-            super().__init__(label, 'probabilities', num_qubits, num_clbits, params)
+        Args:
+            label (str): the snapshot label.
+            num_qubits (int): the number of qubits to snapshot.
+            variance (bool): compute variance of probabilities [Default: False]
 
+        Raises:
+            ExtensionError: if snapshot is invalid.
+        """
+        snapshot_type = 'probabilities_with_variance' if variance else 'probabilities'
+        super().__init__(label, snapshot_type=snapshot_type,
+                         num_qubits=num_qubits)
+
+
+def snapshot_probabilities(self, label, qubits, variance=False):
+    """Take a probability snapshot of the simulator state.
 
-def snapshot_probabilities(self,
-                           label,
-                           qubits=None,
-                           variance=False,
-                           params=None):
-    """Take a snapshot of the internal simulator representation.
-    Works on specified qubits or the full register, and prevents reordering (like barrier).
     Args:
         label (str): a snapshot label to report the result
-        qubits (list or None): the qubits to apply snapshot to [Default: None]
-        variance (bool): set snapshot_type to 'probabilities' or '
-                         probabilities_with_variance' [Default: False]
-        params (list or None): the parameters for snapshot_type [Default: None]
+        qubits (list): the qubits to snapshot.
+        variance (bool): compute variance of probabilities [Default: False]
+
     Returns:
-        QuantumCircuit: with attached command
+        QuantumCircuit: with attached instruction.
+
     Raises:
-        ExtensionError: malformed command
+        ExtensionError: if snapshot is invalid.
     """
     snapshot_register = Snapshot.define_snapshot_register(self, label, qubits)
 
     return self.append(
-        SnapshotProbabilites(label,
-                             num_qubits=len(snapshot_register),
-                             variance=variance,
-                             params=params), snapshot_register)
+        SnapshotProbabilities(label,
+                              num_qubits=len(snapshot_register),
+                              variance=variance),
+        snapshot_register)
 
 
 QuantumCircuit.snapshot_probabilities = snapshot_probabilities

test/terra/backends/qasm_simulator/qasm_snapshot.py

@@ -29,10 +29,12 @@
     snapshot_state_counts_nondeterministic,
     snapshot_state_pre_measure_statevector_nondeterministic,
     snapshot_state_post_measure_statevector_nondeterministic)
-
-def get_snapshots(data, label, snapshot_type):
-        """Format snapshots as list of Numpy arrays"""
-        return data.get("snapshots", {}).get(snapshot_type, {}).get(label, [])
+from test.terra.reference.ref_snapshot_probabilities import (
+    snapshot_probabilities_circuits,
+    snapshot_probabilities_counts,
+    snapshot_probabilities_labels_qubits,
+    snapshot_probabilities_post_meas_probs,
+    snapshot_probabilities_pre_meas_probs)
 
 
 class QasmSnapshotStatevectorTests:
@@ -484,3 +486,81 @@ def test_snapshot_density_matrix_post_measure_nondet(self):
                     target = np.outer(target, target.conj())
                     value = snaps.get(mem)
                     self.assertTrue(np.allclose(value, target))
+
+
+class QasmSnapshotProbabilitiesTests:
+    """QasmSimulator snapshot probabilities tests."""
+
+    SIMULATOR = QasmSimulator()
+    SUPPORTED_QASM_METHODS = [
+        'automatic', 'statevector', 'density_matrix', 'matrix_product_state'
+    ]
+    BACKEND_OPTS = {}
+
+    @staticmethod
+    def probabilitiy_snapshots(data, labels):
+        """Format snapshots as nested dicts"""
+        # Check snapshot entry exists in data
+        output = {}
+        for label in labels:
+            snaps = data.get("snapshots", {}).get("probabilities", {}).get(label, [])
+            output[label] = {snap_dict['memory']: snap_dict['value']
+                             for snap_dict in snaps}
+        return output
+
+    def test_snapshot_probabilities_pre_measure(self):
+        """Test snapshot probabilities before final measurement"""
+        shots = 1000
+        labels = list(snapshot_probabilities_labels_qubits().keys())
+        counts_targets = snapshot_probabilities_counts(shots)
+        prob_targets = snapshot_probabilities_pre_meas_probs()
+
+        circuits = snapshot_probabilities_circuits(post_measure=False)
+
+        qobj = assemble(circuits, self.SIMULATOR, shots=shots)
+        job = self.SIMULATOR.run(qobj, backend_options=self.BACKEND_OPTS)
+        method = self.BACKEND_OPTS.get('method', 'automatic')
+        if method not in QasmSnapshotProbabilitiesTests.SUPPORTED_QASM_METHODS:
+            self.assertRaises(AerError, job.result)
+        else:
+            result = job.result()
+            self.is_completed(result)
+            self.compare_counts(result, circuits, counts_targets, delta=0.1 * shots)
+            # Check snapshots
+            for j, circuit in enumerate(circuits):
+                data = result.data(circuit)
+                all_snapshots = self.probabilitiy_snapshots(data, labels)
+                for label in labels:
+                    snaps = all_snapshots.get(label, {})
+                    self.assertTrue(len(snaps), 1)
+                    for memory, value in snaps.items():
+                        target = prob_targets[j].get(label, {}).get(memory, {})
+                        self.assertDictAlmostEqual(value, target, delta=1e-7)
+
+    def test_snapshot_probabilities_post_measure(self):
+        """Test snapshot probabilities after final measurement"""
+        shots = 1000
+        labels = list(snapshot_probabilities_labels_qubits().keys())
+        counts_targets = snapshot_probabilities_counts(shots)
+        prob_targets = snapshot_probabilities_post_meas_probs()
+
+        circuits = snapshot_probabilities_circuits(post_measure=True)
+
+        qobj = assemble(circuits, self.SIMULATOR, shots=shots)
+        job = self.SIMULATOR.run(qobj, backend_options=self.BACKEND_OPTS)
+        method = self.BACKEND_OPTS.get('method', 'automatic')
+        if method not in QasmSnapshotProbabilitiesTests.SUPPORTED_QASM_METHODS:
+            self.assertRaises(AerError, job.result)
+        else:
+            result = job.result()
+            self.is_completed(result)
+            self.compare_counts(result, circuits, counts_targets, delta=0.1 * shots)
+            # Check snapshots
+            for j, circuit in enumerate(circuits):
+                data = result.data(circuit)
+                all_snapshots = self.probabilitiy_snapshots(data, labels)
+                for label in labels:
+                    snaps = all_snapshots.get(label, {})
+                    for memory, value in snaps.items():
+                        target = prob_targets[j].get(label, {}).get(memory, {})
+                        self.assertDictAlmostEqual(value, target, delta=1e-7)

test/terra/backends/test_qasm_density_matrix_simulator.py

@@ -48,6 +48,7 @@
 from test.terra.backends.qasm_simulator.qasm_snapshot import QasmSnapshotStatevectorTests
 from test.terra.backends.qasm_simulator.qasm_snapshot import QasmSnapshotDensityMatrixTests
 from test.terra.backends.qasm_simulator.qasm_snapshot import QasmSnapshotStabilizerTests
+from test.terra.backends.qasm_simulator.qasm_snapshot import QasmSnapshotProbabilitiesTests
 
 
 class TestQasmDensityMatrixSimulator(common.QiskitAerTestCase,
@@ -75,6 +76,7 @@ class TestQasmDensityMatrixSimulator(common.QiskitAerTestCase,
                                      QasmKrausNoiseTests,
                                      QasmSnapshotStatevectorTests,
                                      QasmSnapshotDensityMatrixTests,
+                                     QasmSnapshotProbabilitiesTests,
                                      QasmSnapshotStabilizerTests):
     """QasmSimulator density_matrix method tests."""
 

test/terra/backends/test_qasm_simulator.py

@@ -44,6 +44,7 @@
 from test.terra.backends.qasm_simulator.qasm_snapshot import QasmSnapshotStatevectorTests
 from test.terra.backends.qasm_simulator.qasm_snapshot import QasmSnapshotDensityMatrixTests
 from test.terra.backends.qasm_simulator.qasm_snapshot import QasmSnapshotStabilizerTests
+from test.terra.backends.qasm_simulator.qasm_snapshot import QasmSnapshotProbabilitiesTests
 # Other tests
 from test.terra.backends.qasm_simulator.qasm_method import QasmMethodTests
 from test.terra.backends.qasm_simulator.qasm_thread_management import QasmThreadManagementTests
@@ -85,6 +86,7 @@ class TestQasmSimulator(common.QiskitAerTestCase,
                         QasmBasicsTests,
                         QasmSnapshotStatevectorTests,
                         QasmSnapshotDensityMatrixTests,
+                        QasmSnapshotProbabilitiesTests,
                         QasmSnapshotStabilizerTests):
     """QasmSimulator automatic method tests."""
 

test/terra/backends/test_qasm_stabilizer_simulator.py

@@ -34,6 +34,7 @@
 from test.terra.backends.qasm_simulator.qasm_snapshot import QasmSnapshotStatevectorTests
 from test.terra.backends.qasm_simulator.qasm_snapshot import QasmSnapshotDensityMatrixTests
 from test.terra.backends.qasm_simulator.qasm_snapshot import QasmSnapshotStabilizerTests
+from test.terra.backends.qasm_simulator.qasm_snapshot import QasmSnapshotProbabilitiesTests
 # Other tests
 from test.terra.backends.qasm_simulator.qasm_method import QasmMethodTests
 
@@ -51,6 +52,7 @@ class TestQasmStabilizerSimulator(common.QiskitAerTestCase,
                                   QasmPauliNoiseTests,
                                   QasmSnapshotStatevectorTests,
                                   QasmSnapshotDensityMatrixTests,
+                                  QasmSnapshotProbabilitiesTests,
                                   QasmSnapshotStabilizerTests):
     """QasmSimulator stabilizer method tests."""
 

test/terra/backends/test_qasm_statevector_simulator.py

@@ -46,6 +46,7 @@
 from test.terra.backends.qasm_simulator.qasm_snapshot import QasmSnapshotStatevectorTests
 from test.terra.backends.qasm_simulator.qasm_snapshot import QasmSnapshotDensityMatrixTests
 from test.terra.backends.qasm_simulator.qasm_snapshot import QasmSnapshotStabilizerTests
+from test.terra.backends.qasm_simulator.qasm_snapshot import QasmSnapshotProbabilitiesTests
 # Other tests
 from test.terra.backends.qasm_simulator.qasm_method import QasmMethodTests
 
@@ -75,6 +76,7 @@ class TestQasmStatevectorSimulator(common.QiskitAerTestCase,
                                    QasmKrausNoiseTests,
                                    QasmSnapshotStatevectorTests,
                                    QasmSnapshotDensityMatrixTests,
+                                   QasmSnapshotProbabilitiesTests,
                                    QasmSnapshotStabilizerTests):
     """QasmSimulator statevector method tests."""
 

test/terra/extensions/test_snapshot_probabilities.py

@@ -0,0 +1,107 @@
+# This code is part of Qiskit.
+#
+# (C) Copyright IBM 2018, 2019.
+#
+# This code is licensed under the Apache License, Version 2.0. You may
+# obtain a copy of this license in the LICENSE.txt file in the root directory
+# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# Any modifications or derivative works of this code must retain this
+# copyright notice, and modified files need to carry a notice indicating
+# that they have been altered from the originals.
+
+
+import unittest
+
+from qiskit import QuantumCircuit, assemble
+from qiskit.extensions.exceptions import ExtensionError
+from qiskit.providers.aer.extensions.snapshot_probabilities import SnapshotProbabilities
+
+
+class TestSnapshotProbabilitiesExtension(unittest.TestCase):
+    """SnapshotProbabilities extension tests"""
+
+    @staticmethod
+    def snapshot_circuit_instr(circ_qubits, label, qubits, variance=False):
+        """Return QobjInstruction for circuit monkey patch method."""
+        circuit = QuantumCircuit(circ_qubits)
+        circuit.snapshot_probabilities(label, qubits, variance)
+        qobj = assemble(circuit)
+        instr = qobj.experiments[0].instructions[0]
+        return instr
+
+    def test_snapshot_label_raises(self):
+        """Test snapshot label must be str"""
+        self.assertRaises(ExtensionError, lambda: SnapshotProbabilities(1, 1))
+
+    def test_snapshot_name(self):
+        """Test snapshot instruction has correct name"""
+        instrs = [
+            SnapshotProbabilities('snap', 1, False).assemble(),
+            SnapshotProbabilities('snap', 1, True).assemble(),
+            self.snapshot_circuit_instr(1, 'snap', [0], False),
+            self.snapshot_circuit_instr(1, 'snap', [0], True)
+        ]
+        for instr in instrs:
+            self.assertTrue(hasattr(instr, 'name'))
+            self.assertEqual(instr.name, 'snapshot')
+
+    def test_snapshot_type(self):
+        """Test snapshot instruction has correct type."""
+        # without variance
+        instrs = [
+            SnapshotProbabilities('snap', 1, False).assemble(),
+            self.snapshot_circuit_instr(1, 'snap', [0], False)
+        ]
+        for instr in instrs:
+            self.assertTrue(hasattr(instr, 'snapshot_type'))
+            self.assertEqual(instr.snapshot_type, 'probabilities')
+        # with variance
+        instrs = [
+            SnapshotProbabilities('snap', 1, True).assemble(),
+            self.snapshot_circuit_instr(1, 'snap', [0], True)
+        ]
+        for instr in instrs:
+            self.assertTrue(hasattr(instr, 'snapshot_type'))
+            self.assertEqual(instr.snapshot_type, 'probabilities_with_variance')
+
+    def test_snapshot_label(self):
+        """Test snapshot instruction has correct label"""
+        for label in ['snap0', 'snap1']:
+            instrs = [
+                SnapshotProbabilities(label, 1, False).assemble(),
+                SnapshotProbabilities(label, 1, True).assemble(),
+                self.snapshot_circuit_instr(1, label, [0], False),
+                self.snapshot_circuit_instr(1, label, [0], True)
+            ]
+            for instr in instrs:
+                self.assertTrue(hasattr(instr, 'label'))
+                self.assertEqual(instr.label, label)
+
+    def test_snapshot_all_qubits(self):
+        """Test snapshot instruction has correct qubits."""
+        for j in range(1, 5):
+            instrs = [
+                SnapshotProbabilities('snap', j, False).assemble(),
+                SnapshotProbabilities('snap', j, True).assemble(),
+                self.snapshot_circuit_instr(j, 'snap', range(j), True),
+                self.snapshot_circuit_instr(j, 'snap', range(j), False)
+            ]
+            for instr in instrs:
+                self.assertTrue(hasattr(instr, 'qubits'))
+                self.assertEqual(instr.qubits, list(range(j)))
+
+    def test_snapshot_specific_qubits(self):
+        """Test snapshot instruction has correct qubits."""
+        for qubits in [[0], [0, 2], [1, 3, 0]]:
+            instrs = [
+                self.snapshot_circuit_instr(5, 'snap', qubits, False),
+                self.snapshot_circuit_instr(5, 'snap', qubits, True)
+            ]
+            for instr in instrs:
+                self.assertTrue(hasattr(instr, 'qubits'))
+                self.assertEqual(instr.qubits, qubits)
+
+
+if __name__ == '__main__':
+    unittest.main()