An initial implementation of the phase noise channel

mrmustard/lab_dev/states/base.py

@@ -42,6 +42,7 @@
 )
 from mrmustard.physics.fock_utils import quadrature_distribution
 from mrmustard.physics.wigner import wigner_discretized
+from ..circuit_components_utils import BtoPS
 from mrmustard.utils.typing import (
     ComplexMatrix,
     ComplexTensor,
@@ -50,7 +51,7 @@
 )
 
 from ..circuit_components import CircuitComponent
-from ..circuit_components_utils import BtoPS
+
 
 __all__ = ["State"]
 
@@ -333,6 +334,7 @@ def phase_space(self, s: float) -> tuple:
             Returns:
                 The covariance matrix, the mean vector and the coefficient of the state in s-parametrized phase space.
         """
+
         if not isinstance(self.ansatz, PolyExpAnsatz):
             raise ValueError("Can calculate phase space only for Bargmann states.")
 

mrmustard/lab_dev/states/dm.py

@@ -32,10 +32,11 @@
 from mrmustard.physics.representations import Representation
 from mrmustard.physics.wires import Wires
 from mrmustard.utils.typing import ComplexMatrix, ComplexVector, ComplexTensor, RealVector
+from ..circuit_components_utils import BtoQ, TraceOut
 
 from .base import State, _validate_operator, OperatorType
 from ..circuit_components import CircuitComponent
-from ..circuit_components_utils import BtoQ, TraceOut
+
 from ..utils import shape_check
 
 __all__ = ["DM"]
@@ -196,6 +197,7 @@ def from_quadrature(
             ValueError: If the given triple has shapes that are inconsistent
                 with the number of modes.
         """
+
         QtoB = BtoQ(modes, phi).inverse()
         Q = DM.from_ansatz(modes, PolyExpAnsatz(*triple))
         return DM.from_ansatz(modes, (Q >> QtoB).ansatz, name)
@@ -302,6 +304,7 @@ def expectation(self, operator: CircuitComponent):
             ValueError: If ``operator`` is defined over a set of modes that is not a subset of the
                 modes of this state.
         """
+
         op_type, msg = _validate_operator(operator)
         if op_type is OperatorType.INVALID_TYPE:
             raise ValueError(msg)
@@ -413,6 +416,7 @@ def __rshift__(self, other: CircuitComponent) -> CircuitComponent:
         Returns a ``DM`` when the wires of the resulting components are compatible with
         those of a ``DM``, a ``CircuitComponent`` otherwise, and a scalar if there are no wires left.
         """
+
         result = super().__rshift__(other)
         if not isinstance(result, CircuitComponent):
             return result  # scalar case handled here

mrmustard/lab_dev/states/ket.py

@@ -31,6 +31,7 @@
 from mrmustard.physics.gaussian import purity
 from mrmustard.physics.representations import Representation
 from mrmustard.physics.wires import Wires
