Merge branch 'develop' into fix_cft

mrmustard/lab_dev/transformations/base.py

@@ -32,7 +32,7 @@
 from mrmustard.physics.ansatz import PolyExpAnsatz, ArrayAnsatz
 from mrmustard.physics.representations import Representation
 from mrmustard.physics.wires import Wires
-from mrmustard.utils.typing import ComplexMatrix, RealMatrix, Vector
+from mrmustard.utils.typing import ComplexTensor, ComplexMatrix, RealMatrix, Vector
 from mrmustard.physics.triples import XY_to_channel_Abc
 from mrmustard.physics.bargmann_utils import au2Symplectic, symplectic2Au, XY_of_channel
 from ..circuit_components import CircuitComponent
@@ -259,6 +259,20 @@ def from_symplectic(cls, modes, S) -> Unitary:
         c = ((-1) ** m * math.det(A_inin @ math.conj(A_inin) - math.eye_like(A_inin))) ** 0.25
         return Unitary.from_bargmann(modes, modes, [A, b, c])
 
+    @classmethod
+    def from_fock(
+        cls,
+        modes_out: Sequence[int],
+        modes_in: Sequence[int],
+        array: ComplexTensor,
+        batched: bool = False,
+        name: str | None = None,
+    ) -> Unitary:
+        r"""
+        Allows initialization of unitaries from their Fock representation.
+        """
+        return Unitary.from_ansatz(modes_in, modes_out, ArrayAnsatz(array, batched), name)
+
     @classmethod
     def random(cls, modes, max_r=1):
         r"""
@@ -478,6 +492,20 @@ def from_XY(
 
         return Channel.from_bargmann(modes_out, modes_in, XY_to_channel_Abc(X, Y, d))
 
+    @classmethod
+    def from_fock(
+        cls,
+        modes_out: Sequence[int],
+        modes_in: Sequence[int],
+        array: ComplexTensor,
+        batched: bool = False,
+        name: str | None = None,
+    ) -> Channel:
+        r"""
+        Allows initialization of unitaries from their Fock representation.
+        """
+        return Channel.from_ansatz(modes_in, modes_out, ArrayAnsatz(array, batched), name)
+
     @classmethod
     def random(cls, modes: Sequence[int], max_r: float = 1.0) -> Channel:
         r"""

mrmustard/physics/ansatz/polyexp_ansatz.py

@@ -750,6 +750,8 @@ def __call__(self, z: Batch[Vector]) -> Scalar | PolyExpAnsatz:
             return self._call_all(z)
 
     def __eq__(self, other: PolyExpAnsatz) -> bool:
+        if not isinstance(other, PolyExpAnsatz):
+            return False
         return self._equal_no_array(other) and np.allclose(self.c, other.c, atol=1e-10)
 
     def __getitem__(self, idx: int | tuple[int, ...]) -> PolyExpAnsatz:

tests/test_lab_dev/test_transformations/test_transformations_base.py

@@ -21,7 +21,7 @@
 
 from mrmustard import math
 from mrmustard.lab_dev.circuit_components import CircuitComponent
-from mrmustard.lab_dev.states import Vacuum
+from mrmustard.lab_dev.states import Vacuum, Coherent
 from mrmustard.lab_dev.transformations import (
     Attenuator,
     Channel,
@@ -31,6 +31,7 @@
     Operation,
     Sgate,
     Unitary,
+    PhaseNoise,
 )
 from mrmustard.physics.wires import Wires
 
@@ -100,6 +101,13 @@ def test_init_from_symplectic(self):
         assert u >> u.dual == Identity([0, 1])
         assert u.dual >> u == Identity([0, 1])
 
+    def test_init_from_fock(self):
+        cutoff = 100
+        eigs = [math.exp(1j * n**2) for n in range(cutoff)]
+        kerr = Unitary.from_fock([0], [0], math.diag(math.astensor(eigs)))
+
+        assert math.allclose((kerr >> kerr.dual).fock_array(), math.eye(cutoff))
+
     def test_inverse_unitary(self):
         gate = Sgate([0], 0.1, 0.2) >> Dgate([0], 0.1, 0.2)
         gate_inv = gate.inverse()
@@ -212,3 +220,20 @@ def test_from_XY(self, nmodes):
         x, y = Channel.from_XY(range(nmodes), range(nmodes), X, Y).XY
         assert math.allclose(x, X)
         assert math.allclose(y, Y)
+
+    def test_from_fock(self):
+        # Here we test our from_fock method by a PhaseNoise example
+        cutoff = 20
+        ph_n = np.zeros((cutoff, cutoff, cutoff, cutoff))
+        sigma = 1
+
+        for m in range(cutoff):
+            for n in range(cutoff):
+                ph_n[m, m, n, n] = math.exp(-0.5 * (m - n) ** 2 * sigma**2)
+
+        phi = Channel.from_fock([0], [0], ph_n)
+        psi = Coherent([0], 2) >> phi
+
+        assert psi.to_fock() == (Coherent([0], 2) >> PhaseNoise([0], sigma)).to_fock(
+            (cutoff, cutoff)
+        )