event_data

  1"""
  2This module defines the EventData class, responsible for managing and storing
  3time-stamped event data collected during an experiment, primarily lick events
  4and motor movements.
  5
  6It utilizes a pandas DataFrame as the data structure to store details
  7about each event, such as trial number, event state (program state for licks/motor direction for motor movemens),
  8duration, and timestamps relative to both the program start and the trial start.
  9
 10It provides methods for initializing the DataFrame, inserting new event rows in a standardized way, and
 11retrieving specific data like lick timestamps for analysis.
 12"""
 13
 14import pandas as pd
 15import logging
 16
 17logger = logging.getLogger(__name__)
 18
 19
 20class EventData:
 21    """
 22    Manages experimental event data using a pandas DataFrame.
 23
 24    This class initializes and maintains a DataFrame to store records of
 25    significant events occurring during an experiment, such as rat licks
 26    or motor movements. It keeps track of lick counts per side and provides
 27    methods to insert new event data rows and query existing data based on
 28    trial number.
 29
 30    Attributes
 31    ----------
 32    - **`side_one_licks`** (*int*): Counter for the total number of licks detected on side one (Port 1) during a given trial.
 33    - **`side_two_licks`** (*int*): Counter for the total number of licks detected on side two (Port 2) during a given trial.
 34    - **`event_dataframe`** (*pd.DataFrame*): The core pandas DataFrame storing event records. Columns include:
 35        - `Trial Number` (*float64*): The 1-indexed trial in which the event occurred.
 36        - `Licked Port` (*float64*): The port number licked (1.0 or 2.0), or NaN for non-lick events.
 37        - `Event Duration` (*float64*): Duration of the event (e.g., lick contact time, motor movement time) in milliseconds. NaN if not applicable.
 38        - `Valve Duration` (*float64*): Duration the valve was open during a lick event (microseconds). NaN if not applicable.
 39        - `Time Stamp` (*float64*): Timestamp relative to the start of the entire program (seconds).
 40        - `Trial Relative Stamp` (*float64*): Timestamp relative to the start of the current trial (seconds).
 41        - `State` (*str*): String describing the experimental state or event type (e.g., "TTC", "SAMPLE", "MOTOR UP").
 42
 43    Methods
 44    -------
 45    - `insert_row_into_df`(...)
 46        Adds a new row representing a single event to the `event_dataframe`.
 47    - `get_lick_timestamps`(...)
 48        Retrieves lists of lick timestamps for a specific trial, separated by port.
 49    """
 50
 51    def __init__(self):
 52        """
 53        Initializes the EventData object.
 54
 55        Sets up the `event_dataframe`, defining columns and expected data types
 56        for storing lick and motor event data. Initializes lick counters
 57        (`side_one_licks`, `side_two_licks`) to zero.
 58        """
 59
 60        self.side_one_licks = 0
 61        self.side_two_licks = 0
 62
 63        self.event_dataframe = pd.DataFrame(
 64            {
 65                "Trial Number": pd.Series(dtype="float64"),
 66                "Licked Port": pd.Series(dtype="float64"),
 67                "Event Duration": pd.Series(dtype="float64"),
 68                "Valve Duration": pd.Series(dtype="float64"),
 69                "Time Stamp": pd.Series(dtype="float64"),
 70                "Trial Relative Stamp": pd.Series(dtype="float64"),
 71                "State": pd.Series(dtype="str"),
 72            }
 73        )
 74
 75        logger.info("Licks dataframe initialized.")
 76
 77    def insert_row_into_df(
 78        self,
 79        trial_num: int,
 80        port: int | None,
 81        duration: float | None,
 82        time_stamp: float,
 83        trial_rel_stamp: float,
 84        state: str,
 85        valve_duration: float | None = None,
 86    ):
 87        """
 88        Inserts a new row representing a single event into the `event_dataframe`.
 89
 90        Appends a new record to the end of the DataFrame. Handles optional parameters
 91        (port, duration, valve_duration) which may not be present for all event types
 92        (e.g., motor movements might not have a 'port').
 93
 94        Parameters
 95        ----------
 96        - **trial_num** (*int*): The 1-indexed trial number during which the event occurred.
 97        - **port** (*int | None*): The port number associated with the event (e.g., 1 or 2 for licks), or None if not applicable.
 98        - **duration** (*float | None*): The duration of the event in milliseconds (e.g., lick contact time), or None if not applicable.
 99        - **time_stamp** (*float*): The timestamp of the event relative to the program start, in seconds.
100        - **trial_rel_stamp** (*float*): The timestamp of the event relative to the start of the current trial, in seconds.
101        - **state** (*str*): A string identifier for the event type (MOTOR) or experimental state (licks) (e.g., "TTC", "SAMPLE", "MOTOR DOWN").
102        - **valve_duration** (*float | None, optional*): The duration the valve was open (microseconds) associated with this event, if applicable. Defaults to None.
103
104        Raises
105        ------
106        - May propagate pandas-related errors if DataFrame operations fail unexpectedly.
107        """
108        # current length in rows of the dataframe, this is where we are going
109        # to insert our next element
110
111        event_df = self.event_dataframe
112        cur_len = len(event_df)
113
114        # optional params, only apply to licks, not motor stamps
115        if port is not None:
116            event_df.loc[cur_len, "Licked Port"] = port
117        if duration is not None:
118            event_df.loc[cur_len, "Event Duration"] = duration
119        if valve_duration is not None:
120            event_df.loc[cur_len, "Valve Duration"] = valve_duration
121
122        event_df.loc[cur_len, "Trial Number"] = trial_num
123        event_df.loc[cur_len, "Time Stamp"] = time_stamp
124        event_df.loc[cur_len, "Trial Relative Stamp"] = trial_rel_stamp
125        event_df.loc[cur_len, "State"] = state
126
127    def get_lick_timestamps(self, logical_trial: int) -> tuple[list, list]:
128        """
129        Retrieves lists of lick timestamps for a specific trial, separated by port.
130
131        Filters the `event_dataframe` based on the provided `logical_trial` number
132        (0-indexed, converted to 1-indexed for filtering) and the 'Licked Port' column.
133        Extracts the 'Time Stamp' (relative to program start) for each port. Used for filling the
134        `RasterizedDataWindow` raster plots.
135
136        Parameters
137        ----------
138        - **logical_trial** (*int*): The 0-indexed trial number for which to retrieve lick timestamps.
139
140        Returns
141        -------
142        - *tuple[list[float], list[float]]*: A tuple containing two lists:
143            - The first list contains timestamps (float, seconds) for licks on Port 1 during the specified trial.
144            - The second list contains timestamps (float, seconds) for licks on Port 2 during the specified trial.
145
146        Raises
147        ------
148        - Propagates potential pandas errors during filtering or data extraction (e.g., `KeyError` if columns are missing). Logs errors if exceptions occur.
149        """
150        # find the 1-indexed trial number, this is how they are stored in the df
151        trial_number = logical_trial + 1
152
153        try:
154            filtered_df_side_one = self.event_dataframe[
155                (self.event_dataframe["Trial Number"] == trial_number)
156                & (self.event_dataframe["Licked Port"].isin([1.0]))
157            ]
158
159            filtered_df_side_two = self.event_dataframe[
160                (self.event_dataframe["Trial Number"] == trial_number)
161                & (self.event_dataframe["Licked Port"].isin([2.0]))
162            ]
163
164            timestamps_side_one = filtered_df_side_one["Time Stamp"].tolist()
165            timestamps_side_two = filtered_df_side_two["Time Stamp"].tolist()
166
167            logger.info(f"Lick timestamps retrieved for trial {trial_number}.")
168
169            return timestamps_side_one, timestamps_side_two
170        except Exception as e:
171            logger.error(f"Error getting lick timestamps for trial {trial_number}: {e}")
172            raise