Add plot window and Y-axis limit controls to GUI and CLI

3 files changed, 207 insertions, 26 deletions

diff --git a/picostream/cli.py b/picostream/cli.py
index 2006645..87c665b 100644
--- a/picostream/cli.py
+++ b/picostream/cli.py
@@ -47,6 +47,8 @@ class Streamer:
         offset_v: float = 0.0,
         max_buffer_seconds: Optional[float] = None,
         is_gui_mode: bool = False,
+        y_min: Optional[float] = None,
+        y_max: Optional[float] = None,
     ) -> None:
         # --- Configuration ---
         self.output_file = output_file
@@ -55,6 +57,8 @@ class Streamer:
         self.is_gui_mode = is_gui_mode
         self.plot_window_s = plot_window_s
         self.max_buffer_seconds = max_buffer_seconds
+        self.y_min = y_min
+        self.y_max = y_max
 
         (
             sample_rate_msps,
@@ -390,6 +394,8 @@ class Streamer:
                 hdf5_path=self.output_file,
                 display_window_seconds=self.plot_window_s,
                 decimation_factor=self.decimation_factor,
+                y_min=self.y_min,
+                y_max=self.y_max,
             )
             plotter.show()
 
@@ -532,6 +538,16 @@ def generate_unique_filename(base_path: str) -> str:
     default=False,
     help="Overwrite existing output file.",
 )
+@click.option(
+    "--y-min",
+    type=float,
+    help="Minimum Y-axis limit in mV for live plot.",
+)
+@click.option(
+    "--y-max",
+    type=float,
+    help="Maximum Y-axis limit in mV for live plot.",
+)
 def main(
     sample_rate: float,
     resolution: str,
@@ -546,9 +562,21 @@ def main(
     offset: float,
     max_buff_sec: Optional[float],
     force: bool,
+    y_min: Optional[float],
+    y_max: Optional[float],
 ) -> None:
     """High-speed data acquisition tool for Picoscope 5000a series."""
     # --- Argument Validation and Processing ---
+    
+    # Validate Y-axis limits
+    if (y_min is None) != (y_max is None):
+        logger.error("Both --y-min and --y-max must be provided together, or neither.")
+        sys.exit(1)
+    
+    if y_min is not None and y_max is not None and y_min >= y_max:
+        logger.error(f"Invalid Y-axis range: y_min ({y_min}) must be less than y_max ({y_max}).")
+        sys.exit(1)
+    
     channel_range_str = VOLTAGE_RANGE_MAP[float(rangev)]
     resolution_bits = int(resolution)
 
@@ -603,6 +631,8 @@ def main(
             downsample_mode=downsample_mode,
             offset_v=offset,
             max_buffer_seconds=max_buff_sec,
+            y_min=y_min,
+            y_max=y_max,
         )
         
         # Update the acquisition command in metadata
diff --git a/picostream/dfplot.py b/picostream/dfplot.py
index e338d58..15196fa 100644
--- a/picostream/dfplot.py
+++ b/picostream/dfplot.py
@@ -38,14 +38,25 @@ class HDF5LivePlotter(QWidget):
         update_interval_ms: int = 50,
         display_window_seconds: float = 0.5,
         decimation_factor: int = 150,
+        y_min: Optional[float] = None,
+        y_max: Optional[float] = None,
     ) -> None:
         """Initializes the HDF5LivePlotter window.
 
-        Args:
-            hdf5_path: Path to the HDF5 file to monitor.
-            update_interval_ms: How often to check the file for updates (in ms).
-            display_window_seconds: The time duration of data to display.
-            decimation_factor: The factor by which to decimate data for plotting.
+        Parameters
+        ----------
+        hdf5_path : str
+            Path to the HDF5 file to monitor.
+        update_interval_ms : int
+            How often to check the file for updates (in ms).
+        display_window_seconds : float
+            The time duration of data to display.
+        decimation_factor : int
+            The factor by which to decimate data for plotting.
+        y_min : float, optional
+            Minimum Y-axis limit in mV. If None, auto-range is used.
+        y_max : float, optional
+            Maximum Y-axis limit in mV. If None, auto-range is used.
         """
         super().__init__()
 
@@ -54,6 +65,8 @@ class HDF5LivePlotter(QWidget):
         self.update_interval_ms: int = update_interval_ms
         self.display_window_seconds: float = display_window_seconds
         self.decimation_factor: int = decimation_factor
