dump platform with updates

src/qibolab/execution_parameters.py

@@ -1,7 +1,6 @@
 from enum import Enum, auto
-from typing import Optional
+from typing import Any, Optional
 
-from qibolab.components import Config
 from qibolab.result import (
     AveragedIntegratedResults,
     AveragedRawWaveformResults,
@@ -52,6 +51,14 @@ class AveragingMode(Enum):
 }
 
 
+ConfigUpdate = dict[str, dict[str, Any]]
+"""Update for component configs.
+
+Maps component name to corresponding update, which in turn is a map from
+config property name that needs an update to its new value.
+"""
+
+
 class ExecutionParameters(Model):
     """Data structure to deal with execution parameters."""
 
@@ -72,12 +79,12 @@ class ExecutionParameters(Model):
     """Data acquisition type."""
     averaging_mode: AveragingMode = AveragingMode.SINGLESHOT
     """Data averaging mode."""
-    configs: list[dict[str, Config]] = []
-    """Configuration for various components (maps component name to respective
-    config).
+    updates: list[ConfigUpdate] = []
+    """List of updates for component configs.
 
-    This takes precedence over platform defaults, and can be only a
-    subset of components (i.e. only the ones that need to be updated).
+    Later entries in the list take precedence over earlier ones (if they
+    happen to update the same thing). These updates will be applied on
+    top of platform defaults.
     """
 
     @property

src/qibolab/platform/platform.py

@@ -1,5 +1,6 @@
 """A platform for executing quantum algorithms."""
 
+import dataclasses
 from collections import defaultdict
 from dataclasses import asdict, dataclass, field
 from typing import Any, Dict, List, Optional, Tuple
@@ -10,7 +11,7 @@
 
 from qibolab.components import Config
 from qibolab.couplers import Coupler
-from qibolab.execution_parameters import ExecutionParameters
+from qibolab.execution_parameters import ConfigUpdate, ExecutionParameters
 from qibolab.instruments.abstract import Controller, Instrument, InstrumentId
 from qibolab.pulses import Delay, PulseSequence
 from qibolab.qubits import Qubit, QubitId, QubitPair, QubitPairId
@@ -59,6 +60,23 @@ def unroll_sequences(
     return total_sequence, readout_map
 
 
+def update_configs(configs: dict[str, Config], updates: list[ConfigUpdate]):
+    """Apply updates to configs in place.
+
+    Args:
+        configs: configs to update. Maps component name to respective config.
+        updates: list of config updates. Later entries in the list take precedence over earlier entries
+                 (if they happen to update the same thing).
+    """
+    for update in updates:
+        for name, changes in update.items():
+            if name not in configs:
+                raise ValueError(
+                    f"Cannot update configuration for unknown component {name}"
+                )
+            configs[name] = dataclasses.replace(configs[name], **changes)
+
+
 @dataclass
 class Settings:
     """Default execution settings read from the runcard."""
@@ -177,30 +195,6 @@ def disconnect(self):
                 instrument.disconnect()
         self.is_connected = False
 
-    def _apply_config_updates(
-        self, updates: list[dict[str, Config]]
-    ) -> dict[str, Config]:
-        """Apply given list of config updates to the default configuration and
-        return the updated config dict.
-
-        Args:
-            updates: list of updates, where each entry is a dict mapping component name to new config. Later entries
-                     in the list override earlier entries (if they happen to update the same thing).
-        """
-        components = self.configs.copy()
-        for update in updates:
-            for name, cfg in update.items():
-                if name not in components:
-                    raise ValueError(
-                        f"Cannot update configuration for unknown component {name}"
-                    )
-                if type(cfg) is not type(components[name]):
-                    raise ValueError(
-                        f"Configuration of component {name} with type {type(components[name])} cannot be updated with configuration of type {type(cfg)}"
-                    )
-                components[name] = cfg
-        return components
-
     @property
     def _controller(self):
         """Identify controller instrument.
