class LinearlyDecreasingN(LinearlyDecreasingAttribute):
memory_decay_rate = traits.Tuple((-1., 0.5), desc="")
memory_decay_rate_time = traits.Float(300, desc="")
def __init__(self, *args, **kwargs):
super(LinearlyDecreasingAssist, self).__init__(*args, **kwargs)
if 'memory_decay_rate' not in self.attrs:
self.attrs.append('memory_decay_rate')
class LinearlyDecreasingHalfLife(LinearlyDecreasingAttribute):
'''
Specific case of LinearlyDecreasingAttribute for a linearly decreasing CLDA half-life
'''
half_life = traits.Tuple((450., 450.),
desc="Initial and final half life for CLDA")
half_life_time = traits.Float(
600, desc="Number of seconds to go from initial to final half life")
def __init__(self, *args, **kwargs):
super(LinearlyDecreasingHalfLife, self).__init__(*args, **kwargs)
if 'half_life' not in self.attrs:
self.attrs.append('half_life')
class LinearlyDecreasingAngAssist(LinearlyDecreasingAttribute):
'''
linearly decreasing angular assist -- for psi and ReHand
'''
rh_assist_level = traits.Tuple(
(0.0, 0.0), desc='level of assist to apply to ang output')
rh_assist_level_time = traits.Float(
600,
desc="Number of seconds to go from initial to minimum assist level")
def __init__(self, *args, **kwargs):
super(LinearlyDecreasingAngAssist, self).__init__(*args, **kwargs)
if 'rh_assist_level' not in self.attrs:
self.attrs.append('rh_assist_level')
class LinearlyDecreasingAssist(LinearlyDecreasingAttribute):
'''
Specific case of LinearlyDecreasingAttribute for a linearly decreasing assist parameter
'''
assist_level = traits.Tuple((0.0, 0.0),
desc="Level of assist to apply to BMI output")
assist_level_time = traits.Float(
600,
desc="Number of seconds to go from initial to minimum assist level")
def __init__(self, *args, **kwargs):
super(LinearlyDecreasingAssist, self).__init__(*args, **kwargs)
if 'assist_level' not in self.attrs:
self.attrs.append('assist_level')
class RedGreen(Sequence, Pygame):
status = dict(
wait = dict(start_trial="pretrial", premature="penalty", stop=None),
pretrial = dict(go="trial", premature="penalty"),
trial = dict(correct="reward", timeout="penalty"),
reward = dict(post_reward="wait"),
penalty = dict(post_penalty="wait"),
)
colors = traits.Array(shape=(2, 3), value=[[255,0,0],[0,255,0]],
desc="Tuple of colors (c1, c2) where c* = [r,g,b] between 0 and 1")
dot_radius = traits.Int(100, desc='dot size')
delay_range = traits.Tuple((0.5, 5.),
desc='delay before switching to second color will be drawn from uniform distribution in this range')
def _while_pretrial(self):
import pygame
self.surf.fill(self.background)
right = [self.next_trial[0] + 1920, self.next_trial[1]]
ts = time.time() - self.start_time
dotsize = (init_dot - self.dot_radius) * (shrinklen - min(ts, shrinklen)) + self.dot_radius
if (np.mod(np.round(ts*1000),freq) < freq/2):
pygame.draw.circle(self.surf, self.colors[0], self.next_trial, int(dotsize))
pygame.draw.circle(self.surf, self.colors[0], right, int(dotsize))
self.flip_wait()
def _while_trial(self):
import pygame
self.surf.fill(self.background)
right = [self.next_trial[0] + 1920, self.next_trial[1]]
ts = time.time() - self.start_time
if (np.mod(np.round(ts*1000),freq) < freq/2):
pygame.draw.circle(self.surf, self.colors[1], self.next_trial, self.dot_radius)
pygame.draw.circle(self.surf, self.colors[1], right, self.dot_radius)
self.flip_wait()
def _start_pretrial(self):
self._wait_time = np.random.rand()*abs(self.delay_range[1]-self.delay_range[0]) + self.delay_range[0]
def _test_correct(self, ts):
return self.event is not None
def _test_go(self, ts):
return ts > self._wait_time + shrinklen
def _test_premature(self, ts):
return self.event is not None
class TestBoundary(Window):
'''
A very simple task that displays a marker at the specified screen locations.
Useful for determining reasonable boundary values for targets.
'''
status = dict(wait=dict(stop=None))
state = "wait"
boundaries = traits.Tuple((-18, 18, -10, 10, -12, 12),
desc="x,y,z boundaries to display")
def __init__(self, **kwargs):
super(TestBoundary, self).__init__(**kwargs)
# Create a small sphere for each of the 6 boundary marks
self.xmin = Sphere(radius=.1, color=(.5, 0, .5, 1))
self.add_model(self.xmin)
self.xmax = Sphere(radius=.1, color=(.5, 0, .5, 1))
self.add_model(self.xmax)
self.ymin = Sphere(radius=.1, color=(.5, 0, .5, 1))
self.add_model(self.ymin)
self.ymax = Sphere(radius=.1, color=(.5, 0, .5, 1))
self.add_model(self.ymax)
self.zmin = Sphere(radius=.1, color=(.5, 0, .5, 1))
self.add_model(self.zmin)
self.zmax = Sphere(radius=.1, color=(.5, 0, .5, 1))
self.add_model(self.zmax)
def _start_wait(self):
self.xmin.translate(self.boundaries[0], 0, 0, reset=True)
self.xmin.attach()
self.xmax.translate(self.boundaries[1], 0, 0, reset=True)
self.xmax.attach()
self.ymin.translate(0, self.boundaries[2], 0, reset=True)
self.ymin.attach()
self.ymax.translate(0, self.boundaries[3], 0, reset=True)
self.ymax.attach()
self.zmin.translate(0, 0, self.boundaries[4], reset=True)
self.zmin.attach()
self.zmax.translate(0, 0, self.boundaries[5], reset=True)
self.zmax.attach()
self.requeue()
def _while_wait(self):
self.draw_world()
class Autostart(traits.HasTraits):
'''
Automatically begins the trial from the wait state,
with a random interval drawn from `rand_start`
'''
rand_start = traits.Tuple((0., 0.), desc="Start interval")
def _start_wait(self):
'''
At the start of the 'wait' state, determine how long to wait before starting the trial
by drawing a sample from the rand_start interval
'''
s, e = self.rand_start
self.wait_time = random.random() * (e - s) + s
super(Autostart, self)._start_wait()
def _test_start_trial(self, ts):
'''
Test if the required random wait time has passed
'''
return ts > self.wait_time and not self.pause
class MouseSpeller(LinearlyDecreasingAssist, BMILoop, Sequence, Window):
## Traits: runtime-configurable parameters
reward_time = traits.Float(.5, desc="Length of juice reward")
target_radius = traits.Float(1.3, desc="Radius of targets in cm")
window_size = traits.Tuple((1920*2, 1080), descr='window size')
hold_time = traits.Float(.2, desc="Length of hold required at targets")
hold_penalty_time = traits.Float(1, desc="Length of penalty time for target hold error")
timeout_time = traits.Float(10, desc="Time allowed to go between targets")
timeout_penalty_time = traits.Float(1, desc="Length of penalty time for timeout error")
max_attempts = traits.Int(10, desc='The number of attempts at a target before\
skipping to the next one')
max_error_penalty = traits.Float(3, desc='Max number of penalties (unrewarded backspaces) for false positive clicks')
background = (0.5, 0.5, 0.5, 1)
state = 'wait'
sequence_generators = ['rand_key_seq_gen', 'mackenzie_soukoreff_corpus']
status = dict(
wait = dict(start_trial="target", stop=None),
target = dict(enter_target="hold", timeout="timeout_penalty", stop=None, false_click="target"),
hold = dict(hold_complete="targ_transition", timeout="timeout_penalty", leave_early="target"),
targ_transition = dict(trial_complete="reward", trial_abort="wait", trial_incomplete="target"),
timeout_penalty = dict(timeout_penalty_end="targ_transition"),
hold_penalty = dict(hold_penalty_end="targ_transition"),
reward = dict(reward_end="wait")
)
trial_end_states = ['reward', 'timeout_penalty', 'hold_penalty']
def __init__(self, *args, **kwargs):
kwargs['instantiate_targets'] = False
self.plant = mouse #plants.CursorPlant() #MouseFakeKF() #plantlist[self.plant_type]
super(MouseSpeller, self).__init__(*args, **kwargs)
## instantiate the keyboard targets
n_target_per_row = [5, 10, 9, 9, 6]
self.targets = dict()
for key in keyboard_spec:
targ = KeyTarget(target_radius=self.target_radius, target_color=np.array([0.25, 0.25, 0.25, 1]))
# targ = VirtualRectangularTarget(target_width=self.target_radius*2, target_color=np.array([0.25, 0.25, 0.25, 1]))
self.targets[key] = targ
p = keyboard_spec[key].pos
p[1] += 2
targ.move_to_position(p)
for model in targ.graphics_models:
self.add_model(model)
# Add graphics models for the plant and targets to the window
if hasattr(self.plant, 'graphics_models'):
for model in self.plant.graphics_models:
self.add_model(model)
self.current_target = None
for attr in self.plant.hdf_attrs:
self.add_dtype(*attr)
self.text_output = ''
def init(self):
self.add_dtype('target', 'f8', (3,))
self.add_dtype('target_index', 'i', (1,))
super(MouseSpeller, self).init()
def _cycle(self):
'''
Calls any update functions necessary and redraws screen. Runs 60x per second.
'''
for key, targ in self.targets.items():
in_targ = targ.pt_inside(self.plant.get_endpoint_pos())
if in_targ and self.plant.click_:
targ.click()
else:
targ.unclick()
if not (self.current_target is None):
self.task_data['target'] = self.current_target.get_position()
else:
self.task_data['target'] = np.ones(3)*np.nan
self.task_data['target_index'] = self.target_index
self.move_plant()
## Save plant status to HDF file
plant_data = self.plant.get_data_to_save()
for key in plant_data:
self.task_data[key] = plant_data[key]
super(MouseSpeller, self)._cycle()
def move_plant(self):
super(MouseSpeller, self).move_plant()
def _test_enter_target(self, ts):
'''
return true if the distance between center of cursor and target is smaller than the cursor radius
'''
ent_targ = self.current_target.pt_inside(self.plant.get_endpoint_pos())
clicked = self.plant.click_
return ent_targ and not clicked
def _test_leave_early(self, ts):
'''
return true if cursor moves outside the exit radius
'''
outside = not self.current_target.pt_inside(self.plant.get_endpoint_pos())
if outside:
self.target_index -= 1
return outside
def _record_char(self, new_char):
new_char = self.targs[self.target_index]
if new_char == 'space': new_char = ' '
if new_char == 'enter': new_char = '\n'
if new_char == 'shift': new_char = ''
if new_char == 'backspace' and len(self.text_output) > 0:
self.text_output = self.text_output[:-1]
elif new_char == 'backspace':
pass
else:
self.text_output += new_char
def _test_hold_complete(self, ts):
hold_complete = self.plant.click_ and self.current_target.pt_inside(self.plant.get_endpoint_pos())
if hold_complete:
self._record_char(self.targs[self.target_index])
return hold_complete
def update_report_stats(self):
'''
see experiment.Experiment.update_report_stats for docs
'''
super(MouseSpeller, self).update_report_stats()
self.reportstats['Text output'] = self.text_output
def _test_timeout(self, ts):
return ts>self.timeout_time
def _test_timeout_penalty_end(self, ts):
return ts>self.timeout_penalty_time
def _test_hold_penalty_end(self, ts):
return ts>self.hold_penalty_time
def _test_trial_complete(self, ts):
return self.target_index == len(self.targs) - 1
def _test_trial_incomplete(self, ts):
return (not self._test_trial_complete(ts)) and (self.tries<self.max_attempts)
def _test_trial_abort(self, ts):
return (not self._test_trial_complete(ts)) and (self.tries==self.max_attempts)
def _test_reward_end(self, ts):
return ts>self.reward_time
def _test_false_click(self, ts):
if self.plant.click_posedge:
for key, targ in self.targets.items():
if targ == self.current_target:
return False
in_targ = targ.pt_inside(self.plant.get_endpoint_pos())
if in_targ:
# add backspace as a target
self.n_false_clicks_this_trial += 1
if self.n_false_clicks_this_trial < self.max_error_penalty:
self.targs.insert(self.target_index, 'backspace')
self.target_index -= 1
self.current_target.reset()
self._record_char(key)
return True
# if the click was outside any of the targets
return False
else:
return False
#### STATE FUNCTIONS ####
def _parse_next_trial(self):
self.targs = list(self.next_trial)
def _start_wait(self):
super(MouseSpeller, self)._start_wait()
self.n_false_clicks_this_trial = 0
self.tries = 0
self.target_index = -1
#hide targets
for key in self.targets:
self.targets[key].reset()
#get target locations for this trial
self._parse_next_trial()
# self.chain_length = len(self.targs)
def _start_target(self):
self.target_index += 1
#move a target to current location (target1 and target2 alternate moving) and set location attribute
target = self.targets[self.targs[self.target_index]]
self.current_target = target
# self.target_location = target.pos
target.cue_trial_start()
def _end_hold(self):
# change current target color to green
self.current_target.reset()
def _start_hold_penalty(self):
#hide targets
for key in self.targets:
self.targets[key].reset()
self.tries += 1
self.target_index = -1
def _start_timeout_penalty(self):
#hide targets
for key in self.targets:
self.targets[key].reset()
self.tries += 1
self.target_index = -1
def _start_targ_transition(self):
self.current_target.reset()
def _start_reward(self):
super(MouseSpeller, self)._start_reward()
self.current_target.cue_trial_end_success()
@staticmethod
def rand_key_seq_gen(length=1000, seq_len=2):
key_sequences = []
keys = keyboard_spec.keys()
n_keys = len(keys)
for k in range(length):
inds = np.random.randint(0, n_keys, seq_len)
key_sequences.append([keys[m] for m in inds])
return key_sequences
@staticmethod
def mackenzie_soukoreff_corpus(length=1000):
trials = []
fh = open('/storage/task_data/MouseSpeller/phrases2.txt')
missing_chars = []
for line in fh:
for char in line:
if char == ' ':
char = 'space'
elif char in ['\n', '\r']:
char = 'enter'
if char in keyboard_spec:
trials.append([char])
elif chr(ord(char) + (ord('a') - ord('A'))) in keyboard_spec:
char = chr(ord(char) + (ord('a') - ord('A')))
trials.append(['shift', char])
elif char not in keyboard_spec:
missing_chars.append(char)
print "missing_chars", missing_chars
return trials
def create_assister(self):
self.assister = FeedbackControllerAssist(mouse_motion_model, style='mixing')
def create_goal_calculator(self):
self.goal_calculator = goal_calculators.ZeroVelocityGoal(self.decoder.ssm)
def get_target_BMI_state(self, *args):
'''
Run the goal calculator to determine the target state of the task
'''
if self.current_target is None:
target_pos = np.zeros(2)
else:
target_pos = self.current_target.get_position()[[0,2]]
opt_click_state = float(self.state == 'hold')
target_pos_state = np.hstack([target_pos, opt_click_state])
data, solution_updated = self.goal_calculator(target_pos_state)
target_state, error = data
return np.tile(np.array(target_state).reshape(-1,1), [1, self.decoder.n_subbins])
class LFP_Mod(BMILoop, Sequence, Window):
background = (0,0,0,1)
plant_visible = traits.Bool(True, desc='Specifies whether entire plant is displayed or just endpoint')
lfp_cursor_rad = traits.Float(.5, desc="length of LFP cursor")
lfp_cursor_color = (.5,0,.5,.75)
lfp_plant_type_options = plantlist.keys()
lfp_plant_type = traits.OptionsList(*plantlist, bmi3d_input_options=plantlist.keys())
window_size = traits.Tuple((1920*2, 1080), desc='window size')
lfp_frac_lims = traits.Tuple((0., 0.35), desc='fraction limits')
xlfp_frac_lims = traits.Tuple((-.7, 1.7), desc = 'x dir fraction limits')
lfp_control_band = traits.Tuple((25, 40), desc='beta power band limits')
lfp_totalpw_band = traits.Tuple((1, 100), desc='total power band limits')
xlfp_control_band = traits.Tuple((0, 5), desc = 'x direction band limits')
n_steps = traits.Int(2, desc='moving average for decoder')
powercap = traits.Float(1, desc="Timeout for total power above this")
zboundaries=(-12,12)
status = dict(
wait = dict(start_trial="lfp_target", stop=None),
lfp_target = dict(enter_lfp_target="lfp_hold", powercap_penalty="powercap_penalty", stop=None),
lfp_hold = dict(leave_early="lfp_target", lfp_hold_complete="reward", powercap_penalty="powercap_penalty"),
powercap_penalty = dict(powercap_penalty_end="lfp_target"),
reward = dict(reward_end="wait")
)
static_states = [] # states in which the decoder is not run
trial_end_states = ['reward']
lfp_cursor_on = ['lfp_target', 'lfp_hold']
#initial state
state = "wait"
#create settable traits
reward_time = traits.Float(.5, desc="Length of juice reward")
lfp_target_rad = traits.Float(3.6, desc="Length of targets in cm")
lfp_hold_time = traits.Float(.2, desc="Length of hold required at lfp targets")
lfp_hold_var = traits.Float(.05, desc="Length of hold variance required at lfp targets")
hold_penalty_time = traits.Float(1, desc="Length of penalty time for target hold error")
powercap_penalty_time = traits.Float(1, desc="Length of penalty time for timeout error")
# max_attempts = traits.Int(10, desc='The number of attempts at a target before\
# skipping to the next one')
session_length = traits.Float(0, desc="Time until task automatically stops. Length of 0 means no auto stop.")
#plant_hide_rate = traits.Float(0.0, desc='If the plant is visible, specifies a percentage of trials where it will be hidden')
lfp_target_color = (123/256.,22/256.,201/256.,.5)
mc_target_color = (1,0,0,.5)
target_index = -1 # Helper variable to keep track of which target to display within a trial
#tries = 0 # Helper variable to keep track of the number of failed attempts at a given trial.
cursor_visible = False # Determines when to hide the cursor.
no_data_count = 0 # Counter for number of missing data frames in a row
sequence_generators = ['lfp_mod_4targ']
def __init__(self, *args, **kwargs):
super(LFP_Mod, self).__init__(*args, **kwargs)
self.cursor_visible = True
print 'INIT FRAC LIMS: ', self.lfp_frac_lims
dec_params = dict(lfp_frac_lims = self.lfp_frac_lims,
xlfp_frac_lims = self.xlfp_frac_lims,
powercap = self.powercap,
zboundaries = self.zboundaries,
lfp_control_band = self.lfp_control_band,
lfp_totalpw_band = self.lfp_totalpw_band,
xlfp_control_band = self.xlfp_control_band,
n_steps = self.n_steps)
self.decoder.filt.init_from_task(**dec_params)
self.decoder.init_from_task(**dec_params)
self.lfp_plant = plantlist[self.lfp_plant_type]
if self.lfp_plant_type == 'inv_cursor_onedimLFP':
print 'MAKE SURE INVERSE GENERATOR IS ON'
self.plant_vis_prev = True
self.current_assist_level = 0
self.learn_flag = False
if hasattr(self.lfp_plant, 'graphics_models'):
for model in self.lfp_plant.graphics_models:
self.add_model(model)
# Instantiate the targets
'''
height and width on kinarm machine are 2.4. Here we make it 2.4/8*12 = 3.6
'''
lfp_target = VirtualSquareTarget(target_radius=self.lfp_target_rad, target_color=self.lfp_target_color)
self.targets = [lfp_target]
# Initialize target location variable
self.target_location_lfp = np.array([-100, -100, -100])
# Declare any plant attributes which must be saved to the HDF file at the _cycle rate
for attr in self.lfp_plant.hdf_attrs:
self.add_dtype(*attr)
def init(self):
self.plant = DummyPlant()
self.add_dtype('lfp_target', 'f8', (3,))
self.add_dtype('target_index', 'i', (1,))
self.add_dtype('powercap_flag', 'i',(1,))
for target in self.targets:
for model in target.graphics_models:
self.add_model(model)
super(LFP_Mod, self).init()
def _cycle(self):
'''
Calls any update functions necessary and redraws screen. Runs 60x per second.
'''
self.task_data['loop_time'] = self.iter_time()
self.task_data['lfp_target'] = self.target_location_lfp.copy()
self.task_data['target_index'] = self.target_index
#self.task_data['internal_decoder_state'] = self.decoder.filt.current_lfp_pos
self.task_data['powercap_flag'] = self.decoder.filt.current_powercap_flag
self.move_plant()
## Save plant status to HDF file, ###ADD BACK
lfp_plant_data = self.lfp_plant.get_data_to_save()
for key in lfp_plant_data:
self.task_data[key] = lfp_plant_data[key]
super(LFP_Mod, self)._cycle()
def move_plant(self):
feature_data = self.get_features()
# Save the "neural features" (e.g. spike counts vector) to HDF file
for key, val in feature_data.items():
self.task_data[key] = val
Bu = None
assist_weight = 0
target_state = np.zeros([self.decoder.n_states, self.decoder.n_subbins])
## Run the decoder
if self.state not in self.static_states:
neural_features = feature_data[self.extractor.feature_type]
self.call_decoder(neural_features, target_state, Bu=Bu, assist_level=assist_weight, feature_type=self.extractor.feature_type)
## Drive the plant to the decoded state, if permitted by the constraints of the plant
self.lfp_plant.drive(self.decoder)
self.task_data['decoder_state'] = decoder_state = self.decoder.get_state(shape=(-1,1))
return decoder_state
def run(self):
'''
See experiment.Experiment.run for documentation.
'''
# Fire up the plant. For virtual/simulation plants, this does little/nothing.
self.lfp_plant.start()
try:
super(LFP_Mod, self).run()
finally:
self.lfp_plant.stop()
##### HELPER AND UPDATE FUNCTIONS ####
def update_cursor_visibility(self):
''' Update cursor visible flag to hide cursor if there has been no good data for more than 3 frames in a row'''
prev = self.cursor_visible
if self.no_data_count < 3:
self.cursor_visible = True
if prev != self.cursor_visible:
self.show_object(self.cursor, show=True)
else:
self.cursor_visible = False
if prev != self.cursor_visible:
self.show_object(self.cursor, show=False)
def update_report_stats(self):
'''
see experiment.Experiment.update_report_stats for docs
'''
super(LFP_Mod, self).update_report_stats()
self.reportstats['Trial #'] = self.calc_trial_num()
self.reportstats['Reward/min'] = np.round(self.calc_events_per_min('reward', 120), decimals=2)
#### TEST FUNCTIONS ####
def _test_powercap_penalty(self, ts):
if self.decoder.filt.current_powercap_flag:
#Turn off power cap flag:
self.decoder.filt.current_powercap_flag = 0
return True
else:
return False
def _test_enter_lfp_target(self, ts):
'''
return true if the distance between center of cursor and target is smaller than the cursor radius in the x and z axis only
'''
cursor_pos = self.lfp_plant.get_endpoint_pos()
dx = np.linalg.norm(cursor_pos[0] - self.target_location_lfp[0])
dz = np.linalg.norm(cursor_pos[2] - self.target_location_lfp[2])
in_targ = False
if dx<= (self.lfp_target_rad/2.) and dz<= (self.lfp_target_rad/2.):
in_targ = True
return in_targ
# #return d <= (self.lfp_target_rad - self.lfp_cursor_rad)
# #If center of cursor enters target at all:
# return d <= (self.lfp_target_rad/2.)
# #New version:
# cursor_pos = self.lfp_plant.get_endpoint_pos()
# d = np.linalg.norm(cursor_pos[2] - self.target_location_lfp[2])
# d <= (self.lfp_target_rad - self.lfp_cursor_rad)
def _test_leave_early(self, ts):
'''
return true if cursor moves outside the exit radius
'''
cursor_pos = self.lfp_plant.get_endpoint_pos()
dx = np.linalg.norm(cursor_pos[0] - self.target_location_lfp[0])
dz = np.linalg.norm(cursor_pos[2] - self.target_location_lfp[2])
out_of_targ = False
if dx > (self.lfp_target_rad/2.) or dz > (self.lfp_target_rad/2.):
out_of_targ = True
#rad = self.lfp_target_rad - self.lfp_cursor_rad
#return d > rad
return out_of_targ
def _test_lfp_hold_complete(self, ts):
return ts>=self.lfp_hold_time_plus_var
# def _test_lfp_timeout(self, ts):
# return ts>self.timeout_time
def _test_powercap_penalty_end(self, ts):
if ts>self.powercap_penalty_time:
self.lfp_plant.turn_on()
return ts>self.powercap_penalty_time
def _test_reward_end(self, ts):
return ts>self.reward_time
def _test_stop(self, ts):
if self.session_length > 0 and (self.get_time() - self.task_start_time) > self.session_length:
self.end_task()
return self.stop
#### STATE FUNCTIONS ####
def _parse_next_trial(self):
self.targs = self.next_trial
def _start_wait(self):
super(LFP_Mod, self)._start_wait()
self.tries = 0
self.target_index = -1
#hide targets
for target in self.targets:
target.hide()
#get target locations for this trial
self._parse_next_trial()
self.chain_length = 1
self.lfp_hold_time_plus_var = self.lfp_hold_time + np.random.uniform(low=-1,high=1)*self.lfp_hold_var
def _start_lfp_target(self):
self.target_index += 1
self.target_index = 0
#only 1 target:
target = self.targets[0]
self.target_location_lfp = self.targs #Just one target.
target.move_to_position(self.target_location_lfp)
target.cue_trial_start()
def _start_lfp_hold(self):
#make next target visible unless this is the final target in the trial
idx = (self.target_index + 1)
if idx < self.chain_length:
target = self.targets[idx % 2]
target.move_to_position(self.targs[idx])
def _end_lfp_hold(self):
# change current target color to green
self.targets[self.target_index % 2].cue_trial_end_success()
def _start_timeout_penalty(self):
#hide targets
for target in self.targets:
target.hide()
self.tries += 1
self.target_index = -1
def _start_reward(self):
super(LFP_Mod, self)._start_reward()
#self.targets[self.target_index % 2].show()
def _start_powercap_penalty(self):
for target in self.targets:
target.hide()
self.lfp_plant.turn_off()
@staticmethod
def lfp_mod_4targ(nblocks=100, boundaries=(-18,18,-12,12), xaxis=-8):
'''Mimics beta modulation task from Kinarm Rig:
In Kinarm rig, the following linear transformations happen:
1. LFP cursor is calculated
2. mapped from fraction limits [0, .35] to [-1, 1] (unit_coordinates)
3. udp sent to kinarm machine and multiplied by 8
4. translated upward in the Y direction by + 2.5
This means, our targets which are at -8, [-0.75, 2.5, 5.75, 9.0]
must be translated down by 2.5 to: -8, [-3.25, 0. , 3.25, 6.5]
then divided by 8: -1, [-0.40625, 0. , 0.40625, 0.8125 ] in unit_coordinates
The radius is 1.2, which is 0.15 in unit_coordinates
Now, we map this to a new system:
- new_zero: (y1+y2) / 2
- new_scale: (y2 - y1) / 2
(([-0.40625, 0. , 0.40625, 0.8125 ]) * new_scale ) + new_zero
new_zero = 0
new_scale = 12
12 * [-0.40625, 0. , 0.40625, 0.8125 ]
= array([-4.875, 0. , 4.875, 9.75 ])
'''
new_zero = (boundaries[3]+boundaries[2]) / 2.
new_scale = (boundaries[3] - boundaries[2]) / 2.
kin_targs = np.array([-0.40625, 0. , 0.40625, 0.8125 ])
lfp_targ_y = (new_scale*kin_targs) + new_zero
for i in range(nblocks):
temp = lfp_targ_y.copy()
np.random.shuffle(temp)
if i==0:
z = temp.copy()
else:
z = np.hstack((z, temp))
#Fixed X axis:
x = np.tile(xaxis,(nblocks*4))
y = np.zeros(nblocks*4)
pairs = np.vstack([x, y, z]).T
return pairs
class ApproachAvoidanceTask(Sequence, Window):
'''
This is for a free-choice task with two targets (left and right) presented to choose from.
The position of the targets may change along the x-axis, according to the target generator,
and each target has a varying probability of reward, also according to the target generator.
The code as it is written is for a joystick.
Notes: want target_index to only write once per trial. if so, can make instructed trials random. else, make new state for instructed trial.
'''
background = (0,0,0,1)
shoulder_anchor = np.array([2., 0., -15.]) # Coordinates of shoulder anchor point on screen
arm_visible = traits.Bool(True, desc='Specifies whether entire arm is displayed or just endpoint')
cursor_radius = traits.Float(.5, desc="Radius of cursor")
cursor_color = (.5,0,.5,1)
joystick_method = traits.Float(1,desc="1: Normal velocity, 0: Position control")
joystick_speed = traits.Float(20, desc="Speed of cursor")
plant_type_options = plantlist.keys()
plant_type = traits.OptionsList(*plantlist, bmi3d_input_options=plantlist.keys())
starting_pos = (5, 0, 5)
# window_size = (1280*2, 1024)
window_size = traits.Tuple((1366*2, 768), desc='window size')
status = dict(
#wait = dict(start_trial="target", stop=None),
wait = dict(start_trial="center", stop=None),
center = dict(enter_center="hold_center", timeout="timeout_penalty", stop=None),
hold_center = dict(leave_early_center = "hold_penalty",hold_center_complete="target", timeout="timeout_penalty", stop=None),
target = dict(enter_targetL="hold_targetL", enter_targetH = "hold_targetH", timeout="timeout_penalty", stop=None),
hold_targetR = dict(leave_early_R="hold_penalty", hold_complete="targ_transition"),
hold_targetL = dict(leave_early_L="hold_penalty", hold_complete="targ_transition"),
targ_transition = dict(trial_complete="check_reward",trial_abort="wait", trial_incomplete="center"),
check_reward = dict(avoid="reward",approach="reward_and_airpuff"),
timeout_penalty = dict(timeout_penalty_end="targ_transition"),
hold_penalty = dict(hold_penalty_end="targ_transition"),
reward = dict(reward_end="wait"),
reward_and_airpuff = dict(reward_and_airpuff_end="wait"),
)
#
target_color = (.5,1,.5,0)
#initial state
state = "wait"
#create settable traits
reward_time_avoid = traits.Float(.2, desc="Length of juice reward for avoid decision")
reward_time_approach_min = traits.Float(.2, desc="Min length of juice for approach decision")
reward_time_approach_max = traits.Float(.8, desc="Max length of juice for approach decision")
target_radius = traits.Float(1.5, desc="Radius of targets in cm")
block_length = traits.Float(100, desc="Number of trials per block")
hold_time = traits.Float(.5, desc="Length of hold required at targets")
hold_penalty_time = traits.Float(1, desc="Length of penalty time for target hold error")
timeout_time = traits.Float(10, desc="Time allowed to go between targets")
timeout_penalty_time = traits.Float(1, desc="Length of penalty time for timeout error")
max_attempts = traits.Int(10, desc='The number of attempts at a target before\
skipping to the next one')
session_length = traits.Float(0, desc="Time until task automatically stops. Length of 0 means no auto stop.")
marker_num = traits.Int(14, desc="The index of the motiontracker marker to use for cursor position")
arm_hide_rate = traits.Float(0.0, desc='If the arm is visible, specifies a percentage of trials where it will be hidden')
target_index = 0 # Helper variable to keep track of whether trial is instructed (1 = 1 choice) or free-choice (2 = 2 choices)
target_selected = 'L' # Helper variable to indicate which target was selected
tries = 0 # Helper variable to keep track of the number of failed attempts at a given trial.
timedout = False # Helper variable to keep track if transitioning from timeout_penalty
reward_counter = 0.0
cursor_visible = False # Determines when to hide the cursor.
no_data_count = 0 # Counter for number of missing data frames in a row
scale_factor = 3.0 #scale factor for converting hand movement to screen movement (1cm hand movement = 3.5cm cursor movement)
starting_dist = 10.0 # starting distance from center target
#color_targets = np.random.randint(2)
color_targets = 1 # 0: yellow low, blue high; 1: blue low, yellow high
stopped_center_hold = False #keep track if center hold was released early
limit2d = 1
color1 = target_colors['purple'] # approach color
color2 = target_colors['lightsteelblue'] # avoid color
reward_color = target_colors['green'] # color of reward bar
airpuff_color = target_colors['red'] # color of airpuff bar
sequence_generators = ['colored_targets_with_probabilistic_reward','block_probabilistic_reward','colored_targets_with_randomwalk_reward','randomwalk_probabilistic_reward']
def __init__(self, *args, **kwargs):
super(ApproachAvoidanceTask, self).__init__(*args, **kwargs)
self.cursor_visible = True
# Add graphics models for the plant and targets to the window
self.plant = plantlist[self.plant_type]
self.plant_vis_prev = True
# Add graphics models for the plant and targets to the window
if hasattr(self.plant, 'graphics_models'):
for model in self.plant.graphics_models:
self.add_model(model)
self.current_pt=np.zeros([3]) #keep track of current pt
self.last_pt=np.zeros([3]) #kee
## Declare cursor
#self.dtype.append(('cursor', 'f8', (3,)))
if 0: #hasattr(self.arm, 'endpt_cursor'):
self.cursor = self.arm.endpt_cursor
else:
self.cursor = Sphere(radius=self.cursor_radius, color=self.cursor_color)
self.add_model(self.cursor)
self.cursor.translate(*self.get_arm_endpoint(), reset=True)
## Instantiate the targets. Target 1 is center target, Target H is target with high probability of reward, Target L is target with low probability of reward.
self.target1 = Sphere(radius=self.target_radius, color=self.target_color) # center target
self.add_model(self.target1)
self.targetR = Sphere(radius=self.target_radius, color=self.target_color) # left target
self.add_model(self.targetH)
self.targetL = Sphere(radius=self.target_radius, color=self.target_color) # right target
self.add_model(self.targetL)
###STOPPED HERE: should define Rect target here and then adapt length during task. Also,
### be sure to change all targetH instantiations to targetR.
# Initialize target location variable.
self.target_location1 = np.array([0,0,0])
self.target_locationR = np.array([-self.starting_dist,0,0])
self.target_locationL = np.array([self.starting_dist,0,0])
self.target1.translate(*self.target_location1, reset=True)
self.targetH.translate(*self.target_locationR, reset=True)
self.targetL.translate(*self.target_locationL, reset=True)
# Initialize colors for high probability and low probability target. Color will not change.
self.targetH.color = self.color_targets*self.color1 + (1 - self.color_targets)*self.color2 # high is magenta if color_targets = 1, juicyorange otherwise
self.targetL.color = (1 - self.color_targets)*self.color1 + self.color_targets*self.color2
#set target colors
self.target1.color = (1,0,0,.5) # center target red
# Initialize target location variable
self.target_location = np.array([0, 0, 0])
# Declare any plant attributes which must be saved to the HDF file at the _cycle rate
for attr in self.plant.hdf_attrs:
self.add_dtype(*attr)
def init(self):
self.add_dtype('targetR', 'f8', (3,))
self.add_dtype('targetL','f8', (3,))
self.add_dtype('reward_scheduleR','f8', (1,))
self.add_dtype('reward_scheduleL','f8', (1,))
self.add_dtype('target_index', 'i', (1,))
super(ApproachAvoidanceTask, self).init()
self.trial_allocation = np.zeros(1000)
def _cycle(self):
''' Calls any update functions necessary and redraws screen. Runs 60x per second. '''
## Run graphics commands to show/hide the arm if the visibility has changed
if self.plant_type != 'cursor_14x14':
if self.arm_visible != self.arm_vis_prev:
self.arm_vis_prev = self.arm_visible
self.show_object(self.arm, show=self.arm_visible)
self.move_arm()
#self.move_plant()
## Save plant status to HDF file
plant_data = self.plant.get_data_to_save()
for key in plant_data:
self.task_data[key] = plant_data[key]
self.update_cursor()
if self.plant_type != 'cursor_14x14':
self.task_data['joint_angles'] = self.get_arm_joints()
super(ApproachAvoidanceTask, self)._cycle()
## Plant functions
def get_cursor_location(self):
# arm returns it's position as if it was anchored at the origin, so have to translate it to the correct place
return self.get_arm_endpoint()
def get_arm_endpoint(self):
return self.plant.get_endpoint_pos()
def set_arm_endpoint(self, pt, **kwargs):
self.plant.set_endpoint_pos(pt, **kwargs)
def set_arm_joints(self, angles):
self.arm.set_intrinsic_coordinates(angles)
def get_arm_joints(self):
return self.arm.get_intrinsic_coordinates()
def update_cursor(self):
'''
Update the cursor's location and visibility status.
'''
pt = self.get_cursor_location()
self.update_cursor_visibility()
if pt is not None:
self.move_cursor(pt)
def move_cursor(self, pt):
''' Move the cursor object to the specified 3D location. '''
# if not hasattr(self.arm, 'endpt_cursor'):
self.cursor.translate(*pt[:3],reset=True)
##
##### HELPER AND UPDATE FUNCTIONS ####
def move_arm(self):
''' Returns the 3D coordinates of the cursor. For manual control, uses
joystick data. If no joystick data available, returns None'''
pt = self.joystick.get()
if len(pt) > 0:
pt = pt[-1][0]
pt[0]=1-pt[0]; #Switch L / R axes
calib = [0.497,0.517] #Sometimes zero point is subject to drift this is the value of the incoming joystick when at 'rest'
if self.joystick_method==0:
pos = np.array([(pt[0]-calib[0]), 0, calib[1]-pt[1]])
pos[0] = pos[0]*36
pos[2] = pos[2]*24
self.current_pt = pos
elif self.joystick_method==1:
vel=np.array([(pt[0]-calib[0]), 0, calib[1]-pt[1]])
epsilon = 2*(10**-2) #Define epsilon to stabilize cursor movement
if sum((vel)**2) > epsilon:
self.current_pt=self.last_pt+self.joystick_speed*vel*(1/60) #60 Hz update rate, dt = 1/60
else:
self.current_pt = self.last_pt
if self.current_pt[0] < -25: self.current_pt[0] = -25
if self.current_pt[0] > 25: self.current_pt[0] = 25
if self.current_pt[-1] < -14: self.current_pt[-1] = -14
if self.current_pt[-1] > 14: self.current_pt[-1] = 14
self.set_arm_endpoint(self.current_pt)
self.last_pt = self.current_pt.copy()
def convert_to_cm(self, val):
return val/10.0
def update_cursor_visibility(self):
''' Update cursor visible flag to hide cursor if there has been no good data for more than 3 frames in a row'''
prev = self.cursor_visible
if self.no_data_count < 3:
self.cursor_visible = True
if prev != self.cursor_visible:
self.show_object(self.cursor, show=True)
self.requeue()
else:
self.cursor_visible = False
if prev != self.cursor_visible:
self.show_object(self.cursor, show=False)
self.requeue()
def calc_n_successfultrials(self):
trialendtimes = np.array([state[1] for state in self.state_log if state[0]=='check_reward'])
return len(trialendtimes)
def calc_n_rewards(self):
rewardtimes = np.array([state[1] for state in self.state_log if state[0]=='reward'])
return len(rewardtimes)
def calc_trial_num(self):
'''Calculates the current trial count: completed + aborted trials'''
trialtimes = [state[1] for state in self.state_log if state[0] in ['wait']]
return len(trialtimes)-1
def calc_targetH_num(self):
'''Calculates the total number of times the high-value target was selected'''
trialtimes = [state[1] for state in self.state_log if state[0] in ['hold_targetH']]
return len(trialtimes) - 1
def calc_rewards_per_min(self, window):
'''Calculates the Rewards/min for the most recent window of specified number of seconds in the past'''
rewardtimes = np.array([state[1] for state in self.state_log if state[0]=='reward'])
if (self.get_time() - self.task_start_time) < window:
divideby = (self.get_time() - self.task_start_time)/sec_per_min
else:
divideby = window/sec_per_min
return np.sum(rewardtimes >= (self.get_time() - window))/divideby
def calc_success_rate(self, window):
'''Calculates the rewarded trials/initiated trials for the most recent window of specified length in sec'''
trialtimes = np.array([state[1] for state in self.state_log if state[0] in ['reward', 'timeout_penalty', 'hold_penalty']])
rewardtimes = np.array([state[1] for state in self.state_log if state[0]=='reward'])
if len(trialtimes) == 0:
return 0.0
else:
return float(np.sum(rewardtimes >= (self.get_time() - window)))/np.sum(trialtimes >= (self.get_time() - window))
def update_report_stats(self):
'''Function to update any relevant report stats for the task. Values are saved in self.reportstats,
an ordered dictionary. Keys are strings that will be displayed as the label for the stat in the web interface,
values can be numbers or strings. Called every time task state changes.'''
super(ApproachAvoidanceTask, self).update_report_stats()
self.reportstats['Trial #'] = self.calc_trial_num()
self.reportstats['Reward/min'] = np.round(self.calc_rewards_per_min(120),decimals=2)
self.reportstats['High-value target selections'] = self.calc_targetH_num()
#self.reportstats['Success rate'] = str(np.round(self.calc_success_rate(120)*100.0,decimals=2)) + '%'
start_time = self.state_log[0][1]
rewardtimes=np.array([state[1] for state in self.state_log if state[0]=='reward'])
if len(rewardtimes):
rt = rewardtimes[-1]-start_time
else:
rt= np.float64("0.0")
sec = str(np.int(np.mod(rt,60)))
if len(sec) < 2:
sec = '0'+sec
self.reportstats['Time Of Last Reward'] = str(np.int(np.floor(rt/60))) + ':' + sec
#### TEST FUNCTIONS ####
def _test_enter_center(self, ts):
#return true if the distance between center of cursor and target is smaller than the cursor radius
d = np.sqrt((self.cursor.xfm.move[0]-self.target_location1[0])**2 + (self.cursor.xfm.move[1]-self.target_location1[1])**2 + (self.cursor.xfm.move[2]-self.target_location1[2])**2)
#print 'TARGET SELECTED', self.target_selected
return d <= self.target_radius - self.cursor_radius
def _test_enter_targetL(self, ts):
if self.target_index == 1 and self.LH_target_on[0]==0:
#return false if instructed trial and this target is not on
return False
else:
#return true if the distance between center of cursor and target is smaller than the cursor radius
d = np.sqrt((self.cursor.xfm.move[0]-self.target_locationL[0])**2 + (self.cursor.xfm.move[1]-self.target_locationL[1])**2 + (self.cursor.xfm.move[2]-self.target_locationL[2])**2)
self.target_selected = 'L'
#print 'TARGET SELECTED', self.target_selected
return d <= self.target_radius - self.cursor_radius
def _test_enter_targetH(self, ts):
if self.target_index ==1 and self.LH_target_on[1]==0:
return False
else:
#return true if the distance between center of cursor and target is smaller than the cursor radius
d = np.sqrt((self.cursor.xfm.move[0]-self.target_locationH[0])**2 + (self.cursor.xfm.move[1]-self.target_locationH[1])**2 + (self.cursor.xfm.move[2]-self.target_locationH[2])**2)
self.target_selected = 'H'
#print 'TARGET SELECTED', self.target_selected
return d <= self.target_radius - self.cursor_radius
def _test_leave_early_center(self, ts):
# return true if cursor moves outside the exit radius (gives a bit of slack around the edge of target once cursor is inside)
d = np.sqrt((self.cursor.xfm.move[0]-self.target_location1[0])**2 + (self.cursor.xfm.move[1]-self.target_location1[1])**2 + (self.cursor.xfm.move[2]-self.target_location1[2])**2)
rad = self.target_radius - self.cursor_radius
return d > rad
def _test_leave_early_L(self, ts):
# return true if cursor moves outside the exit radius (gives a bit of slack around the edge of target once cursor is inside)
d = np.sqrt((self.cursor.xfm.move[0]-self.target_locationL[0])**2 + (self.cursor.xfm.move[1]-self.target_locationL[1])**2 + (self.cursor.xfm.move[2]-self.target_locationL[2])**2)
rad = self.target_radius - self.cursor_radius
return d > rad
def _test_leave_early_H(self, ts):
# return true if cursor moves outside the exit radius (gives a bit of slack around the edge of target once cursor is inside)
d = np.sqrt((self.cursor.xfm.move[0]-self.target_locationH[0])**2 + (self.cursor.xfm.move[1]-self.target_locationH[1])**2 + (self.cursor.xfm.move[2]-self.target_locationH[2])**2)
rad = self.target_radius - self.cursor_radius
return d > rad
def _test_hold_center_complete(self, ts):
return ts>=self.hold_time
def _test_hold_complete(self, ts):
return ts>=self.hold_time
def _test_timeout(self, ts):
return ts>self.timeout_time
def _test_timeout_penalty_end(self, ts):
return ts>self.timeout_penalty_time
def _test_hold_penalty_end(self, ts):
return ts>self.hold_penalty_time
def _test_trial_complete(self, ts):
#return self.target_index==self.chain_length-1
return not self.timedout
def _test_trial_incomplete(self, ts):
return (not self._test_trial_complete(ts)) and (self.tries<self.max_attempts)
def _test_trial_abort(self, ts):
return (not self._test_trial_complete(ts)) and (self.tries==self.max_attempts)
def _test_yes_reward(self,ts):
if self.target_selected == 'H':
#reward_assigned = self.targs[0,1]
reward_assigned = self.rewardH
else:
#reward_assigned = self.targs[1,1]
reward_assigned = self.rewardL
if self.reward_SmallLarge==1:
self.reward_time = reward_assigned*self.reward_time_large + (1 - reward_assigned)*self.reward_time_small # update reward time if using Small/large schedule
reward_assigned = 1 # always rewarded
return bool(reward_assigned)
def _test_no_reward(self,ts):
if self.target_selected == 'H':
#reward_assigned = self.targs[0,1]
reward_assigned = self.rewardH
else:
#reward_assigned = self.targs[1,1]
reward_assigned = self.rewardL
if self.reward_SmallLarge==True:
self.reward_time = reward_assigned*self.reward_time_large + (1 - reward_assigned)*self.reward_time_small # update reward time if using Small/large schedule
reward_assigned = 1 # always rewarded
return bool(not reward_assigned)
def _test_reward_end(self, ts):
time_ended = (ts > self.reward_time)
self.reward_counter = self.reward_counter + 1
return time_ended
def _test_stop(self, ts):
if self.session_length > 0 and (time.time() - self.task_start_time) > self.session_length:
self.end_task()
return self.stop
#### STATE FUNCTIONS ####
def show_object(self, obj, show=False):
'''
Show or hide an object
'''
if show:
obj.attach()
else:
obj.detach()
self.requeue()
def _start_wait(self):
super(ApproachAvoidanceTask, self)._start_wait()
self.tries = 0
self.target_index = 0 # indicator for instructed or free-choice trial
#hide targets
self.show_object(self.target1, False)
self.show_object(self.targetL, False)
self.show_object(self.targetH, False)
#get target positions and reward assignments for this trial
self.targs = self.next_trial
if self.plant_type != 'cursor_14x14' and np.random.rand() < self.arm_hide_rate:
self.arm_visible = False
else:
self.arm_visible = True
#self.chain_length = self.targs.shape[0] #Number of sequential targets in a single trial
#self.task_data['target'] = self.target_locationH.copy()
assign_reward = np.random.randint(0,100,size=2)
self.rewardH = np.greater(self.targs[0,1],assign_reward[0])
#print 'high value target reward prob', self.targs[0,1]
self.rewardL = np.greater(self.targs[1,1],assign_reward[1])
#print 'TARGET GENERATOR', self.targs[0,]
self.task_data['targetH'] = self.targs[0,].copy()
self.task_data['reward_scheduleH'] = self.rewardH.copy()
self.task_data['targetL'] = self.targs[1,].copy()
self.task_data['reward_scheduleL'] = self.rewardL.copy()
self.requeue()
def _start_center(self):
#self.target_index += 1
self.show_object(self.target1, True)
self.show_object(self.cursor, True)
# Third argument in self.targs determines if target is on left or right
# First argument in self.targs determines if location is offset to farther distances
offsetH = (2*self.targs[0,2] - 1)*(self.starting_dist + self.location_offset_allowed*self.targs[0,0]*4.0)
moveH = np.array([offsetH,0,0])
offsetL = (2*self.targs[1,2] - 1)*(self.starting_dist + self.location_offset_allowed*self.targs[1,0]*4.0)
moveL = np.array([offsetL,0,0])
self.targetL.translate(*moveL, reset=True)
#self.targetL.move_to_position(*moveL, reset=True)
##self.targetL.translate(*self.targs[self.target_index], reset=True)
self.show_object(self.targetL, True)
self.target_locationL = self.targetL.xfm.move
self.targetH.translate(*moveH, reset=True)
#self.targetR.move_to_position(*moveR, reset=True)
##self.targetR.translate(*self.targs[self.target_index], reset=True)
self.show_object(self.targetH, True)
self.target_locationH = self.targetH.xfm.move
# Insert instructed trials within free choice trials
if self.trial_allocation[self.calc_trial_num()] == 1:
#if (self.calc_trial_num() % 10) < (self.percentage_instructed_trials/10):
self.target_index = 1 # instructed trial
leftright_coinflip = np.random.randint(0,2)
if leftright_coinflip == 0:
self.show_object(self.targetL, False)
self.LH_target_on = (0, 1)
else:
self.show_object(self.targetH, False)
self.LR_coinflip = 0
self.LH_target_on = (1, 0)
else:
self.target_index = 2 # free-choice trial
self.cursor_visible = True
self.task_data['target_index'] = self.target_index
self.requeue()
def _start_target(self):
#self.target_index += 1
#move targets to current location and set location attribute. Target1 (center target) position does not change.
self.show_object(self.target1, False)
#self.target_location1 = self.target1.xfm.move
self.show_object(self.cursor, True)
self.update_cursor()
self.requeue()
def _start_hold_center(self):
self.show_object(self.target1, True)
self.timedout = False
self.requeue()
def _start_hold_targetL(self):
#make next target visible unless this is the final target in the trial
#if 1 < self.chain_length:
#self.targetL.translate(*self.targs[self.target_index+1], reset=True)
# self.show_object(self.targetL, True)
# self.requeue()
self.show_object(self.targetL, True)
self.timedout = False
self.requeue()
def _start_hold_targetH(self):
#make next target visible unless this is the final target in the trial
#if 1 < self.chain_length:
#self.targetR.translate(*self.targs[self.target_index+1], reset=True)
# self.show_object(self.targetR, True)
# self.requeue()
self.show_object(self.targetH, True)
self.timedout = False
self.requeue()
def _end_hold_center(self):
self.target1.radius = 0.7*self.target_radius # color target green
def _end_hold_targetL(self):
self.targetL.color = (0,1,0,0.5) # color target green
def _end_hold_targetH(self):
self.targetH.color = (0,1,0,0.5) # color target green
def _start_hold_penalty(self):
#hide targets
self.show_object(self.target1, False)
self.show_object(self.targetL, False)
self.show_object(self.targetH, False)
self.timedout = True
self.requeue()
self.tries += 1
#self.target_index = -1
def _start_timeout_penalty(self):
#hide targets
self.show_object(self.target1, False)
self.show_object(self.targetL, False)
self.show_object(self.targetH, False)
self.timedout = True
self.requeue()
self.tries += 1
#self.target_index = -1
def _start_targ_transition(self):
#hide targets
self.show_object(self.target1, False)
self.show_object(self.targetL, False)
self.show_object(self.targetH, False)
self.requeue()
def _start_check_reward(self):
#hide targets
self.show_object(self.target1, False)
self.show_object(self.targetL, False)
self.show_object(self.targetH, False)
self.requeue()
def _start_reward(self):
#super(ApproachAvoidanceTask, self)._start_reward()
if self.target_selected == 'L':
self.show_object(self.targetL, True)
#reward_assigned = self.targs[1,1]
else:
self.show_object(self.targetH, True)
#reward_assigned = self.targs[0,1]
#self.reward_counter = self.reward_counter + float(reward_assigned)
self.requeue()
@staticmethod
def colored_targets_with_probabilistic_reward(length=1000, boundaries=(-18,18,-10,10,-15,15),reward_high_prob=80,reward_low_prob=40):
"""
Generator should return array of ntrials x 2 x 3. The second dimension is for each target.
For example, first is the target with high probability of reward, and the second
entry is for the target with low probability of reward. The third dimension holds three variables indicating
position offset (yes/no), reward probability (fixed in this case), and location (binary returned where the
ouput indicates either left or right).
UPDATE: CHANGED SO THAT THE SECOND DIMENSION CARRIES THE REWARD PROBABILITY RATHER THAN THE REWARD SCHEDULE
"""
position_offsetH = np.random.randint(2,size=(1,length))
position_offsetL = np.random.randint(2,size=(1,length))
location_int = np.random.randint(2,size=(1,length))
# coin flips for reward schedules, want this to be elementwise comparison
#assign_rewardH = np.random.randint(0,100,size=(1,length))
#assign_rewardL = np.random.randint(0,100,size=(1,length))
high_prob = reward_high_prob*np.ones((1,length))
low_prob = reward_low_prob*np.ones((1,length))
#reward_high = np.greater(high_prob,assign_rewardH)
#reward_low = np.greater(low_prob,assign_rewardL)
pairs = np.zeros([length,2,3])
pairs[:,0,0] = position_offsetH
#pairs[:,0,1] = reward_high
pairs[:,0,1] = high_prob
pairs[:,0,2] = location_int
pairs[:,1,0] = position_offsetL
#pairs[:,1,1] = reward_low
pairs[:,1,1] = low_prob
pairs[:,1,2] = 1 - location_int
return pairs
@staticmethod
def block_probabilistic_reward(length=1000, boundaries=(-18,18,-10,10,-15,15),reward_high_prob=80,reward_low_prob=40):
pairs = colored_targets_with_probabilistic_reward(length=length, boundaries=boundaries,reward_high_prob=reward_high_prob,reward_low_prob=reward_low_prob)
return pairs
@staticmethod
def colored_targets_with_randomwalk_reward(length=1000,reward_high_prob=80,reward_low_prob=40,reward_high_span = 20, reward_low_span = 20,step_size_mean = 0, step_size_var = 1):
"""
Generator should return array of ntrials x 2 x 3. The second dimension is for each target.
For example, first is the target with high probability of reward, and the second
entry is for the target with low probability of reward. The third dimension holds three variables indicating
position offset (yes/no), reward probability, and location (binary returned where the
ouput indicates either left or right). The variables reward_high_span and reward_low_span indicate the width
of the range that the high or low reward probability are allowed to span respectively, e.g. if reward_high_prob
is 80 and reward_high_span is 20, then the reward probability for the high value target will be bounded
between 60 and 100 percent.
"""
position_offsetH = np.random.randint(2,size=(1,length))
position_offsetL = np.random.randint(2,size=(1,length))
location_int = np.random.randint(2,size=(1,length))
# define variables for increments: amount of increment and in which direction (i.e. increasing or decreasing)
assign_rewardH = np.random.randn(1,length)
assign_rewardL = np.random.randn(1,length)
assign_rewardH_direction = np.random.randn(1,length)
assign_rewardL_direction = np.random.randn(1,length)
r_0_high = reward_high_prob
r_0_low = reward_low_prob
r_lowerbound_high = r_0_high - (reward_high_span/2)
r_upperbound_high = r_0_high + (reward_high_span/2)
r_lowerbound_low = r_0_low - (reward_low_span/2)
r_upperbound_low = r_0_low + (reward_low_span/2)
reward_high = np.zeros(length)
reward_low = np.zeros(length)
reward_high[0] = r_0_high
reward_low[0] = r_0_low
eps_high = assign_rewardH*step_size_mean + [2*(assign_rewardH_direction > 0) - 1]*step_size_var
eps_low = assign_rewardL*step_size_mean + [2*(assign_rewardL_direction > 0) - 1]*step_size_var
eps_high = eps_high.ravel()
eps_low = eps_low.ravel()
for i in range(1,length):
'''
assign_rewardH_direction = np.random.randn(1)
assign_rewardL_direction = np.random.randn(1)
assign_rewardH = np.random.randn(1)
if assign_rewardH_direction[i-1,] < 0:
eps_high = step_size_mean*assign_rewardH[i-1] - step_size_var
else:
eps_high = step_size_mean*assign_rewardH[i-1] + step_size_var
if assign_rewardL_direction[i] < 0:
eps_low = step_size_mean*assign_rewardL[i] - step_size_var
else:
eps_low = step_size_mean*assign_rewardL[i] + step_size_var
'''
reward_high[i] = reward_high[i-1] + eps_high[i-1]
reward_low[i] = reward_low[i-1] + eps_low[i-1]
reward_high[i] = (r_lowerbound_high < reward_high[i] < r_upperbound_high)*reward_high[i] + (r_lowerbound_high > reward_high[i])*(r_lowerbound_high+ eps_high[i-1]) + (r_upperbound_high < reward_high[i])*(r_upperbound_high - eps_high[i-1])
reward_low[i] = (r_lowerbound_low < reward_low[i] < r_upperbound_low)*reward_low[i] + (r_lowerbound_low > reward_low[i])*(r_lowerbound_low+ eps_low[i-1]) + (r_upperbound_low < reward_low[i])*(r_upperbound_low - eps_low[i-1])
pairs = np.zeros([length,2,3])
pairs[:,0,0] = position_offsetH
pairs[:,0,1] = reward_high
pairs[:,0,2] = location_int
pairs[:,1,0] = position_offsetL
pairs[:,1,1] = reward_low
pairs[:,1,2] = 1 - location_int
return pairs
@staticmethod
def randomwalk_probabilistic_reward(length=1000,reward_high_prob=80,reward_low_prob=40,reward_high_span = 20, reward_low_span = 20,step_size_mean = 0, step_size_var = 1):
pairs = colored_targets_with_randomwalk_reward(length=length,reward_high_prob=reward_high_prob,reward_low_prob=reward_low_prob,reward_high_span = reward_high_span, reward_low_span = reward_low_span,step_size_mean = step_size_mean, step_size_var = step_size_var)
return pairs
class Window(LogExperiment):
'''
Generic stereo window
'''
status = dict(draw=dict(stop=None))
state = "draw"
stop = False
# XPS computer
# window_size = (1920*2, 1080)
window_size = traits.Tuple(monitor_res['monitor_2D'],
desc='Window size, in pixels')
background = traits.Tuple((0., 0., 0., 1.),
desc="Background color (R,G,B,A)")
fullscreen = traits.Bool(True, desc="Fullscreen window")
#Screen parameters, all in centimeters -- adjust for monkey
screen_dist = traits.Float(44.5 + 3, desc="Screen to eye distance (cm)")
screen_half_height = traits.Float(10.75, desc="Screen half height (cm)")
iod = traits.Float(
2.5, desc="Intraocular distance (cm)") # intraocular distance
show_environment = traits.Int(0,
desc="Show wireframe box around environment")
hidden_traits = [
'screen_dist', 'screen_half_height', 'iod', 'show_environment',
'background'
]
def __init__(self, *args, **kwargs):
self.display_start_pos = kwargs.pop('display_start_pos', "0,0")
super(Window, self).__init__(*args, **kwargs)
self.models = []
self.world = None
self.event = None
# os.popen('sudo vbetool dpms on')
self.fov = np.degrees(
np.arctan(self.screen_half_height / self.screen_dist)) * 2
self.screen_cm = [
2 * self.screen_half_height * self.window_size[0] /
self.window_size[1], 2 * self.screen_half_height
]
if self.show_environment:
self.add_model(Box())
def set_os_params(self):
os.environ['SDL_VIDEO_WINDOW_POS'] = self.display_start_pos
os.environ['SDL_VIDEO_X11_WMCLASS'] = "monkey_experiment"
def screen_init(self):
self.set_os_params()
pygame.init()
pygame.display.gl_set_attribute(pygame.GL_DEPTH_SIZE, 24)
flags = pygame.DOUBLEBUF | pygame.HWSURFACE | pygame.OPENGL | pygame.NOFRAME
if self.fullscreen:
flags = flags | pygame.FULLSCREEN
try:
pygame.display.gl_set_attribute(pygame.GL_MULTISAMPLEBUFFERS, 1)
self.screen = pygame.display.set_mode(self.window_size, flags)
except:
pygame.display.gl_set_attribute(pygame.GL_MULTISAMPLEBUFFERS, 0)
self.screen = pygame.display.set_mode(self.window_size, flags)
glEnable(GL_BLEND)
glDepthFunc(GL_LESS)
glEnable(GL_DEPTH_TEST)
glEnable(GL_TEXTURE_2D)
glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
glClearColor(*self.background)
glClearDepth(1.0)
glDepthMask(GL_TRUE)
self.renderer = self._get_renderer()
#this effectively determines the modelview matrix
self.world = Group(self.models)
self.world.init()
#up vector is always (0,0,1), why would I ever need to roll the camera?!
self.set_eye((0, -self.screen_dist, 0), (0, 0))
pygame.mouse.set_visible(False)
def _get_renderer(self):
near = 1
far = 1024
return stereo.MirrorDisplay(self.window_size, self.fov, near, far,
self.screen_dist, self.iod)
def set_eye(self, pos, vec, reset=True):
'''Set the eye's position and direction. Camera starts at (0,0,0), pointing towards positive y'''
self.world.translate(pos[0], pos[2], pos[1], reset=True).rotate_x(-90)
self.world.rotate_y(vec[0]).rotate_x(vec[1])
def add_model(self, model):
if self.world is None:
#world doesn't exist yet, add the model to cache
self.models.append(model)
else:
#We're already running, initialize the model and add it to the world
model.init()
self.world.add(model)
def show_object(self, obj, show=False):
'''
Show or hide an object. This function is an abstraction so that tasks don't need to know about attach/detach
'''
if show:
obj.attach()
else:
obj.detach()
def draw_world(self):
glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
self.renderer.draw(self.world)
pygame.display.flip()
self.renderer.draw_done()
def _get_event(self):
for e in pygame.event.get(pygame.KEYDOWN):
return (e.key, e.type)
def _start_None(self):
pygame.display.quit()
def _test_stop(self, ts):
'''
Stop the task if the escape key is pressed, or if the super _test_stop instructs a stop
'''
super_stop = super(Window, self)._test_stop(ts)
from pygame import K_ESCAPE
key_stop = self.event is not None and self.event[0] == K_ESCAPE
if key_stop:
print("Window closed. Stopping task")
return super_stop or key_stop
def update_report_stats(self):
self.reportstats['FPS'] = round(self.clock.get_fps(), 2)
return super().update_report_stats()
def requeue(self):
self.renderer._queue_render(self.world)
def _cycle(self):
self.requeue()
self.draw_world()
super(Window, self)._cycle(
) # is this order intentional? why not cycle first then draw the screen?
self.event = self._get_event()
class ScreenSync(NIDAQSync):
'''Adds a square in one corner that switches color with every flip.'''
sync_position = {
'TopLeft': (-1, 1),
'TopRight': (1, 1),
'BottomLeft': (-1, -1),
'BottomRight': (1, -1)
}
sync_position_2D = {
'TopLeft': (-1, -1),
'TopRight': (1, -1),
'BottomLeft': (-1, 1),
'BottomRight': (1, 1)
}
sync_corner = traits.OptionsList(tuple(sync_position.keys()),
desc="Position of sync square")
sync_size = traits.Float(1, desc="Sync square size (cm)")
sync_color_off = traits.Tuple((0., 0., 0., 1.),
desc="Sync off color (R,G,B,A)")
sync_color_on = traits.Tuple((1., 1., 1., 1.),
desc="Sync on color (R,G,B,A)")
sync_state_duration = 1 # How long to delay the start of the experiment (seconds)
def __init__(self, *args, **kwargs):
# Create a new "sync" state at the beginning of the experiment
if isinstance(self.status, dict):
self.status["sync"] = dict(start_experiment="wait",
stoppable=False)
else:
from riglib.fsm.fsm import StateTransitions
self.status.states["sync"] = StateTransitions(
start_experiment="wait", stoppable=False)
self.state = "sync"
super().__init__(*args, **kwargs)
self.sync_state = False
if hasattr(self, 'is_pygame_display'):
screen_center = np.divide(self.window_size, 2)
sync_size_pix = self.sync_size * self.window_size[
0] / self.screen_cm[0]
sync_center = [sync_size_pix / 2, sync_size_pix / 2]
from_center = np.multiply(self.sync_position_2D[self.sync_corner],
np.subtract(screen_center, sync_center))
top_left = screen_center + from_center - sync_center
self.sync_rect = pygame.Rect(top_left, np.multiply(sync_center, 2))
else:
from_center = np.multiply(
self.sync_position[self.sync_corner],
np.subtract(self.screen_cm, self.sync_size))
pos = np.array(
[from_center[0] / 2, self.screen_dist, from_center[1] / 2])
self.sync_square = VirtualRectangularTarget(
target_width=self.sync_size,
target_height=self.sync_size,
target_color=self.sync_color_off,
starting_pos=pos)
# self.sync_square = VirtualCircularTarget(target_radius=self.sync_size, target_color=self.sync_color_off, starting_pos=pos)
for model in self.sync_square.graphics_models:
self.add_model(model)
def screen_init(self):
super().screen_init()
if hasattr(self, 'is_pygame_display'):
self.sync = pygame.Surface(self.window_size)
self.sync.fill(TRANSPARENT)
self.sync.set_colorkey(TRANSPARENT)
def _draw_other(self):
# For pygame display
color = self.sync_color_on if self.sync_state else self.sync_color_off
self.sync.fill(255 * np.array(color), rect=self.sync_rect)
self.screen.blit(self.sync, (0, 0))
def init(self):
self.add_dtype('sync_square', bool, (1, ))
super().init()
def _while_sync(self):
'''
Deliberate "startup sequence":
1. Send a clock pulse to denote the start of the FSM loop
2. Turn off the clock and send a single, longer, impulse
to enable measurement of the screen latency
3. Turn the clock back on
'''
# Turn off the clock after the first cycle is synced
if self.cycle_count == 1:
self.sync_every_cycle = False
# Send an impulse to measure latency halfway through the sync state
key_cycle = int(self.fps * self.sync_state_duration / 2)
impulse_duration = 5 # cycles, to make sure it appears on the screen
if self.cycle_count == key_cycle:
self.sync_every_cycle = True
elif self.cycle_count == key_cycle + 1:
self.sync_every_cycle = False
elif self.cycle_count == key_cycle + impulse_duration:
self.sync_every_cycle = True
elif self.cycle_count == key_cycle + impulse_duration + 1:
self.sync_every_cycle = False
def _end_sync(self):
self.sync_every_cycle = True
def _test_start_experiment(self, ts):
return ts > self.sync_state_duration
def _cycle(self):
super()._cycle()
# Update the sync state
if self.sync_every_cycle:
self.sync_state = not self.sync_state
self.task_data['sync_square'] = copy.deepcopy(self.sync_state)
# For OpenGL display, update the graphics
if not hasattr(self, 'is_pygame_display'):
color = self.sync_color_on if self.sync_state else self.sync_color_off
self.sync_square.cube.color = color
class MachineControlExoJointVel_w_Positioner(BMIControlExoJointVel, PositionerTaskController):
'''
Task to automatically go between different joint configurations with the target positioner following
'''
sequence_generators = ['exo_joint_space_targets']
current_assist_level = 1
status = FSMTable(
wait=StateTransitions(start_trial='move'),
move=StateTransitions(reached_config='reward'),
reward=StateTransitions(time_expired='wait'),
)
state = 'wait'
reward_time = traits.Float(3, desc='Time that reward solenoid is open')
config_tolerances_deg = traits.Tuple((10, 10, 10, 7.5, 10), desc='Time that reward solenoid is open')
def __init__(self, *args, **kwargs):
super(MachineControlExoJointVel, self).__init__(*args, **kwargs)
self.config_tolerances = np.deg2rad(np.array(self.config_tolerances_deg))
self.config_tolerances[-1] = np.inf
@staticmethod
def exo_joint_space_targets(n_blocks=10):
# neutral_config = np.array([-0.64732247, 0.79, 0.19634043, 0.97628754, -0.02114062])
target1 = np.array([-1.05, 0.7, 0.4, 1, 0])
target2 = np.array([-0.25, 0.7, 0.4, 1, 0])
target3 = np.array([-0.85, 1.3 , 0.4, 1, 0])
target4 = np.array([-0.65, 1.3, 0.4, 0.2, 0])
trial_target_ls = []
for k in range(n_blocks):
configs = np.random.permutation([target1, target2, target3, target4]) # generate the target sequence for each block
for config in configs:
trial_target_ls.append(dict(target_config=config))
return trial_target_ls
def _test_reached_config(self, *args, **kwargs):
target_endpt = self.plant.kin_chain.endpoint_pos(self._gen_target_config)
current_endpt = self.plant.kin_chain.endpoint_pos(self.plant.joint_angles)
pos_diff = np.linalg.norm(current_endpt - target_endpt)
joint_diff = np.abs(self.plant.joint_angles - self._gen_target_config)
# print pos_diff
# print np.round(np.rad2deg(joint_diff), decimals=2)
# print
return (pos_diff < 3) or np.all(joint_diff < self.config_tolerances)
def load_decoder(self):
self.ssm = StateSpaceJointVelocityActiveExo()
A, B, W = self.ssm.get_ssm_matrices()
filt = MachineOnlyFilter(A, W)
units = []
self.decoder = Decoder(filt, units, self.ssm, binlen=0.1)
self.decoder.n_features = 1
def create_feature_extractor(self):
from riglib.bmi.extractor import DummyExtractor
self.extractor = DummyExtractor()
self._add_feature_extractor_dtype()
def _start_reward(self):
print "trial complete!"
class BMIJointPerturb(BMIControlMulti):
'''
Task where the virtual plant starts in configuration sampled from a discrete set and resets every trial
'''
status = dict(wait=dict(start_trial="premove", stop=None),
premove=dict(premove_complete="target"),
target=dict(enter_target="hold",
timeout="timeout_penalty",
stop=None),
hold=dict(leave_early="hold_penalty",
hold_complete="targ_transition"),
targ_transition=dict(trial_complete="reward",
trial_abort="wait",
trial_incomplete="target",
trial_restart="premove"),
timeout_penalty=dict(timeout_penalty_end="targ_transition"),
hold_penalty=dict(hold_penalty_end="targ_transition"),
reward=dict(reward_end="wait"))
sequence_generators = BMIControlMulti.sequence_generators + [
'tentacle_multi_start_config'
]
pert_angles = traits.Tuple(
(np.pi, -3 * np.pi / 4, 3 * np.pi / 4),
desc="Possible wrist angles for perturbed configurations")
# pert_angles = traits.Tuple((np.pi, -7*np.pi/8, 7*np.pi/8, -3*np.pi/4, 3*np.pi/4), desc="Possible wrist angles for perturbed configurations")
# pert_angles = traits.Tuple((-3*np.pi/4), desc="Possible wrist angles for perturbed configurations")
#np.pi,
#tuple(np.linspace(-np.pi, np.pi, 8)
premove_time = traits.Float(
.1, desc='Time before subject must start doing BMI control')
# static_states = ['premove'] # states in which the decoder is not run
def __init__(self, *args, **kwargs):
super(BMIJointPerturb, self).__init__(*args, **kwargs)
def _parse_next_trial(self):
self.targs, self.curr_pert_angle = self.next_trial
# self.targs, self.curr_pert_angle, self.arm_visible = self.next_trial
# def _start_wait(self):
# self.curr_pert_angle = self.pert_angles[np.random.randint(0, high=len(self.pert_angles))]
# super(BMIJointPerturb, self)._start_wait()
def _while_premove(self):
self.decoder.filt.state.mean = self.calc_perturbed_ik(self.targs[0])
def _end_timeout_penalty(self):
pass
def calc_perturbed_ik(self, endpoint_pos):
distal_angles = np.array([self.curr_pert_angle, -np.pi / 20])
# second angle above used to be self.init_decoder_mean[3,0] for center out version of task
joints = self.plant.perform_ik(endpoint_pos,
distal_angles=-distal_angles)
return np.mat(np.hstack([joints, np.zeros(4), 1]).reshape(-1, 1))
def _test_premove_complete(self, ts):
return ts >= self.premove_time
def _test_hold_complete(self, ts):
if self.target_index == 0:
return True
else:
return ts >= self.hold_time
def _test_trial_incomplete(self, ts):
return (self.target_index < self.chain_length - 1) and (
self.target_index != -1) and (self.tries < self.max_attempts)
def _test_trial_restart(self, ts):
return (self.target_index == -1) and (self.tries < self.max_attempts)
@staticmethod
def tentacle_multi_start_config(nblocks=100,
ntargets=4,
distance=8,
startangle=45):
elbow_angles = np.array([
135, 180, 225
]) * np.pi / 180 # TODO make this a function argument!
startangle = 45 * np.pi / 180
n_configs_per_target = len(elbow_angles)
target_angles = np.arange(startangle, startangle + (2 * np.pi),
2 * np.pi / ntargets)
targets = distance * np.vstack(
[np.cos(target_angles), 0 * target_angles,
np.sin(target_angles)])
seq = []
from itertools import izip
import random
for i in range(nblocks):
target_inds = np.tile(np.arange(ntargets),
(n_configs_per_target, 1)).T.ravel()
config_inds = np.tile(np.arange(n_configs_per_target), ntargets)
sub_seq = []
inds = np.arange(n_configs_per_target * ntargets)
random.shuffle(inds)
for k in inds:
targ_ind = target_inds[k]
config_ind = config_inds[k]
seq_item = (np.vstack([targets[:, targ_ind],
np.zeros(3)]), elbow_angles[config_ind])
seq.append(seq_item)
return seq
class CLDAManipulatedFB(BMIControlManipulatedFB):
'''
BMI task that periodically refits the decoder parameters based on intended
movements toward the targets. Inherits directly from BMIControl. Can be made
to automatically linearly decrease assist level over set time period, or
to provide constant assistance by setting assist_level and assist_min equal.
'''
batch_time = traits.Float(80.0, desc='The length of the batch in seconds')
half_life = traits.Tuple((120., 120.0), desc='Half life of the adaptation in seconds')
decoder_sequence = traits.String('test', desc='signifier to group together sequences of decoders')
#assist_min = traits.Float(0, desc="Assist level to end task at")
#half_life_final = traits.Float(120.0, desc='Half life of the adaptation in seconds')
half_life_decay_time = traits.Float(900.0, desc="Time to go from initial half life to final")
def __init__(self, *args, **kwargs):
super(CLDAManipulatedFB, self).__init__(*args, **kwargs)
#self.assist_start = self.assist_level
self.learn_flag = True
def init(self):
'''
Secondary init function. Decoder has already been created by inclusion
of the 'bmi' feature in the task. Create the 'learner' and 'updater'
components of the CLDA algorithm
'''
# Add CLDA-specific data to save to HDF file
self.dtype.append(('half_life', 'f8', (1,)))
super(CLDAManipulatedFB, self).init()
self.batch_size = int(self.batch_time/self.decoder.binlen)
self.create_learner()
# Create the updater second b/c the update algorithm might need to force
# a particular batch size for the learner
self.create_updater()
# Create the BMI system which combines the decoder, learner, and updater
self.bmi_system = riglib.bmi.BMISystem(self.decoder, self.learner,
self.updater)
def create_learner(self):
self.learner = clda.CursorGoalLearner(self.batch_size)
# Start "learn flag" at True
self.learn_flag = True
homedir = os.getenv('HOME')
f = open(os.path.join(homedir, 'learn_flag_file'), 'w')
f.write('1')
f.close()
def create_updater(self):
clda_input_queue = mp.Queue()
clda_output_queue = mp.Queue()
half_life_start, half_life_end = self.half_life
self.updater = clda.KFSmoothbatch(clda_input_queue, clda_output_queue,self.batch_time, half_life_start)
def update_learn_flag(self):
# Tell the adaptive BMI when to learn (skip parts of the task where we
# assume the subject is not trying to move toward the target)
prev_learn_flag = self.learn_flag
# Open file to read learn flag
try:
homedir = os.getenv('HOME')
f = open(os.path.join(homedir, 'learn_flag_file'))
new_learn_flag = bool(int(f.readline().rstrip('\n')))
except:
new_learn_flag = True
if new_learn_flag and not prev_learn_flag:
print "CLDA enabled"
elif prev_learn_flag and not new_learn_flag:
try:
print "CLDA disabled after %d successful trials" % self.calc_n_rewards()
except:
print "CLDA disabled"
self.learn_flag = new_learn_flag
def call_decoder(self, spike_counts):
half_life_start, half_life_end = self.half_life
current_half_life = self._linear_change(half_life_start, half_life_end, self.half_life_decay_time)
self.task_data['half_life'] = current_half_life
# Get the decoder output
decoder_output, uf = self.bmi_system(spike_counts, self.target_location,
self.state, task_data=self.task_data, assist_level=self.current_assist_level,
target_radius=self.target_radius, speed=self.assist_speed*self.decoder.binlen,
learn_flag=self.learn_flag, half_life=current_half_life)
if uf:
#send msg to hdf file to indicate decoder update
self.hdf.sendMsg("update_bmi")
return decoder_output #self.decoder['hand_px', 'hand_py', 'hand_pz']
def _cycle(self):
self.update_learn_flag()
super(CLDAManipulatedFB, self)._cycle()
def cleanup(self, database, saveid, **kwargs):
super(CLDAManipulatedFB, self).cleanup(database, saveid, **kwargs)
import tempfile, cPickle, traceback, datetime
# Open a log file in case of error b/c errors not visible to console
# at this point
f = open(os.path.join(os.getenv('HOME'), 'Desktop/log'), 'a')
f.write('Opening log file\n')
# save out the parameter history and new decoder unless task was stopped
# before 1st update
try:
f.write('# of paramter updates: %d\n' % len(self.bmi_system.param_hist))
if len(self.bmi_system.param_hist) > 0:
f.write('Starting to save parameter hist\n')
tf = tempfile.NamedTemporaryFile()
# Get the update history for C and Q matrices and save them
#C, Q, m, sd, intended_kin, spike_counts = zip(*self.bmi_system.param_hist)
#np.savez(tf, C=C, Q=Q, mean=m, std=sd, intended_kin=intended_kin, spike_counts=spike_counts)
pickle.dump(self.bmi_system.param_hist, tf)
tf.flush()
# Add the parameter history file to the database entry for this
# session
database.save_data(tf.name, "bmi_params", saveid)
f.write('Finished saving parameter hist\n')
# Save the final state of the decoder as a new decoder
tf2 = tempfile.NamedTemporaryFile(delete=False)
cPickle.dump(self.decoder, tf2)
tf2.flush()
# create suffix for new decoder that has the sequence and the current day
# and time. This suffix will be appended to the name of the
# decoder that we started with and saved as a new decoder.
now = datetime.datetime.now()
decoder_name = self.decoder_sequence + '%02d%02d%02d%02d' % (now.month, now.day, now.hour, now.minute)
database.save_bmi(decoder_name, saveid, tf2.name)
except:
traceback.print_exc(file=f)
f.close()
class RatBMI(BMILoop, LogExperiment):
status = dict(wait=dict(start_trial='feedback_on', stop=None),
feedback_on=dict(baseline_hit='periph_targets', stop=None),
periph_targets=dict(target_hit='check_reward',
timeout='noise_burst',
stop=None),
check_reward=dict(rewarded_target='reward',
unrewarded_target='feedback_pause'),
feedback_pause=dict(end_feedback_pause='wait'),
reward=dict(reward_end='wait'),
noise_burst=dict(noise_burst_end='noise_burst_timeout'),
noise_burst_timeout=dict(noise_burst_timeout_end='wait'))
#Flag for feedback on or not
feedback = False
prev_targ_hit = 't1'
timeout_time = traits.Float(30.)
noise_burst_time = traits.Float(3.)
noise_burst_timeout_time = traits.Float(1.)
reward_time = traits.Float(1., desc='reward time')
#Frequency range:
aud_freq_range = traits.Tuple((1000., 20000.))
plant_type = traits.OptionsList(*plantlist,
desc='',
bmi3d_input_options=plantlist.keys())
#Time to average over:
nsteps = traits.Float(10.)
feedback_pause = traits.Float(3.)
def __init__(self, *args, **kwargs):
super(RatBMI, self).__init__(*args, **kwargs)
if hasattr(self, 'decoder'):
print self.decoder
else:
self.decoder = kwargs['decoder']
dec_params = dict(nsteps=self.nsteps, freq_lim=self.aud_freq_range)
for k, (key, val) in enumerate(dec_params.items()):
print key, val, self.decoder.filt.dec_params[key]
assert self.decoder.filt.dec_params[key] == val
self.decoder.filt.init_from_task(self.decoder.n_units, **dec_params)
self.plant = plantlist[self.plant_type]
def init(self, *args, **kwargs):
self.add_dtype('cursor', 'f8', (2, ))
self.add_dtype('freq', 'f8', (2, ))
super(RatBMI, self).init()
self.decoder.count_max = self.feature_accumulator.count_max
def _cycle(self):
self.rat_cursor = self.decoder.filt.get_mean()
self.task_data['cursor'] = self.rat_cursor
self.task_data['freq'] = self.decoder.filt.F
self.decoder.cnt = self.feature_accumulator.count
self.decoder.feedback = self.feedback
super(RatBMI, self)._cycle()
# def move_plant(self):
# if self.feature_accumulator.count == self.feature_accumulator.count_max:
# print 'self.plant.drive from task.py'
# self.plant.drive(self.decoder)
def _start_wait(self):
return True
def _test_start_trial(self, ts):
return True
def _test_rewarded_target(self, ts):
if self.prev_targ_hit == 't1':
return False
elif self.prev_targ_hit == 't2':
return True
def _test_unrewarded_target(self, ts):
if self.prev_targ_hit == 't1':
return True
elif self.prev_targ_hit == 't2':
return False
def _start_feedback_pause(self):
self.feedback = False
def _test_end_feedback_pause(self, ts):
return ts > self.feedback_pause
def _start_reward(self):
print 'reward!'
def _start_feedback_on(self):
self.feedback = True
def _test_baseline_hit(self, ts):
if self.prev_targ_hit == 't1':
#Must go below baseline:
return self.rat_cursor <= self.decoder.filt.mid
elif self.prev_targ_hit == 't2':
#Must rise above baseline:
return self.rat_cursor >= self.decoder.filt.mid
else:
return False
def _test_target_hit(self, ts):
if self.rat_cursor >= self.decoder.filt.t1:
self.prev_targ_hit = 't1'
self.feedback = False
return True
elif self.rat_cursor <= self.decoder.filt.t2:
self.prev_targ_hit = 't2'
self.feedback = False
return True
else:
return False
def _test_timeout(self, ts):
return ts > self.timeout_time
def _test_noise_burst_end(self, ts):
return ts > self.noise_burst_time
def _test_noise_burst_timeout_end(self, ts):
return ts > self.noise_burst_timeout_time
def _start_noise_burst(self):
self.feedback = False
self.plant.play_white_noise()
def move_plant(self):
super(RatBMI, self).move_plant()
def get_current_assist_level(self):
return 0.
class BMIControlExoEndpt(ExoAssist, BMIControlExoJointVel, PositionerTaskController):
status = dict(
go_to_origin = dict(microcontroller_done='wait', stop=None),
wait = dict(start_trial='init_exo', stop=None),
init_exo = dict(exo_ready='move_target', stop=None),
move_target = dict(microcontroller_done='pause'),
pause = dict(time_expired='reach', stop=None),
reach = dict(force_applied='reward', new_target_set_remotely='move_target', skip='wait', stop=None),
reward = dict(time_expired='wait', stop=None),
)
state = 'go_to_origin'
pause_time = 2
trial_end_states = ['reward']
sequence_generators = ['exo_endpt_targets']
min_force_on_target = traits.Float(1., desc='Force that needs to be applied, in Newtons')
reward_time = traits.Float(3., desc='reward time for solenoid')
config_tolerances_deg = traits.Tuple((10, 10, 10, 7.5, 10), desc='Time that reward solenoid is open')
def __init__(self, *args, **kwargs):
super(BMIControlExoEndpt, self).__init__(*args, **kwargs)
self.config_tolerances = np.deg2rad(np.array(self.config_tolerances_deg))
self.config_tolerances[-1] = np.inf
# fixed target location, for now
self._gen_int_target_pos = np.array([ 27.3264, 47.383, -22.79])
#36.4222245 , 26.27462972, -11.38728596
self.neutral_config = np.array([-0.64732247, 0.79, 0.19634043, 0.97628754, -0.02114062])
self._gen_target_config = np.array([-0.50579373, 1.28357092, 0.66706522, 0.6, -0.02114062])
self.current_assist_level = 1
def init(self):
self.add_dtype('starting_config', 'f8', (5,))
super(BMIControlExoEndpt, self).init()
def init_decoder_state(self):
pass
@staticmethod
def exo_endpt_targets(n_blocks=1):
final_config = np.array([-0.50579373, 1.28357092, 0.66706522, 0.6, -0.02114062])
configs = [final_config]
trial_target_ls = []
init_configs = [
np.array([-0.64732247, 0.79, 0.19634043, 0.97628754, -0.02114062]),
np.array([-0.64732247, 1.5, 0.19634043, 0.97628754, -0.02114062]),
np.array([-1.05, 0.79, 0.19634043, 0.97628754, -0.02114062]),
np.array([-1.05, 0.79, -0.07, 0.97628754, -0.02114062]),
]
for k in range(n_blocks):
for config in init_configs:
trial_target_ls.append(dict(target_config=final_config, init_config=config))
return trial_target_ls
def get_target_BMI_state(self, *args):
'''
Run the goal calculator to determine the target state of the task
'''
if self.state in ['reach', 'reward']:
target_config = self._gen_target_config
elif hasattr(self, '_gen_init_config'):
target_config = self._gen_init_config
else: # before the wait state
target_config = self.neutral_config
target_state = np.hstack([target_config, np.zeros_like(target_config), np.zeros_like(target_config), 1])
target_state = target_state.reshape(-1, 1)
self.print_to_terminal(self.decoder.get_state())
self.print_to_terminal(self.plant.joint_angles)
self.print_to_terminal(target_state[0:5,0].ravel())
self.print_to_terminal('')
return target_state
def _test_exo_ready(self, *args, **kwargs):
target_endpt = self.plant.kin_chain.endpoint_pos(self._gen_init_config)
current_endpt = self.plant.kin_chain.endpoint_pos(self.plant.joint_angles)
pos_diff = np.linalg.norm(current_endpt - target_endpt)
joint_diff = np.abs(self.plant.joint_angles - self._gen_init_config)
return (pos_diff < 3) or np.all(joint_diff < self.config_tolerances)
def _cycle(self):
if hasattr(self, '_gen_init_config'):
starting_config = self._gen_init_config
else: # before the wait state
starting_config = self.neutral_config
self.task_data['starting_config'] = starting_config
super(BMIControlExoEndpt, self)._cycle()
def _end_go_to_origin(self):
steps_actuated = self.pos_uctrl_iface.end_continuous_move(stiff=True)
self.loc = np.zeros(3)
self.steps_from_origin = np.zeros(3)
def _end_move_target(self):
# send command to kill motors
steps_actuated = self.pos_uctrl_iface.end_continuous_move(stiff=True)
self._integrate_steps(steps_actuated, self.pos_uctrl_iface.motor_dir)
def _test_force_applied(self, *args, **kwargs):
return self.plant.force_N > self.min_force_on_target
class Window(LogExperiment):
'''
Generic stereo window
'''
status = dict(draw=dict(stop=None))
state = "draw"
stop = False
window_size = traits.Tuple((1920 * 2, 1080),
descr='window size, in pixels')
# window_size = (1920*2, 1080)
background = (0, 0, 0, 1)
#Screen parameters, all in centimeters -- adjust for monkey
fov = np.degrees(np.arctan(14.65 / (44.5 + 3))) * 2
screen_dist = 44.5 + 3
iod = 2.5 # intraocular distance
show_environment = traits.Int(0)
def __init__(self, *args, **kwargs):
super(Window, self).__init__(*args, **kwargs)
self.models = []
self.world = None
self.event = None
# os.popen('sudo vbetool dpms on')
if self.show_environment:
self.add_model(Box())
def set_os_params(self):
os.environ['SDL_VIDEO_WINDOW_POS'] = config.display_start_pos
os.environ['SDL_VIDEO_X11_WMCLASS'] = "monkey_experiment"
def screen_init(self):
self.set_os_params()
pygame.init()
pygame.display.gl_set_attribute(pygame.GL_DEPTH_SIZE, 24)
flags = pygame.DOUBLEBUF | pygame.HWSURFACE | pygame.OPENGL | pygame.NOFRAME
try:
pygame.display.gl_set_attribute(pygame.GL_MULTISAMPLEBUFFERS, 1)
self.surf = pygame.display.set_mode(self.window_size, flags)
except:
pygame.display.gl_set_attribute(pygame.GL_MULTISAMPLEBUFFERS, 0)
self.surf = pygame.display.set_mode(self.window_size, flags)
glEnable(GL_BLEND)
glDepthFunc(GL_LESS)
glEnable(GL_DEPTH_TEST)
glEnable(GL_TEXTURE_2D)
glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)
glClearColor(*self.background)
glClearDepth(1.0)
glDepthMask(GL_TRUE)
self.renderer = self._get_renderer()
#this effectively determines the modelview matrix
self.world = Group(self.models)
self.world.init()
#up vector is always (0,0,1), why would I ever need to roll the camera?!
self.set_eye((0, -self.screen_dist, 0), (0, 0))
def _get_renderer(self):
near = 1
far = 1024
return stereo.MirrorDisplay(self.window_size, self.fov, near, far,
self.screen_dist, self.iod)
def set_eye(self, pos, vec, reset=True):
'''Set the eye's position and direction. Camera starts at (0,0,0), pointing towards positive y'''
self.world.translate(pos[0], pos[2], pos[1], reset=True).rotate_x(-90)
self.world.rotate_y(vec[0]).rotate_x(vec[1])
def add_model(self, model):
if self.world is None:
#world doesn't exist yet, add the model to cache
self.models.append(model)
else:
#We're already running, initialize the model and add it to the world
model.init()
self.world.add(model)
def show_object(self, obj, show=False):
'''
Show or hide an object. This function is an abstraction so that tasks don't need to know about attach/detach
'''
if show:
obj.attach()
else:
obj.detach()
def draw_world(self):
glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
self.renderer.draw(self.world)
pygame.display.flip()
self.renderer.draw_done()
def _get_event(self):
for e in pygame.event.get(pygame.KEYDOWN):
return (e.key, e.type)
def _start_None(self):
pygame.display.quit()
def _test_stop(self, ts):
'''
Stop the task if the escape key is pressed, or if the super _test_stop instructs a stop
'''
super_stop = super(Window, self)._test_stop(ts)
from pygame import K_ESCAPE
return super_stop or self.event is not None and self.event[
0] == K_ESCAPE
def requeue(self):
self.renderer._queue_render(self.world)
def _cycle(self):
self.requeue()
self.draw_world()
super(Window, self)._cycle()
self.event = self._get_event()
class MovementTraining(Window):
status = dict(wait=dict(stop=None, move_start="movement"),
movement=dict(move_end="reward", move_stop="wait",
stop=None),
reward=dict(reward_end="wait"))
log_exclude = set((("wait", "move_start"), ("movement", "move_stop")))
#initial state
state = "wait"
path = [[0, 0, 0]]
speed = 0
frame_offset = 2
over = 0
inside = 0
#settable traits
movement_distance = traits.Float(
1, desc="Minimum movement distance to trigger reward")
speed_range = traits.Tuple(
20, 30, desc="Range of movement speed in cm/s to trigger reward")
reward_time = traits.Float(14)
#initialize
def __init__(self, **kwargs):
super(MovementTraining, self).__init__(**kwargs)
self.cursor = Sphere(radius=.5, color=(.5, 0, .5, 1))
self.add_model(self.cursor)
def update_cursor(self):
#get data from 13th marker on motion tracker- take average of all data points since last poll
pt = self.motiondata.get()
if len(pt) > 0:
pt = pt[:, 14, :]
# NOTE!!! The marker on the hand was changed from #0 to #14 on
# 5/19/13 after LED #0 broke. All data files saved before this date
# have LED #0 controlling the cursor.
pt = pt[~np.isnan(pt).any(1)]
if len(pt) > 0:
pt = pt.mean(0)
self.path.append(pt)
#ignore y direction
t = pt * .25
t[1] = 0
#move cursor to marker location
self.cursor.translate(*t[:3], reset=True)
else:
self.path.append(self.path[-1])
if len(self.path) > self.frame_offset:
self.path.pop(0)
d = np.sqrt((self.path[-1][0] - self.path[0][0])**2 +
(self.path[-1][1] - self.path[0][1])**2 +
(self.path[-1][2] - self.path[0][2])**2)
self.speed = d / (self.frame_offset / 60)
if self.speed > self.speed_range[0]:
self.over += 1
if self.speed_range[0] < self.speed < self.speed_range[1]:
self.inside += 1
#write to screen
self.draw_world()
def _start_wait(self):
self.over = 0
self.inside = 0
def _while_wait(self):
self.update_cursor()
def _while_movement(self):
self.update_cursor()
def _while_reward(self):
self.update_cursor()
def _test_move_start(self, ts):
return self.over > self.frame_offset
def _test_move_end(self, ts):
return ts > self.movement_distance / self.speed_range[0]
def _test_move_stop(self, ts):
return self.inside > self.frame_offset
def _test_reward_end(self, ts):
return ts > self.reward_time
class BMIResettingObstacles(BMIResetting):
status = dict(
wait=dict(start_trial="premove", stop=None),
premove=dict(premove_complete="target"),
target=dict(enter_target="hold",
timeout="timeout_penalty",
enter_obstacle="obstacle_penalty",
stop=None),
hold=dict(leave_early="hold_penalty", hold_complete="targ_transition"),
targ_transition=dict(trial_complete="reward",
trial_abort="wait",
trial_incomplete="target",
trial_restart="premove"),
timeout_penalty=dict(timeout_penalty_end="targ_transition"),
hold_penalty=dict(hold_penalty_end="targ_transition"),
obstacle_penalty=dict(obstacle_penalty_end="targ_transition"),
reward=dict(reward_end="wait"))
sequence_generators = [
'centerout_2D_discrete_w_obstacle', 'centerout_2D_discrete'
]
obstacle_sizes = traits.Tuple((2, 3), desc='must match generator sizes!')
is_bmi_seed = True
def __init__(self, *args, **kwargs):
super(BMIResettingObstacles, self).__init__(*args, **kwargs)
self.add_obstacles()
def add_obstacles(self):
import target_graphics
#Add obstacle
self.obstacle_list = []
self.obstacle_dict = {}
for i in self.obstacle_sizes:
obstacle = target_graphics.VirtualRectangularTarget(
target_width=i,
target_height=i,
target_color=(0, 0, 1, .5),
starting_pos=np.zeros(3))
self.obstacle_list.append(obstacle)
self.obstacle_dict[i] = len(self.obstacle_list) - 1
for model in obstacle.graphics_models:
self.add_model(model)
def init(self, *args, **kwargs):
self.add_dtype('obstacle_size', 'f8', (1, ))
self.add_dtype('obstacle_location', 'f8', (3, ))
super(BMIResettingObstacles, self).init(*args, **kwargs)
def create_goal_calculator(self):
self.goal_calculator = goal_calculators.Obs_Goal_Calc(self.decoder.ssm)
def create_assister(self, *args, **kwargs):
self.assister = OFCEndpointAssister()
#self.assister = ObstacleAssist(self.decoder.ssm)
def _start_wait(self):
for obs in self.obstacle_list:
obs.hide()
super(BMIResetting, self)._start_wait()
def _start_premove(self):
super(BMIResettingObstacles, self)._start_premove()
def _start_target(self):
print 'start target BMIRes'
self.goal_calculator.clear()
super(BMIResettingObstacles, self)._start_target()
# def _start_target(self):
# super(BMIResettingObstacles, self)._start_target()
def _parse_next_trial(self):
self.targs = self.next_trial[0]
#Width and height of obstacle
self.obstacle_size = self.next_trial[1]
self.obstacle = self.obstacle_list[self.obstacle_dict[
self.obstacle_size]]
self.obstacle_location = self.next_trial[2]
def _start_target(self):
super(BMIResettingObstacles, self)._start_target()
self.obstacle.move_to_position(self.obstacle_location)
self.obstacle.cube.color = (0., 0., 1., .5)
self.obstacle.show()
# print 'targ loc: ', self.target_location.astype(int)
# print 'obstacle_location: ', self.obstacle_location.astype(int)
# print 'self.targs: ', self.targs.astype(int)
def _test_enter_obstacle(self, ts):
cursor_pos = self.plant.get_endpoint_pos()
centered_cursor_pos = np.abs(cursor_pos - self.obstacle_location)
return np.all(centered_cursor_pos < self.obstacle_size / 2.)
def _test_obstacle_penalty_end(self, ts):
self.obstacle.cube.color = (1., 1., 0., .5)
return ts >= self.timeout_penalty_time
def _start_obstacle_penalty(self):
#hide targets
for target in self.targets:
target.hide()
self.tries += 1
self.target_index = -1
def _cycle(self):
self.add_obstacle_data()
super(BMIResettingObstacles, self)._cycle()
def add_obstacle_data(self):
self.task_data['obstacle_size'] = self.obstacle_size
self.task_data['obstacle_location'] = self.obstacle_location
@staticmethod
def centerout_2D_discrete_w_obstacle(nblocks=100,
ntargets=8,
boundaries=(-18, 18, -12, 12),
distance=10,
obstacle_sizes=(2, 3)):
'''
Generates a sequence of 2D (x and z) target pairs with the first target
always at the origin.
Parameters
----------
length : int
The number of target pairs in the sequence.
boundaries: 6 element Tuple
The limits of the allowed target locations (-x, x, -z, z)
distance : float
The distance in cm between the targets in a pair.
obstacle_sizes: tuple of varying sizes
Returns
-------
pairs : [nblocks*ntargets x 2 x 3] array of pairs of target locations
'''
# Choose a random sequence of points on the edge of a circle of radius
# "distance"
theta = []
for i in range(nblocks):
temp_master = []
for o in obstacle_sizes:
angs = np.arange(0, 2 * np.pi, 2 * np.pi / ntargets)
obs_sz = [o] * ntargets
temp = np.vstack((angs, obs_sz)).T
np.random.shuffle(temp)
temp_master.append(temp)
x = np.vstack((temp_master))
np.random.shuffle(x)
theta = theta + [x]
theta = np.vstack(theta)
x = distance * np.cos(theta[:, 0])
y = np.zeros(len(theta[:, 0]))
z = distance * np.sin(theta[:, 0])
pairs = np.zeros([len(theta), 2, 3])
pairs[:, 1, :] = np.vstack([x, y, z]).T
obstacle_location = (pairs[:, 1, :] - pairs[:, 0, :]) * 0.5
return zip(pairs, theta[:, 1], obstacle_location)
class ScreenTargetCapture(TargetCapture, Window):
"""Concrete implementation of TargetCapture task where targets
are acquired by "holding" a cursor in an on-screen target"""
limit2d = 1
sequence_generators = [
'out_2D',
'centerout_2D',
'centeroutback_2D',
'rand_target_chain_2D',
'rand_target_chain_3D',
]
hidden_traits = [
'cursor_color', 'target_color', 'cursor_bounds', 'cursor_radius',
'plant_hide_rate', 'starting_pos'
]
is_bmi_seed = True
# Runtime settable traits
target_radius = traits.Float(2, desc="Radius of targets in cm")
target_color = traits.OptionsList("yellow",
*target_colors,
desc="Color of the target",
bmi3d_input_options=list(
target_colors.keys()))
plant_hide_rate = traits.Float(
0.0,
desc=
'If the plant is visible, specifies a percentage of trials where it will be hidden'
)
plant_type = traits.OptionsList(*plantlist,
bmi3d_input_options=list(plantlist.keys()))
plant_visible = traits.Bool(
True,
desc='Specifies whether entire plant is displayed or just endpoint')
cursor_radius = traits.Float(.5, desc='Radius of cursor in cm')
cursor_color = traits.OptionsList("pink",
*target_colors,
desc='Color of cursor endpoint',
bmi3d_input_options=list(
target_colors.keys()))
cursor_bounds = traits.Tuple(
(-10., 10., 0., 0., -10., 10.),
desc='(x min, x max, y min, y max, z min, z max)')
starting_pos = traits.Tuple((5., 0., 5.),
desc='Where to initialize the cursor')
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# Initialize the plant
if not hasattr(self, 'plant'):
self.plant = plantlist[self.plant_type]
self.plant.set_endpoint_pos(np.array(self.starting_pos))
self.plant.set_bounds(np.array(self.cursor_bounds))
self.plant.set_color(target_colors[self.cursor_color])
self.plant.set_cursor_radius(self.cursor_radius)
self.plant_vis_prev = True
self.cursor_vis_prev = True
# Add graphics models for the plant and targets to the window
if hasattr(self.plant, 'graphics_models'):
for model in self.plant.graphics_models:
self.add_model(model)
# Instantiate the targets
instantiate_targets = kwargs.pop('instantiate_targets', True)
if instantiate_targets:
# Need two targets to have the ability for delayed holds
target1 = VirtualCircularTarget(
target_radius=self.target_radius,
target_color=target_colors[self.target_color])
target2 = VirtualCircularTarget(
target_radius=self.target_radius,
target_color=target_colors[self.target_color])
self.targets = [target1, target2]
# Declare any plant attributes which must be saved to the HDF file at the _cycle rate
for attr in self.plant.hdf_attrs:
self.add_dtype(*attr)
def init(self):
self.add_dtype('trial', 'u4', (1, ))
self.add_dtype('plant_visible', '?', (1, ))
super().init()
def _cycle(self):
'''
Calls any update functions necessary and redraws screen
'''
self.move_effector()
## Run graphics commands to show/hide the plant if the visibility has changed
self.update_plant_visibility()
self.task_data['plant_visible'] = self.plant_visible
## Save plant status to HDF file
plant_data = self.plant.get_data_to_save()
for key in plant_data:
self.task_data[key] = plant_data[key]
# Update the trial index
self.task_data['trial'] = self.calc_trial_num()
super()._cycle()
def move_effector(self):
'''Move the end effector, if a robot or similar is being controlled'''
pass
def run(self):
'''
See experiment.Experiment.run for documentation.
'''
# Fire up the plant. For virtual/simulation plants, this does little/nothing.
self.plant.start()
# Include some cleanup in case the parent class has errors
try:
super().run()
finally:
self.plant.stop()
##### HELPER AND UPDATE FUNCTIONS ####
def update_plant_visibility(self):
''' Update plant visibility'''
if self.plant_visible != self.plant_vis_prev:
self.plant_vis_prev = self.plant_visible
self.plant.set_visibility(self.plant_visible)
#### TEST FUNCTIONS ####
def _test_enter_target(self, ts):
'''
return true if the distance between center of cursor and target is smaller than the cursor radius
'''
cursor_pos = self.plant.get_endpoint_pos()
d = np.linalg.norm(cursor_pos - self.targs[self.target_index])
return d <= (self.target_radius - self.cursor_radius)
def _test_leave_target(self, ts):
'''
return true if cursor moves outside the exit radius
'''
cursor_pos = self.plant.get_endpoint_pos()
d = np.linalg.norm(cursor_pos - self.targs[self.target_index])
rad = self.target_radius - self.cursor_radius
return d > rad
#### STATE FUNCTIONS ####
def _start_wait(self):
super()._start_wait()
if self.calc_trial_num() == 0:
# Instantiate the targets here so they don't show up in any states that might come before "wait"
for target in self.targets:
for model in target.graphics_models:
self.add_model(model)
target.hide()
def _start_target(self):
super()._start_target()
# Show target if it is hidden (this is the first target, or previous state was a penalty)
target = self.targets[self.target_index % 2]
if self.target_index == 0:
target.move_to_position(self.targs[self.target_index])
target.show()
self.sync_event('TARGET_ON', self.gen_indices[self.target_index])
def _start_hold(self):
super()._start_hold()
self.sync_event('CURSOR_ENTER_TARGET',
self.gen_indices[self.target_index])
def _start_delay(self):
super()._start_delay()
# Make next target visible unless this is the final target in the trial
next_idx = (self.target_index + 1)
if next_idx < self.chain_length:
target = self.targets[next_idx % 2]
target.move_to_position(self.targs[next_idx])
target.show()
self.sync_event('TARGET_ON', self.gen_indices[next_idx])
else:
# This delay state should only last 1 cycle, don't sync anything
pass
def _start_targ_transition(self):
super()._start_targ_transition()
if self.target_index == -1:
# Came from a penalty state
pass
elif self.target_index + 1 < self.chain_length:
# Hide the current target if there are more
self.targets[self.target_index % 2].hide()
self.sync_event('TARGET_OFF', self.gen_indices[self.target_index])
def _start_hold_penalty(self):
self.sync_event('HOLD_PENALTY')
super()._start_hold_penalty()
# Hide targets
for target in self.targets:
target.hide()
target.reset()
def _end_hold_penalty(self):
super()._end_hold_penalty()
self.sync_event('TRIAL_END')
def _start_delay_penalty(self):
self.sync_event('DELAY_PENALTY')
super()._start_delay_penalty()
# Hide targets
for target in self.targets:
target.hide()
target.reset()
def _end_delay_penalty(self):
super()._end_delay_penalty()
self.sync_event('TRIAL_END')
def _start_timeout_penalty(self):
self.sync_event('TIMEOUT_PENALTY')
super()._start_timeout_penalty()
# Hide targets
for target in self.targets:
target.hide()
target.reset()
def _end_timeout_penalty(self):
super()._end_timeout_penalty()
self.sync_event('TRIAL_END')
def _start_reward(self):
self.targets[self.target_index % 2].cue_trial_end_success()
self.sync_event('REWARD')
def _end_reward(self):
super()._end_reward()
self.sync_event('TRIAL_END')
# Hide targets
for target in self.targets:
target.hide()
target.reset()
#### Generator functions ####
'''
Note to self: because of the way these get into the database, the parameters don't
have human-readable descriptions like the other traits. So it is useful to define
the descriptions elsewhere, in models.py under Generator.to_json().
Ideally someone should take the time to reimplement generators as their own classes
rather than static methods that belong to a task.
'''
@staticmethod
def static(pos=(0, 0, 0), ntrials=0):
'''Single location, finite (ntrials!=0) or infinite (ntrials==0)'''
if ntrials == 0:
while True:
yield [0], np.array(pos)
else:
for _ in range(ntrials):
yield [0], np.array(pos)
@staticmethod
def out_2D(nblocks=100, ntargets=8, distance=10, origin=(0, 0, 0)):
'''
Generates a sequence of 2D (x and z) targets at a given distance from the origin
Parameters
----------
nblocks : int
The number of ntarget pairs in the sequence.
ntargets : int
The number of equally spaced targets
distance : float
The distance in cm between the center and peripheral targets.
origin : 3-tuple
Location of the central targets around which the peripheral targets span
Returns
-------
[nblocks*ntargets x 1] array of tuples containing trial indices and [1 x 3] target coordinates
'''
rng = np.random.default_rng()
for _ in range(nblocks):
order = np.arange(ntargets) + 1 # target indices, starting from 1
rng.shuffle(order)
for t in range(ntargets):
idx = order[t]
theta = 2 * np.pi * idx / ntargets
pos = np.array(
[distance * np.cos(theta), 0, distance * np.sin(theta)]).T
yield [idx], [pos + origin]
@staticmethod
def centerout_2D(nblocks=100, ntargets=8, distance=10, origin=(0, 0, 0)):
'''
Pairs of central targets at the origin and peripheral targets centered around the origin
Returns
-------
[nblocks*ntargets x 1] array of tuples containing trial indices and [2 x 3] target coordinates
'''
gen = ScreenTargetCapture.out_2D(nblocks, ntargets, distance, origin)
for _ in range(nblocks * ntargets):
idx, pos = next(gen)
targs = np.zeros([2, 3]) + origin
targs[1, :] = pos[0]
indices = np.zeros([2, 1])
indices[1] = idx
yield indices, targs
@staticmethod
def centeroutback_2D(nblocks=100,
ntargets=8,
distance=10,
origin=(0, 0, 0)):
'''
Triplets of central targets, peripheral targets, and central targets
Returns
-------
[nblocks*ntargets x 1] array of tuples containing trial indices and [3 x 3] target coordinates
'''
gen = ScreenTargetCapture.out_2D(nblocks, ntargets, distance, origin)
for _ in range(nblocks * ntargets):
idx, pos = next(gen)
targs = np.zeros([3, 3]) + origin
targs[1, :] = pos[0]
indices = np.zeros([3, 1])
indices[1] = idx
yield indices, targs
@staticmethod
def rand_target_chain_2D(ntrials=100,
chain_length=1,
boundaries=(-12, 12, -12, 12)):
'''
Generates a sequence of 2D (x and z) target pairs.
Parameters
----------
ntrials : int
The number of target chains in the sequence.
chain_length : int
The number of targets in each chain
boundaries: 4 element Tuple
The limits of the allowed target locations (-x, x, -z, z)
Returns
-------
[ntrials x chain_length x 3] array of target coordinates
'''
rng = np.random.default_rng()
idx = 0
for t in range(ntrials):
# Choose a random sequence of points within the boundaries
pts = rng.uniform(size=(chain_length, 3)) * (
(boundaries[1] - boundaries[0]), 0,
(boundaries[3] - boundaries[2]))
pts = pts + (boundaries[0], 0, boundaries[2])
yield idx + np.arange(chain_length), pts
idx += chain_length
@staticmethod
def rand_target_chain_3D(ntrials=100,
chain_length=1,
boundaries=(-12, 12, -10, 10, -12, 12)):
'''
Generates a sequence of 3D target pairs.
Parameters
----------
ntrials : int
The number of target chains in the sequence.
chain_length : int
The number of targets in each chain
boundaries: 6 element Tuple
The limits of the allowed target locations (-x, x, -y, y, -z, z)
Returns
-------
[ntrials x chain_length x 3] array of target coordinates
'''
rng = np.random.default_rng()
idx = 0
for t in range(ntrials):
# Choose a random sequence of points within the boundaries
pts = rng.uniform(size=(chain_length, 3)) * (
(boundaries[1] - boundaries[0]),
(boundaries[3] - boundaries[2]),
(boundaries[5] - boundaries[4]))
pts = pts + (boundaries[0], boundaries[2], boundaries[4])
yield idx + np.arange(chain_length), pts
idx += chain_length
