Add cached evaluation of tensor networks

.github/workflows/build_test.yml

@@ -102,7 +102,7 @@ jobs:
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies with type hints
-      run: python -m pip install "numpy<2.0" "matplotlib>=3.8" pytket spacy torch types-PyYAML types-Pillow ipython types-networkx
+      run: python -m pip install "numpy<2.0" "matplotlib>=3.8" pytket spacy torch types-PyYAML types-Pillow ipython types-networkx opt-einsum
     - name: Install type checker
       run: python -m pip install mypy
     - name: Type check with mypy

lambeq/__init__.py

@@ -30,6 +30,8 @@
         'Sim14Ansatz',
         'Sim15Ansatz',
         'Sim4Ansatz',
+        'Sim9Ansatz',
+        'Sim9CxAnsatz',
         'SpiderAnsatz',
         'StronglyEntanglingAnsatz',
         'Symbol',
@@ -102,6 +104,9 @@
         'PytorchTrainer',
         'QuantumTrainer',
 
+        'TnPathOptimizer',
+        'CachedTnPathOptimizer',
+
         'BinaryCrossEntropyLoss',
         'CrossEntropyLoss',
         'LossFunction',
@@ -112,7 +117,8 @@
 from lambeq import ansatz, core, rewrite, text2diagram, tokeniser, training
 from lambeq.ansatz import (BaseAnsatz, CircuitAnsatz, IQPAnsatz, MPSAnsatz,
                            Sim14Ansatz, Sim15Ansatz, Sim4Ansatz, SpiderAnsatz,
-                           StronglyEntanglingAnsatz, TensorAnsatz)
+                           StronglyEntanglingAnsatz, TensorAnsatz, Sim9Ansatz,
+                           Sim9CxAnsatz)
 from lambeq.core.globals import VerbosityLevel
 from lambeq.core.types import AtomicType
 from lambeq.rewrite import (CoordinationRewriteRule, CurryRewriteRule,
@@ -138,6 +144,7 @@
                              PennyLaneModel, PytorchModel,
                              PytorchQuantumModel, QuantumModel,
                              TketModel, Trainer, PytorchTrainer,
+                             CachedTnPathOptimizer, TnPathOptimizer,
                              QuantumTrainer, BinaryCrossEntropyLoss,
                              CrossEntropyLoss, LossFunction, MSELoss)
 from lambeq.version import (version as __version__,

lambeq/ansatz/__init__.py

@@ -13,11 +13,12 @@
 # limitations under the License.
 
 __all__ = ['BaseAnsatz', 'CircuitAnsatz', 'IQPAnsatz', 'MPSAnsatz',
-           'Sim14Ansatz', 'Sim15Ansatz', 'Sim4Ansatz', 'SpiderAnsatz',
-           'StronglyEntanglingAnsatz', 'TensorAnsatz']
+           'Sim14Ansatz', 'Sim15Ansatz', 'Sim4Ansatz', 'Sim9Ansatz',
+           'Sim9CxAnsatz', 'SpiderAnsatz', 'StronglyEntanglingAnsatz',
+           'TensorAnsatz']
 
 from lambeq.ansatz.base import BaseAnsatz
-from lambeq.ansatz.circuit import (CircuitAnsatz, IQPAnsatz,
-                                   Sim14Ansatz, Sim15Ansatz, Sim4Ansatz,
-                                   StronglyEntanglingAnsatz)
+from lambeq.ansatz.circuit import (CircuitAnsatz, IQPAnsatz, Sim9Ansatz,
+                                   Sim9CxAnsatz, Sim14Ansatz, Sim15Ansatz,
+                                   Sim4Ansatz, StronglyEntanglingAnsatz)
 from lambeq.ansatz.tensor import MPSAnsatz, SpiderAnsatz, TensorAnsatz

lambeq/ansatz/circuit.py

@@ -25,6 +25,8 @@
 __all__ = ['CircuitAnsatz',
            'IQPAnsatz',
            'Sim4Ansatz',
+           'Sim9Ansatz',
+           'Sim9CxAnsatz',
            'Sim14Ansatz',
            'Sim15Ansatz',
            'StronglyEntanglingAnsatz']
@@ -497,3 +499,111 @@ def circuit(self, n_qubits: int, params: np.ndarray) -> Circuit:
             for j in range(n_qubits):
                 circuit = circuit.CX(j, (j+step) % n_qubits)
         return circuit
