9032957364

Committed 10 May 2024 01:25PM UTC coverage: 48.538% (+1.7%) from 46.79%

Build # 9032957364

Build Type

Pull #643

github

Specific Base 74d2ec

Committed by

web-flow

Commit Message

Merge aebf2c9af into ec2d8e4fe

Pull Request Pull Request #643: 8.19.0

Run Details

377 of 1074 new or added lines in 38 files covered. (35.1%)

977 existing lines in 19 files now uncovered.

3253 of 6702 relevant lines covered (48.54%)

0.97 hits per line

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

87.17

/iblrig/online_plots.py

import datetime
import json
import logging
import time
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from pandas.api.types import CategoricalDtype

import one.alf.io
from iblrig.choiceworld import get_subject_training_info
from iblrig.misc import online_std
from iblrig.raw_data_loaders import load_task_jsonable
from iblutil.util import Bunch

NTRIALS_INIT = 2000
NTRIALS_PLOT = 20  # do not edit - this is used also to enforce the completion criteria
CONTRAST_SET = np.array([0, 1 / 16, 1 / 8, 1 / 4, 1 / 2, 1])
# if the mouse does less than 400 trials in the first 45mins it's disengaged
ENGAGED_CRITIERION = {'secs': 45 * 60, 'trial_count': 400}

log = logging.getLogger(__name__)
sns.set_style('darkgrid')


