allow passing qtrl configs to aqt_compile() (#1000)

(these will override pre-uploaded configs)

cirq-superstaq/cirq_superstaq/service.py

@@ -485,6 +485,8 @@ def aqt_compile(
         gate_defs: None | (
             Mapping[str, npt.NDArray[np.complex_] | cirq.Gate | cirq.Operation | None]
         ) = None,
+        pulses: object = None,
+        variables: object = None,
         **kwargs: Any,
     ) -> css.compiler_output.CompilerOutput:
         """Compiles and optimizes the given circuit(s) for the Advanced Quantum Testbed (AQT).
@@ -508,6 +510,8 @@ def aqt_compile(
                 "SWAP/C5C4": cirq.SQRT_ISWAP_INV}` implies `SQRT_ISWAP` for all "SWAP" calibrations
                 except "SWAP/C5C4" (which will instead be mapped to a `SQRT_ISWAP_INV` gate on
                 qubits 4 and 5). Setting any calibration to None will disable that calibration.
+            pulses: Qtrl `PulseManager` or file path for pulse configuration.
+            variables: Qtrl `VariableManager` or file path for variable configuration.
             kwargs: Other desired compile options.
 
         Returns:
@@ -550,6 +554,11 @@ def aqt_compile(
                     val = _to_matrix_gate(val)
                 gate_defs_cirq[key] = val
             options_dict["gate_defs"] = gate_defs_cirq
+        if pulses or variables:
+            options_dict["aqt_configs"] = {
+                "pulses": self._qtrl_config_to_yaml_str(pulses),
+                "variables": self._qtrl_config_to_yaml_str(variables),
+            }
 
         if options_dict:
             request_json["options"] = cirq.to_json(options_dict)

cirq-superstaq/cirq_superstaq/service_test.py

@@ -26,6 +26,7 @@
 import pandas as pd
 import pytest
 import sympy
+import yaml
 from general_superstaq import ResourceEstimate
 
 import cirq_superstaq as css
@@ -279,7 +280,13 @@ def test_service_aqt_compile_single(mock_post_request: mock.MagicMock) -> None:
         "CS2": cirq.unitary(cirq.CZ**0.49),
         "CS3": cirq.unitary(css.CZ3**0.5),
     }
-    out = service.aqt_compile(cirq.Circuit(), gate_defs=gate_defs, atol=1e-3)
+    out = service.aqt_compile(
+        cirq.Circuit(),
+        gate_defs=gate_defs,
+        atol=1e-3,
+        pulses={"foo": "bar"},
+        variables={"abc": 123},
+    )
 
     expected_options = {
         "atol": 1e-3,
@@ -290,6 +297,10 @@ def test_service_aqt_compile_single(mock_post_request: mock.MagicMock) -> None:
             "CS2": cirq.MatrixGate(cirq.unitary(cirq.CZ**0.49)),
             "CS3": cirq.MatrixGate(cirq.unitary(css.CZ3**0.5), qid_shape=(3, 3)),
         },
+        "aqt_configs": {
+            "pulses": yaml.dump({"foo": "bar"}),
+            "variables": yaml.dump({"abc": 123}),
+        },
     }
     mock_post_request.assert_called_with(
         "/aqt_compile",

general-superstaq/general_superstaq/service.py

@@ -191,7 +191,36 @@ def submit_qubo(
         result_dict = self._client.submit_qubo(qubo, target, repetitions, method, max_solutions)
         return gss.serialization.deserialize(result_dict["solution"])
 
-    def aqt_upload_configs(self, pulses: Any, variables: Any) -> str:
+    @staticmethod
+    def _qtrl_config_to_yaml_str(config: object) -> str:
+        if isinstance(config, str):
+            if not os.path.isfile(config):
+                raise ValueError(f"{config!r} is not a valid file path.")
+
+            with open(config) as config_file:
+                return config_file.read()
+
+        config = getattr(config, "_config_raw", config)  # required to serialize qtrl Managers
+        if isinstance(config, dict):
+            try:
+                import yaml
+
+                return yaml.safe_dump(config)
+
+            except ImportError:
+                raise ModuleNotFoundError(
+                    "The PyYAML package is required to upload AQT configurations from dicts. "
+                    "You can install it using 'pip install pyyaml'."
+                )
+            except yaml.YAMLError:
+                pass
+
+        raise ValueError(
+            "Unable to serialize configuration. AQT configs should be qtrl Manager instances "
+            "or valid file paths."
+        )
+
+    def aqt_upload_configs(self, pulses: object, variables: object) -> str:
         """Uploads configs for AQT.
 
         Arguments can be either file paths (in .yaml format) or qtrl Manager instances.
@@ -203,36 +232,8 @@ def aqt_upload_configs(self, pulses: Any, variables: Any) -> str:
         Returns:
             A status of the update (whether or not it failed).
         """
-
-        def _config_to_yaml_str(config: Any) -> str:
-            if isinstance(config, str):
-                if not os.path.isfile(config):
-                    raise ValueError(f"{config!r} is not a valid file path.")
-
-                with open(config) as config_file:
-                    return config_file.read()
-
-            config = getattr(config, "_config_raw", config)  # required to serialize qtrl Managers
-            if isinstance(config, dict):
-                try:
-                    import yaml
-
-                    return yaml.safe_dump(config)
-                except ImportError:
-                    raise ModuleNotFoundError(
-                        "The PyYAML package is required to upload AQT configurations from dicts. "
-                        "You can install it using 'pip install pyyaml'."
-                    )
-                except yaml.YAMLError:
-                    pass
-
-            raise ValueError(
-                "Unable to serialize configuration. AQT configs should be qtrl Manager instances "
-                "or valid file paths."
-            )
-
-        pulses_yaml = _config_to_yaml_str(pulses)
-        variables_yaml = _config_to_yaml_str(variables)
+        pulses_yaml = self._qtrl_config_to_yaml_str(pulses)
+        variables_yaml = self._qtrl_config_to_yaml_str(variables)
 
         return self._client.aqt_upload_configs({"pulses": pulses_yaml, "variables": variables_yaml})
 

qiskit-superstaq/qiskit_superstaq/superstaq_backend.py

@@ -204,6 +204,8 @@ def aqt_compile(
         random_seed: int | None = None,
         atol: float | None = None,
         gate_defs: Mapping[str, str | npt.NDArray[np.complex_] | None] | None = None,
+        pulses: object = None,
+        variables: object = None,
         **kwargs: Any,
     ) -> qss.compiler_output.CompilerOutput:
         """Compiles and optimizes the given circuit(s) for the Advanced Quantum Testbed (AQT).
@@ -226,6 +228,8 @@ def aqt_compile(
                 `<matrix1>` for all "SWAP" calibrations except "SWAP/C5C4" (which will instead be
                 mapped to `<matrix2>` applied to qubits 4 and 5). Setting any calibration to None
                 will disable that calibration.
+            pulses: Qtrl `PulseManager` or file path for pulse configuration.
+            variables: Qtrl `VariableManager` or file path for variable configuration.
             kwargs: Other desired compile options.
 
         Returns:
@@ -251,6 +255,11 @@ def aqt_compile(
             options["atol"] = float(atol)
         if gate_defs is not None:
             options["gate_defs"] = gate_defs
+        if pulses or variables:
+            options["aqt_configs"] = {
+                "pulses": self._provider._qtrl_config_to_yaml_str(pulses),
+                "variables": self._provider._qtrl_config_to_yaml_str(variables),
+            }
 
         request_json = self._get_compile_request_json(circuits, **options)
         circuits_is_list = not isinstance(circuits, qiskit.QuantumCircuit)

qiskit-superstaq/qiskit_superstaq/superstaq_backend_test.py

@@ -9,6 +9,7 @@
 import general_superstaq as gss
 import pytest
 import qiskit
+import yaml
 
 import qiskit_superstaq as qss
 
@@ -131,19 +132,26 @@ def test_aqt_compile(mock_post: MagicMock) -> None:
         },
     )
 
-    out = backend.compile([qc], atol=1e-2)
+    out = backend.compile([qc], atol=1e-2, pulses={"foo": "bar"}, variables={"abc": 123})
     assert out.circuits == [qc]
     assert out.initial_logical_to_physicals == [{0: 1}]
     assert out.final_logical_to_physicals == [{1: 4}]
     assert not hasattr(out, "circuit")
+    expected_options = {
+        "atol": 1e-2,
+        "aqt_configs": {
+            "pulses": yaml.dump({"foo": "bar"}),
+            "variables": yaml.dump({"abc": 123}),
+        },
+    }
     mock_post.assert_called_with(
         f"{provider._client.url}/aqt_compile",
         headers=provider._client.headers,
         verify=provider._client.verify_https,
         json={
             "qiskit_circuits": qss.serialize_circuits(qc),
             "target": "aqt_keysight_qpu",
-            "options": json.dumps({"atol": 1e-2}),
+            "options": json.dumps(expected_options),
         },
     )
 

qiskit-superstaq/qiskit_superstaq/superstaq_provider.py

@@ -188,6 +188,8 @@ def aqt_compile(
         random_seed: int | None = None,
         atol: float | None = None,
         gate_defs: Mapping[str, str | npt.NDArray[np.complex_] | None] | None = None,
+        pulses: object = None,
+        variables: object = None,
         **kwargs: Any,
     ) -> qss.compiler_output.CompilerOutput:
         """Compiles and optimizes the given circuit(s) for the Advanced Quantum Testbed (AQT).
@@ -211,6 +213,8 @@ def aqt_compile(
                 `<matrix1>` for all "SWAP" calibrations except "SWAP/C5C4" (which will instead be
                 mapped to `<matrix2>` applied to qubits 4 and 5). Setting any calibration to None
                 will disable that calibration.
+            pulses: Qtrl `PulseManager` or file path for pulse configuration.
+            variables: Qtrl `VariableManager` or file path for variable configuration.
             kwargs: Other desired compile options.
 
         Returns:
@@ -231,6 +235,8 @@ def aqt_compile(
             random_seed=random_seed,
             atol=atol,
             gate_defs=gate_defs,
+            pulses=pulses,
+            variables=variables,
             **kwargs,
         )
 

qiskit-superstaq/qiskit_superstaq/superstaq_provider_test.py

@@ -12,6 +12,7 @@
 import numpy as np
 import pytest
 import qiskit
+import yaml
 from general_superstaq import ResourceEstimate, testing
 
 import qiskit_superstaq as qss
@@ -72,11 +73,30 @@ def test_aqt_compile(mock_post: MagicMock, fake_superstaq_provider: MockSupersta
     assert out.final_logical_to_physical == {1: 4}
     assert not hasattr(out, "circuits")
 
-    out = fake_superstaq_provider.aqt_compile([qc], atol=1e-2)
+    out = fake_superstaq_provider.aqt_compile(
+        [qc], atol=1e-2, pulses={"foo": "bar"}, variables={"abc": 123}
+    )
     assert out.circuits == [qc]
     assert out.initial_logical_to_physicals == [{0: 1}]
     assert out.final_logical_to_physicals == [{1: 4}]
     assert not hasattr(out, "circuit")
+    expected_options = {
+        "atol": 1e-2,
+        "aqt_configs": {
+            "pulses": yaml.dump({"foo": "bar"}),
+            "variables": yaml.dump({"abc": 123}),
+        },
+    }
+    mock_post.assert_called_with(
+        f"{fake_superstaq_provider._client.url}/aqt_compile",
+        headers=fake_superstaq_provider._client.headers,
+        verify=fake_superstaq_provider._client.verify_https,
+        json={
+            "qiskit_circuits": qss.serialize_circuits(qc),
+            "target": "aqt_keysight_qpu",
+            "options": json.dumps(expected_options),
+        },
+    )
 
     mock_post.return_value.json = lambda: {
         "qiskit_circuits": qss.serialization.serialize_circuits([qc, qc]),