+from ..circuit_components_utils import BtoQ, TraceOut
 from mrmustard.utils.typing import (
     ComplexMatrix,
     ComplexVector,
@@ -43,7 +44,6 @@
 from .base import State, _validate_operator, OperatorType
 from .dm import DM
 from ..circuit_components import CircuitComponent
-from ..circuit_components_utils import BtoQ, TraceOut
 from ..utils import shape_check
 
 __all__ = ["Ket"]
@@ -152,6 +152,7 @@ def from_quadrature(
         phi: float = 0.0,
         name: str | None = None,
     ) -> State:
+
         QtoB = BtoQ(modes, phi).inverse()
         Q = Ket.from_ansatz(modes, PolyExpAnsatz(*triple))
         return Ket.from_ansatz(modes, (Q >> QtoB).ansatz, name)
@@ -263,6 +264,7 @@ def expectation(self, operator: CircuitComponent):
             ValueError: If ``operator`` is defined over a set of modes that is not a subset of the
                 modes of this state.
         """
+
         op_type, msg = _validate_operator(operator)
         if op_type is OperatorType.INVALID_TYPE:
             raise ValueError(msg)
@@ -371,6 +373,10 @@ def __rshift__(self, other: CircuitComponent | Scalar) -> CircuitComponent | Bat
         with those of a ``DM`` or of a ``Ket``. Returns a ``CircuitComponent`` in general,
         and a (batched) scalar if there are no wires left, for convenience.
         """
+        from ..transformations.phasenoise import PhaseNoise
+
+        if isinstance(other, PhaseNoise):
+            return self.dm() >> other
         result = super().__rshift__(other)
         if not isinstance(result, CircuitComponent):
             return result  # scalar case handled here

mrmustard/lab_dev/transformations/__init__.py

@@ -26,6 +26,7 @@
 from .ggate import *
 from .gaussrandnoise import *
 from .identity import *
+from .phasenoise import *
 from .rgate import *
 from .s2gate import *
 from .sgate import *

mrmustard/lab_dev/transformations/phasenoise.py

@@ -0,0 +1,89 @@
+# Copyright 2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+The class representing a Phase noise channel.
+"""
+
+from __future__ import annotations
+from typing import Sequence
+from mrmustard import math
+from mrmustard.physics.representations import Representation
+from mrmustard.physics.ansatz.array_ansatz import ArrayAnsatz
+from mrmustard.lab_dev.circuit_components import CircuitComponent
+from .base import Channel
+from ..utils import make_parameter
+import numpy as np
+
+__all__ = ["PhaseNoise"]
+
+
+class PhaseNoise(Channel):
+    r"""
+    The Phase noise channel.
+
+    This class represents the application of a random phase. The distributiuon of the phase
+    is assumed to be a Gaussian with mean zero, and standard deviation `phase_stdev`.
+
+    Args:
+        modes: The modes the channel is applied to
+        phase_stdev: The standard deviation of the random Gaussian noise.```
+
+    ..details::
+        The Fock representation is connected to the Fourier coefficients of the distribution.
+    """
+
+    short_name = "P~"
+
+    def __init__(
+        self,
+        modes: Sequence[int],
+        phase_stdev: float | Sequence[float],
+        phase_stdev_trainable: bool = False,
+        phase_stdev_bounds: tuple[float | None, float | None] = (0.0, None),
+    ):
+        super().__init__(name="PhaseNoise")
+        self._add_parameter(
+            make_parameter(phase_stdev_trainable, phase_stdev, "phase_stdev", phase_stdev_bounds)
+        )
+        self._representation = self.from_ansatz(
+            modes_in=modes, modes_out=modes, ansatz=None
+        ).representation
+
+    def __custom_rrshift__(self, other: CircuitComponent) -> CircuitComponent:
+        r"""
+        Since PhaseNoise admits a particularly nice form in the Fock basis, we have implemented its right-shift operation separately.
+
+        Args:
+            other: the component other than the PhaseNoise object that is present in the contraction
+
+        Output:
+            the result of the contraction.
+        """
+
+        if (len(other.wires.bra) == 0) or (len(other.wires.ket) == 0):
+            return (other @ other.adjoint) >> self
+        array = math.asnumpy(other.fock_array())
+        for mode in self.modes:
+            for count, _ in enumerate(np.nditer(array)):
+                idx = np.zeros(len(array.shape))
+                temp = count
+                for l in range(len(idx)):
+                    idx[-1 - l] = temp % array.shape[-1 - l]
+                    temp = temp // array.shape[-1 - l]
+                array_index = tuple(idx.astype(int))
+                array[array_index] *= np.exp(
+                    -0.5 * (idx[mode] - idx[other.n_modes + mode]) ** 2 * self.phase_stdev.value**2
+                )
+        return CircuitComponent(Representation(ArrayAnsatz(array, False), other.wires), self.name)

tests/test_lab_dev/test_transformations/test_phasenoise.py

@@ -0,0 +1,44 @@
+# Copyright 2024 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the ``PhaseNoise`` class."""
+
+# pylint: disable=missing-function-docstring, expression-not-assigned
+
+from mrmustard.lab_dev.states import DM, Ket
+from mrmustard.lab_dev.transformations import Dgate, PhaseNoise
+
+
+class TestPhaseNoise:
+    r"""
+    Tests for the ``PhaseNoise`` class.
+    """
+
+    def test_init(self):
+        "Tests the PhaseNoise initialization."
+        ch = PhaseNoise([0, 1], 0.2)
+        assert ch.name == "PhaseNoise"
+        assert ch.phase_stdev.value == 0.2
+        assert ch.modes == [0, 1]
+        assert ch.ansatz == None
+
+    def test_application(self):
+        "Tests application of PhaseNoise on Ket and DM"
+        psi = Ket.random([0, 1]) >> Dgate([0], 0.5, 0.5) >> PhaseNoise([0], 0.2)
+        assert isinstance(psi, DM)
+        assert psi.purity < 1
+
+        rho = DM.random([0, 1]) >> Dgate([0], 0.5, 0.5) >> PhaseNoise([0], 0.2)
+        assert isinstance(rho, DM)
+        assert rho.purity < 1