awg.py

import os
import spcm
import logging
import time
import argparse
import numpy as np
from itertools import chain
from functools import reduce

loglevel = os.getenv("LOGLEVEL", "WARNING").upper()
logging.basicConfig(level=loglevel)

channel_core_map = [
    list(chain(range(8), range(12, 20))),
    list(chain(range(8, 12), range(20, 21))),
]


def setup_channels(channels: spcm.Channels):
    channels.enable(True)
    channels.amp(2.0 * spcm.units.V)
    channels.output_load(50 * spcm.units.ohm)


def setup_dds(card: spcm.Card):
    card.card_mode(spcm.SPC_REP_STD_DDS)
    card.write_setup()


def setup_dds_cores(dds: spcm.DDS, config: dict):
    amplitude = 40 / len(config)

    for core, freq in config.items():
        logging.debug(f"setting core {core} to {freq} MHz and {amplitude} %")

        dds[core].amp(amplitude * spcm.units.percent)
        dds[core].freq(float(freq) * spcm.units.MHz)
    dds.exec_at_trg()
    dds.write_to_card()


def list():
    devices = [
        os.path.join("/dev", device)
        for device in os.listdir("/dev")
        if device.startswith("spcm")
    ]

    with spcm.CardStack(devices) as stack:
        for card in stack.cards:
            print(card)


def static(devices, chanstrings, freqstrings):
    if len(devices) != 1:
        logging.warning("only one device is supported at the moment")

    with spcm.Card(serial_number=devices[0]) as card:
        channels = spcm.Channels(card)
        setup_channels(card, channels)
        setup_dds(card)

        dds = spcm.DDS(card, channels=channels)
        dds.reset()

        dds.cores_on_channel(
            1,
            spcm.SPCM_DDS_CORE8,
            spcm.SPCM_DDS_CORE9,
            spcm.SPCM_DDS_CORE10,
            spcm.SPCM_DDS_CORE11,
            spcm.SPCM_DDS_CORE20,
        )

        setup_dds_cores(
            dds,
            {
                channel_core_map[chan][i]: freq
                for chan, freqsstring in zip(chanstrings, freqstrings)
                for i, freq in enumerate(freqsstring)
            },
        )

        card.start(spcm.M2CMD_CARD_ENABLETRIGGER, spcm.M2CMD_CARD_FORCETRIGGER)

        try:
            while True:
                time.sleep(0.1)
        except KeyboardInterrupt:
            pass


def gridsweep(device, chanstrings, fmin, fmax, size):
    f0 = np.linspace(fmin[0], fmax[0], size)
    f1 = np.linspace(fmin[1], fmax[1], size)

    with spcm.Card(serial_number=device) as card:
        channels = spcm.Channels(card)
        setup_channels(channels)
        setup_dds(card)

        dds = spcm.DDSCommandQueue(card, channels=channels)
        dds.reset()
        dds.cores_on_channel(
            1,
            spcm.SPCM_DDS_CORE8,
            spcm.SPCM_DDS_CORE9,
            spcm.SPCM_DDS_CORE10,
            spcm.SPCM_DDS_CORE11,
            spcm.SPCM_DDS_CORE20,
        )
        dds.data_transfer_mode(spcm.SPCM_DDS_DTM_DMA)
        card.start(spcm.M2CMD_CARD_ENABLETRIGGER)

        dds.trg_src(spcm.SPCM_DDS_TRG_SRC_TIMER)
        dds.trg_timer(0.01)
        setup_dds_cores(
            dds,
            {channel_core_map[0][0]: f0.mean(), channel_core_map[1][0]: f1.mean()},
        )

        # required to not produce an underrun
        dds.trg_timer(0.2)
        for i in range(10):
            setup_dds_cores(
                dds,
                {channel_core_map[0][0]: f0.mean(), channel_core_map[1][0]: f1.mean()},
            )
        dds.exec_at_trg()
        dds.write_to_card()
        dds.mode = dds.WRITE_MODE.WAIT_IF_FULL

        card.cmd(spcm.M2CMD_CARD_FORCETRIGGER)

        try:
            while True:
                for i in range(size):
                    dds.freq(0, f0[::-1][i])
                    for j in range(size):
                        dds.freq(20, f1[:: (-1) ** (1 + i)][j])
                        dds.exec_at_trg()

                if dds.status() & spcm.SPCM_DDS_STAT_QUEUE_UNDERRUN:
                    break

            print("buffer underrun")
        except KeyboardInterrupt:
            pass


def main():
    parser = argparse.ArgumentParser()

    subparsers = parser.add_subparsers(dest="command")
    subparsers.add_parser("list")

    parser_static = subparsers.add_parser("static")
    parser_static.add_argument("--device", type=int, action="append", required=True)
    parser_static.add_argument("--channel", type=int, action="append", required=True)
    parser_static.add_argument("--frequency", type=str, action="append", required=True)

    parser_gridsweep = subparsers.add_parser("gridsweep")
    parser_gridsweep.add_argument("--size", type=int, default=3)
    parser_gridsweep.add_argument("--device", type=int, required=True)
    parser_gridsweep.add_argument("--channel", type=int, action="append", required=True)
    parser_gridsweep.add_argument("--fmin", type=float, action="append", required=True)
    parser_gridsweep.add_argument("--fmax", type=float, action="append", required=True)

    args = parser.parse_args()

    if args.command == "list":
        list()
        return
    if args.command == "static":
        if len(args.device) != len(args.channel) or len(args.device) != len(
            args.frequency
        ):
            print(
                "each device, channel, and frequency must have the same number of arguments"
            )

        logging.debug("static arguments parsed: ", args)
        static(args.device, args.channel, args.frequency)
        return
    if args.command == "gridsweep":
        logging.debug("gridsweep arguments parsed: ", args)
        gridsweep(args.device, args.channel, args.fmin, args.fmax, args.size)
        return

    parser.print_help()


if __name__ == "__main__":
    main()