Skip to content

CoMPASS Level 2 - Utilities

Helper functions for Level 2 analysis and data management.

Overview

This module provides utility functions for:

  • Assigning bouts per session
  • Building phase maps
  • Data organization and manipulation

Functions

compass_labyrinth.compass.level_2.utils

assign_bouts_per_session 

assign_bouts_per_session(
    df: DataFrame, terminal_values=[47], bout_col="Bout_ID"
) -> pd.DataFrame

Assigns bout numbers to a DataFrame based on terminal grid node values, if the specified bout column does not already exist.

Parameters:

  • df (DataFrame) –

    Input DataFrame with 'Grid Number' column.

  • terminal_values (list, default: [47] ) –

    List of grid numbers that signify end-of-bout events.

  • bout_col (str, default: 'Bout_ID' ) –

    Name of the output column to store bout numbers.

Returns:

  • DataFrame –

    DataFrame with the specified bout column added or preserved.

Source code in src/compass_labyrinth/compass/level_2/utils.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def assign_bouts_per_session(
    df: pd.DataFrame,
    terminal_values=[47],
    bout_col="Bout_ID",
) -> pd.DataFrame:
    """
    Assigns bout numbers to a DataFrame based on terminal grid node values,
    if the specified bout column does not already exist.

    Parameters
    -----------
    df : pd.DataFrame
        Input DataFrame with 'Grid Number' column.
    terminal_values : list
        List of grid numbers that signify end-of-bout events.
    bout_col : str
        Name of the output column to store bout numbers.

    Returns
    --------
    pd.DataFrame
        DataFrame with the specified bout column added or preserved.
    """
    if bout_col in df.columns:
        return df  # Already exists; return unchanged

    df = df.reset_index(drop=True).copy()
    bout_num = 1
    in_non_terminal_phase = False
    bout_nums = []

    for i in range(len(df)):
        current = df.iloc[i]["Grid Number"]

        if current not in terminal_values:
            in_non_terminal_phase = True
        else:
            if in_non_terminal_phase:
                in_non_terminal_phase = False  # End bout

        bout_nums.append(bout_num)

        if current in terminal_values and not in_non_terminal_phase:
            if i > 0:
                prev_vals = df.iloc[:i]["Grid Number"]
                last_terminal_idx = prev_vals[prev_vals.isin(terminal_values)].last_valid_index()

                if last_terminal_idx is not None and any(
                    ~df.iloc[last_terminal_idx + 1 : i]["Grid Number"].isin(terminal_values)
                ):
                    if i + 1 < len(df):
                        bout_num += 1

    df[bout_col] = bout_nums
    return df

build_phase_map 

build_phase_map(df: DataFrame, n_phases: int) -> dict
Source code in src/compass_labyrinth/compass/level_2/utils.py 
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def build_phase_map(
    df: pd.DataFrame,
    n_phases: int,
) -> dict:
    phase_map = {}
    sessions = df.Session.unique()
    for sess in sessions:
        df_sess = df[df.Session == sess]
        unique_bouts = df_sess["Bout_ID"].dropna().unique()
        unique_bouts.sort()
        phase_chunks = np.array_split(unique_bouts, n_phases)
        for i, chunk in enumerate(phase_chunks):
            phase_map[(sess, i)] = chunk
    return phase_map