+
+
+class Sim9CxAnsatz(CircuitAnsatz):
+    """Circuit 9 from Sim et al., with CZ gates replaced with CX.
+
+    Ansatz with a layer of H gates, followed by a ladder of CX,
+    followed by a layer of RX.
+
+    Paper at: https://arxiv.org/abs/1905.10876
+
+    """
+    def __init__(self,
+                 ob_map: Mapping[Ty, int],
+                 n_layers: int,
+                 n_single_qubit_params: int = 3,
+                 discard: bool = False) -> None:
+        """Instantiate a Sim9 ansatz with CZ gates replaced with CX gates.
+        Parameters
+        ----------
+        ob_map : dict
+            A mapping from :py:class:`discopy.rigid.Ty` to the number of
+            qubits it uses in a circuit.
+        n_layers : int
+            The number of layers used by the ansatz.
+        n_single_qubit_params : int, default: 3
+            The number of single qubit rotations used by the ansatz.
+        discard : bool, default: False
+            Discard open wires instead of post-selecting.
+        """
+        super().__init__(ob_map,
+                         n_layers,
+                         n_single_qubit_params,
+                         self.circuit,
+                         discard,
+                         [Rx, Rz])
+
+    def params_shape(self, n_qubits: int) -> tuple[int, ...]:
+        return (self.n_layers, n_qubits)
+
+    def circuit(self, n_qubits: int, params: np.ndarray) -> Circuit:
+        if n_qubits == 1:
+            circuit = Rx(params[0]) >> Rz(params[1]) >> Rx(params[2])
+        else:
+            circuit = Id(n_qubits)
+
+            for thetas in params:
+                circuit >>= Id().tensor(*(H for _ in range(n_qubits)))
+                cxs = Id(n_qubits)
+                for q in range(n_qubits - 1):
+                    cxs = cxs.CX(q, q + 1)
+                circuit >>= cxs
+
+                circuit >>= Id().tensor(*map(Rx, thetas))
+        return circuit   # type: ignore[return-value]
+
+
+class Sim9Ansatz(CircuitAnsatz):
+    """Circuit 9 from Sim et al.
+
+    Ansatz with a layer of H gates, followed by a ladder of CZ,
+    followed by a layer of RX.
+
+    Paper at: https://arxiv.org/abs/1905.10876
+
+    """
+    def __init__(self,
+                 ob_map: Mapping[Ty, int],
+                 n_layers: int,
+                 n_single_qubit_params: int = 3,
+                 discard: bool = False) -> None:
+        """Instantiate a Sim 9 ansatz.
+        Parameters
+        ----------
+        ob_map : dict
+            A mapping from :py:class:`discopy.rigid.Ty` to the number of
+            qubits it uses in a circuit.
+        n_layers : int
+            The number of layers used by the ansatz.
+        n_single_qubit_params : int, default: 3
+            The number of single qubit rotations used by the ansatz.
+        discard : bool, default: False
+            Discard open wires instead of post-selecting.
+        """
+        super().__init__(ob_map,
+                         n_layers,
+                         n_single_qubit_params,
+                         self.circuit,
+                         discard,
+                         [Rx, Rz])
+
+    def params_shape(self, n_qubits: int) -> tuple[int, ...]:
+        return (self.n_layers, n_qubits)
+
+    def circuit(self, n_qubits: int, params: np.ndarray) -> Circuit:
+        if n_qubits == 1:
+            circuit = Rx(params[0]) >> Rz(params[1]) >> Rx(params[2])
+        else:
+            circuit = Id(n_qubits)
+
+            for thetas in params:
+                circuit >>= Id().tensor(*(H for _ in range(n_qubits)))
+                czs = Id(n_qubits)
+                for q in range(n_qubits - 1):
+                    czs = czs.CZ(q, q + 1)
+                circuit >>= czs
+
+                circuit >>= Id().tensor(*map(Rx, thetas))
+        return circuit  # type: ignore[return-value]

lambeq/training/__init__.py

@@ -32,6 +32,9 @@
            'PytorchTrainer',
            'QuantumTrainer',
 
+           'TnPathOptimizer',
+           'CachedTnPathOptimizer',
+
            'BinaryCrossEntropyLoss',
            'CrossEntropyLoss',
            'LossFunction',
@@ -57,6 +60,8 @@
 from lambeq.training.trainer import Trainer
 from lambeq.training.pytorch_trainer import PytorchTrainer
 from lambeq.training.quantum_trainer import QuantumTrainer
+from lambeq.training.cached_tn_path_optimizer import (CachedTnPathOptimizer,
+                                                      TnPathOptimizer)
 
 from lambeq.training.loss import (BinaryCrossEntropyLoss,
                                   CrossEntropyLoss,