+        self.y_min: Optional[float] = y_min
+        self.y_max: Optional[float] = y_max
 
         # --- UI State ---
         self.heartbeat_chars: List[str] = ["|", "/", "-", "\\"]
@@ -176,6 +189,40 @@ class HDF5LivePlotter(QWidget):
         # Don't check file or read metadata here - wait for update_from_file()
         # to be called, which will only read metadata once the file exists with valid data
 
+    def set_display_window(self, window_seconds: float) -> None:
+        """Sets the temporal display window width.
+        
+        Parameters
+        ----------
+        window_seconds : float
+            The time duration of data to display in seconds.
+        """
+        self.display_window_seconds = window_seconds
+        logger.info(f"Display window set to {window_seconds}s")
+
+    def set_y_limits(self, y_min: Optional[float], y_max: Optional[float]) -> None:
+        """Sets fixed Y-axis limits or enables auto-ranging.
+        
+        Parameters
+        ----------
+        y_min : float, optional
+            Minimum Y-axis limit in mV. If None, auto-range is used.
+        y_max : float, optional
+            Maximum Y-axis limit in mV. If None, auto-range is used.
+        """
+        self.y_min = y_min
+        self.y_max = y_max
+        
+        has_fixed_limits = y_min is not None and y_max is not None
+        
+        if has_fixed_limits:
+            self.plot_widget.setYRange(y_min, y_max, padding=0)
+            self.plot_widget.disableAutoRange(axis='y')
+            logger.info(f"Y-axis limits set to [{y_min}, {y_max}] mV")
+        else:
+            self.plot_widget.enableAutoRange(axis='y')
+            logger.info("Y-axis auto-ranging enabled")
+
     def start_updates(self) -> None:
         """Starts the plot update timer."""
         self.timer.start(self.update_interval_ms)
@@ -238,6 +285,12 @@ class HDF5LivePlotter(QWidget):
         self.plot_widget.showGrid(x=True, y=True)
         self.plot_widget.setXRange(0, self.display_window_seconds, padding=0)
 
+        # Apply Y-axis limits if provided
+        has_fixed_y_limits = self.y_min is not None and self.y_max is not None
+        if has_fixed_y_limits:
+            self.plot_widget.setYRange(self.y_min, self.y_max, padding=0)
+            self.plot_widget.disableAutoRange(axis='y')
+
         # Plot curve
         self.curve = self.plot_widget.plot(pen="y", width=1)
 
@@ -274,8 +327,10 @@ class HDF5LivePlotter(QWidget):
     def read_metadata(self, hdf5_file: h5py.File) -> None:
         """Reads metadata attributes from the root of an open HDF5 file.
 
-        Args:
-            hdf5_file: An open h5py.File object.
+        Parameters
+        ----------
+        hdf5_file : h5py.File
+            An open h5py.File object.
         """
         if self.metadata_read:
             logger.debug("Metadata already read, skipping")
@@ -511,8 +566,10 @@ class HDF5LivePlotter(QWidget):
 
         This involves decimation, voltage conversion, and updating the plot curve.
 
-        Args:
-            data_window: A NumPy array containing the raw ADC counts for display.
+        Parameters
+        ----------
+        data_window : np.ndarray
+            A NumPy array containing the raw ADC counts for display.
         """
         if len(data_window) == 0:
             return
@@ -578,9 +635,9 @@ class HDF5LivePlotter(QWidget):
         # Update the X-axis range to match the new time axis, creating a "snapshot" effect.
         self.plot_widget.setXRange(time_axis[0], time_axis[-1], padding=0)
 
-        # Auto-scale the Y-axis occasionally. 
-        # TODO put to larger number and allow user to update with key press
-        if self.display_update_count % 5 == 1:
+        # Auto-scale the Y-axis occasionally, but only if fixed limits are not set
+        has_fixed_y_limits = self.y_min is not None and self.y_max is not None
+        if not has_fixed_y_limits and self.display_update_count % 5 == 1:
             self.plot_widget.enableAutoRange(axis="y")
 
     def _format_rate_sps(self, rate_sps: float) -> str:
@@ -596,10 +653,14 @@ class HDF5LivePlotter(QWidget):
     def format_sample_count(self, count: int) -> str:
         """Formats a large integer count into a human-readable string with units.
 
-        Args:
-            count: The integer number to format.
+        Parameters
+        ----------
+        count : int
+            The integer number to format.
 