class DataModel:
    """
    The data model is a pure numpy / pandas container for the choice world task.
    It contains:
    - a psychometrics dataframe that contains the count / choice and response time
    per signed contrast and block contingency
    - a last trials dataframe that contains 20 trials worth of data for the timeline view
    - various counters such as ntrials and water delivered
    """

    task_settings = None
    probability_set = np.array([0.2, 0.5, 0.8])
    ntrials = 0
    ntrials_correct = 0
    ntrials_nan = np.nan
    ntrials_engaged = 0  # trials happening within the first 400s
    percent_correct = np.nan
    percent_error = np.nan
    water_delivered = 0.0
    time_elapsed = 0.0

    def __init__(self, settings_file: Path | str, task_file: Path | str):
        """

        :param task_file: full path to the _iblrig_taskData.raw.jsonable file
        :param settings_file:full path to the _iblrig_taskSettings.raw.json file
        """
        self.session_path = one.alf.files.get_session_path(task_file) or ''
        self.last_trials = pd.DataFrame(
            columns=['correct', 'signed_contrast', 'stim_on', 'play_tone', 'reward_time', 'error_time', 'response_time'],
            index=np.arange(NTRIALS_PLOT),
        )

        if settings_file is None and task_file is not None and task_file.exists():
            settings_file = task_file.parent.joinpath('_iblrig_taskSettings.raw.json')

        if settings_file is not None and settings_file.exists():
            # most of the IBL tasks have a predefined set of probabilities (0.2, 0.5, 0.8), but in the
            # case of the advanced choice world task, the probabilities are defined in the task settings
            with open(settings_file) as fid:
                self.task_settings = json.load(fid)
            self.probability_set = [self.task_settings['PROBABILITY_LEFT']] + self.task_settings.get('BLOCK_PROBABILITY_SET', [])
        else:
            log.warning('Settings file not found - using default settings')

        if task_file is None or not task_file.exists():
            self.psychometrics = pd.DataFrame(
                columns=['count', 'response_time', 'choice', 'response_time_std', 'choice_std'],
                index=pd.MultiIndex.from_product([self.probability_set, np.r_[-np.flipud(CONTRAST_SET[1:]), CONTRAST_SET]]),
            )
            self.psychometrics['count'] = 0
            self.trials_table = pd.DataFrame(columns=['response_time'], index=np.arange(NTRIALS_INIT))
        else:
            trials_table, bpod_data = load_task_jsonable(task_file)
            # here we take the end time of the first trial as reference to avoid factoring in the delay
            self.time_elapsed = bpod_data[-1]['Trial end timestamp'] - bpod_data[0]['Trial end timestamp']
            trials_table['signed_contrast'] = np.sign(trials_table['position']) * trials_table['contrast']
            trials_table['choice'] = trials_table['position'] > 0
            trials_table.loc[~trials_table.trial_correct, 'choice'] = ~trials_table['choice'][~trials_table.trial_correct]
            trials_table['contrast'] = trials_table['contrast'].astype(
                CategoricalDtype(categories=np.unique(np.r_[-CONTRAST_SET, CONTRAST_SET]), ordered=True)
            )
            trials_table['stim_probability_left'] = trials_table['stim_probability_left'].astype(
                CategoricalDtype(categories=self.probability_set, ordered=True)
            )
            self.psychometrics = trials_table.groupby(['stim_probability_left', 'signed_contrast']).agg(
                count=pd.NamedAgg(column='signed_contrast', aggfunc='count'),
                response_time=pd.NamedAgg(column='response_time', aggfunc=np.nanmean),
                choice=pd.NamedAgg(column='choice', aggfunc='mean'),
                response_time_std=pd.NamedAgg(column='response_time', aggfunc=np.nanstd),
                choice_std=pd.NamedAgg(column='choice', aggfunc=np.nanmean),
            )
            self.ntrials = trials_table.shape[0]
            self.ntrials_correct = np.sum(trials_table.trial_correct)
            self.ntrials_nan = self.ntrials if self.ntrials > 0 else np.nan
            self.percent_correct = self.ntrials_correct / self.ntrials_nan * 100
            # agg.water_delivered = trials_table.water_delivered.iloc[-1]
            self.water_delivered = trials_table.reward_amount.sum()
            # init the last trials table
            it = self.last_trials.index[-np.minimum(self.ntrials, NTRIALS_PLOT) :]
            self.last_trials.loc[it, 'correct'] = trials_table.trial_correct.iloc[-NTRIALS_PLOT:].values
            self.last_trials.loc[it, 'signed_contrast'] = trials_table.signed_contrast.iloc[-NTRIALS_PLOT:].values
            self.last_trials.loc[it, 'response_time'] = trials_table.response_time.iloc[-NTRIALS_PLOT:].values
            self.last_trials.loc[it, 'stim_on'] = np.array(
                [bpod_data[i]['States timestamps']['stim_on'][0][0] for i in np.arange(-it.size, 0)]
            )
            self.last_trials.loc[it, 'play_tone'] = np.array(
                [bpod_data[i]['States timestamps']['play_tone'][0][0] for i in np.arange(-it.size, 0)]
            )
            self.last_trials.loc[it, 'reward_time'] = np.array(
                [bpod_data[i]['States timestamps']['reward'][0][0] for i in np.arange(-it.size, 0)]
            )
            self.last_trials.loc[it, 'error_time'] = np.array(
                [bpod_data[i]['States timestamps']['error'][0][0] for i in np.arange(-it.size, 0)]
            )
            # we keep only a single column as buffer
            self.trials_table = trials_table[['response_time']]
        # for the trials plots this is the background image showing green if correct, red if incorrect
        self.rgb_background = np.ones((NTRIALS_PLOT, 1, 3), dtype=np.uint8) * 229
        self.rgb_background[self.last_trials.correct == False, 0, 0] = 255  # noqa
        self.rgb_background[self.last_trials.correct == True, 0, 1] = 255  # noqa
        # keep the last contrasts as a 20 by 2 array
        ileft = np.where(self.last_trials.signed_contrast < 0)[0]  # negative position is left
        iright = np.where(self.last_trials.signed_contrast > 0)[0]
        self.last_contrasts = np.zeros((NTRIALS_PLOT, 2))
        self.last_contrasts[ileft, 0] = np.abs(self.last_trials.signed_contrast[ileft])
        self.last_contrasts[iright, 1] = np.abs(self.last_trials.signed_contrast[iright])

    def update_trial(self, trial_data, bpod_data) -> None:
        # update counters
        self.time_elapsed = bpod_data['Trial end timestamp'] - bpod_data['Bpod start timestamp']
        if self.time_elapsed <= (ENGAGED_CRITIERION['secs']):
            self.ntrials_engaged += 1
        self.ntrials += 1
        self.water_delivered += trial_data.reward_amount
        self.ntrials_correct += trial_data.trial_correct
        signed_contrast = np.sign(trial_data['position']) * trial_data['contrast']
        choice = trial_data.position > 0 if trial_data.trial_correct else trial_data.position < 0
        self.trials_table.at[self.ntrials, 'response_time'] = trial_data.response_time

        # update psychometrics using online statistics method
        indexer = (trial_data.stim_probability_left, signed_contrast)
        if indexer not in self.psychometrics.index:
            self.psychometrics.loc[indexer, :] = np.NaN
            self.psychometrics.loc[indexer, ('count')] = 0
        self.psychometrics.loc[indexer, ('count')] += 1
        self.psychometrics.loc[indexer, ('response_time')], self.psychometrics.loc[indexer, ('response_time_std')] = online_std(
            new_sample=trial_data.response_time,
            new_count=self.psychometrics.loc[indexer, ('count')],
            old_mean=self.psychometrics.loc[indexer, ('response_time')],
            old_std=self.psychometrics.loc[indexer, ('response_time_std')],
        )
        self.psychometrics.loc[indexer, ('choice')], self.psychometrics.loc[indexer, ('choice_std')] = online_std(
            new_sample=float(choice),
            new_count=self.psychometrics.loc[indexer, ('count')],
            old_mean=self.psychometrics.loc[indexer, ('choice')],
            old_std=self.psychometrics.loc[indexer, ('choice_std')],
        )
        # update last trials table
        self.last_trials = self.last_trials.shift(-1)
        i = NTRIALS_PLOT - 1
        self.last_trials.at[i, 'correct'] = trial_data.trial_correct
        self.last_trials.at[i, 'signed_contrast'] = signed_contrast
        self.last_trials.at[i, 'stim_on'] = bpod_data['States timestamps'].get('stim_on', [[np.nan]])[0][0]
        self.last_trials.at[i, 'play_tone'] = bpod_data['States timestamps'].get('play_tone', [[np.nan]])[0][0]
        self.last_trials.at[i, 'reward_time'] = bpod_data['States timestamps'].get('reward', [[np.nan]])[0][0]
        self.last_trials.at[i, 'error_time'] = bpod_data['States timestamps'].get('error', [[np.nan]])[0][0]
        self.last_trials.at[i, 'response_time'] = trial_data.response_time
        # update rgb image
        self.rgb_background = np.roll(self.rgb_background, -1, axis=0)
        self.rgb_background[-1] = np.array([0, 255, 0]) if trial_data.trial_correct else np.array([255, 0, 0])
        # update contrasts
        self.last_contrasts = np.roll(self.last_contrasts, -1, axis=0)
        self.last_contrasts[-1, :] = 0
        self.last_contrasts[-1, int(self.last_trials.signed_contrast.iloc[-1] > 0)] = abs(
            self.last_trials.signed_contrast.iloc[-1]
        )
        self.ntrials_nan = self.ntrials if self.ntrials > 0 else np.nan
        self.percent_correct = self.ntrials_correct / self.ntrials_nan * 100

    def compute_end_session_criteria(self):
        """
        Implements critera to change the color of the figure display, according to the specifications of the task
        :return:
        """
        colour = {'red': '#eb5757', 'green': '#57eb8b', 'yellow': '#ede34e', 'white': '#ffffff'}
        # Within the first part of the session we don't apply response time criterion
        if self.time_elapsed < ENGAGED_CRITIERION['secs']:
            return colour['white']
        # if the mouse has been training for more than 90 minutes subject training too long
        elif self.time_elapsed > (90 * 60):
            return colour['red']
        # the mouse fails to do more than 400 trials in the first 45 mins
        elif self.ntrials_engaged <= ENGAGED_CRITIERION['trial_count']:
            return colour['green']
        # the subject reaction time over the last 20 trials is more than 5 times greater than the overall reaction time
        elif (self.trials_table['response_time'][: self.ntrials].median() * 5) < self.last_trials['response_time'].median():
            return colour['yellow']
        # 90 > time > 45 min and subject's avg response time hasn't significantly decreased
        else:
            return colour['white']


