From 0351323eb05b8f712ee046bd192d19d3c7956316 Mon Sep 17 00:00:00 2001
From: Andrea <andreapasquale97@gmail.com>
Date: Wed, 24 Jan 2024 19:09:12 +0400
Subject: [PATCH] refactor: Introduce folders for platform serialization

---
 src/qibolab/__init__.py                 | 11 +++---
 src/qibolab/dummy.py                    | 18 ++++-----
 src/qibolab/{ => dummy}/dummy.yml       |  0
 src/qibolab/kernels.py                  |  5 ++-
 src/qibolab/qubits.py                   |  2 +-
 src/qibolab/serialize.py                | 49 ++++++++++++++++---------
 tests/dummy_qrc/{ => qblox}/qblox.py    |  7 ++--
 tests/dummy_qrc/{ => qblox}/qblox.yml   |  0
 tests/dummy_qrc/{ => qm}/qm.py          |  7 ++--
 tests/dummy_qrc/{ => qm}/qm.yml         |  0
 tests/dummy_qrc/{ => rfsoc}/rfsoc.py    |  7 ++--
 tests/dummy_qrc/{ => rfsoc}/rfsoc.yml   |  0
 tests/dummy_qrc/{ => zurich}/zurich.py  | 16 ++++----
 tests/dummy_qrc/{ => zurich}/zurich.yml |  0
 tests/test_dummy.py                     |  1 -
 tests/test_platform.py                  | 12 +++---
 16 files changed, 79 insertions(+), 56 deletions(-)
 rename src/qibolab/{ => dummy}/dummy.yml (100%)
 rename tests/dummy_qrc/{ => qblox}/qblox.py (96%)
 rename tests/dummy_qrc/{ => qblox}/qblox.yml (100%)
 rename tests/dummy_qrc/{ => qm}/qm.py (96%)
 rename tests/dummy_qrc/{ => qm}/qm.yml (100%)
 rename tests/dummy_qrc/{ => rfsoc}/rfsoc.py (91%)
 rename tests/dummy_qrc/{ => rfsoc}/rfsoc.yml (100%)
 rename tests/dummy_qrc/{ => zurich}/zurich.py (93%)
 rename tests/dummy_qrc/{ => zurich}/zurich.yml (100%)

diff --git a/src/qibolab/__init__.py b/src/qibolab/__init__.py
index 8853188571..ebc080d299 100644
--- a/src/qibolab/__init__.py
+++ b/src/qibolab/__init__.py
@@ -29,13 +29,14 @@ def get_platforms_path():
     return Path(profiles)
 
 
-def create_platform(name, runcard=None):
+def create_platform(name, path: Path = None) -> Platform:
     """A platform for executing quantum algorithms.
 
     It consists of a quantum processor QPU and a set of controlling instruments.
 
     Args:
         name (str): name of the platform. Options are 'tiiq', 'qili' and 'icarusq'.
+        path (pathlib.Path): path with platform serialization
     Returns:
         The plaform class.
     """
@@ -44,17 +45,17 @@ def create_platform(name, runcard=None):
 
         return create_dummy(with_couplers=name == "dummy_couplers")
 
-    platform = get_platforms_path() / f"{name}.py"
+    platform = get_platforms_path() / f"{name}"
     if not platform.exists():
         raise_error(ValueError, f"Platform {name} does not exist.")
 
-    spec = importlib.util.spec_from_file_location("platform", platform)
+    spec = importlib.util.spec_from_file_location("platform", platform / f"{name}.py")
     module = importlib.util.module_from_spec(spec)
     spec.loader.exec_module(module)
 
-    if runcard is None:
+    if path is None:
         return module.create()
-    return module.create(runcard)
+    return module.create(path)
 
 
 def execute_qasm(circuit: str, platform, runcard=None, initial_state=None, nshots=1000):
diff --git a/src/qibolab/dummy.py b/src/qibolab/dummy.py
index 6321aee56f..7684264b0c 100644
--- a/src/qibolab/dummy.py
+++ b/src/qibolab/dummy.py
@@ -3,8 +3,12 @@
 
 from qibolab.channels import Channel, ChannelMap
 from qibolab.instruments.dummy import DummyInstrument, DummyLocalOscillator
