diff --git a/pslab/instrument/waveform_generator.py b/pslab/instrument/waveform_generator.py index 07c3523..0836206 100644 --- a/pslab/instrument/waveform_generator.py +++ b/pslab/instrument/waveform_generator.py @@ -253,6 +253,11 @@ def load_function( The minimum and maximum x values between which to evaluate 'function'. Should typically correspond to one period. Omit if 'function' is 'sine' or 'tria'. + + Raises + ------ + ValueError + If the function does not return 512 finite real voltage values. """ if function == "sine": @@ -261,7 +266,7 @@ def sine(x): function = sine span = [0, 2 * np.pi] - self._channels[channel].wavetype = "sine" + mode = "sine" elif function == "tria": def tria(x): @@ -269,15 +274,13 @@ def tria(x): function = tria span = [0, 4] - self._channels[channel].wavetype = "tria" + mode = "tria" else: - self._channels[channel].wavetype = "custom" + mode = "custom" x = np.arange(span[0], span[1], (span[1] - span[0]) / 512) y = function(x) - self._load_table( - channel=channel, points=y, mode=self._channels[channel].wavetype - ) + self._load_table(channel=channel, points=y, mode=mode) def load_table(self, channel: str, points: np.ndarray): """Load a custom waveform as a table. @@ -289,13 +292,38 @@ def load_table(self, channel: str, points: np.ndarray): points : np.ndarray Array of voltage values which make up the waveform. Array length must be 512. Values outside the range -3.3 V to 3.3 V will be - clipped. + clipped. All values must be finite real numbers. Complex values and + booleans are rejected. + + Raises + ------ + ValueError + If the table is not a one-dimensional array of 512 finite real + voltage values. Invalid tables are rejected before device writes. """ self._load_table(channel, points, "custom") def _load_table(self, channel, points, mode="custom"): - self._channels[channel].wavetype = mode + points = np.asarray(points) + if points.ndim != 1 or points.size != self._HIGHRES_TABLE_SIZE: + raise ValueError("Waveform table must contain exactly 512 values.") + if points.dtype.kind not in "iufOUS" or ( + points.dtype.kind == "O" + and any(isinstance(value, (bool, np.bool_)) for value in points) + ): + raise ValueError("Waveform voltages must be finite real values.") + if points.dtype.kind in "OUS": + try: + points = points.astype(float) + except (TypeError, ValueError, OverflowError) as error: + raise ValueError( + "Waveform voltages must be finite real values." + ) from error + if not np.isfinite(points).all(): + raise ValueError("Waveform voltages must be finite real values.") + self._channels[channel].waveform_table = points + self._channels[channel].wavetype = mode logger.info(f"Reloaded waveform table for {channel}: {mode}.") self._device.send_byte(CP.WAVEGEN) diff --git a/tests/test_waveform_table_validation.py b/tests/test_waveform_table_validation.py new file mode 100644 index 0000000..c378ca1 --- /dev/null +++ b/tests/test_waveform_table_validation.py @@ -0,0 +1,82 @@ +"""Hardware-independent validation of waveform uploads and command framing.""" + +from unittest.mock import Mock + +import numpy as np +import pytest + +import pslab.protocol as CP +from pslab.connection import ConnectionHandler +from pslab.instrument.waveform_generator import WaveformGenerator + + +@pytest.mark.parametrize("loader", ["table", "function"]) +@pytest.mark.parametrize( + "points", + [ + np.zeros(0), + np.zeros(1), + np.zeros(511), + np.zeros(513), + np.zeros((512, 1)), + np.zeros((2, 256)), + np.full(512, np.nan), + np.full(512, np.inf), + np.full(512, -np.inf), + np.full(512, 1j), + np.full(512, "not-a-voltage"), + np.full(512, None, dtype=object), + np.full(512, {}, dtype=object), + np.resize(np.array([False, True]), 512), + np.full(512, True, dtype=object), + np.array([np.bool_(False)] + [0.5] * 511, dtype=object), + ], +) +def test_invalid_table_does_not_write_or_change_channel(points, loader): + device = Mock(spec=ConnectionHandler) + generator = WaveformGenerator(device) + channel = generator._channels["SI1"] + previous_table = channel.waveform_table.copy() + with pytest.raises(ValueError): + if loader == "table": + generator.load_table("SI1", points) + else: + generator.load_function("SI1", lambda x: points, [0, 1]) + assert device.mock_calls == [] + assert channel.wavetype == "sine" + assert channel.waveform_table == previous_table + + +@pytest.mark.parametrize("channel", ["SI1", "SI2"]) +@pytest.mark.parametrize("dtype", [np.float64, np.int16, object]) +def test_valid_table_preserves_upload_framing(channel, dtype): + device = Mock(spec=ConnectionHandler) + generator = WaveformGenerator(device) + points = np.linspace(-5, 5, 512).astype(dtype) + original = points.copy() + generator.load_table(channel, points) + clipped = np.clip(np.asarray(points, dtype=float), -3.3, 3.3) + expected = np.round((clipped + 3.3) / 6.6 * 511).astype(np.int16).tolist() + assert [call.args[0] for call in device.send_int.call_args_list] == expected + commands = [call.args[0] for call in device.send_byte.call_args_list] + assert commands[:2] == [ + CP.WAVEGEN, + CP.LOAD_WAVEFORM1 if channel == "SI1" else CP.LOAD_WAVEFORM2, + ] + assert ( + commands[2:] + == np.round((clipped[::16] + 3.3) / 6.6 * 63).astype(np.int16).tolist() + ) + device.get_ack.assert_called_once_with() + assert generator._channels[channel].wavetype == "custom" + np.testing.assert_array_equal(points, original) + + +@pytest.mark.parametrize("function", ["sine", "tria"]) +def test_builtin_functions_keep_fixed_upload_sizes(function): + device = Mock(spec=ConnectionHandler) + generator = WaveformGenerator(device) + generator.load_function("SI1", function) + assert device.send_int.call_count == 512 + assert device.send_byte.call_count == 34 + assert generator._channels["SI1"].wavetype == function