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

from __future__ import annotations
from dataclasses import dataclass
from typing import Type

import networkx as nx

from semi_cr.core.lab.devices.pinned import PinnedDevice
from semi_cr.core.lab.station.routing_v1.quantities import (
    RoutedVoltage,
    RoutedCurrent,
    RoutedRFDrive,
    RoutedIQReceiver,
    RoutedRFFrequency,
    RoutedPulseAmplitude,
)
from semi_cr.core.lab.pins.pins import (
    Ohmic,
    Source,
    Drain, 
    Gate,
    Plunger,
    RFInput,
    RFOutput,
)



@dataclass(frozen=True)
class RoutedQuantitySpec:
    name: str
    quantity_class: Type
    modality: str


ELECTRICAL_QUANTITIES = (
    RoutedQuantitySpec("dc_voltage", RoutedVoltage, "dc"),
    RoutedQuantitySpec("ac_voltage", RoutedVoltage, "ac"),
    RoutedQuantitySpec("dc_current", RoutedCurrent, "dc"),
    RoutedQuantitySpec("ac_current", RoutedCurrent, "ac"),
)

RF_INPUT_QUANTITIES = (
    RoutedQuantitySpec("rf_drive", RoutedRFDrive, "rf"),
    RoutedQuantitySpec("iq_receiver", RoutedIQReceiver, "rf"),
    RoutedQuantitySpec("rf_frequency", RoutedRFFrequency, "rf"),
    RoutedQuantitySpec("pulse_amplitude", RoutedPulseAmplitude, "rf"),
)

RF_OUTPUT_QUANTITIES = (
    RoutedQuantitySpec("iq_receiver", RoutedIQReceiver, "rf"),
)

PIN_QUANTITY_SPECS = {
    Source: ELECTRICAL_QUANTITIES,
    Drain: ELECTRICAL_QUANTITIES,
    Ohmic: ELECTRICAL_QUANTITIES,
    Gate: ELECTRICAL_QUANTITIES,
    Plunger: (*ELECTRICAL_QUANTITIES, *RF_INPUT_QUANTITIES),
    RFInput: RF_INPUT_QUANTITIES,
    RFOutput: RF_OUTPUT_QUANTITIES,
}



def quantities_for_pin(pin):
    # Exact type first
    specs = PIN_QUANTITY_SPECS.get(type(pin))

    if specs is not None:
        return specs

    # Then allow inheritance
    for pin_type, specs in PIN_QUANTITY_SPECS.items():
        if isinstance(pin, pin_type):
            return specs

    return ()




def attach_routed_pin_quantities(
    graph: nx.MultiDiGraph,
    device: PinnedDevice,
    replace: bool = False,
) -> None:
    """
    Attach the routed quantities supported by each device pin.

    Parameters
    ----------
    graph:
        Effective station graph used for route resolution.

    device:
        Device whose pins receive routed quantities.

    replace:
        If False, preserve quantities already attached to a pin.
        If True, rebuild them using the supplied graph.
    """

    for pin in device.pins:
        specs = quantities_for_pin(pin)

        for spec in specs:
            if (
                not replace
                and hasattr(pin, spec.name)
            ):
                continue
            quantity = spec.quantity_class(
                graph=graph,
                pin=pin,
                modality=spec.modality,
            )

            setattr(pin, spec.name, quantity)