Add DWTA functions

Signed-off-by: Travis F. Collins <[email protected]>

bench/cli/iiotools.py

@@ -1,4 +1,5 @@
 import click
+from bench.keysight.dwta.data_capture import capture_iq_datafile
 
 try:
     import iio
@@ -9,15 +10,11 @@
 
 @click.group()
 @click.option("--uri", "-u", help="URI of target device/board")
-@click.option("--device", "-d", help="Device driver to use")
-@click.option("--complex", "-x", is_flag=True, help="Use complex mode")
 @click.pass_context
-def cli(ctx, uri, device, complex):
+def cli(ctx, uri):
     """Command line interface for pybench IIO based boards"""
     ctx.obj = {}
     ctx.obj["uri"] = uri
-    ctx.obj["device"] = device
-    ctx.obj["complex"] = complex
 
 
 @cli.command()
@@ -27,17 +24,19 @@ def cli(ctx, uri, device, complex):
 @click.option(
     "--amplitude", "-a", help="Set the amplitude of the DDS in 0->1", required=True
 )
+@click.option("--device", "-d", help="Device driver to use")
 @click.option("--channel", "-c", help="Set the channel of the DDS", required=True)
+@click.option("--complex", "-x", is_flag=True, help="Use complex mode")
 @click.pass_context
-def set_dds(ctx, frequency, amplitude, channel):
+def set_dds(ctx, frequency, amplitude, device, channel, complex):
     """Configure DDS"""
     iioctx = iio.Context(ctx.obj["uri"])
     if not iioctx:
         click.echo("No context")
         return
-    dev = iioctx.find_device(ctx.obj["device"])
+    dev = iioctx.find_device(device)
     if not dev:
-        click.echo(f"Device {ctx.obj['device']} not found")
+        click.echo(f"Device {device} not found")
         return
     dds_channels = [ch.id for ch in dev.channels if "altvoltage" in ch.id]
     # Set all the DDS scales to 0
@@ -49,7 +48,7 @@ def set_dds(ctx, frequency, amplitude, channel):
             return
         chan.attrs["scale"].value = "0"
     # Set the desired DDS scale
-    if ctx.obj["complex"]:
+    if complex:
         # Channels are groups of 4
         ch = int(channel) * 4
         channels = [ch, ch + 1]
@@ -64,11 +63,39 @@ def set_dds(ctx, frequency, amplitude, channel):
             return
         chan.attrs["frequency"].value = frequency
         chan.attrs["scale"].value = amplitude
-        if ctx.obj["complex"]:
+        if complex:
             # i is odd
             if i % 2:
                 chan.attrs["phase"].value = "90000"
             else:
                 chan.attrs["phase"].value = "0"
 
     click.echo(f"Set DDS of channel {channel} to {frequency}Hz and {amplitude} scale")
+
+
+@cli.command()
+@click.option("--filename", "-f", help="Name of file to write data to", required=True)
+@click.option("--device", "-d", help="Name of device to configure", required=True)
+@click.option("--samples", "-s", help="Number of samples to capture", required=True)
+@click.argument("props", nargs=-1)
+@click.pass_context
+def capture_Data(ctx, filename, device, samples, props):
+
+    # Checks
+    samples = int(samples)
+
+    # Parse properties
+    if props:
+        oprops = {}
+        for prop in props:
+            if "=" not in prop:
+                raise ValueError(
+                    f"Invalid property: {prop}. Must be in the form key=value"
+                )
+            k, v = prop.split("=")
+            if v.isdigit():
+                v = int(v)
+            oprops[k] = v
+        props = oprops
+
+    capture_iq_datafile(filename, device, samples, ctx.obj["uri"], **dict(props))

bench/keysight/__init__.py

@@ -1,2 +1,3 @@
 from .e36233a import E36233A
 from .n9040b import N9040B
+from .dwta import utils

bench/keysight/dwta/__init__.py

@@ -0,0 +1,3 @@
+from .utils import data_to_iq_datafile
+
+supported_devices = ["Pluto", "ad9081", "ad9084"]

bench/keysight/dwta/data_capture.py

