from __future__ import annotations
import logging
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any
from semi_cr.core.lab.station.graphing.ids import terminal_id
from semi_cr.core.station.drivers import DCDriver
from semi_cr.core.station.graph import GraphEdge, GraphNode, graph_attributes
from semi_cr.core.station.instruments.qcodes_proxy import QCoDeSInstrumentProxy, QCoDeSInstrumentSpec
from semi_cr.core.station.units import VoltageLike, as_volts
if TYPE_CHECKING:
from collections.abc import Iterable, Mapping
[docs]
@dataclass
class SPIRackD5aTerminal(DCDriver):
parent_ref: SPIRackD5aModule
name: str
channel: int
role: str | None
@property
def module(self) -> SPIRackD5aModule:
return self.parent_ref
@property
def spirack(self) -> SPIRackBlueprint:
return self.parent_ref.parent_ref
@property
def node_id(self) -> str:
return terminal_id(self.module.name, self.name)
[docs]
def graph_nodes(self) -> Iterable[GraphNode]:
yield GraphNode(
id=self.node_id,
obj=self,
kind="terminal",
attrs={
"instrument": self.spirack.name,
"module": self.module.name,
"terminal": self.name,
"role": self.role,
"module_type": self.module.module_type,
"slot": self.module.slot,
},
)
[docs]
def graph_edges(self) -> Iterable[GraphEdge]:
return ()
[docs]
def set_voltage(self, voltage: VoltageLike) -> None:
"""Set the voltage of the terminal. We delegate that logic to the D5a device."""
voltage = as_volts(voltage)
logging.info("Setting voltage of %s to %s V", self.node_id, voltage)
d5a = self.parent_ref
d5a.set_voltage(self.channel, voltage)
[docs]
def get_voltage(self) -> float:
"""Get the voltage of the terminal."""
d5a = self.parent_ref
return d5a.get_voltage(self.channel)
[docs]
@dataclass
class SPIRackD5aModule:
parent_ref: SPIRackBlueprint
name: str
module_type: str
slot: int
channels: list[SPIRackD5aTerminal]
[docs]
@classmethod
def from_instrument_spec(cls, spec: QCoDeSInstrumentSpec, parent_ref: SPIRackBlueprint) -> SPIRackD5aModule:
name, config = spec
module = cls(
parent_ref=parent_ref,
name=name,
module_type=config["type"],
slot=config["slot"],
channels=[],
)
module.channels = [
SPIRackD5aTerminal(
parent_ref=module,
name=channel["name"],
channel=channel["channel"],
role=channel.get("role"),
)
for channel in config.get("channels", [])
]
return module
[docs]
def graph_nodes(self) -> Iterable[GraphNode]:
for channel in self.channels:
yield from channel.graph_nodes()
[docs]
def graph_edges(self) -> Iterable[GraphEdge]:
for channel in self.channels:
yield from channel.graph_edges()
def _live_channel(self, channel: int):
"""Looks up the live QCoDeS D5aChannel; requires the parent SPIRack to be connected."""
instrument = self.parent_ref._qcodes_instrument
if instrument is None:
raise RuntimeError(f"{self.parent_ref.name} is not connected; call connect() first")
name = next(t.name for t in self.channels if t.channel == channel)
return getattr(instrument.modules[self.name], name)
[docs]
def get_voltage(self, channel: int) -> float:
return self._live_channel(channel).get_voltage()
[docs]
def set_voltage(self, channel: int, voltage: float) -> None:
self._live_channel(channel).set_voltage(voltage)
[docs]
@dataclass
class SPIRackStaticTerminal:
parent_ref: SPIRackStaticModule
name: str
@property
def module(self) -> SPIRackStaticModule:
return self.parent_ref
@property
def node_id(self) -> str:
return terminal_id(self.module.name, self.name)
[docs]
def graph_nodes(self) -> Iterable[GraphNode]:
yield GraphNode(
id=self.node_id,
obj=self,
kind="terminal",
attrs={
"instrument": self.module.parent_ref.name,
"module": self.module.name,
"terminal": self.name,
"module_type": self.module.module_type,
"slot": self.module.slot,
**self.module.attributes,
},
)
[docs]
def graph_edges(self) -> Iterable[GraphEdge]:
return ()
def __getattr__(self, name: str) -> Any:
return getattr(self.module, name)
[docs]
@dataclass
class SPIRackStaticModule:
parent_ref: SPIRackBlueprint
name: str
module_type: str
slot: int
connections: list[SPIRackStaticTerminal]
voltage_conversion_factor: float | None = None
gain_v_per_a: float | None = None
attributes: dict[str, Any] = field(default_factory=dict)
[docs]
@classmethod
def from_instrument_spec(cls, spec: QCoDeSInstrumentSpec, parent_ref: SPIRackBlueprint) -> SPIRackStaticModule:
name, config = spec
module = cls(
parent_ref,
name,
config["type"],
config["slot"],
[],
config.get("voltage_conversion_factor"),
config.get("gain_v_per_a"),
graph_attributes(config),
)
module.connections = [SPIRackStaticTerminal(module, connection["name"]) for connection in config.get("connections", [])]
return module
[docs]
def graph_nodes(self) -> Iterable[GraphNode]:
for connection in self.connections:
yield from connection.graph_nodes()
[docs]
def graph_edges(self) -> Iterable[GraphEdge]:
for connection in self.connections:
yield from connection.graph_edges()
[docs]
@dataclass
class SPIRackBlueprint(QCoDeSInstrumentProxy):
"""Static SPIRack blueprint for station graph routing.
This class does not instantiate the live QCoDeS driver. It only turns
configured module channels and connections into terminal nodes and edges.
"""
name: str
d5a_modules: list[SPIRackD5aModule]
static_modules: list[SPIRackStaticModule]
[docs]
@classmethod
def from_instrument_spec(cls, spec: QCoDeSInstrumentSpec) -> SPIRackBlueprint:
name, config = spec
modules: Mapping[str, Mapping[str, Any]] = config.get("init", {}).get("modules", {})
blueprint = cls(
name=name,
qcodes_config=spec,
d5a_modules=[],
static_modules=[],
)
for module_name, module_config in modules.items():
module_type = module_config.get("type")
if module_type == "semi_cr.core.lab.drivers.spirack.QcodesD5aModule":
blueprint.d5a_modules.append(SPIRackD5aModule.from_instrument_spec((module_name, module_config), blueprint))
elif module_type in {
"semi_cr.core.lab.drivers.spirack.QcodesVIModule",
"semi_cr.core.lab.drivers.spirack.QcodesB1aModule",
"semi_cr.core.lab.drivers.spirack.QcodesM1gModule",
}:
blueprint.static_modules.append(SPIRackStaticModule.from_instrument_spec((module_name, module_config), blueprint))
else:
logging.warning("Static SPIRack module %s has unsupported type %s", module_name, module_type)
return blueprint
@property
def _module_specs(self) -> Mapping[str, Mapping[str, Any]]:
"""Traverses the configuration for any non-hydrated submodules. Used by the graph node and edge exporters."""
_, config = self.qcodes_config
return config.get("init", {}).get("modules", {})
[docs]
def graph_nodes(self) -> Iterable[GraphNode]:
"""Returns terminal nodes for configured SPIRack module terminals."""
# The D5a modules we have a driver for already
d5a_module_names = {module.name for module in self.d5a_modules}
for module in self.d5a_modules:
yield from module.graph_nodes()
static_module_names = {module.name for module in self.static_modules}
for module in self.static_modules:
yield from module.graph_nodes()
# The remaining driverless nodes are still created, but with no `obj`
for module_name, module_config in self._module_specs.items():
if module_name in d5a_module_names or module_name in static_module_names:
continue
for connection in module_config.get("connections", []):
connection_name = connection["name"]
yield GraphNode(
id=terminal_id(module_name, connection_name),
obj=self,
kind="terminal",
attrs={
"instrument": self.name,
"module": module_name,
"terminal": connection_name,
"module_type": module_config.get("type"),
"slot": module_config.get("slot"),
},
)
[docs]
def graph_edges(self) -> Iterable[GraphEdge]:
"""Returns graph edges for configured SPIRack module connections."""
for module_name, module_config in self._module_specs.items():
for connection in module_config.get("connections", []):
local = terminal_id(module_name, connection["name"])
for endpoint in connection.get("endpoints", []):
yield GraphEdge(
id=f"{endpoint}->{local}",
source=endpoint,
target=local,
kind="physical",
)