def _get_choice_trajectories(sessions, trans_type, pre_post_trials): '''Evaluates choice trajectories around transitions of specified type. Returns float array of choice trajectories of size (n_transitions, n_trials). Choices are coded such that a choice towards the option which is correct before the transition is 1, the other choice is 0, if the choice trajectory extends past the ends of the blocks before and after the transition analysed, it is padded with nans.''' choice_trajectories = [] n_trans_analysed = 0 n_trials = pre_post_trials[1] - pre_post_trials[0] for session in sessions: blocks = session.blocks selected_transitions = _block_index(blocks)[trans_type] n_trans_analysed +=sum(selected_transitions) start_trials = np.array(blocks['start_trials'][1:])[selected_transitions] # Start trials of blocks following selected transitions. end_trials = np.array(blocks['end_trials'][1:])[selected_transitions] # End trials of blocks following selected transitions. prev_start_trials = np.array(blocks['start_trials'][:-1])[selected_transitions] # Start trials of blocks preceding selected transitions. transition_states = np.array(blocks['transition_states'][:-1])[selected_transitions] # Transition state of blocks following selected transitions. reward_states = np.array(blocks['reward_states'][:-1])[selected_transitions] # Reward state of blocks following selected transitions. for start_trial, end_trial, prev_start_trial, reward_state, transition_state in \ zip(start_trials, end_trials, prev_start_trials, reward_states, transition_states): trial_range = start_trial + np.array(pre_post_trials) if trial_range[0] < prev_start_trial: pad_start = prev_start_trial - trial_range[0] trial_range[0] = prev_start_trial else: pad_start = 0 if trial_range[1] > end_trial: pad_end = trial_range[1] - end_trial trial_range[1] = end_trial else: pad_end = 0 choice_trajectory = session.CTSO['choices'][trial_range[0]:trial_range[1]].astype(bool) choice_trajectory = (choice_trajectory ^ bool(reward_state) ^ bool(transition_state)).astype(float) if pad_start: choice_trajectory = np.hstack((ut.nans(pad_start), choice_trajectory)) if pad_end: choice_trajectory = np.hstack((choice_trajectory, ut.nans(pad_end))) choice_trajectories.append(choice_trajectory) return np.vstack(choice_trajectories)
def _plot_exponential_fit(fit, p_1, pre_post_trials, last_n, col = 'r'): t = np.arange(0,pre_post_trials[1]) p_traj = np.hstack([ut.nans(-pre_post_trials[0]-last_n), np.ones(last_n) * fit['p_0'], \ _exp_choice_traj(fit['tau'], fit['p_0'], p_1, t)]) p.plot(range(pre_post_trials[0], pre_post_trials[1]),p_traj, col, linewidth = 2) p.locator_params(nbins = 4)