from __future__ import annotations
import enum
import logging
from dataclasses import dataclass
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from collections.abc import Iterable
from semi_cr.core.station.instruments.qcodes_proxy import QCoDeSInstrumentSpec
from semi_cr.core.trial.machine import OPXInputRef, OPXOutputMode, OPXOutputRef
from semi_cr.core.lab.station.graphing.ids import terminal_id
from semi_cr.core.station.graph import GraphEdge, GraphNode
from semi_cr.core.station.instruments import LabInstrument
[docs]
@dataclass
class OPX1000LFTerminal:
# We might be tempted to let this inherit from some ACDriver maybe, but let's hold off until we have more than
# one separate driver for AC stuff. By not generalizing too much at this point we can tightly couple to QUA(M)
# for now. I think that's a reasonable thing to do presently (JC / 2026-09-25)
[docs]
class Role(enum.Enum):
Source = enum.auto()
Receiver = enum.auto()
parent_ref: OPX1000LFFEM
name: str
role: Role
@property
def fem(self) -> OPX1000LFFEM:
return self.parent_ref
@property
def node_id(self) -> str:
return terminal_id(f"{self.opx1000.name}/{self.fem.name}", self.name)
@property
def opx1000(self) -> OPX1000Blueprint:
return self.parent_ref.parent_ref
[docs]
def graph_nodes(self) -> Iterable[GraphNode]:
yield GraphNode(
id=self.node_id,
obj=self,
kind="terminal",
attrs={
"instrument": self.opx1000.name,
"module": self.fem.name,
"terminal": self.name,
"role": self.role.name.lower(),
"module_type": "LF",
"slot": self.fem.slot,
},
)
[docs]
def graph_edges(self) -> Iterable[GraphEdge]:
return ()
[docs]
def to_opx_output_ref(
self,
upsampling_mode: str = "pulse",
output_mode: OPXOutputMode = "direct",
) -> OPXOutputRef:
from semi_cr.core.trial.machine import OPXOutputRef # noqa: PLC0415
if self.role == OPX1000LFTerminal.Role.Receiver:
raise ValueError("OPXOutputRef not supported for Receiver terminals.")
if output_mode not in {"direct", "amplified"}:
raise ValueError(f"Unsupported LF-FEM output mode: {output_mode!r}")
return OPXOutputRef(
controller_id=self.fem.controller,
fem_id=self.fem.slot,
port_id=int(self.name.removeprefix("out")),
upsampling_mode=upsampling_mode,
output_mode=output_mode,
)
[docs]
@dataclass
class OPX1000LFFEM:
# We can chop up this model to make OPX1000RFFEM with a common ancestor when we need it.
parent_ref: OPX1000Blueprint
name: str
controller: str
slot: int
connections: list[OPX1000LFTerminal]
[docs]
@classmethod
def from_instrument_spec(cls, spec: QCoDeSInstrumentSpec, parent_ref: OPX1000Blueprint) -> OPX1000LFFEM:
name, config = spec
lffem = cls(
parent_ref=parent_ref,
name=name,
controller=config["controller"],
slot=config["slot"],
connections=[],
)
# Iterate over the `connections` here to represent those. The connectors are so simple that I just parse their
# description here, instead of adding another layer of from_instrument_spec.
for connection in config["connections"]:
if connection["role"] == "source":
role = OPX1000LFTerminal.Role.Source
elif connection["role"] == "receiver":
role = OPX1000LFTerminal.Role.Receiver
else:
raise ValueError("Invalid role specified.")
lffem.connections.append(
OPX1000LFTerminal(
parent_ref=lffem,
name=connection["name"],
role=role,
)
)
return lffem
[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 OPX1000Blueprint(LabInstrument):
"""Static OPX1000 blueprint for station graph routing.
This class does not instantiate the live QCoDeS/QUA driver. It only turns
configured module connections into terminal nodes.
"""
name: str
ip_address: str
port: str
modules: list[OPX1000LFFEM]
[docs]
@classmethod
def from_instrument_spec(cls, spec: QCoDeSInstrumentSpec) -> OPX1000Blueprint:
name, config = spec
init_section = config.get("init", {})
# Just do a quick sanity check. In the config we insist that these two values match:
# instruments:
# opx1000_1: <----- this one
# init:
# name: opx1000_1 <----- and this one
if init_section["name"] != name:
raise ValueError("Invalid configuration: init.name not matching the instrument ID.")
blueprint = cls(
name=name,
ip_address=init_section.get("ip_address"),
port=init_section.get("port"),
modules=[],
)
modules = []
# Iterate over the `modules` section of the config, and passing each name-conf tuple to OPX1000LFFEM's inflater
for module_name, module_spec in init_section.get("modules", []).items():
if module_spec["module_type"] != "LF":
logging.warning("For now, only LF-FEMs are supported, skipping one declared module")
continue
fem_module = OPX1000LFFEM.from_instrument_spec(
(module_name, module_spec),
blueprint,
)
modules.append(fem_module)
blueprint.modules = modules
return blueprint
[docs]
def graph_nodes(self) -> Iterable[GraphNode]:
"""Returns terminal nodes for configured OPX module connections."""
for module in self.modules:
yield from module.graph_nodes()
[docs]
def graph_edges(self) -> Iterable[GraphEdge]:
"""Returns graph edges for this OPX blueprint."""
for module in self.modules:
yield from module.graph_edges()