Batching triples.py and symplectics.py

mrmustard/lab_dev/circuit_components_utils/b_to_q.py

@@ -50,7 +50,9 @@ def __init__(
             modes_in=modes,
             modes_out=modes,
             ansatz=PolyExpAnsatz.from_function(
-                fn=triples.bargmann_to_quadrature_Abc, n_modes=len(modes), phi=self.parameters.phi
+                fn=triples.bargmann_to_quadrature_Abc,
+                n_modes=len(modes),
+                phi=self.parameters.phi,
             ),
         ).representation
         for w in self.representation.wires.input.wires:

mrmustard/lab_dev/states/bargmann_eigenstate.py

@@ -62,6 +62,6 @@ def __init__(
         self._representation = self.from_ansatz(
             modes=(mode,),
             ansatz=PolyExpAnsatz.from_function(
-                fn=triples.bargmann_eigenstate_Abc, x=self.parameters.alpha
+                fn=triples.bargmann_eigenstate_Abc, alpha=self.parameters.alpha
             ),
         ).representation

mrmustard/lab_dev/states/base.py

@@ -174,8 +174,8 @@ def from_bargmann(
             >>> from mrmustard.physics.triples import coherent_state_Abc
             >>> from mrmustard.lab_dev.states.ket import Ket
 
-            >>> modes = (0, 1)
-            >>> triple = coherent_state_Abc(x=[0.1, 0.2])  # parallel coherent states
+            >>> modes = (0,)
+            >>> triple = coherent_state_Abc(x=0.1)
 
             >>> coh = Ket.from_bargmann(modes, triple)
             >>> assert coh.modes == modes

mrmustard/lab_dev/states/quadrature_eigenstate.py

@@ -75,7 +75,9 @@ def __init__(
         self._representation = self.from_ansatz(
             modes=(mode,),
             ansatz=PolyExpAnsatz.from_function(
-                fn=triples.quadrature_eigenstates_Abc, x=self.parameters.x, phi=self.parameters.phi
+                fn=triples.quadrature_eigenstates_Abc,
+                x=self.parameters.x,
+                phi=self.parameters.phi,
             ),
         ).representation
 

mrmustard/lab_dev/states/squeezed_vacuum.py

@@ -66,6 +66,8 @@ def __init__(
         self._representation = self.from_ansatz(
             modes=(mode,),
             ansatz=PolyExpAnsatz.from_function(
-                fn=triples.squeezed_vacuum_state_Abc, r=self.parameters.r, phi=self.parameters.phi
+                fn=triples.squeezed_vacuum_state_Abc,
+                r=self.parameters.r,
+                phi=self.parameters.phi,
             ),
         ).representation

mrmustard/lab_dev/transformations/s2gate.py

@@ -81,6 +81,8 @@ def __init__(
             modes_in=modes,
             modes_out=modes,
             ansatz=PolyExpAnsatz.from_function(
-                fn=triples.twomode_squeezing_gate_Abc, r=self.parameters.r, phi=self.parameters.phi
+                fn=triples.twomode_squeezing_gate_Abc,
+                r=self.parameters.r,
+                phi=self.parameters.phi,
             ),
         ).representation

mrmustard/lab_dev/transformations/sgate.py

@@ -92,6 +92,8 @@ def __init__(
             modes_in=(mode,),
             modes_out=(mode,),
             ansatz=PolyExpAnsatz.from_function(
-                fn=triples.squeezing_gate_Abc, r=self.parameters.r, delta=self.parameters.phi
+                fn=triples.squeezing_gate_Abc,
+                r=self.parameters.r,
+                delta=self.parameters.phi,
             ),
         ).representation

mrmustard/math/backend_manager.py

@@ -387,16 +387,19 @@ def block(self, blocks: list[list[Tensor]], axes=(-2, -1)) -> Tensor:
         """
         return self._apply("block", (blocks, axes))
 
-    def broadcast_to(self, array: Tensor, shape: tuple[int]) -> Tensor:
+    def broadcast_to(self, array: Tensor, shape: tuple[int], dtype=None) -> Tensor:
         r"""Broadcasts an array to a new shape.
 
         Args:
             array: The array to broadcast.
             shape: The shape to broadcast to.
+            dtype: The data type to cast to. If ``None``, the returned array
+                is of the same type as the given one.
 
         Returns:
             The broadcasted array.
         """
+        array = self.astensor(array, dtype)
         return self._apply("broadcast_to", (array, shape))
 
     def cast(self, array: Tensor, dtype=None) -> Tensor:

mrmustard/physics/bargmann_utils.py

@@ -20,7 +20,7 @@
 
 from mrmustard import math, settings
 from mrmustard.physics.husimi import pq_to_aadag, wigner_to_husimi
-from mrmustard.utils.typing import ComplexMatrix, Matrix, Vector, Scalar
+from mrmustard.utils.typing import RealMatrix, ComplexMatrix, Matrix, Vector, Scalar
 
 
 def bargmann_Abc_to_phasespace_cov_means(
@@ -217,33 +217,41 @@ def au2Symplectic(A):
     return math.real(transformation @ S @ math.conj(math.transpose(transformation)))
 
 
-def symplectic2Au(S):
+def symplectic2Au(symplectic: RealMatrix) -> ComplexMatrix:
     r"""
     The inverse of au2Symplectic i.e., returns symplectic, given Au
 
-    S: symplectic in XXPP order
+    symplectic: symplectic in XXPP order
     """
-    m = S.shape[-1]
+    batch_size = symplectic.shape[:-2]
+    batch_shape = batch_size or (1,)
+    batch_dim = len(batch_shape)
+
+    symp_batch = math.broadcast_to(symplectic, batch_shape + symplectic.shape[-2:])
+
+    m = symp_batch.shape[-1]
     m = m // 2
     # the following lines of code transform the quadrature symplectic matrix to
     # the annihilation one
     R = math.rotmat(m)
-    S = R @ S @ math.dagger(R)
+    S = R @ symp_batch @ math.dagger(R)
     # identifying blocks of S
-    S_1 = S[:m, :m]
-    S_2 = S[:m, m:]
 
-    # TODO: broadcasting/batch stuff consider a batch dimension
+    S_1 = S[..., :m, :m]
+    S_2 = S[..., :m, m:]
+
+    perm = tuple(range(len(S_1.shape)))
+    perm = perm[:batch_dim] + perm[batch_dim:][::-1]
 
     # the formula to apply comes here
     A_1 = S_2 @ math.conj(math.inv(S_1))  # use solve for inverse
-    A_2 = math.conj(math.inv(math.transpose(S_1)))
-    A_3 = math.transpose(A_2)
+    A_2 = math.conj(math.inv(math.transpose(S_1, perm)))
+    A_3 = math.transpose(A_2, perm)
     A_4 = -math.conj(math.solve(S_1, S_2))
 
-    A = math.block([[A_1, A_2], [A_3, A_4]])
+    A = math.concat([math.concat([A_1, A_2], -1), math.concat([A_3, A_4], -1)], -2)
 
-    return A
+    return A if batch_size else A[0]
 
 
 def XY_of_channel(A: ComplexMatrix):

mrmustard/physics/symplectics.py

@@ -18,11 +18,14 @@
 """
 from __future__ import annotations
 
+from typing import Iterable
+
 from mrmustard import math
 from mrmustard.utils.typing import Matrix
+from .utils import compute_batch_size
 
 
-def cxgate_symplectic(s: float) -> Matrix:
+def cxgate_symplectic(s: float | Iterable[float]) -> Matrix:
     r"""
     The symplectic matrix of a controlled X gate.
 
@@ -32,12 +35,27 @@ def cxgate_symplectic(s: float) -> Matrix:
     Returns:
         The symplectic matrix of a CX gate.
     """
-    return math.astensor(
-        [[1, 0, 0, 0], [s, 1, 0, 0], [0, 0, 1, -s], [0, 0, 0, 1]], dtype="complex128"
+    batch_size, batch_dim = compute_batch_size(s)
+    batch_shape = batch_size or (1,)
+
+    s_batch = math.broadcast_to(math.cast(s, math.complex128), batch_shape)
+
+    O_matrix = math.zeros(batch_shape, math.complex128)
+    I_matrix = math.ones(batch_shape, math.complex128)
+
+    symplectic = math.stack(
+        [
+            math.stack([I_matrix, O_matrix, O_matrix, O_matrix], batch_dim),
+            math.stack([s_batch, I_matrix, O_matrix, O_matrix], batch_dim),
+            math.stack([O_matrix, O_matrix, I_matrix, -s_batch], batch_dim),
+            math.stack([O_matrix, O_matrix, O_matrix, I_matrix], batch_dim),
+        ],
+        batch_dim,
     )
+    return symplectic if batch_size else symplectic[0]
 
 
-def czgate_symplectic(s: float) -> Matrix:
+def czgate_symplectic(s: float | Iterable[float]) -> Matrix:
     r"""
     The symplectic matrix of a controlled Z gate.
 
@@ -47,7 +65,24 @@ def czgate_symplectic(s: float) -> Matrix:
     Returns:
         The symplectic matrix of a CZ gate.
     """
-    return math.astensor([[1, 0, 0, 0], [0, 1, 0, 0], [0, s, 1, 0], [s, 0, 0, 1]])
+    batch_size, batch_dim = compute_batch_size(s)
+    batch_shape = batch_size or (1,)
+
+    s_batch = math.broadcast_to(math.cast(s, math.complex128), batch_shape)
+
+    O_matrix = math.zeros(batch_shape, math.complex128)
+    I_matrix = math.ones(batch_shape, math.complex128)
+
+    symplectic = math.stack(
+        [
+            math.stack([I_matrix, O_matrix, O_matrix, O_matrix], batch_dim),
+            math.stack([O_matrix, I_matrix, O_matrix, O_matrix], batch_dim),
+            math.stack([O_matrix, s_batch, I_matrix, O_matrix], batch_dim),
+            math.stack([s_batch, O_matrix, O_matrix, I_matrix], batch_dim),
+        ],
+        batch_dim,
+    )
+    return symplectic if batch_size else symplectic[0]
 
 
 def interferometer_symplectic(unitary: Matrix) -> Matrix:
@@ -60,12 +95,22 @@ def interferometer_symplectic(unitary: Matrix) -> Matrix:
     Returns:
         The symplectic matrix of an N-mode interferometer.
     """
-    return math.block(
-        [[math.real(unitary), -math.imag(unitary)], [math.imag(unitary), math.real(unitary)]]
+    batch_size = unitary.shape[:-2]
+    batch_shape = batch_size or (1,)
+    unitary_batch = math.broadcast_to(unitary, batch_shape + unitary.shape[-2:])
+    symplectic = math.concat(
+        [
+            math.concat([math.real(unitary_batch), -math.imag(unitary_batch)], -1),
+            math.concat([math.imag(unitary_batch), math.real(unitary_batch)], -1),
+        ],
+        -2,
     )
+    return symplectic if batch_size else symplectic[0]
 
 
-def mzgate_symplectic(phi_a: float, phi_b: float, internal: bool) -> Matrix:
+def mzgate_symplectic(
+    phi_a: float | Iterable[float], phi_b: float | Iterable[float], internal: bool
+) -> Matrix:
     r"""
     The symplectic matrix of a Mach-Zehnder gate.
 
@@ -81,48 +126,69 @@ def mzgate_symplectic(phi_a: float, phi_b: float, internal: bool) -> Matrix:
     Returns:
         The symplectic matrix of a Mach-Zehnder gate.
     """
-    ca = math.cos(complex(phi_a))
-    sa = math.sin(complex(phi_a))
-    cb = math.cos(complex(phi_b))
-    sb = math.sin(complex(phi_b))
-    cp = math.cos(complex(phi_a + phi_b))
-    sp = math.sin(complex(phi_a + phi_b))
+    batch_size, batch_dim = compute_batch_size(phi_a, phi_b)
+    batch_shape = batch_size or (1,)
+
+    phi_a_batch = math.broadcast_to(phi_a, batch_shape)
+    phi_b_batch = math.broadcast_to(phi_b, batch_shape)
+
+    ca = math.cos(phi_a_batch)
+    sa = math.sin(phi_a_batch)
+    cb = math.cos(phi_b_batch)
+    sb = math.sin(phi_b_batch)
+    cp = math.cos(phi_a_batch + phi_b_batch)
+    sp = math.sin(phi_a_batch + phi_b_batch)
     if internal:
-        return 0.5 * math.astensor(
+        symplectic = math.stack(
             [
-                [ca - cb, -sa - sb, sb - sa, -ca - cb],
-                [-sa - sb, cb - ca, -ca - cb, sa - sb],
-                [sa - sb, ca + cb, ca - cb, -sa - sb],
-                [ca + cb, sb - sa, -sa - sb, cb - ca],
-            ]
+                math.stack([ca - cb, -sa - sb, sb - sa, -ca - cb], batch_dim),
+                math.stack([-sa - sb, cb - ca, -ca - cb, sa - sb], batch_dim),
+                math.stack([sa - sb, ca + cb, ca - cb, -sa - sb], batch_dim),
+                math.stack([ca + cb, sb - sa, -sa - sb, cb - ca], batch_dim),
+            ],
+            batch_dim,
         )
     else:
-        return 0.5 * math.astensor(
+        symplectic = math.stack(
             [
-                [cp - ca, -sb, sa - sp, -1 - cb],
-                [-sa - sp, 1 - cb, -ca - cp, sb],
-                [sp - sa, 1 + cb, cp - ca, -sb],
-                [cp + ca, -sb, -sa - sp, 1 - cb],
-            ]
+                math.stack([cp - ca, -sb, sa - sp, -1 - cb], batch_dim),
+                math.stack([-sa - sp, 1 - cb, -ca - cp, sb], batch_dim),
+                math.stack([sp - sa, 1 + cb, cp - ca, -sb], batch_dim),
+                math.stack([cp + ca, -sb, -sa - sp, 1 - cb], batch_dim),
+            ],
+            batch_dim,
         )
+    symplectic = math.cast(0.5 * symplectic, math.complex128)
+    return symplectic if batch_size else symplectic[0]
 
 
-def pgate_symplectic(n_modes: int, shearing: float) -> Matrix:
+def pgate_symplectic(n_modes: int, shearing: float | Iterable[float]) -> Matrix:
     r"""
     The symplectic matrix of a quadratic phase gate.
 
     Args:
+        n_modes: The number of modes.
         shearing: The shearing parameter.
 
     Returns:
         The symplectic matrix of a phase gate.
     """
-    return math.block(
+    batch_size, _ = compute_batch_size(shearing)
+    batch_shape = batch_size or (1,)
+
+    shearing_batch = math.broadcast_to(shearing, batch_shape)
+
+    I_matrix = math.broadcast_to(math.eye(n_modes), batch_shape + (n_modes, n_modes))
+    O_matrix = math.zeros(batch_shape + (n_modes, n_modes))
+
+    symplectic = math.concat(
         [
-            [math.eye(n_modes), math.zeros((n_modes, n_modes))],
-            [math.eye(n_modes) * shearing, math.eye(n_modes)],
-        ]
+            math.concat([I_matrix, O_matrix], -1),
+            math.concat([math.eye(n_modes) * shearing_batch[..., None, None], I_matrix], -1),
+        ],
+        -2,
     )
+    return symplectic if batch_size else symplectic[0]
 
 
 def realinterferometer_symplectic(orthogonal: Matrix) -> Matrix:
@@ -135,6 +201,14 @@ def realinterferometer_symplectic(orthogonal: Matrix) -> Matrix:
     Returns:
         The symplectic matrix of an N-mode interferometer.
     """
-    return math.block(
-        [[orthogonal, -math.zeros_like(orthogonal)], [math.zeros_like(orthogonal), orthogonal]]
+    batch_size = orthogonal.shape[:-2]
+    batch_shape = batch_size or (1,)
+    orthogonal_batch = math.broadcast_to(orthogonal, batch_shape + orthogonal.shape[-2:])
+    symplectic = math.concat(
+        [
+            math.concat([orthogonal_batch, -math.zeros_like(orthogonal_batch)], -1),
+            math.concat([math.zeros_like(orthogonal_batch), orthogonal_batch], -1),
+        ],
+        -2,
     )
+    return symplectic if batch_size else symplectic[0]