diff --git a/strawberryfields/backends/gaussianbackend/backend.py b/strawberryfields/backends/gaussianbackend/backend.py index f113aa68c..324ead7fb 100644 --- a/strawberryfields/backends/gaussianbackend/backend.py +++ b/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 diff --git a/tests/backend/test_fock_measurement.py b/tests/backend/test_fock_measurement.py index 2080836c8..367e01f4e 100644 --- a/tests/backend/test_fock_measurement.py +++ b/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")