New cft

mrmustard/lab_dev/transformations/cft.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.
+
+"""
+The class representing a complex fourier transform.
+"""
+
+from typing import Sequence
+from mrmustard.lab_dev.transformations.base import Map
+from mrmustard.physics.representations import Bargmann
+from mrmustard.physics import triples
+
+
+class CFT(Map):
+ r"""The Complex Fourier Transformation as a channel.
+ The main use is to convert between Characteristic functions and phase space functions.
+
+ Args:
+ num_modes: number of modes of this channel.
+ """
+
+ def __init__(
+ self,
+ modes: Sequence[int],
+ ):
+ super().__init__(
+ modes_out=modes,
+ modes_in=modes,
+ name="CFT",
+ )
+ self._representation = Bargmann.from_function(
+ fn=triples.complex_fourier_transform_Abc, n_modes=len(modes)
+ )

mrmustard/physics/ansatze.py

@@ -1208,6 +1208,10 @@ def bargmann_Abc_to_phasespace_cov_means(
  Returns:
  The covariance matrix, mean vector and coefficient of the state in phase space.
  """
+ batched = len(A.shape) == 3 and len(b.shape) == 2 and len(c.shape) == 1
+ A = math.atleast_3d(A)
+ b = math.atleast_2d(b)
+ c = math.atleast_1d(c)
  num_modes = A.shape[-1] // 2
  Omega = math.cast(math.transpose(math.J(num_modes)), dtype=math.complex128)
  W = math.transpose(math.conj(math.rotmat(num_modes)))
@@ -1219,4 +1223,6 @@ def bargmann_Abc_to_phasespace_cov_means(
  1j * math.matvec(Omega @ W, bvec) * math.sqrt(settings.HBAR, dtype=math.complex128)
  for bvec in b
  ]
- return math.astensor(cov), math.astensor(mean), coeff
+ if batched:
+ return math.astensor(cov), math.astensor(mean), coeff
+ return cov[0], mean[0], coeff[0]

mrmustard/physics/triples.py

@@ -677,10 +677,32 @@ def displacement_map_s_parametrized_Abc(s: int, n_modes: int) -> Union[Matrix, V
 
  A = math.astensor(math.asnumpy(A)[order_list, :][:, order_list])
  b = _vacuum_B_vector(4 * n_modes)
- c = 1.0 + 0j
+ c = 1.0 / (2 * np.pi) ** n_modes + 0.0j
  return math.astensor(A), b, c
 
 
+def complex_fourier_transform_Abc(n_modes: int) -> Union[Matrix, Vector, Scalar]:
+ r"""
+ The ``(A, b, c)`` triple of the complex Fourier transform between two pairs of complex variables.
+ Given a function :math:`f(z^*, z)`, the complex Fourier transform is defined as
+ :math:
+ \hat{f} (y^*, y) = \int_{\mathbb{C}} \frac{d^2 z}{\pi} e^{yz^* - y^*z} f(z^*, z).
+ The indices of this triple correspond to the variables :math:`(y^*, z^*, y, z)`.
+
+ Args:
+ n_modes: the number of modes for this map.
+
+ Returns:
+ The ``(A, b, c)`` triple of the complex fourier transform.
+ """
+ O2n = math.zeros((2 * n_modes, 2 * n_modes))
+ Omega = math.J(n_modes)
+ A = math.block([[O2n, -Omega], [Omega, O2n]])
+ b = _vacuum_B_vector(4 * n_modes)
+ c = 1.0 + 0j
+ return A, b, c
+
+
 # ~~~~~~~~~~~~~~~~
 # Kraus operators
 # ~~~~~~~~~~~~~~~~

tests/test_lab_dev/test_states/test_states_base.py

@@ -144,7 +144,7 @@ def test_to_from_fock(self, modes):
  @pytest.mark.parametrize("modes", [[0], [0, 1], [3, 19, 2]])
  def test_to_from_phase_space(self, modes):
  cov, means, coeff = Coherent([0], x=1, y=2).phase_space(s=0)
- assert math.allclose(coeff[0], 1.0)
+ assert math.allclose(coeff[0], 1.0 / (2 * np.pi))
  assert math.allclose(cov[0], np.eye(2) * settings.HBAR / 2)
  assert math.allclose(means[0], np.array([1.0, 2.0]) * np.sqrt(2 * settings.HBAR))
  n_modes = len(modes)
@@ -552,7 +552,7 @@ def test_to_from_fock(self, modes):
  def test_to_from_phase_space(self):
  state0 = Coherent([0], x=1, y=2) >> Attenuator([0], 1.0)
  cov, means, coeff = state0.phase_space(s=0) # batch = 1
- assert coeff[0] == 1.0
+ assert coeff[0] == 1.0 / (2 * np.pi)
  assert math.allclose(cov[0], np.eye(2) * settings.HBAR / 2)
  assert math.allclose(means[0], np.array([1.0, 2.0]) * np.sqrt(settings.HBAR * 2))
 

tests/test_lab_dev/test_transformations/test_cft.py

@@ -0,0 +1,54 @@
+# 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 ``CFT`` class."""
+
+import numpy as np
+from mrmustard import math, settings
+from mrmustard.lab_dev import BtoPS, DisplacedSqueezed
+from mrmustard.lab_dev.transformations.cft import CFT
+from mrmustard.physics.wigner import wigner_discretized
+
+
+class TestCFT:
+ r"""
+ Tests the CFT gate
+ """
+
+ def test_init(self):
+ "tests the initialization of the CFT gate"
+ cft = CFT([0])
+ assert cft.name == "CFT"
+ assert cft.modes == [0]
+
+ def test_wigner_function(self):
+ r"""
+ Tests that the characteristic function is converted to the Wigner function
+ for a single-mode squeezed state.
+ """
+ settings.HBAR = 2
+ settings.HBAR = 2
+ state = DisplacedSqueezed([0], r=0.5, phi=0.0, x=1.0, y=0.0)
+
+ state2 = DisplacedSqueezed([0], r=0.5, phi=np.pi, x=0.0, y=1.0)
+ dm = math.sum(state2.to_fock(100).dm().representation.array, axes=[0])
+ xvec = np.linspace(-5, 5, 100)
+ pvec = np.linspace(-5, 5, 100)
+ wigner, _, _ = wigner_discretized(dm, xvec, pvec)
+
+ Wigner = (state >> CFT([0]) >> BtoPS([0], s=0)).representation.ansatz
+ X, Y = np.meshgrid(xvec, pvec)
+ Z = np.array([X - 1j * Y, X + 1j * Y]).transpose((1, 2, 0))
+ assert math.allclose((np.real(Wigner(Z))), (np.real(wigner)), atol=1e-8)
+ assert math.allclose(np.max(np.real(Wigner(Z))), np.max(np.real(wigner)), atol=1e-8)

tests/test_physics/test_ansatz.py

@@ -233,15 +233,15 @@ def test_bargmann_Abc_to_phasespace_cov_means(self):
  state_means = np.array([0.2, 0.3])
  state = DM.from_bargmann([0], wigner_to_bargmann_rho(state_cov, state_means))
  state_after = state >> BtoPS(modes=[0], s=0) # pylint: disable=protected-access
- A1, b1, c1 = state_after.bargmann_triple(batched=True)
+ A1, b1, c1 = state_after.bargmann_triple()
  (
  new_state_cov,
  new_state_means,
  new_state_coeff,
  ) = bargmann_Abc_to_phasespace_cov_means(A1, b1, c1)
- assert np.allclose(state_cov, new_state_cov[0])
- assert np.allclose(state_means, new_state_means[0])
- assert np.allclose(1.0, new_state_coeff[0])
+ assert np.allclose(state_cov, new_state_cov)
+ assert np.allclose(state_means, new_state_means)
+ assert np.allclose(1.0 / (2 * np.pi), new_state_coeff)
 
  state_cov = np.array(
  [
@@ -256,27 +256,25 @@ def test_bargmann_Abc_to_phasespace_cov_means(self):
  state = DM.from_bargmann(modes=[0, 1], triple=(A, b, c))
 
  state_after = state >> BtoPS(modes=[0, 1], s=0) # pylint: disable=protected-access
- A1, b1, c1 = state_after.bargmann_triple(batched=True)
+ A1, b1, c1 = state_after.bargmann_triple()
  (
  new_state_cov1,
  new_state_means1,
  new_state_coeff1,
  ) = bargmann_Abc_to_phasespace_cov_means(A1, b1, c1)
 
- A22, b22, c22 = (state >> BtoPS([0], 0) >> BtoPS([1], 0)).bargmann_triple(
- batched=True
- ) # pylint: disable=protected-access
+ A22, b22, c22 = (state >> BtoPS([0], 0) >> BtoPS([1], 0)).bargmann_triple()
  (
  new_state_cov22,
  new_state_means22,
  new_state_coeff22,
  ) = bargmann_Abc_to_phasespace_cov_means(A22, b22, c22)
- assert math.allclose(new_state_cov22[0], state_cov)
- assert math.allclose(new_state_cov1[0], state_cov)
- assert math.allclose(new_state_means1[0], state_means)
- assert math.allclose(new_state_means22[0], state_means)
- assert math.allclose(new_state_coeff1[0], 1.0)
- assert math.allclose(new_state_coeff22[0], 1.0)
+ assert math.allclose(new_state_cov22, state_cov)
+ assert math.allclose(new_state_cov1, state_cov)
+ assert math.allclose(new_state_means1, state_means)
+ assert math.allclose(new_state_means22, state_means)
+ assert math.allclose(new_state_coeff1, 1 / (2 * np.pi) ** 2)
+ assert math.allclose(new_state_coeff22, 1 / (2 * np.pi) ** 2)
 
 
 class TestPolyExpAnsatz:

tests/test_physics/test_triples.py

@@ -317,19 +317,19 @@ def test_displacement_gate_s_parametrized_Abc(self):
  A1_correct = np.array([[0, -0.5, -1, 0], [-0.5, 0, 0, 1], [-1, 0, 0, 1], [0, 1, 1, 0]])
  assert math.allclose(A1, A1_correct[[0, 3, 1, 2], :][:, [0, 3, 1, 2]])
  assert math.allclose(b1, math.zeros(4))
- assert math.allclose(c1, 1)
+ assert math.allclose(c1, 1 / (2 * np.pi))
 
  A2, b2, c2 = triples.displacement_map_s_parametrized_Abc(s=1, n_modes=1)
  A2_correct = np.array([[0, 0, -1, 0], [0, 0, 0, 1], [-1, 0, 0, 1], [0, 1, 1, 0]])
  assert math.allclose(A2, A2_correct[[0, 3, 1, 2], :][:, [0, 3, 1, 2]])
  assert math.allclose(b2, math.zeros(4))
- assert math.allclose(c2, 1)
+ assert math.allclose(c2, 1 / (2 * np.pi))
 
  A3, b3, c3 = triples.displacement_map_s_parametrized_Abc(s=-1, n_modes=1)
  A3_correct = np.array([[0, -1, -1, 0], [-1, 0, 0, 1], [-1, 0, 0, 1], [0, 1, 1, 0]])
  assert math.allclose(A3, A3_correct[[0, 3, 1, 2], :][:, [0, 3, 1, 2]])
  assert math.allclose(b3, math.zeros(4))
- assert math.allclose(c3, 1)
+ assert math.allclose(c3, 1 / (2 * np.pi))
 
  @pytest.mark.parametrize("eta", [0.0, 0.1, 0.5, 0.9, 1.0])
  def test_attenuator_kraus_Abc(self, eta):