@@ -225,7 +219,7 @@ def _execute(self, sequences, options, integration_setup, *sweepers):
         for instrument in self.instruments.values():
             if isinstance(instrument, Controller):
                 new_result = instrument.play(
-                    options.configs,
+                    options.updates,
                     sequences,
                     options,
                     integration_setup,
@@ -283,7 +277,8 @@ def execute(
             time *= len(sweep.values)
         log.info(f"Minimal execution time: {time}")
 
-        configs = self._apply_config_updates(options.configs)
+        configs = self.configs.copy()
+        update_configs(configs, options.updates)
 
         # for components that represent aux external instruments (e.g. lo) to the main control instrument
         # set the config directly

src/qibolab/serialize.py

@@ -8,11 +8,12 @@
 import json
 from dataclasses import asdict, fields
 from pathlib import Path
-from typing import Tuple, Union
+from typing import Optional, Tuple, Union
 
 from pydantic import ConfigDict, TypeAdapter
 
 from qibolab.couplers import Coupler
+from qibolab.execution_parameters import ConfigUpdate
 from qibolab.kernels import Kernels
 from qibolab.native import (
     FixedSequenceFactory,
@@ -27,6 +28,7 @@
     QubitMap,
     QubitPairMap,
     Settings,
+    update_configs,
 )
 from qibolab.pulses import Pulse, PulseSequence
 from qibolab.pulses.pulse import PulseLike
@@ -307,23 +309,30 @@ def dump_component_configs(component_configs) -> dict:
     return {name: asdict(cfg) for name, cfg in component_configs.items()}
 
 
-def dump_runcard(platform: Platform, path: Path):
+def dump_runcard(
+    platform: Platform, path: Path, updates: Optional[list[ConfigUpdate]] = None
+):
     """Serializes the platform and saves it as a json runcard file.
 
     The file saved follows the format explained in :ref:`Using runcards <using_runcards>`.
 
     Args:
         platform (qibolab.platform.Platform): The platform to be serialized.
         path (pathlib.Path): Path that the json file will be saved.
+        updates: List if updates for platform configs.
+                 Later entries in the list take precedence over earlier ones (if they happen to update the same thing).
     """
 
+    configs = platform.configs.copy()
+    update_configs(configs, updates or [])
+
     settings = {
         "nqubits": platform.nqubits,
         "settings": asdict(platform.settings),
         "qubits": list(platform.qubits),
         "topology": [list(pair) for pair in platform.ordered_pairs],
         "instruments": dump_instruments(platform.instruments),
-        "components": dump_component_configs(platform.configs),
+        "components": dump_component_configs(configs),
     }
 
     if platform.couplers:
@@ -363,13 +372,17 @@ def dump_kernels(platform: Platform, path: Path):
         kernels.dump(path)
 
 
-def dump_platform(platform: Platform, path: Path):
+def dump_platform(
+    platform: Platform, path: Path, updates: Optional[list[ConfigUpdate]] = None
+):
     """Platform serialization as runcard (json) and kernels (npz).
 
     Args:
         platform (qibolab.platform.Platform): The platform to be serialized.
         path (pathlib.Path): Path where json and npz will be dumped.
+        updates: List if updates for platform configs.
+                 Later entries in the list take precedence over earlier ones (if they happen to update the same thing).
     """
 
     dump_kernels(platform=platform, path=path)
-    dump_runcard(platform=platform, path=path)
+    dump_runcard(platform=platform, path=path, updates=updates)

tests/test_platform.py

@@ -14,13 +14,15 @@
 
 from qibolab import create_platform
 from qibolab.backends import QibolabBackend
+from qibolab.components import IqConfig, OscillatorConfig
 from qibolab.dummy import create_dummy
 from qibolab.dummy.platform import FOLDER
 from qibolab.execution_parameters import ExecutionParameters
 from qibolab.instruments.qblox.controller import QbloxController
 from qibolab.kernels import Kernels
 from qibolab.platform import Platform, unroll_sequences
 from qibolab.platform.load import PLATFORMS
+from qibolab.platform.platform import update_configs
 from qibolab.pulses import Delay, Gaussian, Pulse, PulseSequence, PulseType, Rectangular
 from qibolab.serialize import (
     PLATFORM,
@@ -97,6 +99,38 @@ def test_platform_sampling_rate(platform):
     assert platform.sampling_rate >= 1
 
 
+def test_update_configs(platform):
+    drive_name = "q0/drive"
+    pump_name = "twpa_pump"
+    configs = {
+        drive_name: IqConfig(4.1e9),
+        pump_name: OscillatorConfig(3e9, -5),
+    }
+
+    updated = update_configs(configs, [{drive_name: {"frequency": 4.2e9}}])
+    assert updated is None
+    assert configs[drive_name].frequency == 4.2e9
+
+    update_configs(
+        configs, [{drive_name: {"frequency": 4.3e9}, pump_name: {"power": -10}}]
+    )
+    assert configs[drive_name].frequency == 4.3e9
+    assert configs[pump_name].frequency == 3e9
+    assert configs[pump_name].power == -10
+
+    update_configs(
+        configs,
+        [{drive_name: {"frequency": 4.4e9}}, {drive_name: {"frequency": 4.5e9}}],
+    )
+    assert configs[drive_name].frequency == 4.5e9
+
+    with pytest.raises(ValueError, match="unknown component"):
+        update_configs(configs, [{"non existent": {"property": 1.0}}])
+
+    with pytest.raises(TypeError, match="prprty"):
+        update_configs(configs, [{pump_name: {"prprty": 0.7}}])
+
+
 def test_dump_runcard(platform, tmp_path):
     dump_runcard(platform, tmp_path)
     final_runcard = load_runcard(tmp_path)
@@ -105,6 +139,7 @@ def test_dump_runcard(platform, tmp_path):
     else:
         target_path = pathlib.Path(__file__).parent / "dummy_qrc" / f"{platform.name}"
         target_runcard = load_runcard(target_path)
+
     # for the characterization section the dumped runcard may contain
     # some default ``Qubit`` parameters
     target_char = target_runcard.pop("characterization")["single_qubit"]
@@ -120,6 +155,20 @@ def test_dump_runcard(platform, tmp_path):
     assert final_instruments == target_instruments
 
 
+def test_dump_runcard_with_updates(platform, tmp_path):
+    qubit = next(iter(platform.qubits.values()))
+    frequency = platform.config(qubit.drive.name).frequency + 1.5e9
+    smearing = platform.config(qubit.acquisition.name).smearing + 10
+    update = {
+        qubit.drive.name: {"frequency": frequency},
+        qubit.acquisition.name: {"smearing": smearing},
+    }
+    dump_runcard(platform, tmp_path, [update])
+    final_runcard = load_runcard(tmp_path)
+    assert final_runcard["components"][qubit.drive.name]["frequency"] == frequency
+    assert final_runcard["components"][qubit.acquisition.name]["smearing"] == smearing
+
+
 @pytest.mark.parametrize("has_kernels", [False, True])
 def test_kernels(tmp_path, has_kernels):
     """Test dumping and loading of `Kernels`."""