Add parametric waveform control and TQASM export support

Author: EverestYan

tensorcircuit/__init__.py

@@ -24,8 +24,9 @@
 )  # prerun of set hooks
 from . import gates
 from . import basecircuit
+from . import waveforms
 from .gates import Gate
-from .circuit import Circuit, expectation
+from .circuit import Circuit, expectation, Param
 from .mpscircuit import MPSCircuit
 from .densitymatrix import DMCircuit as DMCircuit_reference
 from .densitymatrix import DMCircuit2

tensorcircuit/abstractcircuit.py

@@ -79,6 +79,14 @@ class AbstractCircuit:
     mpogates = mpogates
     gate_aliases = gate_aliases
 
+    def __init__(self, nqubits: int, *args, **kwargs):
+        self._nqubits = nqubits
+        self._qir = []
+        self._extra_qir = []
+        self.inputs = None
+        self.circuit_param = {}
+        self.is_mps = False
+
     def apply_general_gate(
         self,
         gate: Union[Gate, QuOperator],
@@ -114,6 +122,7 @@ def apply(self: "AbstractCircuit", *index: int, **vars: Any) -> None:
             gate_dict = {
                 "gatef": gatef,
                 "index": index,
+                "target": list(index),
                 "name": localname,
                 "split": split,
                 "mpo": mpo,

tensorcircuit/basecircuit.py

@@ -40,6 +40,9 @@ class BaseCircuit(AbstractCircuit):
 
     is_mps = False
 
+    def __init__(self, nqubits: int, *args, **kwargs):
+        super().__init__(nqubits, *args, **kwargs)
+
     @staticmethod
     def all_zero_nodes(n: int, d: int = 2, prefix: str = "qb-") -> List[tn.Node]:
         l = [0.0 for _ in range(d)]

tensorcircuit/circuit.py

@@ -60,6 +60,11 @@ def __init__(
             ``max_singular_values`` and ``max_truncation_err``.
         :type split: Optional[Dict[str, Any]]
         """
+        super().__init__(
+            nqubits=nqubits, inputs=inputs, mps_inputs=mps_inputs, split=split
+        )
+        self.calibrations = []
+
         self.inputs = inputs
         self.mps_inputs = mps_inputs
         self.split = split
@@ -111,6 +116,43 @@ def __init__(
         # self._qcode += str(self._nqubits) + "\n"
         self._qir: List[Dict[str, Any]] = []
         self._extra_qir: List[Dict[str, Any]] = []
+    
+    def add_calibration(
+        self, name: str, parameters: List[str], instructions: List[Dict]
+    ) -> None:
+        self.calibrations.append({
+            "name": name,
+            "parameters": parameters,
+            "instructions": instructions
+        })
+
+
+    def to_tqasm(self) -> str:
+        qasm_lines = []
+        qasm_lines.append("OPENQASM 3;")
+        qasm_lines.append(f"qubit[{self._nqubits}] q;")
+
+        for gate in self._qir:
+            gname = gate["name"]
+            targets = ", ".join(f"q[{i}]" for i in gate["target"])
+            qasm_lines.append(f"{gname} {targets};")
+
+        for cal in getattr(self, "calibrations", []):
+            pname = ", ".join(cal["parameters"])
+            qasm_lines.append(f"\ndefcal {cal['name']}({pname}) {{")
+            for inst in cal["instructions"]:
+                if inst["type"] == "frame":
+                    qasm_lines.append(f"  frame {inst['frame']} = newframe({inst['qubit']});")
+                elif inst["type"] == "play":
+                    args_str = ", ".join(str(x) for x in inst["args"])
+                    wf_type = inst["waveform_type"]
+                    qasm_lines.append(f"  play({inst['frame']}, {wf_type}({args_str}));")
+            qasm_lines.append("}")
+
+        return "\n".join(qasm_lines)
+    
+    def calibrate(self, name: str, parameters: List["Param"]) -> "DefcalBuilder":
+        return DefcalBuilder(self, name, parameters)
 
     def replace_mps_inputs(self, mps_inputs: QuOperator) -> None:
         """
@@ -998,3 +1040,46 @@ def expectation(
         else:
             den = 1.0
         return num / den
+
+class Param:
+    def __init__(self, name: str):
+        self.name = name
+    
+class DefcalBuilder:
+    def __init__(self, circuit, name: str, parameters: List["Param"]):
+        self.circuit = circuit
+        self.name = name
+        self.parameters = parameters
+        self.instructions = []
+
+    def new_frame(self, frame_name: str, param: "Param"):
+        self.instructions.append({
+            "type": "frame",
+            "frame": frame_name,
+            "qubit": param.name,
+        })
+        return self
+
+    def play(self, frame_name: str, waveform: Any, start_time: int = None):
+        if not hasattr(waveform, "__dataclass_fields__"):
+            raise TypeError("Unsupported waveform type")
+
+        waveform_type = waveform.qasm_name()
+        args = list(waveform.__dict__.values())
+        if start_time is not None:
+            args = [start_time] + args
+
+        self.instructions.append({
+            "type": "play",
+            "frame": frame_name,
+            "waveform_type": waveform_type,
+            "args": args,
+        })
+        return self
+
+    def build(self):
+        self.circuit.add_calibration(
+            name=self.name,
+            parameters=[p.name for p in self.parameters],
+            instructions=self.instructions,
+        )

tensorcircuit/waveforms.py

@@ -0,0 +1,85 @@
+from typing import List, Union
+from dataclasses import dataclass
+
+ParamType = Union[float, str]
+
+__all__ = [
+    "Gaussian",
+    "GaussianSquare",
+    "Drag",
+    "Constant",
+    "Sine",
+    "Cosine",
+    "CosineDrag",
+    "Flattop",
+]
+
+@dataclass
+class Gaussian:
+    amp: ParamType
+    duration: int
+    sigma: ParamType
+    def qasm_name(self) -> str:
+        return "gaussian"
+
+@dataclass
+class GaussianSquare:
+    amp: ParamType
+    duration: int
+    sigma: ParamType
+    width: ParamType
+    def qasm_name(self) -> str:
+        return "gaussian_square"
+
+@dataclass
+class Drag:
+    amp: ParamType
+    duration: int
+    sigma: ParamType
+    beta: ParamType
+    def qasm_name(self) -> str:
+        return "drag"
+
+@dataclass
+class Constant:
+    amp: ParamType
+    duration: int
+    def qasm_name(self) -> str:
+        return "constant"
+
+
+@dataclass
+class Sine:
+    amp: ParamType
+    frequency: ParamType
+    duration: int
+    def qasm_name(self) -> str:
+        return "sine"
+
+
+@dataclass
+class Cosine:
+    amp: ParamType
+    frequency: ParamType
+    duration: int
+    def qasm_name(self) -> str:
+        return "cosine"
+
+
+@dataclass
+class CosineDrag:
+    amp: ParamType
+    duration: int
+    phase: ParamType
+    alpha: ParamType
+    def qasm_name(self) -> str:
+        return "cosine_drag"
+
+
+@dataclass
+class Flattop:
+    amp: ParamType
+    width: ParamType
+    duration: int
+    def qasm_name(self) -> str:
+        return "flattop"

tests/test_parametric_waveform.py

@@ -0,0 +1,54 @@
+import sys
+import os
+
+# Add the directory containing your module to Python's search path
+module_path = "/Users/mac/tensorcircuit"
+sys.path.insert(0, module_path)
+
+from tensorcircuit import Circuit, Param, gates, waveforms
+from tensorcircuit.cloud.apis import submit_task, get_device
+import re
+
+print("✅ TEST FILE LOADED")
+
+def test_parametric_waveform():
+    qc = Circuit(2)
+
+    param0 = Param("param0")
+    param1 = Param("param1")
+
+    builder = qc.calibrate("my_gate", [param0, param1])
+    builder.new_frame("f0", param0)
+    builder.play("f0", waveforms.CosineDrag(10, 0.2, "0.5*PI", 0.01))
+    builder.new_frame("f1", param1)
+    builder.play("f1", waveforms.CosineDrag(20, 0.01, "0", 0))
+    builder.build()
+
+    tqasm_code = qc.to_tqasm()
+    print(tqasm_code)  
+
+    assert "OPENQASM 3;" in tqasm_code
+    assert "defcal my_gate(param0, param1)" in tqasm_code
+    assert "frame f0 = newframe(param0);" in tqasm_code
+    assert "play(f0, cosine_drag(10, 0.2, 0.5*PI, 0.01));" in tqasm_code
+    assert "frame f1 = newframe(param1);" in tqasm_code
+    assert "play(f1, cosine_drag(20, 0.01, 0, 0));" in tqasm_code
+    assert re.search(r"defcal my_gate\([^\)]*\)\s*\{", tqasm_code)
+    
+    return 
+
+
+    device_name = "tianxuan_s1" 
+    d = get_device(device_name)
+    t = submit_task(
+    circuit=qc,
+    shots=8192,
+    device=d,
+    enable_qos_gate_decomposition=False,
+    enable_qos_qubit_mapping=False,
+    )
+    rf = t.results()
+    print(rf)
+
+
+test_parametric_waveform()

tests/test_waveform_types_to_tqasm.py

@@ -0,0 +1,48 @@
+import sys
+import os
+
+# Add the directory containing your module to Python's search path
+module_path = "/Users/mac/tensorcircuit"
+sys.path.insert(0, module_path)
+
+from tensorcircuit import Circuit, Param, waveforms
+
+def test_waveform_types_to_tqasm_():
+    test_cases = [
+        (waveforms.CosineDrag, [10, 0.2, "PI/2", 0.01], "cosine_drag"),
+        (waveforms.Gaussian, [0.5, 20, 5], "gaussian"),
+        (waveforms.GaussianSquare, [0.5, 20, 5, 10], "gaussian_square"),
+        (waveforms.Drag, [0.4, 30, 10, 0.2], "drag"),
+        (waveforms.Flattop, [1.0, 5, 20], "flattop"),
+        (waveforms.Constant, [0.7, 15], "constant"),
+        (waveforms.Sine, [0.3, "2*PI", 10], "sine"),
+        (waveforms.Cosine, [0.6, "PI", 8], "cosine"),
+    ]
+
+    failed = False
+
+    for waveform_class, args, expected_name in test_cases:
+        qc = Circuit(1)
+        p = Param("q0")
+
+        builder = qc.calibrate("test_gate", [p])
+        builder.new_frame("f0", p)
+        builder.play("f0", waveform_class(*args))
+        builder.build()
+
+        qasm = qc.to_tqasm()
+        print(f"--- {expected_name} ---\n{qasm}\n")
+
+        if f"play(f0, {expected_name}(" not in qasm:
+            print(f"❌ FAILED: waveform {waveform_class.__name__} expected '{expected_name}', but not found in QASM output.")
+            failed = True
+        else:
+            print(f"✅ PASSED: waveform {expected_name}")
+
+    if failed:
+        raise AssertionError("Some waveform QASM outputs were incorrect.")
+    else:
+        print("All waveform QASM outputs passed.")
+
+if __name__ == "__main__":
+    test_waveform_types_to_tqasm_()