Source code for atomiq.components.optimizers

from __future__ import annotations

from atomiq.components.primitives import Component, Parametrizable, Measurable, Switchable
from atomiq.helper import replace_member

from artiq.experiment import kernel, delay
from artiq.language.types import TStr, TFloat, TBool
from artiq.language.units import us




class Optimizer(Component):

    kernel_invariants = {"actor", "monitor", "actor_name"}

    def __init__(self,
                 actor_component: Parametrizable,
                 actor_name: TStr, monitor_component: Measurable,
                 monitor_name: TStr,
                 *args, **kwargs):
        Component.__init__(self, *args, **kwargs)

        self.actor = actor_component
        self.monitor = monitor_component
        self.actor_name = actor_name

        # TODO: This is bit hacky since we modify the objects we act on. Better solutions welcome..
        replace_member(self.monitor, "_measure_for_optimizer", f"get_{monitor_name}")
        replace_member(self.actor, "_set_for_optimizer", f"set_{actor_name}")



    def optimize(self):
        raise NotImplementedError("Implement the optimzer")






class BoundOptimizer(Optimizer):
    """An Optimizer with restrictions on the values of the actor.

    Args:
        actor_min: The minumum value the actor is allowed to take during optimization
        actor_max: The maxiumum value the actor is allowed to take during optimization
    """

    kernel_invariants = {"actor_min", "actor_max"}

    def __init__(self, actor_min: TFloat, actor_max: TFloat, *args, **kwargs):
        Optimizer.__init__(self, *args, **kwargs)

        self.actor_min = actor_min
        self.actor_max = actor_max





class BisectionOptimizer(BoundOptimizer):
    """An optimizer that uses bisection to achieve the optimization target on the monitor

    Args:
        timestep: The time the individual steps in the algorithm take
        epsilon: Acceptable relative difference from the target value (default 0.01, i.e. 1% of the target value)
        switch_actor: If True, the actor (needs to be Switchable) is switched on before the measurement of the monitor
                      and switched off after (default False)
    """

    kernel_invariants = {"timestep", "epsilon", "max_steps", "switch_actor"}

    def __init__(self, timestep=325*us, epsilon=0.01, max_steps=20, switch_actor: TBool = False, *args, **kwargs):
        BoundOptimizer.__init__(self, *args, **kwargs)

        self.timestep = timestep
        self.epsilon = epsilon
        self.max_steps = max_steps

        @kernel
        def dummy_function(self):
            self.experiment.log.warning(self.identifier + ": cannot switch actor "+self.actor.identifier)

        self.switch_actor = switch_actor
        if not isinstance(self.actor, Switchable):
            self._actor_on = dummy_function.__get__(self)
            self._actor_off = dummy_function.__get__(self)

    @kernel
    def _actor_on(self):
        self.actor.on()

    @kernel
    def _actor_off(self):
        self.actor.off()



    @kernel(flags={"fast-math"})
    def optimize(self, target: TFloat) -> TFloat:

        soft_max = self.actor_max
        soft_min = self.actor_min

        # initial values for the bisection
        trial = (self.actor_max - self.actor_min) / 2
        measurement = 0.0

        # Do the bisection iteration
        for i in range(self.max_steps):
            delay(self.timestep / 2)

            # calculate and set
            trial = soft_min / 2 + soft_max / 2
            if trial < self.actor_min:
                trial = self.actor_min
            if trial > self.actor_max:
                trial = self.actor_max
            # self.log.info(trial)
            self.actor._set_for_optimizer(trial)
            delay(self.timestep)

            # measure
            if self.switch_actor:
                self._actor_on()
            delay(self.timestep)
            measurement = self.monitor._measure_for_optimizer()
            delay(self.timestep)
            if self.switch_actor:
                self._actor_off()
            self.experiment.log.debug("Bisection measurement result: {0}V ", [measurement])

            # compare
            difference = measurement - target

            if abs(difference / target) < self.epsilon:
                # self.log.info('eps')
                break

            # change borders for next iteration
            if difference > 0.0:
                soft_max = trial
            else:
                soft_min = trial
            delay(self.timestep / 2)

        # delay to avoid collisions
        delay(0.03*us)

        return trial