from __future__ import annotations
from dataclasses import dataclass
from typing import Any
from qcodes.instrument import Instrument, InstrumentChannel, InstrumentModule
from qcodes.instrument.channel import ChannelList
from qcodes.parameters import ManualParameter
from qm import QuantumMachinesManager
from quam.components.ports.analog_outputs import LFFEMAnalogOutputPort
from quam.components.ports.analog_inputs import LFFEMAnalogInputPort
from semi_cr.core.lab.station.graphing.models import Node, Edge
from semi_cr.core.lab.station.graphing.operations import build_nx_graph
from semi_cr.core.lab.station.routing_v1.models import ConnectionRoutable
from semi_cr.core.lab.station.graphing.ids import (
instrument_node_id,
instrument_module_node_id,
terminal_node_id,
)
from semi_cr.core.lab.station.graphing.runtime_builder import build_connection_graph_for_instrument_module
from semi_cr.core.lab.station.routing_v1.rf.models import (
OPXRFSource,
OPXIQReceiver,
)
from semi_cr.core.lab.devices.pinned import infer_type
import networkx as nx
MODULE_TEMPLATES = {
"LF": {
"n_outputs": 8,
"n_inputs": 2,
"n_digital_outputs": 8,
},
"MW": {
"n_outputs": 4,
"n_inputs": 2,
"n_digital_outputs": 8,
},
}
def _normalize_channels(
channels: list[int] | dict[int, dict[str, Any]] | None,
max_channel: int,
) -> dict[int, dict[str, Any]]:
"""
Accepts:
None
[1, 2, 3]
{1: {...}, 2: {...}}
Returns:
{1: {...}, 2: {...}}
"""
if channels is None:
return {}
if isinstance(channels, list):
result = {ch: {} for ch in channels}
else:
result = dict(channels)
for ch in result:
if ch < 1 or ch > max_channel:
raise ValueError(f"Invalid channel {ch}. Valid range is 1 to {max_channel}.")
return result
[docs]
class OPXChannel(InstrumentChannel):
def __init__(
self,
parent: InstrumentModule,
name: str,
channel_number: int,
direction: str,
**kwargs: Any,
) -> None:
super().__init__(parent, name, **kwargs)
self.channel_number = channel_number
self.direction = direction
self.add_parameter(
"enabled",
parameter_class=ManualParameter,
initial_value=True,
vals=None,
)
@property
def number(self) -> int:
return self.channel_number
@property
def terminal_name(self) -> str:
return {
"output": f"out{self.channel_number}",
"input": f"in{self.channel_number}",
"digital_output": f"dig_out{self.channel_number}",
}[self.direction]
def __repr__(self) -> str:
return (
f"<{self.__class__.__name__}: {self.short_name} "
f"({self.direction} {self.channel_number})>"
)
[docs]
class OPXOutputChannel(OPXChannel):
def __init__(
self,
parent: InstrumentModule,
name: str,
channel_number: int,
intermediate_frequency: float = 10e6, # Hz
sampling_rate: float = 2e9, # Gbits/s
**kwargs: Any,
) -> None:
super().__init__(
parent,
name,
channel_number=channel_number,
direction="output",
**kwargs,
)
self.intermediate_frequency = intermediate_frequency
self.sampling_rate = sampling_rate
[docs]
def to_config(self) -> dict[str, Any]:
cfg = {}
if self.intermediate_frequency is not None:
cfg["intermediate_frequency"] = self.intermediate_frequency
if self.sampling_rate is not None:
cfg["sampling_rate"] = self.sampling_rate
return cfg
[docs]
class OPXDigitalOutputChannel(OPXChannel):
def __init__(
self,
parent: InstrumentModule,
name: str,
channel_number: int,
trigger_id: str | None = None,
port: int | None = None,
intermediate_frequency: float | None = None,
delay: int = 0,
buffer: int = 0,
**kwargs: Any,
) -> None:
super().__init__(
parent,
name,
channel_number=channel_number,
direction="digital_output",
**kwargs,
)
self.trigger_id = trigger_id
self.port = port
self.intermediate_frequency = intermediate_frequency
self.delay = delay
self.buffer = buffer
@property
def port_id(self) -> int:
if self.port is not None:
return self.port
return self.channel_number
[docs]
def to_config(self) -> dict[str, Any]:
cfg = {
"delay": self.delay,
"buffer": self.buffer,
}
if self.trigger_id is not None:
cfg["id"] = self.trigger_id
if self.port is not None:
cfg["port"] = self.port
if self.intermediate_frequency is not None:
cfg["intermediate_frequency"] = self.intermediate_frequency
return cfg
def _channels_from_connections(
connections: tuple[dict[str, Any]] | None,
prefix: str,
) -> list[int]:
channel_numbers: list[int] = []
for connection in connections or []:
name = str(connection["name"])
if not name.startswith(prefix):
continue
suffix = name.removeprefix(prefix)
if suffix.isdigit():
channel_numbers.append(int(suffix))
return channel_numbers
[docs]
class FEMModule(InstrumentModule, ConnectionRoutable):
def __init__(
self,
# parent: InstrumentModule,
parent: Instrument,
name: str,
controller: str,
slot: int,
module_type: str,
connections: list[dict[str, Any]] | None = None,
inputs: list[int] | dict[int, dict[str, Any]] | None = None,
outputs: list[int] | dict[int, dict[str, Any]] | None = None,
digital_outputs: list[int] | dict[int, dict[str, Any]] | None = None,
**kwargs: Any,
) -> None:
super().__init__(parent, name, **kwargs)
self._init_connections(connections)
if module_type not in MODULE_TEMPLATES:
raise ValueError(
f"Unknown module_type {module_type!r}. "
f"Valid types are {list(MODULE_TEMPLATES)}."
)
self.controller = controller
self.slot = slot
self.module_type = module_type
# self._connections = connections or []
template = MODULE_TEMPLATES[module_type]
inputs = inputs if inputs is not None else _channels_from_connections(
self._connections,
prefix="in",
)
outputs = outputs if outputs is not None else _channels_from_connections(
self._connections,
prefix="out",
)
digital_outputs = (
digital_outputs
if digital_outputs is not None
else _channels_from_connections(
self._connections,
prefix="dig_out",
)
)
input_configs = _normalize_channels(
inputs,
max_channel=template["n_inputs"],
)
output_configs = _normalize_channels(
outputs,
max_channel=template["n_outputs"],
)
digital_output_configs = _normalize_channels(
digital_outputs,
template["n_digital_outputs"],
)
self.add_submodule(
"outputs",
ChannelList(self, "outputs", OPXOutputChannel, snapshotable=True),
)
for channel_number, channel_config in output_configs.items():
channel_name = f"out_chan{channel_number}"
ch = OPXOutputChannel(
self,
channel_name,
channel_number=channel_number,
**channel_config,
)
self.add_submodule(channel_name, ch)
self.outputs.append(ch)
self.outputs.lock()
self.add_submodule(
"inputs",
ChannelList(self, "inputs", OPXInputChannel, snapshotable=True),
)
for channel_number, channel_config in input_configs.items():
channel_name = f"in_chan{channel_number}"
ch = OPXInputChannel(
self,
channel_name,
channel_number=channel_number,
**channel_config,
)
self.add_submodule(channel_name, ch)
self.inputs.append(ch)
self.inputs.lock()
self.add_submodule(
"digital_outputs",
ChannelList(
self,
"digital_outputs",
OPXDigitalOutputChannel,
snapshotable=True,
),
)
for channel_number, channel_config in digital_output_configs.items():
channel_name = f"dig_out_chan{channel_number}"
ch = OPXDigitalOutputChannel(
self,
channel_name,
channel_number=channel_number,
**channel_config,
)
self.add_submodule(channel_name, ch)
self.digital_outputs.append(ch)
self.digital_outputs.lock()
[docs]
def to_config(self) -> dict[str, Any]:
return {
"controller": self.controller,
"slot": self.slot,
"module_type": self.module_type,
"outputs": {
ch.channel_number: ch.to_config()
for ch in self.outputs
if ch.enabled()
},
"inputs": {
ch.channel_number: ch.to_config()
for ch in self.inputs
if ch.enabled()
},
"digital_outputs": {
ch.channel_number: ch.to_config()
for ch in self.digital_outputs
if ch.enabled()
},
}
@property
def owner_name(self) -> str:
parent_name = self.parent.name
module_name = self.short_name
return f"{parent_name}/{module_name}"
[docs]
def terminal_id(
self,
terminal_name: str,
) -> str:
return terminal_node_id(
self.owner_name,
terminal_name,
)
[docs]
def get_terminal(
self,
terminal_name: str,
) -> OPXChannel:
for collection in (
self.outputs,
self.inputs,
self.digital_outputs,
):
for channel in collection:
if channel.terminal_name == terminal_name:
return channel
raise KeyError(
f"Unknown terminal {terminal_name!r} "
f"on module {self.owner_name!r}."
)
def __repr__(self) -> str:
return (
f"<{self.__class__.__name__}: {self.short_name} "
f"type={self.module_type} slot={self.slot}>"
)
@property
def graph(self) -> nx.MultiDiGraph:
parent_name = self.parent.name
owner_name = self.owner_name
G = build_connection_graph_for_instrument_module(
module=self,
connections=self._connections,
parent_name=parent_name,
owner_name=owner_name,
)
# Attach RF source backend objects to output terminal nodes
for channel in self.outputs:
terminal_id = terminal_node_id(owner_name, channel.terminal_name)
lf_port = LFFEMAnalogOutputPort(
controller_id=self.controller,
fem_id=self.slot,
port_id=channel.channel_number,
sampling_rate=channel.sampling_rate,
)
rf_source = OPXRFSource(
opx=self.parent,
module=self,
terminal_id=terminal_id,
port=lf_port,
intermediate_frequency=channel.intermediate_frequency,
)
G.nodes[terminal_id]["obj"] = rf_source
G.nodes[terminal_id]["backend_kind"] = "rf_source"
# Attach IQ receiver backend objects to input terminal nodes
for channel in self.inputs:
terminal_id = terminal_node_id(owner_name, channel.terminal_name)
lf_port = LFFEMAnalogInputPort(
controller_id=self.controller,
fem_id=self.slot,
port_id=channel.channel_number,
sampling_rate=channel.sampling_rate,
)
iq_receiver = OPXIQReceiver(
opx=self.parent,
module=self,
terminal_id=terminal_id,
port=lf_port,
time_of_flight=channel.time_of_flight,
)
G.nodes[terminal_id]["obj"] = iq_receiver
G.nodes[terminal_id]["backend_kind"] = "iq_receiver"
# Digital trigger outputs
for channel in self.digital_outputs:
terminal_id = self.terminal_id(
channel.terminal_name
)
G.nodes[terminal_id]["obj"] = channel
G.nodes[terminal_id]["backend_kind"] = "digital_trigger"
return G
[docs]
class OPXModules:
def __init__(self):
self._modules: dict[str, FEMModule] = {}
[docs]
def add(self, name: str, module: FEMModule) -> None:
self._modules[name] = module
def __getitem__(self, name: str) -> FEMModule:
return self._modules[name]
def __iter__(self):
return iter(self._modules.values())
[docs]
def items(self):
return self._modules.items()
[docs]
def values(self):
return self._modules.values()
[docs]
def keys(self):
return self._modules.keys()
[docs]
def to_config(self) -> dict[str, dict[str, Any]]:
return {
name: module.to_config()
for name, module in self.items()
}
[docs]
@dataclass
class TriggerConfig:
id: str | None = None
port: int | None = None
intermediate_frequency: float | None = None
delay: float = 0
buffer: int = 0
[docs]
def to_config(self) -> dict[str, Any]:
cfg = {
"delay": self.delay,
"buffer": self.buffer,
}
if self.id is not None:
cfg["id"] = self.id
if self.port is not None:
cfg["port"] = self.port
if self.intermediate_frequency is not None:
cfg["intermediate_frequency"] = self.intermediate_frequency
return cfg
[docs]
class OPX1000(Instrument, ConnectionRoutable):
def __init__(
self,
name: str,
ip_address: str,
port: int,
modules: dict[str, dict[str, Any]] | None = None,
**kwargs: Any,
) -> None:
super().__init__(name, metadata={}, label="opx1000", **kwargs)
self.ip_address = ip_address
self.port = port
self.qmm = QuantumMachinesManager(
host=ip_address,
port=port,
)
self.modules = OPXModules()
reserved = {"type"}
for module_name, module_cfg in (modules or {}).items():
module_class = infer_type(module_cfg["type"])
module_kwargs = {
key: value
for key, value in module_cfg.items()
if key not in reserved
}
module = module_class(
parent=self,
name=module_name,
**module_kwargs,
)
self.modules.add(module_name, module)
self.add_submodule(module_name, module)
[docs]
def get_idn(self) -> dict[str, str | None]:
return {
"vendor": "Quantum Machines",
"model": self.name,
"serial": None,
"firmware": None,
}
@property
def graph(self) -> nx.MultiDiGraph:
nodes: list[Node] = []
edges: list[Edge] = []
opx_id = instrument_node_id(self.name)
nodes.append(Node(
id=opx_id,
kind="instrument",
name=self.name,
obj=self,
))
G = build_nx_graph(nodes, edges)
for module in self.modules:
module_id = instrument_module_node_id(self.name, module.name)
G.add_node(
module_id,
kind="instrument_module",
name=module.name,
obj=module,
)
G.add_edge(opx_id, module_id, kind="contains")
G = nx.compose(G, module.graph)
return G