Merge branch 'develop' of https://github.com/XanaduAI/MrMustard into …

…wires-final

mrmustard/lab/abstract/state.py

@@ -760,9 +760,10 @@ def mikkel_plot(
     plt.subplots_adjust(wspace=0.05, hspace=0.05)
 
     # Wigner function
+
     ax[1][0].contourf(X, P, W, 120, cmap=plot_args["cmap"], vmin=-abs(W).max(), vmax=abs(W).max())
-    ax[1][0].set_xlabel("$x$", fontsize=12)
-    ax[1][0].set_ylabel("$p$", fontsize=12)
+    ax[1][0].set_xlabel("x", fontsize=12)
+    ax[1][0].set_ylabel("p", fontsize=12)
     ax[1][0].get_xaxis().set_ticks(plot_args["xticks"])
     ax[1][0].xaxis.set_ticklabels(plot_args["xtick_labels"])
     ax[1][0].get_yaxis().set_ticks(plot_args["yticks"])
@@ -779,7 +780,7 @@ def mikkel_plot(
     ax[0][0].xaxis.set_ticklabels([])
     ax[0][0].get_yaxis().set_ticks([])
     ax[0][0].tick_params(direction="in")
-    ax[0][0].set_ylabel("Prob($x$)", fontsize=12)
+    ax[0][0].set_ylabel("Prob(x)", fontsize=12)
     ax[0][0].set_xlim(xbounds)
     ax[0][0].set_ylim([0, 1.1 * max(ProbX)])
     ax[0][0].grid(plot_args["grid"])
@@ -791,14 +792,14 @@ def mikkel_plot(
     ax[1][1].get_yaxis().set_ticks(plot_args["yticks"])
     ax[1][1].yaxis.set_ticklabels([])
     ax[1][1].tick_params(direction="in")
-    ax[1][1].set_xlabel("Prob($p$)", fontsize=12)
+    ax[1][1].set_xlabel("Prob(p)", fontsize=12)
     ax[1][1].set_xlim([0, 1.1 * max(ProbP)])
     ax[1][1].set_ylim(ybounds)
     ax[1][1].grid(plot_args["grid"])
 
     # Density matrix
     ax[0][1].matshow(abs(rho), cmap=plot_args["cmap"], vmin=-abs(rho).max(), vmax=abs(rho).max())
-    ax[0][1].set_title(r"abs($\rho$)", fontsize=12)
+    ax[0][1].set_title("abs(ρ)", fontsize=12)
     ax[0][1].tick_params(direction="in")
     ax[0][1].get_xaxis().set_ticks([])
     ax[0][1].get_yaxis().set_ticks([])

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)
 

mrmustard/math/backend_tensorflow.py

@@ -18,14 +18,14 @@
 
 from typing import Callable, List, Optional, Sequence, Tuple, Union
 
-import logging
+import os
 import numpy as np
+import tensorflow_probability as tfp
 
-logging.getLogger("tensorflow").setLevel(logging.ERROR)
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
 import tensorflow as tf
-import tensorflow_probability as tfp
 
-logging.getLogger("tensorflow").setLevel(logging.INFO)
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"
 
 
 from mrmustard.math.lattice.strategies.compactFock.inputValidation import (

mrmustard/training/callbacks.py

@@ -76,13 +76,13 @@ def rolling_cost_cb(optimizer, cost, **kwargs):
 from pathlib import Path
 from typing import Callable, Optional, Mapping, Sequence, Union
 
-import logging
+import os
 import numpy as np
 
-logging.getLogger("tensorflow").setLevel(logging.ERROR)
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
 import tensorflow as tf
 
-logging.getLogger("tensorflow").setLevel(logging.INFO)
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "0"
 
 
 @dataclass

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