updates to paulitwodesign, structure of twolocal circuits, and tests …

…for twolocal

test/layers/test_nlocal.py

@@ -29,9 +29,9 @@ def compare_tq_to_qiskit(tq_circuit, qiskit_circuit, instance_info = ""):
         for op in tq_circuit.op_history
     ]
 
-    # create tuples (NOTE: WILL LOSE ORDER SO NOT ENTIRELY CORRECT)
-    tq_ops_tuple = {tuple(op) for op in tq_ops}
-    qiskit_ops_tuple = {tuple(op) for op in qiskit_ops}
+    # create tuples, preserving order
+    tq_ops_tuple = [tuple(op) for op in tq_ops]
+    qiskit_ops_tuple = [tuple(op) for op in qiskit_ops]
 
     # assert if they are the same
     assert len(tq_ops) == len(
@@ -43,51 +43,54 @@ def compare_tq_to_qiskit(tq_circuit, qiskit_circuit, instance_info = ""):
 
 ## TEST TWOLOCAL
 
-# iterate through different parameters to test
-for entanglement_type in ("linear", "circular", "full"):
-    for n_wires in (3, 5, 10):
-        for reps in range(1, 5):
-            # create the TQ circuit
-            tq_two = tq.layer.TwoLocal(
-                [tq.RY, tq.RZ],
-                [tq.CZ],
-                arch={"n_wires": n_wires},
-                entanglement_layer=entanglement_type,
-                reps=reps,
-            )
-            qdev = tq.QuantumDevice(n_wires, record_op=True)
-            tq_two(qdev)
+def test_twolocal():
+    # iterate through different parameters to test
+    for entanglement_type in ("linear", "circular", "full"):
+        for n_wires in (3, 5, 10):
+            for reps in range(1, 5):
+                # create the TQ circuit
+                tq_two = tq.layer.TwoLocal(
+                    n_wires,
+                    ["ry", "rz"],
+                    "cz",
+                    entanglement_layer=entanglement_type,
+                    reps=reps,
+                )
+                qdev = tq.QuantumDevice(n_wires, record_op=True)
+                tq_two(qdev)
 
-            # create the qiskit circuit
-            qiskit_two = TwoLocal(
-                n_wires,
-                ["ry", "rz"],
-                "cz",
-                entanglement_type,
-                reps=reps,
-                insert_barriers=False,
-            )
-
-            # compare the circuits
-            test_info = f"{entanglement_type} with {n_wires} wires and {reps} reps"
-            compare_tq_to_qiskit(qdev, qiskit_two)
+                # create the qiskit circuit
+                qiskit_two = TwoLocal(
+                    n_wires,
+                    ["ry", "rz"],
+                    "cz",
+                    entanglement_type,
+                    reps=reps,
+                    insert_barriers=False,
+                )
+                
+                # compare the circuits
+                test_info = f"{entanglement_type} with {n_wires} wires and {reps} reps"
+                compare_tq_to_qiskit(qdev, qiskit_two)
 
 
 ## TEST OTHER CIRCUITS
 
-tq_to_qiskit = {
-    "EfficientSU2": (tq.layer.EfficientSU2, EfficientSU2),
-    "ExcitationPreserving": (tq.layer.ExcitationPreserving, ExcitationPreserving),
-    "RealAmplitudes": (tq.layer.RealAmplitudes, RealAmplitudes),
-}
+def test_twolocal_variants():
+    tq_to_qiskit = {
+        "EfficientSU2": (tq.layer.EfficientSU2, EfficientSU2),
+        "ExcitationPreserving": (tq.layer.ExcitationPreserving, ExcitationPreserving),
+        "RealAmplitudes": (tq.layer.RealAmplitudes, RealAmplitudes),
+        "PauliTwo": (tq.layer.PauliTwoDesign, PauliTwoDesign),
+    }
 
-# run all the tests
-for circuit_name in tq_to_qiskit:
-    tq_instance, qiskit_instance = tq_to_qiskit[circuit_name]
-    for n_wires in range(2, 5):
-        tq_circuit = tq_instance({"n_wires": n_wires})
-        circuit = qiskit_instance(n_wires)
-        qdev = tq.QuantumDevice(n_wires, record_op=True)
-        tq_circuit(qdev)
-        compare_tq_to_qiskit(qdev, circuit, f"{circuit_name} with {n_wires} wires")
+    # run all the tests
+    for circuit_name in tq_to_qiskit:
+        tq_instance, qiskit_instance = tq_to_qiskit[circuit_name]
+        for n_wires in range(2, 5):
+            tq_circuit = tq_instance(n_wires)
+            circuit = qiskit_instance(n_wires)
+            qdev = tq.QuantumDevice(n_wires, record_op=True)
+            tq_circuit(qdev)
+            compare_tq_to_qiskit(qdev, circuit, f"{circuit_name} with {n_wires} wires")
 

torchquantum/layer/layers.py

@@ -230,7 +230,7 @@ def forward(self, q_device: tq.QuantumDevice):
 
 class RandomOp1All(tq.QuantumModule):
     def __init__(
-        self, n_wires: int, op_types=(tq.RX, tq.RY, tq.RZ), op_ratios=None, seed=None
+        self, n_wires: int, op_types=(tq.RX, tq.RY, tq.RZ), op_ratios=None, has_params=True, trainable=True, seed=None
     ):
         """Layer adding a random gate to all wires
 
@@ -246,26 +246,18 @@ def __init__(
         self.op_ratios = op_ratios
         self.seed = seed
         self.gate_all = nn.ModuleList()
+
         if seed is not None:
             np.random.seed(seed)
-        self.build_random_layer()
 
-    def build_random_layer(self):
         for k in range(self.n_wires):
             op = np.random.choice(self.op_types, p=self.op_ratios)
-            self.gate_all.append(op())
+            self.gate_all.append(op(has_params=has_params, trainable=trainable))
 
     @tq.static_support
-    def forward(self, q_device: tq.QuantumDevice, x):
-        # op on all wires, assert the number of gate is the same as the number
-        # of wires in the device.
-        assert self.n_gate == q_device.n_wires, (
-            f"Number of gates ({self.n_wires}) is different from number "
-            f"of wires ({q_device.n_wires})!"
-        )
-
+    def forward(self, q_device: tq.QuantumDevice):
         for k in range(self.n_wires):
-            self.gate_all[k](q_device, wires=k, params=x[:, k])
+            self.gate_all[k](q_device, wires=k)
 
 
 class RandomLayer(tq.QuantumModule):

torchquantum/layer/nlocal/efficient_su2.py

@@ -42,14 +42,14 @@ class EfficientSU2(TwoLocal):
 
     def __init__(
         self,
-        arch: dict = None,
+        n_wires: int = None,
         entanglement_layer: str = "reverse_linear",
         reps: int = 3,
         skip_final_rotation_layer: bool = False,
     ):
         # construct circuit with rotation layers of RY and RZ and entanglement with CX
         super().__init__(
-            arch=arch,
+            n_wires = n_wires,
             rotation_ops=[tq.RY, tq.RZ],
             entanglement_ops=[tq.CNOT],
             entanglement_layer=entanglement_layer,

torchquantum/layer/nlocal/excitation_preserving.py

@@ -42,14 +42,14 @@ class ExcitationPreserving(TwoLocal):
 
     def __init__(
         self,
-        arch: dict = None,
+        n_wires: int = 1,
         entanglement_layer: str = "full",
         reps: int = 3,
         skip_final_rotation_layer: bool = False,
     ):
         # construct circuit with rotation layers of RZ and entanglement with RXX and RYY
         super().__init__(
-            arch=arch,
+            n_wires = n_wires,
             rotation_ops=[tq.RZ],
             entanglement_ops=[tq.RXX, tq.RYY],
             entanglement_layer=entanglement_layer,

torchquantum/layer/nlocal/pauli_two.py

@@ -44,7 +44,7 @@ class PauliTwoDesign(TwoLocal):
 
     def __init__(
         self,
-        arch: dict = None,
+        n_wires: int = 1,
         entanglement_layer: str = "reverse_linear",
         reps: int = 3,
         skip_final_rotation_layer: bool = False,
@@ -54,7 +54,7 @@ def __init__(
         self.seed = seed
         # construct circuit with entanglement with CX
         super().__init__(
-            arch=arch,
+            n_wires=n_wires,
             entanglement_ops=[tq.CNOT],
             entanglement_layer=entanglement_layer,
             reps=reps,

torchquantum/layer/nlocal/real_amplitudes.py

@@ -42,14 +42,14 @@ class RealAmplitudes(TwoLocal):
 
     def __init__(
         self,
-        arch: dict = None,
+        n_wires: int = 1,
         entanglement_layer: str = "reverse_linear",
         reps: int = 3,
         skip_final_rotation_layer: bool = False,
     ):
         # construct circuit with rotation layers of RY and entanglement with CX
         super().__init__(
-            arch=arch,
+            n_wires=n_wires,
             rotation_ops=[tq.RY],
             entanglement_ops=[tq.CNOT],
             entanglement_layer=entanglement_layer,

torchquantum/layer/nlocal/two_local.py

@@ -28,7 +28,9 @@
     Op2QAllLayer,
     Op2QDenseLayer,
 )
+from torchquantum.operator import op_name_dict
 from .nlocal import NLocal
+from collections.abc import Iterable
 
 __all__ = [
     "TwoLocal",
@@ -52,9 +54,9 @@ class TwoLocal(NLocal):
 
     def __init__(
         self,
+        n_wires: int = 1,
         rotation_ops: list = None,
         entanglement_ops: list = None,
-        arch: dict = None,
         rotation_layer: tq.QuantumModule = Op1QAllLayer,
         entanglement_layer: str = "linear",
         reps: int = 1,
@@ -74,9 +76,27 @@ def __init__(
         elif entanglement_layer == "full":
             entanglement_layer = Op2QDenseLayer
 
+        # handling different input types for the rotation ops
+        if isinstance(rotation_ops, str):
+            rotation_ops = [op_name_dict[rotation_ops]]
+        elif isinstance(rotation_ops, Iterable):
+            if all(isinstance(rot, str) for rot in rotation_ops):
+                rotation_ops = [op_name_dict[rot] for rot in rotation_ops]
+        else:
+            rotation_ops = [rotation_ops]
+
+        # handling different input types for the entanglment ops
+        if isinstance(entanglement_ops, str):
+            entanglement_ops = [op_name_dict[entanglement_ops]]
+        elif isinstance(entanglement_ops, Iterable):
+            if all(isinstance(op, str) for op in entanglement_ops):
+                entanglement_ops = [op_name_dict[op] for op in entanglement_ops]
+        else:
+            entanglement_ops = [entanglement_ops]
+
         # initialize
         super().__init__(
-            arch=arch,
+            arch={"n_wires": n_wires},
             rotation_ops=rotation_ops,
             rotation_layer=rotation_layer,
             rotation_layer_params={"has_params": True, "trainable": True},