Source code for atomiq.components.electronics.adc

"""
All analog voltages have to go through an ADC at some point to enter the digital world. Often ADCs appear as
multi-channel devices that sample several channels at once. Thus the ADC itself and its channels are represented
by seperate classes. The :class:`ADC` class provides a :func:`ADC.measure` function to sample all channels at once
while the :class:`ADCChannel` provides a :func:`ADCChannel.measure` method to extract only a single channel from the
ADC.
"""

from __future__ import annotations

from atomiq.components.primitives import Component, Measurable

from artiq.experiment import kernel
from artiq.language.types import TFloat, TArray, TInt32, TStr, TBool

import numpy as np




class ADC(Component, Measurable):
    """An analog-to-ditigal converter (ADC) with multiple channels

    Many ADCs come with multiple channels that are always sampled together. This class represents such ADCs

    Args:
        num_chan: Number of channels the ADC has.
    """

    kernel_invariants = {"num_chan", "invalid_return"}

    def __init__(self, num_chan: TInt32, *args, **kwargs):
        Component.__init__(self, *args, **kwargs)

        self.num_chan = num_chan
        self.invalid_return = np.array([-1.0])

        Measurable.__init__(self, [f"voltage{i}" for i in range(self.num_chan)])

    def _build(self):
        self.current_sample = np.zeros(self.num_chan, dtype=float)
        self.current_sample_sum = np.zeros(self.num_chan, dtype=float)

    @kernel
    def _measure(self, target_arr: TArray(TFloat)) -> TArray(TFloat):
        raise NotImplementedError("Implement measure method")



    @kernel(flags={"fast-math"})
    def measure(self, samples: TInt32 = 1, channel: TStr = "") -> TArray(TFloat):
        """
        Measures all channels and returns the mean for a given number of subsequent samples.

        This method advances the timeline by samples * sampling_time

        Args:
            samples: number of samples to take and average
        """

        # current_sample_sum = np.zeros(self.num_chan, dtype=float)

        for i in range(samples):
            self.current_sample_sum += self._measure(self.current_sample)

        self.current_sample_sum /= float(samples)

        if channel == "":
            return self.current_sample_sum
        else:
            for i in range(len(self.channels)):
                if self.channels[i] == channel:
                    return self.current_sample_sum[i:i+1]
            return self.invalid_return





class ADCChannel(Component, Measurable):

    kernel_invariants = {"adc_device", "channel"}

    def __init__(self, adc_device: ADC, channel: TInt32, *args, **kwargs):
        """A single channel of a (possibly multichannel) ADC

        Args:
            adc_device: The ADC the channel belongs to
            channel: The number of the channel
        """
        Component.__init__(self, *args, **kwargs)

        self.adc_device = adc_device
        self.channel = channel

        Measurable.__init__(self, ["voltage"])



    @kernel
    def measure(self, samples: TInt32 = 1, cached: TBool = False, channel: TStr = "") -> TFloat:
        if not cached:
            self.adc_device.measure(samples)

        return self.adc_device.current_sample[self.channel]