class PenaltyAudio(traits.HasTraits): ''' Play a sound in any penalty state. Have to define a new _start method for each different penalty state that might occur. ''' files = list(reversed([f for f in os.listdir(audio_path) if '.wav' in f])) penalty_sound = traits.OptionsList( files, desc="File in riglib/audio to play on each penalty") def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.penalty_player = AudioPlayer(self.penalty_sound) def _start_hold_penalty(self): if hasattr(super(), '_start_hold_penalty'): super()._start_hold_penalty() self.penalty_player.play() def _start_delay_penalty(self): if hasattr(super(), '_start_delay_penalty'): super()._start_delay_penalty() self.penalty_player.play() def _start_reach_penalty(self): if hasattr(super(), '_start_reach_penalty'): super()._start_reach_penalty() self.penalty_player.play() def _start_timeout_penalty(self): if hasattr(super(), '_start_timeout_penalty'): super()._start_timeout_penalty() self.penalty_player.play()
class EndPostureFeedbackController(BMILoop, traits.HasTraits): ssm_type_options = bmi_ssm_options ssm_type = traits.OptionsList(*bmi_ssm_options, bmi3d_input_options=bmi_ssm_options) def load_decoder(self): self.ssm = StateSpaceEndptVel2D() 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): self.extractor = DummyExtractor() self._add_feature_extractor_dtype()
class RewardAudio(traits.HasTraits): ''' Play a sound in any reward state. Need to add other reward states you want to be included. ''' files = [f for f in os.listdir(audio_path) if '.wav' in f] reward_sound = traits.OptionsList( files, desc="File in riglib/audio to play on each reward") def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.reward_player = AudioPlayer(self.reward_sound) def _start_reward(self): if hasattr(super(), '_start_reward'): super()._start_reward() self.reward_player.play()
class EndPostureFeedbackController(BMILoop, traits.HasTraits): ssm_type_options = bmi_ssm_options ssm_type = traits.OptionsList(*bmi_ssm_options, bmi3d_input_options=bmi_ssm_options) def load_decoder(self): from db.namelist import bmi_state_space_models # from config import config # with open(os.path.join(config.log_dir, 'EndPostureFeedbackController'), 'w') as fh: # fh.write('%s' % self.ssm_type) self.ssm = bmi_state_space_models[self.ssm_type] 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): self.extractor = DummyExtractor() self._add_feature_extractor_dtype()
class ArmPlant(Window): ''' This task creates a RobotArm object and allows it to move around the screen based on either joint or endpoint positions. There is a spherical cursor at the end of the arm. The links of the arm can be visible or hidden. ''' background = (0, 0, 0, 1) 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) arm_class = traits.OptionsList(*plantlist, bmi3d_input_options=plantlist.keys()) starting_pos = (5, 0, 5) def __init__(self, *args, **kwargs): super(ArmPlant, self).__init__(*args, **kwargs) self.cursor_visible = True # Initialize the arm self.arm = ik.test_3d self.arm_vis_prev = True if self.arm_class == 'CursorPlant': pass else: self.dtype.append(('joint_angles', 'f8', (self.arm.num_joints, ))) self.dtype.append(('arm_visible', 'f8', (1, ))) self.add_model(self.arm) ## Declare cursor self.dtype.append(('cursor', 'f8', (3, ))) self.cursor = Sphere(radius=self.cursor_radius, color=self.cursor_color) self.add_model(self.cursor) self.cursor.translate(*self.arm.get_endpoint_pos(), reset=True) def _cycle(self): ''' Calls any update functions necessary and redraws screen. Runs 60x per second by default. ''' ## Run graphics commands to show/hide the arm if the visibility has changed if self.arm_class != 'CursorPlant': 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.update_cursor() if self.cursor_visible: self.task_data['cursor'] = self.cursor.xfm.move.copy() else: #if the cursor is not visible, write NaNs into cursor location saved in file self.task_data['cursor'] = np.array([np.nan, np.nan, np.nan]) if self.arm_class != 'CursorPlant': if self.arm_visible: self.task_data['arm_visible'] = 1 else: self.task_data['arm_visible'] = 0 super(ArmPlant, self)._cycle() ## Functions to move the cursor using keyboard/mouse input def get_mouse_events(self): import pygame events = [] for btn in pygame.event.get( (pygame.MOUSEBUTTONDOWN, pygame.MOUSEBUTTONUP)): events = events + [btn.button] return events def get_key_events(self): import pygame return pygame.key.get_pressed() def move_arm(self): ''' allows use of keyboard keys to test movement of arm. Use QW/OP for joint movements, arrow keys for endpoint movements ''' import pygame keys = self.get_key_events() joint_speed = (np.pi / 6) / 60 hand_speed = .2 x, y, z = self.arm.get_endpoint_pos() if keys[pygame.K_RIGHT]: x = x - hand_speed self.arm.set_endpoint_pos(np.array([x, 0, z])) if keys[pygame.K_LEFT]: x = x + hand_speed self.arm.set_endpoint_pos(np.array([x, 0, z])) if keys[pygame.K_DOWN]: z = z - hand_speed self.arm.set_endpoint_pos(np.array([x, 0, z])) if keys[pygame.K_UP]: z = z + hand_speed self.arm.set_endpoint_pos(np.array([x, 0, z])) if self.arm.num_joints == 2: xz, xy = self.get_arm_joints() e = np.array([xz[0], xy[0]]) s = np.array([xz[1], xy[1]]) if keys[pygame.K_q]: s = s - joint_speed self.set_arm_joints([e[0], s[0]], [e[1], s[1]]) if keys[pygame.K_w]: s = s + joint_speed self.set_arm_joints([e[0], s[0]], [e[1], s[1]]) if keys[pygame.K_o]: e = e - joint_speed self.set_arm_joints([e[0], s[0]], [e[1], s[1]]) if keys[pygame.K_p]: e = e + joint_speed self.set_arm_joints([e[0], s[0]], [e[1], s[1]]) if self.arm.num_joints == 4: jts = self.get_arm_joints() keyspressed = [ keys[pygame.K_q], keys[pygame.K_w], keys[pygame.K_e], keys[pygame.K_r] ] for i in range(self.arm.num_joints): if keyspressed[i]: jts[i] = jts[i] + joint_speed self.set_arm_joints(jts) def get_cursor_location(self): return self.arm.get_endpoint_pos() def set_arm_endpoint(self, pt, **kwargs): self.arm.set_endpoint_pos(pt, **kwargs) def set_arm_joints(self, angle_xz, angle_xy): self.arm.set_intrinsic_coordinates(angle_xz, angle_xy) 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() 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)
class FreeChoiceFA(FactorBMIBase): ''' Task where the virtual plant starts in configuration sampled from a discrete set and resets every trial ''' sequence_generators = [ 'centerout_2D_discrete_w_free_choice', 'centerout_2D_discrete_w_free_choice_v2', 'centerout_2D_discrete_w_free_choices_evenly_spaced' ] #sequence_generators = ['centerout_2D_discrete'] input_type_list = [ 'shared', 'private', 'shared_scaled', 'private_scaled', 'all', 'all_scaled_by_shar', 'sc_shared+unsc_priv', 'sc_shared+sc_priv', 'main_shared', 'main_sc_shared', 'main_sc_private', 'main_sc_shar+unsc_priv', 'main_sc_shar+sc_priv', 'pca', 'split' ] input_type_0 = traits.OptionsList(*input_type_list, bmi3d_input_options=input_type_list) color_0 = traits.OptionsList(*target_colors.keys(), bmi3d_input_options=target_colors.keys()) input_type_1 = traits.OptionsList(*input_type_list, bmi3d_input_options=input_type_list) color_1 = traits.OptionsList(*target_colors.keys(), bmi3d_input_options=target_colors.keys()) input_type_2 = traits.OptionsList(*input_type_list, bmi3d_input_options=input_type_list) color_2 = traits.OptionsList(*target_colors.keys(), bmi3d_input_options=target_colors.keys()) input_type_3 = traits.OptionsList(*input_type_list, bmi3d_input_options=input_type_list) color_3 = traits.OptionsList(*target_colors.keys(), bmi3d_input_options=target_colors.keys()) choice_assist = traits.Float(0.) target_assist = traits.Float(0.) choice_target_rad = traits.Float(2.) status = dict(wait=dict(start_trial="targ_transition", stop=None), pre_choice_orig=dict(enter_orig='choice_target', timeout='timeout_penalty', stop=None), choice_target=dict(enter_choice_target='targ_transition', timeout='timeout_penalty', stop=None), 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", make_choice='pre_choice_orig'), timeout_penalty=dict(timeout_penalty_end="targ_transition"), hold_penalty=dict(hold_penalty_end="targ_transition"), reward=dict(reward_end="wait")) hidden_traits = [ 'arm_hide_rate', 'arm_visible', 'hold_penalty_time', 'rand_start', 'reset', 'window_size', 'assist_level', 'assist_level_time', 'plant_hide_rate', 'plant_visible', 'show_environment', 'trials_per_reward' ] def __init__(self, *args, **kwargs): super(FreeChoiceFA, self).__init__(*args, **kwargs) seq_params = eval(kwargs.pop('seq_params', '{}')) print 'SEQ PARAMS: ', seq_params, type(seq_params) self.choice_per_n_blocks = seq_params.pop('blocks_per_free_choice', 1) self.n_free_choices = seq_params.pop('n_free_choices', 2) self.n_targets = seq_params.pop('ntargets', 8) self.input_type_dict = dict() self.input_type_dict[0] = self.input_type_0 self.input_type_dict[0, 'color'] = target_colors[self.color_0] self.input_type_dict[1] = self.input_type_1 self.input_type_dict[1, 'color'] = target_colors[self.color_1] self.input_type_dict[2] = self.input_type_2 self.input_type_dict[2, 'color'] = target_colors[self.color_2] self.input_type_dict[3] = self.input_type_3 self.input_type_dict[3, 'color'] = target_colors[self.color_3] # Instantiate the choice targets self.choices_targ_list = [] for c in range(self.n_free_choices): self.choices_targ_list.append( target_graphics.VirtualCircularTarget( target_radius=self.choice_target_rad, target_color=self.input_type_dict[c, 'color'])) for c in self.choices_targ_list: for model in c.graphics_models: self.add_model(model) self.subblock_cnt = 0 self.subblock_end = self.choice_per_n_blocks * self.n_targets self.choice_made = 0 self.choice_ts = 0 self.chosen_input_ix = -1 self.choice_locs = np.zeros((self.n_free_choices, 3)) def init(self): self.add_dtype('trial_type', np.str_, 16) self.add_dtype('choice_ix', 'f8', (1, )) self.add_dtype('choice_targ_loc', 'f8', (self.n_free_choices, 3)) super(FreeChoiceFA, self).init() def _start_pre_choice_orig(self): target = self.targets[0] target.move_to_position(np.array([0., 0., 0.])) target.cue_trial_start() self.chosen_input_ix = -1 def _start_timeout_penalty(self): #hide targets for target in self.targets: target.hide() for target in self.choices_targ_list: target.hide() self.tries += 1 self.target_index = -1 def _test_enter_orig(self, ts): cursor_pos = self.plant.get_endpoint_pos() d = np.linalg.norm(cursor_pos) return d <= self.target_radius def update_level(self): pass def create_goal_calculator(self): self.goal_calculator = Choice_Goal_Calc(self.decoder.ssm) def get_target_BMI_state(self, *args): ''' Run the goal calculator to determine the target state of the task ''' target_state = self.goal_calculator(self.targs, self.choice_locs, self.choice_asst_ix, self.target_index, self.state) return np.array(target_state).reshape(-1, 1) def _parse_next_trial(self): print 'parse next: ', self.next_trial[2][0] pairs = self.next_trial[0] self.targs = pairs[:, :, 1] self.choice_locs = pairs[:, :, 0] self.choice_asst_ix = self.next_trial[1][0] self.choice_instructed = self.next_trial[2][0] if self.subblock_cnt >= self.subblock_end: self.choice_made = 0 self.subblock_cnt = 0 def _test_make_choice(self, ts): return not self.choice_made def _cycle(self): self.task_data['trial_type'] = self.choice_instructed self.task_data['choice_ix'] = self.chosen_input_ix self.task_data['choice_targ_loc'] = self.choice_locs super(FreeChoiceFA, self)._cycle() def _start_choice_target(self): self.choice_ts = 0 if self.choice_instructed == 'Free': for ic, c in enumerate(self.choices_targ_list): #move a target to current location (target1 and target2 alternate moving) and set location attribute c.move_to_position(self.choice_locs[ic, :]) c.sphere.color = self.input_type_dict[ic, 'color'] c.show() elif self.choice_instructed == 'Instructed': ic = self.choice_asst_ix c = self.choices_targ_list[ic] c.move_to_position(self.choice_locs[ic, :]) c.sphere.color = self.input_type_dict[ic, 'color'] c.show() target = self.targets[0] target.hide() self.choice_ts = 0 def _start_target(self): super(FreeChoiceFA, self)._start_target() self.current_assist_level = self.target_assist for ic, c in enumerate(self.choices_targ_list): c.hide() def _test_enter_choice_target(self, ts): cursor_pos = self.plant.get_endpoint_pos() enter_targ = 0 for ic, c in enumerate(self.choice_locs): d = np.linalg.norm(cursor_pos - c) if d <= self.choice_target_rad: #NOTE, gets in if CENTER of cursor is in target (not entire cursor) enter_targ += 1 #Set chosen as new input: self.chosen_input_ix = ic self.decoder.filt.FA_input = self.input_type_dict[ic] print 'trial: ', self.decoder.filt.FA_input, self.choice_instructed #Declare that choice has been made: self.choice_made = 1 #Change color of cursor: sph = self.plant.graphics_models[0] sph.color = self.input_type_dict[ic, 'color'] return enter_targ > 0 def _test_trial_incomplete(self, ts): if self.choice_made == 0: return False else: return (not self._test_trial_complete(ts)) and (self.tries < self.max_attempts) def _start_reward(self): self.subblock_cnt += 1 super(FreeChoiceFA, self)._start_reward() @staticmethod def centerout_2D_discrete_w_free_choice(nblocks=100, ntargets=8, boundaries=(-18, 18, -12, 12), distance=10, n_free_choices=2, blocks_per_free_choice=1, percent_instructed=50.): return True @staticmethod def centerout_2D_discrete_w_free_choice_v2(nblocks=100, ntargets=8, boundaries=(-18, 18, -12, 12), distance=10, n_free_choices=2, blocks_per_free_choice=1, percent_instructed=50.): ''' Generates a sequence of 2D (x and z) target pairs with the first target always at the origin and a sequence of 2D (x and z) target locations for nblocks of free choices where the location of each choice changes. 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. n_free_choices: number of choices. Returns ------- ([nblocks x ntargets x 2 x 3], [nblocks x n_free_choices x 3]) array of 1) pairs of target locations and 2) set of free choices ''' # Choose a random sequence of points on the edge of a circle of radius # "distance" theta = [] theta_choice = [] ix_choice_assist = [] ix_choice_instructed = [] for i in range(nblocks): temp_ = [] for j in range(blocks_per_free_choice): temp = np.arange(0, 2 * np.pi, 2 * np.pi / ntargets) np.random.shuffle(temp) temp_ = temp_ + list(temp) theta.append(temp_) temp2 = np.arange(0, np.pi / 2., np.pi / 2. / n_free_choices) + ( np.pi / 4.) + (np.pi / 2.) * np.random.randint(0, 2) temp3 = np.random.randint(0, n_free_choices) temp4 = np.random.rand() if temp4 < percent_instructed / 100.: ix_choice_instructed.append('Instructed') else: ix_choice_instructed.append('Free') np.random.shuffle(temp2) theta_choice.append(temp2) ix_choice_assist.append(temp3) theta = np.vstack(theta) theta_choice = np.vstack(theta_choice) #nblocks x n_free_choices ix_choice_assist = np.array(ix_choice_assist) ix_choice_instructed = np.array(ix_choice_instructed) #### calculate targets: x = distance * np.cos(theta) y = np.zeros((nblocks, ntargets * blocks_per_free_choice)) z = distance * np.sin(theta) pairs = np.zeros([nblocks, ntargets * blocks_per_free_choice, 2, 3]) pairs[:, :, 1, :] = np.dstack([x, y, z]) #### calculate free choices: x = distance * np.cos(theta_choice) y = np.zeros((nblocks, n_free_choices)) z = distance * np.sin(theta_choice) choice = np.zeros((nblocks, n_free_choices, 3)) choice = np.dstack((x, y, z)) g = [] for i in range(nblocks): chz = choice[i, :, :] chz_assist = ix_choice_assist[i] type_chz = ix_choice_instructed[i] for j in range(ntargets * blocks_per_free_choice): tg = pairs[i, j, :, :] g.append((np.dstack((chz, tg)), [chz_assist], [type_chz])) return g @staticmethod def centerout_2D_discrete_w_free_choices_evenly_spaced( nblocks=100, ntargets=8, boundaries=(-18, 18, -12, 12), distance=10, n_free_choices=2, blocks_per_free_choice=1, percent_instructed=50., choice_targ_ang=30.): ''' Generates a sequence of 2D (x and z) target pairs with the first target always at the origin and a sequence of 2D (x and z) target locations for nblocks of free choices where the location of each choice changes -- specifically the location of the free choices are opposite the previous target, and spaced at 30 degree angle offsets 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. n_free_choices: number of choices. Returns ------- ([nblocks x ntargets x 2 x 3], [nblocks x n_free_choices x 3]) array of 1) pairs of target locations and 2) set of free choices ''' # Choose a random sequence of points on the edge of a circle of radius # "distance" theta = [] theta_choice = [] ix_choice_assist = [] ix_choice_instructed = [] last_targ_ang_ = 0. for i in range(nblocks): temp_ = [] for j in range(blocks_per_free_choice): temp = np.arange(0, 2 * np.pi, 2 * np.pi / ntargets) np.random.shuffle(temp) temp_ = temp_ + list(temp) theta.append(temp_) ang = np.array([ -choice_targ_ang * (np.pi / 180), choice_targ_ang * (np.pi / 180.) ]) temp2 = ang + np.pi + last_targ_ang_ last_targ_ang_ = temp_[-1] temp3 = np.random.randint(0, n_free_choices) temp4 = np.random.rand() if temp4 < percent_instructed / 100.: ix_choice_instructed.append('Instructed') else: ix_choice_instructed.append('Free') np.random.shuffle(temp2) theta_choice.append(temp2) ix_choice_assist.append(temp3) theta = np.vstack(theta) theta_choice = np.vstack(theta_choice) #nblocks x n_free_choices ix_choice_assist = np.array(ix_choice_assist) ix_choice_instructed = np.array(ix_choice_instructed) #### calculate targets: x = distance * np.cos(theta) y = np.zeros((nblocks, ntargets * blocks_per_free_choice)) z = distance * np.sin(theta) pairs = np.zeros([nblocks, ntargets * blocks_per_free_choice, 2, 3]) pairs[:, :, 1, :] = np.dstack([x, y, z]) #### calculate free choices: x = distance * np.cos(theta_choice) y = np.zeros((nblocks, n_free_choices)) z = distance * np.sin(theta_choice) choice = np.zeros((nblocks, n_free_choices, 3)) choice = np.dstack((x, y, z)) g = [] for i in range(nblocks): chz = choice[i, :, :] chz_assist = ix_choice_assist[i] type_chz = ix_choice_instructed[i] for j in range(ntargets * blocks_per_free_choice): tg = pairs[i, j, :, :] g.append((np.dstack((chz, tg)), [chz_assist], [type_chz])) return g
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 LFP_Mod_plus_MC_reach(LFP_Mod_plus_MC_hold): mc_cursor_radius = traits.Float(.5, desc="Radius of cursor") mc_target_radius = traits.Float(3, desc="Radius of MC target") mc_cursor_color = (.5,0,.5,1) mc_plant_type_options = plantlist.keys() mc_plant_type = traits.OptionsList(*plantlist, bmi3d_input_options=plantlist.keys()) origin_hold_time = traits.Float(.2, desc="Hold time in center") mc_periph_holdtime = traits.Float(.2, desc="Hold time in center") mc_timeout_time = traits.Float(10, desc="Time allowed to go between targets") exclude_parent_traits = ['goal_cache_block'] #redefine this to NOT include marker_num, marker_count marker_num = traits.Int(14,desc='Index') marker_count = traits.Int(16,desc='Num of markers') scale_factor = 3.0 #scale factor for converting hand movement to screen movement (1cm hand movement = 3.5cm cursor movement) wait_flag = 1 # 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. limit2d = 1 # state_file = open("/home/helene/preeya/tot_pw.txt","w") state_cnt = 0 status = dict( wait = dict(start_trial="origin", stop=None), origin = dict(enter_origin="origin_hold", stop=None), origin_hold = dict(origin_hold_complete="lfp_target",leave_origin="hold_penalty", stop=None), lfp_target = dict(enter_lfp_target="lfp_hold", leave_origin="hold_penalty", powercap_penalty="powercap_penalty", stop=None), lfp_hold = dict(leave_early="lfp_target", lfp_hold_complete="mc_target", leave_origin="hold_penalty",powercap_penalty="powercap_penalty"), mc_target = dict(enter_mc_target='mc_hold',mc_timeout="timeout_penalty", stop=None), mc_hold = dict(leave_periph_early='hold_penalty',mc_hold_complete="reward"), powercap_penalty = dict(powercap_penalty_end="origin"), timeout_penalty = dict(timeout_penalty_end="wait"), hold_penalty = dict(hold_penalty_end="origin"), reward = dict(reward_end="wait"), ) static_states = ['origin'] # states in which the decoder is not run trial_end_states = ['reward', 'timeout_penalty'] lfp_cursor_on = ['lfp_target', 'lfp_hold', 'reward'] sequence_generators = ['lfp_mod_plus_MC_reach', 'lfp_mod_plus_MC_reach_INV'] def __init__(self, *args, **kwargs): # import pickle # decoder = pickle.load(open('/storage/decoders/cart20141216_03_cart_new2015_2.pkl')) # self.decoder = decoder super(LFP_Mod_plus_MC_reach, self).__init__(*args, **kwargs) mc_origin = VirtualCircularTarget(target_radius=self.mc_target_radius, target_color=RED) mc_periph = VirtualCircularTarget(target_radius=self.mc_target_radius, target_color=RED) lfp_target = VirtualSquareTarget(target_radius=self.lfp_target_rad, target_color=self.lfp_target_color) self.targets = [lfp_target, mc_origin, mc_periph] # #Should be unnecessary: # for target in self.targets: # for model in target.graphics_models: # self.add_model(model) # self.lfp_plant = plantlist[self.lfp_plant_type] # if hasattr(self.lfp_plant, 'graphics_models'): # for model in self.lfp_plant.graphics_models: # self.add_model(model) # self.mc_plant = plantlist[self.mc_plant_type] # if hasattr(self.mc_plant, 'graphics_models'): # for model in self.mc_plant.graphics_models: # self.add_model(model) def _parse_next_trial(self): t = self.next_trial self.lfp_targ = t['lfp'] self.mc_targ_orig = t['origin'] self.mc_targ_periph = t['periph'] def _start_mc_target(self): #Turn off LFP things self.lfp_plant.turn_off() self.targets[0].hide() self.targets[1].hide() target = self.targets[2] #MC target self.target_location_mc = self.mc_targ_periph target.move_to_position(self.target_location_mc) target.cue_trial_start() def _test_enter_mc_target(self,ts): cursor_pos = self.mc_plant.get_endpoint_pos() d = np.linalg.norm(cursor_pos - self.target_location_mc) return d <= (self.mc_target_radius - self.mc_cursor_radius) def _test_mc_timeout(self, ts): return ts>self.mc_timeout_time def _test_leave_periph_early(self, ts): cursor_pos = self.mc_plant.get_endpoint_pos() d = np.linalg.norm(cursor_pos - self.target_location_mc) rad = self.mc_target_radius - self.mc_cursor_radius return d > rad def _test_mc_hold_complete(self, ts): return ts>self.mc_periph_holdtime def _timeout_penalty_end(self, ts): print 'timeout', ts #return ts > 1. return True def _end_mc_hold(self): self.targets[2].cue_trial_end_success() # def _cycle(self): # if self.state_cnt < 3600*3: # self.state_cnt +=1 # s = "%s\n" % self.state # self.state_file.write(str(s)) # if self.state_cnt == 3600*3: # self.state_file.close() # super(LFP_Mod_plus_MC_reach, self)._cycle() def _start_reward(self): super(LFP_Mod_plus_MC_reach, self)._start_reward() lfp_targ = self.targets[0] mc_orig = self.targets[1] lfp_targ.hide() mc_orig.hide() @staticmethod def lfp_mod_plus_MC_reach(nblocks=100, boundaries=(-18,18,-12,12), xaxis=-8, target_distance=6, n_mc_targets=4, mc_target_angle_offset=0,**kwargs): 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) lfp = np.vstack([x, y, z]).T origin = np.zeros(( lfp.shape )) theta = [] for i in range(nblocks*4): temp = np.arange(0, 2*np.pi, 2*np.pi/float(n_mc_targets)) np.random.shuffle(temp) theta = theta + [temp] theta = np.hstack(theta) theta = theta + (mc_target_angle_offset*(np.pi/180.)) x = target_distance*np.cos(theta) y = np.zeros(len(theta)) z = target_distance*np.sin(theta) periph = np.vstack([x, y, z]).T it = iter([dict(lfp=lfp[i,:], origin=origin[i,:], periph=periph[i,:]) for i in range(lfp.shape[0])]) if ('return_arrays' in kwargs.keys()) and kwargs['return_arrays']==True: return lfp, origin, periph else: return it @staticmethod def lfp_mod_plus_MC_reach_INV(nblocks=100, boundaries=(-18,18,-12,12), xaxis=-8, target_distance=6, n_mc_targets=4, mc_target_angle_offset=0): kw = dict(return_arrays=True) lfp, origin, periph = LFP_Mod_plus_MC_reach.lfp_mod_plus_MC_reach(nblocks=nblocks, boundaries=boundaries, xaxis=xaxis, target_distance=target_distance, n_mc_targets=n_mc_targets, mc_target_angle_offset=mc_target_angle_offset,**kw) #Invert LFP: lfp[:,2] = -1.0*lfp[:,2] it = iter([dict(lfp=lfp[i,:], origin=origin[i,:], periph=periph[i,:]) for i in range(lfp.shape[0])]) return it
class LFP_Mod_plus_MC_hold(LFP_Mod): mc_cursor_radius = traits.Float(.5, desc="Radius of cursor") mc_target_radius = traits.Float(3, desc="Radius of MC target") mc_cursor_color = (.5,0,.5,1) mc_plant_type_options = plantlist.keys() mc_plant_type = traits.OptionsList(*plantlist, bmi3d_input_options=plantlist.keys()) origin_hold_time = traits.Float(.2, desc="Hold time in center") exclude_parent_traits = ['goal_cache_block'] #redefine this to NOT include marker_num, marker_count marker_num = traits.Int(14,desc='Index') marker_count = traits.Int(16,desc='Num of markers') joystick_method = traits.Float(1,desc="1: Normal velocity, 0: Position control") joystick_speed = traits.Float(20, desc="Radius of cursor") move_while_in_center = traits.Float(1, desc="1 = update plant while in lfp_target, lfp_hold, 0 = don't update in these states") scale_factor = 3.0 #scale factor for converting hand movement to screen movement (1cm hand movement = 3.5cm cursor movement) wait_flag = 1 # 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. limit2d = 1 status = dict( wait = dict(start_trial="origin", stop=None), origin = dict(enter_origin="origin_hold", stop=None), origin_hold = dict(origin_hold_complete="lfp_target",leave_origin="hold_penalty", stop=None), lfp_target = dict(enter_lfp_target="lfp_hold", leave_origin="hold_penalty", powercap_penalty="powercap_penalty", stop=None), lfp_hold = dict(leave_early="lfp_target", lfp_hold_complete="reward", leave_origin="hold_penalty", powercap_penalty="powercap_penalty",stop=None), powercap_penalty = dict(powercap_penalty_end="origin"), hold_penalty = dict(hold_penalty_end="origin",stop=None), reward = dict(reward_end="wait") ) static_states = ['origin'] # states in which the decoder is not run trial_end_states = ['reward'] lfp_cursor_on = ['lfp_target', 'lfp_hold', 'reward'] sequence_generators = ['lfp_mod_4targ_plus_mc_orig'] def __init__(self, *args, **kwargs): super(LFP_Mod_plus_MC_hold, self).__init__(*args, **kwargs) if self.move_while_in_center>0: self.no_plant_update_states = [] else: self.no_plant_update_states = ['lfp_target', 'lfp_hold'] mc_origin = VirtualCircularTarget(target_radius=self.mc_target_radius, target_color=RED) lfp_target = VirtualSquareTarget(target_radius=self.lfp_target_rad, target_color=self.lfp_target_color) self.targets = [lfp_target, mc_origin] self.mc_plant = plantlist[self.mc_plant_type] if hasattr(self.mc_plant, 'graphics_models'): for model in self.mc_plant.graphics_models: self.add_model(model) # Declare any plant attributes which must be saved to the HDF file at the _cycle rate for attr in self.mc_plant.hdf_attrs: self.add_dtype(*attr) self.target_location_mc = np.array([-100, -100, -100]) self.manual_control_type = None self.current_pt=np.zeros([3]) #keep track of current pt self.last_pt=np.zeros([3]) def init(self): self.add_dtype('mc_targ', 'f8', (3,)) ###ADD BACK super(LFP_Mod_plus_MC_hold, self).init() def _cycle(self): ''' Calls any update functions necessary and redraws screen. Runs 60x per second. ''' self.task_data['mc_targ'] = self.target_location_mc.copy() mc_plant_data = self.mc_plant.get_data_to_save() for key in mc_plant_data: self.task_data[key] = mc_plant_data[key] super(LFP_Mod_plus_MC_hold, self)._cycle() def _parse_next_trial(self): t = self.next_trial self.lfp_targ = t['lfp'] self.mc_targ_orig = t['origin'] def _start_origin(self): if self.wait_flag: self.origin_hold_time_store = self.origin_hold_time self.origin_hold_time = 3 self.wait_flag = 0 else: self.origin_hold_time = self.origin_hold_time_store #only 1 target: target = self.targets[1] #Origin self.target_location_mc = self.mc_targ_orig #Origin target.move_to_position(self.target_location_mc) target.cue_trial_start() #Turn off lfp things self.lfp_plant.turn_off() self.targets[0].hide() def _start_lfp_target(self): #only 1 target: target = self.targets[0] #LFP target self.target_location_lfp = self.lfp_targ #LFP target target.move_to_position(self.target_location_lfp) target.cue_trial_start() self.lfp_plant.turn_on() def _start_lfp_hold(self): #make next target visible unless this is the final target in the trial pass def _start_hold_penalty(self): #hide targets for target in self.targets: target.hide() self.tries += 1 self.target_index = -1 #Turn off lfp things self.lfp_plant.turn_off() self.targets[0].hide() def _end_origin(self): self.targets[1].cue_trial_end_success() def _test_enter_origin(self, ts): cursor_pos = self.mc_plant.get_endpoint_pos() d = np.linalg.norm(cursor_pos - self.target_location_mc) return d <= (self.mc_target_radius - self.mc_cursor_radius) # def _test_origin_timeout(self, ts): # return ts>self.timeout_time def _test_leave_origin(self, ts): if self.manual_control_type == 'joystick': if hasattr(self,'touch'): if self.touch <0.5: self.last_touch_zero_event = time.time() return True cursor_pos = self.mc_plant.get_endpoint_pos() d = np.linalg.norm(cursor_pos - self.target_location_mc) return d > (self.mc_target_radius - self.mc_cursor_radius) def _test_origin_hold_complete(self,ts): return ts>=self.origin_hold_time # def _test_enter_lfp_target(self, ts): # ''' # return true if the distance between center of cursor and target is smaller than the cursor radius # ''' # cursor_pos = self.lfp_plant.get_endpoint_pos() # cursor_pos = [cursor_pos[0], cursor_pos[2]] # targ_loc = np.array([self.target_location_lfp[0], self.target_location_lfp[2]]) # d = np.linalg.norm(cursor_pos - targ_loc) # return 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() # d = np.linalg.norm(cursor_pos - self.target_location_lfp) # rad = self.lfp_target_rad - self.lfp_cursor_rad # return d > rad def _test_hold_penalty_end(self, ts): return ts>self.hold_penalty_time def _end_lfp_hold(self): # change current target color to green self.targets[0].cue_trial_end_success() def move_plant(self): if self.state in self.lfp_cursor_on: 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 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 #Sets the plant configuration based on motiontracker data. For manual control, uses #motiontracker data. If no motiontracker data available, returns None''' #get data from motion tracker- take average of all data points since last poll if self.state in self.no_plant_update_states: pt = np.array([0, 0, 0]) print 'no update' else: if self.manual_control_type == 'motiondata': pt = self.motiondata.get() if len(pt) > 0: pt = pt[:, self.marker_num, :] conds = pt[:, 3] inds = np.nonzero((conds>=0) & (conds!=4))[0] if len(inds) > 0: pt = pt[inds,:3] #scale actual movement to desired amount of screen movement pt = pt.mean(0) * self.scale_factor #Set y coordinate to 0 for 2D tasks if self.limit2d: #pt[1] = 0 pt[2] = pt[1].copy() pt[1] = 0 pt[1] = pt[1]*2 # Return cursor location self.no_data_count = 0 pt = pt * mm_per_cm #self.convert_to_cm(pt) else: #if no usable data self.no_data_count += 1 pt = None else: #if no new data self.no_data_count +=1 pt = None elif self.manual_control_type == 'joystick': pt = self.joystick.get() #if touch sensor on: try: self.touch = pt[-1][0][2] except: pass 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: #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' # calib = [ 0.487, 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 pt = self.current_pt #self.plant.set_endpoint_pos(self.current_pt) self.last_pt = self.current_pt.copy() elif self.manual_control_type == None: pt = None try: pt0 = self.motiondata.get() self.manual_control_type='motiondata' except: print 'not motiondata' try: pt0 = self.joystick.get() self.manual_control_type = 'joystick' except: print 'not joystick data' # Set the plant's endpoint to the position determined by the motiontracker, unless there is no data available if self.manual_control_type is not None: if pt is not None and len(pt)>0: self.mc_plant.set_endpoint_pos(pt) @staticmethod def lfp_mod_4targ_plus_mc_orig(nblocks=100, boundaries=(-18,18,-12,12), xaxis=-8): ''' See lfp_mod_4targ for lfp target explanation ''' 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) lfp = np.vstack([x, y, z]).T origin = np.zeros(( lfp.shape )) it = iter([dict(lfp=lfp[i,:], origin=origin[i,:]) for i in range(lfp.shape[0])]) return it
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 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 Optitrack(traits.HasTraits): ''' Enable reading of raw motiontracker data from Optitrack system Requires the natnet library from https://github.com/leoscholl/python_natnet To be used as a feature with the ManualControl task for the time being. However, ideally this would be implemented as a decoder :) ''' optitrack_feature = traits.OptionsList(("rigid body", "skeleton", "marker")) smooth_features = traits.Int(1, desc="How many features to average") scale = traits.Float(DEFAULT_SCALE, desc="Control scale factor") offset = traits.Array(value=DEFAULT_OFFSET, desc="Control offset") hidden_traits = ['optitrack_feature', 'smooth_features'] def init(self): ''' Secondary init function. See riglib.experiment.Experiment.init() Prior to starting the task, this 'init' sets up the DataSource for interacting with the motion tracker system and registers the source with the SinkRegister so that the data gets saved to file as it is collected. ''' # Start the natnet client and recording import natnet now = datetime.now() local_path = "C:/Users/Orsborn Lab/Documents" session = "OptiTrack/Session " + now.strftime("%Y-%m-%d") take = now.strftime("Take %Y-%m-%d %H:%M:%S") logger = Logger(take) try: client = natnet.Client.connect(logger=logger) if self.saveid is not None: take += " (%d)" % self.saveid client.set_session(os.path.join(local_path, session)) client.set_take(take) self.filename = os.path.join(session, take + '.tak') client._send_command_and_wait("LiveMode") time.sleep(0.1) if client.start_recording(): self.optitrack_status = 'recording' else: self.optitrack_status = 'streaming' except natnet.DiscoveryError: self.optitrack_status = 'Optitrack couldn\'t be started, make sure Motive is open!' client = optitrack.SimulatedClient() self.client = client # Create a source to buffer the motion tracking data from riglib import source self.motiondata = source.DataSource(optitrack.make(optitrack.System, self.client, self.optitrack_feature, 1)) # Save to the sink from riglib import sink sink_manager = sink.SinkManager.get_instance() sink_manager.register(self.motiondata) super().init() def run(self): ''' Code to execute immediately prior to the beginning of the task FSM executing, or after the FSM has finished running. See riglib.experiment.Experiment.run(). This 'run' method starts the motiondata source and stops it after the FSM has finished running ''' if not self.optitrack_status in ['recording', 'streaming']: import io self.terminated_in_error = True self.termination_err = io.StringIO() self.termination_err.write(self.optitrack_status) self.termination_err.seek(0) self.state = None super().run() else: self.motiondata.start() try: super().run() finally: print("Stopping optitrack") self.client.stop_recording() self.motiondata.stop() def _start_None(self): ''' Code to run before the 'None' state starts (i.e., the task stops) ''' #self.client.stop_recording() self.motiondata.stop() super()._start_None() def join(self): ''' See riglib.experiment.Experiment.join(). Re-join the motiondata source process before cleaning up the experiment thread ''' print("Joining optitrack datasource") self.motiondata.join() super().join() def cleanup(self, database, saveid, **kwargs): ''' Save the optitrack recorded file into the database ''' super_result = super().cleanup(database, saveid, **kwargs) print("Saving optitrack file to database...") try: database.save_data(self.filename, "optitrack", saveid, False, False) # Make sure you actually have an "optitrack" system added! except Exception as e: print(e) return False print("...done.") return super_result def _get_manual_position(self): ''' Overridden method to get input coordinates based on motion data''' # Get data from optitrack datasource data = self.motiondata.get() # List of (list of features) if len(data) == 0: # Data is not being streamed return recent = data[-self.smooth_features:] # How many recent coordinates to average averaged = np.nanmean(recent, axis=0) # List of averaged features if np.isnan(averaged).any(): # No usable coords return return averaged*100 # convert meters to centimeters
class ScreenTargetCapture(TargetCapture, Window): """Concrete implementation of TargetCapture task where targets are acquired by "holding" a cursor in an on-screen target""" background = (0, 0, 0, 1) cursor_color = (.5, 0, .5, 1) plant_type = traits.OptionsList(*plantlist, desc='', bmi3d_input_options=list(plantlist.keys())) starting_pos = (5, 0, 5) target_color = (1, 0, 0, .5) 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) limit2d = 1 sequence_generators = ['centerout_2D_discrete'] is_bmi_seed = True _target_color = RED # Runtime settable traits reward_time = traits.Float(.5, desc="Length of juice reward") target_radius = traits.Float(2, desc="Radius of targets in cm") 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') 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_options = list(plantlist.keys()) 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") def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.cursor_visible = True # Initialize the plant if not hasattr(self, 'plant'): 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) # Instantiate the targets instantiate_targets = kwargs.pop('instantiate_targets', True) if instantiate_targets: target1 = VirtualCircularTarget(target_radius=self.target_radius, target_color=self._target_color) target2 = VirtualCircularTarget(target_radius=self.target_radius, target_color=self._target_color) self.targets = [target1, target2] for target in self.targets: for model in target.graphics_models: self.add_model(model) # 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('target', 'f8', (3, )) self.add_dtype('target_index', 'i', (1, )) super().init() def _cycle(self): ''' Calls any update functions necessary and redraws screen. Runs 60x per second. ''' self.task_data['target'] = self.target_location.copy() self.task_data['target_index'] = self.target_index ## Run graphics commands to show/hide the plant if the visibility has changed if self.plant_type != 'CursorPlant': if self.plant_visible != self.plant_vis_prev: self.plant_vis_prev = self.plant_visible self.plant.set_visibility(self.plant_visible) # self.show_object(self.plant, show=self.plant_visible) self.move_effector() ## 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()._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() try: super().run() finally: self.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) #### 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.target_location) return d <= (self.target_radius - self.cursor_radius) def _test_leave_early(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.target_location) rad = self.target_radius - self.cursor_radius return d > rad #### STATE FUNCTIONS #### def _start_wait(self): super()._start_wait() # hide targets for target in self.targets: target.hide() def _start_target(self): super()._start_target() # move one of the two targets to the new target location target = self.targets[self.target_index % 2] target.move_to_position(self.target_location) target.cue_trial_start() def _start_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_hold(self): # change current target color to green self.targets[self.target_index % 2].cue_trial_end_success() def _start_hold_penalty(self): super()._start_hold_penalty() # hide targets for target in self.targets: target.hide() def _start_timeout_penalty(self): super()._start_timeout_penalty() # hide targets for target in self.targets: target.hide() def _start_targ_transition(self): #hide targets for target in self.targets: target.hide() def _start_reward(self): self.targets[self.target_index % 2].show() #### Generator functions #### @staticmethod def centerout_2D_discrete(nblocks=100, ntargets=8, boundaries=(-18, 18, -12, 12), distance=10): ''' 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. 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 = np.arange(0, 2 * np.pi, 2 * np.pi / ntargets) np.random.shuffle(temp) theta = theta + [temp] theta = np.hstack(theta) x = distance * np.cos(theta) y = np.zeros(len(theta)) z = distance * np.sin(theta) pairs = np.zeros([len(theta), 2, 3]) pairs[:, 1, :] = np.vstack([x, y, z]).T return pairs
class FactorBMIBase(BMIResetting): #Choose task to use trials from (no assist) #TODO make offline function to quickly assess optimal number of factors sequence_generators = [ 'centerout_2D_discrete', 'generate_catch_trials', 'all_shar_trials' ] input_type_list = [ 'shared', 'private', 'shared_scaled', 'private_scaled', 'all', 'all_scaled_by_shar', 'sc_shared+unsc_priv', 'sc_shared+sc_priv', 'main_shared', 'main_sc_shared', 'main_sc_private', 'main_sc_shar+unsc_priv', 'main_sc_shar+sc_priv', 'pca', 'split' ] #, 'priv_shar_concat'] input_type = traits.OptionsList(*input_type_list, bmi3d_input_options=input_type_list) def init(self): nU = self.decoder.n_units #Add datatypes for 1) trial type 2) input types: self.decoder.filt.FA_input_dict = {} self.input_types = [i + '_input' for i in self.input_type_list] + ['task_input'] for k in self.input_types: if k == 'split_input': nUn = self.decoder.trained_fa_dict['fa_main_shar_n_dim'] + nU #nUn = 2*nU elif k == 'task_input' and self.input_type == 'split': nUn = self.decoder.trained_fa_dict['fa_main_shar_n_dim'] + nU #nUn = 2*nU else: nUn = nU self.decoder.filt.FA_input_dict[k] = np.zeros((nUn, 1)) self.decoder.filt.FA_input_dict[k][:] = np.nan self.add_dtype(k, 'f8', (nUn, 1)) #Add datatype for 'trial-type': self.add_dtype('fa_input', np.str_, 16) super(FactorBMIBase, self).init() def init_decoder_state(self): try: fa_dict = self.decoder.trained_fa_dict self.decoder.filt.FA_kwargs = fa_dict except: raise Exception( 'Must run riglib.train.add_fa_dict_to_decoder and resave with dbq.save' ) #Check if isinstance of FAKalmanFilter or KalmanFilter from riglib.bmi.kfdecoder import KalmanFilter, FAKalmanFilter if isinstance(self.decoder.filt, KalmanFilter): self.decoder.filt.__class__ = FAKalmanFilter #Add FA elements to dict: print 'adding ', self.input_type, ' as input_type to FA decoder' import time time.sleep(2.) self.decoder.filt.FA_input = self.input_type super(FactorBMIBase, self).init_decoder_state() def _cycle(self): for k in self.input_types: try: self.task_data[k] = self.decoder.filt.FA_input_dict[k] except: print k, self.decoder.filt.FA_input_dict[ k].shape, self.task_data[k].shape self.task_data['fa_input'] = self.decoder.filt.FA_input super(FactorBMIBase, self)._cycle() def _parse_next_trial(self): if type(self.next_trial[1]) is str: self.targs = self.next_trial[0] else: self.targs = self.next_trial #print 'trial: ', self.decoder.filt.FA_input, self.targs ### FA param saving functions ###: def cleanup_hdf(self): ''' Re-open the HDF file and save any extra task data kept in RAM ''' super(FactorBMIBase, self).cleanup_hdf() try: self.write_FA_data_to_hdf_table(self.h5file.name, self.decoder.filt.FA_kwargs) print 'writing FA params to HDF file' except: print 'error in writing FA params to hdf file' import traceback traceback.print_exc() @staticmethod def write_FA_data_to_hdf_table(hdf_fname, FA_dict, ignore_none=False): import tables compfilt = tables.Filters(complevel=5, complib="zlib", shuffle=True) h5file = tables.openFile(hdf_fname, mode='a') fa_grp = h5file.createGroup(h5file.root, "fa_params", "Parameters for FA model used") for key in FA_dict: if isinstance(FA_dict[key], np.ndarray): h5file.createArray(fa_grp, key, FA_dict[key]) else: try: h5file.createArray(fa_grp, key, np.array([FA_dict[key]])) except: print 'cannot save: ', key, 'from FA in hdf file' h5file.close() @classmethod def generate_FA_matrices(self, training_task_entry, plot=False, hdf=None, dec=None, bin_spk=None): import utils.fa_decomp as pa if bin_spk is None: if training_task_entry is not None: from db import dbfunctions as dbfn te = dbfn.TaskEntry(training_task_entry) hdf = te.hdf dec = te.decoder bin_spk, targ_pos, targ_ix, z, zz = self.extract_trials_all( hdf, dec) #Zscore is in time x neurons zscore_X, mu = self.zscore_spks(bin_spk) # #Find optimal number of factors: LL, psv = pa.find_k_FA(zscore_X, iters=3, max_k=10, plot=False) #Np.nanmean: nan_ix = np.isnan(LL) samp = np.sum(nan_ix == False, axis=0) ll = np.nansum(LL, axis=0) LL_new = np.divide(ll, samp) num_factors = 1 + (np.argmax(LL_new)) print 'optimal LL factors: ', num_factors FA = skdecomp.FactorAnalysis(n_components=num_factors) #Samples x features: FA.fit(zscore_X) #FA matrices: U = np.mat(FA.components_).T i = np.diag_indices(U.shape[0]) Psi = np.mat(np.zeros((U.shape[0], U.shape[0]))) Psi[i] = FA.noise_variance_ A = U * U.T B = np.linalg.inv(U * U.T + Psi) mu_vect = np.array([mu[0, :]]).T #Size = N x 1 sharL = A * B #Calculate shared / priv scaling: bin_spk_tran = bin_spk.T mu_mat = np.tile(np.array([mu[0, :]]).T, (1, bin_spk_tran.shape[1])) demn = bin_spk_tran - mu_mat shared_bin_spk = (sharL * demn) priv_bin_spk = bin_spk_tran - mu_mat - shared_bin_spk #Scaling: eps = 1e-15 x_var = np.var(np.mat(bin_spk_tran), axis=1) + eps pr_var = np.var(priv_bin_spk, axis=1) + eps sh_var = np.var(shared_bin_spk, axis=1) + eps priv_scalar = np.sqrt(np.divide(x_var, pr_var)) shared_scalar = np.sqrt(np.divide(x_var, sh_var)) if plot: tmp = np.diag(U.T * U) plt.plot(np.arange(1, num_factors + 1), np.cumsum(tmp) / np.sum(tmp), '.-') plt.plot([0, num_factors + 1], [.9, .9], '-') #Get main shared space: u, s, v = np.linalg.svd(A) s_red = np.zeros_like(s) s_hd = np.zeros_like(s) ix = np.nonzero(np.cumsum(s**2) / float(np.sum(s**2)) > .90)[0] if len(ix) > 0: n_dim_main_shared = ix[0] + 1 else: n_dim_main_shared = len(s) if n_dim_main_shared < 2: n_dim_main_shared = 2 print "main shared: n_dim: ", n_dim_main_shared, np.cumsum(s) / float( np.sum(s)) s_red[:n_dim_main_shared] = s[:n_dim_main_shared] s_hd[n_dim_main_shared:] = s[n_dim_main_shared:] main_shared_A = u * np.diag(s_red) * v hd_shared_A = u * np.diag(s_hd) * v main_shared_B = np.linalg.inv(main_shared_A + hd_shared_A + Psi) uut_psi_inv = main_shared_B.copy() u_svd = u[:, :n_dim_main_shared] main_sharL = main_shared_A * main_shared_B main_shar = main_sharL * demn main_shar_var = np.var(main_shar, axis=1) + eps main_shar_scal = np.sqrt(np.divide(x_var, main_shar_var)) main_priv = demn - main_shar main_priv_var = np.var(main_priv, axis=1) + eps main_priv_scal = np.sqrt(np.divide(x_var, main_priv_var)) # #Get PCA decomposition: #LL, ax = pa.FA_all_targ_ALLms(hdf, iters=2, max_k=20, PCA_instead=True) #num_PCs = 1+(np.argmax(np.mean(LL, axis=0))) # Main PCA space: # Get cov matrix: cov_pca = np.cov(zscore_X.T) eig_val, eig_vec = np.linalg.eig(cov_pca) tot_var = sum(eig_val) cum_var_exp = np.cumsum( [i / tot_var for i in sorted(eig_val, reverse=True)]) n_PCs = np.nonzero(cum_var_exp > 0.9)[0][0] + 1 proj_mat = eig_vec[:, :n_PCs] proj_trans = np.mat(proj_mat) * np.mat(proj_mat.T) #PC matrices: return dict(fa_sharL=sharL, fa_mu=mu_vect, fa_shar_var_sc=shared_scalar, fa_priv_var_sc=priv_scalar, U=U, Psi=Psi, training_task_entry=training_task_entry, FA_iterated_power=FA.iterated_power, FA_score=FA.score(zscore_X), FA_LL=np.array(FA.loglike_), fa_main_shared=main_sharL, fa_main_shared_sc=main_shar_scal, fa_main_private_sc=main_priv_scal, fa_main_shar_n_dim=n_dim_main_shared, sing_vals=s, own_pc_trans=proj_trans, FA_model=FA, uut_psi_inv=uut_psi_inv, u_svd=u_svd) @classmethod def zscore_spks(self, proc_spks): '''Assumes a time x units matrix''' mu = np.tile(np.mean(proc_spks, axis=0), (proc_spks.shape[0], 1)) zscore_X = proc_spks - mu return zscore_X, mu @classmethod def extract_trials_all(self, hdf, dec, neural_bins=100, time_cutoff=None, hdf_ix=False): ''' Summary: method to extract all time points from trials Input param: hdf: task file input Input param: rew_ix: rows in the hdf file corresponding to reward times Input param: neural_bins: ms per bin Input param: time_cutoff: time in minutes, only extract trials before this time Input param: hdf_ix: bool, whether to return hdf row corresponding to time of decoder update (and hence end of spike bin) Output param: bin_spk -- binned spikes in time x units targ_i_all -- target location at each update targ_ix -- target index trial_ix -- trial number reach_time -- reach time for trial hdf_ix -- end bin in units of hdf rows ''' rew_ix = np.array([ t[1] for it, t in enumerate(hdf.root.task_msgs[:]) if t[0] == 'reward' ]) if time_cutoff is not None: it_cutoff = time_cutoff * 60 * 60 else: it_cutoff = len(hdf.root.task) #Get Go cue and go_ix = np.array([ hdf.root.task_msgs[it - 3][1] for it, t in enumerate(hdf.root.task_msgs[:]) if t[0] == 'reward' ]) go_ix = go_ix[go_ix < it_cutoff] rew_ix = rew_ix[go_ix < it_cutoff] targ_i_all = np.array([[-1, -1]]) trial_ix_all = np.array([-1]) reach_tm_all = np.array([-1]) hdf_ix_all = np.array([-1]) bin_spk = np.zeros((1, hdf.root.task[0]['spike_counts'].shape[0])) - 1 drives_neurons = dec.drives_neurons drives_neurons_ix0 = np.nonzero(drives_neurons)[0][0] update_bmi_ix = np.nonzero( np.diff( np.squeeze(hdf.root.task[:]['internal_decoder_state'] [:, drives_neurons_ix0, 0])))[0] + 1 for ig, (g, r) in enumerate(zip(go_ix, rew_ix)): spk_i = hdf.root.task[g:r]['spike_counts'][:, :, 0] #Sum spikes in neural_bins: bin_spk_i, nbins, hdf_ix_i = self._bin_spks( spk_i, g, r, update_bmi_ix) bin_spk = np.vstack((bin_spk, bin_spk_i)) targ_i_all = np.vstack( (targ_i_all, np.tile(hdf.root.task[g + 1]['target'][[0, 2]], (bin_spk_i.shape[0], 1)))) trial_ix_all = np.hstack( (trial_ix_all, np.zeros((bin_spk_i.shape[0])) + ig)) reach_tm_all = np.hstack( (reach_tm_all, np.zeros( (bin_spk_i.shape[0])) + ((r - g) * 1000. / 60.))) hdf_ix_all = np.hstack((hdf_ix_all, hdf_ix_i)) targ_ix = self._get_target_ix(targ_i_all[1:, :]) if hdf_ix: return bin_spk[1:, :], targ_i_all[1:, :], targ_ix, trial_ix_all[ 1:], reach_tm_all[1:], hdf_ix_all[1:] else: return bin_spk[1:, :], targ_i_all[ 1:, :], targ_ix, trial_ix_all[1:], reach_tm_all[1:] @classmethod def _get_target_ix(self, targ_pos): #Target Index: b = np.ascontiguousarray(targ_pos).view( np.dtype((np.void, targ_pos.dtype.itemsize * targ_pos.shape[1]))) _, idx = np.unique(b, return_index=True) unique_targ = targ_pos[idx, :] #Order by theta: theta = np.arctan2(unique_targ[:, 1], unique_targ[:, 0]) thet_i = np.argsort(theta) unique_targ = unique_targ[thet_i, :] targ_ix = np.zeros((targ_pos.shape[0]), ) for ig, (x, y) in enumerate(targ_pos): targ_ix[ig] = np.nonzero( np.sum(targ_pos[ig, :] == unique_targ, axis=1) == 2)[0] return targ_ix @classmethod def _bin_spks(self, spk_i, g_ix, r_ix, update_bmi_ix): #Need to use 'update_bmi_ix' from ReDecoder to get bin edges correctly: trial_inds = np.arange(g_ix, r_ix + 1) end_bin = np.array([ (j, i) for j, i in enumerate(trial_inds) if np.logical_and(i in update_bmi_ix, i >= (g_ix + 5)) ]) nbins = len(end_bin) bin_spk_i = np.zeros((nbins, spk_i.shape[1])) hdf_ix_i = [] for ib, (i_ix, hdf_ix) in enumerate(end_bin): #Inclusive of EndBin bin_spk_i[ib, :] = np.sum(spk_i[i_ix - 5:i_ix + 1, :], axis=0) hdf_ix_i.append(hdf_ix) return bin_spk_i, nbins, np.array(hdf_ix_i) @staticmethod def generate_catch_trials(nblocks=5, ntargets=8, distance=10, perc_shar=10, perc_priv=10): ''' Generates a sequence of 2D (x and z) target pairs with the first target always at the origin and a second field indicating the extractor type (full, shared, priv) 1 shared / 1 private for nblocks: multiples of 80 perc_shar, perc_priv: multiples of 10, please ''' assert (not np.mod(perc_shar, 10)) and (not np.mod(perc_priv, 10)) #Make blocks of 80 trials: theta = [] for i in range(10): temp = np.arange(0, 2 * np.pi, 2 * np.pi / ntargets) np.random.shuffle(temp) theta = theta + [temp] theta = np.hstack(theta) #Each target has correct % of private and correct % of shared targets trial_type = np.empty(len(theta), dtype='S10') for i in temp: targ_ix = np.nonzero(theta == i)[0] trial_ix = np.arange(len(targ_ix)) tmp_trial = np.array(['all'] * len(targ_ix), dtype='S10') n_trial_shar = np.floor(perc_shar / 100. * float(len(targ_ix))) n_trial_priv = np.floor(perc_priv / 100. * float(len(targ_ix))) tmp_trial[:int(n_trial_shar)] = ['shared'] tmp_trial[int(n_trial_shar):int(n_trial_shar + n_trial_priv)] = ['private'] np.random.shuffle(tmp_trial) trial_type[targ_ix] = tmp_trial #Make Target set: x = distance * np.cos(theta) y = np.zeros(len(theta)) z = distance * np.sin(theta) pairs = np.zeros([len(theta), 2, 3]) pairs[:, 1, :] = np.vstack([x, y, z]).T Pairs = np.tile(pairs, [nblocks, 1, 1]) Trial_type = np.tile(trial_type, [nblocks]) #Will yield a tuple where target location is in next_trial[0], trial_type is in next_trial[1] return zip(Pairs, Trial_type) @staticmethod def all_shar_trials(nblocks=5, ntargets=8, distance=10): ''' Generates a sequence of 2D (x and z) target pairs with the first target always at the origin and a second field indicating the extractor type (always shared) ''' #Make blocks of 80 trials: theta = [] for i in range(10): temp = np.arange(0, 2 * np.pi, 2 * np.pi / ntargets) np.random.shuffle(temp) theta = theta + [temp] theta = np.hstack(theta) #Each target has correct % of private and correct % of shared targets trial_type = np.empty(len(theta), dtype='S10') trial_type[:] = 'shared' #Make Target set: x = distance * np.cos(theta) y = np.zeros(len(theta)) z = distance * np.sin(theta) pairs = np.zeros([len(theta), 2, 3]) pairs[:, 1, :] = np.vstack([x, y, z]).T Pairs = np.tile(pairs, [nblocks, 1, 1]) Trial_type = np.tile(trial_type, [nblocks]) #Will yield a tuple where target location is in next_trial[0], trial_type is in next_trial[1] return zip(Pairs, Trial_type)
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 ManualControlMixin(traits.HasTraits): '''Target capture task where the subject operates a joystick to control a cursor. Targets are captured by having the cursor dwell in the screen target for the allotted time''' # Settable Traits wait_time = traits.Float(2., desc="Time between successful trials") velocity_control = traits.Bool(False, desc="Position or velocity control") random_rewards = traits.Bool(False, desc="Add randomness to reward") rotation = traits.OptionsList(*rotations, desc="Control rotation matrix", bmi3d_input_options=list(rotations.keys())) scale = traits.Float(1.0, desc="Control scale factor") offset = traits.Array(value=[0,0,0], desc="Control offset") is_bmi_seed = True def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.current_pt=np.zeros([3]) #keep track of current pt self.last_pt=np.zeros([3]) #keep track of last pt to calc. velocity self.no_data_count = 0 self.reportstats['Input quality'] = "100 %" if self.random_rewards: self.reward_time_base = self.reward_time def init(self): self.add_dtype('manual_input', 'f8', (3,)) super().init() def _test_start_trial(self, ts): return ts > self.wait_time and not self.pause def _test_trial_complete(self, ts): if self.target_index==self.chain_length-1 : if self.random_rewards: if not self.rand_reward_set_flag: #reward time has not been set for this iteration self.reward_time = np.max([2*(np.random.rand()-0.5) + self.reward_time_base, self.reward_time_base/2]) #set randomly with min of base / 2 self.rand_reward_set_flag =1 #print self.reward_time, self.rand_reward_set_flag return self.target_index==self.chain_length-1 def _test_reward_end(self, ts): #When finished reward, reset flag. if self.random_rewards: if ts > self.reward_time: self.rand_reward_set_flag = 0 #print self.reward_time, self.rand_reward_set_flag, ts return ts > self.reward_time def _transform_coords(self, coords): ''' Returns transformed coordinates based on rotation, offset, and scale traits ''' offset = np.array( [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [self.offset[0], self.offset[1], self.offset[2], 1]] ) scale = np.array( [[self.scale, 0, 0, 0], [0, self.scale, 0, 0], [0, 0, self.scale, 0], [0, 0, 0, 1]] ) old = np.concatenate((np.reshape(coords, -1), [1])) new = np.linalg.multi_dot((old, offset, scale, rotations[self.rotation])) return new[0:3] def _get_manual_position(self): ''' Fetches joystick position ''' if not hasattr(self, 'joystick'): return pt = self.joystick.get() if len(pt) == 0: return pt = pt[-1] # Use only the latest coordinate if len(pt) == 2: pt = np.concatenate((np.reshape(pt, -1), [0])) return [pt] def move_effector(self): ''' Sets the 3D coordinates of the cursor. For manual control, uses motiontracker / joystick / mouse data. If no data available, returns None ''' # Get raw input and save it as task data raw_coords = self._get_manual_position() # array of [3x1] arrays if raw_coords is None or len(raw_coords) < 1: self.no_data_count += 1 self.update_report_stats() self.task_data['manual_input'] = np.empty((3,)) return self.task_data['manual_input'] = raw_coords.copy() # Transform coordinates coords = self._transform_coords(raw_coords) if self.limit2d: coords[1] = 0 # Set cursor position if not self.velocity_control: self.current_pt = coords else: epsilon = 2*(10**-2) # Define epsilon to stabilize cursor movement if sum((coords)**2) > epsilon: # Add the velocity (units/s) to the position (units) self.current_pt = coords / self.fps + self.last_pt else: self.current_pt = self.last_pt self.plant.set_endpoint_pos(self.current_pt) self.last_pt = self.current_pt.copy() def update_report_stats(self): super().update_report_stats() quality = 1 - self.no_data_count / max(1, self.cycle_count) self.reportstats['Input quality'] = "{} %".format(int(100*quality)) @classmethod def get_desc(cls, params, log_summary): duration = round(log_summary['runtime'] / 60, 1) return "{}/{} succesful trials in {} min".format( log_summary['n_success_trials'], log_summary['n_trials'], duration)
class BMIControlMulti(BMILoop, LinearlyDecreasingAssist, manualcontrolmultitasks.ManualControlMulti): ''' Target capture task with cursor position controlled by BMI output. Cursor movement can be assisted toward target by setting assist_level > 0. ''' background = (.5, .5, .5, 1) # Set the screen background color to grey reset = traits.Int( 0, desc='reset the decoder state to the starting configuration') ordered_traits = [ 'session_length', 'assist_level', 'assist_level_time', 'reward_time', 'timeout_time', 'timeout_penalty_time' ] exclude_parent_traits = ['marker_count', 'marker_num', 'goal_cache_block'] static_states = [] # states in which the decoder is not run hidden_traits = [ 'arm_hide_rate', 'arm_visible', 'hold_penalty_time', 'rand_start', 'reset', 'target_radius', 'window_size' ] is_bmi_seed = False cursor_color_adjust = traits.OptionsList( *target_colors.keys(), bmi3d_input_options=target_colors.keys()) def __init__(self, *args, **kwargs): super(BMIControlMulti, self).__init__(*args, **kwargs) def init(self, *args, **kwargs): sph = self.plant.graphics_models[0] sph.color = target_colors[self.cursor_color_adjust] sph.radius = self.cursor_radius self.plant.cursor_radius = self.cursor_radius self.plant.cursor.radius = self.cursor_radius super(BMIControlMulti, self).init(*args, **kwargs) def move_effector(self, *args, **kwargs): pass def create_assister(self): # Create the appropriate type of assister object start_level, end_level = self.assist_level kwargs = dict(decoder_binlen=self.decoder.binlen, target_radius=self.target_radius) if hasattr(self, 'assist_speed'): kwargs['assist_speed'] = self.assist_speed from db import namelist if self.decoder.ssm == namelist.endpt_2D_state_space and isinstance( self.decoder, ppfdecoder.PPFDecoder): self.assister = OFCEndpointAssister() elif self.decoder.ssm == namelist.endpt_2D_state_space: self.assister = SimpleEndpointAssister(**kwargs) elif (self.decoder.ssm == namelist.tentacle_2D_state_space) or ( self.decoder.ssm == namelist.joint_2D_state_space): # kin_chain = self.plant.kin_chain # A, B, W = self.decoder.ssm.get_ssm_matrices(update_rate=self.decoder.binlen) # Q = np.mat(np.diag(np.hstack([kin_chain.link_lengths, np.zeros_like(kin_chain.link_lengths), 0]))) # R = 10000*np.mat(np.eye(B.shape[1])) # fb_ctrl = LQRController(A, B, Q, R) # self.assister = FeedbackControllerAssist(fb_ctrl, style='additive') self.assister = TentacleAssist(ssm=self.decoder.ssm, kin_chain=self.plant.kin_chain, update_rate=self.decoder.binlen) else: raise NotImplementedError( "Cannot assist for this type of statespace: %r" % self.decoder.ssm) print self.assister def create_goal_calculator(self): from db import namelist if self.decoder.ssm == namelist.endpt_2D_state_space: self.goal_calculator = goal_calculators.ZeroVelocityGoal( self.decoder.ssm) elif self.decoder.ssm == namelist.joint_2D_state_space: self.goal_calculator = goal_calculators.PlanarMultiLinkJointGoal( self.decoder.ssm, self.plant.base_loc, self.plant.kin_chain, multiproc=False, init_resp=None) elif self.decoder.ssm == namelist.tentacle_2D_state_space: shoulder_anchor = self.plant.base_loc chain = self.plant.kin_chain q_start = self.plant.get_intrinsic_coordinates() x_init = np.hstack([q_start, np.zeros_like(q_start), 1]) x_init = np.mat(x_init).reshape(-1, 1) cached = True if cached: goal_calc_class = goal_calculators.PlanarMultiLinkJointGoalCached multiproc = False else: goal_calc_class = goal_calculators.PlanarMultiLinkJointGoal multiproc = True self.goal_calculator = goal_calc_class( namelist.tentacle_2D_state_space, shoulder_anchor, chain, multiproc=multiproc, init_resp=x_init) else: raise ValueError("Unrecognized decoder state space!") def get_target_BMI_state(self, *args): ''' Run the goal calculator to determine the target state of the task ''' if isinstance(self.goal_calculator, goal_calculators.PlanarMultiLinkJointGoalCached): task_eps = np.inf else: task_eps = 0.5 ik_eps = task_eps / 10 data, solution_updated = self.goal_calculator( self.target_location, verbose=False, n_particles=500, eps=ik_eps, n_iter=10, q_start=self.plant.get_intrinsic_coordinates()) target_state, error = data if isinstance(self.goal_calculator, goal_calculators.PlanarMultiLinkJointGoal ) and error > task_eps and solution_updated: self.goal_calculator.reset() return np.array(target_state).reshape(-1, 1) def _end_timeout_penalty(self): if self.reset: self.decoder.filt.state.mean = self.init_decoder_mean self.hdf.sendMsg("reset") def move_effector(self): pass
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