Add cutoff keyword argument to measure_fock in gaussian backend

strawberryfields/backends/gaussianbackend/backend.py

@@ -15,7 +15,17 @@
 """Gaussian backend"""
 import warnings
 
-from numpy import empty, concatenate, array, identity, sqrt, vstack, zeros_like, allclose, ix_
+from numpy import (
+ empty,
+ concatenate,
+ array,
+ identity,
+ sqrt,
+ vstack,
+ zeros_like,
+ allclose,
+ ix_,
+)
 from thewalrus.samples import hafnian_sample_state, torontonian_sample_state
 from thewalrus.symplectic import xxpp_to_xpxp
 
@@ -132,7 +142,8 @@ def measure_homodyne(self, phi, mode, shots=1, select=None, **kwargs):
  )
 
  raise NotImplementedError(
- "Gaussian backend currently does not support " "shots != 1 for homodyne measurement"
+ "Gaussian backend currently does not support "
+ "shots != 1 for homodyne measurement"
  )
 
  # phi is the rotation of the measurement operator, hence the minus
@@ -149,7 +160,6 @@ def measure_homodyne(self, phi, mode, shots=1, select=None, **kwargs):
  return array([[qs * sqrt(2 * self.circuit.hbar) / 2]])
 
  def measure_heterodyne(self, mode, shots=1, select=None, **kwargs):
-
  if shots != 1:
  if select is not None:
  raise NotImplementedError(
@@ -180,7 +190,9 @@ def prepare_gaussian_state(self, r, V, modes):
  # make sure number of modes matches shape of r and V
  N = len(modes)
  if len(r) != 2 * N:
- raise ValueError("Length of means vector must be twice the number of modes.")
+ raise ValueError(
+ "Length of means vector must be twice the number of modes."
+ )
  if V.shape != (2 * N, 2 * N):
  raise ValueError(
  "Shape of covariance matrix must be [2N, 2N], where N is the number of modes."
@@ -231,6 +243,14 @@ def measure_fock(self, modes, shots=1, select=None, **kwargs):
  # check we are sampling from a gaussian state with zero mean
  if allclose(mu, zeros_like(mu)):
  samples = hafnian_sample_state(reduced_cov, shots)
+ elif "cutoff_dim" in kwargs:
+ cutoff = kwargs.get("cutoff_dim")
+ if not isinstance(cutoff, int):
+ raise ValueError(f"cutoff must be an integer, got {type(cutoff)}.")
+ else:
+ samples = hafnian_sample_state(
+ reduced_cov, shots, cutoff=kwargs.get("cutoff_dim")
+ )
  else:
  samples = hafnian_sample_state(reduced_cov, shots, mean=reduced_mean)
 
@@ -257,7 +277,9 @@ def measure_threshold(self, modes, shots=1, select=None, **kwargs):
  modes_idxs = concatenate([x_idxs, p_idxs])
  reduced_cov = cov[ix_(modes_idxs, modes_idxs)]
  reduced_mean = mean[modes_idxs]
- samples = torontonian_sample_state(mu=reduced_mean, cov=reduced_cov, samples=shots)
+ samples = torontonian_sample_state(
+ mu=reduced_mean, cov=reduced_cov, samples=shots
+ )
 
  return samples
 

tests/backend/test_fock_measurement.py

@@ -37,6 +37,45 @@ def measure_fock_gaussian_warning(self, setup_backend):
  "Fock measurement has not been updated.",
  ):
  backend.measure_fock([0, 1], shots=5)
+
+ def test_fock_measurements(self, setup_backend, cutoff, batch_size, pure, tol):
+ """Tests if Fock measurements results on a variety of multi-mode Fock states are correct."""
+ state_preps = [n for n in range(cutoff)] + [
+ cutoff - n for n in range(cutoff)
+ ] # [0, 1, 2, ..., cutoff-1, cutoff, cutoff-1, ..., 2, 1]
+
+ singletons = [(0,), (1,), (2,)]
+ pairs = [(0, 1), (0, 2), (1, 2)]
+ triples = [(0, 1, 2)]
+ mode_choices = singletons + pairs + triples
+
+ backend = setup_backend(3)
+
+ for idx in range(NUM_REPEATS):
+ backend.reset(pure=pure)
+
+ n = [
+ state_preps[idx % cutoff],
+ state_preps[(idx + 1) % cutoff],
+ state_preps[(idx + 2) % cutoff],
+ ]
+ n = np.array(n)
+ meas_modes = np.array(
+ mode_choices[idx % len(mode_choices)]
+ ) # cycle through mode choices
+
+ backend.prepare_fock_state(n[0], 0)
+ backend.prepare_fock_state(n[1], 1)
+ backend.prepare_fock_state(n[2], 2)
+
+ meas_result = backend.measure_fock(meas_modes)
+ ref_result = n[meas_modes]
+
+ if batch_size is not None:
+ ref_result = np.array([[i] * batch_size for i in ref_result]).T.reshape(
+ (batch_size, 1, ref_result.shape[0])
+ )
+ assert np.allclose(meas_result, ref_result, atol=tol, rtol=0)
 
 
 @pytest.mark.backends("fock", "tf")