tests

mrmustard/physics/bargmann.py

@@ -133,7 +133,10 @@ def complex_gaussian_integral(
         measure: the exponent of the measure (default is -1: Bargmann measure)
 
     Returns:
-        The ``(A,b,c)`` triple of the result of the integral
+        The ``(A,b,c)`` triple of the result of the integral.
+
+    Raises:
+        ValueError: If ``idx_z`` and ``idx_zconj`` have different lengths.
     """
     A, b, c = Abc
     if len(idx_z) != len(idx_zconj):

mrmustard/physics/representations.py

@@ -257,9 +257,6 @@ def trace(self, idx_z: tuple[int, ...], idx_zconj: tuple[int, ...]) -> Bargmann:
             raise NotImplementedError(
                 "Partial trace is only supported for ansatze with polynomial of degree ``0``."
             )
-        if len(idx_z) != len(idx_zconj):
-            msg = f"The number of indices to trace over must be the same for ``z`` and ``z*`` (got {len(idx_z)} and {len(idx_zconj)})."
-            raise ValueError(msg)
         A, b, c = [], [], []
         for Abci in zip(self.A, self.b, self.c):
             Aij, bij, cij = bargmann.complex_gaussian_integral(Abci, idx_z, idx_zconj, measure=-1.0)

mrmustard/physics/triples.py

@@ -95,7 +95,7 @@ def vacuum_state_Abc(n_modes: int) -> Union[Matrix, Vector, Scalar]:
     """
     A = _vacuum_A_matrix(n_modes)
     b = _vacuum_B_vector(n_modes)
-    c = 1.0 + 0.j
+    c = 1.0 + 0.0j
 
     return A, b, c
 
@@ -246,7 +246,7 @@ def rotation_gate_Abc(theta: Union[Scalar, Iterable]) -> Union[Matrix, Vector, S
     A = math.astensor([[0, 1], [1, 0]], math.complex128)
     A = np.kron(A, math.exp(1j * theta) * math.eye(n_modes, math.complex128))
     b = _vacuum_B_vector(n_modes)
-    c = 1.0 + 0.j
+    c = 1.0 + 0.0j
 
     return A, b, c
 

tests/random.py

@@ -82,10 +82,14 @@ def Abc_triple(n: int):
     A = 0.5 * (A + A.T)  # make it symmetric
 
     # complex vector b
-    b = np.random.uniform(min_magnitude, max_magnitude, (n,))
+    b = np.random.uniform(min_magnitude, max_magnitude, (n,)) + 1.0j * np.random.uniform(
+        min_magnitude, max_magnitude, (n,)
+    )
 
     # complex scalar c
-    c = np.random.uniform(min_magnitude, max_magnitude, (1,))
+    c = np.random.uniform(min_magnitude, max_magnitude, (1,)) + 1.0j * np.random.uniform(
+        min_magnitude, max_magnitude, (1,)
+    )
 
     return A, b, c
 

tests/test_physics/test_bargmann.py

@@ -75,6 +75,7 @@ def test_bargmann_numpy_transformation():
     transformation = Ggate(1)
     assert all(isinstance(thing, np.ndarray) for thing in transformation.bargmann(numpy=True))
 
+
 def test_join_Abc():
     """Tests the ``join_Abc`` method."""
     A1, b1, c1 = triples.vacuum_state_Abc(2)
@@ -90,6 +91,7 @@ def test_join_Abc():
     c12 = math.outer(c1, c2)
     return A12, b12, c12
 
+
 def test_complex_gaussian_integral():
     """Tests the ``complex_gaussian_integral`` method."""
     A1, b1, c1 = triples.vacuum_state_Abc(2)
@@ -113,13 +115,21 @@ def test_complex_gaussian_integral():
     assert np.allclose(res3[1], 0)
     assert np.allclose(res3[2], 1)
 
+
 def test_complex_gaussian_integral_error():
     """Tests the error of the ``complex_gaussian_integral`` method."""
     A1, b1, c1 = triples.vacuum_state_Abc(2)
     A2, b2, c2 = triples.displacement_gate_Abc(x=[0.1, 0.2], y=0.3)
 
     with pytest.raises(ValueError, match="idx_z and idx_zconj must have the same length"):
-        complex_gaussian_integral(join_Abc((A1, b1, c1), (A2, b2, c2)), [0, 1], [4,])
+        complex_gaussian_integral(
+            join_Abc((A1, b1, c1), (A2, b2, c2)),
+            [0, 1],
+            [
+                4,
+            ],
+        )
+
 
 def test_contract_two_Abc():
     """Tests the error of the ``contract_two_Abc`` method."""
@@ -128,15 +138,15 @@ def test_contract_two_Abc():
 
     res1 = contract_two_Abc((A1, b1, c1), (A2, b2, c2), [], [])
     assert np.allclose(res1[0], math.block_diag(A1, A2))
-    assert np.allclose(res1[1], [0, 0, 0.1+0.3j, 0.2+0.3j, -0.1+0.3j, -0.2+0.3j])
+    assert np.allclose(res1[1], [0, 0, 0.1 + 0.3j, 0.2 + 0.3j, -0.1 + 0.3j, -0.2 + 0.3j])
     assert np.allclose(res1[2], c1 * c2)
 
     res2 = contract_two_Abc((A1, b1, c1), (A2, b2, c2), [0, 1], [2, 3])
     assert np.allclose(res2[0], math.zeros((2, 2)))
-    assert np.allclose(res2[1], [0.1+0.3j, 0.2+0.3j])
+    assert np.allclose(res2[1], [0.1 + 0.3j, 0.2 + 0.3j])
     assert np.allclose(res2[2], c1 * c2)
 
     res3 = contract_two_Abc((A1, b1, c1), (A2, b2, c2), [0, 1], [0, 1])
     assert np.allclose(res3[0], math.zeros((2, 2)))
-    assert np.allclose(res3[1], [-0.1+0.3j, -0.2+0.3j])
+    assert np.allclose(res3[1], [-0.1 + 0.3j, -0.2 + 0.3j])
     assert np.allclose(res3[2], c1 * c2)

tests/test_physics/test_representations.py

@@ -18,7 +18,7 @@
 import pytest
 
 from mrmustard import math
-from mrmustard.physics.bargmann import contract_two_Abc
+from mrmustard.physics.bargmann import contract_two_Abc, complex_gaussian_integral
 from mrmustard.physics.representations import Bargmann, Fock
 from ..random import Abc_triple
 
@@ -142,12 +142,14 @@ def test_sub(self, n):
         assert np.allclose(bargmann_add.b, math.concat([bargmann1.b, bargmann2.b], axis=0))
         assert np.allclose(bargmann_add.c, math.concat([bargmann1.c, -bargmann2.c], axis=0))
 
-    # @ Filippo??
-    # def test_trace(self):
-    #     bargmann = Bargmann(*Abc_triple(4)).trace([0], [2])
-    #     assert np.allclose(bargmann.A.shape, (1, 2, 2))
-    #     assert np.allclose(bargmann.b.shape, (1, 2))
-    #     assert np.allclose(bargmann.c.shape, (1,))
+    def test_trace(self):
+        triple = Abc_triple(4)
+        bargmann = Bargmann(*triple).trace([0], [2])
+        A, b, c = complex_gaussian_integral(triple, [0], [2])
+
+        assert np.allclose(bargmann.A, A)
+        assert np.allclose(bargmann.b, b)
+        assert np.allclose(bargmann.c, c)
 
     def test_reorder(self):
         triple = Abc_triple(3)
@@ -171,7 +173,6 @@ def test_matmul(self):
         assert np.allclose(res1.A, exp1[0])
         assert np.allclose(res1.b, exp1[1])
         assert np.allclose(res1.c, exp1[2])
-
 
 
 class TestFockRepresentation: