Merge pull request #275 from king-p3nguin/qiskit2tq-parameterexpression

01110011011101010110010001101111 · web-flow · commit 91c80def7558 · 2024-06-13T04:42:20.000-05:00
diff --git a/test/plugin/test_qiskit2tq.py b/test/plugin/test_qiskit2tq.py
@@ -0,0 +1,174 @@
+"""
+MIT License
+
+Copyright (c) 2020-present TorchQuantum Authors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+"""
+
+import random
+
+import numpy as np
+import pytest
+import torch
+import torch.optim as optim
+from qiskit import QuantumCircuit
+from qiskit.circuit import Parameter, ParameterVector
+from torch.optim.lr_scheduler import CosineAnnealingLR
+
+import torchquantum as tq
+from torchquantum.plugin import qiskit2tq
+
+seed = 42
+random.seed(seed)
+np.random.seed(seed)
+torch.manual_seed(seed)
+
+
+class TQModel(tq.QuantumModule):
+    def __init__(self, init_params=None):
+        super().__init__()
+        self.n_wires = 2
+        self.rx = tq.RX(has_params=True, trainable=True, init_params=[init_params[0]])
+        self.u3_0 = tq.U3(has_params=True, trainable=True, init_params=init_params[1:4])
+        self.u3_1 = tq.U3(
+            has_params=True,
+            trainable=True,
+            init_params=torch.tensor(
+                [
+                    init_params[4] + init_params[2],
+                    init_params[5] * init_params[3],
+                    init_params[6] * init_params[1],
+                ]
+            ),
+        )
+        self.cu3_0 = tq.CU3(
+            has_params=True,
+            trainable=True,
+            init_params=torch.tensor(
+                [
+                    torch.sin(init_params[7]),
+                    torch.abs(torch.sin(init_params[8])),
+                    torch.abs(torch.sin(init_params[9]))
+                    * torch.exp(init_params[2] + init_params[3]),
+                ]
+            ),
+        )
+
+    def forward(self, q_device: tq.QuantumDevice):
+        q_device.reset_states(1)
+        self.rx(q_device, wires=0)
+        self.u3_0(q_device, wires=0)
+        self.u3_1(q_device, wires=1)
+        self.cu3_0(q_device, wires=[0, 1])
+
+
+def get_qiskit_ansatz():
+    ansatz = QuantumCircuit(2)
+    ansatz_param = Parameter("Θ")  # parameter
+    ansatz.rx(ansatz_param, 0)
+    ansatz_param_vector = ParameterVector("φ", 9)  # parameter vector
+    ansatz.u(ansatz_param_vector[0], ansatz_param_vector[1], ansatz_param_vector[2], 0)
+    ansatz.u(
+        ansatz_param_vector[3] + ansatz_param_vector[1],  # parameter expression
+        ansatz_param_vector[4] * ansatz_param_vector[2],
+        ansatz_param_vector[5] / ansatz_param_vector[0],
+        1,
+    )
+    ansatz.cu(
+        np.sin(ansatz_param_vector[6]),  # numpy functions
+        np.abs(np.sin(ansatz_param_vector[7])),  # nested numpy functions
+        # complex expression
+        np.abs(np.sin(ansatz_param_vector[8]))
+        * np.exp(ansatz_param_vector[1] + ansatz_param_vector[2]),
+        0.0,
+        0,
+        1,
+    )
+    return ansatz
+
+
+def train_step(target_state, device, model, optimizer):
+    model(device)
+    result_state = device.get_states_1d()[0]
+
+    # compute the state infidelity
+    loss = 1 - torch.dot(result_state, target_state).abs() ** 2
+
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
+
+    infidelity = loss.item()
+    target_state_vector = target_state.detach().cpu().numpy()
+    result_state_vector = result_state.detach().cpu().numpy()
+    print(
+        f"infidelity (loss): {infidelity}, \n target state : "
+        f"{target_state_vector}, \n "
+        f"result state : {result_state_vector}\n"
+    )
+    return infidelity, target_state_vector, result_state_vector
+
+
+def train(init_params, backend):
+    device = torch.device("cpu")
+
+    if backend == "qiskit":
+        ansatz = get_qiskit_ansatz()
+        model = qiskit2tq(ansatz, initial_parameters=init_params).to(device)
+    elif backend == "torchquantum":
+        model = TQModel(init_params).to(device)
+
+    print(f"{backend} model:", model)
+
+    n_epochs = 10
+    optimizer = optim.Adam(model.parameters(), lr=1e-2, weight_decay=0)
+    scheduler = CosineAnnealingLR(optimizer, T_max=n_epochs)
+
+    q_device = tq.QuantumDevice(n_wires=2)
+    target_state = torch.tensor([0, 1, 0, 0], dtype=torch.complex64)
+
+    result_list = []
+    for epoch in range(1, n_epochs + 1):
+        print(f"Epoch {epoch}, LR: {optimizer.param_groups[0]['lr']}")
+        result_list.append(train_step(target_state, q_device, model, optimizer))
+        scheduler.step()
+
+    return result_list
+
+
+@pytest.mark.parametrize(
+    "init_params",
+    [
+        torch.nn.init.uniform_(torch.ones(10), -np.pi, np.pi),
+        torch.nn.init.uniform_(torch.ones(10), -np.pi, np.pi),
+        torch.nn.init.uniform_(torch.ones(10), -np.pi, np.pi),
+    ],
+)
+def test_qiskit2tq(init_params):
+    qiskit_result = train(init_params, "qiskit")
+    tq_result = train(init_params, "torchquantum")
+    for qi_tensor, tq_tensor in zip(qiskit_result, tq_result):
+        torch.testing.assert_close(qi_tensor[0], tq_tensor[0])
+        torch.testing.assert_close(qi_tensor[1], tq_tensor[1])
+        torch.testing.assert_close(qi_tensor[2], tq_tensor[2])
+
+
+if __name__ == "__main__":
+    test_qiskit2tq(torch.nn.init.uniform_(torch.ones(10), -np.pi, np.pi))
diff --git a/torchquantum/layer/layers/module_from_ops.py b/torchquantum/layer/layers/module_from_ops.py
@@ -22,16 +22,16 @@
 SOFTWARE.
 """
 
+from typing import Iterable
+
+import numpy as np
 import torch
 import torch.nn as nn
+from torchpack.utils.logging import logger
+
 import torchquantum as tq
 import torchquantum.functional as tqf
-import numpy as np
-
-
-from typing import Iterable
 from torchquantum.plugin.qiskit import QISKIT_INCOMPATIBLE_FUNC_NAMES
-from torchpack.utils.logging import logger
 
 __all__ = [
     "QuantumModuleFromOps",
@@ -61,6 +61,6 @@ def forward(self, q_device: tq.QuantumDevice):
             None
 
         """
