Fix math.Categorical (#342)

**Context:**
```
probs = np.array([1e-6 for _ in range(300)])
results = [math.Categorical(probs, "") for _ in range(100)]
assert len(set(results)) > 1
```
When using the numpy backend, the code above fails: `results` contains
100 identical values.
This is because `math.Categorical` uses `np.random.multinomial`, which
assumes that the probabilities sum up to `1` (if not, the last
probability is increased such that the resulting probabilities sum up to
`1`)

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

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)
 

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))

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