-        Returns:
+        Returns
+        -------
+        str
             A formatted string (e.g., "1.23M", "2.34G").
         """
         if count >= 1_000_000_000:
@@ -622,11 +683,15 @@ class HDF5LivePlotter(QWidget):
         draw vertical lines for each min-max pair. For non-decimated data, it
         generates a linearly spaced time axis.
 
-        Args:
-            n_samples: The number of points for the time axis. This should be
-                the number of points *after* decimation.
+        Parameters
+        ----------
+        n_samples : int
+            The number of points for the time axis. This should be
+            the number of points *after* decimation.
 
-        Returns:
+        Returns
+        -------
+        np.ndarray
             A NumPy array representing the time axis in seconds.
         """
         if n_samples == 0:
@@ -730,6 +795,16 @@ def main() -> None:
         default=150,
         help="Decimation factor for plotting. [default: 150]",
     )
+    parser.add_argument(
+        "--y-min",
+        type=float,
+        help="Minimum Y-axis limit in mV.",
+    )
+    parser.add_argument(
+        "--y-max",
+        type=float,
+        help="Maximum Y-axis limit in mV.",
+    )
     args = parser.parse_args()
 
     logger.info("Plotter process starting")
@@ -738,6 +813,8 @@ def main() -> None:
             hdf5_path=args.hdf5_path,
             display_window_seconds=args.window,
             decimation_factor=args.decimation,
+            y_min=args.y_min,
+            y_max=args.y_max,
         )
         plotter.show()
         sys.exit(app.exec_())
diff --git a/picostream/main.py b/picostream/main.py
index 81d027f..1a0b4d8 100644
--- a/picostream/main.py
+++ b/picostream/main.py
@@ -129,14 +129,53 @@ class PicoStreamMainWindow(QMainWindow):
         )
         form_layout.addRow(self.live_only_checkbox, self.max_buffer_input)
 
+        # Plot settings
+        self.plot_window_input = QDoubleSpinBox()
+        self.plot_window_input.setRange(0.01, 10.0)
+        self.plot_window_input.setValue(0.5)
+        self.plot_window_input.setSingleStep(0.1)
+        self.plot_window_input.setDecimals(2)
+        self.plot_window_input.setSuffix(" s")
+        form_layout.addRow("Plot Window:", self.plot_window_input)
+
+        self.y_axis_auto_checkbox = QCheckBox("Auto Y-axis")
+        self.y_axis_auto_checkbox.setChecked(True)
+        self.y_min_input = QDoubleSpinBox()
+        self.y_min_input.setRange(-100000, 100000)
+        self.y_min_input.setValue(-1000)
+        self.y_min_input.setSuffix(" mV")
+        self.y_min_input.setEnabled(False)
+        self.y_max_input = QDoubleSpinBox()
+        self.y_max_input.setRange(-100000, 100000)
+        self.y_max_input.setValue(1000)
+        self.y_max_input.setSuffix(" mV")
+        self.y_max_input.setEnabled(False)
+        
+        self.y_axis_auto_checkbox.stateChanged.connect(
+            lambda state: self.y_min_input.setEnabled(state == 0)
+        )
+        self.y_axis_auto_checkbox.stateChanged.connect(
+            lambda state: self.y_max_input.setEnabled(state == 0)
+        )
+        
+        form_layout.addRow(self.y_axis_auto_checkbox)
+        y_limits_layout = QHBoxLayout()
+        y_limits_layout.addWidget(QLabel("Min:"))
+        y_limits_layout.addWidget(self.y_min_input)
+        y_limits_layout.addWidget(QLabel("Max:"))
+        y_limits_layout.addWidget(self.y_max_input)
+        form_layout.addRow("Y-axis Limits:", y_limits_layout)
+
         settings_layout.addLayout(form_layout)
 
         self.start_button = QPushButton("Start Acquisition")
         self.stop_button = QPushButton("Stop Acquisition")
         self.stop_button.setEnabled(False)
+        self.apply_plot_settings_button = QPushButton("Apply Plot Settings")
 
         settings_layout.addWidget(self.start_button)
         settings_layout.addWidget(self.stop_button)
+        settings_layout.addWidget(self.apply_plot_settings_button)
         settings_layout.addStretch()
 
         # Right panel for the plot
@@ -148,12 +187,10 @@ class PicoStreamMainWindow(QMainWindow):
         # Connect signals
         self.start_button.clicked.connect(self.start_acquisition)
         self.stop_button.clicked.connect(self.stop_acquisition)
