Source code for atomiq.components.basics.calibration

"""
In all experiments, calibrations are ubiquious. Examples are

  * voltage - power relation on a photodiode
  * Relation between the current through a coil and the created magentic field
  * RF power in an AOM and light power in the diffracted order
  * The current-voltage relation for a voltage-controlled current supply
  * ...

In atomiq calibrations are comoponents just like every other piece of your experiment. Every calibration inherits
from the abstract :class:`Calibration` class. A special subclass of calibration functions are invertable
calibrations that can be analytically inverted. The most frequently used example is a linear calibration function.
Invertable calibrations inherit from :class:`InvertableCalibration`.
"""
from __future__ import annotations

from atomiq.components.primitives import Component

from artiq.experiment import portable
from artiq.language.types import TList, TFloat, TStr, TBool




@portable
def exp(x: TFloat) -> TFloat:
    return 2.7182818**x





@portable
def ln(x: TFloat, n: TFloat = 10000.0) -> TFloat:
    return n * ((x ** (1/n)) - 1)





class Calibration(Component):
    """An abstract Calibration

    This is an abstract class to describe a calibration.

    Args:
        input_unit: A string determining the input unit (e.g. 'mW', 'V', or 'uA')
        output_unit: A string determining the output unit (e.g. 'mW', 'V', or 'uA')
    """
    kernel_invariants = {"input_unit", "output_unit"}

    def __init__(self, input_unit: TStr, output_unit: TStr, *args, **kwargs):
        Component.__init__(self, *args, **kwargs)

        self.input_unit = input_unit
        self.output_unit = output_unit



    @portable
    def transform(self, input_value: TFloat) -> TFloat:
        """Transform a value according to the calibration

        Args:
            input_value: value to be transformed

        Returns:
            TFloat: transformed value
        """
        raise NotImplementedError(self.identifier + " does not have transform() method")






class InvertableCalibration(Calibration):


    @portable
    def transform_inv(self, input_value: TFloat) -> TFloat:
        """Perform inverse transform of a value according to the calibration

        Args:
            input_value: value to be inversely transformed. Must be given in units of the output unit

        Returns:
            TFloat: transformed value. The returned value is in units of the input unit
        """
        raise NotImplementedError(self.identifier + " does not have transform() method")






class DummyCalibration(Calibration):
    def __init__(self, *args, **kwargs):
        Calibration.__init__(self, "", "", *args, **kwargs)



    @portable
    def transform(self, input_value: TFloat) -> TFloat:
        return input_value






class SplineCalibration(Calibration):
    """Calibration via data points

    Data points are interpolated with linear splines

    Args:
        calibration_points: List of tuples (x, y) containing the calibration data. The data must be ordered
            monotonously in `x`
    """

    kernel_invariants = {"calibration_points"}

    def __init__(self, calibration_points: TList(TList(TFloat)), *args, **kwargs):
        Calibration.__init__(self, *args, **kwargs)

        self.calibration_points = calibration_points



    @portable(flags={"fast-math"})
    def transform(self, input_value: TFloat, invert: TBool = False) -> TFloat:
        input_idx = 0 if not invert else 1
        output_idx = 1 if not invert else 0
        search_direction = 0 if self.calibration_points[0][input_idx] < self.calibration_points[-1][input_idx] else 1

        assert self.calibration_points[0][input_idx] < input_value < self.calibration_points[-1][input_idx] \
            or self.calibration_points[0][input_idx] > input_value > self.calibration_points[-1][input_idx], \
            "input value outside of calibration range!"

        # find the two neighboring calibration points
        point_below, point_above = [0., 0.], [1., 1.]
        for i in range(len(self.calibration_points)):
            if (search_direction == 0 and self.calibration_points[i][input_idx] < input_value) \
                    or (search_direction == 1 and self.calibration_points[i][input_idx] > input_value):
                continue
            else:
                point_above = self.calibration_points[i]
                point_below = self.calibration_points[i-1]
                break

        # interpolate
        m = (point_above[output_idx] - point_below[output_idx]) / (point_above[input_idx] - point_below[input_idx])
        b = point_below[output_idx] - m * point_below[input_idx]
        return m*input_value + b






