Stage1 mvp batch dimension

.github/CHANGELOG.md

@@ -3,20 +3,26 @@
 ### New features
 
 * A general factory method and a duck style type-checker.
-  [#256](https://github.com/XanaduAI/MrMustard/pull/256)
+  [(#256)](https://github.com/XanaduAI/MrMustard/pull/256)
+
+* Data classes for the representations project. Includes `WavefunctionArrayData`, `GaussianData`, 
+  `ABCData` and `SymplecticData`.
+  [(#258)](https://github.com/XanaduAI/MrMustard/pull/258)
 
 ### Breaking changes
 
 ### Improvements
 
 ### Bug fixes
 
-* Fixed a bug about the variable names in functions (apply_kraus_to_ket, apply_kraus_to_dm, apply_choi_to_ket, apply_choi_to_dm). [(#271)](https://github.com/XanaduAI/MrMustard/pull/271)
+* Fixed a bug about the variable names in functions (apply_kraus_to_ket, apply_kraus_to_dm, apply_choi_to_ket, apply_choi_to_dm).
+  [(#271)](https://github.com/XanaduAI/MrMustard/pull/271)
 
 ### Documentation
 
 ### Contributors
-[Yuan Yao](https://github.com/sylviemonet), [Richard A. Wolf](https://github.com/ryk-wolf)
+
+[Richard A. Wolf](https://github.com/ryk-wolf), [Filippo Miatto](https://github.com/ziofil), [Yuan Yao](https://github.com/sylviemonet)
 
 
 # Release 0.5.0 (current release)
@@ -105,6 +111,9 @@
 * More robust implementation of cutoffs for States.
 [(#239)](https://github.com/XanaduAI/MrMustard/pull/239)
 
+* Dependencies and versioning are now managed using Poetry.
+[(#257)](https://github.com/XanaduAI/MrMustard/pull/257)
+
 ### Bug fixes
 
 * Fixed a bug that would make two progress bars appear during an optimization
@@ -126,6 +135,7 @@ cutoff of the first detector is equal to 1, the resulting density matrix is now
 [Robbe De Prins](https://github.com/rdprins), [Gabriele Gullì](https://github.com/ggulli),
 [Richard A. Wolf](https://github.com/ryk-wolf)
 
+
 ---
 
 # Release 0.4.1

mrmustard/__init__.py

@@ -43,6 +43,7 @@ def __init__(self):
         self.EQ_TRANSFORMATION_CUTOFF = 3  # 3 is enough to include a full step of the rec relations
         self.EQ_TRANSFORMATION_RTOL_FOCK = 1e-3
         self.EQ_TRANSFORMATION_RTOL_GAUSS = 1e-6
+        self.EQUALITY_PRECISION_DECIMALS = 6  #  decimals for (A,b,c) equality check
         # for the detectors
         self.PNR_INTERNAL_CUTOFF = 50
         self.HOMODYNE_SQUEEZING = 10.0

mrmustard/lab/representations/data/__init__.py

@@ -0,0 +1,23 @@
+# Copyright 2021 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.
+
+""" This package contains the modules implementing base classes for datas.
+"""
+from .gaussian_data import GaussianData
+from .abc_data import ABCData
+from .array_data import ArrayData
+from .symplectic_data import SymplecticData
+from .wavefunctionarray_data import WavefunctionArrayData
+
+__all__ = ["GaussianData", "ABCData", "ArrayData", "SymplecticData", "WavefunctionArrayData"]

mrmustard/lab/representations/data/abc_data.py

@@ -0,0 +1,170 @@
+# Copyright 2023 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.
+
+from __future__ import annotations
+
+from itertools import product
+from typing import Optional, Union
+import numpy as np
+
+from mrmustard.lab.representations.data.matvec_data import MatVecData
+from mrmustard.math import Math
+from mrmustard.typing import Batch, ComplexMatrix, ComplexVector, Scalar
+
+math = Math()
+
+
+class ABCData(MatVecData):
+    r"""Exponential of quadratic polynomial for the Bargmann representation.
+
+    Quadratic Gaussian data is made of: quadratic coefficients, linear coefficients, constant.
+    Each of these has a batch dimension, and the batch dimension is the same for all of them.
+    They are the parameters of the function `c * exp(x^T A x / 2 + x^T b)`.
+
+    Note that if constants are not provided, they will all be initialized at 1.
+
+    Args:
+        A (Batch[Matrix]):          series of quadratic coefficient
+        b (Batch[Vector]):          series of linear coefficients
+        c (Optional[Batch[Scalar]]):series of constants
+    """
+
+    def __init__(
+        self, A: Batch[ComplexMatrix], b: Batch[ComplexVector], c: Optional[Batch[Scalar]] = None
+    ) -> None:
+        super().__init__(mat=A, vec=b, coeffs=c)
+        self._contract_idxs = []
+
+    def value(self, x: ComplexVector) -> Scalar:
+        r"""Value of this function at x.
+
+        Args:
+            x (Vector): point at which the function is evaluated
+
+        Returns:
+            Scalar: value of the function
+        """
+        val = 0.0
+        for A, b, c in zip(self.A, self.b, self.c):
+            val += math.exp(0.5 * math.sum(x * math.matvec(A, x)) + math.sum(x * b)) * c
+        return val
+
+    def __call__(self, x: ComplexVector) -> Scalar:
+        return self.value(x)
+
+    @property
+    def A(self) -> Batch[ComplexMatrix]:
+        return self.mat
+
+    @property
+    def b(self) -> Batch[ComplexVector]:
+        return self.vec
+
+    @property
+    def c(self) -> Batch[Scalar]:
+        return self.coeffs
+
+    def __mul__(self, other: Union[Scalar, ABCData]) -> ABCData:
+        if isinstance(other, ABCData):
+            new_a = [A1 + A2 for A1, A2 in product(self.A, other.A)]
+            new_b = [b1 + b2 for b1, b2 in product(self.b, other.b)]
+            new_c = [c1 * c2 for c1, c2 in product(self.c, other.c)]
+            return self.__class__(A=new_a, b=new_b, c=new_c)
+        else:
+            try:  # scalar
+                return self.__class__(self.A, self.b, other * self.c)
+            except Exception as e:  # Neither same object type nor a scalar case
+                raise TypeError(f"Cannot multiply {self.__class__} and {other.__class__}.") from e
+
+    def __and__(self, other: ABCData) -> ABCData:
+        As = [math.block_diag(a1, a2) for a1 in self.A for a2 in other.A]
+        bs = [math.concat([b1, b2], axis=-1) for b1 in self.b for b2 in other.b]
+        cs = [c1 * c2 for c1 in self.c for c2 in other.c]
+        return self.__class__(As, bs, cs)
+
+    def conj(self):
+        new = self.__class__(math.conj(self.A), math.conj(self.b), math.conj(self.c))
+        new._contract_idxs = self._contract_idxs
+        return new
+
+    def __matmul__(self, other: ABCData) -> ABCData:
+        r"""Implements the contraction of (A,b,c) triples across the marked indices."""
+        # Useful for the future, but not for this PR
+        # raise NotImplementedError()
+
+        graph = self & other
+        newA = graph.A
+        newb = graph.b
+        newc = graph.c
+        for n, (i, j) in enumerate(zip(self._contract_idxs, other._contract_idxs)):
+            i = i - np.sum(np.array(self._contract_idxs[:n]) < i)
+            j = j + self.dim - n - np.sum(np.array(other._contract_idxs[:n]) < j)
+            noij = list(range(i)) + list(range(i + 1, j)) + list(range(j + 1, newA.shape[-1]))
+            Abar = math.gather(math.gather(newA, noij, axis=1), noij, axis=2)
+            bbar = math.gather(newb, noij, axis=1)
+            D = math.gather(
+                math.concat([newA[..., i][..., None], newA[..., j][..., None]], axis=-1),
+                noij,
+                axis=1,
+            )
+            M = math.concat(
+                [
+                    math.concat(
+                        [
+                            newA[:, i, i][:, None, None],
+                            newA[:, j, i][:, None, None] - 1,
+                        ],
+                        axis=-1,
+                    ),
+                    math.concat(
+                        [
+                            newA[:, i, j][:, None, None] - 1,
+                            newA[:, j, j][:, None, None],
+                        ],
+                        axis=-1,
+                    ),
+                ],
+                axis=-2,
+            )
+            Minv = math.inv(M)
+            b_ = math.concat([newb[:, i][:, None], newb[:, j][:, None]], axis=-1)
+
+            newA = Abar - math.einsum("bij,bjk,blk", D, Minv, D)
+            newb = bbar - math.einsum("bij,bjk,bk", D, Minv, b_)
+            newc = (
+                newc
+                * math.exp(-math.einsum("bi,bij,bj", b_, Minv, b_) / 2)
+                / math.sqrt(-math.det(M))
+            )
+        return self.__class__(newA, newb, newc)
+
+    def __getitem__(self, idx: int | tuple[int, ...]) -> ABCData:
+        idx = (idx,) if isinstance(idx, int) else idx
+        for i in idx:
+            if i > self.dim:
+                raise IndexError(
+                    f"Index {i} out of bounds for {self.__class__.__qualname__} of dimension {self.dim}."
+                )
+        new = self.__class__(self.A, self.b, self.c)
+        new._contract_idxs = idx
+        return new
+
+    def transpose(self, order: tuple[int, ...] | list[int]) -> ABCData:
+        new = self.__class__(
+            A=math.gather(math.gather(self.A, order, -1), order, -2),
+            b=math.gather(self.b, order, -1),
+            c=self.c,
+        )
+        new._contract_idxs = self._contract_idxs
+        return new

mrmustard/lab/representations/data/array_data.py

@@ -0,0 +1,71 @@
+# Copyright 2023 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.
+
+from __future__ import annotations
+
+import numpy as np
+
+from mrmustard.lab.representations.data.data import Data
+from mrmustard.math import Math
+from mrmustard.typing import Scalar, Vector
+
+math = Math()
+
+
+class ArrayData(Data):
+    """Contains array-like data for certain Representation objects.
+
+    Args:
+        array (Vector) : data to be contained in the class
+    """
+
+    def __init__(self, array: Vector) -> None:
+        self.array = array
+
+    @property
+    def cutoffs(self):
+        return self.array.shape
+
+    def __neg__(self) -> Data:
+        return self.__class__(array=-self.array)
+
+    def __eq__(self, other: ArrayData) -> bool:
+        try:
+            return np.allclose(self.array, other.array)
+        except AttributeError as e:
+            raise TypeError(f"Cannot compare {self.__class__} and {other.__class__}.") from e
+
+    def __add__(self, other: ArrayData) -> ArrayData:
+        try:
+            return self.__class__(array=self.array + other.array)
+        except AttributeError as e:
+            raise TypeError(f"Cannot add/subtract {self.__class__} and {other.__class__}.") from e
+
+    def __truediv__(self, other: Scalar) -> ArrayData:
+        try:
+            return self.__class__(array=self.array / other)
+        except TypeError as e:
+            raise TypeError("Can only divide by a scalar.") from e
+
+    def __mul__(self, other: Scalar) -> ArrayData:
+        try:
+            return self.__class__(array=self.array * other)
+        except TypeError as e:
+            raise TypeError("Can only multiply by a scalar.") from e
+
+    def __and__(self, other: ArrayData) -> ArrayData:
+        try:
+            return self.__class__(array=np.outer(self.array, other.array))
+        except AttributeError as e:
+            raise TypeError(f"Cannot tensor product {self.__class__} and {other.__class__}.") from e

mrmustard/lab/representations/data/data.py

@@ -0,0 +1,58 @@
+# Copyright 2023 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.
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Union
+
+from mrmustard.typing import Scalar
+
+
+class Data(ABC):
+    r"""Abstract parent class for types of data encoding a quantum state's representation."""
+
+    @abstractmethod
+    def __neg__(self) -> Data:
+        ...
+
+    @abstractmethod
+    def __eq__(self, other: Data) -> bool:
+        ...
+
+    @abstractmethod
+    def __add__(self, other: Data) -> Data:
+        ...
+
+    def __sub__(self, other: Data) -> Data:
+        try:
+            return self.__add__(-other)
+        except AttributeError as e:
+            raise TypeError(f"Cannot subtract {self.__class__} and {other.__class__}.") from e
+
+    # @abstractmethod
+    # def __call__(self, dom: Any) -> Scalar:
+    #     r"""Evaluate the function at a point in the domain."""
+    #     ...
+
+    @abstractmethod
+    def __truediv__(self, other: Union[Scalar, Data]) -> Data:
+        ...
+
+    @abstractmethod
+    def __mul__(self, other: Union[Scalar, Data]) -> Data:
+        ...
+
+    def __rmul__(self, other: Scalar) -> Data:
+        return self.__mul__(other=other)