14196118657

Committed 01 Apr 2025 12:52PM UTC coverage: 47.634% (+0.8%) from 46.79%

Build # 14196118657

Build Type

Pull #795

github

Specific Base cfb5bd

Committed by

web-flow

Commit Message

Merge 5ba5d5f25 into 58cf64236

Pull Request Pull Request #795: fixes for habituation CW

Run Details

11 of 12 new or added lines in 1 file covered. (91.67%)

1083 existing lines in 22 files now uncovered.

4288 of 9002 relevant lines covered (47.63%)

0.95 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

36.36

/iblrig/raw_data_loaders.py

import json
import logging
import re
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd
from pandas.core.dtypes.concat import union_categoricals

log = logging.getLogger(__name__)
RE_PATTERN_EVENT = re.compile(r'^(\D+\d)_?(.+)$')


def load_task_jsonable(jsonable_file: str | Path, offset: int | None = None) -> tuple[pd.DataFrame, list[Any]]:
    """
    Reads in a task data jsonable file and returns a trials dataframe and a bpod data list.

    Parameters
    ----------
    - jsonable_file (str): full path to jsonable file.
    - offset (int or None): The offset to start reading from (default: None).

    Returns
    -------
    - tuple: A tuple containing:
        - trials_table (pandas.DataFrame): A DataFrame with the trial info in the same format as the Session trials table.
        - bpod_data (list): timing data for each trial
    """
    trials_table = []
    with open(jsonable_file) as f:
        if offset is not None:
            f.seek(offset, 0)
        for line in f:
            trials_table.append(json.loads(line))

    # pop-out the bpod data from the table
    bpod_data = []
    for td in trials_table:
        bpod_data.append(td.pop('behavior_data'))

    trials_table = pd.DataFrame(trials_table)
    return trials_table, bpod_data


def bpod_session_data_to_dataframe(bpod_data: list[dict[str, Any]], trials: int | list[int] | slice | None = None):
    """
    Convert Bpod session data into a single Pandas DataFrame.

    Parameters
    ----------
    bpod_data : list of dict
        A list of dictionaries as returned by load_task_jsonable, where each dictionary contains data for a single trial.
    trials : int, list of int, slice, or None, optional
        Specifies which trials to include in the DataFrame. All trials are included by default.

    Returns
    -------
    pd.DataFrame
        A Pandas DataFrame containing event data from the specified trials, with the following columns:

        *  Time : datetime.timedelta
              timestamp of the event (datetime.timedelta)
        *  Type : str (categorical)
              type of the event (TrialStart, StateStart, InputEvent, etc.)
        *  Trial : int
              index of the trial, zero-based
        *  State : str (categorical)
              name of the state
        *  Event : str (categorical)
              name of the event (only for type InputEvent)
        *  Channel : str (categorical)
              name of the event's channel (only for a subset of InputEvents)
        *  Value : int
              value of the event (only for a subset of InputEvents)
    """
    # define trial index
    if trials is None:
        trials = range(len(bpod_data))
    elif isinstance(trials, int):
        return bpod_trial_data_to_dataframe(bpod_data[trials], trials)
    elif isinstance(trials, slice):
        trials = range(len(bpod_data))[trials]

    # loop over requested trials
    dataframes = []
    for trial in trials:
        dataframes.append(bpod_trial_data_to_dataframe(bpod_data[trial], trial))

    # combine trials into a single dataframe
    categories_type = union_categoricals([df['Type'] for df in dataframes])
    categories_state = union_categoricals([df['State'] for df in dataframes])
    categories_event = union_categoricals([df['Event'] for df in dataframes])
    categories_channel = union_categoricals([df['Channel'] for df in dataframes])
    for df in dataframes:
        df['Type'] = df['Type'].cat.set_categories(categories_type.categories)
        df['State'] = df['State'].cat.set_categories(categories_state.categories)
        df['Event'] = df['Event'].cat.set_categories(categories_event.categories)
        df['Channel'] = df['Channel'].cat.set_categories(categories_channel.categories)
    return pd.concat(dataframes)