-        self.q_device = q_device
+        q_device.reset_states(1)
         for op in self.ops:
-            op(q_device)
+            op(q_device, wires=op.wires)
diff --git a/torchquantum/plugin/qiskit/qiskit_plugin.py b/torchquantum/plugin/qiskit/qiskit_plugin.py
@@ -21,15 +21,20 @@
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
 """
+from __future__ import annotations
 
-from typing import Iterable
+from typing import Iterable, Optional
 
 import numpy as np
 import qiskit
 import qiskit.circuit.library.standard_gates as qiskit_gate
+import symengine
+import sympy
 import torch
 from qiskit import ClassicalRegister, QuantumCircuit, transpile
-from qiskit.circuit import Parameter
+from qiskit.circuit import CircuitInstruction, Parameter
+from qiskit.circuit.parameter import ParameterExpression
+from qiskit.circuit.parametervector import ParameterVectorElement
 from qiskit_aer import AerSimulator
 from torchpack.utils.logging import logger
 
@@ -691,7 +696,7 @@ def op_history2qiskit_expand_params(n_wires, op_history, bsz):
 
 # construct a tq QuantumModule object according to the qiskit QuantumCircuit
 # object
-def qiskit2tq_Operator(circ: QuantumCircuit):
+def qiskit2tq_Operator(circ: QuantumCircuit, initial_parameters=None):
     if getattr(circ, "_layout", None) is not None:
         try:
             p2v_orig = circ._layout.final_layout.get_physical_bits().copy()
@@ -711,14 +716,23 @@ def qiskit2tq_Operator(circ: QuantumCircuit):
         for p in range(circ.num_qubits):
             p2v[p] = p
 
+    if initial_parameters is None:
+        initial_parameters = torch.nn.init.uniform_(
+            torch.ones(len(circ.parameters)), -np.pi, np.pi
+        )
+
+    param_to_index = {}
+    for i, param in enumerate(circ.parameters):
+        param_to_index[param] = i
+
     ops = []
     for gate in circ.data:
         op_name = gate[0].name
         wires = [circ.find_bit(qb).index for qb in gate.qubits]
         wires = [p2v[wire] for wire in wires]
-        # sometimes the gate.params is ParameterExpression class
-        init_params = (
-            list(map(float, gate[0].params)) if len(gate[0].params) > 0 else None
+
+        init_params = qiskit2tq_translate_qiskit_params(
+            gate, initial_parameters, param_to_index
         )
 
         if op_name in [
@@ -780,8 +794,53 @@ def qiskit2tq_Operator(circ: QuantumCircuit):
     return ops
 
 
-def qiskit2tq(circ: QuantumCircuit):
-    ops = qiskit2tq_Operator(circ)
+def qiskit2tq_translate_qiskit_params(
+    circuit_instruction: CircuitInstruction, initial_parameters, param_to_index
+):
+    parameters = []
+    for p in circuit_instruction.operation.params:
+        if isinstance(p, Parameter) or isinstance(p, ParameterVectorElement):
+            parameters.append(initial_parameters[param_to_index[p]])
+        elif isinstance(p, ParameterExpression):
+            if len(p.parameters) == 0:
+                parameters.append(float(p))
+                continue
+
+            expr = p.sympify().simplify()
+            if isinstance(expr, symengine.Expr):  # qiskit uses symengine if available
+                expr = expr._sympy_()  # sympy.Expr
+
+            for free_symbol in expr.free_symbols:
+                # replace names: theta[0] -> theta_0
+                # ParameterVector creates symbols with brackets like theta[0]
+                # but sympy.lambdify does not allow brackets in symbol names
+                free_symbol.name = free_symbol.name.replace("[", "_").replace("]", "")
+
+            parameter_list = list(p.parameters)
+            sympy_symbols = [param._symbol_expr for param in parameter_list]
+            # replace names again: theta[0] -> theta_0
+            sympy_symbols = [
+                sympy.Symbol(str(symbol).replace("[", "_").replace("]", ""))
+                for symbol in sympy_symbols
+            ]
+            lam_f = sympy.lambdify(sympy_symbols, expr, modules="math")
+            parameters.append(
+                lam_f(
+                    *[
+                        initial_parameters[param_to_index[param]]
+                        for param in parameter_list
+                    ]
+                )
+            )
+        else:  # non-parameterized gate
+            parameters.append(p)
+    return parameters
+
+
+def qiskit2tq(
+    circ: QuantumCircuit, initial_parameters: Optional[list[torch.nn.Parameter]] = None
+):
+    ops = qiskit2tq_Operator(circ, initial_parameters)
     return tq.QuantumModuleFromOps(ops)
 
 
