sampler: port to nac3

artiq/coredevice/sampler.py

@@ -1,13 +1,24 @@
-from artiq.language.core import kernel, delay, portable
+from numpy import int32
+
+from artiq.language.core import nac3, kernel, portable, Kernel, KernelInvariant
 from artiq.language.units import ns
 
-from artiq.coredevice import spi2 as spi
+from artiq.coredevice.core import Core
+from artiq.coredevice.spi2 import *
+from artiq.coredevice.ttl import TTLOut
 
 
-SPI_CONFIG = (0*spi.SPI_OFFLINE | 0*spi.SPI_END |
+# NAC3TODO work around https://git.m-labs.hk/M-Labs/nac3/issues/189
-              0*spi.SPI_INPUT | 0*spi.SPI_CS_POLARITY |
+@nac3
-              0*spi.SPI_CLK_POLARITY | 0*spi.SPI_CLK_PHASE |
+class ValueError(Exception):
-              0*spi.SPI_LSB_FIRST | 0*spi.SPI_HALF_DUPLEX)
+    pass
+
+
+
+SPI_CONFIG = (0*SPI_OFFLINE | 0*SPI_END |
+              0*SPI_INPUT | 0*SPI_CS_POLARITY |
+              0*SPI_CLK_POLARITY | 0*SPI_CLK_PHASE |
+              0*SPI_LSB_FIRST | 0*SPI_HALF_DUPLEX)
 
 
 SPI_CS_ADC = 0  # no CS, SPI_END does not matter, framing is done with CNV
@@ -15,26 +26,28 @@ SPI_CS_PGIA = 1  # separate SPI bus, CS used as RCLK
 
 
 @portable
-def adc_mu_to_volt(data, gain=0):
+def adc_mu_to_volt(data: int32, gain: int32 = 0) -> float:
     """Convert ADC data in machine units to Volts.
 
     :param data: 16 bit signed ADC word
     :param gain: PGIA gain setting (0: 1, ..., 3: 1000)
     :return: Voltage in Volts
     """
+    volt_per_lsb = 0.
     if gain == 0:
-        volt_per_lsb = 20./(1 << 16)
+        volt_per_lsb = 20./float(1 << 16)
     elif gain == 1:
-        volt_per_lsb = 2./(1 << 16)
+        volt_per_lsb = 2./float(1 << 16)
     elif gain == 2:
-        volt_per_lsb = .2/(1 << 16)
+        volt_per_lsb = .2/float(1 << 16)
     elif gain == 3:
-        volt_per_lsb = .02/(1 << 16)
+        volt_per_lsb = .02/float(1 << 16)
     else:
         raise ValueError("invalid gain")
-    return data*volt_per_lsb
+    return float(data)*volt_per_lsb
 
 
+@nac3
 class Sampler:
     """Sampler ADC.
 
@@ -50,7 +63,12 @@ class Sampler:
         between experiments.
     :param core_device: Core device name
     """
-    kernel_invariants = {"bus_adc", "bus_pgia", "core", "cnv", "div"}
+    core: KernelInvariant[Core]
+    bus_adc: KernelInvariant[SPIMaster]
+    bus_pgia: KernelInvariant[SPIMaster]
+    cnv: KernelInvariant[TTLOut]
+    div: KernelInvariant[int32]
+    gains: Kernel[int32]
 
     def __init__(self, dmgr, spi_adc_device, spi_pgia_device, cnv_device,
                  div=8, gains=0x0000, core_device="core"):
@@ -69,13 +87,13 @@ class Sampler:
 
         Sets up SPI channels.
         """
-        self.bus_adc.set_config_mu(SPI_CONFIG | spi.SPI_INPUT | spi.SPI_END,
+        self.bus_adc.set_config_mu(SPI_CONFIG | SPI_INPUT | SPI_END,
                                    32, self.div, SPI_CS_ADC)
-        self.bus_pgia.set_config_mu(SPI_CONFIG | spi.SPI_END,
+        self.bus_pgia.set_config_mu(SPI_CONFIG | SPI_END,
                                     16, self.div, SPI_CS_PGIA)
 
     @kernel
-    def set_gain_mu(self, channel, gain):
+    def set_gain_mu(self, channel: int32, gain: int32):
         """Set instrumentation amplifier gain of a channel.
 
         The four gain settings (0, 1, 2, 3) corresponds to gains of
@@ -91,21 +109,21 @@ class Sampler:
         self.gains = gains
 
     @kernel
-    def get_gains_mu(self):
+    def get_gains_mu(self) -> int32:
         """Read the PGIA gain settings of all channels.
 
         :return: The PGIA gain settings in machine units.
         """
-        self.bus_pgia.set_config_mu(SPI_CONFIG | spi.SPI_END | spi.SPI_INPUT,
+        self.bus_pgia.set_config_mu(SPI_CONFIG | SPI_END | SPI_INPUT,
                                     16, self.div, SPI_CS_PGIA)
         self.bus_pgia.write(self.gains << 16)
-        self.bus_pgia.set_config_mu(SPI_CONFIG | spi.SPI_END,
+        self.bus_pgia.set_config_mu(SPI_CONFIG | SPI_END,
                                     16, self.div, SPI_CS_PGIA)
         self.gains = self.bus_pgia.read() & 0xffff
         return self.gains
 
     @kernel
-    def sample_mu(self, data):
+    def sample_mu(self, data: list[int32]):
         """Acquire a set of samples.
 
         Perform a conversion and transfer the samples.
@@ -119,8 +137,8 @@ class Sampler:
             The `data` list will always be filled with the last item
             holding to the sample from channel 7.
         """
-        self.cnv.pulse(30*ns)  # t_CNVH
+        self.cnv.pulse(30.*ns)  # t_CNVH
-        delay(450*ns)  # t_CONV
+        self.core.delay(450.*ns)  # t_CONV
         mask = 1 << 15
         for i in range(len(data)//2):
             self.bus_adc.write(0)
@@ -131,7 +149,7 @@ class Sampler:
             data[i - 1] = -(val & mask) + (val & ~mask)
 
     @kernel
-    def sample(self, data):
+    def sample(self, data: list[float]):
         """Acquire a set of samples.
 
         .. seealso:: :meth:`sample_mu`
@@ -139,7 +157,7 @@ class Sampler:
         :param data: List of floating point data samples to fill.
         """
         n = len(data)
-        adc_data = [0]*n
+        adc_data = [0 for _ in range(n)]
         self.sample_mu(adc_data)
         for i in range(n):
             channel = i + 8 - len(data)