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)
