Source code for semi_cr.core.lab.station.routing_v1.configurable

from dataclasses import dataclass
import networkx as nx



@dataclass(frozen=True)
class RouteConfiguration:
    instrument_name: str
    channel: int
    polarity: str


# Pure Graph operations



def extract_route_configurations(
    graph: nx.MultiDiGraph,
    path: list[str],
) -> list[RouteConfiguration]:

    configurations: set[RouteConfiguration] = set()

    required_attrs = {
        "instrument_name",
        "channel",
        "polarity",
    }

    for source, target in zip(path, path[1:]):

        edge_data = graph.get_edge_data(
            source,
            target,
        )

        if edge_data is None:
            raise RuntimeError(
                f"No edge found between "
                f"{source!r} and {target!r}."
            )

        for attrs in edge_data.values():

            if attrs.get("kind") != "configurable":
                continue

            missing = required_attrs - attrs.keys()

            if missing:
                raise ValueError(
                    f"Configurable edge "
                    f"{source!r} -> {target!r} "
                    f"is missing attributes: "
                    f"{sorted(missing)}"
                )

            configurations.add(
                RouteConfiguration(
                    instrument_name=str(
                        attrs["instrument_name"]
                    ),
                    channel=int(
                        attrs["channel"]
                    ),
                    polarity=str(
                        attrs["polarity"]
                    ),
                )
            )

    return list(configurations)




def find_active_configurations(
    graph: nx.MultiDiGraph,
    instrument_name: str,
    polarity: str | None = None,
) -> list[RouteConfiguration]:

    found: set[RouteConfiguration] = set()

    for _, _, attrs in graph.edges(data=True):

        if attrs.get("kind") != "configurable":
            continue

        if not attrs.get("active", False):
            continue

        if (
            attrs.get("instrument_name")
            != instrument_name
        ):
            continue

        if (
            polarity is not None
            and attrs.get("polarity") != polarity
        ):
            continue

        found.add(
            RouteConfiguration(
                instrument_name=str(
                    attrs["instrument_name"]
                ),
                channel=int(attrs["channel"]),
                polarity=str(attrs["polarity"]),
            )
        )

    return list(found)


def _set_route_configuration_active(
    graph: nx.MultiDiGraph,
    configuration: RouteConfiguration,
    active: bool,
) -> None:

    matched = False

    for source, target, key, attrs in graph.edges(
        keys=True,
        data=True,
    ):

        if attrs.get("kind") != "configurable":
            continue

        if (
            attrs.get("instrument_name")
            != configuration.instrument_name
        ):
            continue

        if int(attrs.get("channel", -1)) != configuration.channel:
            continue

        if attrs.get("polarity") != configuration.polarity:
            continue

        attrs["active"] = active
        matched = True

    if not matched:
        raise LookupError(
            "No configurable graph edge matches "
            f"{configuration!r}."
        )


# Hardware-facing functions



def apply_route_configuration(
    context,
    configuration: RouteConfiguration,
) -> None:

    instrument_name = (
        configuration.instrument_name
    )

    try:
        instrument = context.instruments[
            instrument_name
        ]
    except KeyError as exc:
        raise KeyError(
            f"Routing instrument "
            f"{instrument_name!r} is not connected."
        ) from exc

    apply_configuration = getattr(
        instrument,
        "apply_route_configuration",
        None,
    )

    if not callable(apply_configuration):
        raise TypeError(
            f"Instrument {instrument_name!r} "
            "does not support configurable routing."
        )

    apply_configuration(
        channel=configuration.channel,
        polarity=configuration.polarity,
    )





def activate_route_configuration(
    context,
    configuration: RouteConfiguration,
) -> None:

    # instrument_name = configuration.instrument_config_name

    # try:
    #     instrument = context.instruments[instrument_name]
    # except KeyError as exc:
    #     raise KeyError(
    #         f"Routing instrument {instrument_name!r} is not connected."
    #     ) from exc

    # apply_configuration = getattr(
    #     instrument,
    #     "apply_route_configuration",
    #     None,
    # )

    # if not callable(apply_configuration):
    #     raise TypeError(
    #         f"Instrument {instrument_name!r} "
    #         "does not support configurable routing."
    #     )

    # Find currently active routes on this same bank.
    active_configurations = find_active_configurations(
        context.graph,
        instrument_name=configuration.instrument_name,
        polarity=configuration.polarity,
    )

    # Hardware first
    apply_route_configuration(
        context,
        configuration,
    )

    # Graph state only changes after successful hardware configuration.
    for old_configuration in active_configurations:

        if old_configuration == configuration:
            continue

        _set_route_configuration_active(
            context.graph,
            old_configuration,
            active=False,
        )

    _set_route_configuration_active(
        context.graph,
        configuration,
        active=True,
    )




def deactivate_route_configuration(
    context,
    configuration: RouteConfiguration,
) -> None:

    instrument_name = configuration.instrument_name

    try:
        instrument = context.instruments[instrument_name]
    except KeyError as exc:
        raise KeyError(
            f"Routing instrument {instrument_name!r} is not connected."
        ) from exc

    deactivate_configuration = getattr(
        instrument,
        "deactivate_route_configuration",
        None,
    )

    if not callable(deactivate_configuration):
        raise TypeError(
            f"Instrument {instrument_name!r} "
            "does not support configurable routing."
        )

    # Hardware first.
    deactivate_configuration(
        channel=configuration.channel,
        polarity=configuration.polarity,
    )

    # Graph only changes if hardware succeeded.
    _set_route_configuration_active(
        context.graph,
        configuration,
        active=False,
    )