sampler quad_dist

mrmustard/lab_dev/sampler.py

@@ -188,7 +188,5 @@ def __init__(
         self._step = step
 
     def probabilities(self, state, atol=1e-4):
-        q_state = state.dm() >> BtoQ(state.modes, phi=self.povms[0].phi.value[0])
-        z = [x * 2 for x in product(self.meas_outcomes, repeat=len(state.modes))]
-        probs = q_state.representation(z) * math.sqrt(settings.HBAR)
+        probs = state.quadrature_distribution(self.meas_outcomes, self.povms[0].phi.value[0])
         return self._validate_probs(probs, self._step ** len(state.modes), atol)

mrmustard/lab_dev/states/base.py

@@ -334,7 +334,7 @@ def phase_space(self, s: float) -> tuple:
             *new_state.bargmann_triple(batched=True), batched=True
         )
 
-    def quadrature_distribution(self, quad: Vector, phi: float = 0.0) -> tuple | ComplexTensor:
+    def quadrature_distribution(self, quad: Vector, phi: float = 0.0) -> ComplexTensor:
         r"""
         The (discretized) quadrature distribution of the State.
 
@@ -345,18 +345,19 @@ def quadrature_distribution(self, quad: Vector, phi: float = 0.0) -> tuple | Com
         Returns:
             The quadrature distribution.
         """
+        quad = math.astensor(quad)
         if len(quad.shape) != 1 and len(quad.shape) != self.n_modes:
             raise ValueError(
                 f"The dimensionality of quad should be 1, or match the number of modes."
             )
 
         if len(quad.shape) == 1:
-            quad = math.transpose(math.astensor([quad] * self.n_modes))
+            quad = math.astensor(list(product(quad, repeat=len(self.modes))))
 
         if isinstance(self, Ket):
             return math.abs(self.quadrature(quad, phi)) ** 2
         else:
-            quad = math.tile(math.astensor(quad), (1, 2))
+            quad = math.tile(quad, (1, 2))
             return self.quadrature(quad, phi)
 
     def visualize_2d(