Klaus 65 extend jpp table feature column code

.dockerignore

@@ -0,0 +1,8 @@
+.github/
+.venv/
+.temp/
+vm/
+tests/
+*.csv/
+.python-version/
+dist/

.gitignore

@@ -7,3 +7,4 @@ __pycache__/
 .venv/
 .pytest_cache/
 .ruff_cache/
+.temp/

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"

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[

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.
 

tests/cli/__init__.py

@@ -0,0 +1 @@
+"""Tests for the CLI module."""

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"]