+from qibolab.kernels import Kernels
 from qibolab.platform import Platform
-from qibolab.serialize import load_qubits, load_runcard, load_settings
+from qibolab.serialize import KERNELS_FILE, load_qubits, load_runcard, load_settings
+
+DUMMY_FOLDER = pathlib.Path(__file__).parent / "dummy"
+DUMMY_RUNCARD = "dummy.yml"
 
 
 def remove_couplers(runcard):
@@ -21,11 +25,6 @@ def remove_couplers(runcard):
     return runcard
 
 
-def load_dummy_runcard():
-    """Loads the runcard YAML of the dummy platform."""
-    return load_runcard(pathlib.Path(__file__).parent / "dummy.yml")
-
-
 def create_dummy(with_couplers: bool = True):
     """Create a dummy platform using the dummy instrument.
 
@@ -40,8 +39,9 @@ def create_dummy(with_couplers: bool = True):
     twpa_pump.frequency = 1e9
     twpa_pump.power = 10
 
-    runcard = load_dummy_runcard()
-    extras_folder = pathlib.Path(__file__).parent / "dummy"
+    runcard = load_runcard(DUMMY_FOLDER / "dummy.yml")
+    kernels = Kernels.load(DUMMY_FOLDER / KERNELS_FILE)
+
     if not with_couplers:
         runcard = remove_couplers(runcard)
 
@@ -65,7 +65,7 @@ def create_dummy(with_couplers: bool = True):
     channels["readout"].attenuation = 0
     channels["twpa"].local_oscillator = twpa_pump
 
-    qubits, couplers, pairs = load_qubits(runcard, extras_folder)
+    qubits, couplers, pairs = load_qubits(runcard, kernels)
     settings = load_settings(runcard)
 
     # map channels to qubits
diff --git a/src/qibolab/dummy.yml b/src/qibolab/dummy/dummy.yml
similarity index 100%
rename from src/qibolab/dummy.yml
rename to src/qibolab/dummy/dummy.yml
diff --git a/src/qibolab/kernels.py b/src/qibolab/kernels.py
index 654f97575c..5fa8f61ceb 100644
--- a/src/qibolab/kernels.py
+++ b/src/qibolab/kernels.py
@@ -5,6 +5,8 @@
 
 from qibolab.qubits import QubitId
 
+KERNELS_FILE = "kernels.npz"
+
 
 class Kernels(dict[QubitId, np.ndarray]):
     """A dictionary subclass for handling Qubit Kernels.
@@ -30,5 +32,6 @@ def dump(self, path: Path):
         (numpy arrays) are kept as is.
         """
         np.savez(
-            path, **{json.dumps(qubit_id): value for qubit_id, value in self.items()}
+            path / KERNELS_FILE,
+            **{json.dumps(qubit_id): value for qubit_id, value in self.items()}
         )
diff --git a/src/qibolab/qubits.py b/src/qibolab/qubits.py
index 4359667c5b..139c442c20 100644
--- a/src/qibolab/qubits.py
+++ b/src/qibolab/qubits.py
@@ -15,7 +15,7 @@
 
 Not all channels are required to operate a qubit.
 """
-EXCLUDED_FIELDS = CHANNEL_NAMES + ("name", "native_gates", "kernels", "_flux")
+EXCLUDED_FIELDS = CHANNEL_NAMES + ("name", "native_gates", "kernel", "_flux")
 """Qubit dataclass fields that are excluded by the ``characterization``
 property."""
 
