Amplitude Look-Up Table (LUT) Support for Waveform Processing

.gitignore

@@ -86,6 +86,9 @@ ipython_config.py
 #   For a library or package, you might want to ignore these files since the code is
 #   intended to run in multiple environments; otherwise, check them in:
 # .python-version
+.conda
+bootstrap_requirements.txt
+environment.yml
 
 # pipenv
 #   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.

docs/examples/index.rst

@@ -9,4 +9,5 @@ Broadbean Examples
     Making_output_for_Tektronix_AWG70000A.ipynb
     Example_Write_Read_JSON.ipynb
     Filter_compensation.ipynb
+    Amplitude_LUT_Example.ipynb
     Subsequences.ipynb

src/broadbean/ripasso.py

@@ -6,9 +6,11 @@
 #
 
 import logging
+from typing import TypedDict
 
 import numpy as np
 from numpy.fft import fft, fftfreq, ifft
+from numpy.typing import ArrayLike
 
 log = logging.getLogger(__name__)
 
@@ -17,7 +19,7 @@ class MissingFrequenciesError(Exception):
     pass
 
 
-def _rcFilter(SR, npts, f_cut, kind="HP", order=1, DCgain=0):
+def _rcFilter(SR: int, npts: int, f_cut: float, kind="HP", order=1, DCgain=0):
     """
     Nth order (RC circuit) filter
     made with frequencies matching the fft output
@@ -44,7 +46,9 @@ def _rcFilter(SR, npts, f_cut, kind="HP", order=1, DCgain=0):
     return tf**order
 
 
-def applyRCFilter(signal, SR, kind, f_cut, order, DCgain=0):
+def applyRCFilter(
+    signal: np.ndarray, SR: int, kind: str, f_cut: float, order: int, DCgain: float = 0
+):
     """
     Apply a simple RC-circuit filter
     to signal and return the filtered signal.
@@ -80,7 +84,9 @@ def applyRCFilter(signal, SR, kind, f_cut, order, DCgain=0):
     return output
 
 
-def applyInverseRCFilter(signal, SR, kind, f_cut, order, DCgain=1):
+def applyInverseRCFilter(
+    signal: np.ndarray, SR: int, kind: str, f_cut: float, order: int, DCgain: float = 1
+):
     """
     Apply the inverse of an RC-circuit filter to a signal and return the
     compensated signal.
@@ -125,7 +131,13 @@ def applyInverseRCFilter(signal, SR, kind, f_cut, order, DCgain=1):
     return output
 
 
-def applyCustomTransferFunction(signal, SR, tf_freqs, tf_amp, invert=False):
+def applyCustomTransferFunction(
+    signal: np.ndarray,
+    SR: int,
+    tf_freqs: np.ndarray,
+    tf_amp: np.ndarray,
+    invert: bool = False,
+):
     """
     Apply custom transfer function
 
@@ -198,3 +210,26 @@ def applyCustomTransferFunction(signal, SR, tf_freqs, tf_amp, invert=False):
     signal_filtered = np.real(signal_filtered)
 
     return signal_filtered
+
+
+class AmplitudeLUT(TypedDict):
+    LUT_input: ArrayLike
+    LUT_output: ArrayLike
+
+
+def applyAmplitudeLUT(signal: np.ndarray, lut: AmplitudeLUT):
+    """
+    Apply an amplitude LUT to the signal.
+
+    Args:
+        signal (np.array): The input signal.
+        lut (AmplitudeLUT): A lookup table for amplitude compensation. The LUT_input
+            field should contain the input values and
+            the LUT_output field should contain the corresponding output values.
+
+    Returns:
+        np.array:
+            The signal after applying the amplitude LUT.
+    """
+
+    return np.interp(signal, lut["LUT_input"], lut["LUT_output"])

src/broadbean/sequence.py

@@ -11,7 +11,7 @@
 
 from broadbean.blueprint import BluePrint
 from broadbean.element import Element  # TODO: change import to element.py
-from broadbean.ripasso import applyInverseRCFilter
+from broadbean.ripasso import applyAmplitudeLUT, applyInverseRCFilter
 
 from .broadbean import (
     PulseAtoms,
@@ -385,6 +385,30 @@ def setChannelFilterCompensation(
             "tau": tau,
         }
 
+    def setAmplitudeLUT(
+        self, channel: int | str, lut_input: list[float], lut_output: list[float]
+    ) -> None:
+        """
+        Set an amplitude lookup table for a channel. This is used when making
+        output for .awg files. The LUT is applied after all other waveform
+        processing.
+
+        Args:
+            channel: The channel number/name
+            lut_input: The input levels for the LUT
+            lut_output: The output levels for the LUT
+
+        Example:
+            >>> lut_input = [-1.0, -0.5, 0.0, 0.5, 1.0]
+            >>> lut_output = [-0.8, -0.3, 0.0, 0.3, 0.8]
+            >>> seq.setAmplitudeLUT(1, lut_input, lut_output)
+        """
+
+        self._awgspecs[f"channel{channel}_amplitude_LUT"] = {
+            "LUT_input": lut_input,
+            "LUT_output": lut_output,
+        }
+
     def addElement(self, position: int, element: Element) -> None:
         """
         Add an element to the sequence. Overwrites previous values.
@@ -841,6 +865,19 @@ def forge(
                                 output[pos1]["content"][pos2]["data"][channame]["wfm"]
                             ) = postfilter
 
+                        # Apply amplitude LUT if present
+                        lut_key = f"channel{channame}_amplitude_LUT"
+                        if lut_key in self._awgspecs.keys():
+                            lut = self._awgspecs[lut_key]
+                            output[pos1]["content"][pos2]["data"][channame]["wfm"] = (
+                                applyAmplitudeLUT(
+                                    output[pos1]["content"][pos2]["data"][channame][
+                                        "wfm"
+                                    ],
+                                    lut,
+                                )
+                            )
+
         return output
 
     def _prepareForOutputting(self) -> list[dict[int, Any]]:
@@ -954,6 +991,15 @@ def _prepareForOutputting(self) -> list[dict[int, Any]]:
                     )
                     elements[pos][chan]["wfm"] = postfilter
 
+            # Apply amplitude LUT if present
+            lut_key = f"channel{chan}_amplitude_LUT"
+            if lut_key in self._awgspecs.keys():
+                lut = self._awgspecs[lut_key]
+                for pos in range(seqlen):
+                    elements[pos][chan]["wfm"] = applyAmplitudeLUT(
+                        elements[pos][chan]["wfm"], lut
+                    )
+
         return elements
 
     def outputForSEQXFile(