Added frequency response measurement initial steps

9 files changed, 80 insertions, 15 deletions

diff --git a/lab_control/frequency_response.py b/lab_control/frequency_response.py
index 6508d80..9f6c79e 100644
--- a/lab_control/frequency_response.py
+++ b/lab_control/frequency_response.py
@@ -1,14 +1,17 @@
 from pathlib import Path
+import time
 
 from lab_control.function_generator import FunctionGenerator
 from lab_control.oscilloscope import Oscilloscope
 from lab_control.measurement import Measurement, getLinearRange
 
+
 class FrequencyResponseMeasurement(Measurement):
     def __init__(self):
         self.minFrequency = 20e0
         self.maxFrequency = 16e3
         self.steps = 50
+        self.inputAmplitude = 0.1
         self.functionGeneratorChannel = 1
         self.oscilloscopeChannel = 1
         self.measurementDone = False
@@ -18,8 +21,15 @@ class FrequencyResponseMeasurement(Measurement):
         frequencies = getLinearRange(self.minFrequency, self.maxFrequency, self.steps)
         self.data = []
 
+        # Initial set up
+        fg.setAmplitude(self.functionGeneratorChannel, self.inputAmplitude)
+        fg.setFunction(self.functionGeneratorChannel, FunctionGenerator.SINE)
+        fg.setOn(self.functionGeneratorChannel)
+        osc.setVoltsPerDivision(self.oscilloscopeChannel, self.inputAmplitude / 6)
+
         for f in frequencies:
             fg.setFrequency(self.functionGeneratorChannel, f)
+            #time.sleep(0.2)
             response = osc.measureAmplitude(self.oscilloscopeChannel)
             self.data.append((f, response))
 
diff --git a/lab_control/function_generator.py b/lab_control/function_generator.py
index 20f1b8c..49d04e3 100644
--- a/lab_control/function_generator.py
+++ b/lab_control/function_generator.py
@@ -5,6 +5,11 @@ class FunctionGenerator:
     This interface specifies the common API for all
     supported function generator devices.
     """
+
+    SINE       = 0
+    SQUARE     = 1
+    TRIANGULAR = 3
+
     def setOn(self, channel: int) -> None:
         """ Enable channel. """
 
diff --git a/lab_control/jds6600.py b/lab_control/jds6600.py
index 4ab4955..7af522b 100644
--- a/lab_control/jds6600.py
+++ b/lab_control/jds6600.py
@@ -22,9 +22,6 @@ class JDS6600(FunctionGenerator):
     and offer an API to control the device.
     """
     AVAILABLE_CHANNELS = [1, 2]
-    SINE       = 0
-    SQUARE     = 1
-    TRIANGULAR = 3
 
     def __init__(self, portName):
         super().__init__()
diff --git a/lab_control/oscilloscope.py b/lab_control/oscilloscope.py
index 167879c..04bd338 100644
--- a/lab_control/oscilloscope.py
+++ b/lab_control/oscilloscope.py
@@ -16,3 +16,7 @@ class Oscilloscope:
 
     def measureFrequency(self, channel: int) -> float:
         """ Return frequency measurement on specific channel. """
+
+    def setVoltsPerDivision(self, channel: int, volts: float) -> None:
+        """ Sets the vertical scale of specific channel. """
+
diff --git a/lab_control/sds1000xe.py b/lab_control/sds1000xe.py
index 67fc29b..7f6d552 100644
--- a/lab_control/sds1000xe.py
+++ b/lab_control/sds1000xe.py
@@ -4,6 +4,7 @@ Implements partial support for Siglent SDS1000X-E series oscilloscopes.
 
 import socket
 import re
+
 from lab_control.oscilloscope import Oscilloscope
 
 def _checkChannel(channel):
