diff --git a/mrmustard/math/backend_manager.py b/mrmustard/math/backend_manager.py index be58ebd31..f77dc510b 100644 --- a/mrmustard/math/backend_manager.py +++ b/mrmustard/math/backend_manager.py @@ -1187,9 +1187,8 @@ def Categorical(self, probs: Tensor, name: str): """Categorical distribution over integers. Args: - probs (Tensor): tensor representing the probabilities of a set of Categorical - distributions. - name (str): name prefixed to operations created by this class + probs: The unnormalized probabilities of a set of Categorical distributions. + name: The name prefixed to operations created by this class. Returns: tfp.distributions.Categorical: instance of ``tfp.distributions.Categorical`` class diff --git a/mrmustard/math/backend_numpy.py b/mrmustard/math/backend_numpy.py index dada0eb2d..dc1064342 100644 --- a/mrmustard/math/backend_numpy.py +++ b/mrmustard/math/backend_numpy.py @@ -367,8 +367,8 @@ def __init__(self, probs): self._probs = probs def sample(self): - array = np.random.multinomial(1, pvals=probs) - return np.where(array == 1)[0][0] + idx = [i for i, _ in enumerate(probs)] + return np.random.choice(idx, p=probs / sum(probs)) return Generator(probs) diff --git a/tests/test_lab/test_detectors.py b/tests/test_lab/test_detectors.py index ee0cfdd22..15f32c932 100644 --- a/tests/test_lab/test_detectors.py +++ b/tests/test_lab/test_detectors.py @@ -14,7 +14,6 @@ import numpy as np import pytest -import tensorflow as tf from hypothesis import given from hypothesis import strategies as st from hypothesis.extra.numpy import arrays @@ -297,16 +296,16 @@ def test_homodyne_on_2mode_squeezed_vacuum_with_displacement(self, s, X, d): ], ) @pytest.mark.parametrize("gaussian_state", [True, False]) + @pytest.mark.parametrize("normalization", [1, 1 / 3]) def test_sampling_mean_and_var( - self, state, kwargs, mean_expected, var_expected, gaussian_state + self, state, kwargs, mean_expected, var_expected, gaussian_state, normalization ): """Tests that the mean and variance estimates of many homodyne measurements are in agreement with the expected values for the states""" state = state(**kwargs) - tf.random.set_seed(123) if not gaussian_state: - state = State(dm=state.dm(cutoffs=[40])) + state = State(dm=state.dm(cutoffs=[40]) * normalization) detector = Homodyne(0.0) results = np.zeros((self.N_MEAS, 2)) diff --git a/tests/test_math/test_backend_manager.py b/tests/test_math/test_backend_manager.py index 64935f3c0..44017d63b 100644 --- a/tests/test_math/test_backend_manager.py +++ b/tests/test_math/test_backend_manager.py @@ -587,3 +587,11 @@ def test_sum(self): arr = 4 * np.eye(3) res = math.asnumpy(math.sum(arr)) assert np.allclose(res, 12) + + def test_categorical(self): + r""" + Tests the ``Categorical`` method. + """ + probs = np.array([1e-6 for _ in range(300)]) + results = [math.Categorical(probs, "") for _ in range(100)] + assert len(set(results)) > 1