diff --git a/src/qibolab/serialize.py b/src/qibolab/serialize.py
index 0705187876..31f38b98b6 100644
--- a/src/qibolab/serialize.py
+++ b/src/qibolab/serialize.py
@@ -37,7 +37,7 @@ def load_settings(runcard: dict) -> Settings:
 
 
 def load_qubits(
-    runcard: dict, extras_folder: Path = None
+    runcard: dict, kernels: Kernels = None
 ) -> Tuple[QubitMap, CouplerMap, QubitPairMap]:
     """Load qubits and pairs from the runcard.
 
@@ -52,8 +52,7 @@ def load_qubits(
         for q, char in runcard["characterization"]["single_qubit"].items()
     }
 
-    if extras_folder is not None:
-        kernels = Kernels.load(path=extras_folder / KERNELS_FILE)
+    if kernels is not None:
         for q in kernels:
             qubits[q].kernel = kernels[q]
 
@@ -151,11 +150,6 @@ def dump_characterization(qubits: QubitMap, couplers: CouplerMap = None) -> dict
         "single_qubit": {q: qubit.characterization for q, qubit in qubits.items()},
     }
 
-    kernels = Kernels()
-    for qubit in qubits.values():
-        if qubit.kernel is not None:
-            kernels[qubit.name] = qubit.kernel
-
     if couplers:
         characterization["coupler"] = {
             c.name: {"sweetspot": c.sweetspot} for c in couplers.values()
@@ -188,16 +182,6 @@ def dump_runcard(platform: Platform, path: Path):
         path (pathlib.Path): Path that the yaml file will be saved.
     """
 
-    kernels = Kernels()
-    for qubit in platform.qubits.values():
-        if qubit.kernel is not None:
-            kernels[qubit.name] = qubit.kernel
-            qubit.kernel = None
-    name = platform.name
-    if platform.name == "dummy_couplers":
-        name = "dummy"
-    Kernels(kernels).dump(Path(__file__).parent / name / KERNELS_FILE)
-
     settings = {
         "nqubits": platform.nqubits,
         "settings": asdict(platform.settings),
@@ -223,3 +207,32 @@ def dump_runcard(platform: Platform, path: Path):
     path.write_text(
         yaml.dump(settings, sort_keys=False, indent=4, default_flow_style=None)
     )
+
+
+def dump_kernels(platform: Platform, path: Path):
+    """Creates Kernels instance from platform and dumps as npz.
+
+    Args:
+        platform (qibolab.platform.Platform): The platform to be serialized.
+        path (pathlib.Path): Path that the kernels file will be saved.
+    """
+
+    # create and dump kernels
+    kernels = Kernels()
+    for qubit in platform.qubits.values():
+        if qubit.kernel is not None:
+            kernels[qubit.name] = qubit.kernel
+
+    kernels.dump(path / KERNELS_FILE)
+
+
+def dump_platform(platform: Platform, path: Path):
+    """Platform serialization as runcard (yaml) and kernels (npz).
+
+    Args:
+        platform (qibolab.platform.Platform): The platform to be serialized.
+        path (pathlib.Path): Path where yaml and npz will be dumped.
+    """
+
+    dump_kernels(platform=platform, path=path)
+    dump_runcard(platform=platform, path=path)
diff --git a/tests/dummy_qrc/qblox.py b/tests/dummy_qrc/qblox/qblox.py
similarity index 96%
rename from tests/dummy_qrc/qblox.py
rename to tests/dummy_qrc/qblox/qblox.py
index 2355180730..fe47b0fc00 100644
--- a/tests/dummy_qrc/qblox.py
+++ b/tests/dummy_qrc/qblox/qblox.py
@@ -17,17 +17,18 @@
 NAME = "qblox"
 ADDRESS = "192.168.0.6"
 TIME_OF_FLIGHT = 500
-RUNCARD = pathlib.Path(__file__).parent / "qblox.yml"
+FOLDER = pathlib.Path(__file__).parent
+RUNCARD = "qblox.yml"
 
 
-def create(runcard_path=RUNCARD):
+def create(folder: pathlib.Path = FOLDER):
     """QuantWare 5q-chip controlled using qblox cluster.
 
     Args:
         runcard_path (str): Path to the runcard file.
     """
 
-    runcard = load_runcard(runcard_path)
+    runcard = load_runcard(folder / RUNCARD)
     modules = {}
 
     # DEBUG: debug folder = report folder
diff --git a/tests/dummy_qrc/qblox.yml b/tests/dummy_qrc/qblox/qblox.yml
similarity index 100%
rename from tests/dummy_qrc/qblox.yml
rename to tests/dummy_qrc/qblox/qblox.yml
diff --git a/tests/dummy_qrc/qm.py b/tests/dummy_qrc/qm/qm.py
similarity index 96%
rename from tests/dummy_qrc/qm.py
rename to tests/dummy_qrc/qm/qm.py
index 745a77f6e5..2998230b3f 100644
--- a/tests/dummy_qrc/qm.py
+++ b/tests/dummy_qrc/qm/qm.py
@@ -11,10 +11,11 @@
     load_settings,
 )
 