def bpod_trial_data_to_dataframe(bpod_trial_data: dict[str, Any], trial: int) -> pd.DataFrame:
    """
    Convert a single Bpod trial's data into a Pandas DataFrame.

    Parameters
    ----------
    bpod_trial_data : dict
        A dictionary containing data for a single trial, including timestamps and events.
    trial : int
        An integer representing the trial index.

    Returns
    -------
    pd.DataFrame
        A Pandas DataFrame containing event data from the specified trial, with the following columns:

        *  Time : datetime.timedelta
              timestamp of the event (datetime.timedelta)
        *  Type : str (categorical)
              type of the event (TrialStart, StateStart, InputEvent, etc.)
        *  Trial : int
              index of the trial, zero-based
        *  State : str (categorical)
              name of the state
        *  Event : str (categorical)
              name of the event (only for type InputEvent)
        *  Channel : str (categorical)
              name of the event's channel (only for a subset of InputEvents)
        *  Value : int
              value of the event (only for a subset of InputEvents)
    """
    trial_start = bpod_trial_data['Trial start timestamp']
    trial_end = bpod_trial_data['Trial end timestamp']

    state_times = bpod_trial_data['States timestamps'].items()
    event_times = bpod_trial_data['Events timestamps'].items()

    # convert bpod data to list of tuples
    event_list = [(0, 'TrialStart', pd.NA, pd.NA)]
    event_list += [(t, 'StateStart', state, pd.NA) for state, times in state_times for t, _ in times if not np.isnan(t)]
    event_list += [(t, 'InputEvent', pd.NA, event) for event, times in event_times for t in times]
    event_list += [(t, 'StateEnd', state, pd.NA) for state, times in state_times for _, t in times if not np.isnan(t)]
    event_list += [(trial_end - trial_start, 'TrialEnd', pd.NA, pd.NA)]
    event_list = sorted(event_list)

    # create dataframe with TimedeltaIndex
    df = pd.DataFrame(data=event_list, columns=['Time', 'Type', 'State', 'Event'])
    df.Time = pd.to_timedelta(df.Time + trial_start, unit='seconds')
    df.set_index('Time', inplace=True)
    df.rename_axis(index=None, inplace=True)

    # cast types
    df['Type'] = df['Type'].astype('category')
    df['State'] = df['State'].astype('category').ffill()
    df['Event'] = df['Event'].astype('category')
    df.insert(2, 'Trial', pd.to_numeric(pd.Series(trial, index=df.index), downcast='unsigned'))

    # deduce channels and values from event names
    df[['Channel', 'Value']] = df['Event'].str.extract(RE_PATTERN_EVENT, expand=True)
    df['Channel'] = df['Channel'].astype('category')
    df['Value'] = df['Value'].replace({'Low': 0, 'High': 1, 'Out': 0, 'In': 1})
    df['Value'] = pd.to_numeric(df['Value'], errors='coerce', downcast='unsigned', dtype_backend='numpy_nullable')

    return df

