Merging Representation and Ansatz

.github/workflows/tests_docs.yml

@@ -39,6 +39,5 @@ jobs:
       - name: Run tests
         run: |
           python -m pytest --doctest-modules mrmustard/math/parameter_set.py
-          python -m pytest --doctest-modules mrmustard/physics/ansatze.py
-          python -m pytest --doctest-modules mrmustard/physics/representations.py
+          python -m pytest --doctest-modules mrmustard/physics/representations
           python -m pytest --doctest-modules mrmustard/lab_dev

doc/code/physics.rst

@@ -4,7 +4,6 @@ mrmustard.physics
 .. toctree::
     :maxdepth: 1
 
-    physics/ansatze
     physics/representations
 
 .. toctree::

doc/code/physics/utils/bargmann_calculations.rst

@@ -1,8 +1,8 @@
 Calculations on Bargmann objects
 ================================
 
-.. currentmodule:: mrmustard.physics.bargmann
+.. currentmodule:: mrmustard.physics.bargmann_utils
 
-.. automodapi:: mrmustard.physics.bargmann
+.. automodapi:: mrmustard.physics.bargmann_utils
     :no-heading:
     :include-all-objects:

doc/code/physics/utils/fock_calculations.rst

@@ -1,8 +1,8 @@
 Calculations on Fock objects
 ============================
 
-.. currentmodule:: mrmustard.physics.fock
+.. currentmodule:: mrmustard.physics.fock_utils
 
-.. automodapi:: mrmustard.physics.fock
+.. automodapi:: mrmustard.physics.fock_utils
     :no-heading:
     :include-all-objects:

mrmustard/lab/abstract/state.py

@@ -28,7 +28,7 @@
 
 from mrmustard import math, settings
 from mrmustard.math.parameters import Constant, Variable