-RUNCARD = pathlib.Path(__file__).parent / "qm.yml"
+FOLDER = pathlib.Path(__file__).parent
+RUNCARD = "qm.yml"
 
 
-def create(runcard_path=RUNCARD):
+def create(folder: pathlib.Path = FOLDER):
     """Dummy platform using Quantum Machines (QM) OPXs and Rohde Schwarz local
     oscillators.
 
@@ -69,7 +70,7 @@ def create(runcard_path=RUNCARD):
     channels["L4-26"].local_oscillator = local_oscillators[5]
 
     # create qubit objects
-    runcard = load_runcard(runcard_path)
+    runcard = load_runcard(FOLDER / RUNCARD)
     qubits, couplers, pairs = load_qubits(runcard)
 
     # assign channels to qubits
diff --git a/tests/dummy_qrc/qm.yml b/tests/dummy_qrc/qm/qm.yml
similarity index 100%
rename from tests/dummy_qrc/qm.yml
rename to tests/dummy_qrc/qm/qm.yml
diff --git a/tests/dummy_qrc/rfsoc.py b/tests/dummy_qrc/rfsoc/rfsoc.py
similarity index 91%
rename from tests/dummy_qrc/rfsoc.py
rename to tests/dummy_qrc/rfsoc/rfsoc.py
index 5bab3f5f1b..f214923546 100644
--- a/tests/dummy_qrc/rfsoc.py
+++ b/tests/dummy_qrc/rfsoc/rfsoc.py
@@ -12,10 +12,11 @@
     load_settings,
 )
 
-RUNCARD = pathlib.Path(__file__).parent / "rfsoc.yml"
+FOLDER = pathlib.Path(__file__).parent
+RUNCARD = "rfsoc.yml"
 
 
-def create(runcard_path=RUNCARD):
+def create(folder: pathlib.Path = FOLDER):
     """Dummy platform using QICK project on the RFSoC4x2 board.
 
     Used in ``test_instruments_rfsoc.py``.
@@ -34,7 +35,7 @@ def create(runcard_path=RUNCARD):
     lo_era = ERA("ErasynthLO", "192.168.0.212", ethernet=True)
     channels["L3-18_ro"].local_oscillator = lo_era
 
-    runcard = load_runcard(runcard_path)
+    runcard = load_runcard(FOLDER / RUNCARD)
     qubits, couplers, pairs = load_qubits(runcard)
 
     # assign channels to qubits
diff --git a/tests/dummy_qrc/rfsoc.yml b/tests/dummy_qrc/rfsoc/rfsoc.yml
similarity index 100%
rename from tests/dummy_qrc/rfsoc.yml
rename to tests/dummy_qrc/rfsoc/rfsoc.yml
diff --git a/tests/dummy_qrc/zurich.py b/tests/dummy_qrc/zurich/zurich.py
similarity index 93%
rename from tests/dummy_qrc/zurich.py
rename to tests/dummy_qrc/zurich/zurich.py
index 1ff4fe6549..c61aad30ab 100644
--- a/tests/dummy_qrc/zurich.py
+++ b/tests/dummy_qrc/zurich/zurich.py
@@ -9,6 +9,7 @@
 from qibolab.channels import Channel, ChannelMap
 from qibolab.instruments.dummy import DummyLocalOscillator as LocalOscillator
 from qibolab.instruments.zhinst import Zurich
