10568073180

Committed 26 Aug 2024 10:13PM UTC coverage: 47.538% (+0.7%) from 46.79%

Build # 10568073180

Build Type

Pull #711

github

Specific Base eeff82

Committed by

web-flow

Commit Message

Merge 599c9edfb into ad41db25f

Pull Request Pull Request #711: 8.23.2

Run Details

121 of 135 new or added lines in 8 files covered. (89.63%)

1025 existing lines in 22 files now uncovered.

4084 of 8591 relevant lines covered (47.54%)

0.95 hits per line

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

75.74

/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])  # used as a default if instantiated without settings
PROBABILITY_SET = np.array([0.2, 0.5, 0.8])  # used as a default if instantiated without settings
# 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 = PROBABILITY_SET
    contrast_set = CONTRAST_SET
    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 | None):
        self.session_path = one.alf.files.get_session_path(settings_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 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', [])
            self.contrast_set = self.task_settings.get('CONTRAST_SET', CONTRAST_SET)
            log.info(f'Settings from file: contrast set {self.contrast_set}, probability set {self.probability_set}')
        else:
            log.warning(
                f'Settings file not found - using default probabilities {self.probability_set} and contrasts {CONTRAST_SET}'
            )

        # instantiate the psychometrics dataframe
        signed_contrasts = np.r_[-np.flipud(np.unique(np.abs(self.contrast_set))[1:]), np.unique(np.abs(self.contrast_set))]
        self.psychometrics = pd.DataFrame(
            columns=['count', 'response_time', 'choice', 'response_time_std', 'choice_std'],
            index=pd.MultiIndex.from_product([self.probability_set, signed_contrasts]),
        )
        self.psychometrics['count'] = 0
        self.trials_table = pd.DataFrame(columns=['response_time'], index=np.arange(NTRIALS_INIT))

        # 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):
        """Implement critera to change the color of the figure display, according to the specifications of the task."""
        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(settings)
    Or it can be instantiated empty, and then run on a file during acquisition.
    Use ctrl + Z to interrupt
    >>> OnlinePlots().run(jsonable_file)
    """

    def __init__(self, settings_file=None):
        settings_file = Path(settings_file) if settings_file is not None else None
        self.data = DataModel(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()
        if plt.rcParams['backend'] != 'agg':
            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):
        """
        Update, both, the data model and the graphics for an upcoming trial.

        Parameters
        ----------
        trial_data : pandas.DataFrame
            pandas record
        bpod_data : dict
            doct 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
            if self.data.psychometrics.loc[p]['count'].sum() == 0:
                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:
        self._session_string = ''
        try:
            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'
                )
        except FileNotFoundError:
            # there is a chance that people run this directly on another computer for looking at a session,
            # in which case the iblrig_settings.yaml are not necessarily on the machine
            pass

    def run(self, file_jsonable: Path | str) -> None:
        """
        Watch a jsonable file in conjunction with an iblrigv8 running task (for online use).

        Parameters
        ----------
        file_jsonable : Path or str
            The sessions jsonable file
        """
        file_jsonable = Path(file_jsonable)
        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 = file_jsonable.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(file_jsonable, offset=self.real_time.fseek)
                new_size = file_jsonable.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()

    def display_full_jsonable(self, jsonable_file: Path | str):
        trials_table, bpod_data = load_task_jsonable(jsonable_file)
        # here we take the end time of the first trial as reference to avoid factoring in the delay
        self.data.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.data.probability_set, ordered=True)
        )
        self.data.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.data.ntrials = trials_table.shape[0]
        self.data.ntrials_correct = np.sum(trials_table.trial_correct)
        self.data.ntrials_nan = self.data.ntrials if self.data.ntrials > 0 else np.nan
        self.data.percent_correct = self.data.ntrials_correct / self.data.ntrials_nan * 100
        # agg.water_delivered = trials_table.water_delivered.iloc[-1]
        self.data.water_delivered = trials_table.reward_amount.sum()
        # init the last trials table
        it = self.data.last_trials.index[-np.minimum(self.data.ntrials, NTRIALS_PLOT) :]
        self.data.last_trials.loc[it, 'correct'] = trials_table.trial_correct.iloc[-NTRIALS_PLOT:].values
        self.data.last_trials.loc[it, 'signed_contrast'] = trials_table.signed_contrast.iloc[-NTRIALS_PLOT:].values
        self.data.last_trials.loc[it, 'response_time'] = trials_table.response_time.iloc[-NTRIALS_PLOT:].values
        self.data.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.data.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.data.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.data.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.data.trials_table = trials_table[['response_time']]
        self.update_graphics()

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

int-brain-lab / iblrig / 10568073180

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