import time
from collections.abc import Sequence
from typing import Any
import networkx as nx
from qcodes.instrument import Instrument, InstrumentChannel, InstrumentModule
from spirack import (
D5a_module,
SPI_rack,
)
from semi_cr.core.lab.devices.pinned import infer_type
from semi_cr.core.lab.station.graphing.ids import qcodes_node_id, terminal_node_id, instrument_node_id, instrument_module_node_id
from semi_cr.core.lab.station.graphing.runtime_builder import build_connection_graph_for_instrument_module
from semi_cr.core.lab.station.graphing.models import (
Edge,
Node,
)
from semi_cr.core.lab.station.graphing.operations import build_nx_graph
from semi_cr.core.lab.station.routing_v1.models import ConnectionRoutable, Routable
[docs]
class D5aChannel(InstrumentChannel, Routable):
def __init__(
self,
parent,
name: str,
d5a: D5a_module,
channel: int,
role: str | None = None,
enabled: bool = True,
):
super().__init__(parent, name)
self.d5a = d5a
self._channel_index = channel
self._role = role
self._is_enabled = enabled
self.add_parameter(
"channel_index",
initial_cache_value=self._channel_index,
get_cmd=None,
set_cmd=False,
)
self.add_parameter(
"channel_role",
initial_cache_value=self._role,
get_cmd=None,
set_cmd=False,
)
self.add_parameter(
"channel_enabled",
initial_cache_value=self._is_enabled,
get_cmd=None,
set_cmd=False,
)
self.add_parameter(
"voltage",
unit="V",
label=f"D5a channel {self._channel_index} voltage",
# get_cmd=lambda: self._voltage_cache,
get_cmd=self.get_voltage,
set_cmd=self.set_voltage,
get_parser=float,
set_parser=float,
)
self.add_parameter(
"voltage_readback",
unit="V",
label=f"D5a channel {self._channel_index} voltage readback",
get_cmd=self.get_voltage,
set_cmd=False,
get_parser=float,
)
[docs]
def set_voltage(self, voltage: float) -> None:
self.d5a.set_voltage(self._channel_index, voltage)
# self._voltage_cache = float(voltage)
[docs]
def get_voltage(self) -> float:
return float(self.d5a.get_settings(self._channel_index)[0])
@property
def graph(self) -> nx.MultiDiGraph:
terminal_id = terminal_node_id(
self.parent.short_name,
self.short_name,
)
channel_id = qcodes_node_id(self)
attrs = {
"role": self._role,
"quantity": "voltage",
"modality": "dc",
"channel": self._channel_index,
"number": self._channel_index,
"enabled": self._is_enabled,
"module_type": "D5a",
}
nodes = [
Node(
id=channel_id,
kind="instrument_channel",
name=self.short_name,
obj=self,
attrs=attrs,
),
Node(
id=terminal_id,
kind="terminal",
name=f"{self.parent.short_name}[{self.short_name}]",
obj=self,
attrs=attrs,
),
]
edges = [
Edge(
source=channel_id,
target=terminal_id,
kind="represents",
)
]
return build_nx_graph(nodes, edges)
[docs]
class QcodesD5aModule(InstrumentModule, ConnectionRoutable):
def __init__(
self,
parent,
name: str,
spirack,
slot: int,
num_dacs: int = 16,
reset_voltages: bool = False,
connections: Sequence[dict[str, Any]] | None = None,
channels: dict[str, Any] | None = None,
):
super().__init__(parent, name)
self._slot = slot
self._channel_names: list[str] = []
if connections is None:
connections = [{"name": channel_cfg["name"]} for channel_cfg in (channels or []) if channel_cfg.get("enabled", True)]
self._init_connections(connections)
self.d5a = D5a_module(
spirack,
module=self._slot,
reset_voltages=reset_voltages,
num_dacs=num_dacs,
)
self.add_parameter(
"module_slot",
initial_cache_value=self._slot, # use initial_cache_value for read-only parameters
get_cmd=None,
set_cmd=False,
)
for channel_cfg in channels or []:
channel = D5aChannel(
parent=self,
name=channel_cfg["name"],
d5a=self.d5a,
channel=channel_cfg["channel"],
role=channel_cfg.get("role"),
enabled=channel_cfg.get("enabled", True),
)
self.add_submodule(channel_cfg["name"], channel)
self._channel_names.append(channel_cfg["name"])
@property
def channels(self) -> list[D5aChannel]:
return [getattr(self, name) for name in self._channel_names]
@property
def graph(self) -> nx.MultiDiGraph:
module_id = qcodes_node_id(self)
nodes = [
Node(
id=module_id,
kind="instrument_module",
name=self.short_name,
obj=self,
attrs={
"slot": self._slot,
"module_type": getattr(self, "module_type", None),
},
)
]
graph = build_nx_graph(nodes, [])
for channel in self.channels:
if not channel._is_enabled:
continue
graph = nx.compose(graph, channel.graph)
channel_id = qcodes_node_id(channel)
graph.add_edge(
module_id,
channel_id,
kind="contains",
)
return graph
[docs]
class QcodesM1gModule(InstrumentModule, ConnectionRoutable):
def __init__(
self,
parent,
name: str,
slot: int,
gain_v_per_a: float,
connections: Sequence[dict[str, Any]] | None = None,
spirack=None,
**kwargs,
):
super().__init__(parent, name)
self._init_connections(connections)
self._slot = slot
self.add_parameter(
"module_slot",
initial_cache_value=self._slot,
get_cmd=None,
set_cmd=False,
)
self.add_parameter(
"gain_v_per_a",
unit="V/A",
initial_cache_value=gain_v_per_a,
get_cmd=None,
set_cmd=False,
)
[docs]
class QcodesVIModuleVoltageSource(InstrumentModule, ConnectionRoutable):
"""
Passive voltage-source section of a VI module.
It has no hardware commands. It only describes how a voltage from a
controllable source, e.g. D5a.ch0, is converted to the voltage delivered
at the VI module output.
"""
def __init__(
self,
parent,
name: str,
conversion_factor: float,
connections: Sequence[dict[str, Any]] | None = None,
):
super().__init__(parent, name)
self._init_connections(connections)
self.add_parameter(
"voltage_conversion_factor",
unit="",
initial_cache_value=conversion_factor,
get_cmd=None,
set_cmd=False,
)
[docs]
class QcodesVIModule(InstrumentModule, ConnectionRoutable):
def __init__(
self,
parent,
name: str,
slot: int,
voltage_conversion_factor: float,
connections: Sequence[dict[str, Any]] | None = None,
spirack=None,
**kwargs,
):
super().__init__(parent, name)
self._init_connections(connections)
self._slot = slot
self.add_parameter(
"module_slot",
initial_cache_value=self._slot,
get_cmd=None,
set_cmd=False,
)
voltage_source_connections = [connection for connection in self._connections if connection["name"] in {"V+", "V-"}]
self.add_submodule(
"voltage_source",
QcodesVIModuleVoltageSource(
parent=self,
name="voltage_source",
conversion_factor=voltage_conversion_factor,
connections=voltage_source_connections,
),
)
@property
def graph(self) -> nx.MultiDiGraph:
vi_id = instrument_module_node_id(
self.parent.short_name,
self.short_name,
)
graph = build_connection_graph_for_instrument_module(
self,
self._connections,
parent_id=instrument_node_id(self.parent.short_name),
)
if hasattr(self, "voltage_source"):
voltage_source_id = instrument_module_node_id(
vi_id,
"voltage_source",
)
voltage_source_graph = build_nx_graph(
nodes=[
Node(
id=voltage_source_id,
kind="instrument_module",
name="voltage_source",
obj=self.voltage_source,
)
],
edges=[
Edge(
source=vi_id,
target=voltage_source_id,
kind="contains",
)
],
)
graph = nx.compose(graph, voltage_source_graph)
for connection in self._connections:
if connection["name"] in {"V+", "V-"}:
connection_id = instrument_module_node_id(
vi_id,
str(connection["name"]),
)
graph.add_edge(
voltage_source_id,
connection_id,
kind="uses",
)
return graph
[docs]
class QcodesB1aModule(InstrumentModule, ConnectionRoutable):
"""
Passive Break IN/OUT module.
No SPI commands.
"""
def __init__(
self,
parent,
name: str,
slot: int,
connections: Sequence[dict[str, Any]] | None = None,
spirack=None,
**kwargs,
):
super().__init__(parent, name)
self._init_connections(connections)
self._slot = slot
self.add_parameter(
"module_slot",
initial_cache_value=self._slot,
get_cmd=None,
set_cmd=False,
)
[docs]
class QcodesR1Module(InstrumentModule, ConnectionRoutable):
"""
Passive R1 resistance reference module.
Contains a manual knob selecting one of several resistors.
Supports metadata for 2-probe and 4-probe wiring.
"""
RESISTANCES_OHM = {
"0": 0.0,
"10": 10,
"100": 100,
"1k": 1e3,
"10k": 10e3,
"100k": 100e3,
"1M": 1e6,
"10M": 10e6,
"100M": 100e6,
"1G": 1e9,
}
def __init__(
self,
parent,
name: str,
slot: int,
selected_position: str,
connections: Sequence[dict[str, Any]] | None = None,
spirack=None,
**kwargs,
):
super().__init__(parent, name)
self._init_connections(connections)
self._slot = slot
self._selected_position = selected_position
self.add_parameter(
"module_slot",
initial_cache_value=self._slot,
get_cmd=None,
set_cmd=False,
)
self.add_parameter(
"knob_selected_position",
initial_cache_value=self._selected_position,
get_cmd=None,
set_cmd=False,
)
self.add_parameter(
"nominal_resistance",
unit="Ohm",
initial_cache_value=float(self.RESISTANCES_OHM[selected_position]),
get_cmd=None,
set_cmd=False,
)
[docs]
class SPIRackModules:
def __init__(self):
self._modules: dict[str, InstrumentModule] = {}
[docs]
def add(self, name: str, module: InstrumentModule) -> None:
self._modules[name] = module
def __getitem__(self, name: str) -> InstrumentModule:
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]
class QcodesSPIRack(Instrument):
def __init__(
self,
name: str,
port: str,
baud: int,
modules: dict[str, dict[str, Any]] | None = None,
timeout: float = 1.0,
):
# Calls Instrument.__init__(name)
super().__init__(name)
self.spirack = SPI_rack(
port=port,
baud=baud,
timeout=timeout,
)
reserved = {"type"}
self.modules = SPIRackModules()
try:
self.spirack.unlock()
# FIXME: This shouldn't be necessary but the SPI rack acts up if we don't wait a little bit after unlocking
time.sleep(0.01)
for module_name, module_cfg in (modules or {}).items():
module_class = infer_type(module_cfg["type"])
kwargs = {
"parent": self,
"name": module_name,
"spirack": self.spirack,
**{key: value for key, value in module_cfg.items() if key not in reserved},
}
module = module_class(**kwargs)
self.modules.add(module_name, module)
# Also register with QCoDeS containment
self.add_submodule(module_name, module)
except Exception:
self.close()
raise
[docs]
def close(self):
try:
self.spirack.close()
finally:
super().close()
[docs]
def get_idn(self) -> dict[str, str | None]:
return {
"vendor": "QuTech",
"model": self.name,
"serial": None,
"firmware": None,
}
[docs]
def add_d5a_module(
self,
name: str,
slot: int,
num_dacs: int = 16,
) -> QcodesD5aModule:
d5a = QcodesD5aModule(
parent=self,
name=name,
spirack=self.spirack,
slot=slot,
num_dacs=num_dacs,
)
self.add_submodule(name, d5a)
return d5a
[docs]
def add_m1g_module(
self,
name: str,
slot: int,
gain_v_per_a: float,
) -> "QcodesM1gModule":
mod = QcodesM1gModule(
parent=self,
name=name,
slot=slot,
gain_v_per_a=gain_v_per_a,
)
self.add_submodule(name, mod)
return mod
[docs]
def add_vi_module(
self,
name: str,
slot: int,
voltage_conversion_factor: float,
) -> "QcodesVIModule":
mod = QcodesVIModule(
parent=self,
name=name,
slot=slot,
voltage_conversion_factor=voltage_conversion_factor,
)
self.add_submodule(name, mod)
return mod
[docs]
def add_b1a_module(
self,
name: str,
slot: int,
) -> "QcodesB1aModule":
mod = QcodesB1aModule(
parent=self,
name=name,
slot=slot,
)
self.add_submodule(name, mod)
return mod
[docs]
def add_r1_module(
self,
name: str,
slot: int,
selected_position: str,
) -> "QcodesR1Module":
mod = QcodesR1Module(
parent=self,
name=name,
slot=slot,
selected_position=selected_position,
)
self.add_submodule(name, mod)
return mod
@property
def graph(self) -> nx.MultiDiGraph:
rack_name = getattr(self, "short_name", self.name)
rack_id = instrument_node_id(rack_name)
# rack_id = instrument_node_id(self.name)
nodes = [
Node(
id=rack_id,
kind="instrument",
name=self.name,
obj=self,
)
]
edges: list[Edge] = []
for module in self.modules:
edges.append(
Edge(
source=rack_id,
target=instrument_module_node_id(
rack_name,
module.name,
),
kind="contains",
)
)
graph = build_nx_graph(nodes, edges)
for module in self.modules:
if hasattr(module, "graph"):
graph = nx.compose(graph, module.graph)
return graph