1	import json	2✔
2	import logging	2✔
3	import re	2✔
4	from pathlib import Path	2✔
5	from typing import Any	2✔
6
7	import numpy as np	2✔
8	import pandas as pd	2✔
9	from pandas.core.dtypes.concat import union_categoricals	2✔
10
11	log = logging.getLogger(__name__)	2✔
12	RE_PATTERN_EVENT = re.compile(r'^(\D+\d)_?(.+)$')	2✔
13
14
15	def load_task_jsonable(jsonable_file: str \| Path, offset: int \| None = None) -> tuple[pd.DataFrame, list[Any]]:	2✔
16	"""
17	Reads in a task data jsonable file and returns a trials dataframe and a bpod data list.
18
19	Parameters
20	----------
21	- jsonable_file (str): full path to jsonable file.
22	- offset (int or None): The offset to start reading from (default: None).
23
24	Returns
25	-------
26	- tuple: A tuple containing:
27	- trials_table (pandas.DataFrame): A DataFrame with the trial info in the same format as the Session trials table.
28	- bpod_data (list): timing data for each trial
29	"""
30	trials_table = []	2✔
31	with open(jsonable_file) as f:	2✔
32	if offset is not None:	2✔
33	f.seek(offset, 0)	2✔
34	for line in f:	2✔
35	trials_table.append(json.loads(line))	2✔
36
37	# pop-out the bpod data from the table
38	bpod_data = []	2✔
39	for td in trials_table:	2✔
40	bpod_data.append(td.pop('behavior_data'))	2✔
41
42	trials_table = pd.DataFrame(trials_table)	2✔
43	return trials_table, bpod_data	2✔
44
45
46	def bpod_session_data_to_dataframe(bpod_data: list[dict[str, Any]], trials: int \| list[int] \| slice \| None = None):	2✔
47	"""
48	Convert Bpod session data into a single Pandas DataFrame.
49
50	Parameters
51	----------
52	bpod_data : list of dict
53	A list of dictionaries as returned by load_task_jsonable, where each dictionary contains data for a single trial.
54	trials : int, list of int, slice, or None, optional
55	Specifies which trials to include in the DataFrame. All trials are included by default.
56
57	Returns
58	-------
59	pd.DataFrame
60	A Pandas DataFrame containing event data from the specified trials, with the following columns:
61
62	* Time : datetime.timedelta
63	timestamp of the event (datetime.timedelta)
64	* Type : str (categorical)
65	type of the event (TrialStart, StateStart, InputEvent, etc.)
66	* Trial : int
67	index of the trial, zero-based
68	* State : str (categorical)
69	name of the state
70	* Event : str (categorical)
71	name of the event (only for type InputEvent)
72	* Channel : str (categorical)
73	name of the event's channel (only for a subset of InputEvents)
74	* Value : int
75	value of the event (only for a subset of InputEvents)
76	"""
77	# define trial index
78	if trials is None:	×
79	trials = range(len(bpod_data))	×
80	elif isinstance(trials, int):	×
UNCOV 81	return bpod_trial_data_to_dataframe(bpod_data[trials], trials)	×
UNCOV 82	elif isinstance(trials, slice):	×
UNCOV 83	trials = range(len(bpod_data))[trials]	×
84
85	# loop over requested trials
UNCOV 86	dataframes = []	×
UNCOV 87	for trial in trials:	×
UNCOV 88	dataframes.append(bpod_trial_data_to_dataframe(bpod_data[trial], trial))	×
89
90	# combine trials into a single dataframe
UNCOV 91	categories_type = union_categoricals([df['Type'] for df in dataframes])	×
UNCOV 92	categories_state = union_categoricals([df['State'] for df in dataframes])	×
UNCOV 93	categories_event = union_categoricals([df['Event'] for df in dataframes])	×
UNCOV 94	categories_channel = union_categoricals([df['Channel'] for df in dataframes])	×
UNCOV 95	for df in dataframes:	×
UNCOV 96	df['Type'] = df['Type'].cat.set_categories(categories_type.categories)	×
UNCOV 97	df['State'] = df['State'].cat.set_categories(categories_state.categories)	×
UNCOV 98	df['Event'] = df['Event'].cat.set_categories(categories_event.categories)	×
99	df['Channel'] = df['Channel'].cat.set_categories(categories_channel.categories)	×
100	return pd.concat(dataframes)	×
101
102
103	def bpod_trial_data_to_dataframe(bpod_trial_data: dict[str, Any], trial: int) -> pd.DataFrame:	2✔
104	"""
105	Convert a single Bpod trial's data into a Pandas DataFrame.
106
107	Parameters
108	----------
109	bpod_trial_data : dict
110	A dictionary containing data for a single trial, including timestamps and events.
111	trial : int
112	An integer representing the trial index.
113
114	Returns
115	-------
116	pd.DataFrame
117	A Pandas DataFrame containing event data from the specified trial, with the following columns:
118
119	* Time : datetime.timedelta
120	timestamp of the event (datetime.timedelta)
121	* Type : str (categorical)
122	type of the event (TrialStart, StateStart, InputEvent, etc.)
123	* Trial : int
124	index of the trial, zero-based
125	* State : str (categorical)
126	name of the state
127	* Event : str (categorical)
128	name of the event (only for type InputEvent)
129	* Channel : str (categorical)
130	name of the event's channel (only for a subset of InputEvents)
131	* Value : int
132	value of the event (only for a subset of InputEvents)
133	"""
UNCOV 134	trial_start = bpod_trial_data['Trial start timestamp']	×
UNCOV 135	trial_end = bpod_trial_data['Trial end timestamp']	×
136
UNCOV 137	state_times = bpod_trial_data['States timestamps'].items()	×
UNCOV 138	event_times = bpod_trial_data['Events timestamps'].items()	×
139
140	# convert bpod data to list of tuples
UNCOV 141	event_list = [(0, 'TrialStart', pd.NA, pd.NA)]	×
UNCOV 142	event_list += [(t, 'StateStart', state, pd.NA) for state, times in state_times for t, _ in times if not np.isnan(t)]	×
UNCOV 143	event_list += [(t, 'InputEvent', pd.NA, event) for event, times in event_times for t in times]	×
UNCOV 144	event_list += [(t, 'StateEnd', state, pd.NA) for state, times in state_times for _, t in times if not np.isnan(t)]	×
UNCOV 145	event_list += [(trial_end - trial_start, 'TrialEnd', pd.NA, pd.NA)]	×
UNCOV 146	event_list = sorted(event_list)	×
147
148	# create dataframe with TimedeltaIndex
149	df = pd.DataFrame(data=event_list, columns=['Time', 'Type', 'State', 'Event'])	×
UNCOV 150	df.Time = pd.to_timedelta(df.Time + trial_start, unit='seconds')	×
UNCOV 151	df.set_index('Time', inplace=True)	×
UNCOV 152	df.rename_axis(index=None, inplace=True)	×
153
154	# cast types
UNCOV 155	df['Type'] = df['Type'].astype('category')	×
UNCOV 156	df['State'] = df['State'].astype('category').ffill()	×
UNCOV 157	df['Event'] = df['Event'].astype('category')	×
UNCOV 158	df.insert(2, 'Trial', pd.to_numeric(pd.Series(trial, index=df.index), downcast='unsigned'))	×
159
160	# deduce channels and values from event names
UNCOV 161	df[['Channel', 'Value']] = df['Event'].str.extract(RE_PATTERN_EVENT, expand=True)	×
UNCOV 162	df['Channel'] = df['Channel'].astype('category')	×
UNCOV 163	df['Value'] = df['Value'].replace({'Low': 0, 'High': 1, 'Out': 0, 'In': 1})	×
UNCOV 164	df['Value'] = pd.to_numeric(df['Value'], errors='coerce', downcast='unsigned', dtype_backend='numpy_nullable')	×
165
166	return df	×

int-brain-lab / iblrig / 14196118657

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous