Source code for atomiq.components

"""
The atomiq components are a mechanism to put the complexity of the experimental setup into a configuration and to allow
to work with the setup at a higher abstraction level. The key ideas behind it are:

  * Have abstract, generic classes to represent different types of lab equipment
    (e.g. :class:`~electronics.rfsource.RFSource`, :class:`~electronics.voltagesource.DAC`,
    :class:`~optoelectronics.lightmodulator.AOM`, :class:`~coil.Coil`, :class:`~laser.Laser`, etc)
  * Derive more specific classes from the generic ones. For example a channel of the Sinara Urukul is an RFSource but
    also a standalone AOM driver that is controlled by analog voltages is an RFSource. Thus, both,
    :class:`~sinara.urukul.UrukulChannel` and :class:`~electronics.rfsource.VoltageControlledRFSource` inherit from
    :class:`~electronics.rfsource.RFSource`.
  * Allow to reference other components in a component definition and recursively build all need components.

These ideas allow to move hardware specific information (like what devices are used and what device is plugged to which
other device) into a configuration and to work with generic software objects that represent the actual hardware in the
lab (like AOMs, Coils, Lasers, etc). It allows to define your experiment hardware in a simple dictionary and to
dynamically link them together.

To get the idea, lets start with a simple example:

.. literalinclude:: ../examples/aom/components.py
   :language: python

Each entry in the components dict is identified by a name that can be chosen to your liking.

.. note::
   Since the component dict is flat, it might make sense to think about a proper naming scheme like
   `<component_type>_<component_name>` (e.g. `aom_cooler`, `pd_repumper`, `coil_north`, `dds_cooler`, etc.).

Every component needs to define a `classname` which defines the class the software object should have. Atomiq
automatically builds these objects from the defined class, if you define that you require them in your experiment
code. The arguments field takes the parameters that are needed to create the object. Look for the constructor method
of the defined class to find what arguments are required/supported. These arguments are the basic means of adding
default configuration for your components and to interconnect them.

Typically atomiq components rely on other components to do their job. It is thus necessary to reference to other
components when writing the configuration for a certain component. Such interconnections between components are
provided by starting an argument either with "&" (which refers to another component) or with "@" with refers to
a device in your ARTIQ `device_db.py`. When creating the objects from the configuration, atomiq detects these
references and recursively builds the components required by the component you need.

.. warning::
   Please note that a restart of your ARTIQ master is required if you change your components definition since the
   `device_db` is only read upon starting the master.

Once you defined your components you can request atomiq to build it for you by adding it to the class-level components
list of your experiment class

Atomiq comes with a lot of the most used classes already. See the following list to see what exists.
"""
import artiq
import artiq.master.worker_db
from atomiq.helper import get_class_by_name


import logging
import random
import string

logging.basicConfig()
logger = logging.getLogger(__name__)


components = None

suservo_replacements = []


def _create_device_replace(desc, device_mgr, *args, **kwargs):
    global suservo_replacements
    if desc in suservo_replacements:
        desc["class"] += "_Suservo"
        for i, existing in enumerate(device_mgr.active_devices):
            # Avoid double init due to change in class name from original class as in ddb
            if desc == existing[0]:
                del device_mgr.active_devices[i]
                return existing[1]
    return artiq.master.worker_db._create_device_original(desc, device_mgr, *args, **kwargs)

artiq.master.worker_db._create_device_original = artiq.master.worker_db._create_device
artiq.master.worker_db._create_device = _create_device_replace




def build_object_from_device_db(parent, object_id):
    global suservo_replacements
    logger.debug(f"building device_db object {object_id} with parent {parent}")

    # workaroud for https://forum.m-labs.hk/d/273-using-su-servo-mode-and-freestanding-urukul-in-the-same-experiment/3
    # this should not be necessary anymore once we run nac3
    def fixup_suservo_class(node):
        for classname in ["CPLD", "AD9910"]:
            if "class" in node and node["class"] == classname:
                suservo_replacements.append(node)
                logger.debug(f"fixing class for {node}")

    if "compiler" in dir(artiq):
        ddb = parent.get_device_db()
        if object_id in ddb:
            entry = ddb[object_id]
            logger.info(f"checking {object_id}")
            if "class" in entry and entry["class"] == "SUServo":
                for node in entry["arguments"]["cpld_devices"]:
                    fixup_suservo_class(ddb[node])
                for node in entry["arguments"]["dds_devices"]:
                    fixup_suservo_class(ddb[node])

    return parent.get_device(object_id)





def build_object(classname, arg_dict):
    logger.debug(f"building object of class {classname}")

    # fixup for artiq compilers prior to nac3 where the type inferral complains when two objects of the
    # same class differ in the types of the attributes. Here we create unique dummy classes for each object
    # to work around the issue. When we are running on the legacy compiler artiq.compiler submodule should exist.
    if "compiler" in dir(artiq):
        ephemeral_classname = f"{classname}_{''.join(random.choice(string.ascii_lowercase) for i in range(8))}"
        logger.debug(f"apply workaround for old artiq type inferral {classname} -> {ephemeral_classname}")
        target_class = type(ephemeral_classname, (get_class_by_name(classname),), {})
    else:
        target_class = get_class_by_name(classname)

    return target_class(**arg_dict)





class ComponentFactory():


    @staticmethod
    def produce(name, parent):
        global components
        if components is None:
            ddb = parent.get_device_db()
            components = ddb["components"] if "components" in ddb else {}

        logger.debug(f"Try to produce {name}")

        if name in components:
            return ComponentFactory._produce_from_dict(name, parent, components[name])
        else:
            raise KeyError(f"component '{name}' required by {parent.__class__.__name__} does not exist in component"\
                           "definition dict")


    @staticmethod
    def _produce_from_dict(identifier, parent, dictionary):
        global components

        if "obj" in dictionary:
            # component was already built. Return it
            return dictionary["obj"]
        else:
            # Build component. If arguments are themselfes components, build them recursively
            if "classname" not in dictionary:
                raise KeyError(f"component '{identifier}' required by {parent.__class__.__name__} does not have the"\
                        "mandatory argument `classname`")
            arg_dict = {"identifier": identifier, "parent": parent}
            if "arguments" in dictionary and type(dictionary["arguments"]) is dict:
                for argname, argval in dictionary["arguments"].items():
                    if type(argval) is str and argval.startswith("&"):
                        obj = ComponentFactory.produce(argval[1:], parent)
                        arg_dict[argname] = obj
                    elif type(argval) is str and argval.startswith("@"):
                        arg_dict[argname] = build_object_from_device_db(parent, argval[1:])
                    else:
                        arg_dict[argname] = argval

            try:
                obj = build_object(dictionary["classname"], arg_dict)
                dictionary["obj"] = obj
                return obj
            except Exception as e:
                # Make a more user friendly output that names the component that cannot be built
                logger.error(f"Cannot build component {identifier} due to following error: {e}")
                raise e