Update max measurements check (#662)

* update max modes check * update tests * update changelog * Update strawberryfields/program.py Co-authored-by: Mikhail Andrenkov <Mandrenkov@users.noreply.github.com> * update error msg * add link to changelog * add test Co-authored-by: Mikhail Andrenkov <Mandrenkov@users.noreply.github.com>
XanaduAI · Dec 2, 2021 · e8c1403 · e8c1403
1 parent 5c2aeb3
commit e8c1403
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 8 deletions.
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
@@ -10,6 +10,31 @@
   a `TDMProgram` is compiled using the "TDM" compiler.
   [(#659)](https://github.com/XanaduAI/strawberryfields/pull/659)
 
+* Updates `Program.assert_max_number_of_measurements` to expect the maximum number
+  of measurements from the device specification as a flat dictionary entry instead
+  of a nested one.
+  [(#662)](https://github.com/XanaduAI/strawberryfields/pull/662)
+
+  ```python
+  "modes": {
+      "pnr_max": 20,
+      "homodyne_max": 1000,
+      "heterodyne_max": 1000,
+  }
+  ```
+
+  instead of
+
+  ```python
+  "modes": {
+      "max": {
+          "pnr": 20,
+          "homodyne": 1000,
+          "heterodyne": 1000,
+      }
+  }
+  ```
+
 <h3>Documentation</h3>
 
 <h3>Contributors</h3>

diff --git a/strawberryfields/program.py b/strawberryfields/program.py
@@ -526,13 +526,18 @@ def assert_max_number_of_measurements(self, device):
         num_pnr, num_homodyne, num_heterodyne = 0, 0, 0
 
         try:
-            max_pnr = device.modes["max"]["pnr"]
-            max_homodyne = device.modes["max"]["homodyne"]
-            max_heterodyne = device.modes["max"]["heterodyne"]
+            max_pnr = device.modes["pnr_max"]
+            max_homodyne = device.modes["homodyne_max"]
+            max_heterodyne = device.modes["heterodyne_max"]
         except (KeyError, TypeError) as e:
+            device_modes = device.modes
+            if isinstance(device.modes, dict):
+                device_modes = set(device.modes.keys())
+
             raise KeyError(
-                "Device specification must contain an entry for the maximum allowed number "
-                "of measurments. Have you specified the correct target?"
+                "Expected keys for the maximum allowed number of PNR ('pnr_max'), homodyne "
+                "('homodyne_max'), and heterodyne ('heterodyne_max') measurements. Got keys "
+                f"{device_modes}"
             ) from e
 
         for c in self.circuit:
@@ -605,6 +610,8 @@ def compile(self, *, device=None, compiler=None, **kwargs):
                 The default is False.
             warn_connected (bool): If True, the user is warned if the quantum circuit is not weakly
                 connected. The default is True.
+            shots (int): Number of times the program measurement evaluation is repeated. Passed
+                along to the compiled program's ``run_options``.
 
         Returns:
             Program: compiled program

diff --git a/tests/frontend/test_program.py b/tests/frontend/test_program.py
@@ -245,7 +245,7 @@ def test_assert_max_number_of_measurements(self, measure_op, measure_name):
         # set maximum number of measurements to 2, and measure 3 in prog below
         device_dict = {
             "target": "simulon_gaussian",
-            "modes": {"max": {"pnr": 2, "homodyne": 2, "heterodyne": 2}},
+            "modes": {"pnr_max": 2, "homodyne_max": 2, "heterodyne_max": 2},
             "layout": "",
             "gate_parameters": {},
             "compiler": ["gaussian"],
@@ -260,6 +260,28 @@ def test_assert_max_number_of_measurements(self, measure_op, measure_name):
         with pytest.raises(program.CircuitError, match=f"contains 3 {measure_name} measurements"):
             prog.assert_max_number_of_measurements(spec)
 
+    def test_keyerror_assert_max_number_of_measurements(self):
+        """Check that the correct error is raised when calling `prog.assert_number_of_measurements`
+        with an incorrect device spec modes entry."""
+        # set maximum number of measurements to 2, and measure 3 in prog below
+        device_dict = {
+            "target": "simulon_gaussian",
+            "modes": {"max": {"pnr": 2, "homodyne": 2, "heterodyne": 2}},
+            "layout": "",
+            "gate_parameters": {},
+            "compiler": ["gaussian"],
+        }
+        spec = sf.DeviceSpec(spec=device_dict)
+
+        prog = sf.Program(3)
+        with prog.context as q:
+            for reg in q:
+                ops.MeasureFock() | reg
+
+        match = "Expected keys for the maximum allowed number of PNR"
+        with pytest.raises(KeyError, match=match):
+            prog.assert_max_number_of_measurements(spec)
+
     def test_assert_max_number_of_measurements_wrong_entry(self):
         """Check that the correct error is raised when calling `prog.assert_number_of_measurements`
         with the incorrect type of device spec mode entry."""
@@ -277,7 +299,7 @@ def test_assert_max_number_of_measurements_wrong_entry(self):
             ops.S2gate(0.6) | [q[0], q[1]]
             ops.S2gate(0.6) | [q[1], q[2]]
 
-        with pytest.raises(KeyError, match="Have you specified the correct target?"):
+        with pytest.raises(KeyError, match="Expected keys for the maximum allowed number of PNR"):
             prog.assert_max_number_of_measurements(spec)
 
     def test_has_post_selection(self):
@@ -565,7 +587,7 @@ class DummyCompiler(Compiler):
         # set maximum number of measurements to 2, and measure 3 in prog below
         device_dict = {
             "target": "simulon_gaussian",
-            "modes": {"max": {"pnr": 2, "homodyne": 2, "heterodyne": 2}},
+            "modes": {"pnr_max": 2, "homodyne_max": 2, "heterodyne_max": 2},
             "layout": "",
             "gate_parameters": {},
             "compiler": ["gaussian"],