diff --git a/mrmustard/lab/abstract/state.py b/mrmustard/lab/abstract/state.py index 25b2ab177..faf311790 100644 --- a/mrmustard/lab/abstract/state.py +++ b/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([]) 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/mrmustard/math/backend_tensorflow.py b/mrmustard/math/backend_tensorflow.py index dde3268ea..636b7459a 100644 --- a/mrmustard/math/backend_tensorflow.py +++ b/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 ( diff --git a/mrmustard/training/callbacks.py b/mrmustard/training/callbacks.py index 63a0bfead..16f8e7e9d 100644 --- a/mrmustard/training/callbacks.py +++ b/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 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