12279337432

Committed 11 Dec 2024 03:15PM UTC coverage: 47.031% (+0.2%) from 46.79%

Build # 12279337432

Build Type

Pull #751

github

Specific Base d4edef

Committed by

web-flow

Commit Message

Merge eea51f2f7 into 2f9d65d86

Pull Request Pull Request #751: Fiber trajectory GUI

Run Details

0 of 114 new or added lines in 1 file covered. (0.0%)

1076 existing lines in 22 files now uncovered.

4246 of 9028 relevant lines covered (47.03%)

0.94 hits per line

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

75.85

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

int-brain-lab / iblrig / 12279337432

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