+from qibolab.kernels import KERNELS_FILE, Kernels
 from qibolab.serialize import (
     load_instrument_settings,
     load_qubits,
@@ -16,16 +17,16 @@
     load_settings,
 )
 
-RUNCARD = pathlib.Path(__file__).parent / "zurich.yml"
-FOLDER = pathlib.Path(__file__).parent / "zurich"
+RUNCARD = "zurich.yml"
+FOLDER = pathlib.Path(__file__).parent
 N_QUBITS = 5
 
 
-def create(runcard_path=RUNCARD):
-    """IQM 5q-chip controlled Zurich Instrumetns (Zh) SHFQC, HDAWGs and PQSC.
+def create(path: pathlib.Path = FOLDER):
+    """IQM 5q-chip controlled Zurich Instruments (Zh) SHFQC, HDAWGs and PQSC.
 
     Args:
-        runcard_path (str): Path to the runcard file.
+        path (str): Path to configuration folder.
     """
 
     device_setup = DeviceSetup("EL_ZURO")
@@ -171,8 +172,9 @@ def create(runcard_path=RUNCARD):
         channels[ch].local_oscillator = local_oscillators[lo]
 
     # create qubit objects
-    runcard = load_runcard(runcard_path)
-    qubits, couplers, pairs = load_qubits(runcard, FOLDER)
+    runcard = load_runcard(FOLDER / RUNCARD)
+    kernels = Kernels.load(FOLDER / KERNELS_FILE)
+    qubits, couplers, pairs = load_qubits(runcard, kernels)
     settings = load_settings(runcard)
 
     # assign channels to qubits and sweetspots(operating points)
diff --git a/tests/dummy_qrc/zurich.yml b/tests/dummy_qrc/zurich/zurich.yml
similarity index 100%
rename from tests/dummy_qrc/zurich.yml
rename to tests/dummy_qrc/zurich/zurich.yml
diff --git a/tests/test_dummy.py b/tests/test_dummy.py
index ffc02b4227..8109833b92 100644
--- a/tests/test_dummy.py
+++ b/tests/test_dummy.py
@@ -165,7 +165,6 @@ def test_dummy_single_sweep_coupler(
         sweeper = Sweeper(parameter, parameter_range, couplers=[platform.couplers[0]])
     else:
         sweeper = Sweeper(parameter, parameter_range, pulses=[coupler_pulse])
-    print(sweeper)
     options = ExecutionParameters(
         nshots=nshots,
         averaging_mode=average,
diff --git a/tests/test_platform.py b/tests/test_platform.py
index 09c0193a1c..706a0c6570 100644
--- a/tests/test_platform.py
+++ b/tests/test_platform.py
@@ -11,7 +11,7 @@
 
 from qibolab import create_platform
 from qibolab.backends import QibolabBackend
-from qibolab.dummy import load_dummy_runcard
+from qibolab.dummy import DUMMY_FOLDER, DUMMY_RUNCARD
 from qibolab.execution_parameters import ExecutionParameters
 from qibolab.instruments.qblox.controller import QbloxController
 from qibolab.instruments.rfsoc.driver import RFSoC
@@ -60,15 +60,17 @@ def test_platform_pickle(platform):
     assert new_platform.is_connected == platform.is_connected
 
 
-def test_dump_runcard(platform):
-    path = pathlib.Path(__file__).parent / "test.yml"
+def test_dump_runcard(platform, tmp_path):
+    path = tmp_path / "test.yml"
     dump_runcard(platform, path)
     final_runcard = load_runcard(path)
     if platform.name == "dummy" or platform.name == "dummy_couplers":
-        target_runcard = load_dummy_runcard()
+        target_runcard = load_runcard(DUMMY_FOLDER / DUMMY_RUNCARD)
     else:
         target_path = (
-            pathlib.Path(__file__).parent / "dummy_qrc" / f"{platform.name}.yml"
+            pathlib.Path(__file__).parent
+            / "dummy_qrc"
+            / f"{platform.name}/{platform.name}.yml"
         )
         target_runcard = load_runcard(target_path)
     # for the characterization section the dumped runcard may contain