@@ -15,7 +16,7 @@ class SDS1000XE(Oscilloscope):
     and port and offer an API to control the device.
     """
     PORT = 5025
-    TIMEOUT = 1.0
+    TIMEOUT = 5.0
     AVAILABLE_CHANNELS = range(1, 5)
 
     def __init__(self, address):
@@ -24,7 +25,6 @@ class SDS1000XE(Oscilloscope):
         self._socket.connect((address, SDS1000XE.PORT))
         self._socket.settimeout(SDS1000XE.TIMEOUT)
 
-
     def _measure(self, channel: int, code: str) -> float:
         _checkChannel(channel)
         pattern = r"C(?P<responseChannel>\d):PAVA .+,(?P<rawMeasurement>[\d.E+-]+)\w+"
@@ -52,3 +52,9 @@ class SDS1000XE(Oscilloscope):
 
     def measureFrequency(self, channel: int) -> float:
         return self._measure(channel, "FREQ")
+
+    def setVoltsPerDivision(self, channel: int, volts: float) -> None:
+        _checkChannel(channel)
+        query = f"C{channel}:VDIV {volts:.2E}V\r\n"
+        self._socket.sendall(query.encode())
+        # no response expected
diff --git a/lab_control/test/frequency_response_test.py b/lab_control/test/frequency_response_test.py
index 2997413..61fd021 100644
--- a/lab_control/test/frequency_response_test.py
+++ b/lab_control/test/frequency_response_test.py
@@ -3,6 +3,7 @@ import pytest
 from lab_control.test.mock_lab import MockLab
 from lab_control.frequency_response import FrequencyResponseMeasurement
 from lab_control.measurement import getLinearRange
+from lab_control.function_generator import FunctionGenerator
 
 @pytest.fixture
 def mockLab():
@@ -14,7 +15,6 @@ def uut(mockLab):
 
 def prepareRampResponse(uut, mockLab):
     # Expect a ramp response from 0.5 to 1.5 * input amplitude
-    inputAmplitude = 2.0
     minScale = 0.5
     maxScale = 1.5
 
@@ -25,15 +25,14 @@ def prepareRampResponse(uut, mockLab):
 
     mockLab.connectChannels(uut.functionGeneratorChannel, uut.oscilloscopeChannel)
     mockLab.setTestFunction(uut.oscilloscopeChannel, testFunction)
-    mockLab.setAmplitude(uut.functionGeneratorChannel, inputAmplitude)
-    mockLab.setOn(uut.functionGeneratorChannel)
 
-    return [(f, testFunction(f) * inputAmplitude) for f in getLinearRange(uut.minFrequency, uut.maxFrequency, uut.steps)]
+    return [(f, testFunction(f) * uut.inputAmplitude) for f in getLinearRange(uut.minFrequency, uut.maxFrequency, uut.steps)]
 
 def test_frequencyResponseDefaults(uut):
     assert uut.minFrequency == 20e0
     assert uut.maxFrequency == 16e3
     assert uut.steps == 50
+    assert uut.inputAmplitude == 0.1
     assert uut.functionGeneratorChannel == 1
     assert uut.oscilloscopeChannel == 1
 
@@ -43,14 +42,26 @@ def test_frequencyResponseRamp(uut, mockLab):
     uut.steps = 11
     uut.functionGeneratorChannel = 1
     uut.oscilloscopeChannel = 1
+    uut.inputAmplitude = 2.0
+
+    fg = mockLab.getFunctionGeneratorChannel(uut.functionGeneratorChannel)
+    osc = mockLab.getOscilloscopeChannel(uut.oscilloscopeChannel)
 
     expectedData = prepareRampResponse(uut, mockLab)
 
     assert not uut.measurementDone
-    assert uut.data == None
+    assert uut.data is None
+    assert fg.amplitude is None
+    assert fg.function is None
+    assert not fg.on
+    assert osc.voltsPerDiv is None
     
     uut.measure(mockLab, mockLab)
 
+    assert fg.amplitude == uut.inputAmplitude
+    assert fg.function == FunctionGenerator.SINE
+    assert fg.on
+    assert osc.voltsPerDiv == uut.inputAmplitude / 6
     assert uut.measurementDone
     assert uut.data == expectedData
 
diff --git a/lab_control/test/mock_lab.py b/lab_control/test/mock_lab.py
index 40893ca..75ba1bf 100644
--- a/lab_control/test/mock_lab.py
+++ b/lab_control/test/mock_lab.py
@@ -9,11 +9,13 @@ class MockLab(FunctionGenerator, Oscilloscope):
             self.on = False
             self.frequency = None
             self.amplitude = None
+            self.function = None
 
     class OscChannelState:
         def __init__(self):
             self.testFunction = None
             self.connectedChannel = None
+            self.voltsPerDiv = None
 
     def __init__(self):
         self.fgChannels = [MockLab.FGChannelState() for i in range(0, 2)]
@@ -32,7 +34,7 @@ class MockLab(FunctionGenerator, Oscilloscope):
         self.fgChannels[channel - 1].amplitude = amplitude
 
     def setFunction(self, channel: int, function: int) -> None:
-        pass
+        self.fgChannels[channel - 1].function = function
 
     def measureAmplitude(self, channel: int) -> float:
         fgChannel = self.oscChannels[channel - 1].connectedChannel
@@ -50,8 +52,17 @@ class MockLab(FunctionGenerator, Oscilloscope):
     def measureFrequency(self, channel: int) -> float:
         pass
 
+    def setVoltsPerDivision(self, channel: int, volts: float) -> None:
+        self.oscChannels[channel - 1].voltsPerDiv = volts
+
     def setTestFunction(self, channel: int, f: Callable[[float], float]) -> None:
         self.oscChannels[channel - 1].testFunction = f
 
     def connectChannels(self, fg: int, osc: int) -> None:
         self.oscChannels[osc - 1].connectedChannel = self.fgChannels[fg - 1]
+
+    def getFunctionGeneratorChannel(self, channel: int):
+        return self.fgChannels[channel - 1]
+
+    def getOscilloscopeChannel(self, channel: int):
+        return self.oscChannels[channel - 1]
diff --git a/lab_control/test/mock_sds1000xe_device.py b/lab_control/test/mock_sds1000xe_device.py
index b4dd222..68c2471 100644
--- a/lab_control/test/mock_sds1000xe_device.py
+++ b/lab_control/test/mock_sds1000xe_device.py
@@ -24,8 +24,8 @@ class MockSDS1000XEDevice:
         self._mainThread = threading.Thread(target=self._mainLoop)
         self._mainThread.start()
 
-        # Mock measured values
-        self._channels = [{"AMPL" : None} for i in range(0, 4)]
+        # Mock internal values
+        self._channels = [{"AMPL" : None, "VDIV" : None} for i in range(0, 4)]
 
     def _mainLoop(self) -> None:
         self._clientSocket, _ = _serverSocket.accept()
@@ -35,7 +35,7 @@ class MockSDS1000XEDevice:
             while not self._stopFlag:
                 try:
                     request = self._clientSocket.recv(4096).decode()
-                    response = self._handleRequest(request)
+                    response = self._handleRequest(request.strip())
                     if response is not None:
                         self._clientSocket.send(response.encode())
                 except TimeoutError as e:
@@ -44,7 +44,7 @@ class MockSDS1000XEDevice:
             self._clientSocket.close()
     
     def _handleRequest(self, request: str) -> str:
-        m = re.search(r"C(?P<channel>\d):(?P<opcode>\w+)\?\s(?P<arg>\w+)", request)
+        m = re.search(r"C(?P<channel>\d):(?P<opcode>\w+)\??\s(?P<arg>.+)", request)
         if not m:
             return None
         
@@ -61,6 +61,10 @@ class MockSDS1000XEDevice:
             else:
                 response = f"C{m.group('channel')}:PAVA {arg},{value:.6E}{unit}"
                 return response
+        elif opcode == "VDIV":
+            arg = float(m.group("arg").rstrip("V"))
+            self._channels[channelIndex]["VDIV"] = arg
+            return None
 
     def stop(self) -> None:
         self._stopFlag = True
@@ -78,3 +82,6 @@ class MockSDS1000XEDevice:
     def setFrequency(self, channel: int, value: float) -> None:
         self._channels[channel - 1]["FREQ"] = value
 
+    def getVoltsPerDivision(self, channel: int) -> float:
+        return self._channels[channel - 1]["VDIV"]
+
diff --git a/lab_control/test/sds1000xe_test.py b/lab_control/test/sds1000xe_test.py
index 6f3de80..3774d52 100644
--- a/lab_control/test/sds1000xe_test.py
+++ b/lab_control/test/sds1000xe_test.py
@@ -1,4 +1,5 @@
 import pytest
+import time
 
 from lab_control.sds1000xe import SDS1000XE
 from lab_control.test.mock_sds1000xe_device import MockSDS1000XEDevice
@@ -48,3 +49,16 @@ def test_invalidChannel(uut, mockDevice):
             with pytest.raises(AssertionError):
                 m(t)
 
+def test_setVoltsPerDivision(uut, mockDevice):
+    testValues = [5e-3, 50e-3, 1e0, 5e0, 10e0, 100e0]
+
+    for channel in SDS1000XE.AVAILABLE_CHANNELS:
+        assert mockDevice.getVoltsPerDivision(channel) is None
+        
+        for value in testValues:
+            uut.setVoltsPerDivision(channel, value)
+
+            time.sleep(0.1) # Allow time for the mock to receive and process the request
+
+            assert mockDevice.getVoltsPerDivision(channel) == value
+