-from mrmustard.physics import bargmann, fock, gaussian
+from mrmustard.physics import bargmann_utils, fock_utils, gaussian
 from mrmustard.physics.wigner import wigner_discretized
 from mrmustard.utils.typing import (
     ComplexMatrix,
@@ -157,7 +157,7 @@ def purity(self) -> float:
             if self.is_gaussian:
                 self._purity = gaussian.purity(self.cov)
             else:
-                self._purity = fock.purity(self._dm)
+                self._purity = fock_utils.purity(self._dm)
         return self._purity
 
     @property
@@ -187,15 +187,15 @@ def number_stdev(self) -> RealVector:
             return math.sqrt(math.diag_part(self.number_cov))
 
         return math.sqrt(
-            fock.number_variances(self.fock, is_dm=len(self.fock.shape) == self.num_modes * 2)
+            fock_utils.number_variances(self.fock, is_dm=len(self.fock.shape) == self.num_modes * 2)
         )
 
     @property
     def cutoffs(self) -> list[int]:
         r"""Returns the Hilbert space dimension of each mode."""
         if self._cutoffs is None:
             if self._ket is None and self._dm is None:
-                self._cutoffs = fock.autocutoffs(
+                self._cutoffs = fock_utils.autocutoffs(
                     self.cov, self.means, settings.AUTOSHAPE_PROBABILITY
                 )
             else:
@@ -223,7 +223,7 @@ def shape(self) -> list[int]:
     def fock(self) -> ComplexTensor:
         r"""Returns the Fock representation of the state."""
         if self._dm is None and self._ket is None:
-            _fock = fock.wigner_to_fock_state(
+            _fock = fock_utils.wigner_to_fock_state(
                 self.cov,
                 self.means,
                 shape=self.shape,
@@ -243,7 +243,7 @@ def number_means(self) -> RealVector:
         if self.is_gaussian:
             return gaussian.number_means(self.cov, self.means)
 
-        return fock.number_means(tensor=self.fock, is_dm=self.is_mixed)
+        return fock_utils.number_means(tensor=self.fock, is_dm=self.is_mixed)
 
     @property
     def number_cov(self) -> RealMatrix:
@@ -258,7 +258,7 @@ def norm(self) -> float:
         r"""Returns the norm of the state."""
         if self.is_gaussian:
             return self._norm
-        return fock.norm(self.fock, not self.is_hilbert_vector)
+        return fock_utils.norm(self.fock, not self.is_hilbert_vector)
 
     @property
     def probability(self) -> float:
@@ -296,7 +296,7 @@ def ket(
             cutoffs = [c if c is not None else self.cutoffs[i] for i, c in enumerate(cutoffs)]
 
         if self.is_gaussian:
-            self._ket = fock.wigner_to_fock_state(
+            self._ket = fock_utils.wigner_to_fock_state(
                 self.cov,
                 self.means,
                 shape=cutoffs,
@@ -308,12 +308,12 @@ def ket(
             if self._ket is None:
                 # if state is pure and has a density matrix, calculate the ket
                 if self.is_pure:
-                    self._ket = fock.dm_to_ket(self._dm)
+                    self._ket = fock_utils.dm_to_ket(self._dm)
             current_cutoffs = [int(s) for s in self._ket.shape]
             if cutoffs != current_cutoffs:
                 paddings = [(0, max(0, new - old)) for new, old in zip(cutoffs, current_cutoffs)]
                 if any(p != (0, 0) for p in paddings):
-                    padded = fock.math.pad(self._ket, paddings, mode="constant")
+                    padded = fock_utils.math.pad(self._ket, paddings, mode="constant")
                 else:
                     padded = self._ket
                 return padded[tuple(slice(s) for s in cutoffs)]
@@ -336,16 +336,16 @@ def dm(self, cutoffs: list[int] | None = None) -> ComplexTensor:
         if self.is_pure:
             ket = self.ket(cutoffs=cutoffs)
             if ket is not None:
-                return fock.ket_to_dm(ket)
+                return fock_utils.ket_to_dm(ket)
         else:
             if self.is_gaussian:
-                self._dm = fock.wigner_to_fock_state(
+                self._dm = fock_utils.wigner_to_fock_state(
                     self.cov, self.means, shape=cutoffs + cutoffs, return_dm=True
                 )
             elif cutoffs != (current_cutoffs := list(self._dm.shape[: self.num_modes])):
                 paddings = [(0, max(0, new - old)) for new, old in zip(cutoffs, current_cutoffs)]
                 if any(p != (0, 0) for p in paddings):
-                    padded = fock.math.pad(self._dm, paddings + paddings, mode="constant")
+                    padded = fock_utils.math.pad(self._dm, paddings + paddings, mode="constant")
                 else:
                     padded = self._dm
                 return padded[tuple(slice(s) for s in cutoffs + cutoffs)]
@@ -366,10 +366,10 @@ def fock_probabilities(self, cutoffs: Sequence[int]) -> RealTensor:
         if self._fock_probabilities is None:
             if self.is_mixed:
                 dm = self.dm(cutoffs=cutoffs)
-                self._fock_probabilities = fock.dm_to_probs(dm)
+                self._fock_probabilities = fock_utils.dm_to_probs(dm)
             else:
                 ket = self.ket(cutoffs=cutoffs)
-                self._fock_probabilities = fock.ket_to_probs(ket)
+                self._fock_probabilities = fock_utils.ket_to_probs(ket)
         return self._fock_probabilities
 
     def primal(self, other: State | Transformation) -> State:
@@ -430,9 +430,9 @@ def _project_onto_fock(self, other: State) -> State | float:
 
         # return the probability (norm) of the state when there are no modes left
         return (
-            fock.math.abs(out_fock) ** 2
+            fock_utils.math.abs(out_fock) ** 2
             if other.is_pure and self.is_pure
-            else fock.math.abs(out_fock)
+            else fock_utils.math.abs(out_fock)
         )
 
     def _contract_with_other(self, other):
@@ -444,7 +444,7 @@ def _contract_with_other(self, other):
         else:
             # matching other's cutoffs
             self_cutoffs = [other.cutoffs[other.indices(m)] for m in self.modes]
-            out_fock = fock.contract_states(
+            out_fock = fock_utils.contract_states(
                 stateA=(other.ket(other_cutoffs) if other.is_pure else other.dm(other_cutoffs)),
                 stateB=(self.ket(self_cutoffs) if self.is_pure else self.dm(self_cutoffs)),
                 a_is_dm=other.is_mixed,
@@ -499,7 +499,7 @@ def __and__(self, other: State) -> State:
             if self.is_mixed or other.is_mixed:
                 self_fock = self.dm()
                 other_fock = other.dm()
-                dm = fock.math.tensordot(self_fock, other_fock, [[], []])
+                dm = fock_utils.math.tensordot(self_fock, other_fock, [[], []])
                 # e.g. self has shape [1,3,1,3] and other has shape [2,2]
                 # we want self & other to have shape [1,3,2,1,3,2]
                 # before transposing shape is [1,3,1,3]+[2,2]
@@ -519,7 +519,7 @@ def __and__(self, other: State) -> State:
             self_fock = self.ket()
             other_fock = other.ket()
             return State(
-                ket=fock.math.tensordot(self_fock, other_fock, [[], []]),
+                ket=fock_utils.math.tensordot(self_fock, other_fock, [[], []]),
                 modes=self.modes + [m + max(self.modes) + 1 for m in other.modes],
             )
         cov = gaussian.join_covs([self.cov, other.cov])
@@ -551,9 +551,9 @@ def bargmann(self, numpy=False) -> tuple[ComplexMatrix, ComplexVector, complex]
         """
         if self.is_gaussian:
             if self.is_pure:
-                A, B, C = bargmann.wigner_to_bargmann_psi(self.cov, self.means)
+                A, B, C = bargmann_utils.wigner_to_bargmann_psi(self.cov, self.means)
             else:
-                A, B, C = bargmann.wigner_to_bargmann_rho(self.cov, self.means)
+                A, B, C = bargmann_utils.wigner_to_bargmann_rho(self.cov, self.means)
         else:
             return None
         if numpy:
@@ -582,7 +582,7 @@ def get_modes(self, item) -> State:
             means, _ = gaussian.partition_means(self.means, item_idx)
             return State(cov=cov, means=means, modes=item)
 
-        fock_partitioned = fock.trace(self.dm(self.cutoffs), keep=item_idx)
+        fock_partitioned = fock_utils.trace(self.dm(self.cutoffs), keep=item_idx)
         return State(dm=fock_partitioned, modes=item)
 
     def __eq__(self, other) -> bool:  # pylint: disable=too-many-return-statements
@@ -736,8 +736,8 @@ def mikkel_plot(
     if plot_args["ytick_labels"] is None:
         plot_args["ytick_labels"] = plot_args["yticks"]
 
-    q, ProbX = fock.quadrature_distribution(rho)
-    p, ProbP = fock.quadrature_distribution(rho, np.pi / 2)
+    q, ProbX = fock_utils.quadrature_distribution(rho)
+    p, ProbP = fock_utils.quadrature_distribution(rho, np.pi / 2)
 
     xvec = np.linspace(*xbounds, plot_args["resolution"])
     pvec = np.linspace(*ybounds, plot_args["resolution"])

mrmustard/lab/abstract/transformation.py

@@ -27,7 +27,7 @@
 from mrmustard.math.parameter_set import ParameterSet
 from mrmustard.math.parameters import Constant, Variable
 from mrmustard.math.tensor_networks import Tensor
-from mrmustard.physics import bargmann, fock, gaussian
+from mrmustard.physics import bargmann_utils, fock_utils, gaussian
 from mrmustard.utils.typing import RealMatrix, RealVector
 
 from .state import State
@@ -172,9 +172,9 @@ def d_vector_dual(self) -> RealVector | None:
     def bargmann(self, numpy=False):
         X, Y, d = self.XYd(allow_none=False)
         if self.is_unitary:
-            A, B, C = bargmann.wigner_to_bargmann_U(X, d)
+            A, B, C = bargmann_utils.wigner_to_bargmann_U(X, d)
         else:
-            A, B, C = bargmann.wigner_to_bargmann_Choi(X, Y, d)
+            A, B, C = bargmann_utils.wigner_to_bargmann_Choi(X, Y, d)
         if numpy:
             return math.asnumpy(A), math.asnumpy(B), math.asnumpy(C)
         return A, B, C
@@ -208,11 +208,11 @@ def choi(
             U = self.U(shape[: self.num_modes])
             Udual = self.U(shape[self.num_modes :])
             if dual:
-                return fock.U_to_choi(U=Udual, Udual=U)
-            return fock.U_to_choi(U=U, Udual=Udual)
+                return fock_utils.U_to_choi(U=Udual, Udual=U)
+            return fock_utils.U_to_choi(U=U, Udual=Udual)
 
         X, Y, d = self.XYd(allow_none=False)
-        choi = fock.wigner_to_fock_Choi(X, Y, d, shape=shape)
+        choi = fock_utils.wigner_to_fock_Choi(X, Y, d, shape=shape)
         if dual:
             n = len(shape) // 4
             N0 = list(range(0, n))
@@ -382,8 +382,10 @@ def _transform_fock(self, state: State, dual=False) -> State:
         op_idx = [state.modes.index(m) for m in self.modes]
         U = self.U(cutoffs=[state.cutoffs[i] for i in op_idx])
         if state.is_hilbert_vector:
-            return State(ket=fock.apply_kraus_to_ket(U, state.ket(), op_idx), modes=state.modes)
-        return State(dm=fock.apply_kraus_to_dm(U, state.dm(), op_idx), modes=state.modes)
+            return State(
+                ket=fock_utils.apply_kraus_to_ket(U, state.ket(), op_idx), modes=state.modes
+            )
+        return State(dm=fock_utils.apply_kraus_to_dm(U, state.dm(), op_idx), modes=state.modes)
 
     def U(
         self,
@@ -411,7 +413,7 @@ def U(
             raise ValueError(f"len(cutoffs) must be {self.num_modes} (got {len(cutoffs)})")
         shape = shape or tuple(cutoffs) * 2
         X, _, d = self.XYd(allow_none=False)
-        return fock.wigner_to_fock_U(X, d, shape=shape)
+        return fock_utils.wigner_to_fock_U(X, d, shape=shape)
 
     def __eq__(self, other):
         r"""Returns ``True`` if the two transformations are equal."""
@@ -453,8 +455,10 @@ def _transform_fock(self, state: State, dual: bool = False) -> State:
         op_idx = [state.modes.index(m) for m in self.modes]
         choi = self.choi(cutoffs=[state.cutoffs[i] for i in op_idx], dual=dual)
         if state.is_hilbert_vector:
-            return State(dm=fock.apply_choi_to_ket(choi, state.ket(), op_idx), modes=state.modes)
-        return State(dm=fock.apply_choi_to_dm(choi, state.dm(), op_idx), modes=state.modes)
+            return State(
+                dm=fock_utils.apply_choi_to_ket(choi, state.ket(), op_idx), modes=state.modes
+            )
+        return State(dm=fock_utils.apply_choi_to_dm(choi, state.dm(), op_idx), modes=state.modes)
 
     def value(self, shape: tuple[int]):
         return self.choi(shape=shape)

mrmustard/lab/detectors.py

@@ -21,7 +21,7 @@
 from typing import Iterable
 
 from mrmustard import settings
-from mrmustard.physics import fock, gaussian
+from mrmustard.physics import fock_utils, gaussian
 from mrmustard.utils.typing import RealMatrix, RealVector
 
 from mrmustard import math
@@ -407,7 +407,7 @@ def _measure_fock(self, other) -> State | float:
         reduced_state = other.get_modes(self.modes)
 
         # build pdf and sample homodyne outcome
-        x_outcome, probability = fock.sample_homodyne(
+        x_outcome, probability = fock_utils.sample_homodyne(
             state=reduced_state.ket() if reduced_state.is_pure else reduced_state.dm(),
             quadrature_angle=self.quadrature_angle,
         )
@@ -427,7 +427,7 @@ def _measure_fock(self, other) -> State | float:
         other_cutoffs = [
             None if m not in self.modes else other.cutoffs[other.indices(m)] for m in other.modes
         ]
-        out_fock = fock.contract_states(
+        out_fock = fock_utils.contract_states(
             stateA=(other.ket(other_cutoffs) if other.is_pure else other.dm(other_cutoffs)),
             stateB=self.state.ket(self_cutoffs),
             a_is_dm=other.is_mixed,