fix: bound three-channel captures by the physical buffer

pslab/cli.py

@@ -83,7 +83,7 @@ def oscilloscope(
     ----------
     device : :class:`Handler`
         Serial interface for communicating with the PSLab device.
-    channels : {1, 2, 4}
+    channels : {1, 2, 3, 4}
         Number of channels to sample from simultaneously. By default, samples are
         captured from CH1, CH2, CH3 and MIC.
     duration : float
@@ -99,7 +99,8 @@ def oscilloscope(
         additional to the number of channels that were used to capture samples.
     """
     scope = Oscilloscope(device)
-    max_samples = CP.MAX_SAMPLES // channels
+    capture_channels = 4 if channels == 3 else channels
+    max_samples = CP.MAX_SAMPLES // capture_channels
     min_timegap = scope._lookup_mininum_timegap(channels)
     max_duration = max_samples * min_timegap * 1e-6
     active_channels = ([scope._channel_one_map] + scope._CH234)[:channels]

pslab/instrument/oscilloscope.py

@@ -60,8 +60,9 @@ def capture(
             from. If channel is an integer, the oscilloscope will sample the
             first one, two, three, or four channels in the aforementioned list.
         samples : int
-            Number of samples to fetch. Maximum 10000 divided by number of
-            channels.
+            Number of samples per channel. Maximum 10000 for one channel,
+            5000 for two channels, and 2500 for three or four channels.
+            Three-channel captures also sample MIC internally.
         timegap : float
             Time gap between samples in microseconds. Will be rounded to the
             closest 1 / 8 µs. The minimum time gap depends on the type of
@@ -114,9 +115,11 @@ def capture(
         ValueError
             If :channels: is not 1, 2, 3, 4, or one of CH1, CH2, CH3, MIC, CAP,
             RES, VOL, or
-            :samples: > 10000 / :channels:, or
+            :samples: exceeds the buffer capacity (2500 for three channels), or
             :timegap: is too low.
         """
+        self._check_capture_size(1 if isinstance(channels, str) else channels, samples)
+
         if isinstance(channels, str):
             self._channel_one_map = channels
             channels = 1
@@ -132,7 +135,7 @@ def capture(
             if trigger != self._trigger_voltage:
                 self.configure_trigger(voltage=trigger)
 
-        self._check_args(channels, samples, timegap)
+        self._check_args(channels, timegap)
         timegap = int(timegap * 8) / 8
 
         for channel in ("CH1", "CH2"):
@@ -155,16 +158,19 @@ def capture(
         else:
             return x
 
-    def _check_args(self, channels: int, samples: int, timegap: float):
+    def _check_capture_size(self, channels: int, samples: int):
         if channels not in (1, 2, 3, 4):
             raise ValueError("Number of channels to sample must be 1, 2, 3, or 4.")
 
-        max_samples = CP.MAX_SAMPLES // channels
+        # Three-channel captures use CAPTURE_FOUR and discard MIC.
+        capture_channels = 4 if channels == 3 else channels
+        max_samples = CP.MAX_SAMPLES // capture_channels
         if not 0 < samples <= max_samples:
             e1 = f"Cannot collect more than {max_samples} when sampling from "
             e2 = f"{channels} channels."
             raise ValueError(e1 + e2)
 
+    def _check_args(self, channels: int, timegap: float):
         min_timegap = self._lookup_mininum_timegap(channels)
         if timegap < min_timegap:
             raise ValueError(f"timegap must be at least {min_timegap}.")

tests/test_cli.py

@@ -18,7 +18,9 @@
 
 import pslab.protocol as CP
 from pslab import cli
+from pslab.connection import ConnectionHandler
 from pslab.instrument.analog import AnalogOutput
+from pslab.instrument.oscilloscope import Oscilloscope
 from pslab.instrument.waveform_generator import WaveformGenerator
 
 LA_CHANNELS = 4
@@ -90,6 +92,31 @@ def test_oscilloscope(scope):
         assert len(value) > SAMPLES
 
 
+@pytest.mark.parametrize(
+    "channels,capacity", [(1, 10000), (2, 5000), (3, 2500), (4, 2500)]
+)
+@pytest.mark.parametrize("duration", [0.001, 0.1])
+def test_oscilloscope_capture_capacity(mocker, channels, capacity, duration):
+    instrument = Oscilloscope(mocker.Mock(spec=ConnectionHandler))
+    timegap = instrument._lookup_mininum_timegap(channels)
+    samples = min(capacity, round(duration * 1e6 / timegap))
+    voltages = [np.full(samples, channel) for channel in range(4)]
+    mocker.patch("pslab.cli.Oscilloscope", return_value=instrument)
+    mocker.patch.object(instrument, "progress", return_value=(True, samples))
+    mocker.patch.object(instrument, "fetch_data", return_value=voltages)
+    mocker.patch("pslab.cli.time.time", side_effect=[0, duration + 1])
+    mocker.patch("pslab.cli.time.sleep")
+    capture = mocker.spy(instrument, "capture")
+
+    headers, values = cli.oscilloscope(mocker.Mock(), channels, duration)
+
+    capture.assert_called_once_with(channels, samples, timegap)
+    assert headers == ["Timestamp", "CH1", "CH2", "CH3", "MIC"][: channels + 1]
+    np.testing.assert_array_equal(values[0], timegap * np.arange(samples))
+    for channel in range(channels):
+        np.testing.assert_array_equal(values[channel + 1], voltages[channel])
+
+
 def test_collect_csv_stdout(collect, capsys):
     cli.cmdline(["collect", "logic_analyzer", "--channels", str(LA_CHANNELS)])
     output = list(csv.reader(capsys.readouterr().out.splitlines()))

tests/test_oscilloscope_capacity.py

@@ -0,0 +1,92 @@
+"""Test oscilloscope buffer limits without connected hardware."""
+
+from unittest.mock import Mock, call
+
+import numpy as np
+import pytest
+
+import pslab.protocol as CP
+from pslab.connection import ConnectionHandler
+from pslab.instrument.oscilloscope import Oscilloscope
+
+
+@pytest.fixture
+def scope():
+    """Return an oscilloscope using a mock connection."""
+    device = Mock(spec=ConnectionHandler)
+    instrument = Oscilloscope(device)
+    device.reset_mock()
+    return instrument
+
+
+@pytest.mark.parametrize(
+    "channels,capacity", [(1, 10000), (2, 5000), (3, 2500), (4, 2500)]
+)
+@pytest.mark.parametrize("offset", [1, 833])
+def test_capture_rejects_samples_beyond_buffer(scope, channels, capacity, offset):
+    with pytest.raises(ValueError, match=f"Cannot collect more than {capacity}"):
+        scope.capture(channels, capacity + offset, timegap=2, block=False)
+
+    assert scope._device.mock_calls == []
+
+
+@pytest.mark.parametrize(
+    "channels,capacity", [(1, 10000), (2, 5000), (3, 2500), (4, 2500)]
+)
+def test_capture_accepts_full_buffer(scope, channels, capacity):
+    timestamps = scope.capture(channels, capacity, timegap=2, block=False)[0]
+
+    np.testing.assert_array_equal(timestamps, 2 * np.arange(capacity))
+    assert call(capacity) in scope._device.send_int.call_args_list
+    for channel in scope._channels.values():
+        if channel.buffer_idx is not None:
+            assert channel.buffer_idx + channel.samples_in_buffer <= CP.MAX_SAMPLES
+
+
+def test_three_channel_capture_uses_four_channel_command(scope, monkeypatch):
+    voltages = [np.zeros(2500) for _ in range(4)]
+    monkeypatch.setattr(scope, "progress", lambda: (True, 2500))
+    monkeypatch.setattr(scope, "fetch_data", lambda: voltages)
+    monkeypatch.setattr("pslab.instrument.oscilloscope.time.sleep", lambda _: None)
+
+    result = scope.capture(3, 2500, timegap=2)
+
+    assert len(result) == 4
+    assert all(result[i + 1] is voltages[i] for i in range(3))
+    scope._device.send_byte.assert_any_call(CP.CAPTURE_FOUR)
+    assert scope._channels["MIC"].buffer_idx == 7500
+    assert scope._channels["MIC"].samples_in_buffer == 2500
+
+
+@pytest.mark.parametrize("trigger", [None, False, 0.5])
+def test_oversized_capture_preserves_trigger_configuration(scope, trigger):
+    scope.configure_trigger(channel="CH2", voltage=1.1)
+    scope._device.reset_mock()
+    previous = (scope.trigger_enabled, scope.trigger_channel, scope.trigger_voltage)
+
+    with pytest.raises(ValueError, match="Cannot collect more than 2500"):
+        scope.capture(3, 2501, timegap=2, trigger=trigger, trigger_channel="CH3")
+
+    assert (
+        scope.trigger_enabled,
+        scope.trigger_channel,
+        scope.trigger_voltage,
+    ) == previous
+    assert scope._device.mock_calls == []
+
+
+def test_oversized_named_channel_capture_preserves_channel_map(scope):
+    with pytest.raises(ValueError, match="Cannot collect more than 10000"):
+        scope.capture("VOL", 10001, timegap=2, block=False)
+
+    assert scope._channel_one_map == "CH1"
+    assert scope._device.mock_calls == []
+
+
+def test_valid_capture_can_disable_trigger_for_shorter_timegap(scope):
+    scope.configure_trigger(voltage=1.1)
+
+    timestamps = scope.capture(1, 2, timegap=0.5, trigger=False, block=False)[0]
+
+    assert not scope.trigger_enabled
+    np.testing.assert_array_equal(timestamps, [0, 0.5])