homodyne fix

XanaduAI · Sep 4, 2024 · 8f135f9 · 8f135f9
1 parent 74d0500
commit 8f135f9
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/mrmustard/lab_dev/sampler.py b/mrmustard/lab_dev/sampler.py
@@ -216,7 +216,7 @@ def __init__(
         num: int = 1000,
     ) -> None:
         meas_outcomes, step = np.linspace(*bounds, num, retstep=True)
-        super().__init__(meas_outcomes, BtoQ(modes, phi=phi))
+        super().__init__(list(product(meas_outcomes, repeat=len(modes))), BtoQ(modes, phi=phi))
         self._step = step
 
     def probabilities(self, state=None, atol=1e-4):
@@ -225,7 +225,7 @@ def probabilities(self, state=None, atol=1e-4):
             trace_modes = self._trace_modes(state)
             dm_state = state.dm() >> TraceOut(trace_modes) if trace_modes else state.dm()
             q_state = dm_state >> self.meas_ops
-            z = [[x] * q_state.representation.ansatz.num_vars for x in self.meas_outcomes]
+            z = [x * 2 for x in self.meas_outcomes]
             probs = q_state.representation(z) * math.sqrt(settings.HBAR)
-            return self._validate_probs(probs, self._step, atol)
+            return self._validate_probs(probs, self._step**2, atol)
         return self.prob_dist
diff --git a/tests/test_lab_dev/test_sampler.py b/tests/test_lab_dev/test_sampler.py
@@ -112,7 +112,9 @@ class TestHomodyneSampler:
     def test_init(self):
         sampler = HomodyneSampler([0, 1], bounds=(-5, 5), num=100)
         assert sampler.meas_ops == BtoQ([0, 1])
-        assert math.allclose(sampler.meas_outcomes, np.linspace(-5, 5, 100))
+        assert math.allclose(
+            sampler.meas_outcomes, list(product(np.linspace(-5, 5, 100), repeat=2))
+        )
         assert sampler.prob_dist is None
 
     def test_probabilties(self):