class OnlinePlots:
    """
    Full object to implement the online plots
    Either the object is instantiated in a static mode from an existing jsonable file and it will produce the figure
    >>> oplt = OnlinePlots(task_file)
    Or it can be instantiated empty, and then run on a file during acquisition.
    Use ctrl + Z to interrupt
    >>> OnlinePlots().run(task_file)
    """

    def __init__(self, task_file=None, settings_file=None):
        task_file = Path(task_file) if task_file is not None else None
        settings_file = Path(settings_file) if settings_file is not None else None
        self.data = DataModel(task_file=task_file, settings_file=settings_file)

        # create figure and axes
        h = Bunch({})
        h.fig = plt.figure(constrained_layout=True, figsize=(10, 8))
        self._set_session_string()
        h.fig_title = h.fig.suptitle(f'{self._session_string}')
        nc = 9
        hc = nc // 2
        h.gs = h.fig.add_gridspec(2, nc)
        h.ax_trials = h.fig.add_subplot(h.gs[:, :hc])
        h.ax_psych = h.fig.add_subplot(h.gs[0, hc : nc - 1])
        h.ax_performance = h.fig.add_subplot(h.gs[0, nc - 1])
        h.ax_reaction = h.fig.add_subplot(h.gs[1, hc : nc - 1])
        h.ax_water = h.fig.add_subplot(h.gs[1, nc - 1])

        h.ax_psych.set(title='psychometric curve', xlim=[-1, 1], ylim=[0, 1])
        h.ax_reaction.set(title='reaction times', xlim=[-1, 1], ylim=[0.1, 100], yscale='log', xlabel='signed contrast')
        xticks = np.arange(-1, 1.1, 0.25)
        xticklabels = np.array([f'{x:g}' for x in xticks])
        xticklabels[1::2] = ''
        h.ax_psych.set_xticks(xticks, xticklabels)
        h.ax_reaction.set_xticks(xticks, xticklabels)

        h.ax_trials.set(yticks=[], title='trials timeline', xlim=[-5, 30], xlabel='time (s)')
        h.ax_trials.set_xticks(h.ax_trials.get_xticks(), [''] + h.ax_trials.get_xticklabels()[1::])
        h.ax_performance.set(xticks=[], xlim=[-0.6, 0.6], ylim=[0, 100], title='performance')
        h.ax_water.set(xticks=[], xlim=[-0.6, 0.6], ylim=[0, 1000], title='reward')

        # create psych curves
        h.curve_psych = {}
        h.curve_reaction = {}
        for p in self.data.probability_set:
            h.curve_psych[p] = h.ax_psych.plot(
                self.data.psychometrics.loc[p].index,
                self.data.psychometrics.loc[p]['choice'],
                '.-',
                zorder=10,
                clip_on=False,
                label=f'p = {p}',
            )
            h.curve_reaction[p] = h.ax_reaction.plot(
                self.data.psychometrics.loc[p].index, self.data.psychometrics.loc[p]['response_time'], '.-', label=f'p = {p}'
            )
        h.ax_psych.legend()
        h.ax_reaction.legend()

        # create the two bars on the right side
        h.bar_correct = h.ax_performance.bar(0, self.data.percent_correct, label='correct', color='k')
        h.bar_water = h.ax_water.bar(0, self.data.water_delivered, label='water delivered', color='b')

        # create the trials timeline view in a single axis
        xpos = np.tile([[-3.75, -1.25]], (NTRIALS_PLOT, 1)).T.flatten()
        ypos = np.tile(np.arange(NTRIALS_PLOT), 2)
        h.im_trials = h.ax_trials.imshow(
            self.data.rgb_background, alpha=0.2, extent=[-10, 50, -0.5, NTRIALS_PLOT - 0.5], aspect='auto', origin='lower'
        )
        kwargs = dict(markersize=10, markeredgewidth=2)
        h.lines_trials = {
            'stim_on': h.ax_trials.plot(
                self.data.last_trials.stim_on, np.arange(NTRIALS_PLOT), '|', color='b', **kwargs, label='stimulus on'
            ),
            'reward_time': h.ax_trials.plot(
                self.data.last_trials.reward_time, np.arange(NTRIALS_PLOT), '|', color='g', **kwargs, label='reward'
            ),
            'error_time': h.ax_trials.plot(
                self.data.last_trials.error_time, np.arange(NTRIALS_PLOT), '|', color='r', **kwargs, label='error'
            ),
            'play_tone': h.ax_trials.plot(
                self.data.last_trials.play_tone, np.arange(NTRIALS_PLOT), '|', color='m', **kwargs, label='tone'
            ),
        }
        h.scatter_contrast = h.ax_trials.scatter(
            xpos, ypos, s=250, c=self.data.last_contrasts.T.flatten(), alpha=1, marker='o', vmin=0.0, vmax=1, cmap='Greys'
        )
        h.ax_trials.legend()

        xticks = np.arange(-1, 1.1, 0.25)
        xticklabels = np.array([f'{x:g}' for x in xticks])
        xticklabels[1::2] = ''
        h.ax_psych.set_xticks(xticks, xticklabels)

        self.h = h
        self.update_titles()
        plt.show(block=False)
        plt.draw()

    def update_titles(self):
        protocol = (self.data.task_settings['PYBPOD_PROTOCOL'] if self.data.task_settings else '').replace('_', r'\_')
        spacer = r'\ \ ·\ \ '
        main_title = (
            r'$\mathbf{' + protocol + rf'{spacer}{self.data.ntrials}\ trials{spacer}time\ elapsed:\ '
            rf'{str(datetime.timedelta(seconds=int(self.data.time_elapsed)))}' + r'}$'
        )
        self.h.fig_title.set_text(main_title + '\n' + self._session_string)
        self.h.ax_water.title.set_text(f'total reward\n{self.data.water_delivered:.1f}μL')
        self.h.ax_performance.title.set_text(f'performance\n{self.data.percent_correct:.0f}%')

    def update_trial(self, trial_data, bpod_data):
        """
        This method updates both the data model and the graphics for an upcoming trial
        :param trial data: pandas record
        :param bpod data: dict interpreted from the bpod json dump
        :return:
        """
        self.data.update_trial(trial_data, bpod_data)
        self.update_graphics(pupdate=trial_data.stim_probability_left)

    def update_graphics(self, pupdate: float | None = None):
        background_color = self.data.compute_end_session_criteria()
        h = self.h
        h.fig.set_facecolor(background_color)
        self.update_titles()
        for p in self.data.probability_set:
            if pupdate is not None and p != pupdate:
                continue
            # update psychometric curves
            iok = ~np.isnan(self.data.psychometrics.loc[p]['choice'].values.astype(np.float32))
            xval = self.data.psychometrics.loc[p].index[iok]
            h.curve_psych[p][0].set(xdata=xval, ydata=self.data.psychometrics.loc[p]['choice'][iok])
            h.curve_reaction[p][0].set(xdata=xval, ydata=self.data.psychometrics.loc[p]['response_time'][iok])
            # update the last trials plot
            self.h.im_trials.set_array(self.data.rgb_background)
            for k in ['stim_on', 'reward_time', 'error_time', 'play_tone']:
                h.lines_trials[k][0].set(xdata=self.data.last_trials[k])
            self.h.scatter_contrast.set_array(self.data.last_contrasts.T.flatten())
            # update barplots
            self.h.bar_correct[0].set(height=self.data.percent_correct)
            self.h.bar_water[0].set(height=self.data.water_delivered)

    def _set_session_string(self) -> None:
        if isinstance(self.data.task_settings, dict):
            training_info, _ = get_subject_training_info(
                subject_name=self.data.task_settings['SUBJECT_NAME'],
                task_name=self.data.task_settings['PYBPOD_PROTOCOL'],
                lab=self.data.task_settings['ALYX_LAB'],
            )
            self._session_string = (
                f'subject: {self.data.task_settings["SUBJECT_NAME"]}  ·  '
                f'weight: {self.data.task_settings["SUBJECT_WEIGHT"]}g  ·  '
                f'training phase: {training_info["training_phase"]}  ·  '
                f'stimulus gain: {self.data.task_settings["STIM_GAIN"]}  ·  '
                f'reward amount: {self.data.task_settings["REWARD_AMOUNT_UL"]}µl'
            )
        else:
            self._session_string = ''

    def run(self, task_file: Path | str) -> None:
        """
        This methods is for online use, it will watch for a file in conjunction with an iblrigv8 running task
        :param task_file:
        :return:
        """
        task_file = Path(task_file)
        self._set_session_string()
        self.update_titles()
        self.h.fig.canvas.flush_events()
        self.real_time = Bunch({'fseek': 0, 'time_last_check': 0})
        flag_file = task_file.parent.joinpath('new_trial.flag')

        while True:
            self.h.fig.canvas.draw_idle()
            self.h.fig.canvas.flush_events()
            time.sleep(0.4)
            if not plt.fignum_exists(self.h.fig.number):
                break
            if flag_file.exists():
                trial_data, bpod_data = load_task_jsonable(task_file, offset=self.real_time.fseek)
                new_size = task_file.stat().st_size
                for i in np.arange(len(bpod_data)):
                    self.update_trial(trial_data.iloc[i], bpod_data[i])
                self.real_time.fseek = new_size
                self.real_time.time_last_check = time.time()
                flag_file.unlink()