class InvertableSplineCalibration(InvertableCalibration, SplineCalibration):
    """Calibration via data points that can be inverted

    Data points are interpolated with linear splines. For the inversion to work, both `x` and `y` of the calibration
    data must be monotonous.
    """
    def __init__(self, *args, **kwargs):
        InvertableCalibration.__init__(self, *args, **kwargs)
        SplineCalibration.__init__(self, *args, **kwargs)



    @portable(flags={"fast-math"})
    def transform_inv(self, input_value: TFloat) -> TFloat:
        return self.transform(input_value, invert=True)






class PolynomialCalibration(Calibration):
    """Calibration described by a polynomial

    The calibration is given by the function

    $$f(x) = \\sum_i c_i x^i$$

    Args:
        coefficients: List of coefficients $c_i$ of the polynomial, start from the lowest order.
    """
    kernel_invariants = {"coefficients"}

    def __init__(self, *args, coefficients: TList(TFloat), **kwargs):
        Calibration.__init__(self, *args, **kwargs)
        self.coefficients = coefficients



    @portable(flags={"fast-math"})
    def transform(self, input_value: TFloat) -> TFloat:
        # very non-pythonic so the artiq compiler takes it...
        sum = 0.0
        i = 0
        while i < len(self.coefficients):
            sum += self.coefficients[i] * input_value**i
            i += 1
        return sum






class LinearCalibration(InvertableCalibration):
    """Linear calibration

    The calibration is given by the function

    $$f(x) = ax + b$$

    Args:
        input_unit: Unit of the input
        output_unit: Unit of the output
        a: Calibration coefficient a
        b: Calibration coefficient b
    """

    kernel_invariants = {"a", "b"}

    def __init__(self, a: TFloat, b: TFloat, *args, **kwargs):
        InvertableCalibration.__init__(self, *args, **kwargs)
        self.a = a
        self.b = b



    @portable(flags={"fast-math"})
    def transform(self, input_value: TFloat) -> TFloat:
        return self.a*input_value + self.b




    @portable(flags={"fast-math"})
    def transform_inv(self, output_value: TFloat) -> TFloat:
        return (output_value - self.b)/self.a






class SigmoidCalibration(Calibration):
    """Sigmoid calibration

    ```
    output = A / ( 1 + e^k*(input - x_offset)) + y_offset
    ```

    Args:
        input_unit: Unit of the input
        output_unit: Unit of the output
        A: Amplitude of the sigmoid
        k: stretching of the sigmoid
        x_offset: offset on the x axis
        y_offset: offset on the y axis

    """

    kernel_invariants = {"A", "k", "x_offset", "y_offset"}

    def __init__(self,  *args, A: TFloat, k: TFloat, x_offset: TFloat = 0, y_offset: TFloat = 0, **kwargs):
        Calibration.__init__(self, *args, **kwargs)

        self.A = A
        self.k = k
        self.x_offset = x_offset
        self.y_offset = y_offset



    @portable(flags={"fast-math"})
    def transform(self, input_value: TFloat) -> TFloat:
        return self.A/(1+exp(-self.k*(input_value - self.x_offset))) + self.y_offset






class InvSigmoidCalibration(Calibration):
    """Inverse sigmoid calibration

    ```
    output = -ln( A / (input - y_offset) - 1) / k + x_offset
    ```

    The paremeters are defined in a way that they match the sigmoid definition.

    Args:
        input_unit: Unit of the input
        output_unit: Unit of the output
        A: Amplitude of the sigmoid
        k: stretching of the sigmoid
        x_offset: offset on the x axis
        y_offset: offset on the y axis

    """

    kernel_invariants = {"A", "k", "x_offset", "y_offset"}

    def __init__(self,  *args, A: TFloat, k: TFloat, x_offset: TFloat = 0, y_offset: TFloat = 0, **kwargs):
        Calibration.__init__(self,  *args, **kwargs)

        self.A = A
        self.k = k
        self.x_offset = x_offset
        self.y_offset = y_offset



    @portable(flags={"fast-math"})
    def transform(self, input_value: TFloat) -> TFloat:
        return -1 * ln(self.A/(input_value - self.y_offset) - 1)/self.k + self.x_offset