added Gaussian Ket and DM. Tests to be added

mrmustard/lab_dev/states/__init__.py

@@ -23,6 +23,7 @@
 from .bargmann_eigenstate import BargmannEigenstate
 from .coherent import Coherent
 from .displaced_squeezed import DisplacedSqueezed
+from .gstate import Gket, Gdm
 from .number import Number
 from .quadrature_eigenstate import QuadratureEigenstate
 from .squeezed_vacuum import SqueezedVacuum

mrmustard/lab_dev/states/gstate.py

@@ -0,0 +1,116 @@
+# 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.
+
+"""
+Classes that are initialized by a symplectic matrix and possibly thermal states if we have a DM.
+"""
+
+from __future__ import annotations
+
+from typing import Sequence
+
+from mrmustard import math, settings
+from mrmustard.math.parameters import update_symplectic
+from mrmustard.physics.ansatz import PolyExpAnsatz
+from mrmustard.physics import triples
+from mrmustard.utils.typing import RealMatrix, Vector
+from .ket import Ket
+from .dm import DM
+from ..utils import make_parameter, reshape_params
+
+__all__ = ["Gket", "Gdm"]
+
+
+class Gket(Ket):
+    r"""
+    The `N`-mode pure state described by a Gaussian gate that acts on Vacuum.
+    """
+
+    short_name = "Gk"
+
+    def __init__(
+        self,
+        modes: Sequence[int],
+        symplectic: RealMatrix = None,
+        symplectic_trainable: bool = False,
+    ) -> None:
+        super().__init__(name="Gket")
+        m = len(modes)
+        if symplectic is None:
+            symplectic = math.random_symplectic(m)
+
+        self._add_parameter(
+            make_parameter(
+                symplectic_trainable,
+                symplectic,
+                "symplectic",
+                (None, None),
+                update_symplectic,
+            )
+        )
+
+        self._representation = self.from_ansatz(
+            modes=modes,
+            ansatz=PolyExpAnsatz.from_function(
+                fn=triples.gket_state_Abc, symplectic=self.symplectic
+            ),
+        ).representation
+
+
+class Gdm(DM):
+    r"""
+    The `N`-mode mixed state described by a Gaussian gate that acts on a given thermal state.
+    """
+
+    short_name = "Gd"
+
+    def __init__(
+        self,
+        modes: Sequence[int],
+        betas: float | Sequence[float],
+        symplectic: RealMatrix = None,
+        symplectic_trainable: bool = False,
+        betas_trainable: bool = False,
+    ) -> None:
+        super().__init__(name="Gdm")
+        m = len(modes)
+
+        if symplectic is None:
+            symplectic = math.random_symplectic(m)
+
+        self._add_parameter(
+            make_parameter(
+                symplectic_trainable,
+                symplectic,
+                "symplectic",
+                (None, None),
+                update_symplectic,
+            )
+        )
+
+        self._add_parameter(
+            make_parameter(
+                betas_trainable,
+                math.astensor(betas),
+                "betas",
+                (0, None),
+            )
+        )
+
+        self._representation = self.from_ansatz(
+            modes=modes,
+            ansatz=PolyExpAnsatz.from_function(
+                fn=triples.gdm_state_Abc, betas=self.betas, symplectic=self.symplectic
+            ),
+        ).representation

mrmustard/lab_dev/states/ket.py

@@ -172,7 +172,7 @@ def random(cls, modes: Sequence[int], max_r: float = 1.0) -> Ket:
         S = math.random_symplectic(m, max_r)
         transformation = (
             1
-            / np.sqrt(2)
+            / math.sqrt(complex(2))
             * math.block(
                 [
                     [

mrmustard/physics/triples.py

@@ -24,6 +24,7 @@
 from mrmustard import math, settings
 from mrmustard.utils.typing import Matrix, Vector, Scalar, RealMatrix, ComplexMatrix
 from mrmustard.physics.gaussian_integrals import complex_gaussian_integral_2
+from .bargmann_utils import symplectic2Au
 
 
 #  ~~~~~~~~~
@@ -232,6 +233,78 @@ def two_mode_squeezed_vacuum_state_Abc(
     return A, b, c
 
 
+def gket_state_Abc(symplectic: RealMatrix):
+    r"""
+    The A,b,c parameters of a Gaussian Ket (Gket) state. This is simply a Gaussian acted on the vacuum.
+
+    Args:
+        symplectic: the symplectic representation of the Gaussian
+
+    Returns:
+        The ``(A,b,c)`` triple of the Gket state.
+    """
+
+    m = symplectic.shape[-1] // 2  # num of modes
+
+    transformation = (
+        1
+        / math.sqrt(complex(2))
+        * math.block(
+            [
+                [
+                    math.eye(m, dtype=math.complex128),
+                    math.eye(m, dtype=math.complex128),
+                ],
+                [
+                    -1j * math.eye(m, dtype=math.complex128),
+                    1j * math.eye(m, dtype=math.complex128),
+                ],
+            ]
+        )
+    )
+
+    S = math.conj(math.transpose(transformation)) @ symplectic @ transformation
+    S_1 = S[:m, :m]
+    S_2 = S[:m, m:]
+    A = math.transpose(math.solve(math.dagger(S_1), math.transpose(S_2)))
+    b = math.zeros(m, dtype=A.dtype)
+    c = (math.det(A @ math.conj(A) - math.eye_like(A))) ** 0.25
+    return A, b, c
+
+
+def gdm_state_Abc(betas: Vector, symplectic: RealMatrix):
+    r"""
+    The A,b,c parameters of a Gaussian mixed state that is defined by the action of a Guassian on a thermal state
+
+    Args:
+        betas: the list of betas corresponding to the temperatures of the initial thermal state
+        symplectic: the symplectic matrix of the Gaussian
+
+    Returns:
+        The ``(A,b,c)`` triple of the resulting Gaussian DM state.
+    """
+    betas = math.astensor(betas)
+    m = len(betas)
+    Au = symplectic2Au(symplectic)
+
+    Au00 = Au[:m, :m]
+    Au01 = Au[:m, m:]
+    Au11 = Au[m:, m:]
+
+    D = math.diag(betas)
+
+    # a few auxilarry matrices that help us with computation (consistant naming with Gaussian integrations)
+    M_prime = math.block([[math.conj(Au00), -math.inv(D)], [-math.inv(D), Au00]])
+    D_prime = math.block([[math.conj(Au01), math.zeros((m, m))], [math.zeros((m, m)), Au01]])
+    A_prime = math.block([[math.conj(Au11), math.zeros((m, m))], [math.zeros((m, m)), Au11]])
+
+    A_new = A_prime - D_prime @ math.inv(M_prime) @ D_prime
+    b = math.zeros(2 * m, dtype=A_new.dtype)
+    c = complex(1)  # TODO: update with proper c s.t. tr(rho)=1
+
+    return A_new, b, c
+
+
 def sauron_state_Abc(n: int, epsilon: float):
     r"""
     The A,b,c parametrization of Sauron states. These are Fock states written as a linear superposition of a