Introducing trainable-encoding layers

src/qiboml/interfaces/pytorch.py

@@ -15,7 +15,7 @@
 class QuantumModel(torch.nn.Module):
 
     encoding: QuantumEncoding
-    circuit: Circuit
+    trainable_circuit: Circuit
     decoding: QuantumDecoding
     differentiation_rule: DifferentiationRule = None
 
@@ -24,11 +24,12 @@ def __post_init__(
     ):
         super().__init__()
 
-        circuit = self.encoding.circuit
-        params = [p for param in self.circuit.get_parameters() for p in param]
+        circuit = self.encoding.circuit + self.trainable_circuit
+        params = [p for param in circuit.get_parameters() for p in param]
         params = torch.as_tensor(self.backend.to_numpy(x=params)).ravel()
         params.requires_grad = True
-        self.circuit_parameters = torch.nn.Parameter(params)
+        # all trainable parameters
+        self.trainable_parameters = torch.nn.Parameter(params)
 
     def forward(self, x: torch.Tensor):
         if (
@@ -39,15 +40,15 @@ def forward(self, x: torch.Tensor):
             x = QuantumModelAutoGrad.apply(
                 x,
                 self.encoding,
-                self.circuit,
+                self.trainable_circuit,
                 self.decoding,
                 self.backend,
                 self.differentiation_rule,
                 *list(self.parameters())[0],
             )
         else:
-            self.circuit.set_parameters(list(self.parameters())[0])
-            x = self.encoding(x) + self.circuit
+            x = self.encoding(x) + self.trainable_circuit
+            x.set_parameters(list(self.parameters())[0])
             x = self.decoding(x)
         return x
 
@@ -67,6 +68,10 @@ def backend(
     def output_shape(self):
         return self.decoding.output_shape
 
+    def draw(self):
+        circuit = self.encoding.circuit + self.trainable_circuit
+        circuit.draw()
+
 
 class QuantumModelAutoGrad(torch.autograd.Function):
 

src/qiboml/models/ansatze.py

@@ -1,23 +1,37 @@
 import random
+from typing import List, Optional
 
 import numpy as np
 from qibo import Circuit, gates
 
 
-def ReuploadingCircuit(
-    nqubits: int, qubits: list[int] = None, nlayers: int = 1
-) -> Circuit:
+def entangling_circuit(nqubits: int, entangling_gate: gates.Gate = gates.CNOT):
+    """Construct entangling layer."""
+    circuit = Circuit(nqubits)
+    for q in range(nqubits):
+        circuit.add(entangling_gate(q0=q % nqubits, q1=(q + 1) % nqubits))
+    return circuit
+
+
+def layered_ansatz(
+    nqubits: int,
+    nlayers: int = 1,
+    qubits: list[int] = None,
+    gates_list: Optional[List[gates.Gate]] = [
+        gates.RY,
+        gates.RZ,
+    ],  # TODO: this has to be a circuit
+    entanglement: bool = True,
+):
     if qubits is None:
         qubits = list(range(nqubits))
 
     circuit = Circuit(nqubits)
-
     for _ in range(nlayers):
         for q in qubits:
-            circuit.add(gates.RY(q, theta=random.random() * np.pi, trainable=True))
-            circuit.add(gates.RZ(q, theta=random.random() * np.pi, trainable=True))
-        for i, q in enumerate(qubits[:-2]):
-            circuit.add(gates.CNOT(q0=q, q1=qubits[i + 1]))
-        circuit.add(gates.CNOT(q0=qubits[-1], q1=qubits[0]))
+            for gate in gates_list:
+                circuit.add(gate(q, theta=random.random(), trainable=True))
+        if entanglement:
+            circuit += entangling_circuit(nqubits, gates.CNOT)
 
     return circuit

src/qiboml/models/encoding.py

@@ -6,6 +6,7 @@
 from qibo.config import raise_error
 
 from qiboml import ndarray
+from qiboml.models.ansatze import layered_ansatz
 
 
 @dataclass
@@ -66,3 +67,63 @@ def __call__(self, x: ndarray) -> Circuit:
         for bit in ones:
             circuit.add(gates.X(self.qubits[bit]))
         return circuit
+
+
+@dataclass
+class ReuploadingEncoding(QuantumEncoding):
+    """
+    Implementing reuploading scheme alternating encoding U and training V layers.
+    It follows the scheme V - U - V - U - V (in case of 2 layers), namely upload `nlayers` times x and
+    and each encoding layer is preceded and followed by a trainable layer V.
+    The chosen default V layer is a
+    `qiboml.models.ansatze.layered_ansatz(nqubits, 1, qubits, [RY, RZ], True)`.
+    """
+
+    # big TODO: make this model more flexible enabling e.g. lambda function of data and params
+
+    # TODO: rm this and raise error when calling the class
+    data_shape: tuple = (1,)
+    nlayers: int = 1
+    trainable_ansatz: Circuit = None
+    encoding_gate: gates.Gate = gates.RX
+
+    def __post_init__(
+        self,
+    ):
+
+        super().__post_init__()
+
+        data_dim = np.prod(self.data_shape, axis=0)
+
+        if int(data_dim) % len(self.qubits) != 0:
+            raise_error(
+                ValueError,
+                f"The data dimension has to be equal to the length of the chosen {self.qubits} subset of the {self.nqubits} system.",
+            )
+
+        # TODO: use deepcopy to repeat the call creating single elements
+        if self.trainable_ansatz is None:
+            self._circuit += layered_ansatz(nqubits=self.nqubits, qubits=self.qubits)
+        for _ in range(self.nlayers):
+            for q in self.qubits:
+                self._circuit.add(self.encoding_gate(q=q, theta=0.0, trainable=False))
+            if self.trainable_ansatz is None:
+                self._circuit += layered_ansatz(
+                    nqubits=self.nqubits, qubits=self.qubits
+                )
+
+    def _set_phases(self, x: ndarray):
+        encoding_gates = [
+            g for g in self._circuit.parametrized_gates if g.trainable == False
+        ]
+        data_length = len(x.ravel())
+        for l in range(self.nlayers):
+            for gate, phase in zip(
+                encoding_gates[data_length * l : data_length * l + data_length],
+                x.ravel(),
+            ):
+                gate.parameters = phase
+
+    def __call__(self, x: ndarray) -> Circuit:
+        self._set_phases(x)
+        return self._circuit