1	import datetime	2✔
2	import json	2✔
3	import logging	2✔
4	import time	2✔
5	from pathlib import Path	2✔
6
7	import matplotlib.pyplot as plt	2✔
8	import numpy as np	2✔
9	import pandas as pd	2✔
10	import seaborn as sns	2✔
11	from pandas.api.types import CategoricalDtype	2✔
12
13	import one.alf.io	2✔
14	from iblrig.choiceworld import get_subject_training_info	2✔
15	from iblrig.misc import online_std	2✔
16	from iblrig.raw_data_loaders import load_task_jsonable	2✔
17	from iblutil.util import Bunch	2✔
18
19	NTRIALS_INIT = 2000	2✔
20	NTRIALS_PLOT = 20 # do not edit - this is used also to enforce the completion criteria	2✔
21	CONTRAST_SET = np.array([0, 1 / 16, 1 / 8, 1 / 4, 1 / 2, 1])	2✔
22	# if the mouse does less than 400 trials in the first 45mins it's disengaged
23	ENGAGED_CRITIERION = {'secs': 45 * 60, 'trial_count': 400}	2✔
24
25	log = logging.getLogger(__name__)	2✔
26	sns.set_style('darkgrid')	2✔
27
28
29	class DataModel:	2✔
30	"""
31	The data model is a pure numpy / pandas container for the choice world task.
32	It contains:
33	- a psychometrics dataframe that contains the count / choice and response time
34	per signed contrast and block contingency
35	- a last trials dataframe that contains 20 trials worth of data for the timeline view
36	- various counters such as ntrials and water delivered
37	"""
38
39	task_settings = None	2✔
40	probability_set = np.array([0.2, 0.5, 0.8])	2✔
41	ntrials = 0	2✔
42	ntrials_correct = 0	2✔
43	ntrials_nan = np.nan	2✔
44	ntrials_engaged = 0 # trials happening within the first 400s	2✔
45	percent_correct = np.nan	2✔
46	percent_error = np.nan	2✔
47	water_delivered = 0.0	2✔
48	time_elapsed = 0.0	2✔
49
50	def __init__(self, settings_file: Path \| str, task_file: Path \| str):	2✔
51	"""
52
53	:param task_file: full path to the _iblrig_taskData.raw.jsonable file
54	:param settings_file:full path to the _iblrig_taskSettings.raw.json file
55	"""
56	self.session_path = one.alf.files.get_session_path(task_file) or ''	2✔
57	self.last_trials = pd.DataFrame(	2✔
58	columns=['correct', 'signed_contrast', 'stim_on', 'play_tone', 'reward_time', 'error_time', 'response_time'],
59	index=np.arange(NTRIALS_PLOT),
60	)
61
62	if settings_file is None and task_file is not None and task_file.exists():	2✔
63	settings_file = task_file.parent.joinpath('_iblrig_taskSettings.raw.json')	2✔
64
65	if settings_file is not None and settings_file.exists():	2✔
66	# most of the IBL tasks have a predefined set of probabilities (0.2, 0.5, 0.8), but in the
67	# case of the advanced choice world task, the probabilities are defined in the task settings
NEW 68	with open(settings_file) as fid:	×
NEW 69	self.task_settings = json.load(fid)	×
NEW 70	self.probability_set = [self.task_settings['PROBABILITY_LEFT']] + self.task_settings.get('BLOCK_PROBABILITY_SET', [])	×
71	else:
72	log.warning('Settings file not found - using default settings')	2✔
73
74	if task_file is None or not task_file.exists():	2✔
75	self.psychometrics = pd.DataFrame(	2✔
76	columns=['count', 'response_time', 'choice', 'response_time_std', 'choice_std'],
77	index=pd.MultiIndex.from_product([self.probability_set, np.r_[-np.flipud(CONTRAST_SET[1:]), CONTRAST_SET]]),
78	)
79	self.psychometrics['count'] = 0	2✔
80	self.trials_table = pd.DataFrame(columns=['response_time'], index=np.arange(NTRIALS_INIT))	2✔
81	else:
82	trials_table, bpod_data = load_task_jsonable(task_file)	2✔
83	# here we take the end time of the first trial as reference to avoid factoring in the delay
84	self.time_elapsed = bpod_data[-1]['Trial end timestamp'] - bpod_data[0]['Trial end timestamp']	2✔
85	trials_table['signed_contrast'] = np.sign(trials_table['position']) * trials_table['contrast']	2✔
86	trials_table['choice'] = trials_table['position'] > 0	2✔
87	trials_table.loc[~trials_table.trial_correct, 'choice'] = ~trials_table['choice'][~trials_table.trial_correct]	2✔
88	trials_table['contrast'] = trials_table['contrast'].astype(	2✔
89	CategoricalDtype(categories=np.unique(np.r_[-CONTRAST_SET, CONTRAST_SET]), ordered=True)
90	)
91	trials_table['stim_probability_left'] = trials_table['stim_probability_left'].astype(	2✔
92	CategoricalDtype(categories=self.probability_set, ordered=True)
93	)
94	self.psychometrics = trials_table.groupby(['stim_probability_left', 'signed_contrast']).agg(	2✔
95	count=pd.NamedAgg(column='signed_contrast', aggfunc='count'),
96	response_time=pd.NamedAgg(column='response_time', aggfunc=np.nanmean),
97	choice=pd.NamedAgg(column='choice', aggfunc='mean'),
98	response_time_std=pd.NamedAgg(column='response_time', aggfunc=np.nanstd),
99	choice_std=pd.NamedAgg(column='choice', aggfunc=np.nanmean),
100	)
101	self.ntrials = trials_table.shape[0]	2✔
102	self.ntrials_correct = np.sum(trials_table.trial_correct)	2✔
103	self.ntrials_nan = self.ntrials if self.ntrials > 0 else np.nan	2✔
104	self.percent_correct = self.ntrials_correct / self.ntrials_nan * 100	2✔
105	# agg.water_delivered = trials_table.water_delivered.iloc[-1]
106	self.water_delivered = trials_table.reward_amount.sum()	2✔
107	# init the last trials table
108	it = self.last_trials.index[-np.minimum(self.ntrials, NTRIALS_PLOT) :]	2✔
109	self.last_trials.loc[it, 'correct'] = trials_table.trial_correct.iloc[-NTRIALS_PLOT:].values	2✔
110	self.last_trials.loc[it, 'signed_contrast'] = trials_table.signed_contrast.iloc[-NTRIALS_PLOT:].values	2✔
111	self.last_trials.loc[it, 'response_time'] = trials_table.response_time.iloc[-NTRIALS_PLOT:].values	2✔
112	self.last_trials.loc[it, 'stim_on'] = np.array(	2✔
113	[bpod_data[i]['States timestamps']['stim_on'][0][0] for i in np.arange(-it.size, 0)]
114	)
115	self.last_trials.loc[it, 'play_tone'] = np.array(	2✔
116	[bpod_data[i]['States timestamps']['play_tone'][0][0] for i in np.arange(-it.size, 0)]
117	)
118	self.last_trials.loc[it, 'reward_time'] = np.array(	2✔
119	[bpod_data[i]['States timestamps']['reward'][0][0] for i in np.arange(-it.size, 0)]
120	)
121	self.last_trials.loc[it, 'error_time'] = np.array(	2✔
122	[bpod_data[i]['States timestamps']['error'][0][0] for i in np.arange(-it.size, 0)]
123	)
124	# we keep only a single column as buffer
125	self.trials_table = trials_table[['response_time']]	2✔
126	# for the trials plots this is the background image showing green if correct, red if incorrect
127	self.rgb_background = np.ones((NTRIALS_PLOT, 1, 3), dtype=np.uint8) * 229	2✔
128	self.rgb_background[self.last_trials.correct == False, 0, 0] = 255 # noqa	2✔
129	self.rgb_background[self.last_trials.correct == True, 0, 1] = 255 # noqa	2✔
130	# keep the last contrasts as a 20 by 2 array
131	ileft = np.where(self.last_trials.signed_contrast < 0)[0] # negative position is left	2✔
132	iright = np.where(self.last_trials.signed_contrast > 0)[0]	2✔
133	self.last_contrasts = np.zeros((NTRIALS_PLOT, 2))	2✔
134	self.last_contrasts[ileft, 0] = np.abs(self.last_trials.signed_contrast[ileft])	2✔
135	self.last_contrasts[iright, 1] = np.abs(self.last_trials.signed_contrast[iright])	2✔
136
137	def update_trial(self, trial_data, bpod_data) -> None:	2✔
138	# update counters
139	self.time_elapsed = bpod_data['Trial end timestamp'] - bpod_data['Bpod start timestamp']	2✔
140	if self.time_elapsed <= (ENGAGED_CRITIERION['secs']):	2✔
141	self.ntrials_engaged += 1	2✔
142	self.ntrials += 1	2✔
143	self.water_delivered += trial_data.reward_amount	2✔
144	self.ntrials_correct += trial_data.trial_correct	2✔
145	signed_contrast = np.sign(trial_data['position']) * trial_data['contrast']	2✔
146	choice = trial_data.position > 0 if trial_data.trial_correct else trial_data.position < 0	2✔
147	self.trials_table.at[self.ntrials, 'response_time'] = trial_data.response_time	2✔
148
149	# update psychometrics using online statistics method
150	indexer = (trial_data.stim_probability_left, signed_contrast)	2✔
151	if indexer not in self.psychometrics.index:	2✔
NEW 152	self.psychometrics.loc[indexer, :] = np.NaN	×
NEW 153	self.psychometrics.loc[indexer, ('count')] = 0	×
154	self.psychometrics.loc[indexer, ('count')] += 1	2✔
155	self.psychometrics.loc[indexer, ('response_time')], self.psychometrics.loc[indexer, ('response_time_std')] = online_std(	2✔
156	new_sample=trial_data.response_time,
157	new_count=self.psychometrics.loc[indexer, ('count')],
158	old_mean=self.psychometrics.loc[indexer, ('response_time')],
159	old_std=self.psychometrics.loc[indexer, ('response_time_std')],
160	)
161	self.psychometrics.loc[indexer, ('choice')], self.psychometrics.loc[indexer, ('choice_std')] = online_std(	2✔
162	new_sample=float(choice),
163	new_count=self.psychometrics.loc[indexer, ('count')],
164	old_mean=self.psychometrics.loc[indexer, ('choice')],
165	old_std=self.psychometrics.loc[indexer, ('choice_std')],
166	)
167	# update last trials table
168	self.last_trials = self.last_trials.shift(-1)	2✔
169	i = NTRIALS_PLOT - 1	2✔
170	self.last_trials.at[i, 'correct'] = trial_data.trial_correct	2✔
171	self.last_trials.at[i, 'signed_contrast'] = signed_contrast	2✔
172	self.last_trials.at[i, 'stim_on'] = bpod_data['States timestamps'].get('stim_on', [[np.nan]])[0][0]	2✔
173	self.last_trials.at[i, 'play_tone'] = bpod_data['States timestamps'].get('play_tone', [[np.nan]])[0][0]	2✔
174	self.last_trials.at[i, 'reward_time'] = bpod_data['States timestamps'].get('reward', [[np.nan]])[0][0]	2✔
175	self.last_trials.at[i, 'error_time'] = bpod_data['States timestamps'].get('error', [[np.nan]])[0][0]	2✔
176	self.last_trials.at[i, 'response_time'] = trial_data.response_time	2✔
177	# update rgb image
178	self.rgb_background = np.roll(self.rgb_background, -1, axis=0)	2✔
179	self.rgb_background[-1] = np.array([0, 255, 0]) if trial_data.trial_correct else np.array([255, 0, 0])	2✔
180	# update contrasts
181	self.last_contrasts = np.roll(self.last_contrasts, -1, axis=0)	2✔
182	self.last_contrasts[-1, :] = 0	2✔
183	self.last_contrasts[-1, int(self.last_trials.signed_contrast.iloc[-1] > 0)] = abs(	2✔
184	self.last_trials.signed_contrast.iloc[-1]
185	)
186	self.ntrials_nan = self.ntrials if self.ntrials > 0 else np.nan	2✔
187	self.percent_correct = self.ntrials_correct / self.ntrials_nan * 100	2✔
188
189	def compute_end_session_criteria(self):	2✔
190	"""
191	Implements critera to change the color of the figure display, according to the specifications of the task
192	:return:
193	"""
194	colour = {'red': '#eb5757', 'green': '#57eb8b', 'yellow': '#ede34e', 'white': '#ffffff'}	2✔
195	# Within the first part of the session we don't apply response time criterion
196	if self.time_elapsed < ENGAGED_CRITIERION['secs']:	2✔
197	return colour['white']	2✔
198	# if the mouse has been training for more than 90 minutes subject training too long
199	elif self.time_elapsed > (90 * 60):	2✔
UNCOV 200	return colour['red']	×
201	# the mouse fails to do more than 400 trials in the first 45 mins
202	elif self.ntrials_engaged <= ENGAGED_CRITIERION['trial_count']:	2✔
UNCOV 203	return colour['green']	×
204	# the subject reaction time over the last 20 trials is more than 5 times greater than the overall reaction time
205	elif (self.trials_table['response_time'][: self.ntrials].median() * 5) < self.last_trials['response_time'].median():	2✔
206	return colour['yellow']	2✔
207	# 90 > time > 45 min and subject's avg response time hasn't significantly decreased
208	else:
209	return colour['white']	2✔
210
211
212	class OnlinePlots:	2✔
213	"""
214	Full object to implement the online plots
215	Either the object is instantiated in a static mode from an existing jsonable file and it will produce the figure
216	>>> oplt = OnlinePlots(task_file)
217	Or it can be instantiated empty, and then run on a file during acquisition.
218	Use ctrl + Z to interrupt
219	>>> OnlinePlots().run(task_file)
220	"""
221
222	def __init__(self, task_file=None, settings_file=None):	2✔
223	task_file = Path(task_file) if task_file is not None else None	2✔
224	settings_file = Path(settings_file) if settings_file is not None else None	2✔
225	self.data = DataModel(task_file=task_file, settings_file=settings_file)	2✔
226
227	# create figure and axes
228	h = Bunch({})	2✔
229	h.fig = plt.figure(constrained_layout=True, figsize=(10, 8))	2✔
230	self._set_session_string()	2✔
231	h.fig_title = h.fig.suptitle(f'{self._session_string}')	2✔
232	nc = 9	2✔
233	hc = nc // 2	2✔
234	h.gs = h.fig.add_gridspec(2, nc)	2✔
235	h.ax_trials = h.fig.add_subplot(h.gs[:, :hc])	2✔
236	h.ax_psych = h.fig.add_subplot(h.gs[0, hc : nc - 1])	2✔
237	h.ax_performance = h.fig.add_subplot(h.gs[0, nc - 1])	2✔
238	h.ax_reaction = h.fig.add_subplot(h.gs[1, hc : nc - 1])	2✔
239	h.ax_water = h.fig.add_subplot(h.gs[1, nc - 1])	2✔
240
241	h.ax_psych.set(title='psychometric curve', xlim=[-1, 1], ylim=[0, 1])	2✔
242	h.ax_reaction.set(title='reaction times', xlim=[-1, 1], ylim=[0.1, 100], yscale='log', xlabel='signed contrast')	2✔
243	xticks = np.arange(-1, 1.1, 0.25)	2✔
244	xticklabels = np.array([f'{x:g}' for x in xticks])	2✔
245	xticklabels[1::2] = ''	2✔
246	h.ax_psych.set_xticks(xticks, xticklabels)	2✔
247	h.ax_reaction.set_xticks(xticks, xticklabels)	2✔
248
249	h.ax_trials.set(yticks=[], title='trials timeline', xlim=[-5, 30], xlabel='time (s)')	2✔
250	h.ax_trials.set_xticks(h.ax_trials.get_xticks(), [''] + h.ax_trials.get_xticklabels()[1::])	2✔
251	h.ax_performance.set(xticks=[], xlim=[-0.6, 0.6], ylim=[0, 100], title='performance')	2✔
252	h.ax_water.set(xticks=[], xlim=[-0.6, 0.6], ylim=[0, 1000], title='reward')	2✔
253
254	# create psych curves
255	h.curve_psych = {}	2✔
256	h.curve_reaction = {}	2✔
257	for p in self.data.probability_set:	2✔
258	h.curve_psych[p] = h.ax_psych.plot(	2✔
259	self.data.psychometrics.loc[p].index,
260	self.data.psychometrics.loc[p]['choice'],
261	'.-',
262	zorder=10,
263	clip_on=False,
264	label=f'p = {p}',
265	)
266	h.curve_reaction[p] = h.ax_reaction.plot(	2✔
267	self.data.psychometrics.loc[p].index, self.data.psychometrics.loc[p]['response_time'], '.-', label=f'p = {p}'
268	)
269	h.ax_psych.legend()	2✔
270	h.ax_reaction.legend()	2✔
271
272	# create the two bars on the right side
273	h.bar_correct = h.ax_performance.bar(0, self.data.percent_correct, label='correct', color='k')	2✔
274	h.bar_water = h.ax_water.bar(0, self.data.water_delivered, label='water delivered', color='b')	2✔
275
276	# create the trials timeline view in a single axis
277	xpos = np.tile([[-3.75, -1.25]], (NTRIALS_PLOT, 1)).T.flatten()	2✔
278	ypos = np.tile(np.arange(NTRIALS_PLOT), 2)	2✔
279	h.im_trials = h.ax_trials.imshow(	2✔
280	self.data.rgb_background, alpha=0.2, extent=[-10, 50, -0.5, NTRIALS_PLOT - 0.5], aspect='auto', origin='lower'
281	)
282	kwargs = dict(markersize=10, markeredgewidth=2)	2✔
283	h.lines_trials = {	2✔
284	'stim_on': h.ax_trials.plot(
285	self.data.last_trials.stim_on, np.arange(NTRIALS_PLOT), '\|', color='b', **kwargs, label='stimulus on'
286	),
287	'reward_time': h.ax_trials.plot(
288	self.data.last_trials.reward_time, np.arange(NTRIALS_PLOT), '\|', color='g', **kwargs, label='reward'
289	),
290	'error_time': h.ax_trials.plot(
291	self.data.last_trials.error_time, np.arange(NTRIALS_PLOT), '\|', color='r', **kwargs, label='error'
292	),
293	'play_tone': h.ax_trials.plot(
294	self.data.last_trials.play_tone, np.arange(NTRIALS_PLOT), '\|', color='m', **kwargs, label='tone'
295	),
296	}
297	h.scatter_contrast = h.ax_trials.scatter(	2✔
298	xpos, ypos, s=250, c=self.data.last_contrasts.T.flatten(), alpha=1, marker='o', vmin=0.0, vmax=1, cmap='Greys'
299	)
300	h.ax_trials.legend()	2✔
301
302	xticks = np.arange(-1, 1.1, 0.25)	2✔
303	xticklabels = np.array([f'{x:g}' for x in xticks])	2✔
304	xticklabels[1::2] = ''	2✔
305	h.ax_psych.set_xticks(xticks, xticklabels)	2✔
306
307	self.h = h	2✔
308	self.update_titles()	2✔
309	plt.show(block=False)	2✔
310	plt.draw()	2✔
311
312	def update_titles(self):	2✔
313	protocol = (self.data.task_settings['PYBPOD_PROTOCOL'] if self.data.task_settings else '').replace('_', r'\_')	2✔
314	spacer = r'\ \ ·\ \ '	2✔
315	main_title = (	2✔
316	r'$\mathbf{' + protocol + rf'{spacer}{self.data.ntrials}\ trials{spacer}time\ elapsed:\ '
317	rf'{str(datetime.timedelta(seconds=int(self.data.time_elapsed)))}' + r'}$'
318	)
319	self.h.fig_title.set_text(main_title + '\n' + self._session_string)	2✔
320	self.h.ax_water.title.set_text(f'total reward\n{self.data.water_delivered:.1f}μL')	2✔
321	self.h.ax_performance.title.set_text(f'performance\n{self.data.percent_correct:.0f}%')	2✔
322
323	def update_trial(self, trial_data, bpod_data):	2✔
324	"""
325	This method updates both the data model and the graphics for an upcoming trial
326	:param trial data: pandas record
327	:param bpod data: dict interpreted from the bpod json dump
328	:return:
329	"""
330	self.data.update_trial(trial_data, bpod_data)	2✔
331	self.update_graphics(pupdate=trial_data.stim_probability_left)	2✔
332
333	def update_graphics(self, pupdate: float \| None = None):	2✔
334	background_color = self.data.compute_end_session_criteria()	2✔
335	h = self.h	2✔
336	h.fig.set_facecolor(background_color)	2✔
337	self.update_titles()	2✔
338	for p in self.data.probability_set:	2✔
339	if pupdate is not None and p != pupdate:	2✔
340	continue	2✔
341	# update psychometric curves
342	iok = ~np.isnan(self.data.psychometrics.loc[p]['choice'].values.astype(np.float32))	2✔
343	xval = self.data.psychometrics.loc[p].index[iok]	2✔
344	h.curve_psych[p][0].set(xdata=xval, ydata=self.data.psychometrics.loc[p]['choice'][iok])	2✔
345	h.curve_reaction[p][0].set(xdata=xval, ydata=self.data.psychometrics.loc[p]['response_time'][iok])	2✔
346	# update the last trials plot
347	self.h.im_trials.set_array(self.data.rgb_background)	2✔
348	for k in ['stim_on', 'reward_time', 'error_time', 'play_tone']:	2✔
349	h.lines_trials[k][0].set(xdata=self.data.last_trials[k])	2✔
350	self.h.scatter_contrast.set_array(self.data.last_contrasts.T.flatten())	2✔
351	# update barplots
352	self.h.bar_correct[0].set(height=self.data.percent_correct)	2✔
353	self.h.bar_water[0].set(height=self.data.water_delivered)	2✔
354
355	def _set_session_string(self) -> None:	2✔
356	if isinstance(self.data.task_settings, dict):	2✔
UNCOV 357	training_info, _ = get_subject_training_info(	×
358	subject_name=self.data.task_settings['SUBJECT_NAME'],
359	task_name=self.data.task_settings['PYBPOD_PROTOCOL'],
360	lab=self.data.task_settings['ALYX_LAB'],
361	)
UNCOV 362	self._session_string = (	×
363	f'subject: {self.data.task_settings["SUBJECT_NAME"]} · '
364	f'weight: {self.data.task_settings["SUBJECT_WEIGHT"]}g · '
365	f'training phase: {training_info["training_phase"]} · '
366	f'stimulus gain: {self.data.task_settings["STIM_GAIN"]} · '
367	f'reward amount: {self.data.task_settings["REWARD_AMOUNT_UL"]}µl'
368	)
369	else:
370	self._session_string = ''	2✔
371
372	def run(self, task_file: Path \| str) -> None:	2✔
373	"""
374	This methods is for online use, it will watch for a file in conjunction with an iblrigv8 running task
375	:param task_file:
376	:return:
377	"""
UNCOV 378	task_file = Path(task_file)	×
UNCOV 379	self._set_session_string()	×
UNCOV 380	self.update_titles()	×
UNCOV 381	self.h.fig.canvas.flush_events()	×
UNCOV 382	self.real_time = Bunch({'fseek': 0, 'time_last_check': 0})	×
UNCOV 383	flag_file = task_file.parent.joinpath('new_trial.flag')	×
384
UNCOV 385	while True:	×
UNCOV 386	self.h.fig.canvas.draw_idle()	×
UNCOV 387	self.h.fig.canvas.flush_events()	×
UNCOV 388	time.sleep(0.4)	×
UNCOV 389	if not plt.fignum_exists(self.h.fig.number):	×
UNCOV 390	break	×
UNCOV 391	if flag_file.exists():	×
UNCOV 392	trial_data, bpod_data = load_task_jsonable(task_file, offset=self.real_time.fseek)	×
UNCOV 393	new_size = task_file.stat().st_size	×
UNCOV 394	for i in np.arange(len(bpod_data)):	×
UNCOV 395	self.update_trial(trial_data.iloc[i], bpod_data[i])	×
UNCOV 396	self.real_time.fseek = new_size	×
UNCOV 397	self.real_time.time_last_check = time.time()	×
UNCOV 398	flag_file.unlink()	×

int-brain-lab / iblrig / 9032957364

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