diff --git a/docs/howtos/head_plot_gui_integration.md b/docs/howtos/head_plot_gui_integration.md new file mode 100644 index 0000000..4f6a1cc --- /dev/null +++ b/docs/howtos/head_plot_gui_integration.md @@ -0,0 +1,161 @@ +# Head Plot GUI Integration + +## Overview +Head plots are now integrated into the GraphViewerScene and appear automatically in the GUI when calculations complete, following the same pattern as dot plots. Both use the DataFrame directly for optimal performance. + +## Implementation Details + +### Configuration Flag +Head plots are controlled by the `show_head_plot` flag in the config's `shown_outputs` section: + +```json +{ + "shown_outputs": { + "show_head_plot": true + } +} +``` + +### Integration Points + +#### GraphViewerScene.py +- **HEAD_PLOT_SPEC**: Configuration dict defining: + - `flag`: Config key to enable head plots (`show_head_plot`) + - `title_prefix`: Display name prefix for generated plots + - `required_df_cols`: Required DataFrame columns (`HeadYaw`, `LF_Angle`, `RF_Angle`) + - `settings_key`: Config section for head plot settings (`head_plot_settings`) + - `default_settings`: Fallback values for plot parameters + +- **build_head_plot_graphs()**: New function that: + 1. Checks if `show_head_plot` flag is enabled + 2. Validates required DataFrame columns are present + 3. Extracts data directly from DataFrame using `_as_numeric_array()` helper + 4. Gets time ranges from config (defaults to full dataset if not specified) + 5. Calls `plot_head()` with `ctx=None` (no file output in GUI mode) + 6. Returns dict of Plotly figures with descriptive names + +- **set_data()**: Updated to: + - Build both dot plots and head plots from the same DataFrame + - Combine warnings from both builders + - Display all graphs in the sidebar list + +### Data Flow (Optimized) +1. **Calculation Scene** completes and emits payload: + ```python + { + "results_df": pandas.DataFrame, # Contains all necessary columns + "config": dict + } + ``` + +2. **GraphViewerScene.set_data()** receives payload and: + - Calls `build_dot_plot_graphs(results_df, config)` + - Calls `build_head_plot_graphs(results_df, config)` ← **Optimized: uses DataFrame directly** + - Merges results and populates graph list + +3. **User selects graph** from sidebar → rendered as PNG via kaleido + +### Performance Optimization +**Key improvement**: Head plots now extract data directly from the DataFrame instead of reloading the CSV file: +- ✅ **No redundant file I/O**: Eliminates CSV re-read +- ✅ **No redundant parsing**: Avoids re-creating GraphDataLoader/GraphDataBundle +- ✅ **Lower memory footprint**: No duplicate data structures +- ✅ **Faster GUI response**: Data already in memory from calculations +- ✅ **Consistent pattern**: Both plot types use DataFrame extraction + +### Head Plot Behavior +- When `split_plots_by_bout=true` (default): + - Multiple figures generated, one per time range + - Named: "Head Plot: head_plot_range_[start,_end]" +- When `split_plots_by_bout=false`: + - Single combined figure + - Named: "Head Plot" + +### Error Handling +Graceful degradation with warning messages: +- **Flag disabled**: Head plot not built (silent) +- **Missing DataFrame columns**: Warning tooltip listing missing columns (e.g., "HeadYaw") +- **Empty DataFrame**: Warning: "Head plot skipped: DataFrame is empty." +- **Plot generation failure**: Warning with exception message + +### Settings Integration +Head plot settings from config are merged with defaults: +```python +{ + "head_plot_settings": { + "plot_draw_offset": 15, + "ignore_synchronized_fin_peaks": True, + "sync_fin_peaks_range": 3, + "fin_peaks_for_right_fin": False, + "split_plots_by_bout": True + }, + "graph_cutoffs": { + "left_fin_angle": 10, + "right_fin_angle": 10 + } +} +``` + +## Testing +Four test cases in `tests/unit/ui/test_graph_viewer_scene.py`: +1. **test_set_data_generates_requested_dot_plots**: Verifies dot plots still work +2. **test_set_data_skips_head_plot_when_flag_disabled**: Confirms head plots aren't generated when flag is false +3. **test_set_data_warns_when_head_plot_missing_required_columns**: Validates warning when DataFrame columns are missing +4. **test_set_data_generates_head_plot_when_all_data_present**: Confirms head plot generation with complete data + +## Usage +1. Enable head plots in your config JSON: + ```json + "shown_outputs": { "show_head_plot": true } + ``` +2. Run calculations in the GUI +3. Head plots appear automatically in the graph viewer sidebar alongside dot plots +4. Click to view static PNG renderings + +## Dependencies +- **plotly**: For figure generation +- **kaleido**: For PNG export (required for GUI rendering) +- **plot_head**: Core plotting function from `src.core.graphs.plots.headplot` +- **pandas/numpy**: For DataFrame manipulation and array operations + +## Alignment with Dotplot Pattern +Following `dotplot_gui_plan.md`: +- ✅ Emit richer payloads (DataFrame contains all necessary data) +- ✅ Interpret config flags inside GraphViewerScene +- ✅ Extract data directly from DataFrame (no file I/O on GUI path) +- ✅ Call plotting function directly with `ctx=None` +- ✅ Populate scene with returned figures +- ✅ Surface runtime issues via tooltips and empty-state messages +- ✅ Keep responsive by avoiding redundant data loading + +## Architecture Decision: DataFrame-Based Approach + +### Why Extract from DataFrame? +The initial implementation reloaded the CSV file via `GraphDataLoader`, but this was redundant since: +- `leftFinAngles` = DataFrame column `LF_Angle` +- `rightFinAngles` = DataFrame column `RF_Angle` +- `headYaw` = DataFrame column `HeadYaw` + +### Comparison + +| Aspect | Old Approach (CSV Reload) | New Approach (DataFrame) | +|--------|---------------------------|--------------------------| +| Data Source | Reload CSV via GraphDataLoader | Extract from existing DataFrame | +| File I/O | Re-reads entire CSV | None - data in memory | +| Memory | Duplicate loader/bundle objects | Single DataFrame reference | +| Performance | Slower (disk I/O + parsing) | Faster (memory access only) | +| Coupling | Requires csv_path in payload | Only needs DataFrame | +| Consistency | Different from dot plots | Same pattern as dot plots | + +### Benefits of Current Implementation +1. **Performance**: No redundant file I/O or parsing +2. **Simplicity**: Both plot types use same DataFrame extraction pattern +3. **Maintainability**: Less coupling to file system paths +4. **Memory efficiency**: No duplicate data structures +5. **GUI responsiveness**: Data already available from calculations + +## Future Enhancements +- Add integration test with real calculation pipeline +- Support dynamic plot refresh on config change +- Add export-to-file button in GUI for saving plots manually +- Consider adding time range editor for interactive bout selection diff --git a/docs/howtos/ui/HEAD_PLOT_INTEGRATION.md b/docs/howtos/ui/HEAD_PLOT_INTEGRATION.md new file mode 100644 index 0000000..e69de29 diff --git a/src/core/graphs/plots/headplot.py b/src/core/graphs/plots/headplot.py new file mode 100644 index 0000000..437b29e --- /dev/null +++ b/src/core/graphs/plots/headplot.py @@ -0,0 +1,207 @@ +from __future__ import annotations + +import math +import os +import webbrowser +from dataclasses import dataclass +from typing import Sequence, Optional + +import plotly.graph_objects as go + +from ..io import OutputContext + + +@dataclass(frozen=True) +class HeadPlotResult: + """Metadata returned after rendering a head plot.""" + figures: list[go.Figure] + labels: list[str] + output_base: Optional[str] + + +def rotate_point(origin: tuple[float, float], point: tuple[float, float], angle_deg: float) -> tuple[float, float]: + """Rotate a point around an origin by a given angle in degrees.""" + ox, oy = origin + px, py = point + angle_rad = math.radians(angle_deg) + qx = ox + math.cos(angle_rad) * (px - ox) - math.sin(angle_rad) * (py - oy) + qy = oy + math.sin(angle_rad) * (px - ox) + math.cos(angle_rad) * (py - oy) + return qx, qy + + +def plot_head( + head_yaw: Sequence[float], + left_fin_values: Sequence[float], + right_fin_values: Sequence[float], + time_ranges: Sequence[tuple[int, int]], + head_settings: dict, + cutoffs: dict, + ctx: Optional[OutputContext] = None, + open_plot: bool = False, +) -> HeadPlotResult: + """ + Plot head frames at fin peaks with optional tabbed HTML output. + """ + left_fin_cutoff = cutoffs["left_fin_angle"] + right_fin_cutoff = cutoffs["right_fin_angle"] + + draw_offset = head_settings["plot_draw_offset"] + remove_synced_peaks = head_settings["ignore_synchronized_fin_peaks"] + sync_time_range = head_settings["sync_fin_peaks_range"] + use_right_side_finpeaks = head_settings["fin_peaks_for_right_fin"] + split_by_bout = head_settings["split_plots_by_bout"] + + head_rows_by_bout = [] + fin_values = right_fin_values if use_right_side_finpeaks else left_fin_values + cutoff = right_fin_cutoff if use_right_side_finpeaks else left_fin_cutoff + opposing_fin_values = left_fin_values if use_right_side_finpeaks else right_fin_values + opposing_cutoff = left_fin_cutoff if use_right_side_finpeaks else right_fin_cutoff + + for start_index, end_index in time_ranges: + current_head_rows = [start_index] + on_peak = False + current_max_val = 0 + current_max_index = 0 + for i in range(start_index, end_index + 1): + if not on_peak and fin_values[i] > cutoff: + current_max_val = fin_values[i] + current_max_index = i + on_peak = True + elif on_peak and fin_values[i] > cutoff: + if fin_values[i] > current_max_val: + current_max_val = fin_values[i] + current_max_index = i + elif on_peak and fin_values[i] <= cutoff: + valid_point = True + if remove_synced_peaks: + for j in range(max(0, current_max_index - sync_time_range), min(current_max_index + sync_time_range, end_index)): + if opposing_fin_values[j] > opposing_cutoff: + valid_point = False + if valid_point: + current_head_rows.append(current_max_index) + on_peak = False + head_rows_by_bout.append(current_head_rows) + + head_x_points = [-0.5, 0, 0.5, 0, 0, 0] + head_y_points = [-1, 0, -1, 0, -5, -10] + origin = (head_x_points[1], head_y_points[1]) + + figures = [] + labels = [] + output_base = None + + if split_by_bout: + for bout_num, head_rows in enumerate(head_rows_by_bout): + fig = go.Figure() + fig.update_yaxes(scaleanchor="x", scaleratio=1, visible=False) + fig.update_xaxes(constrain="domain", visible=False) + fig.update_layout(showlegend=False, title="Head Plot") + + current_offset = 0 + for row in head_rows: + new_head_x = [] + new_head_y = [] + for px, py in zip(head_x_points, head_y_points): + nx, ny = rotate_point(origin, (px, py), -head_yaw[row]) + new_head_x.append(nx + current_offset) + new_head_y.append(ny) + current_offset += draw_offset + fig.add_trace(go.Scatter(x=new_head_x, y=new_head_y, mode='lines', line=dict(color="black"))) + figures.append(fig) + start, end = time_ranges[bout_num] + label = f"head_plot_range_[{start},_{end}]" + labels.append(label) + + if ctx: + html_path = os.path.join(ctx.output_folder, f"{label}.html") + png_path = os.path.join(ctx.output_folder, f"{label}.png") + fig.write_html(html_path) + try: + fig.write_image(png_path) + except Exception as e: + print(f"Warning: Could not write PNG for {label}: {e}") + print("Hint: Upgrade plotly and kaleido: pip install -U plotly kaleido") + output_base = "head_plot" + + if ctx: + tabs_html = f""" + + + + + +
+ """ + for i, label in enumerate(labels): + active_class = "active" if i == 0 else "" + tabs_html += f'
{label}
' + tabs_html += '
' + for i, fig in enumerate(figures): + fig_html = fig.to_html(include_plotlyjs=(i == 0), full_html=False) + active_style = "active" if i == 0 else "" + tabs_html += f'
{fig_html}
' + tabs_html += """ + + + + """ + tabs_path = os.path.join(ctx.output_folder, "head_plots_tabbed.html") + with open(tabs_path, 'w', encoding='utf-8') as f: + f.write(tabs_html) + if open_plot: + webbrowser.open(tabs_path) + else: + fig = go.Figure() + fig.update_layout(showlegend=False, title="Head Plot") + current_offset = 0 + for head_rows in head_rows_by_bout: + for row in head_rows: + new_head_x = [] + new_head_y = [] + for px, py in zip(head_x_points, head_y_points): + nx, ny = rotate_point(origin, (px, py), -head_yaw[row]) + new_head_x.append(nx + current_offset) + new_head_y.append(ny) + current_offset += draw_offset + fig.add_trace(go.Scatter(x=new_head_x, y=new_head_y, mode='lines', line=dict(color="black"))) + if ctx: + html_path = os.path.join(ctx.output_folder, "head_plot_plotly.html") + png_path = os.path.join(ctx.output_folder, "head_plot.png") + fig.write_html(html_path) + try: + fig.write_image(png_path) + except Exception as e: + print(f"Warning: Could not write PNG: {e}") + print("Hint: Upgrade plotly and kaleido: pip install -U plotly kaleido") + output_base = "head_plot" + if open_plot: + fig.show() + + return HeadPlotResult(figures=figures, labels=labels, output_base=output_base) + + + + + + + + + diff --git a/src/core/graphs/tests/test_generate_headplot.py b/src/core/graphs/tests/test_generate_headplot.py new file mode 100755 index 0000000..b4b9603 --- /dev/null +++ b/src/core/graphs/tests/test_generate_headplot.py @@ -0,0 +1,182 @@ +#!/usr/bin/env python3 +""" +Simple script to generate head plots (HTML + PNG) from enriched CSV data. + +Usage (from repo root): + python -m src.core.graphs.tests.test_generate_headplot + +Common options: + --csv PATH Enriched CSV (default: data/samples/csv/calculated_data_enriched.csv) + --config PATH Config JSON (default: data/samples/jsons/BaseConfig.json) + --out-dir DIR Output directory (default: results) + --output-dir DIR Alias for --out-dir + --open Open resulting HTML in a browser + +Example: + python -m src.core.graphs.tests.test_generate_headplot \ + --csv data/samples/csv/calculated_data_enriched.csv \ + --config data/samples/jsons/BaseConfig.json \ + --output-dir results/test_headplot_run +""" + +import os +import sys +from pathlib import Path + + +def _find_repo_root(start: Path) -> Path: + # Walk parents looking for sentinel files that identify project root + for p in [start] + list(start.parents): + if (p / "README.md").exists() and (p / "app.py").exists(): + return p + # Fallback: ascend until we hit directory named 'src' + for p in start.parents: + if p.name == "src": + return p.parent + # Last resort: four levels up + return start.parents[4] + + +REPO_ROOT = _find_repo_root(Path(__file__).resolve()) +SRC_ROOT = REPO_ROOT / "src" +if str(SRC_ROOT) not in sys.path: + sys.path.insert(0, str(SRC_ROOT)) + +from core.graphs.data_loader import GraphDataLoader +from core.graphs.loader_bundle import GraphDataBundle +from core.graphs.io import get_output_context +from core.graphs.plots.headplot import plot_head + + +def main(): + import argparse + parser = argparse.ArgumentParser(description="Generate head plot (HTML + PNG) from enriched CSV data") + parser.add_argument("--csv", default="data/samples/csv/calculated_data_enriched.csv", + help="Path to enriched CSV (default: data/samples/csv/calculated_data_enriched.csv)") + parser.add_argument("--config", default="data/samples/jsons/BaseConfig.json", + help="Path to config JSON (default: data/samples/jsons/BaseConfig.json)") + parser.add_argument("--out-dir", dest="out_dir", default="results", + help="Output directory (default: results)") + parser.add_argument("--output-dir", dest="out_dir", + help="Alias for --out-dir") + parser.add_argument("--open", action="store_true", + help="Open the plot in browser after generation") + + args = parser.parse_args() + + # Make paths absolute if they're relative + csv_path = args.csv if os.path.isabs(args.csv) else os.path.join(REPO_ROOT, args.csv) + config_path = args.config if os.path.isabs(args.config) else os.path.join(REPO_ROOT, args.config) + + print(f"Repository root: {REPO_ROOT}") + print(f"Loading data from: {csv_path}") + print(f"Using config: {config_path}") + + # Check files exist + if not os.path.exists(csv_path): + # Fallback attempt: use canonical sample path + alt = REPO_ROOT / "data/samples/csv/calculated_data_enriched.csv" + if alt.exists(): + print(f"CSV not found at provided path; falling back to {alt}") + csv_path = str(alt) + else: + print(f"ERROR: CSV file not found: {csv_path}") + sys.exit(1) + if not os.path.exists(config_path): + print(f"ERROR: Config file not found: {config_path}") + sys.exit(1) + + # Load data through GraphDataLoader + try: + loader = GraphDataLoader(csv_path, config_path) + bundle = GraphDataBundle.from_loader(loader) + except Exception as e: + print(f"ERROR loading data: {e}") + sys.exit(1) + + # Extract data needed for head plot + input_vals = bundle.input_values + calc_vals = bundle.calculated_values + + head_yaw = calc_vals.get("headYaw") + left_fin_angles = calc_vals.get("leftFinAngles") + right_fin_angles = calc_vals.get("rightFinAngles") + time_ranges = input_vals.get("timeRanges") + + if head_yaw is None: + print("ERROR: 'headYaw' not found in calculated values") + print(f"Available keys: {list(calc_vals.keys())}") + sys.exit(1) + if left_fin_angles is None or right_fin_angles is None: + print("ERROR: Fin angle data not found") + print(f"Available keys: {list(calc_vals.keys())}") + sys.exit(1) + if not time_ranges: + print("WARNING: No time ranges found, using full dataset") + time_ranges = [(0, len(head_yaw) - 1)] + + # Get output context + ctx = get_output_context(bundle.config, base_path=args.out_dir) + + # Get config settings for head plot + config = bundle.config + head_settings = config.get("head_settings", {}) + cutoffs = config.get("graph_cutoffs", {}) + + # Set defaults if missing + if "left_fin_angle" not in cutoffs: + cutoffs["left_fin_angle"] = 10 + if "right_fin_angle" not in cutoffs: + cutoffs["right_fin_angle"] = 10 + + default_head_settings = { + "plot_draw_offset": 15, + "ignore_synchronized_fin_peaks": True, + "sync_fin_peaks_range": 3, + "fin_peaks_for_right_fin": False, + "split_plots_by_bout": True + } + head_settings = {**default_head_settings, **head_settings} + + print(f"\nGenerating head plot...") + print(f"Output directory: {ctx.output_folder}") + print(f"Note: PNG generation requires compatible kaleido/plotly versions.") + print(f" HTML files will be generated regardless.") + + # Generate the plot + try: + result = plot_head( + head_yaw=head_yaw, + left_fin_values=left_fin_angles, + right_fin_values=right_fin_angles, + time_ranges=time_ranges, + head_settings=head_settings, + cutoffs=cutoffs, + ctx=ctx, + open_plot=args.open + ) + + print(f"\n✓ Head plot generated successfully!") + print(f" Output folder: {ctx.output_folder}") + print(f" Number of figures: {len(result.figures)}") + if result.labels: + print(f" Figure labels: {', '.join(result.labels)}") + print(f"\nFiles created:") + if head_settings.get("split_plots_by_bout"): + print(f" - head_plots_tabbed.html (interactive tabbed view)") + for label in result.labels: + print(f" - {label}.html") + print(f" - {label}.png") + else: + print(f" - head_plot_plotly.html") + print(f" - head_plot.png") + + except Exception as e: + print(f"ERROR generating plot: {e}") + import traceback + traceback.print_exc() + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/src/ui/scenes/GraphViewerScene.py b/src/ui/scenes/GraphViewerScene.py index a32fe72..9e36698 100644 --- a/src/ui/scenes/GraphViewerScene.py +++ b/src/ui/scenes/GraphViewerScene.py @@ -18,6 +18,8 @@ import plotly.graph_objs as go import plotly.io as pio +from src.core.graphs.plots import render_dot_plot +from src.core.graphs.plots.headplot import plot_head from src.core.graphs.loader_bundle import GraphDataBundle from src.core.graphs.plots import render_dot_plot, render_fin_tail, render_spines @@ -70,6 +72,20 @@ }, ) +HEAD_PLOT_SPEC: Dict[str, Any] = { + "flag": "show_head_plot", + "title_prefix": "Head Plot", + "required_df_cols": ["HeadYaw", "LF_Angle", "RF_Angle"], + "settings_key": "head_plot_settings", + "default_settings": { + "plot_draw_offset": 15, + "ignore_synchronized_fin_peaks": True, + "sync_fin_peaks_range": 3, + "fin_peaks_for_right_fin": False, + "split_plots_by_bout": True, + }, +} + class GraphViewerScene(QWidget): """ Simple static graph viewer: @@ -145,7 +161,7 @@ def add_graph(self, name: str, graph: GraphSource): self.list.setCurrentRow(0) def set_data(self, data): - """Consume calculation payload and build the requested dot plots.""" + """Consume calculation payload and build the requested dot plots and head plots.""" self._data = data self._graphs.clear() self.list.clear() @@ -160,16 +176,25 @@ def set_data(self, data): results_df = data.get("results_df") config = data.get("config") + csv_path = data.get("csv_path") if not isinstance(results_df, pd.DataFrame): - self._show_empty_state("Dot plots require a pandas DataFrame payload.") + self._show_empty_state("Graphs require a pandas DataFrame payload.") return if not isinstance(config, dict): self._show_empty_state("Missing config dictionary; cannot determine requested plots.") return + + + graphs: Dict[str, GraphSource] = {} warnings: List[str] = [] + + # Build head plots (extracts data directly from DataFrame) + head_graphs, head_warnings = build_head_plot_graphs(results_df, config) + all_graphs.update(head_graphs) + all_warnings.extend(head_warnings) dot_graphs, dot_warnings = build_dot_plot_graphs(results_df, config) graphs.update(dot_graphs) @@ -183,16 +208,16 @@ def set_data(self, data): graphs.update(spine_graphs) warnings.extend(spine_warnings) - tooltip = "\n".join(warnings) if warnings else "" + tooltip = "\n".join(all_warnings) if all_warnings else "" self.list.setToolTip(tooltip) self.image_label.setToolTip(tooltip) - if not graphs: - message = warnings[0] if warnings else "No dot plots were requested in the config." + if not all_graphs: + message = all_warnings[0] if all_warnings else "No plots were requested in the config." self._show_empty_state(message) return - self.set_graphs(graphs) + self.set_graphs(all_graphs) # Internal functions def _on_selection_changed(self): @@ -345,6 +370,14 @@ def build_dot_plot_graphs( return graphs, warnings +def build_head_plot_graphs( + results_df: pd.DataFrame, config: Dict[str, Any] +) -> Tuple[Dict[str, GraphSource], List[str]]: + """ + Build head plot figures if requested in config flags. + Extracts data directly from the DataFrame. + + Returns the generated figures and any warnings explaining skipped plots. def build_fin_tail_graphs( results_df: pd.DataFrame, config: Dict[str, Any] ) -> Tuple[Dict[str, GraphSource], List[str]]: diff --git a/tests/unit/ui/test_graph_viewer_scene.py b/tests/unit/ui/test_graph_viewer_scene.py index a62ce44..1e132df 100644 --- a/tests/unit/ui/test_graph_viewer_scene.py +++ b/tests/unit/ui/test_graph_viewer_scene.py @@ -44,10 +44,126 @@ def fake_pixmap(self, fig): }, "video_parameters": {"recorded_framerate": 50}, } - payload = {"results_df": df, "config": config, "csv_path": "demo.csv"} + payload = {"results_df": df, "config": config} scene.set_data(payload) assert len(scene._graphs) == 2 assert "Tail Distance vs Left Fin Angle" in scene._graphs assert "Tail Distance vs Left Fin Angle (Moving)" in scene._graphs + + +def test_set_data_skips_head_plot_when_flag_disabled(qt_app, monkeypatch): + """Ensure head plot is not generated when show_head_plot is False.""" + scene = GraphViewerScene() + + def fake_pixmap(self, fig): + pix = QPixmap(10, 10) + pix.fill(Qt.white) + return pix + + monkeypatch.setattr(GraphViewerScene, "_figure_to_pixmap", fake_pixmap) + + df = pd.DataFrame( + { + "Tail_Distance": [0.0, 0.05, 0.15], + "LF_Angle": [1.0, 1.1, 1.4], + "RF_Angle": [0.9, 1.0, 1.2], + } + ) + config = { + "shown_outputs": { + "show_tail_left_fin_angle_dot_plot": True, + "show_head_plot": False, + }, + "video_parameters": {"recorded_framerate": 50}, + } + payload = {"results_df": df, "config": config} + + scene.set_data(payload) + + # Only dot plot should be present + assert len(scene._graphs) == 1 + assert "Tail Distance vs Left Fin Angle" in scene._graphs + assert not any("Head Plot" in name for name in scene._graphs.keys()) + + +def test_set_data_warns_when_head_plot_missing_required_columns(qt_app, monkeypatch): + """Ensure head plot generates a warning when required DataFrame columns are missing.""" + scene = GraphViewerScene() + + def fake_pixmap(self, fig): + pix = QPixmap(10, 10) + pix.fill(Qt.white) + return pix + + monkeypatch.setattr(GraphViewerScene, "_figure_to_pixmap", fake_pixmap) + + # DataFrame missing HeadYaw column + df = pd.DataFrame( + { + "Tail_Distance": [0.0, 0.05, 0.15], + "LF_Angle": [1.0, 1.1, 1.4], + "RF_Angle": [0.9, 1.0, 1.2], + } + ) + config = { + "shown_outputs": { + "show_tail_left_fin_angle_dot_plot": True, + "show_head_plot": True, + }, + "video_parameters": {"recorded_framerate": 50}, + } + payload = {"results_df": df, "config": config} + + scene.set_data(payload) + + # Should still have the dot plot + assert len(scene._graphs) == 1 + assert "Tail Distance vs Left Fin Angle" in scene._graphs + # Check that tooltip contains warning about missing columns + tooltip = scene.list.toolTip().lower() + assert "missing dataframe columns" in tooltip and "headyaw" in tooltip + + +def test_set_data_generates_head_plot_when_all_data_present(qt_app, monkeypatch): + """Ensure head plot is generated when flag is enabled and all required data is present.""" + scene = GraphViewerScene() + + def fake_pixmap(self, fig): + pix = QPixmap(10, 10) + pix.fill(Qt.white) + return pix + + monkeypatch.setattr(GraphViewerScene, "_figure_to_pixmap", fake_pixmap) + + # DataFrame with all required columns for head plot + df = pd.DataFrame( + { + "Tail_Distance": [0.0, 0.05, 0.15], + "LF_Angle": [45.0, 50.0, 55.0], + "RF_Angle": [40.0, 45.0, 50.0], + "HeadYaw": [10.0, 15.0, 20.0], + } + ) + config = { + "shown_outputs": { + "show_tail_left_fin_angle_dot_plot": True, + "show_head_plot": True, + }, + "video_parameters": {"recorded_framerate": 50}, + "graph_cutoffs": { + "left_fin_angle": 30, + "right_fin_angle": 30, + }, + } + payload = {"results_df": df, "config": config} + + scene.set_data(payload) + + # Should have both dot plot and head plot + assert len(scene._graphs) >= 1 # At least the dot plot + assert "Tail Distance vs Left Fin Angle" in scene._graphs + # Head plot should be present (may be named with range or just "Head Plot") + has_head_plot = any("Head Plot" in name for name in scene._graphs.keys()) + assert has_head_plot