diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..e10a425 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,8 @@ +.github/ +.venv/ +.temp/ +vm/ +tests/ +*.csv/ +.python-version/ +dist/ diff --git a/.gitignore b/.gitignore index 843fd61..71bf71a 100644 --- a/.gitignore +++ b/.gitignore @@ -7,3 +7,4 @@ __pycache__/ .venv/ .pytest_cache/ .ruff_cache/ +.temp/ diff --git a/pyproject.toml b/pyproject.toml index 00d5bc2..44e2108 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "jabs-postprocess" -version = "0.1.0" +version = "0.2.0" description = "A python library for JABS postprocessing utilities." readme = "README.md" license = "LicenseRef-PLATFORM-LICENSE-AGREEMENT-FOR-NON-COMMERCIAL-USE" diff --git a/src/jabs_postprocess/analysis_utils/__init__.py b/src/jabs_postprocess/analysis_utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/jabs_postprocess/cli/main.py b/src/jabs_postprocess/cli/main.py index 4f6863e..880f512 100644 --- a/src/jabs_postprocess/cli/main.py +++ b/src/jabs_postprocess/cli/main.py @@ -13,6 +13,7 @@ generate_behavior_tables, heuristic_classify as heuristic_classify_func, ) +from jabs_postprocess.utils.project_utils import BoutTable from jabs_postprocess.utils.metadata import ( DEFAULT_INTERPOLATE, DEFAULT_MIN_BOUT, @@ -239,11 +240,24 @@ def generate_tables( ), ] = None, overwrite: Annotated[bool, typer.Option(help="Overwrites output files")] = False, + add_statistics: Annotated[ + bool, + typer.Option( + help="Add bout statistics (count, duration stats, latency) to behavior tables", + ), + ] = True, ): """Generate behavior tables from JABS predictions. This command transforms behavior predictions from a JABS project into tabular format, creating both bout-level and summary tables. + + The --add-statistics option adds additional columns with bout-level statistics: + - total_bout_count: Number of behavior bouts per animal + - avg_bout_duration: Average bout duration per animal + - bout_duration_std: Standard deviation of bout durations + - bout_duration_var: Variance of bout durations + - latency_to_first_bout: Frame number of first behavior bout """ # Convert Path to string feature_folder = feature_folder if feature_folder else None @@ -267,10 +281,29 @@ def generate_tables( overwrite=overwrite, ) + # Add bout statistics if requested + if add_statistics: + typer.echo("Adding bout statistics to generated tables...") + for behavior_name, (bout_file, summary_file) in zip( + behavior, results, strict=True + ): + try: + # Load bout table and add statistics + bout_table = BoutTable.from_file(bout_file) + bout_table.add_bout_statistics() + bout_table.to_file(bout_file, overwrite=True) + typer.echo(f" Added statistics to {bout_file}") + except Exception as e: + typer.echo( + f" Warning: Failed to add statistics to {bout_file}: {str(e)}" + ) + for behavior_name, (bout_file, summary_file) in zip(behavior, results, strict=True): typer.echo(f"Generated tables for {behavior_name}:") typer.echo(f" Bout table: {bout_file}") typer.echo(f" Summary table: {summary_file}") + if add_statistics: + typer.echo(" ✓ Includes bout statistics") @app.command() @@ -383,6 +416,83 @@ def merge_tables( raise typer.Exit(1) +@app.command() +def add_bout_statistics( + input_tables: Annotated[ + List[Path], + typer.Option(help="Paths to bout table files to add statistics to"), + ], + output_suffix: Annotated[ + str, + typer.Option(help="Suffix to add to output filenames (before .csv)"), + ] = "_with_stats", + overwrite: Annotated[ + bool, typer.Option(help="Overwrites input files instead of creating new ones") + ] = False, +): + """Add bout statistics to existing behavior tables. + + This command adds bout-level statistics to existing bout table files: + - total_bout_count: Number of behavior bouts per animal + - avg_bout_duration: Average bout duration per animal + - bout_duration_std: Standard deviation of bout durations + - bout_duration_var: Variance of bout durations + - latency_to_first_bout: Frame number of first behavior bout + + By default, creates new files with '_with_stats' suffix. Use --overwrite to modify files in-place. + """ + if not input_tables: + typer.echo("Error: No input tables provided.") + raise typer.Exit(1) + + # Validate all input files exist + for table_path in input_tables: + if not table_path.exists(): + typer.echo(f"Error: Input table not found: {table_path}") + raise typer.Exit(1) + + successful_count = 0 + for table_path in input_tables: + try: + # Load bout table and add statistics + bout_table = BoutTable.from_file(table_path) + bout_table.add_bout_statistics() + + # Determine output path + if overwrite: + output_path = table_path + else: + # Add suffix before .csv extension + stem = table_path.stem + output_path = table_path.parent / f"{stem}{output_suffix}.csv" + + # Save enhanced table + bout_table.to_file(output_path, overwrite=True) + + typer.echo(f"✓ Added statistics to: {table_path}") + if not overwrite: + typer.echo(f" Output saved to: {output_path}") + + successful_count += 1 + + except Exception as e: + typer.echo(f"✗ Failed to process {table_path}: {str(e)}") + + if successful_count > 0: + typer.echo( + f"\nSuccessfully processed {successful_count} out of {len(input_tables)} tables." + ) + typer.echo("Added statistics columns:") + typer.echo(" - total_bout_count: Number of behavior bouts per animal") + typer.echo(" - avg_bout_duration: Average bout duration per animal") + typer.echo(" - bout_duration_std: Standard deviation of bout durations") + typer.echo(" - bout_duration_var: Variance of bout durations") + typer.echo(" - latency_to_first_bout: Frame number of first behavior bout") + else: + typer.echo("Error: No tables were successfully processed.") + raise typer.Exit(1) + + @app.command() def merge_multiple_tables( table_folder: Annotated[ diff --git a/src/jabs_postprocess/utils/__init__.py b/src/jabs_postprocess/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/jabs_postprocess/utils/project_utils.py b/src/jabs_postprocess/utils/project_utils.py index 58e7761..00d203c 100644 --- a/src/jabs_postprocess/utils/project_utils.py +++ b/src/jabs_postprocess/utils/project_utils.py @@ -675,6 +675,11 @@ def __init__(self, settings: ClassifierSettings, data: pd.DataFrame): "closest_id", "closest_lixit", "closest_corner", + "total_bout_count", + "avg_bout_duration", + "bout_duration_std", + "bout_duration_var", + "latency_to_first_bout", ] self._check_fields() @@ -778,6 +783,77 @@ def add_bout_features(self, feature_file: Path): except (KeyError, ValueError): pass + def add_bout_statistics(self): + """Adds bout-level statistics as new columns to the table. + + This method calculates aggregate statistics per behavior per animal and adds them + as new columns to each bout row. This is different from add_bout_features which + summarizes per-frame features over individual bouts. + + Added columns: + - total_bout_count: Total number of behavior bouts for this animal + - avg_bout_duration: Average bout duration for this behavior for this animal + - bout_duration_std: Standard deviation of bout durations for this animal + - bout_duration_var: Variance of bout durations for this animal + - latency_to_first_bout: Frame number of first behavior bout (if any) + """ + + # Group by animal and calculate statistics for behavior bouts only + behavior_bouts = self._data[self._data["is_behavior"] == 1] + + if len(behavior_bouts) == 0: + # No behavior bouts, add columns with default values + self._data["total_bout_count"] = 0 + self._data["avg_bout_duration"] = np.nan + self._data["bout_duration_std"] = np.nan + self._data["bout_duration_var"] = np.nan + self._data["latency_to_first_bout"] = np.nan + return + + # Calculate statistics per animal + stats_by_animal = ( + behavior_bouts.groupby("animal_idx") + .agg( + { + "duration": ["count", "mean", "std", "var"], + "start": "min", # First bout start time + } + ) + .round(2) + ) + + # Flatten column names + stats_by_animal.columns = [ + "total_bout_count", + "avg_bout_duration", + "bout_duration_std", + "bout_duration_var", + "latency_to_first_bout", + ] + + # Merge statistics back to the main table + self._data = self._data.merge( + stats_by_animal, left_on="animal_idx", right_index=True, how="left" + ) + + # Fill NaN values for animals with no behavior bouts + self._data["total_bout_count"] = self._data["total_bout_count"].fillna(0) + self._data[ + [ + "avg_bout_duration", + "bout_duration_std", + "bout_duration_var", + "latency_to_first_bout", + ] + ] = self._data[ + [ + "avg_bout_duration", + "bout_duration_std", + "bout_duration_var", + "latency_to_first_bout", + ] + ].fillna(np.nan) + def to_summary_table(self, bin_size_minutes: int = 60): """Converts bout information into binned summary table. diff --git a/tests/cli/__init__.py b/tests/cli/__init__.py new file mode 100644 index 0000000..c1e6e8d --- /dev/null +++ b/tests/cli/__init__.py @@ -0,0 +1 @@ +"""Tests for the CLI module.""" diff --git a/tests/cli/conftest.py b/tests/cli/conftest.py new file mode 100644 index 0000000..73dfb5a --- /dev/null +++ b/tests/cli/conftest.py @@ -0,0 +1,65 @@ +"""Fixtures for the CLI tests.""" + +from typer.testing import CliRunner +import pytest + + +@pytest.fixture +def runner(): + """Create a Typer CLI test runner.""" + return CliRunner() + + +@pytest.fixture +def mock_project_folder(tmp_path): + """Create a mock project folder for testing.""" + project_folder = tmp_path / "test_project" + project_folder.mkdir() + return project_folder + + +@pytest.fixture +def mock_feature_folder(tmp_path): + """Create a mock feature folder for testing.""" + feature_folder = tmp_path / "features" + feature_folder.mkdir() + return feature_folder + + +@pytest.fixture +def mock_table_folder(tmp_path): + """Create a mock table folder with CSV files for testing.""" + table_folder = tmp_path / "tables" + table_folder.mkdir() + + # Create mock CSV files + files = [] + for i in range(3): + file_path = table_folder / f"behavior_{i}_bout.csv" + file_path.write_text(f"Behavior,Animal,Frame\nbehavior_{i},mouse1,100") + files.append(file_path) + + return table_folder, files + + +@pytest.fixture +def nonexistent_folder(tmp_path): + """Create path to nonexistent folder for testing.""" + return tmp_path / "nonexistent" + + +@pytest.fixture +def mock_bout_table_files(tmp_path): + """Create mock bout table files for testing.""" + files = [] + for i in range(6): # Create enough files for all test cases + file_path = tmp_path / f"test_bout_{i}.csv" + file_path.write_text("mock,csv,content") + files.append(file_path) + return files + + +@pytest.fixture +def nonexistent_files(tmp_path): + """Create paths to nonexistent files for testing.""" + return [tmp_path / "nonexistent1.csv", tmp_path / "nonexistent2.csv"] diff --git a/tests/cli/test_add_bout_statistics.py b/tests/cli/test_add_bout_statistics.py new file mode 100644 index 0000000..3e27c17 --- /dev/null +++ b/tests/cli/test_add_bout_statistics.py @@ -0,0 +1,346 @@ +"""Unit tests for the add_bout_statistics CLI command. + +This test module validates the functionality of the add_bout_statistics CLI command. +The command adds bout-level statistics to existing bout table files, including count, +duration statistics, and latency metrics. + +Key functionality tested: +1. Basic statistics addition to single and multiple files +2. Output file handling (overwrite vs new files with suffix) +3. Input validation and error handling +4. File existence checks and error reporting +5. Success/failure counting and reporting +""" + +from unittest.mock import MagicMock, patch +import pytest + +from jabs_postprocess.cli.main import app + + +class TestAddBoutStatistics: + """Test class for the add_bout_statistics CLI command.""" + + @pytest.mark.parametrize("file_count", [1, 2, 5]) + @pytest.mark.parametrize("overwrite", [True, False]) + @patch("jabs_postprocess.cli.main.BoutTable") + def test_add_bout_statistics_basic( + self, + mock_bout_table_class, + runner, + mock_bout_table_files, + file_count, + overwrite, + ): + """Test basic bout statistics addition functionality. + + Args: + mock_bout_table_class: Mock BoutTable class + runner: CLI test runner + mock_bout_table_files: List of mock bout table files + file_count: Number of files to process + overwrite: Whether to overwrite original files + """ + # Arrange + input_files = mock_bout_table_files[:file_count] + mock_bout_table = MagicMock() + mock_bout_table_class.from_file.return_value = mock_bout_table + + cmd_args = ["add-bout-statistics"] + for file_path in input_files: + cmd_args.extend(["--input-tables", str(file_path)]) + + if overwrite: + cmd_args.append("--overwrite") + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code == 0 + + # Verify BoutTable operations + assert mock_bout_table_class.from_file.call_count == file_count + assert mock_bout_table.add_bout_statistics.call_count == file_count + assert mock_bout_table.to_file.call_count == file_count + + # Verify output messages + assert ( + f"Successfully processed {file_count} out of {file_count} tables" + in result.stdout + ) + assert "total_bout_count: Number of behavior bouts per animal" in result.stdout + assert "avg_bout_duration: Average bout duration per animal" in result.stdout + + # Verify file operations for each input file + for i, input_file in enumerate(input_files): + # Check that from_file was called with correct path + from_file_calls = [ + call[0][0] for call in mock_bout_table_class.from_file.call_args_list + ] + assert input_file in from_file_calls + + # Check output path logic + to_file_calls = [ + call[0][0] for call in mock_bout_table.to_file.call_args_list + ] + if overwrite: + assert input_file in to_file_calls + else: + # Should create new file with suffix + expected_output = ( + input_file.parent / f"{input_file.stem}_with_stats.csv" + ) + assert expected_output in to_file_calls + + @pytest.mark.parametrize("output_suffix", ["_custom", "_stats_v2", "_enhanced"]) + @patch("jabs_postprocess.cli.main.BoutTable") + def test_add_bout_statistics_custom_suffix( + self, mock_bout_table_class, runner, mock_bout_table_files, output_suffix + ): + """Test bout statistics addition with custom output suffix. + + Args: + mock_bout_table_class: Mock BoutTable class + runner: CLI test runner + mock_bout_table_files: List of mock bout table files + output_suffix: Custom suffix for output files + """ + # Arrange + input_file = mock_bout_table_files[0] + mock_bout_table = MagicMock() + mock_bout_table_class.from_file.return_value = mock_bout_table + + cmd_args = [ + "add-bout-statistics", + "--input-tables", + str(input_file), + "--output-suffix", + output_suffix, + ] + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code == 0 + + # Verify output file path uses custom suffix + to_file_calls = mock_bout_table.to_file.call_args_list + assert len(to_file_calls) == 1 + output_path = to_file_calls[0][0][0] + expected_output = input_file.parent / f"{input_file.stem}{output_suffix}.csv" + assert output_path == expected_output + + @pytest.mark.parametrize("missing_file_count", [1, 2]) + def test_add_bout_statistics_missing_files( + self, runner, nonexistent_files, missing_file_count + ): + """Test error handling for missing input files. + + Args: + runner: CLI test runner + nonexistent_files: List of nonexistent file paths + missing_file_count: Number of missing files to test + """ + # Arrange + missing_files = nonexistent_files[:missing_file_count] + + cmd_args = ["add-bout-statistics"] + for file_path in missing_files: + cmd_args.extend(["--input-tables", str(file_path)]) + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code == 1 + assert "Error: Input table not found" in result.stdout + + @patch("jabs_postprocess.cli.main.BoutTable") + def test_add_bout_statistics_processing_errors( + self, mock_bout_table_class, runner, mock_bout_table_files + ): + """Test error handling during bout table processing. + + Args: + mock_bout_table_class: Mock BoutTable class + runner: CLI test runner + mock_bout_table_files: List of mock bout table files + """ + # Arrange + input_files = mock_bout_table_files[:2] + + # First file succeeds, second fails + mock_bout_table_success = MagicMock() + mock_bout_table_class.from_file.side_effect = [ + mock_bout_table_success, + Exception("Processing error"), + ] + + cmd_args = ["add-bout-statistics"] + for file_path in input_files: + cmd_args.extend(["--input-tables", str(file_path)]) + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code == 0 # Should not exit with error if some succeed + assert "Added statistics to:" in result.stdout + assert "Failed to process" in result.stdout + assert "Successfully processed 1 out of 2 tables" in result.stdout + + @patch("jabs_postprocess.cli.main.BoutTable") + def test_add_bout_statistics_all_processing_errors( + self, mock_bout_table_class, runner, mock_bout_table_files + ): + """Test error handling when all processing fails. + + Args: + mock_bout_table_class: Mock BoutTable class + runner: CLI test runner + mock_bout_table_files: List of mock bout table files + """ + # Arrange + input_file = mock_bout_table_files[0] + mock_bout_table_class.from_file.side_effect = Exception("Processing error") + + cmd_args = [ + "add-bout-statistics", + "--input-tables", + str(input_file), + ] + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code == 1 + assert "Failed to process" in result.stdout + assert "Error: No tables were successfully processed" in result.stdout + + def test_add_bout_statistics_no_input_tables(self, runner): + """Test error handling when no input tables are provided. + + Args: + runner: CLI test runner + """ + # Arrange & Act + result = runner.invoke(app, ["add-bout-statistics"]) + + # Assert - Typer will error before reaching our code due to missing required option + assert result.exit_code != 0 + assert "Missing option" in result.stdout + + @patch("jabs_postprocess.cli.main.BoutTable") + def test_add_bout_statistics_output_messages( + self, mock_bout_table_class, runner, mock_bout_table_files + ): + """Test that appropriate output messages are displayed. + + Args: + mock_bout_table_class: Mock BoutTable class + runner: CLI test runner + mock_bout_table_files: List of mock bout table files + """ + # Arrange + input_files = mock_bout_table_files[:2] + mock_bout_table = MagicMock() + mock_bout_table_class.from_file.return_value = mock_bout_table + + cmd_args = ["add-bout-statistics"] + for file_path in input_files: + cmd_args.extend(["--input-tables", str(file_path)]) + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code == 0 + + # Check for specific output messages + for input_file in input_files: + assert f"Added statistics to: {input_file}" in result.stdout + assert ( + f"Output saved to: {input_file.parent / (input_file.stem + '_with_stats.csv')}" + in result.stdout + ) + + # Check for statistics description + expected_stats = [ + "total_bout_count: Number of behavior bouts per animal", + "avg_bout_duration: Average bout duration per animal", + "bout_duration_std: Standard deviation of bout durations", + "bout_duration_var: Variance of bout durations", + "latency_to_first_bout: Frame number of first behavior bout", + ] + + for stat_desc in expected_stats: + assert stat_desc in result.stdout + + @patch("jabs_postprocess.cli.main.BoutTable") + def test_add_bout_statistics_overwrite_no_output_message( + self, mock_bout_table_class, runner, mock_bout_table_files + ): + """Test that output path message is not shown when overwriting. + + Args: + mock_bout_table_class: Mock BoutTable class + runner: CLI test runner + mock_bout_table_files: List of mock bout table files + """ + # Arrange + input_file = mock_bout_table_files[0] + mock_bout_table = MagicMock() + mock_bout_table_class.from_file.return_value = mock_bout_table + + cmd_args = [ + "add-bout-statistics", + "--input-tables", + str(input_file), + "--overwrite", + ] + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code == 0 + assert f"Added statistics to: {input_file}" in result.stdout + assert "Output saved to:" not in result.stdout + + @patch("jabs_postprocess.cli.main.BoutTable") + def test_add_bout_statistics_mixed_success_failure( + self, mock_bout_table_class, runner, mock_bout_table_files + ): + """Test handling of mixed success and failure scenarios. + + Args: + mock_bout_table_class: Mock BoutTable class + runner: CLI test runner + mock_bout_table_files: List of mock bout table files + """ + # Arrange + input_files = mock_bout_table_files[:3] + + # Setup mixed success/failure pattern + def side_effect(file_path): + if "test_bout_1" in str(file_path): + raise Exception("Middle file error") + return MagicMock() + + mock_bout_table_class.from_file.side_effect = side_effect + + cmd_args = ["add-bout-statistics"] + for file_path in input_files: + cmd_args.extend(["--input-tables", str(file_path)]) + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code == 0 + assert "Successfully processed 2 out of 3 tables" in result.stdout + assert "Added statistics to:" in result.stdout + assert "Failed to process" in result.stdout diff --git a/tests/cli/test_generate_tables.py b/tests/cli/test_generate_tables.py new file mode 100644 index 0000000..391ab0e --- /dev/null +++ b/tests/cli/test_generate_tables.py @@ -0,0 +1,337 @@ +"""Unit tests for the generate_tables CLI command. + +This test module validates the functionality of the generate_tables CLI command. +The command transforms behavior predictions from a JABS project into tabular format, +creating both bout-level and summary tables. + +Key functionality tested: +1. Basic table generation without statistics +2. Table generation with bout statistics +3. Parameter validation and error handling +4. File output and overwrite behavior +5. Integration with underlying processing modules +""" + +from unittest.mock import MagicMock, patch +import pytest + +from jabs_postprocess.cli.main import app + + +class TestGenerateTables: + """Test class for the generate_tables CLI command.""" + + @pytest.mark.parametrize("behavior_count", [1, 2, 3]) + @pytest.mark.parametrize("add_statistics", [True, False, None]) + @patch("jabs_postprocess.cli.main.generate_behavior_tables") + @patch("jabs_postprocess.cli.main.BoutTable") + def test_generate_tables_basic( + self, + mock_bout_table_class, + mock_generate_module, + runner, + mock_project_folder, + behavior_count, + add_statistics, + ): + """Test basic table generation functionality. + + Args: + mock_bout_table_class: Mock BoutTable class + mock_generate_module: Mock generate_behavior_tables module + runner: CLI test runner + mock_project_folder: Mock project directory + behavior_count: Number of behaviors to test + add_statistics: Whether to add bout statistics + """ + # Arrange + behaviors = [f"behavior{i}" for i in range(behavior_count)] + out_prefix = "test_output" + + # Mock the process_multiple_behaviors return value + mock_results = [ + (f"bout_{i}.csv", f"summary_{i}.csv") for i in range(behavior_count) + ] + mock_generate_module.process_multiple_behaviors.return_value = mock_results + + # Mock BoutTable for statistics addition + mock_bout_table = MagicMock() + mock_bout_table_class.from_file.return_value = mock_bout_table + + # Prepare CLI arguments + cmd_args = [ + "generate-tables", + "--project-folder", + str(mock_project_folder), + "--out-prefix", + out_prefix, + ] + + for behavior in behaviors: + cmd_args.extend(["--behavior", behavior]) + + if add_statistics: + cmd_args.append("--add-statistics") + elif add_statistics is False: + cmd_args.append("--no-add-statistics") + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code == 0 + + # Verify process_multiple_behaviors was called with correct parameters + mock_generate_module.process_multiple_behaviors.assert_called_once() + call_args = mock_generate_module.process_multiple_behaviors.call_args + + assert call_args.kwargs["project_folder"] == mock_project_folder + assert call_args.kwargs["out_prefix"] == out_prefix + assert len(call_args.kwargs["behaviors"]) == behavior_count + + for i, behavior_config in enumerate(call_args.kwargs["behaviors"]): + assert behavior_config["behavior"] == f"behavior{i}" + assert "interpolate_size" in behavior_config + assert "stitch_gap" in behavior_config + assert "min_bout_length" in behavior_config + + # Verify statistics handling + if add_statistics or add_statistics is None: + assert mock_bout_table_class.from_file.call_count == behavior_count + assert mock_bout_table.add_bout_statistics.call_count == behavior_count + assert mock_bout_table.to_file.call_count == behavior_count + elif add_statistics is False: + mock_bout_table_class.from_file.assert_not_called() + + @pytest.mark.parametrize( + "interpolate_size,stitch_gap,min_bout_length,out_bin_size", + [ + (None, None, None, 60), + (30, 5, 10, 120), + (50, None, 15, 30), + ], + ) + @patch("jabs_postprocess.cli.main.generate_behavior_tables") + def test_generate_tables_with_parameters( + self, + mock_generate_module, + runner, + mock_project_folder, + mock_feature_folder, + interpolate_size, + stitch_gap, + min_bout_length, + out_bin_size, + ): + """Test table generation with various parameter combinations. + + Args: + mock_generate_module: Mock generate_behavior_tables module + runner: CLI test runner + mock_project_folder: Mock project directory + mock_feature_folder: Mock feature directory + interpolate_size: Interpolation size parameter + stitch_gap: Stitch gap parameter + min_bout_length: Minimum bout length parameter + out_bin_size: Output bin size parameter + """ + # Arrange + behavior = "test_behavior" + mock_generate_module.process_multiple_behaviors.return_value = [ + ("bout.csv", "summary.csv") + ] + + cmd_args = [ + "generate-tables", + "--project-folder", + str(mock_project_folder), + "--behavior", + behavior, + "--feature-folder", + str(mock_feature_folder), + "--out-bin-size", + str(out_bin_size), + ] + + if interpolate_size is not None: + cmd_args.extend(["--interpolate-size", str(interpolate_size)]) + if stitch_gap is not None: + cmd_args.extend(["--stitch-gap", str(stitch_gap)]) + if min_bout_length is not None: + cmd_args.extend(["--min-bout-length", str(min_bout_length)]) + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code == 0 + + call_args = mock_generate_module.process_multiple_behaviors.call_args + assert call_args.kwargs["feature_folder"] == mock_feature_folder + assert call_args.kwargs["out_bin_size"] == out_bin_size + + behavior_config = call_args.kwargs["behaviors"][0] + assert behavior_config["interpolate_size"] == interpolate_size + assert behavior_config["stitch_gap"] == stitch_gap + assert behavior_config["min_bout_length"] == min_bout_length + + @pytest.mark.parametrize("overwrite", [True, False]) + @patch("jabs_postprocess.cli.main.generate_behavior_tables") + def test_generate_tables_overwrite_option( + self, mock_generate_module, runner, mock_project_folder, overwrite + ): + """Test table generation with overwrite option. + + Args: + mock_generate_module: Mock generate_behavior_tables module + runner: CLI test runner + mock_project_folder: Mock project directory + overwrite: Whether to enable overwrite + """ + # Arrange + behavior = "test_behavior" + mock_generate_module.process_multiple_behaviors.return_value = [ + ("bout.csv", "summary.csv") + ] + + cmd_args = [ + "generate-tables", + "--project-folder", + str(mock_project_folder), + "--behavior", + behavior, + ] + + if overwrite: + cmd_args.append("--overwrite") + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code == 0 + + call_args = mock_generate_module.process_multiple_behaviors.call_args + assert call_args.kwargs["overwrite"] == overwrite + + @patch("jabs_postprocess.cli.main.generate_behavior_tables") + @patch("jabs_postprocess.cli.main.BoutTable") + def test_generate_tables_statistics_error_handling( + self, mock_bout_table_class, mock_generate_module, runner, mock_project_folder + ): + """Test error handling when adding bout statistics fails. + + Args: + mock_bout_table_class: Mock BoutTable class + mock_generate_module: Mock generate_behavior_tables module + runner: CLI test runner + mock_project_folder: Mock project directory + """ + # Arrange + behavior = "test_behavior" + mock_generate_module.process_multiple_behaviors.return_value = [ + ("bout.csv", "summary.csv") + ] + + # Mock BoutTable to raise an exception + mock_bout_table_class.from_file.side_effect = Exception("Test error") + + cmd_args = [ + "generate-tables", + "--project-folder", + str(mock_project_folder), + "--behavior", + behavior, + "--add-statistics", + ] + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code == 0 # Command should still succeed + assert "Warning: Failed to add statistics" in result.stdout + + @patch("jabs_postprocess.cli.main.generate_behavior_tables") + @patch("jabs_postprocess.cli.main.BoutTable") + def test_generate_tables_output_messages( + self, mock_bout_table_class, mock_generate_module, runner, mock_project_folder + ): + """Test that appropriate output messages are displayed. + + Args: + mock_bout_table_class: Mock BoutTable class + mock_generate_module: Mock generate_behavior_tables module + runner: CLI test runner + mock_project_folder: Mock project directory + """ + # Arrange + behaviors = ["behavior1", "behavior2"] + mock_generate_module.process_multiple_behaviors.return_value = [ + ("bout1.csv", "summary1.csv"), + ("bout2.csv", "summary2.csv"), + ] + + # Mock BoutTable for statistics addition + mock_bout_table = MagicMock() + mock_bout_table_class.from_file.return_value = mock_bout_table + + cmd_args = [ + "generate-tables", + "--project-folder", + str(mock_project_folder), + ] + + for behavior in behaviors: + cmd_args.extend(["--behavior", behavior]) + + cmd_args.append("--add-statistics") + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code == 0 + assert "Generated tables for behavior1:" in result.stdout + assert "Generated tables for behavior2:" in result.stdout + assert "Bout table: bout1.csv" in result.stdout + assert "Summary table: summary1.csv" in result.stdout + assert "Includes bout statistics" in result.stdout + + def test_generate_tables_missing_required_args(self, runner): + """Test that missing required arguments cause appropriate errors. + + Args: + runner: CLI test runner + """ + # Arrange & Act + result = runner.invoke(app, ["generate-tables"]) + + # Assert + assert result.exit_code != 0 + assert "Missing option" in result.stdout + + @patch("jabs_postprocess.cli.main.generate_behavior_tables") + def test_generate_tables_no_behaviors( + self, mock_generate_module, runner, mock_project_folder + ): + """Test behavior when no behaviors are specified. + + Args: + mock_generate_module: Mock generate_behavior_tables module + runner: CLI test runner + mock_project_folder: Mock project directory + """ + # Arrange + cmd_args = [ + "generate-tables", + "--project-folder", + str(mock_project_folder), + ] + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code != 0 + assert "Missing option" in result.stdout diff --git a/tests/cli/test_merge_multiple_tables.py b/tests/cli/test_merge_multiple_tables.py new file mode 100644 index 0000000..9cde757 --- /dev/null +++ b/tests/cli/test_merge_multiple_tables.py @@ -0,0 +1,525 @@ +"""Unit tests for the merge_multiple_tables CLI command. + +This test module validates the functionality of the merge_multiple_tables CLI command. +The command scans a folder for behavior table files, groups them by behavior name, and +merges each group separately. This is useful for combining results from multiple +experiments. + +Key functionality tested: +1. Basic table merging with auto-detected behaviors +2. Behavior filtering and selection +3. File pattern matching and table discovery +4. Output file generation and overwrite behavior +5. Error handling for missing folders and invalid files +6. Table grouping by behavior name extraction +""" + +from pathlib import Path +from unittest.mock import patch +import pytest +import pandas as pd + +from jabs_postprocess.cli.main import app + + +class TestMergeMultipleTables: + """Test class for the merge_multiple_tables CLI command.""" + + @pytest.mark.parametrize("behavior_count", [1, 2, 4]) + @pytest.mark.parametrize("overwrite", [True, False]) + @patch("jabs_postprocess.cli.main.generate_behavior_tables") + @patch("pandas.read_csv") + def test_merge_multiple_tables_basic( + self, + mock_read_csv, + mock_generate_module, + runner, + mock_table_folder, + behavior_count, + overwrite, + ): + """Test basic table merging functionality. + + Args: + mock_read_csv: Mock pandas read_csv function + mock_generate_module: Mock generate_behavior_tables module + runner: CLI test runner + mock_table_folder: Mock table folder with files + behavior_count: Number of behaviors to simulate + overwrite: Whether to enable overwrite + """ + # Arrange + table_folder, files = mock_table_folder + + # Mock pandas read_csv to return behavior names + def mock_csv_reader(file_path, nrows=None): + file_name = Path(file_path).name + if "behavior_0" in file_name: + return pd.DataFrame({"Behavior": ["behavior_0"]}) + elif "behavior_1" in file_name: + return pd.DataFrame({"Behavior": ["behavior_1"]}) + else: + return pd.DataFrame({"Behavior": ["behavior_2"]}) + + mock_read_csv.side_effect = mock_csv_reader + + # Mock merge function return value + mock_results = { + f"behavior_{i}": (f"merged_bout_{i}.csv", f"merged_bin_{i}.csv") + for i in range(behavior_count) + } + mock_generate_module.merge_multiple_behavior_tables.return_value = mock_results + + cmd_args = [ + "merge-multiple-tables", + "--table-folder", + str(table_folder), + ] + + if overwrite: + cmd_args.append("--overwrite") + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code == 0 + + # Verify merge function was called + mock_generate_module.merge_multiple_behavior_tables.assert_called_once() + call_args = mock_generate_module.merge_multiple_behavior_tables.call_args + + assert call_args.kwargs["output_prefix"] == "merged_behavior" + assert call_args.kwargs["overwrite"] == overwrite + assert "table_groups" in call_args.kwargs + + # Verify output messages + assert ( + f"Successfully merged tables for {behavior_count} behaviors:" + in result.stdout + ) + + @pytest.mark.parametrize("table_pattern", ["*.csv", "*_bout.csv", "behavior_*.csv"]) + @patch("jabs_postprocess.cli.main.generate_behavior_tables") + @patch("pandas.read_csv") + def test_merge_multiple_tables_custom_pattern( + self, + mock_read_csv, + mock_generate_module, + runner, + mock_table_folder, + table_pattern, + ): + """Test table merging with custom file patterns. + + Args: + mock_read_csv: Mock pandas read_csv function + mock_generate_module: Mock generate_behavior_tables module + runner: CLI test runner + mock_table_folder: Mock table folder with files + table_pattern: File pattern to test + """ + # Arrange + table_folder, files = mock_table_folder + mock_read_csv.return_value = pd.DataFrame({"Behavior": ["test_behavior"]}) + mock_generate_module.merge_multiple_behavior_tables.return_value = { + "test_behavior": ("bout.csv", "bin.csv") + } + + cmd_args = [ + "merge-multiple-tables", + "--table-folder", + str(table_folder), + "--table-pattern", + table_pattern, + ] + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code == 0 + mock_generate_module.merge_multiple_behavior_tables.assert_called_once() + + @pytest.mark.parametrize( + "selected_behaviors", [["behavior_0"], ["behavior_1", "behavior_2"]] + ) + @patch("jabs_postprocess.cli.main.generate_behavior_tables") + @patch("pandas.read_csv") + def test_merge_multiple_tables_behavior_filtering( + self, + mock_read_csv, + mock_generate_module, + runner, + mock_table_folder, + selected_behaviors, + ): + """Test table merging with behavior filtering. + + Args: + mock_read_csv: Mock pandas read_csv function + mock_generate_module: Mock generate_behavior_tables module + runner: CLI test runner + mock_table_folder: Mock table folder with files + selected_behaviors: List of behaviors to filter for + """ + # Arrange + table_folder, files = mock_table_folder + + # Mock pandas to return different behaviors for different files + def mock_csv_reader(file_path, nrows=None): + file_name = Path(file_path).name + if "behavior_0" in file_name: + return pd.DataFrame({"Behavior": ["behavior_0"]}) + elif "behavior_1" in file_name: + return pd.DataFrame({"Behavior": ["behavior_1"]}) + else: + return pd.DataFrame({"Behavior": ["behavior_2"]}) + + mock_read_csv.side_effect = mock_csv_reader + + # Mock results for selected behaviors only + mock_results = { + behavior: (f"{behavior}_bout.csv", f"{behavior}_bin.csv") + for behavior in selected_behaviors + } + mock_generate_module.merge_multiple_behavior_tables.return_value = mock_results + + cmd_args = [ + "merge-multiple-tables", + "--table-folder", + str(table_folder), + ] + + for behavior in selected_behaviors: + cmd_args.extend(["--behaviors", behavior]) + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code == 0 + + call_args = mock_generate_module.merge_multiple_behavior_tables.call_args + table_groups = call_args.kwargs["table_groups"] + + # Verify only selected behaviors are included + assert set(table_groups.keys()) == set(selected_behaviors) + + @pytest.mark.parametrize( + "output_prefix", ["custom_prefix", "experiment_1", "merged_data"] + ) + @patch("jabs_postprocess.cli.main.generate_behavior_tables") + @patch("pandas.read_csv") + def test_merge_multiple_tables_custom_prefix( + self, + mock_read_csv, + mock_generate_module, + runner, + mock_table_folder, + output_prefix, + ): + """Test table merging with custom output prefix. + + Args: + mock_read_csv: Mock pandas read_csv function + mock_generate_module: Mock generate_behavior_tables module + runner: CLI test runner + mock_table_folder: Mock table folder with files + output_prefix: Custom output prefix to test + """ + # Arrange + table_folder, files = mock_table_folder + mock_read_csv.return_value = pd.DataFrame({"Behavior": ["test_behavior"]}) + mock_generate_module.merge_multiple_behavior_tables.return_value = { + "test_behavior": ("bout.csv", "bin.csv") + } + + cmd_args = [ + "merge-multiple-tables", + "--table-folder", + str(table_folder), + "--output-prefix", + output_prefix, + ] + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code == 0 + + call_args = mock_generate_module.merge_multiple_behavior_tables.call_args + assert call_args.kwargs["output_prefix"] == output_prefix + + def test_merge_multiple_tables_nonexistent_folder(self, runner, nonexistent_folder): + """Test error handling for nonexistent table folder. + + Args: + runner: CLI test runner + nonexistent_folder: Path to nonexistent folder + """ + # Arrange & Act + result = runner.invoke( + app, + [ + "merge-multiple-tables", + "--table-folder", + str(nonexistent_folder), + ], + ) + + # Assert + assert result.exit_code == 1 + assert f"Error: Table folder not found: {nonexistent_folder}" in result.stdout + + def test_merge_multiple_tables_empty_folder(self, runner, mock_project_folder): + """Test error handling for folder with no matching files. + + Args: + runner: CLI test runner + mock_project_folder: Path to empty folder + """ + # Arrange & Act + result = runner.invoke( + app, + [ + "merge-multiple-tables", + "--table-folder", + str(mock_project_folder), + ], + ) + + # Assert + assert result.exit_code == 1 + assert "Error: No table files found matching pattern" in result.stdout + + @patch("pandas.read_csv") + def test_merge_multiple_tables_invalid_csv_files( + self, mock_read_csv, runner, mock_table_folder + ): + """Test handling of invalid CSV files that cannot be read. + + Args: + mock_read_csv: Mock pandas read_csv function + runner: CLI test runner + mock_table_folder: Mock table folder with files + """ + # Arrange + table_folder, files = mock_table_folder + mock_read_csv.side_effect = pd.errors.EmptyDataError("No data") + + cmd_args = [ + "merge-multiple-tables", + "--table-folder", + str(table_folder), + ] + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code == 1 + assert "Error: No valid behavior tables found to merge" in result.stdout + assert "Warning: Could not read behavior from" in result.stdout + + @patch("jabs_postprocess.cli.main.generate_behavior_tables") + @patch("pandas.read_csv") + def test_merge_multiple_tables_file_exists_error( + self, mock_read_csv, mock_generate_module, runner, mock_table_folder + ): + """Test handling of FileExistsError during merge operation. + + Args: + mock_read_csv: Mock pandas read_csv function + mock_generate_module: Mock generate_behavior_tables module + runner: CLI test runner + mock_table_folder: Mock table folder with files + """ + # Arrange + table_folder, files = mock_table_folder + mock_read_csv.return_value = pd.DataFrame({"Behavior": ["test_behavior"]}) + mock_generate_module.merge_multiple_behavior_tables.side_effect = ( + FileExistsError("Output file already exists") + ) + + cmd_args = [ + "merge-multiple-tables", + "--table-folder", + str(table_folder), + ] + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code == 1 + assert "Error: Output file already exists" in result.stdout + assert "Use --overwrite to force overwrite" in result.stdout + + @patch("jabs_postprocess.cli.main.generate_behavior_tables") + @patch("pandas.read_csv") + def test_merge_multiple_tables_unexpected_error( + self, mock_read_csv, mock_generate_module, runner, mock_table_folder + ): + """Test handling of unexpected errors during merge operation. + + Args: + mock_read_csv: Mock pandas read_csv function + mock_generate_module: Mock generate_behavior_tables module + runner: CLI test runner + mock_table_folder: Mock table folder with files + """ + # Arrange + table_folder, files = mock_table_folder + mock_read_csv.return_value = pd.DataFrame({"Behavior": ["test_behavior"]}) + mock_generate_module.merge_multiple_behavior_tables.side_effect = Exception( + "Unexpected error" + ) + + cmd_args = [ + "merge-multiple-tables", + "--table-folder", + str(table_folder), + ] + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code == 1 + assert "Unexpected error: Unexpected error" in result.stdout + + @patch("jabs_postprocess.cli.main.generate_behavior_tables") + @patch("pandas.read_csv") + def test_merge_multiple_tables_output_messages( + self, mock_read_csv, mock_generate_module, runner, mock_table_folder + ): + """Test that appropriate output messages are displayed. + + Args: + mock_read_csv: Mock pandas read_csv function + mock_generate_module: Mock generate_behavior_tables module + runner: CLI test runner + mock_table_folder: Mock table folder with files + """ + # Arrange + table_folder, files = mock_table_folder + mock_read_csv.return_value = pd.DataFrame({"Behavior": ["test_behavior"]}) + + mock_results = { + "behavior_1": ("bout1.csv", "bin1.csv"), + "behavior_2": ("bout2.csv", None), # No bin file + } + mock_generate_module.merge_multiple_behavior_tables.return_value = mock_results + + cmd_args = [ + "merge-multiple-tables", + "--table-folder", + str(table_folder), + ] + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code == 0 + assert "Successfully merged tables for 2 behaviors:" in result.stdout + assert "behavior_1:" in result.stdout + assert "Bout table: bout1.csv" in result.stdout + assert "Bin table: bin1.csv" in result.stdout + assert "behavior_2:" in result.stdout + assert "Bout table: bout2.csv" in result.stdout + + @patch("pandas.read_csv") + def test_merge_multiple_tables_missing_behavior_column( + self, mock_read_csv, runner, mock_table_folder + ): + """Test handling of CSV files without Behavior column. + + Args: + mock_read_csv: Mock pandas read_csv function + runner: CLI test runner + mock_table_folder: Mock table folder with files + """ + # Arrange + table_folder, files = mock_table_folder + mock_read_csv.return_value = pd.DataFrame( + {"Animal": ["mouse1"], "Frame": [100]} + ) # No Behavior column + + cmd_args = [ + "merge-multiple-tables", + "--table-folder", + str(table_folder), + ] + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code == 1 + assert "Warning: Could not read behavior from" in result.stdout + assert "Error: No valid behavior tables found to merge" in result.stdout + + @patch("jabs_postprocess.cli.main.generate_behavior_tables") + @patch("pandas.read_csv") + def test_merge_multiple_tables_behavior_grouping( + self, mock_read_csv, mock_generate_module, runner, mock_table_folder + ): + """Test that files are properly grouped by behavior name. + + Args: + mock_read_csv: Mock pandas read_csv function + mock_generate_module: Mock generate_behavior_tables module + runner: CLI test runner + mock_table_folder: Mock table folder with files + """ + # Arrange + table_folder, files = mock_table_folder + + # Create additional files with same behavior names + extra_files = [] + for i in range(2): + file_path = table_folder / f"extra_behavior_0_{i}.csv" + file_path.write_text("Behavior,Animal,Frame\nbehavior_0,mouse1,100") + extra_files.append(file_path) + + # Mock pandas to return behaviors based on filename + def mock_csv_reader(file_path, nrows=None): + file_name = Path(file_path).name + if "behavior_0" in file_name: + return pd.DataFrame({"Behavior": ["behavior_0"]}) + elif "behavior_1" in file_name: + return pd.DataFrame({"Behavior": ["behavior_1"]}) + else: + return pd.DataFrame({"Behavior": ["behavior_2"]}) + + mock_read_csv.side_effect = mock_csv_reader + + mock_results = { + "behavior_0": ("bout0.csv", "bin0.csv"), + "behavior_1": ("bout1.csv", "bin1.csv"), + "behavior_2": ("bout2.csv", "bin2.csv"), + } + mock_generate_module.merge_multiple_behavior_tables.return_value = mock_results + + cmd_args = [ + "merge-multiple-tables", + "--table-folder", + str(table_folder), + ] + + # Act + result = runner.invoke(app, cmd_args) + + # Assert + assert result.exit_code == 0 + + call_args = mock_generate_module.merge_multiple_behavior_tables.call_args + table_groups = call_args.kwargs["table_groups"] + + # Verify that behavior_0 has multiple files (original + extras) + assert len(table_groups["behavior_0"]) == 3 # 1 original + 2 extras + assert len(table_groups["behavior_1"]) == 1 + assert len(table_groups["behavior_2"]) == 1 diff --git a/tests/utils/test_project_utils.py b/tests/utils/test_project_utils.py index ac65586..7e8a3bf 100644 --- a/tests/utils/test_project_utils.py +++ b/tests/utils/test_project_utils.py @@ -23,11 +23,12 @@ """ import numpy as np +import pandas as pd import pytest # Assuming Bouts is in jabs_utils.project_utils # Adjust the import path if your project structure is different -from jabs_postprocess.utils.project_utils import Bouts +from jabs_postprocess.utils.project_utils import Bouts, BoutTable, ClassifierSettings def test_bouts_to_vector_empty_bouts_uses_min_frames(): @@ -232,3 +233,285 @@ def test_bouts_to_vector_various_scenarios( assert np.all(vector_output == fill_state), ( "Empty initial bouts should result in vector of fill_state" ) + + +# Tests for BoutTable.add_bout_statistics method + + +class TestBoutTableAddBoutStatistics: + """Test suite for BoutTable.add_bout_statistics method.""" + + def test_add_bout_statistics_with_behavior_bouts(self): + """Test add_bout_statistics with multiple animals having behavior bouts.""" + # Arrange + data = pd.DataFrame( + { + "animal_idx": [0, 0, 0, 1, 1, 2], + "video_name": ["vid1"] * 6, + "start": [100, 300, 800, 150, 600, 50], + "duration": [50, 75, 30, 60, 45, 100], + "is_behavior": [1, 1, 1, 1, 1, 1], # All behavior bouts + "exp_prefix": ["exp1"] * 6, + "time": ["2024-01-01 12:00:00"] * 6, + } + ) + + settings = ClassifierSettings("grooming", 0, 0, 0) + bout_table = BoutTable(settings, data) + + # Act + bout_table.add_bout_statistics() + + # Assert + result_data = bout_table.data + + # Check that new columns exist + expected_columns = [ + "total_bout_count", + "avg_bout_duration", + "bout_duration_std", + "bout_duration_var", + "latency_to_first_bout", + ] + for col in expected_columns: + assert col in result_data.columns, f"Column {col} should be added" + + # Check animal 0 statistics (3 bouts: 50, 75, 30) + animal_0_data = result_data[result_data["animal_idx"] == 0].iloc[0] + assert animal_0_data["total_bout_count"] == 3 + assert abs(animal_0_data["avg_bout_duration"] - 51.67) < 0.1 # (50+75+30)/3 + assert animal_0_data["latency_to_first_bout"] == 100 + + # Check animal 1 statistics (2 bouts: 60, 45) + animal_1_data = result_data[result_data["animal_idx"] == 1].iloc[0] + assert animal_1_data["total_bout_count"] == 2 + assert abs(animal_1_data["avg_bout_duration"] - 52.5) < 0.1 # (60+45)/2 + assert animal_1_data["latency_to_first_bout"] == 150 + + # Check animal 2 statistics (1 bout: 100) + animal_2_data = result_data[result_data["animal_idx"] == 2].iloc[0] + assert animal_2_data["total_bout_count"] == 1 + assert animal_2_data["avg_bout_duration"] == 100 + assert animal_2_data["latency_to_first_bout"] == 50 + + def test_add_bout_statistics_no_behavior_bouts(self): + """Test add_bout_statistics when there are no behavior bouts.""" + # Arrange + data = pd.DataFrame( + { + "animal_idx": [0, 0, 1], + "video_name": ["vid1"] * 3, + "start": [100, 300, 150], + "duration": [50, 75, 60], + "is_behavior": [0, -1, 0], # No behavior bouts + "exp_prefix": ["exp1"] * 3, + "time": ["2024-01-01 12:00:00"] * 3, + } + ) + + settings = ClassifierSettings("grooming", 0, 0, 0) + bout_table = BoutTable(settings, data) + + # Act + bout_table.add_bout_statistics() + + # Assert + result_data = bout_table.data + + # All statistics should be 0 or NaN for no behavior bouts + assert (result_data["total_bout_count"] == 0).all() + assert result_data["avg_bout_duration"].isna().all() + assert result_data["bout_duration_std"].isna().all() + assert result_data["bout_duration_var"].isna().all() + assert result_data["latency_to_first_bout"].isna().all() + + def test_add_bout_statistics_mixed_behavior_states(self): + """Test add_bout_statistics with mixed behavior states per animal.""" + # Arrange + data = pd.DataFrame( + { + "animal_idx": [0, 0, 0, 0, 1, 1, 1], + "video_name": ["vid1"] * 7, + "start": [100, 200, 300, 400, 150, 250, 350], + "duration": [50, 25, 75, 40, 60, 30, 80], + "is_behavior": [1, 0, 1, -1, 1, 1, 0], # Mixed states + "exp_prefix": ["exp1"] * 7, + "time": ["2024-01-01 12:00:00"] * 7, + } + ) + + settings = ClassifierSettings("grooming", 0, 0, 0) + bout_table = BoutTable(settings, data) + + # Act + bout_table.add_bout_statistics() + + # Assert + result_data = bout_table.data + + # Animal 0: 2 behavior bouts (50, 75) + animal_0_rows = result_data[result_data["animal_idx"] == 0] + assert (animal_0_rows["total_bout_count"] == 2).all() + assert abs(animal_0_rows["avg_bout_duration"].iloc[0] - 62.5) < 0.1 # (50+75)/2 + assert (animal_0_rows["latency_to_first_bout"] == 100).all() + + # Animal 1: 2 behavior bouts (60, 30) + animal_1_rows = result_data[result_data["animal_idx"] == 1] + assert (animal_1_rows["total_bout_count"] == 2).all() + assert abs(animal_1_rows["avg_bout_duration"].iloc[0] - 45.0) < 0.1 # (60+30)/2 + assert (animal_1_rows["latency_to_first_bout"] == 150).all() + + def test_add_bout_statistics_single_bout_per_animal(self): + """Test add_bout_statistics with single bout per animal (variance should be NaN).""" + # Arrange + data = pd.DataFrame( + { + "animal_idx": [0, 1, 2], + "video_name": ["vid1"] * 3, + "start": [100, 200, 300], + "duration": [50, 75, 90], + "is_behavior": [1, 1, 1], + "exp_prefix": ["exp1"] * 3, + "time": ["2024-01-01 12:00:00"] * 3, + } + ) + + settings = ClassifierSettings("grooming", 0, 0, 0) + bout_table = BoutTable(settings, data) + + # Act + bout_table.add_bout_statistics() + + # Assert + result_data = bout_table.data + + # Each animal should have 1 bout, std and var should be NaN for single values + assert (result_data["total_bout_count"] == 1).all() + assert (result_data["avg_bout_duration"] == result_data["duration"]).all() + assert ( + result_data["bout_duration_std"].isna().all() + ) # std of single value is NaN + assert ( + result_data["bout_duration_var"].isna().all() + ) # var of single value is NaN + assert (result_data["latency_to_first_bout"] == result_data["start"]).all() + + def test_add_bout_statistics_preserves_existing_columns(self): + """Test that add_bout_statistics preserves existing columns and data.""" + # Arrange + data = pd.DataFrame( + { + "animal_idx": [0, 1], + "video_name": ["vid1", "vid1"], + "start": [100, 200], + "duration": [50, 75], + "is_behavior": [1, 1], + "exp_prefix": ["exp1", "exp1"], + "time": ["2024-01-01 12:00:00", "2024-01-01 12:00:00"], + "distance": [10.5, 15.2], # Use existing optional column instead + } + ) + + settings = ClassifierSettings("grooming", 0, 0, 0) + bout_table = BoutTable(settings, data) + + # Act + bout_table.add_bout_statistics() + + # Assert + result_data = bout_table.data + + # Original columns should be preserved + original_columns = [ + "animal_idx", + "video_name", + "start", + "duration", + "is_behavior", + "exp_prefix", + "time", + "distance", + ] + for col in original_columns: + assert col in result_data.columns, ( + f"Original column {col} should be preserved" + ) + + # Original data should be unchanged + assert result_data["distance"].tolist() == [10.5, 15.2] + assert result_data["start"].tolist() == [100, 200] + assert result_data["duration"].tolist() == [50, 75] + + def test_add_bout_statistics_variance_calculation(self): + """Test that variance and standard deviation are calculated correctly.""" + # Arrange - Create data with known variance + data = pd.DataFrame( + { + "animal_idx": [0, 0, 0], + "video_name": ["vid1"] * 3, + "start": [100, 200, 300], + "duration": [10, 20, 30], # Simple values for easy calculation + "is_behavior": [1, 1, 1], + "exp_prefix": ["exp1"] * 3, + "time": ["2024-01-01 12:00:00"] * 3, + } + ) + + settings = ClassifierSettings("grooming", 0, 0, 0) + bout_table = BoutTable(settings, data) + + # Act + bout_table.add_bout_statistics() + + # Assert + result_data = bout_table.data.iloc[0] + + # Manual calculation: durations = [10, 20, 30], mean = 20 + # variance = ((10-20)^2 + (20-20)^2 + (30-20)^2) / (3-1) = (100 + 0 + 100) / 2 = 100 + # std = sqrt(100) = 10 + expected_mean = 20.0 + expected_var = 100.0 + expected_std = 10.0 + + assert abs(result_data["avg_bout_duration"] - expected_mean) < 0.1 + assert abs(result_data["bout_duration_var"] - expected_var) < 0.1 + assert abs(result_data["bout_duration_std"] - expected_std) < 0.1 + + def test_add_bout_statistics_empty_dataframe(self): + """Test add_bout_statistics with completely empty DataFrame.""" + # Arrange + data = pd.DataFrame( + { + "animal_idx": [], + "video_name": [], + "start": [], + "duration": [], + "is_behavior": [], + "exp_prefix": [], + "time": [], + } + ) + + settings = ClassifierSettings("grooming", 0, 0, 0) + bout_table = BoutTable(settings, data) + + # Act + bout_table.add_bout_statistics() + + # Assert + result_data = bout_table.data + + # Should handle empty dataframe gracefully + expected_columns = [ + "total_bout_count", + "avg_bout_duration", + "bout_duration_std", + "bout_duration_var", + "latency_to_first_bout", + ] + for col in expected_columns: + assert col in result_data.columns, ( + f"Column {col} should be added even for empty data" + ) + + # No rows to check values, but shouldn't crash diff --git a/uv.lock b/uv.lock index d237644..9026d45 100644 --- a/uv.lock +++ b/uv.lock @@ -340,7 +340,7 @@ wheels = [ [[package]] name = "jabs-postprocess" -version = "0.1.0" +version = "0.2.0" source = { editable = "." } dependencies = [ { name = "black" },