-        # Disconnect the automatic path update - we'll handle it manually in start_acquisition
-        # self.output_file_input.textChanged.connect(self.plotter.set_hdf5_path)
+        self.apply_plot_settings_button.clicked.connect(self.apply_plot_settings)
 
         self.load_settings()
 
-
     def select_output_file(self) -> None:
         """Open a dialog to select the output HDF5 file."""
         file_name, _ = QFileDialog.getSaveFileName(
@@ -162,6 +199,18 @@ class PicoStreamMainWindow(QMainWindow):
         if file_name:
             self.output_file_input.setText(file_name)
 
+    def apply_plot_settings(self) -> None:
+        """Apply plot settings to the live plotter."""
+        window_seconds = self.plot_window_input.value()
+        self.plotter.set_display_window(window_seconds)
+        
+        if self.y_axis_auto_checkbox.isChecked():
+            self.plotter.set_y_limits(None, None)
+        else:
+            y_min = self.y_min_input.value()
+            y_max = self.y_max_input.value()
+            self.plotter.set_y_limits(y_min, y_max)
+
     def start_acquisition(self) -> None:
         """Start the background data acquisition."""
         self.start_button.setEnabled(False)
@@ -182,11 +231,14 @@ class PicoStreamMainWindow(QMainWindow):
                 logger.info(f"Removed existing file: {output_file}")
             except OSError as e:
                 self.on_acquisition_error(f"Failed to remove old file: {e}")
-                self.on_acquisition_finished()  # Reset UI state
+                self.on_acquisition_finished()
                 return
 
         # Now set the new file path (plotter is stopped and cleared)
         self.plotter.set_hdf5_path(output_file)
+        
+        # Apply plot settings before starting
+        self.apply_plot_settings()
 
         settings = {
             "sample_rate_msps": self.sample_rate_input.value(),
@@ -239,15 +291,13 @@ class PicoStreamMainWindow(QMainWindow):
     def on_acquisition_error(self, err_msg: str) -> None:
         """Handle acquisition error."""
         logger.error(f"Acquisition error: {err_msg}")
-        # In a real app, this would show a QMessageBox.
-        # The UI state is reset in on_acquisition_finished.
 
     def closeEvent(self, event: QCloseEvent) -> None:
         """Handle window close event."""
         self.save_settings()
         if self.thread and self.thread.isRunning():
             self.stop_acquisition()
-            self.thread.wait()  # Wait for the thread to finish
+            self.thread.wait()
         event.accept()
 
     def save_settings(self) -> None:
@@ -261,6 +311,13 @@ class PicoStreamMainWindow(QMainWindow):
         self.settings.setValue("offset_v", self.offset_v_input.value())
         self.settings.setValue("live_only_mode", self.live_only_checkbox.isChecked())
         self.settings.setValue("max_buffer_seconds", self.max_buffer_input.value())
+        self.settings.setValue("plot_window", self.plot_window_input.value())
+        self.settings.setValue("y_axis_auto", self.y_axis_auto_checkbox.isChecked())
+        
+        # Only save Y-axis limits if they're actually being used (not in auto mode)
+        if not self.y_axis_auto_checkbox.isChecked():
+            self.settings.setValue("y_min", self.y_min_input.value())
+            self.settings.setValue("y_max", self.y_max_input.value())
 
     def load_settings(self) -> None:
         """Load settings."""
@@ -281,6 +338,23 @@ class PicoStreamMainWindow(QMainWindow):
             self.settings.value("max_buffer_seconds", 1.0, type=float)
         )
         self.max_buffer_input.setEnabled(live_only)
+        
+        self.plot_window_input.setValue(self.settings.value("plot_window", 0.5, type=float))
+        y_axis_auto = self.settings.value("y_axis_auto", True, type=bool)
+        self.y_axis_auto_checkbox.setChecked(y_axis_auto)
+        
+        # Only load Y-axis limits if they exist in settings (were saved in manual mode)
+        # Otherwise use sensible defaults
+        if self.settings.contains("y_min") and self.settings.contains("y_max"):
+            self.y_min_input.setValue(self.settings.value("y_min", type=float))
+            self.y_max_input.setValue(self.settings.value("y_max", type=float))
+        else:
+            # Use sensible defaults if no saved values
+            self.y_min_input.setValue(-1000.0)
+            self.y_max_input.setValue(1000.0)
+            
+        self.y_min_input.setEnabled(not y_axis_auto)
+        self.y_max_input.setEnabled(not y_axis_auto)
 
 
 def main() -> None: