Merge pull request #35 from KumarLabJax/KLAUS-118-add-bout-table-outputs-to-jabs-features

bergsalex · web-flow · commit 5dea0dd3cccf · 2025-08-22T16:57:19.000-04:00
Add behavior table merging functionality for JABS-postprocess
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "jabs-postprocess"
-version = "0.3.0"
+version = "0.4.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/cli/main.py b/src/jabs_postprocess/cli/main.py
@@ -2,6 +2,7 @@
 from pathlib import Path
 from typing import Annotated, List, Optional
 
+import pandas as pd
 import numpy as np
 import typer
 
@@ -329,5 +330,144 @@ def heuristic_classify(
     )
 
 
+@app.command()
+def merge_tables(
+    input_tables: Annotated[
+        List[Path],
+        typer.Option(
+            help="Paths to behavior table files to merge (must be same behavior and table type)"
+        ),
+    ],
+    output_prefix: Annotated[
+        str,
+        typer.Option(help="File prefix for merged output table"),
+    ] = "merged_behavior",
+    overwrite: Annotated[bool, typer.Option(help="Overwrites output files")] = False,
+):
+    """Merge multiple behavior tables of the same type and behavior.
+
+    This command merges behavior tables that contain the same behavior data,
+    combining them into a single consolidated table while preserving header information.
+    """
+    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)
+
+    try:
+        output_file, _ = generate_behavior_tables.merge_behavior_tables(
+            input_tables=input_tables,
+            output_prefix=output_prefix,
+            overwrite=overwrite,
+        )
+
+        typer.echo(f"Successfully merged {len(input_tables)} tables:")
+        for table in input_tables:
+            typer.echo(f"  - {table}")
+        typer.echo(f"Output saved to: {output_file}")
+
+    except FileExistsError as e:
+        typer.echo(f"Error: {str(e)}")
+        typer.echo("Use --overwrite to force overwrite.")
+        raise typer.Exit(1)
+    except ValueError as e:
+        typer.echo(f"Error: {str(e)}")
+        raise typer.Exit(1)
+    except Exception as e:
+        typer.echo(f"Unexpected error: {str(e)}")
+        raise typer.Exit(1)
+
+
+@app.command()
+def merge_multiple_tables(
+    table_folder: Annotated[
+        Path,
+        typer.Option(
+            help="Folder containing behavior table files to merge, grouped by behavior"
+        ),
+    ],
+    behaviors: Annotated[
+        Optional[List[str]],
+        typer.Option(help="Specific behaviors to merge (default: auto-detect all)"),
+    ] = None,
+    table_pattern: Annotated[
+        str,
+        typer.Option(help="File pattern to match behavior tables"),
+    ] = "*.csv",
+    output_prefix: Annotated[
+        str,
+        typer.Option(help="File prefix for merged output tables"),
+    ] = "merged_behavior",
+    overwrite: Annotated[bool, typer.Option(help="Overwrites output files")] = False,
+):
+    """Merge multiple sets of behavior tables, automatically grouping by behavior.
+
+    This command scans a folder for behavior table files, groups them by behavior name,
+    and merges each group separately. Useful for combining results from multiple experiments.
+    """
+    if not table_folder.exists():
+        typer.echo(f"Error: Table folder not found: {table_folder}")
+        raise typer.Exit(1)
+
+    # Find all table files matching the pattern
+    table_files = list(table_folder.glob(table_pattern))
+    if not table_files:
+        typer.echo(
+            f"Error: No table files found matching pattern '{table_pattern}' in {table_folder}"
+        )
+        raise typer.Exit(1)
+
+    # Group tables by behavior (extract from filename or header)
+    table_groups = {}
+    for table_file in table_files:
+        try:
+            header_data = pd.read_csv(table_file, nrows=1)
+            behavior_name = header_data["Behavior"][0]
+
+            # Filter by requested behaviors if specified
+            if behaviors and behavior_name not in behaviors:
+                continue
+
+            if behavior_name not in table_groups:
+                table_groups[behavior_name] = []
+            table_groups[behavior_name].append(table_file)
+
+        except (KeyError, pd.errors.EmptyDataError, Exception):
+            typer.echo(f"Warning: Could not read behavior from {table_file}, skipping.")
+            continue
+
+    if not table_groups:
+        typer.echo("Error: No valid behavior tables found to merge.")
+        raise typer.Exit(1)
+
+    try:
+        results = generate_behavior_tables.merge_multiple_behavior_tables(
+            table_groups=table_groups,
+            output_prefix=output_prefix,
+            overwrite=overwrite,
+        )
+
+        typer.echo(f"Successfully merged tables for {len(results)} behaviors:")
+        for behavior_name, (bout_file, bin_file) in results.items():
+            typer.echo(f"  {behavior_name}:")
+            if bout_file:
+                typer.echo(f"    Bout table: {bout_file}")
+            if bin_file:
+                typer.echo(f"    Bin table: {bin_file}")
+
+    except FileExistsError as e:
+        typer.echo(f"Error: {str(e)}")
+        typer.echo("Use --overwrite to force overwrite.")
+        raise typer.Exit(1)
+    except Exception as e:
+        typer.echo(f"Unexpected error: {str(e)}")
+        raise typer.Exit(1)
+
+
 if __name__ == "__main__":
     app()
diff --git a/src/jabs_postprocess/generate_behavior_tables.py b/src/jabs_postprocess/generate_behavior_tables.py
@@ -3,7 +3,13 @@
 from typing import Dict, List, Optional, Tuple
 
 from pathlib import Path
-from jabs_postprocess.utils.project_utils import ClassifierSettings, JabsProject
+import pandas as pd
+from jabs_postprocess.utils.project_utils import (
+    ClassifierSettings,
+    JabsProject,
+    BoutTable,
+    BinTable,
+)
 
 
 def process_behavior_tables(
@@ -111,3 +117,137 @@ def process_multiple_behaviors(
         results.append((bout_path, bin_path))
 
     return results
+
+
+def merge_behavior_tables(
+    input_tables: List[Path],
+    output_prefix: str = "merged_behavior",
+    overwrite: bool = False,
+) -> Tuple[str, str]:
+    """Merge multiple behavior tables for the same behavior.
+
+    Args:
+            input_tables: List of paths to behavior table files to merge
+            output_prefix: Prefix for output filenames
+            overwrite: Whether to overwrite existing files
+
+    Returns:
+            Tuple[str, str]: (merged_bout_table_path, merged_bin_table_path) - Paths to the created files
+
+    Raises:
+            FileNotFoundError: If any input table file doesn't exist
+            ValueError: If tables have different behaviors or incompatible headers
+            FileExistsError: If output files exist and overwrite is False
+    """
+    if not input_tables:
+        raise ValueError("No input tables provided")
+
+    # Validate all input files exist
+    for table_path in input_tables:
+        if not Path(table_path).exists():
+            raise FileNotFoundError(f"Input table not found: {table_path}")
+
+    # Read the first table to determine if it's a bout or bin table and get behavior info
+    first_table = BoutTable.from_file(input_tables[0])
+    behavior_name = first_table.settings.behavior
+    table_type = "bout"
+
+    # Try to determine table type by checking columns
+    if "bout_behavior" in first_table.data.columns:
+        # This is likely a bin table
+        first_table = BinTable.from_file(input_tables[0])
+        table_type = "bin"
+
+    # Load all tables and validate they're compatible
+    tables = []
+    for table_path in input_tables:
+        if table_type == "bout":
+            table = BoutTable.from_file(table_path)
+        else:
+            table = BinTable.from_file(table_path)
+
+        # Validate same behavior
+        if table.settings.behavior != behavior_name:
+            raise ValueError(
+                f"Incompatible behaviors: {behavior_name} vs {table.settings.behavior} in {table_path}"
+            )
+
+        tables.append(table)
+
+    # Merge the tables using the existing combine_data method
+    if table_type == "bout":
+        merged_table = BoutTable.combine_data(tables)
+        output_file = f"{output_prefix}_{behavior_name}_bouts_merged.csv"
+    else:
+        merged_table = BinTable.combine_data(tables)
+        output_file = f"{output_prefix}_{behavior_name}_summaries_merged.csv"
+
+    # Write the merged table
+    merged_table.to_file(output_file, overwrite)
+
+    return (
+        output_file,
+        output_file,
+    )  # Return same file for both since we only merged one type
+
+
+def merge_multiple_behavior_tables(
+    table_groups: Dict[str, List[Path]],
+    output_prefix: str = "merged_behavior",
+    overwrite: bool = False,
+) -> Dict[str, Tuple[str, str]]:
+    """Merge multiple sets of behavior tables grouped by behavior.
+
+    Args:
+            table_groups: Dictionary mapping behavior names to lists of table file paths
+            output_prefix: Prefix for output filenames
+            overwrite: Whether to overwrite existing files
+
+    Returns:
+            Dictionary mapping behavior names to (bout_table_path, bin_table_path) tuples
+
+    Raises:
+            ValueError: If any behavior group is empty
+            FileNotFoundError: If any input table file doesn't exist
+            FileExistsError: If output files exist and overwrite is False
+    """
+    results = {}
+
+    for behavior_name, table_paths in table_groups.items():
+        if not table_paths:
+            raise ValueError(f"No tables provided for behavior: {behavior_name}")
+
+        # Group tables by type (bout vs bin) for this behavior
+        bout_tables = []
+        bin_tables = []
+
+        for table_path in table_paths:
+            data_sample = pd.read_csv(table_path, skiprows=2, nrows=1)
+
+            if data_sample.empty:
+                continue  # Skip empty tables
+
+            # Check if it's a bin table (has bout_behavior column) or bout table
+            full_data = pd.read_csv(table_path, skiprows=2)
+            if "bout_behavior" in full_data.columns:
+                bin_tables.append(table_path)
+            else:
+                bout_tables.append(table_path)
+
+        # Merge bout tables if any exist
+        bout_output = None
+        if bout_tables:
+            bout_output, _ = merge_behavior_tables(
+                bout_tables, f"{output_prefix}_{behavior_name}_bouts", overwrite
+            )
+
+        # Merge bin tables if any exist
+        bin_output = None
+        if bin_tables:
+            bin_output, _ = merge_behavior_tables(
+                bin_tables, f"{output_prefix}_{behavior_name}_summaries", overwrite
+            )
+
+        results[behavior_name] = (bout_output, bin_output)
+
+    return results
diff --git a/src/jabs_postprocess/utils/project_utils.py b/src/jabs_postprocess/utils/project_utils.py
@@ -526,7 +526,7 @@ def data(self):
         return self._data
 
     @classmethod
-    def combine_data(cls, data_list: List(Table)):
+    def combine_data(cls, data_list: List[Table]):
         """Combines multiple data tables together.
 
         Args:
@@ -1138,7 +1138,7 @@ def from_no_prediction(
         return cls(settings, bout_df, video_metadata)
 
     @classmethod
-    def combine_data(cls, data_list: List(Table)):
+    def combine_data(cls, data_list: List[Table]):
         """Combines multiple prediction tables together.
 
         Args:
diff --git a/tests/test_generate_behavior_tables.py b/tests/test_generate_behavior_tables.py
diff --git a/uv.lock b/uv.lock