@@ -0,0 +1,45 @@
+"""Keysight DWTA data capture module."""
+
+import adi
+
+from .utils import data_to_iq_datafile
+from ..dwta import supported_devices
+
+
+def capture_iq_datafile(
+    filename: str, device_name: str, samples: int, uri: str, **kwargs
+) -> None:
+    """Capture IQ data to a file.
+
+    Args:
+        filename: The filename to write to.
+        device_name: The device name.
+        samples: The number of samples to capture.
+        uri: The URI of the device.
+        *args: Additional arguments.
+        **kwargs: Additional keyword arguments.
+    """
+    # Checks
+    assert (
+        device_name in supported_devices
+    ), f"Device not supported: {device_name}.\n Supported devices: {supported_devices}"
+    assert samples > 0, "Number of samples must be greater than 0"
+    if "ip:" not in uri and "usb:" not in uri:
+        raise ValueError(f"URI must start with 'ip:' or 'usb:'. Got: {uri}")
+
+    # Create device
+    device = getattr(adi, device_name)(uri)
+
+    # Assume kwargs are attributes to set
+    for key, value in kwargs.items():
+        print(f"Setting {key} to {value}")
+        if not hasattr(device, key):
+            raise AttributeError(f"Device does not have attribute: {key}")
+        setattr(device, key, value)
+
+    # Capture data
+    device.rx_buffer_size = samples
+    data = device.rx()
+
+    # Write data to file
+    data_to_iq_datafile(device, data, filename, device_check=False)

bench/keysight/dwta/utils.py

@@ -0,0 +1,50 @@
+""""Utilities for the Keysight DWTA."""
+
+from typing import List, Union
+import numpy as np
+
+import adi
+
+
+def data_to_iq_datafile(
+    device, data: np.ndarray, filename: str, device_check: bool = True
+) -> None:
+    """Write data to a file in IQ data format.
+
+    Args:
+        device: Device object from pyadi-iio library.
+        data: The data to write.
+        filename: The filename to write to.
+        device_check: Check if specific class is supported. Default is True.
+    """
+    # Checks
+    assert data.dtype == np.complex128, "Data must be complex"
+    assert device._complex_data, "Device does not support complex data"
+    if device_check:
+        supported_devices = ["Pluto", "ad9081"]
+        assert any(
+            dev in device.__class__.__name__ for dev in supported_devices
+        ), f"Device not supported. Supported devices: {supported_devices}"
+
+    # Get sample rate and carrier frequencies
+    if hasattr(device, "sample_rate"):
+        sample_rate = device.sample_rate
+    elif hasattr(device, "rx_sample_rate"):
+        sample_rate = device.rx_sample_rate
+    else:
+        raise AttributeError("Device does not have sample rate attribute")
+
+    if type(device) in [adi.Pluto]:
+        center_frequency = int(device.rx_lo)  # FIXME
+    else:
+        raise NotImplementedError("Center frequency not implemented for this device")
+
+    assert device._complex_data, "Device does not support complex data"
+
+    with open(filename, "w") as f:
+        f.write(f"SampleRate={sample_rate}\n")
+        f.write(f"CenterFrequency={center_frequency}\n")
+        for d in data:
+            f.write(f"{d.real},{d.imag}\n")
+
+    print(f"Data written to {filename}")

pyproject.toml

@@ -50,7 +50,8 @@ pybenchiio = "bench.cli.iiotools:cli"
 pybench = "bench.cli.visatools:cli"
 
 [project.optional-dependencies]
-cli = ["click", "pylibiio>=0.23.1"]
+cli = ["click", "pyadi-iio>=0.0.17"]
+web = ["fastapi[standard]", "uvicorn", "pydantic", "jinja2", "starlette"]
 
 [project.entry-points.pytest11]
 pybench = "bench.plugin"
@@ -72,7 +73,7 @@ ignore_missing_imports="true"
 
 [tool.pytest.ini_options]
 minversion = "6.0"
-# addopts = "--custom-hw-map=test/emu/hardware_map.yml --emu-xml-dir=test/emu/devices"
+addopts = "--adi-hw-map"
 testpaths = [
     "tests",
 ]

requirements_dev.txt

@@ -8,7 +8,7 @@ scapy
 scipy
 pytest-cov
 coveralls
-# pytest-libiio>=0.0.18
+pytest-libiio>=0.0.18
 bump2version
 pytest-html==3.2.0
 plotly-express